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
This commit is contained in:
ioanbudea-OAI
2026-10-27 00:01:05 +00:00
committed by copyberry
parent e977cb394d
commit 1f281290e3
6 changed files with 319 additions and 155 deletions
+14 -46
View File
@@ -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<RetainedContextOrder>, GuardianRootMessage)],
worker: &AgentPath,
tool_namespace: Option<&str>,
) -> Option<BTreeSet<usize>> {
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<RetainedContextOrder> {
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 {
@@ -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::<HashMap<_, _>>();
.collect::<Vec<_>>();
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("<user_action>"))
.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::<Vec<_>>();
let mut messages = codex_guardian_context::prepare_root_user_messages(&user_candidates)
.into_iter()
.map(|message| (message.order, message.message))
.collect::<Vec<_>>();
// 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::<Vec<_>>();
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::<Vec<_>>();
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<String> {
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::<String>()
}
};
Some(guardian_truncate_text(&text, GUARDIAN_MAX_ROOT_MESSAGE_TOKENS).0)
}
+11 -1
View File
@@ -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>,
@@ -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;
@@ -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<RetainedContextOrder>,
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<String>,
pub order: Option<RetainedContextOrder>,
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<RetainedContextOrder>, GuardianRootMessage)],
boundaries: &[RetainedContextOrder],
) -> Option<BTreeSet<usize>> {
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<crate::RootMessageCandidate> {
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<crate::RootMessageCandidate>,
pub missing_context: bool,
}
struct LiveRootMessage {
source_id: Option<String>,
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::<HashMap<_, _>>();
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::<String>(),
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::<Vec<_>>();
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::<Vec<_>>();
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<String> {
crate::retained_assistant_message(message).map(|_| message.text.clone())
}
@@ -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<String> {
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<Trus
}
for (order, entry) in source_order_labels(context) {
if let RetainedContextEntry::VerifiedAnswer(answer) = entry {
let text = answer_text(answer);
let text = render_original_verified_answer(answer);
items.push(TrustedUserAnswer {
answer_bytes: Some(text.len()),
text: format!("Retained source order: {order}\n{text}"),