From e64530972398604464a6aa0d1ace7b9399e9e942 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 19 Oct 2026 21:29:10 +0000 Subject: [PATCH] Deduplicate native agent messages in Guardian review context (#55391) ## Why Full-transcript fallback can select encrypted agent messages already present in the reviewer's history. These duplicates should not consume the input budget or displace new evidence. ## What changed - Reconcile native agent messages with the effective reviewer history and earlier messages in the pending input before budgeting, including after compaction. - Deduplicate messages with nonempty IDs while preserving messages without IDs. Allow delivery again when the prior copy is absent from the current history. - Reject conflicting authors, recipients, or content under the same ID without modifying the pending context. ## Testing Add coverage for duplicate admission, conflicting payloads, redelivery after history removal, and retained reviewer requests. Extend the history eviction regression to verify that full-transcript fallback reuses the reviewer and preserves a single encrypted agent message. GitOrigin-RevId: 97e90f6359a7cf5a75d1c5967fc98cd885afb4e4 --- codex-rs/core/src/guardian/input_budget.rs | 17 ++-- codex-rs/core/tests/suite/guardian_history.rs | 56 +++++++++-- .../src/async_scorer/request_tests.rs | 52 +++++++++- .../src/async_scorer/transcript.rs | 3 +- .../src/async_scorer/transcript_tests.rs | 15 ++- codex-rs/guardian-context/src/enforcement.rs | 68 +++++++++++++ .../guardian-context/src/enforcement_tests.rs | 98 +++++++++++++++++++ codex-rs/guardian-context/src/lib.rs | 5 + 8 files changed, 294 insertions(+), 20 deletions(-) diff --git a/codex-rs/core/src/guardian/input_budget.rs b/codex-rs/core/src/guardian/input_budget.rs index 3a5f348ac8..923aacf1cb 100644 --- a/codex-rs/core/src/guardian/input_budget.rs +++ b/codex-rs/core/src/guardian/input_budget.rs @@ -85,12 +85,14 @@ pub(crate) async fn check_pending(session: &Session, turn: &TurnContext) -> Code let maximum = effective_input_token_limit(turn.model_info(), turn.config.model_context_window) .saturating_sub(super::request_budget::INPUT_TOKEN_MARGIN); let mut context = pending.0.clone(); - context.retain_new_instructions( - &session - .clone_history() - .await - .for_prompt_annotated(&turn.model_info().input_modalities), - ); + let history = session + .clone_history() + .await + .for_prompt_annotated(&turn.model_info().input_modalities); + context.retain_new_instructions(&history); + context + .retain_new_agent_messages(&history) + .map_err(|error| CodexErr::InvalidRequest(error.to_string()))?; if context.estimated_tokens().saturating_add(minimum_prefix) > maximum { context .enforce_budget( @@ -153,6 +155,9 @@ pub(crate) async fn finalize( // Use the history that will actually reach the model. Recompute after every // compaction retry; evidence removed by compaction must be delivered again. context.retain_new_instructions(&history); + context + .retain_new_agent_messages(&history) + .map_err(|error| CodexErr::InvalidRequest(error.to_string()))?; let prompt = build_prompt( history .into_iter() diff --git a/codex-rs/core/tests/suite/guardian_history.rs b/codex-rs/core/tests/suite/guardian_history.rs index adc2c5bb56..fdb2f0a858 100644 --- a/codex-rs/core/tests/suite/guardian_history.rs +++ b/codex-rs/core/tests/suite/guardian_history.rs @@ -12,12 +12,14 @@ use codex_history::InitialHistory; use codex_history::ResumedHistory; use codex_history::RolloutItem; use codex_login::CodexAuth; +use codex_protocol::AgentPath; use codex_protocol::TranscriptFormat; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::mcp::ClientMcpExtensions; use codex_protocol::models::ImageReference; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::Op; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::request_user_input::RequestUserInputAnswer; @@ -408,7 +410,7 @@ async fn guardian_history_uses_deltas_between_eviction_batches( .build_with_auto_env(&server) .await?; let plan = r#"{"plan":[{"step":"inspect the repository","status":"completed"}]}"#; - let mut inspection: Vec<_> = (0..130) + let mut inspection: Vec<_> = (0..155) .map(|index| ev_function_call(&format!("initial-{index}"), "update_plan", plan)) .collect(); inspection.push(ev_completed("inspection")); @@ -436,8 +438,10 @@ async fn guardian_history_uses_deltas_between_eviction_batches( let mut responses = Vec::new(); for index in 0..3 { if index == 1 { - // Force eviction after the first review has saved its transcript cursor. - let mut traffic: Vec<_> = (0..130) + // Evict the older prefix after a completed review, while retaining the + // recent encrypted reply within the transcript's recent-entry limit. + // Full fallback must not admit that reply twice. + let mut traffic: Vec<_> = (0..10) .map(|index| ev_function_call(&format!("later-{index}"), "update_plan", plan)) .collect(); traffic.push(ev_completed("more-inspection")); @@ -468,7 +472,22 @@ async fn guardian_history_uses_deltas_between_eviction_batches( } responses.push(sse(vec![ev_completed("done")])); let reviews = mount_sse_sequence(&server, responses).await; - test.submit_text_turn("Continue inspecting.").await?; + test.codex + .submit(Op::InterAgentCommunication { + communication: InterAgentCommunication::new_encrypted( + AgentPath::root(), + AgentPath::root().join("worker").expect("worker path"), + Vec::new(), + "opaque-parent-restriction".into(), + /*trigger_turn*/ true, + ), + start_options: Default::default(), + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; let requests = reviews.requests(); let guardian_requests = requests .iter() @@ -479,7 +498,13 @@ async fn guardian_history_uses_deltas_between_eviction_batches( .map(|request| { request .message_input_text_groups("user") - .last() + .into_iter() + .rfind(|group| { + group.iter().any(|text| { + text.contains(">>> TRANSCRIPT START\n") + || text.contains(">>> TRANSCRIPT DELTA START\n") + }) + }) .expect("Guardian review prompt") .join("") }) @@ -515,9 +540,24 @@ async fn guardian_history_uses_deltas_between_eviction_batches( assert!(prompts[1].contains(restriction)); } assert!(prompts[2].contains("review-2")); - assert_eq!( - guardian_requests[1].body_json()["client_metadata"]["thread_id"], - guardian_requests[2].body_json()["client_metadata"]["thread_id"] + let expected_messages = guardian_requests[0].inputs_of_type("agent_message"); + assert_eq!(expected_messages.len(), 1); + assert!(expected_messages[0]["id"].is_string()); + for request in &guardian_requests { + assert_eq!(request.inputs_of_type("agent_message"), expected_messages); + assert_eq!( + request.body_json()["client_metadata"]["thread_id"], + guardian_requests[0].body_json()["client_metadata"]["thread_id"], + "the completed reviewer is reused across full-transcript fallback", + ); + } + let history = test.codex.conversation_history_snapshot().await; + assert!( + history.review_items().any(|item| { + item.id() + .is_some_and(|id| Some(id.as_str()) == expected_messages[0]["id"].as_str()) + }), + "full fallback must still select the encrypted reply after eviction" ); test.codex.shutdown_and_wait().await?; Ok(()) diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/request_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/request_tests.rs index 9938bd22c3..734bca0839 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/request_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/request_tests.rs @@ -63,6 +63,17 @@ fn collect(history: &History, images: &[ContentItem]) -> CollectedTranscript { #[test] fn preparation_preserves_committed_prefix_and_selects_only_new_entries_and_images() { + let native = ResponseItem::AgentMessage { + id: Some(ResponseItemId::from_server("amsg_parent".to_owned())), + author: "/root".to_owned(), + recipient: "/root/worker".to_owned(), + content: vec![ + codex_protocol::models::AgentMessageInputContent::EncryptedContent { + encrypted_content: "opaque-parent-restriction".to_owned(), + }, + ], + internal_chat_message_metadata_passthrough: None, + }; let mut images = [ ImageReference::Inline { image_url: "data:image/png;base64,Zmlyc3Q=".to_owned(), @@ -90,7 +101,11 @@ fn preparation_preserves_committed_prefix_and_selects_only_new_entries_and_image .collect::>() }; let mut history = History { - items: vec![responses::user_message_item("Inspect README.md.")], + items: vec![ + responses::user_message_item("Inspect README.md."), + native.clone(), + native.clone(), + ], version: 1, }; if let ResponseItem::Message { content, .. } = &mut history.items[0] { @@ -106,6 +121,15 @@ fn preparation_preserves_committed_prefix_and_selects_only_new_entries_and_image vec![images[1].clone(), images[0].clone()] ); assert!(first.input.iter().all(|item| item.id().is_some())); + assert_eq!( + first + .input + .iter() + .filter(|item| matches!(item.item, ResponseItem::AgentMessage { .. })) + .map(|item| item.item.clone()) + .collect::>(), + vec![native.clone()] + ); let mut committed = first.input; // Keep the actual returned item, including the server ID and metadata. let output: ResponseItem = @@ -117,6 +141,10 @@ fn preparation_preserves_committed_prefix_and_selects_only_new_entries_and_image history .items .push(responses::user_message_item("Inspect LICENSE next.")); + history.items.push(native.clone()); + let mut next_native = native.clone(); + next_native.set_id(Some(ResponseItemId::from_server("amsg_next".to_owned()))); + history.items.push(next_native.clone()); // Source detail can differ from the normalized image already in classifier history. if let ContentItem::InputImage { detail, .. } = &mut images[0] { *detail = Some(ImageDetail::High); @@ -140,7 +168,16 @@ fn preparation_preserves_committed_prefix_and_selects_only_new_entries_and_image input_images(&second.input[committed.len()..]), vec![images[2].clone()] ); - assert_eq!(second.cursor.transcript_entry_count, 2); + assert_eq!( + second + .input + .iter() + .filter(|item| matches!(item.item, ResponseItem::AgentMessage { .. })) + .map(|item| item.item.clone()) + .collect::>(), + vec![native.clone(), next_native.clone()] + ); + assert_eq!(second.cursor.transcript_entry_count, 6); let delta = serde_json::to_string( &second.input[committed.len()..] .iter() @@ -149,7 +186,7 @@ fn preparation_preserves_committed_prefix_and_selects_only_new_entries_and_image ) .unwrap(); assert!(delta.contains("TRANSCRIPT DELTA START")); - assert!(delta.contains("[2] user: Inspect LICENSE next.")); + assert!(delta.contains("[4] user: Inspect LICENSE next.")); assert!(!delta.contains("[1] user: Inspect README.md.")); assert!(delta.contains("Do not write files.")); assert!(delta.contains("read_file")); @@ -162,6 +199,15 @@ fn preparation_preserves_committed_prefix_and_selects_only_new_entries_and_image .prepare_retained(&evidence, &mut ConversationState::default(), INPUT_BUDGET) .unwrap(); assert_eq!(input_images(&rebuilt.input), input_images(&second.input)); + assert_eq!( + rebuilt + .input + .iter() + .filter(|item| matches!(item.item, ResponseItem::AgentMessage { .. })) + .map(|item| item.item.clone()) + .collect::>(), + vec![native, next_native] + ); } #[test] 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 4a1337b33a..12222eb9bf 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs @@ -195,7 +195,8 @@ impl CollectedTranscript { } }; let transcript = self.profile.prepare_transcript(entries, offset); - let context = self.context.compose(presentation, transcript)?; + let mut context = self.context.compose(presentation, transcript)?; + context.retain_new_agent_messages(reviewer_history)?; Ok((context, next)) } } 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 64635534a2..3b1567b0e7 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 @@ -3,6 +3,7 @@ use codex_extension_api::ConversationHistorySnapshot; use codex_extension_api::ResponseItem; use codex_guardian_context::ContextTarget; use codex_protocol::AgentPath; +use codex_protocol::ResponseItemId; use codex_protocol::models::AgentMessageInputContent; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputContentItem; @@ -1258,7 +1259,7 @@ fn transcript_omits_media_payloads_and_keeps_readable_content() { #[test] fn transcript_preserves_encrypted_agent_messages_and_omits_other_encrypted_fields() { - let items = vec![ + let mut items = vec![ ResponseItem::AgentMessage { id: None, author: "worker".to_string(), @@ -1303,6 +1304,9 @@ fn transcript_preserves_encrypted_agent_messages_and_omits_other_encrypted_field internal_chat_message_metadata_passthrough: None, }, ]; + let mut native = items[0].clone(); + native.set_id(Some(ResponseItemId::from_server("amsg_parent".to_owned()))); + items.extend([native.clone(), native.clone()]); let context = TranscriptConfig::default() .build_context(ContextInput { @@ -1319,7 +1323,14 @@ fn transcript_preserves_encrypted_agent_messages_and_omits_other_encrypted_field }) .expect("collect transcript") .into_messages(); - assert_eq!(context[1], items[0]); + assert_eq!( + context + .iter() + .filter(|item| matches!(item, ResponseItem::AgentMessage { .. })) + .cloned() + .collect::>(), + vec![items[0].clone(), native] + ); let text = context .iter() .filter_map(|item| match item { diff --git a/codex-rs/guardian-context/src/enforcement.rs b/codex-rs/guardian-context/src/enforcement.rs index 9f17891fe9..ed3213be19 100644 --- a/codex-rs/guardian-context/src/enforcement.rs +++ b/codex-rs/guardian-context/src/enforcement.rs @@ -3,12 +3,16 @@ //! hosts may shorten historical instructions after compaction cannot make room. //! Every reduction reserves an omission notice and preserves source order. //! JSON recovery shortens only the body; line mode shortens the whole rendered entry. +//! Native agent messages are reconciled with current history before budgeting. +use std::collections::HashMap; use std::collections::HashSet; +use codex_history::ResponseItemEnvelope; use codex_protocol::models::ContentItem; use codex_protocol::models::ImageDetail; use codex_protocol::models::ImageReference; +use codex_protocol::models::ResponseItem; use codex_protocol::protocol::TruncationPolicy; use crate::ComposedContext; @@ -92,6 +96,70 @@ pub enum HistoryTruncation { } impl ComposedContext { + /// Omits native agent messages already present in the effective reviewer history + /// or earlier in this input. Hosts must repeat this check after compaction; + /// missing evidence may be delivered again. Identity excludes host annotations. + /// Conflicting payloads under one ID reject the input without changing it. + pub fn retain_new_agent_messages( + &mut self, + reviewer_history: &[ResponseItemEnvelope], + ) -> Result<(), SectionError> { + if !self.sections.iter().any(|section| { + matches!( + §ion.delivery, + SectionDelivery::Message(message) + if matches!(message.content.item, ResponseItem::AgentMessage { .. }) + ) + }) { + return Ok(()); + } + let mut seen = HashMap::new(); + let mut duplicates = HashSet::new(); + let pending = + self.sections + .iter() + .enumerate() + .filter_map(|(index, section)| match §ion.delivery { + SectionDelivery::Message(message) => Some((Some(index), &message.content.item)), + SectionDelivery::UserContent(_) => None, + }); + for (index, item) in reviewer_history + .iter() + .map(|item| (None, &item.item)) + .chain(pending) + { + let ResponseItem::AgentMessage { + id: Some(id), + author, + recipient, + content, + .. + } = item + else { + continue; + }; + if id.is_empty() { + continue; + } + let payload = (author, recipient, content); + if let Some(previous) = seen.insert(id, payload) { + if previous != payload { + return Err(SectionError::ConflictingAgentMessage); + } + if let Some(index) = index { + duplicates.insert(index); + } + } + } + let mut index = 0; + self.sections.retain(|_| { + let retain = !duplicates.contains(&index); + index += 1; + retain + }); + Ok(()) + } + /// Applies host image admission before aggregate selection, preserving section /// identity and each retained item's selection policy. pub fn retain_images( diff --git a/codex-rs/guardian-context/src/enforcement_tests.rs b/codex-rs/guardian-context/src/enforcement_tests.rs index c8505f795e..ecd4b20d7b 100644 --- a/codex-rs/guardian-context/src/enforcement_tests.rs +++ b/codex-rs/guardian-context/src/enforcement_tests.rs @@ -9,6 +9,104 @@ use codex_protocol::models::ResponseItem; use pretty_assertions::assert_eq; use std::borrow::Cow; +#[test] +fn native_agent_admission_uses_current_history_and_preserves_identity() { + let original = ResponseItem::AgentMessage { + id: Some(codex_protocol::ResponseItemId::from_server( + "amsg_original".to_owned(), + )), + author: "/root".to_owned(), + recipient: "/root/worker".to_owned(), + content: vec![ + codex_protocol::models::AgentMessageInputContent::EncryptedContent { + encrypted_content: "opaque-restriction".to_owned(), + }, + ], + internal_chat_message_metadata_passthrough: None, + }; + let mut new = original.clone(); + new.set_id(Some(codex_protocol::ResponseItemId::from_server( + "amsg_new".to_owned(), + ))); + let mut anonymous = original.clone(); + anonymous.set_id(/*new_id*/ None); + let mut empty_id = original.clone(); + empty_id.set_id(Some(codex_protocol::ResponseItemId::from_server( + String::new(), + ))); + let context = ComposedContext { + sections: [ + original.clone(), + original.clone(), + new.clone(), + new.clone(), + anonymous.clone(), + anonymous.clone(), + empty_id.clone(), + empty_id.clone(), + ] + .into_iter() + .map(|item| SectionOutput { + id: "conversation_transcript", + delivery: SectionDelivery::Message(Budgeted::required(Box::new( + ResponseItemEnvelope::new(item), + ))), + }) + .collect(), + truncations: Vec::new(), + }; + let mut delivered = original.clone(); + delivered.set_turn_id_if_missing("reviewer-turn"); + // Removing the old copy through compaction or rollback permits delivery again. + for history in [vec![ResponseItemEnvelope::new(delivered)], Vec::new()] { + let mut pending = context.clone(); + pending.retain_new_agent_messages(&history).unwrap(); + let mut expected = if history.is_empty() { + vec![original.clone()] + } else { + Vec::new() + }; + expected.extend([ + new.clone(), + anonymous.clone(), + anonymous.clone(), + empty_id.clone(), + empty_id.clone(), + ]); + // Already delivered copies must not consume admission budget or evict new evidence. + let admitted = pending + .enforce_budget( + RequestBudget { + max_input_tokens: expected.iter().map(crate::estimate_input_tokens).sum(), + existing_context_tokens: 0, + }, + "evidence omitted".to_owned(), + HistoryTruncation::Preserve, + ) + .unwrap(); + assert_eq!(admitted.into_messages(), expected); + } + + let mut conflicting = context; + let SectionDelivery::Message(message) = &mut conflicting.sections[1].delivery else { + panic!("native message") + }; + let ResponseItem::AgentMessage { content, .. } = &mut message.content.item else { + panic!("agent message") + }; + *content = vec![ + codex_protocol::models::AgentMessageInputContent::EncryptedContent { + encrypted_content: "different-evidence".to_owned(), + }, + ]; + let before = conflicting.clone().into_messages(); + assert_eq!( + conflicting.retain_new_agent_messages(&[]), + Err(SectionError::ConflictingAgentMessage) + ); + assert_eq!(conflicting.into_messages(), before); +} + fn text(value: &str) -> ContentItem { ContentItem::InputText { text: value.to_owned(), diff --git a/codex-rs/guardian-context/src/lib.rs b/codex-rs/guardian-context/src/lib.rs index af09a72acc..5c9ac85273 100644 --- a/codex-rs/guardian-context/src/lib.rs +++ b/codex-rs/guardian-context/src/lib.rs @@ -246,6 +246,8 @@ pub trait SectionContributor: Send + Sync { /// A section could not provide the evidence needed for a valid review context. #[derive(Clone, Debug, Eq, PartialEq)] pub enum SectionError { + /// One native agent-message ID identifies different authors, recipients or content. + ConflictingAgentMessage, /// Evidence required by this contributor for the current input is missing. MissingRequiredEvidence { section: &'static str }, /// A section cannot be delivered by the requested consumer. @@ -257,6 +259,9 @@ pub enum SectionError { impl std::fmt::Display for SectionError { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { + Self::ConflictingAgentMessage => { + formatter.write_str("conflicting native agent messages share an ID") + } Self::MissingRequiredEvidence { section } => { write!(formatter, "missing required evidence for section {section}") }