Preserve Guardian transcript boundaries in sampling input (#38555)

## What changed

- Represent Guardian sampling input as an ordered list of text entries.
- Send each transcript entry and approval-request delimiter as a separate input
  content item while preserving the existing transcript byte limit.

## Testing

- Update transcript and sampler tests to verify entry boundaries, multi-item
  request serialization, and bounded transcript behavior.

GitOrigin-RevId: 433e2e6ca18ce9288025715f5956295bfe852e98
This commit is contained in:
felixxia-oai authored and copyberry committed 2026-08-14 11:02:01 +00:00
1 parent a44f266a4a
commit db5b4a4cfb
6 files changed
+101 -63

No files matched your search

+10 -8
View File
@@ -127,14 +127,16 @@ impl ToolLifecycleContributor for GuardianV2Extension {
return; return;
} }
}; };
let classification_input = format!( let mut classification_input = vec![">>> TRANSCRIPT START\n".to_owned()];
">>> TRANSCRIPT START\n{transcript}>>> TRANSCRIPT END\n\n\ classification_input.extend(transcript);
The Codex agent has requested the following action:\n\ classification_input.extend([
>>> APPROVAL REQUEST START\n\ ">>> TRANSCRIPT END\n\n".to_owned(),
Planned action JSON:\n\ "The Codex agent has requested the following action:\n".to_owned(),
{planned_action}\n\ ">>> APPROVAL REQUEST START\n".to_owned(),
>>> APPROVAL REQUEST END\n" "Planned action JSON:\n".to_owned(),
); format!("{planned_action}\n"),
">>> APPROVAL REQUEST END\n".to_owned(),
]);
let result: Result<(), String> = async { let result: Result<(), String> = async {
let output = sampler let output = sampler
.sample(LunaSamplingRequest { .sample(LunaSamplingRequest {
+21 -18
View File
@@ -162,24 +162,27 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result<
json!({"type": "number", "minimum": 0.0, "maximum": 1.0}) json!({"type": "number", "minimum": 0.0, "maximum": 1.0})
); );
assert_eq!( assert_eq!(
request["input"][2]["content"][0]["text"], request["input"][2]["content"],
concat!( json!([
">>> TRANSCRIPT START\n", {"type": "input_text", "text": ">>> TRANSCRIPT START\n"},
"[1] user: Inspect the repository guidelines.\n", {"type": "input_text", "text": "[1] user: Inspect the repository guidelines.\n"},
"[2] reasoning: Find the repository documentation.\n", {"type": "input_text", "text": "[2] reasoning: Find the repository documentation.\n"},
"[3] tool list_dir call: {\"path\":\".\"}\n", {"type": "input_text", "text": "[3] tool list_dir call: {\"path\":\".\"}\n"},
"[4] tool list_dir result: README.md\n", {"type": "input_text", "text": "[4] tool list_dir result: README.md\n"},
"[5] tool read_file call: {\"path\":\"README.md\"}\n", {"type": "input_text", "text": "[5] tool read_file call: {\"path\":\"README.md\"}\n"},
">>> TRANSCRIPT END\n\n", {"type": "input_text", "text": ">>> TRANSCRIPT END\n\n"},
"The Codex agent has requested the following action:\n", {
">>> APPROVAL REQUEST START\n", "type": "input_text",
"Planned action JSON:\n", "text": "The Codex agent has requested the following action:\n"
"{\n", },
" \"path\": \"README.md\",\n", {"type": "input_text", "text": ">>> APPROVAL REQUEST START\n"},
" \"tool\": \"read_file\"\n", {"type": "input_text", "text": "Planned action JSON:\n"},
"}\n", {
">>> APPROVAL REQUEST END\n", "type": "input_text",
) "text": "{\n \"path\": \"README.md\",\n \"tool\": \"read_file\"\n}\n"
},
{"type": "input_text", "text": ">>> APPROVAL REQUEST END\n"},
])
); );
let score = tokio::time::timeout(Duration::from_secs(5), async { let score = tokio::time::timeout(Duration::from_secs(5), async {
loop { loop {
+7 -5
View File
@@ -66,8 +66,8 @@ pub struct LunaSamplerConfig {
pub struct LunaSamplingRequest { pub struct LunaSamplingRequest {
/// Trusted instructions describing the requested classification. /// Trusted instructions describing the requested classification.
pub instructions: String, pub instructions: String,
/// Untrusted input that the model should classify. /// Ordered untrusted input entries that the model should classify.
pub input: String, pub input: Vec<String>,
/// Strict JSON schema constraining the model response. /// Strict JSON schema constraining the model response.
pub output_schema: Value, pub output_schema: Value,
/// Reasoning budget explicitly selected for this request. /// Reasoning budget explicitly selected for this request.
@@ -284,9 +284,11 @@ impl LunaSampler {
ResponseItem::Message { ResponseItem::Message {
id: None, id: None,
role: "user".to_owned(), role: "user".to_owned(),
content: vec![ContentItem::InputText { content: request
text: request.input, .input
}], .into_iter()
.map(|text| ContentItem::InputText { text })
.collect(),
phase: None, phase: None,
internal_chat_message_metadata_passthrough: None, internal_chat_message_metadata_passthrough: None,
}, },
+14 -4
View File
@@ -69,7 +69,7 @@ fn sampler_config(base_url: String) -> LunaSamplerConfig {
fn sample_request(turn_id: &str) -> LunaSamplingRequest { fn sample_request(turn_id: &str) -> LunaSamplingRequest {
LunaSamplingRequest { LunaSamplingRequest {
instructions: "Return a risk score.".to_owned(), instructions: "Return a risk score.".to_owned(),
input: "The user requested a README summary.".to_owned(), input: vec!["The user requested a README summary.".to_owned()],
output_schema: json!({ output_schema: json!({
"type": "object", "type": "object",
"properties": { "score": { "type": "number" } }, "properties": { "score": { "type": "number" } },
@@ -157,7 +157,10 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ
let first = sampler let first = sampler
.sample(LunaSamplingRequest { .sample(LunaSamplingRequest {
instructions: "Return a risk score.".to_owned(), instructions: "Return a risk score.".to_owned(),
input: "The user requested a README summary.".to_owned(), input: vec![
"The user requested a README summary.".to_owned(),
"The assistant inspected README.md.".to_owned(),
],
output_schema: schema.clone(), output_schema: schema.clone(),
reasoning_effort: ReasoningEffort::None, reasoning_effort: ReasoningEffort::None,
turn_id: "turn-1".to_owned(), turn_id: "turn-1".to_owned(),
@@ -178,7 +181,7 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ
let second = sampler let second = sampler
.sample(LunaSamplingRequest { .sample(LunaSamplingRequest {
instructions: "Return a risk score.".to_owned(), instructions: "Return a risk score.".to_owned(),
input: "The user requested a source review.".to_owned(), input: vec!["The user requested a source review.".to_owned()],
output_schema: schema, output_schema: schema,
reasoning_effort: ReasoningEffort::Medium, reasoning_effort: ReasoningEffort::Medium,
turn_id: "turn-2".to_owned(), turn_id: "turn-2".to_owned(),
@@ -189,6 +192,13 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_structured_requ
assert_eq!(second, r#"{"score":0.75}"#); assert_eq!(second, r#"{"score":0.75}"#);
let requests = server.single_connection(); let requests = server.single_connection();
assert_eq!(requests.len(), 2); assert_eq!(requests.len(), 2);
assert_eq!(
requests[0].body_json()["input"][2]["content"],
json!([
{"type": "input_text", "text": "The user requested a README summary."},
{"type": "input_text", "text": "The assistant inspected README.md."},
])
);
for (index, request) in requests.iter().enumerate() { for (index, request) in requests.iter().enumerate() {
let request = request.body_json(); let request = request.body_json();
assert_eq!(request["type"], "response.create"); assert_eq!(request["type"], "response.create");
@@ -254,7 +264,7 @@ async fn sampler_returns_complete_json_before_terminal_response_events() -> Resu
Duration::from_secs(2), Duration::from_secs(2),
sampler.sample(LunaSamplingRequest { sampler.sample(LunaSamplingRequest {
instructions: "Return a risk score.".to_owned(), instructions: "Return a risk score.".to_owned(),
input: "The user requested a README summary.".to_owned(), input: vec!["The user requested a README summary.".to_owned()],
output_schema: json!({ output_schema: json!({
"type": "object", "type": "object",
"properties": { "score": { "type": "number" } }, "properties": { "score": { "type": "number" } },
+26 -9
View File
@@ -1,4 +1,5 @@
use std::collections::HashMap; use std::collections::HashMap;
use std::collections::VecDeque;
use codex_extension_api::ResponseItem; use codex_extension_api::ResponseItem;
use codex_protocol::models::ContentItem; use codex_protocol::models::ContentItem;
@@ -36,8 +37,12 @@ impl Default for TranscriptConfig {
} }
impl TranscriptConfig { impl TranscriptConfig {
pub(crate) fn build<'a>(&self, items: impl IntoIterator<Item = &'a ResponseItem>) -> String { pub(crate) fn build<'a>(
let mut transcript = String::new(); &self,
items: impl IntoIterator<Item = &'a ResponseItem>,
) -> Vec<String> {
let mut transcript = VecDeque::new();
let mut transcript_bytes = 0;
let mut tool_names_by_call_id = HashMap::new(); let mut tool_names_by_call_id = HashMap::new();
let mut entry_number = 0; let mut entry_number = 0;
@@ -173,18 +178,30 @@ impl TranscriptConfig {
continue; continue;
} }
entry_number += 1; entry_number += 1;
transcript.push_str(&format!("[{entry_number}] {role}: {text}\n")); let entry = format!("[{entry_number}] {role}: {text}\n");
transcript_bytes += entry.len();
transcript.push_back(entry);
if transcript.len() > MAX_TRANSCRIPT_BYTES { while transcript_bytes > MAX_TRANSCRIPT_BYTES {
let mut first_retained_byte = transcript.len() - MAX_TRANSCRIPT_BYTES; let Some(first_entry) = transcript.front_mut() else {
while !transcript.is_char_boundary(first_retained_byte) { break;
first_retained_byte += 1; };
let bytes_to_remove = transcript_bytes - MAX_TRANSCRIPT_BYTES;
if first_entry.len() <= bytes_to_remove {
transcript_bytes -= first_entry.len();
transcript.pop_front();
} else {
let mut first_retained_byte = bytes_to_remove;
while !first_entry.is_char_boundary(first_retained_byte) {
first_retained_byte += 1;
}
first_entry.drain(..first_retained_byte);
transcript_bytes -= first_retained_byte;
} }
transcript.drain(..first_retained_byte);
} }
} }
transcript transcript.into()
} }
} }
@@ -62,13 +62,12 @@ fn transcript_keeps_conversation_and_configured_sources() {
let transcript = TranscriptConfig::default().build(&items); let transcript = TranscriptConfig::default().build(&items);
assert_eq!( assert_eq!(
transcript, transcript,
concat!( vec![
"[1] user: Inspect the workspace.\n", "[1] user: Inspect the workspace.\n",
"[2] tool exec_command call: {}\n", "[2] tool exec_command call: {}\n",
"[3] tool exec_command result: Workspace inspected.\n", "[3] tool exec_command result: Workspace inspected.\n",
"[4] reasoning: Review the current files.\n", "[4] reasoning: Review the current files.\nPlaintext reasoning.\n",
"Plaintext reasoning.\n", ]
)
); );
let output_and_reasoning = TranscriptConfig { let output_and_reasoning = TranscriptConfig {
@@ -78,12 +77,11 @@ fn transcript_keeps_conversation_and_configured_sources() {
let transcript = output_and_reasoning.build(&items); let transcript = output_and_reasoning.build(&items);
assert_eq!( assert_eq!(
transcript, transcript,
concat!( vec![
"[1] user: Inspect the workspace.\n", "[1] user: Inspect the workspace.\n",
"[2] tool exec_command result: Workspace inspected.\n", "[2] tool exec_command result: Workspace inspected.\n",
"[3] reasoning: Review the current files.\n", "[3] reasoning: Review the current files.\nPlaintext reasoning.\n",
"Plaintext reasoning.\n", ]
)
); );
let calls_only = TranscriptConfig { let calls_only = TranscriptConfig {
@@ -93,10 +91,10 @@ fn transcript_keeps_conversation_and_configured_sources() {
let transcript = calls_only.build(&items); let transcript = calls_only.build(&items);
assert_eq!( assert_eq!(
transcript, transcript,
concat!( vec![
"[1] user: Inspect the workspace.\n", "[1] user: Inspect the workspace.\n",
"[2] tool exec_command call: {}\n", "[2] tool exec_command call: {}\n",
) ]
); );
} }
@@ -125,8 +123,12 @@ fn transcript_retains_the_most_recent_bounded_content() {
let transcript = TranscriptConfig::default().build(&items); let transcript = TranscriptConfig::default().build(&items);
assert!(transcript.len() <= MAX_TRANSCRIPT_BYTES); assert!(transcript.iter().map(String::len).sum::<usize>() <= MAX_TRANSCRIPT_BYTES);
assert!(transcript.contains("latest response")); assert!(
transcript
.iter()
.any(|entry| entry.contains("latest response"))
);
} }
#[test] #[test]
@@ -154,7 +156,10 @@ fn transcript_keeps_only_manual_approval_developer_messages() {
]; ];
let transcript = TranscriptConfig::default().build(&items); let transcript = TranscriptConfig::default().build(&items);
assert_eq!(transcript, format!("[1] developer: {approval_text}\n")); assert_eq!(
transcript,
vec![format!("[1] developer: {approval_text}\n")]
);
} }
#[test] #[test]
@@ -224,10 +229,10 @@ fn transcript_omits_media_payloads_and_keeps_readable_content() {
let transcript = TranscriptConfig::default().build(&items); let transcript = TranscriptConfig::default().build(&items);
assert_eq!( assert_eq!(
transcript, transcript,
concat!( vec![
"[1] user: Review this screenshot.\n", "[1] user: Review this screenshot.\n",
"[2] tool result: Screenshot captured.\n", "[2] tool result: Screenshot captured.\n",
) ]
); );
} }
@@ -280,11 +285,10 @@ fn transcript_omits_encrypted_messages_arguments_and_tool_outputs() {
let transcript = TranscriptConfig::default().build(&items); let transcript = TranscriptConfig::default().build(&items);
assert_eq!( assert_eq!(
transcript, transcript,
concat!( vec![
"[1] assistant: Agent message from worker:\n", "[1] assistant: Agent message from worker:\nThe workspace is ready.\n",
"The workspace is ready.\n",
"[2] tool exec_command call: {}\n", "[2] tool exec_command call: {}\n",
"[3] tool exec_command result: Command completed.\n", "[3] tool exec_command result: Command completed.\n",
) ]
); );
} }