From 1f281290e3e565c9d3287f1a226dbb00a4c01272 Mon Sep 17 00:00:00 2001 From: ioanbudea-OAI Date: Mon, 26 Oct 2026 23:54:22 +0000 Subject: [PATCH] Centralize Guardian root context preparation (#57312) ## What changed Move root message preparation and handoff context selection into `codex-guardian-context`. Keep source collection and handoff boundary detection in core, while shared helpers handle text limits, whole-record verified answer admission, delivered assistant message lookup, and retained assistant fallback. GitOrigin-RevId: a3eb35cddb3b0c5b93243b9d9c553397a7a61a14 --- .../core/src/agent/control/root_handoff.rs | 60 ++--- .../src/agent/control/user_authorization.rs | 175 ++++++-------- codex-rs/guardian-context/src/lib.rs | 12 +- .../guardian-context/src/preparation/mod.rs | 2 + .../guardian-context/src/preparation/root.rs | 218 ++++++++++++++++++ .../guardian-context/src/verified_answers.rs | 7 +- 6 files changed, 319 insertions(+), 155 deletions(-) create mode 100644 codex-rs/guardian-context/src/preparation/root.rs diff --git a/codex-rs/core/src/agent/control/root_handoff.rs b/codex-rs/core/src/agent/control/root_handoff.rs index d6d49430f9..1cbe32ca6f 100644 --- a/codex-rs/core/src/agent/control/root_handoff.rs +++ b/codex-rs/core/src/agent/control/root_handoff.rs @@ -12,14 +12,23 @@ use codex_protocol::ToolName; use codex_protocol::models::ExecutedToolCallArguments; use codex_protocol::models::ResponseItem; -const ROOT_CONTEXT_WINDOW: usize = 3; - pub(super) fn selected_message_indices( history: &[ResponseItemEnvelope], messages: &[(Option, GuardianRootMessage)], worker: &AgentPath, tool_namespace: Option<&str>, ) -> Option> { + codex_guardian_context::select_root_handoff_indices( + messages, + &handoff_boundaries(history, worker, tool_namespace), + ) +} + +pub(super) fn handoff_boundaries( + history: &[ResponseItemEnvelope], + worker: &AgentPath, + tool_namespace: Option<&str>, +) -> Vec { let configured_handoff_names = ["spawn_agent", "send_message", "followup_task"].map(|name| { ( codex_tools::code_mode_name_for_tool_name(&ToolName::new( @@ -29,7 +38,7 @@ pub(super) fn selected_message_indices( name, ) }); - let mut selected = BTreeSet::new(); + let mut boundaries = Vec::new(); let mut preceding_order = None; for envelope in history { if let Some(order) = envelope.metadata.as_ref().and_then(|metadata| { @@ -78,51 +87,10 @@ pub(super) fn selected_message_indices( if (direct || nested) && let Some(boundary) = preceding_order { - selected.extend( - messages - .iter() - .enumerate() - .rev() - .filter_map(|(index, (order, message))| { - if matches!(message, GuardianRootMessage::UserInput(_)) { - return None; - } - order.filter(|order| *order <= boundary).map(|_| index) - }) - .take(ROOT_CONTEXT_WINDOW), - ); + boundaries.push(boundary); } } - if !selected.is_empty() { - // Include the latest ordinary messages even without a new handoff to this branch. - selected.extend( - messages - .iter() - .enumerate() - .rev() - .filter_map(|(index, (order, message))| { - (order.is_some() && !matches!(message, GuardianRootMessage::UserInput(_))) - .then_some(index) - }) - .take(ROOT_CONTEXT_WINDOW), - ); - // Later user inputs can restrict an earlier handoff even when they are not recent. - // Leave all user inputs, verified answers, and unordered legacy evidence to the shared cap. - selected.extend( - messages - .iter() - .enumerate() - .filter_map(|(index, (order, message))| { - (order.is_none() - || matches!( - message, - GuardianRootMessage::User(_) | GuardianRootMessage::UserInput(_) - )) - .then_some(index) - }), - ); - } - (!selected.is_empty()).then_some(selected) + boundaries } fn is_handoff(name: &str, args: &serde_json::Value, worker: &AgentPath) -> bool { diff --git a/codex-rs/core/src/agent/control/user_authorization.rs b/codex-rs/core/src/agent/control/user_authorization.rs index 749921aefb..c97d8b9d05 100644 --- a/codex-rs/core/src/agent/control/user_authorization.rs +++ b/codex-rs/core/src/agent/control/user_authorization.rs @@ -10,7 +10,6 @@ //! Handoff projection preserves saved heartbeat instructions and the active turn's skills. use std::borrow::Cow; -use std::collections::HashMap; use std::collections::VecDeque; use std::collections::hash_map::DefaultHasher; use std::hash::Hash; @@ -24,17 +23,18 @@ use crate::context::ContextualUserFragment; use crate::context::GuardianReviewEvidence; use crate::context::UserGoalUpdate; use crate::context::is_contextual_user_fragment; -use crate::context::render_retained_assistant_context; use crate::event_mapping::parse_turn_item; -use crate::guardian::GUARDIAN_MAX_ROOT_MESSAGE_TOKENS; -use crate::guardian::guardian_truncate_text; +use codex_guardian_context::RootAssistantCandidate; +use codex_guardian_context::RootAssistantContent; +use codex_guardian_context::RootDeliveredMessage; +use codex_guardian_context::RootMessageCandidate; +use codex_guardian_context::RootRetainedAssistant; use codex_history::ReconciledRetainedContext; use codex_history::RetainedContextEntry; use codex_history::RetainedContextOrder; use codex_history::RetainedUserMessage; use codex_protocol::AgentPath; use codex_protocol::ThreadId; -use codex_protocol::items::AgentMessageContent; use codex_protocol::items::TurnItem; use codex_protocol::models::ContentItem; use codex_protocol::models::MessagePhase; @@ -75,7 +75,7 @@ impl LocalAgentControl { let history = root_history.conversation_history_snapshot(); // Join calls to host-confirmed outputs in this snapshot. Older outputs without // captured text cannot establish what was sent, including after a hook rewrite. - let delivered_messages = root_history + let delivered_candidates = root_history .annotated_items() .iter() .rev() @@ -92,10 +92,9 @@ impl LocalAgentControl { .as_ref()? .delivered_assistant_message .as_deref()?; - Some((call_id.as_str(), text)) + Some(RootDeliveredMessage { call_id, text }) }) - .take(MAX_ROOT_MESSAGES) - .collect::>(); + .collect::>(); let root_evidence = root_thread .session .services @@ -141,7 +140,7 @@ impl LocalAgentControl { ); let mut latest_user_turn_id = reconciled.latest_user_turn_id.clone(); let mut missing_root_instructions = reconciled.missing_user_messages; - let mut messages = reconciled + let user_candidates = reconciled .ordered_entries() .filter_map(|(order, entry)| match entry { RetainedContextEntry::UserMessage(message) => { @@ -176,23 +175,24 @@ impl LocalAgentControl { return None; } (!is_summary_message(&text) && !text.trim_start().starts_with("")) - .then(|| { - ( - Some(order), - GuardianRootMessage::User( - guardian_truncate_text(&text, GUARDIAN_MAX_ROOT_MESSAGE_TOKENS) - .0, - ), - ) + .then(|| RootMessageCandidate { + order: Some(order), + message: GuardianRootMessage::User(text.to_string()), }) } - RetainedContextEntry::VerifiedAnswer(answer) => { - codex_guardian_context::render_verified_answer(answer) - .map(|text| (Some(order), GuardianRootMessage::UserInput(text))) - } + RetainedContextEntry::VerifiedAnswer(answer) => Some(RootMessageCandidate { + order: Some(order), + message: GuardianRootMessage::UserInput( + codex_guardian_context::render_original_verified_answer(answer), + ), + }), RetainedContextEntry::AssistantMessage(_) => None, }) .collect::>(); + let mut messages = codex_guardian_context::prepare_root_user_messages(&user_candidates) + .into_iter() + .map(|message| (message.order, message.message)) + .collect::>(); // Old opt-out checkpoints have genuine instructions but no retained source order. // Keep them separately bounded and explicitly unordered; never replace newer facts. let mut legacy_messages = VecDeque::new(); @@ -210,10 +210,7 @@ impl LocalAgentControl { latest_user_turn_id = latest_legacy_turn_id; } } - let mut missing_assistant_context = retained_context.has_omitted_assistant_messages(); - // Prefer a renderable retained original over a shortened checkpoint copy. - // Otherwise preserve the bounded live evidence, including confirmed messaging sends. - let live_assistant_messages = root_history + let assistant_candidates = root_history .annotated_items() .iter() .filter(|envelope| { @@ -223,45 +220,41 @@ impl LocalAgentControl { .is_some_and(|metadata| metadata.compaction_output) }) .filter_map(|envelope| { - let text = root_assistant_text(&envelope.item, &delivered_messages)?; - let id = envelope - .item - .id() - .map(codex_protocol::ResponseItemId::as_str); - let retained = reconciled.ordered_entries().find_map(|(order, entry)| { - let RetainedContextEntry::AssistantMessage(message) = entry else { - return None; - }; - id.filter(|id| message.message_id.as_deref() == Some(*id)) - .map(|_| (order, message)) - }); - if let Some((order, message)) = retained - && let Some(text) = render_retained_assistant_context(message) - { - return Some((id, Some(order), GuardianRootMessage::Assistant(text))); - } - let order = envelope - .metadata - .as_ref() - .filter(|metadata| !metadata.inherited_user_message) - .and_then(|metadata| metadata.user_input_order) - .map(RetainedContextOrder::Local); - let order = retained.map(|(order, _)| order).or(order); - let message = if order.is_some() { - GuardianRootMessage::Assistant(text) - } else { - GuardianRootMessage::UnorderedAssistant(text) + let content = match &envelope.item { + ResponseItem::FunctionCall { call_id, .. } => { + RootAssistantContent::DeliveredCall { + call_id: call_id.clone(), + } + } + ResponseItem::Message { + role, + content, + phase, + .. + } if role == "assistant" && *phase != Some(MessagePhase::Commentary) => { + RootAssistantContent::Text(content) + } + _ => return None, }; - Some((id, order, message)) + Some(RootAssistantCandidate { + message_id: envelope.item.id().map(|id| id.as_str().to_owned()), + order: envelope + .metadata + .as_ref() + .filter(|metadata| !metadata.inherited_user_message) + .and_then(|metadata| metadata.user_input_order) + .map(RetainedContextOrder::Local), + content, + }) }) .collect::>(); - let mut assistant_messages = reconciled + let retained_assistant_candidates = reconciled .ordered_entries() .filter_map(|(order, entry)| { let RetainedContextEntry::AssistantMessage(message) = entry else { return None; }; - if message.phase == Some(MessagePhase::Commentary) + let standalone_eligible = !(message.phase == Some(MessagePhase::Commentary) || message.message_id.as_deref().is_some_and(|id| { root_history .raw_items() @@ -271,29 +264,27 @@ impl LocalAgentControl { id: Some(source_id), phase: Some(MessagePhase::Commentary), .. } if source_id.as_str() == id) }) - }) - { - return None; - } - if message.message_id.as_deref().is_some_and(|id| { - live_assistant_messages - .iter() - .any(|(live_id, _, _)| *live_id == Some(id)) - }) { - return None; - } - let rendered = - render_retained_assistant_context(message).map(GuardianRootMessage::Assistant); - missing_assistant_context |= rendered.is_none(); - rendered.map(|message| (Some(order), message)) + })); + Some(RootRetainedAssistant { + order, + message, + standalone_eligible, + }) }) .collect::>(); - assistant_messages.extend( - live_assistant_messages - .into_iter() - .map(|(_, order, message)| (order, message)), + let assistants = codex_guardian_context::prepare_root_assistant_messages( + &assistant_candidates, + &retained_assistant_candidates, + &delivered_candidates, + retained_context.has_omitted_assistant_messages(), + ); + let mut missing_assistant_context = assistants.missing_context; + messages.extend( + assistants + .messages + .into_iter() + .map(|message| (message.order, message.message)), ); - messages.extend(assistant_messages); // Unknown order cannot establish recency; prune these assistants first. messages.sort_by_key(|(order, _)| *order); let mut review_context_revision = history.guardian_review_context_revision(); @@ -364,7 +355,8 @@ impl LocalAgentControl { ); } let mut authorization_version = root_evidence.authorization_version(history.as_ref()); - if !authorization_version.retained_context_complete { + let missing_verified_answers = !authorization_version.retained_context_complete; + if missing_verified_answers { messages.insert( /*index*/ 0, GuardianRootMessage::IncompleteVerifiedAnswers, @@ -397,30 +389,3 @@ impl LocalAgentControl { }) } } - -fn root_assistant_text( - item: &ResponseItem, - delivered_messages: &HashMap<&str, &str>, -) -> Option { - let text = match item { - ResponseItem::FunctionCall { call_id, .. } => { - (*delivered_messages.get(call_id.as_str())?).to_owned() - } - _ => { - let Some(TurnItem::AgentMessage(message)) = parse_turn_item(item) else { - return None; - }; - if message.phase == Some(MessagePhase::Commentary) { - return None; - } - message - .content - .iter() - .map(|content| match content { - AgentMessageContent::Text { text } => text.as_str(), - }) - .collect::() - } - }; - Some(guardian_truncate_text(&text, GUARDIAN_MAX_ROOT_MESSAGE_TOKENS).0) -} diff --git a/codex-rs/guardian-context/src/lib.rs b/codex-rs/guardian-context/src/lib.rs index 9419fc54ac..771587bf76 100644 --- a/codex-rs/guardian-context/src/lib.rs +++ b/codex-rs/guardian-context/src/lib.rs @@ -28,6 +28,11 @@ pub use action::PlannedAction; pub use action::PlannedActionKind; pub use action::action_for_review; pub use authorization::GuardianRootMessage; +pub use preparation::root::RootAssistantCandidate; +pub use preparation::root::RootAssistantContent; +pub use preparation::root::RootDeliveredMessage; +pub use preparation::root::RootMessageCandidate; +pub use preparation::root::RootRetainedAssistant; pub use section::ContextSection; pub use entry::ConversationTranscriptEntry; @@ -48,6 +53,7 @@ mod verified_answers; pub use verified_answers::RenderedVerifiedAnswers; pub use verified_answers::TrustedUserAnswer; pub use verified_answers::collect_verified_answer_candidates; +pub use verified_answers::render_original_verified_answer; pub use verified_answers::render_verified_answer; pub use verified_answers::render_verified_answers; @@ -86,6 +92,10 @@ pub use composition::CollectedContext; pub use composition::ComposedContext; pub use composition::ContextPresentation; pub use composition::PreparedTranscript; +pub use preparation::root::PreparedRootAssistants; +pub use preparation::root::prepare_root_assistant_messages; +pub use preparation::root::prepare_root_user_messages; +pub use preparation::root::select_root_handoff_indices; pub use profile::ContextProfile; mod authorization; mod entry; @@ -164,7 +174,7 @@ pub struct SectionInput<'a> { pub transcript: &'a ConversationTranscriptConfig, /// Bounded root evidence resolved by the host; empty when not applicable. pub root_conversation: &'a [GuardianRootMessage], - /// Bounded, role-labeled answers selected from the host-owned context snapshot. + /// Role-labeled sync answer candidates or already-bounded answers from other consumers. pub trusted_user_answers: &'a [TrustedUserAnswer], /// Exact action JSON; sync composition bounds the reason and descriptions. pub planned_action: Option<&'a PlannedAction>, diff --git a/codex-rs/guardian-context/src/preparation/mod.rs b/codex-rs/guardian-context/src/preparation/mod.rs index 66dfb4f80d..0806099f09 100644 --- a/codex-rs/guardian-context/src/preparation/mod.rs +++ b/codex-rs/guardian-context/src/preparation/mod.rs @@ -2,6 +2,8 @@ //! Collection keeps captured sources intact. This owner applies text and image policy; //! delivery checks repeat against current reviewer history after compaction. +pub(crate) mod root; + use crate::CollectedContext; use crate::ContextSection; use crate::GuardianRootMessage; diff --git a/codex-rs/guardian-context/src/preparation/root.rs b/codex-rs/guardian-context/src/preparation/root.rs new file mode 100644 index 0000000000..1114cfa769 --- /dev/null +++ b/codex-rs/guardian-context/src/preparation/root.rs @@ -0,0 +1,218 @@ +//! Fixed root request selection, including legacy source matching and omission notices. +//! Collection and persistent storage keep their existing bounds and source ownership. +//! Assistant candidates borrow one immutable host snapshot until synchronous preparation ends. + +use crate::GuardianRootMessage; +use codex_history::RetainedContextOrder; +use codex_history::RetainedUserMessage; +use codex_protocol::models::ContentItem; +use std::collections::BTreeSet; +use std::collections::HashMap; + +#[derive(Clone, PartialEq)] +pub struct RootMessageCandidate { + pub order: Option, + pub message: GuardianRootMessage, +} + +#[derive(Clone, PartialEq)] +pub enum RootAssistantContent<'a> { + Text(&'a [ContentItem]), + DeliveredCall { call_id: String }, +} + +#[derive(Clone, PartialEq)] +pub struct RootAssistantCandidate<'a> { + pub message_id: Option, + pub order: Option, + pub content: RootAssistantContent<'a>, +} + +#[derive(Clone, PartialEq)] +pub struct RootRetainedAssistant<'a> { + pub order: RetainedContextOrder, + pub message: &'a RetainedUserMessage, + /// Host checks exclude commentary from standalone context, before request selection. + pub standalone_eligible: bool, +} + +#[derive(Clone, PartialEq)] +pub struct RootDeliveredMessage<'a> { + pub call_id: &'a str, + pub text: &'a str, +} + +const ROOT_CONTEXT_WINDOW: usize = 3; + +/// Preserves the current handoff windows, latest context and all user restrictions. +pub fn select_root_handoff_indices( + messages: &[(Option, GuardianRootMessage)], + boundaries: &[RetainedContextOrder], +) -> Option> { + let mut selected = BTreeSet::new(); + for boundary in boundaries { + selected.extend( + messages + .iter() + .enumerate() + .rev() + .filter_map(|(index, (order, message))| { + if matches!(message, GuardianRootMessage::UserInput(_)) { + return None; + } + order.filter(|order| order <= boundary).map(|_| index) + }) + .take(ROOT_CONTEXT_WINDOW), + ); + } + if !selected.is_empty() { + // Include the latest ordinary messages even without a new handoff to this branch. + selected.extend( + messages + .iter() + .enumerate() + .rev() + .filter_map(|(index, (order, message))| { + (order.is_some() && !matches!(message, GuardianRootMessage::UserInput(_))) + .then_some(index) + }) + .take(ROOT_CONTEXT_WINDOW), + ); + // Later user inputs can restrict an earlier handoff even when they are not recent. + // Leave all user inputs, verified answers, and unordered legacy evidence to the shared cap. + selected.extend( + messages + .iter() + .enumerate() + .filter_map(|(index, (order, message))| { + (order.is_none() + || matches!( + message, + GuardianRootMessage::User(_) | GuardianRootMessage::UserInput(_) + )) + .then_some(index) + }), + ); + } + (!selected.is_empty()).then_some(selected) +} + +const MAX_ROOT_MESSAGE_TOKENS: usize = 900; + +/// Bounds root user text and admits verified answers only as whole records. +pub fn prepare_root_user_messages( + input: &[crate::RootMessageCandidate], +) -> Vec { + input + .iter() + .filter_map(|candidate| { + if let GuardianRootMessage::UserInput(text) = &candidate.message + && !crate::verified_answers::answer_fits(text.len()) + { + return None; + } + let message = if let GuardianRootMessage::User(text) = &candidate.message { + GuardianRootMessage::User(crate::truncate_text(text, MAX_ROOT_MESSAGE_TOKENS)) + } else { + candidate.message.clone() + }; + Some(crate::RootMessageCandidate { + order: candidate.order, + message, + }) + }) + .collect() +} + +const MAX_ROOT_MESSAGES: usize = 16; + +/// Root assistant evidence after delivery lookup, original admission and live fallback. +pub struct PreparedRootAssistants { + pub messages: Vec, + pub missing_context: bool, +} + +struct LiveRootMessage { + source_id: Option, + candidate: RootMessageCandidate, +} + +/// Prepares assistant evidence without changing which sources were captured. +pub fn prepare_root_assistant_messages( + assistant_messages: &[crate::RootAssistantCandidate<'_>], + retained_assistants: &[crate::RootRetainedAssistant<'_>], + delivered_messages: &[crate::RootDeliveredMessage<'_>], + mut missing_assistant_context: bool, +) -> PreparedRootAssistants { + let delivered = delivered_messages + .iter() + .take(MAX_ROOT_MESSAGES) + .map(|message| (message.call_id, message.text)) + .collect::>(); + let live = assistant_messages + .iter() + .filter_map(|source| { + let text = match &source.content { + RootAssistantContent::Text(content) => content + .iter() + .filter_map(|item| match item { + ContentItem::InputText { text } | ContentItem::OutputText { text, .. } => { + Some(text.as_str()) + } + _ => None, + }) + .collect::(), + RootAssistantContent::DeliveredCall { call_id } => { + delivered.get(call_id.as_str())?.to_string() + } + }; + let retained = retained_assistants.iter().find(|retained| { + source + .message_id + .as_ref() + .is_some_and(|id| retained.message.message_id.as_ref() == Some(id)) + }); + let order = retained.map(|retained| retained.order).or(source.order); + let text = retained + .and_then(|retained| root_retained_assistant(retained.message)) + .unwrap_or_else(|| crate::truncate_text(&text, MAX_ROOT_MESSAGE_TOKENS)); + let message = if order.is_some() { + GuardianRootMessage::Assistant(text) + } else { + GuardianRootMessage::UnorderedAssistant(text) + }; + Some(LiveRootMessage { + source_id: source.message_id.clone(), + candidate: RootMessageCandidate { order, message }, + }) + }) + .collect::>(); + let mut messages = retained_assistants + .iter() + .filter_map(|source| { + if !source.standalone_eligible + || source.message.message_id.as_ref().is_some_and(|id| { + live.iter() + .any(|message| message.source_id.as_ref() == Some(id)) + }) + { + return None; + } + let text = root_retained_assistant(source.message); + missing_assistant_context |= text.is_none(); + text.map(|text| RootMessageCandidate { + order: Some(source.order), + message: GuardianRootMessage::Assistant(text), + }) + }) + .collect::>(); + messages.extend(live.into_iter().map(|message| message.candidate)); + PreparedRootAssistants { + messages, + missing_context: missing_assistant_context, + } +} + +fn root_retained_assistant(message: &codex_history::RetainedUserMessage) -> Option { + crate::retained_assistant_message(message).map(|_| message.text.clone()) +} diff --git a/codex-rs/guardian-context/src/verified_answers.rs b/codex-rs/guardian-context/src/verified_answers.rs index e8d10304c4..d7aa15815b 100644 --- a/codex-rs/guardian-context/src/verified_answers.rs +++ b/codex-rs/guardian-context/src/verified_answers.rs @@ -44,11 +44,12 @@ pub fn render_verified_answers(context: &RetainedContext) -> RenderedVerifiedAns /// Selects one complete response, preserving both sides of every question/answer pair. /// Returning None means the caller must account for missing evidence, not partial permission. pub fn render_verified_answer(answer: &VerifiedAnswer) -> Option { - let text = answer_text(answer); + let text = render_original_verified_answer(answer); answer_fits(text.len()).then_some(text) } -fn answer_text(answer: &VerifiedAnswer) -> String { +/// Renders original question/answer bodies without making a request-selection decision. +pub fn render_original_verified_answer(answer: &VerifiedAnswer) -> String { answer .questions .iter() @@ -105,7 +106,7 @@ pub fn collect_verified_answer_candidates(context: &RetainedContext) -> Vec