diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs index 42c8d67a24..a531225b47 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs @@ -35,6 +35,7 @@ use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ImageDetail; use codex_protocol::models::InternalChatMessageMetadataPassthrough; +use codex_protocol::models::MessagePhase; use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::openai_models::GuardianV2ModelConfig; use codex_protocol::openai_models::GuardianV2TranscriptModelConfig; @@ -2101,6 +2102,74 @@ async fn contributor_bounds_configured_policy_in_luna_developer_instructions() - Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn contributor_preserves_final_assistant_messages_after_tool_eviction() -> Result<()> { + skip_if_no_network!(Ok(())); + + let mut history = vec![ + responses::user_message_item("Find a flight to New York."), + ResponseItem::Message { + id: None, + role: "assistant".to_owned(), + content: vec![ContentItem::OutputText { + text: "I found a $450 flight. Should I book it?".to_owned(), + }], + phase: Some(MessagePhase::FinalAnswer), + internal_chat_message_metadata_passthrough: None, + }, + responses::user_message_item("Yes."), + ResponseItem::Message { + id: None, + role: "assistant".to_owned(), + content: vec![ContentItem::OutputText { + text: "Searching airline websites.".to_owned(), + }], + phase: Some(MessagePhase::Commentary), + internal_chat_message_metadata_passthrough: None, + }, + ]; + history.extend((0..6).map(|index| ResponseItem::FunctionCall { + id: None, + name: "exec_command".to_owned(), + namespace: None, + arguments: format!("booking step {index}"), + encrypted_function_args: None, + call_id: format!("call-{index}"), + internal_chat_message_metadata_passthrough: None, + })); + let configuration = "[features.guardianv2]\nenabled = true\n\n[features.guardianv2.transcript]\nmax_recent_non_user_entries = 4\n"; + + let (request, _test, _registry) = sample_configured_conversation_history( + history, + r#"{"path":"README.md"}"#, + Some(TEST_GUARDIAN_POLICY), + configuration, + /*model_defaults*/ None, + ) + .await?; + let entries = request["input"][2]["content"] + .as_array() + .expect("Luna request should contain separate transcript text items") + .iter() + .filter_map(|entry| entry["text"].as_str()) + .filter(|entry| entry.starts_with('[')) + .collect::>(); + + assert_eq!( + entries, + vec![ + "[1] user: Find a flight to New York.\n", + "[2] assistant: I found a $450 flight. Should I book it?\n", + "[3] user: Yes.\n", + "[8] tool exec_command call: booking step 3\n", + "[9] tool exec_command call: booking step 4\n", + "[10] tool exec_command call: booking step 5\n", + ] + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn contributor_sends_compacted_conversation_history_to_luna() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs b/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs index f51c373ee0..eb89ec23db 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs @@ -6,9 +6,11 @@ pub(crate) use codex_features::GuardianV2TranscriptSource as TranscriptSource; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::ImageDetail; +use codex_protocol::models::MessagePhase; use codex_protocol::models::ReasoningItemContent; use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::models::plaintext_agent_message_content; +use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::TruncationPolicy; use super::truncation::TruncationObservation; @@ -18,6 +20,7 @@ pub(crate) const MAX_TOOL_ENTRY_TOKENS: usize = 1_000; pub(crate) const MAX_MESSAGE_TRANSCRIPT_TOKENS: usize = 10_000; pub(crate) const MAX_TOOL_TRANSCRIPT_TOKENS: usize = 10_000; pub(crate) const MAX_RECENT_NON_USER_ENTRIES: usize = 40; +const MIN_RECENT_TOOL_ENTRIES: usize = 5; const MAX_TRANSCRIPT_IMAGES: usize = 4; const MAX_TRANSCRIPT_IMAGE_BYTES: usize = 8 * 1024 * 1024; const MANUAL_APPROVAL_DEVELOPER_PREFIX: &str = @@ -26,6 +29,7 @@ const MANUAL_APPROVAL_DEVELOPER_PREFIX: &str = #[derive(Clone, Copy, Debug, Eq, PartialEq)] enum TranscriptEntryKind { User, + ProtectedMessage, Message, Tool, } @@ -57,7 +61,14 @@ fn retained_tokens( .iter() .filter_map(|&index| { let entry = &entries[index]; - (entry.kind == kind).then_some(entry.tokens) + let same_token_pool = match kind { + TranscriptEntryKind::ProtectedMessage | TranscriptEntryKind::Message => matches!( + entry.kind, + TranscriptEntryKind::ProtectedMessage | TranscriptEntryKind::Message + ), + TranscriptEntryKind::User | TranscriptEntryKind::Tool => entry.kind == kind, + }; + same_token_pool.then_some(entry.tokens) }) .sum() } @@ -347,16 +358,29 @@ impl TranscriptConfig { continue; } - let kind = match role.as_str() { - "user" => TranscriptEntryKind::User, - role if role.starts_with("tool ") => TranscriptEntryKind::Tool, + let kind = match (role.as_str(), item) { + ("user", _) => TranscriptEntryKind::User, + ("developer", ResponseItem::Message { .. }) => { + TranscriptEntryKind::ProtectedMessage + } + ( + "assistant", + ResponseItem::Message { + phase: None | Some(MessagePhase::FinalAnswer), + content, + .. + }, + ) if !InterAgentCommunication::is_message_content(content) => { + TranscriptEntryKind::ProtectedMessage + } + (role, _) if role.starts_with("tool ") => TranscriptEntryKind::Tool, _ => TranscriptEntryKind::Message, }; let token_cap = match kind { TranscriptEntryKind::Tool => self.max_tool_entry_tokens, - TranscriptEntryKind::User | TranscriptEntryKind::Message => { - self.max_message_entry_tokens - } + TranscriptEntryKind::User + | TranscriptEntryKind::ProtectedMessage + | TranscriptEntryKind::Message => self.max_message_entry_tokens, }; let original_bytes = text.len(); let text = truncate_entry(&text, token_cap); @@ -409,11 +433,11 @@ impl TranscriptConfig { let available_message_tokens = self .max_message_transcript_tokens .saturating_sub(message_tokens); - let mut retained_non_user = VecDeque::new(); + let mut retained_non_user: VecDeque = VecDeque::new(); - // Replay the transcript as an append-only bounded buffer. When a new entry would overflow - // a budget, remove half of the existing entries in that pool before appending it. This - // keeps the retained prefix stable between overflow points without storing sampler state. + // Replay non-user entries as an append-only bounded buffer. Evict ordinary evidence + // before protected conversation messages, removing half of whichever pool overflows so + // the retained prefix stays stable between overflow points without storing sampler state. for (index, entry) in entries.iter().enumerate() { if entry.kind == TranscriptEntryKind::User { continue; @@ -421,45 +445,139 @@ impl TranscriptConfig { let token_budget = match entry.kind { TranscriptEntryKind::Tool => self.max_tool_transcript_tokens, - TranscriptEntryKind::Message => available_message_tokens, + TranscriptEntryKind::ProtectedMessage | TranscriptEntryKind::Message => { + available_message_tokens + } TranscriptEntryKind::User => unreachable!("user entries were selected separately"), }; - // An entry that cannot fit on its own must not evict retained evidence. - if entry.tokens > token_budget || self.max_recent_non_user_entries == 0 { + let protected_tokens = if entry.kind == TranscriptEntryKind::Message { + retained_non_user + .iter() + .filter_map(|&retained_index| { + let retained = &entries[retained_index]; + (retained.kind == TranscriptEntryKind::ProtectedMessage) + .then_some(retained.tokens) + }) + .sum() + } else { + 0 + }; + + // An entry that cannot fit beside its non-evictable context must not evict evidence. + if protected_tokens + entry.tokens > token_budget + || self.max_recent_non_user_entries == 0 + { continue; } + let eviction_order: &[TranscriptEntryKind] = match entry.kind { + TranscriptEntryKind::ProtectedMessage => &[ + TranscriptEntryKind::Message, + TranscriptEntryKind::ProtectedMessage, + ], + TranscriptEntryKind::Message => &[TranscriptEntryKind::Message], + TranscriptEntryKind::Tool => &[TranscriptEntryKind::Tool], + TranscriptEntryKind::User => unreachable!("user entries were selected separately"), + }; + + for &eviction_kind in eviction_order { + if retained_tokens(&entries, &retained_non_user, entry.kind) + .saturating_add(entry.tokens) + <= token_budget + { + break; + } + + let retained_kind_count = retained_non_user + .iter() + .filter(|&&retained_index| entries[retained_index].kind == eviction_kind) + .count(); + if retained_kind_count == 0 { + continue; + } + + evict_oldest_matching( + &entries, + &mut retained_non_user, + retained_kind_count.div_ceil(2), + |kind| kind == eviction_kind, + ); + while retained_tokens(&entries, &retained_non_user, entry.kind) + .saturating_add(entry.tokens) + > token_budget + && retained_non_user + .iter() + .any(|&retained_index| entries[retained_index].kind == eviction_kind) + { + evict_oldest_matching( + &entries, + &mut retained_non_user, + /*count*/ 1, + |kind| kind == eviction_kind, + ); + } + } + if retained_tokens(&entries, &retained_non_user, entry.kind) .saturating_add(entry.tokens) > token_budget { - let retained_kind_count = retained_non_user - .iter() - .filter(|&&retained_index| entries[retained_index].kind == entry.kind) - .count(); - evict_oldest_matching( - &entries, - &mut retained_non_user, - retained_kind_count.div_ceil(2), - |kind| kind == entry.kind, - ); - while retained_tokens(&entries, &retained_non_user, entry.kind) - .saturating_add(entry.tokens) - > token_budget - { - evict_oldest_matching( - &entries, - &mut retained_non_user, - /*count*/ 1, - |kind| kind == entry.kind, - ); - } + continue; } if retained_non_user.len().saturating_add(1) > self.max_recent_non_user_entries { - let entries_to_evict = retained_non_user.len().div_ceil(2); - evict_oldest_matching(&entries, &mut retained_non_user, entries_to_evict, |_| true); + let retained_tool_count = retained_non_user + .iter() + .filter(|&&retained_index| { + entries[retained_index].kind == TranscriptEntryKind::Tool + }) + .count(); + // Keep one slot available for protected context in smaller configured windows. + let minimum_existing_tool_entries = MIN_RECENT_TOOL_ENTRIES + .min(self.max_recent_non_user_entries.saturating_sub(1)) + .saturating_sub(usize::from(entry.kind == TranscriptEntryKind::Tool)); + let mut removable_tool_count = + retained_tool_count.saturating_sub(minimum_existing_tool_entries); + let expendable_message_count = retained_non_user + .iter() + .filter(|&&retained_index| { + entries[retained_index].kind == TranscriptEntryKind::Message + }) + .count(); + let ordinary_evidence_count = expendable_message_count + retained_tool_count; + let removable_count = expendable_message_count + removable_tool_count; + + if removable_count > 0 { + evict_oldest_matching( + &entries, + &mut retained_non_user, + ordinary_evidence_count.div_ceil(2).min(removable_count), + |kind| match kind { + TranscriptEntryKind::Message => true, + TranscriptEntryKind::Tool if removable_tool_count > 0 => { + removable_tool_count -= 1; + true + } + TranscriptEntryKind::User + | TranscriptEntryKind::ProtectedMessage + | TranscriptEntryKind::Tool => false, + }, + ); + } else if matches!( + entry.kind, + TranscriptEntryKind::ProtectedMessage | TranscriptEntryKind::Tool + ) { + let entries_to_evict = retained_non_user.len().div_ceil(2); + evict_oldest_matching( + &entries, + &mut retained_non_user, + entries_to_evict, + |kind| kind == TranscriptEntryKind::ProtectedMessage, + ); + } else { + continue; + } } retained_non_user.push_back(index); @@ -476,7 +594,9 @@ impl TranscriptConfig { .filter_map(|(index, entry)| { let component = match entry.kind { TranscriptEntryKind::User => "transcript_user", - TranscriptEntryKind::Message => "transcript_message", + TranscriptEntryKind::ProtectedMessage | TranscriptEntryKind::Message => { + "transcript_message" + } TranscriptEntryKind::Tool => "transcript_tool", }; let retained_bytes = if included[index] { diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/transcript_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/transcript_tests.rs index cf3fd83a98..c2e6ea3b7f 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/transcript_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/transcript_tests.rs @@ -1,11 +1,15 @@ use codex_extension_api::ResponseItem; +use codex_protocol::AgentPath; use codex_protocol::models::AgentMessageInputContent; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::MessagePhase; use codex_protocol::models::ReasoningItemContent; use codex_protocol::models::ReasoningItemReasoningSummary; +use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::TruncationPolicy; +use core_test_support::responses::user_message_item; use pretty_assertions::assert_eq; use super::MANUAL_APPROVAL_DEVELOPER_PREFIX; @@ -17,6 +21,16 @@ use super::TranscriptConfig; use super::TranscriptSource; use super::truncate_entry; +fn assistant_message(text: impl Into, phase: MessagePhase) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { text: text.into() }], + phase: Some(phase), + internal_chat_message_metadata_passthrough: None, + } +} + #[test] fn transcript_keeps_conversation_and_configured_sources() { let items = vec![ @@ -205,6 +219,277 @@ fn transcript_preserves_first_and_latest_user_messages_and_recent_history() { ); } +#[test] +fn transcript_preserves_user_restrictions_before_final_assistant_messages() { + let first_user_entry = "[1] user: Find a flight to New York.\n"; + let restriction_entry = "[2] user: Never book anything above $300.\n"; + let latest_user_entry = "[4] user: Yes.\n"; + let message_budget = [first_user_entry, restriction_entry, latest_user_entry] + .iter() + .map(|entry| TruncationPolicy::Bytes(entry.len()).token_budget()) + .sum(); + let items = vec![ + user_message_item("Find a flight to New York."), + user_message_item("Never book anything above $300."), + assistant_message( + "I found a $450 flight. Should I book it?", + MessagePhase::FinalAnswer, + ), + user_message_item("Yes."), + ]; + + let transcript = TranscriptConfig { + max_message_transcript_tokens: message_budget, + ..TranscriptConfig::default() + } + .build(&items) + .entries; + + assert_eq!( + transcript, + vec![first_user_entry, restriction_entry, latest_user_entry] + ); +} + +#[test] +fn transcript_preserves_recent_tool_evidence_when_protected_messages_fill_entry_limit() { + let mut items = vec![user_message_item("Inspect the workspace.")]; + items.extend((0..4).map(|index| { + assistant_message(format!("final answer {index}"), MessagePhase::FinalAnswer) + })); + items.push(ResponseItem::FunctionCall { + id: None, + name: "exec_command".to_string(), + namespace: None, + arguments: "recent evidence".to_string(), + encrypted_function_args: None, + call_id: "call-1".to_string(), + internal_chat_message_metadata_passthrough: None, + }); + + let transcript = TranscriptConfig { + max_recent_non_user_entries: 4, + ..TranscriptConfig::default() + } + .build(&items); + + assert_eq!( + transcript.entries, + vec![ + "[1] user: Inspect the workspace.\n", + "[4] assistant: final answer 2\n", + "[5] assistant: final answer 3\n", + "[6] tool exec_command call: recent evidence\n", + ] + ); + assert_eq!( + transcript + .truncations + .iter() + .map(|observation| (observation.component, observation.retained_bytes)) + .collect::>(), + vec![("transcript_message", 0); 2] + ); +} + +#[test] +fn transcript_reserves_five_recent_tool_entries_from_protected_messages() { + let mut items = vec![user_message_item("Inspect the workspace.")]; + items.extend((0..39).map(|index| { + assistant_message(format!("final answer {index}"), MessagePhase::FinalAnswer) + })); + for index in 0..3 { + let call_id = format!("call-{index}"); + items.push(ResponseItem::FunctionCall { + id: None, + name: "exec_command".to_string(), + namespace: None, + arguments: format!("click {index}"), + encrypted_function_args: None, + call_id: call_id.clone(), + internal_chat_message_metadata_passthrough: None, + }); + items.push(ResponseItem::FunctionCallOutput { + id: None, + call_id: Some(call_id), + name: None, + namespace: None, + output: FunctionCallOutputPayload::from_text(format!("page state {index}")), + internal_chat_message_metadata_passthrough: None, + }); + } + items.extend((0..20).map(|index| { + assistant_message( + format!("new final answer {index}"), + MessagePhase::FinalAnswer, + ) + })); + + let transcript = TranscriptConfig::default().build(&items).entries; + let tool_entries = transcript + .iter() + .filter(|entry| entry.contains("tool exec_command ")) + .map(String::as_str) + .collect::>(); + + assert_eq!( + tool_entries, + vec![ + "[42] tool exec_command result: page state 0\n", + "[43] tool exec_command call: click 1\n", + "[44] tool exec_command result: page state 1\n", + "[45] tool exec_command call: click 2\n", + "[46] tool exec_command result: page state 2\n", + ] + ); +} + +#[test] +fn rejected_commentary_does_not_evict_retained_message_evidence() { + let protected_entry = + "[1] assistant: Previously approved a limited action after confirming authorization.\n"; + let retained_entry = "[2] assistant: ok\n"; + let message_budget = [protected_entry, retained_entry] + .iter() + .map(|entry| TruncationPolicy::Bytes(entry.len()).token_budget()) + .sum(); + let items = vec![ + assistant_message( + "Previously approved a limited action after confirming authorization.", + MessagePhase::FinalAnswer, + ), + assistant_message("ok", MessagePhase::Commentary), + assistant_message( + "Additional commentary cannot fit.", + MessagePhase::Commentary, + ), + ]; + + let transcript = TranscriptConfig { + max_message_transcript_tokens: message_budget, + ..TranscriptConfig::default() + } + .build(&items) + .entries; + + assert_eq!(transcript, vec![protected_entry, retained_entry]); +} + +#[test] +fn transcript_evicts_protected_messages_in_cacheable_chunks() { + let user_entry = "[1] user: Inspect the workspace.\n"; + let assistant_entry = "[2] assistant: final answer 0\n"; + let message_budget = TruncationPolicy::Bytes(user_entry.len()).token_budget() + + 4 * TruncationPolicy::Bytes(assistant_entry.len()).token_budget(); + let configurations = [ + TranscriptConfig { + max_recent_non_user_entries: 4, + ..TranscriptConfig::default() + }, + TranscriptConfig { + max_message_transcript_tokens: message_budget, + ..TranscriptConfig::default() + }, + ]; + + for config in configurations { + let build_transcript = |message_count| { + let mut items = vec![user_message_item("Inspect the workspace.")]; + items.extend((0..message_count).map(|index| { + assistant_message(format!("final answer {index}"), MessagePhase::FinalAnswer) + })); + config.build(&items).entries + }; + + assert_eq!( + build_transcript(5), + vec![ + user_entry, + "[4] assistant: final answer 2\n", + "[5] assistant: final answer 3\n", + "[6] assistant: final answer 4\n", + ] + ); + assert_eq!( + build_transcript(6), + vec![ + user_entry, + "[4] assistant: final answer 2\n", + "[5] assistant: final answer 3\n", + "[6] assistant: final answer 4\n", + "[7] assistant: final answer 5\n", + ] + ); + assert_eq!( + build_transcript(7), + vec![ + user_entry, + "[6] assistant: final answer 4\n", + "[7] assistant: final answer 5\n", + "[8] assistant: final answer 6\n", + ] + ); + } +} + +#[test] +fn transcript_does_not_protect_legacy_inter_agent_instructions() { + let communication = InterAgentCommunication::new( + AgentPath::root(), + AgentPath::root() + .join("worker") + .expect("worker path should be valid"), + Vec::new(), + "Review the pending action.".to_string(), + /*trigger_turn*/ true, + ); + let items = vec![ + ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: serde_json::to_string(&communication) + .expect("legacy inter-agent communication should serialize"), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }, + ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: "I found a $450 flight. Should I book it?".to_string(), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }, + ResponseItem::FunctionCall { + id: None, + name: "exec_command".to_string(), + namespace: None, + arguments: "booking details".to_string(), + encrypted_function_args: None, + call_id: "call-1".to_string(), + internal_chat_message_metadata_passthrough: None, + }, + ]; + + let transcript = TranscriptConfig { + max_recent_non_user_entries: 2, + ..TranscriptConfig::default() + } + .build(&items) + .entries; + + assert_eq!( + transcript, + vec![ + "[2] assistant: I found a $450 flight. Should I book it?\n", + "[3] tool exec_command call: booking details\n", + ] + ); +} + #[test] fn transcript_reserves_separate_budget_for_recent_tool_evidence() { let mut items = (0..8)