Centralize Guardian root context preparation (#57425)

## Why

Verified answers omitted because of request-size limits must be marked as missing. Legacy checkpoint messages also need to match retained excerpts after truncation so workers receive only one copy.

## What changed

Move root message preparation into `codex-guardian-context::prepare_root_context`, including handoff selection, shared message limits, legacy matching, and incomplete-context markers. Pass original legacy text into preparation and mark verified answers incomplete when oversized records are excluded.

## Testing

Add a handoff regression scenario for an ID-less checkpoint with a retained excerpt and an oversized verified answer. Assert that the worker receives one capped user message, both incomplete-evidence markers, and an incomplete authorization version.

GitOrigin-RevId: bd3f5fe164ae2fddba440e9269fd63b52e8b12f8
This commit is contained in:
ioanbudea-OAI
2026-10-27 10:29:14 +00:00
committed by copyberry
parent 124e4e5bfa
commit 5c6ef3d1a6
6 changed files with 278 additions and 183 deletions
@@ -2,9 +2,6 @@
//! History position approximates dispatch order; absent calls leave the existing context unchanged.
//! User inputs bypass relevance filtering; only assistant context is narrowed.
use std::collections::BTreeSet;
use crate::codex_thread::GuardianRootMessage;
use codex_history::ResponseItemEnvelope;
use codex_history::RetainedContextOrder;
use codex_protocol::AgentPath;
@@ -12,18 +9,6 @@ use codex_protocol::ToolName;
use codex_protocol::models::ExecutedToolCallArguments;
use codex_protocol::models::ResponseItem;
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,
@@ -10,7 +10,6 @@
//! Handoff projection preserves saved heartbeat instructions and the active turn's skills.
use std::borrow::Cow;
use std::collections::VecDeque;
use std::collections::hash_map::DefaultHasher;
use std::hash::Hash;
use std::hash::Hasher;
@@ -26,6 +25,7 @@ use crate::context::is_contextual_user_fragment;
use crate::event_mapping::parse_turn_item;
use codex_guardian_context::RootAssistantCandidate;
use codex_guardian_context::RootAssistantContent;
use codex_guardian_context::RootConversationCandidates;
use codex_guardian_context::RootDeliveredMessage;
use codex_guardian_context::RootMessageCandidate;
use codex_guardian_context::RootRetainedAssistant;
@@ -41,8 +41,6 @@ use codex_protocol::models::MessagePhase;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::MultiAgentVersion;
const MAX_ROOT_MESSAGES: usize = 16;
impl LocalAgentControl {
/// Returns bounded root conversation and authorization state for a MultiAgent V2 worker.
pub(crate) async fn root_user_authorization(
@@ -189,27 +187,15 @@ impl LocalAgentControl {
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();
let legacy_capacity = MAX_ROOT_MESSAGES.saturating_sub(messages.len());
if missing_root_instructions && legacy_capacity > 0 {
let mut latest_legacy_turn_id = None;
for message in reconciled.unmatched_user_messages(root_history.legacy_user_messages()) {
latest_legacy_turn_id = (!message.turn_id.is_empty()).then_some(message.turn_id);
legacy_messages.push_back(GuardianRootMessage::User(message.text));
if legacy_messages.len() > legacy_capacity {
legacy_messages.pop_front();
}
}
if latest_user_turn_id.is_none() {
latest_user_turn_id = latest_legacy_turn_id;
}
}
let original_missing_root_instructions = missing_root_instructions;
let retained_user_sources = reconciled
.ordered_entries()
.filter_map(|(_, entry)| match entry {
RetainedContextEntry::UserMessage(message) => Some(message.clone()),
RetainedContextEntry::AssistantMessage(_)
| RetainedContextEntry::VerifiedAnswer(_) => None,
})
.collect();
let assistant_candidates = root_history
.annotated_items()
.iter()
@@ -272,26 +258,9 @@ impl LocalAgentControl {
})
})
.collect::<Vec<_>>();
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)),
);
// 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();
if let Some(worker) = worker_path
&& let Some(selected) = super::root_handoff::selected_message_indices(
let handoff_boundaries = if let Some(worker) = worker_path {
super::root_handoff::handoff_boundaries(
root_history.annotated_items(),
&messages,
&worker,
root_thread
.config()
@@ -300,80 +269,37 @@ impl LocalAgentControl {
.tool_namespace
.as_deref(),
)
{
// A handoff can change the selection without adding a new root user message.
} else {
Vec::new()
};
let request_candidates = RootConversationCandidates {
user_messages: user_candidates,
assistant_messages: assistant_candidates,
retained_assistants: retained_assistant_candidates,
delivered_messages: delivered_candidates,
retained_user_sources,
handoff_boundaries,
missing_root_instructions: original_missing_root_instructions,
missing_assistant_context: retained_context.has_omitted_assistant_messages(),
missing_verified_answers: !retained_context.verified_answers_complete(),
};
let prepared = codex_guardian_context::prepare_root_context(
&request_candidates,
root_history.original_legacy_user_messages(),
);
let mut authorization_version = root_evidence.authorization_version(history.as_ref());
if prepared.missing_root_instructions {
authorization_version.retained_context_complete = false;
}
let mut review_context_revision = history.guardian_review_context_revision();
if let Some(selected) = &prepared.handoff_selection {
let mut revision = DefaultHasher::new();
review_context_revision.hash(&mut revision);
selected.hash(&mut revision);
review_context_revision = revision.finish();
missing_assistant_context |=
messages.iter().enumerate().any(|(index, (_, message))| {
!selected.contains(&index)
&& matches!(
message,
GuardianRootMessage::Assistant(_)
| GuardianRootMessage::UnorderedAssistant(_)
)
});
messages = messages
.into_iter()
.enumerate()
.filter_map(|(index, message)| selected.contains(&index).then_some(message))
.collect();
}
// Apply one shared cap so user records and confirmed questions compete
// while keeping omissions explicit for authorization and assistant context.
let removed = messages.len().saturating_sub(MAX_ROOT_MESSAGES);
missing_root_instructions |= messages[..removed].iter().any(|(_, message)| {
matches!(
message,
GuardianRootMessage::User(_) | GuardianRootMessage::UserInput(_)
)
});
missing_assistant_context |= messages[..removed].iter().any(|(_, message)| {
matches!(
message,
GuardianRootMessage::Assistant(_) | GuardianRootMessage::UnorderedAssistant(_)
)
});
messages.drain(..removed);
// Keep the existing presentation of unordered evidence after ordered records.
messages.sort_by_key(|(order, _)| (order.is_none(), *order));
legacy_messages.drain(
..legacy_messages
.len()
.saturating_sub(MAX_ROOT_MESSAGES - messages.len()),
);
let mut messages = messages
.into_iter()
.map(|(_, message)| message)
.collect::<Vec<_>>();
if missing_assistant_context {
messages.insert(
/*index*/ 0,
GuardianRootMessage::IncompleteAssistantContext,
);
}
let mut authorization_version = root_evidence.authorization_version(history.as_ref());
let missing_verified_answers = !authorization_version.retained_context_complete;
if missing_verified_answers {
messages.insert(
/*index*/ 0,
GuardianRootMessage::IncompleteVerifiedAnswers,
);
}
if missing_root_instructions {
authorization_version.retained_context_complete = false;
messages.insert(
/*index*/ 0,
GuardianRootMessage::IncompleteRootInstructions,
);
}
messages.insert(/*index*/ 0, GuardianRootMessage::RetainedContextScope);
if !legacy_messages.is_empty() {
legacy_messages.push_front(GuardianRootMessage::LegacyContextScope);
legacy_messages.extend(messages);
messages = legacy_messages.into_iter().collect();
if latest_user_turn_id.is_none() {
latest_user_turn_id = prepared.latest_legacy_turn_id;
}
let trusted_skill_paths = latest_user_turn_id
.as_deref()
@@ -384,7 +310,7 @@ impl LocalAgentControl {
history_reset_version: root_history.reset_version,
authorization_version,
review_context_revision,
messages,
messages: prepared.messages,
trusted_skill_paths,
})
}
@@ -37,6 +37,17 @@ impl ContextManager {
/// Bounded instruction candidates from the legacy backup; their ordering is unknown.
pub(crate) fn legacy_user_messages(&self) -> impl Iterator<Item = RetainedUserMessage> + '_ {
self.original_legacy_user_messages().map(|mut message| {
message.text =
guardian_truncate_text(&message.text, GUARDIAN_MAX_ROOT_MESSAGE_TOKENS).0;
message
})
}
/// Original legacy candidates for request preparation; storage remains independently bounded.
pub(crate) fn original_legacy_user_messages(
&self,
) -> impl Iterator<Item = RetainedUserMessage> + '_ {
self.guardian_history_items()
.into_iter()
.flatten()
@@ -54,7 +65,6 @@ impl ContextManager {
{
return None;
}
let text = guardian_truncate_text(&text, GUARDIAN_MAX_ROOT_MESSAGE_TOKENS).0;
Some(RetainedUserMessage {
phase: None,
origin: codex_history::UserInputOrigin::from_message(item),
@@ -290,6 +290,41 @@ pub(crate) async fn handoff_scenario() -> Result<Vec<ResponsesRequest>> {
);
assert_eq!(root_messages(&leaf).await, expected);
// An old ID-less checkpoint stores an excerpt beside the original legacy backup.
// The worker must receive one capped copy and preserve both evidence gaps.
let original = "Never publish project details. ".repeat(200);
let excerpt = codex_guardian_context::truncate_text(&original, /*max_tokens*/ 900);
let history = serde_json::from_value(json!([{"type":"compacted", "payload": {
"message":"Old ID-less checkpoint.", "replacement_history":[],
"guardian_history":[{"type":"message", "role":"user",
"content":[{"type":"input_text", "text":original}]}],
"retained_context":{
"user_messages":[{"turn_id":"", "text":excerpt, "complete":false}],
"user_messages_incomplete":true, "incomplete":false,
"verified_answers":[{"turn_id":"large", "call_id":"large", "questions":[{
"question":"Publish?", "answer":"Only if conditions hold. ".repeat(200)
}]}]
}
}}]))?;
test.codex =
crate::suite::guardian_checkpoint_migration::resume(&test, &test.codex, history).await?;
let snapshot = leaf.guardian_root_snapshot().await.expect("migrated root");
assert_eq!(
(
snapshot.messages,
snapshot.authorization_version.retained_context_complete
),
(
vec![
GuardianRootMessage::RetainedContextScope,
GuardianRootMessage::IncompleteRootInstructions,
GuardianRootMessage::IncompleteVerifiedAnswers,
GuardianRootMessage::User(excerpt),
],
false
)
);
let shutdown = test
.thread_manager
.shutdown_all_threads_bounded(Duration::from_secs(/*secs*/ 10))
+3 -4
View File
@@ -30,6 +30,7 @@ pub use action::action_for_review;
pub use authorization::GuardianRootMessage;
pub use preparation::root::RootAssistantCandidate;
pub use preparation::root::RootAssistantContent;
pub use preparation::root::RootConversationCandidates;
pub use preparation::root::RootDeliveredMessage;
pub use preparation::root::RootMessageCandidate;
pub use preparation::root::RootRetainedAssistant;
@@ -92,10 +93,8 @@ 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 preparation::root::PreparedRootContext;
pub use preparation::root::prepare_root_context;
pub use profile::ContextProfile;
mod authorization;
mod entry;
+190 -50
View File
@@ -8,6 +8,7 @@ use codex_history::RetainedUserMessage;
use codex_protocol::models::ContentItem;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::collections::VecDeque;
#[derive(Clone, PartialEq)]
pub struct RootMessageCandidate {
@@ -42,11 +43,33 @@ pub struct RootDeliveredMessage<'a> {
pub text: &'a str,
}
/// Unselected evidence from the same snapshot as the root authorization metadata.
#[derive(Clone, PartialEq)]
pub struct RootConversationCandidates<'a> {
pub user_messages: Vec<RootMessageCandidate>,
pub assistant_messages: Vec<RootAssistantCandidate<'a>>,
pub retained_assistants: Vec<RootRetainedAssistant<'a>>,
/// Newest first, preserving host-confirmed delivery lookup order.
pub delivered_messages: Vec<RootDeliveredMessage<'a>>,
/// Retained identities used for legacy matching after request-size preparation.
pub retained_user_sources: Vec<RetainedUserMessage>,
pub handoff_boundaries: Vec<RetainedContextOrder>,
pub missing_root_instructions: bool,
pub missing_assistant_context: bool,
pub missing_verified_answers: bool,
}
impl std::fmt::Debug for RootConversationCandidates<'_> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("RootConversationCandidates(..)")
}
}
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)],
fn select_root_handoff_indices(
messages: &[RootMessageCandidate],
boundaries: &[RetainedContextOrder],
) -> Option<BTreeSet<usize>> {
let mut selected = BTreeSet::new();
@@ -56,11 +79,14 @@ pub fn select_root_handoff_indices(
.iter()
.enumerate()
.rev()
.filter_map(|(index, (order, message))| {
if matches!(message, GuardianRootMessage::UserInput(_)) {
.filter_map(|(index, message)| {
if matches!(message.message, GuardianRootMessage::UserInput(_)) {
return None;
}
order.filter(|order| order <= boundary).map(|_| index)
message
.order
.filter(|order| order <= boundary)
.map(|_| index)
})
.take(ROOT_CONTEXT_WINDOW),
);
@@ -72,64 +98,35 @@ pub fn select_root_handoff_indices(
.iter()
.enumerate()
.rev()
.filter_map(|(index, (order, message))| {
(order.is_some() && !matches!(message, GuardianRootMessage::UserInput(_)))
.then_some(index)
.filter_map(|(index, message)| {
(message.order.is_some()
&& !matches!(message.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.extend(messages.iter().enumerate().filter_map(|(index, message)| {
(message.order.is_none()
|| matches!(
message.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 PreparedRootAssistants {
messages: Vec<crate::RootMessageCandidate>,
missing_context: bool,
}
struct LiveRootMessage {
@@ -138,7 +135,7 @@ struct LiveRootMessage {
}
/// Prepares assistant evidence without changing which sources were captured.
pub fn prepare_root_assistant_messages(
fn prepare_root_assistant_messages(
assistant_messages: &[crate::RootAssistantCandidate<'_>],
retained_assistants: &[crate::RootRetainedAssistant<'_>],
delivered_messages: &[crate::RootDeliveredMessage<'_>],
@@ -216,3 +213,146 @@ pub fn prepare_root_assistant_messages(
fn root_retained_assistant(message: &codex_history::RetainedUserMessage) -> Option<String> {
crate::retained_assistant_message(message).map(|_| message.text.clone())
}
/// Bounded root view plus the existing cache and skill-scope selection metadata.
pub struct PreparedRootContext {
pub messages: Vec<GuardianRootMessage>,
pub missing_root_instructions: bool,
pub handoff_selection: Option<BTreeSet<usize>>,
pub latest_legacy_turn_id: Option<String>,
}
/// Applies the fixed root request policy without discarding the captured sources.
pub fn prepare_root_context(
input: &RootConversationCandidates<'_>,
legacy_candidates: impl IntoIterator<Item = RetainedUserMessage>,
) -> PreparedRootContext {
let mut missing_root_instructions = input.missing_root_instructions;
let mut missing_verified_answers = input.missing_verified_answers;
// Verified answers stay whole; ordinary root user text keeps its existing cap.
let mut messages = input
.user_messages
.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::<Vec<_>>();
missing_verified_answers |= messages.len() != input.user_messages.len();
let legacy_capacity = MAX_ROOT_MESSAGES.saturating_sub(messages.len());
let mut legacy_messages = VecDeque::new();
let mut latest_legacy_turn_id = None;
if missing_root_instructions && legacy_capacity > 0 {
for candidate in legacy_candidates {
let text = crate::truncate_text(&candidate.text, MAX_ROOT_MESSAGE_TOKENS);
// Legacy records without IDs historically match after the request-size cap.
if input.retained_user_sources.iter().any(|source| {
if let Some(id) = &source.message_id {
candidate.message_id.as_ref() == Some(id)
} else {
candidate.turn_id == source.turn_id && text == source.text
}
}) {
continue;
}
latest_legacy_turn_id =
(!candidate.turn_id.is_empty()).then(|| candidate.turn_id.clone());
legacy_messages.push_back(GuardianRootMessage::User(text));
if legacy_messages.len() > legacy_capacity {
legacy_messages.pop_front();
}
}
}
let assistants = prepare_root_assistant_messages(
&input.assistant_messages,
&input.retained_assistants,
&input.delivered_messages,
input.missing_assistant_context,
);
let mut missing_assistant_context = assistants.missing_context;
messages.extend(assistants.messages);
messages.sort_by_key(|message| message.order);
let handoff_selection = select_root_handoff_indices(&messages, &input.handoff_boundaries);
if let Some(selected) = &handoff_selection {
messages = messages
.into_iter()
.enumerate()
.filter_map(|(index, message)| {
if selected.contains(&index) {
Some(message)
} else {
missing_assistant_context |= matches!(
message.message,
GuardianRootMessage::Assistant(_)
| GuardianRootMessage::UnorderedAssistant(_)
);
None
}
})
.collect();
}
let removed = messages.len().saturating_sub(MAX_ROOT_MESSAGES);
for message in &messages[..removed] {
missing_root_instructions |= matches!(
message.message,
GuardianRootMessage::User(_) | GuardianRootMessage::UserInput(_)
);
missing_assistant_context |= matches!(
message.message,
GuardianRootMessage::Assistant(_) | GuardianRootMessage::UnorderedAssistant(_)
);
}
messages.drain(..removed);
messages.sort_by_key(|message| (message.order.is_none(), message.order));
legacy_messages.drain(
..legacy_messages
.len()
.saturating_sub(MAX_ROOT_MESSAGES - messages.len()),
);
let mut messages = messages
.into_iter()
.map(|message| message.message)
.collect::<Vec<_>>();
if missing_assistant_context {
messages.insert(
/*index*/ 0,
GuardianRootMessage::IncompleteAssistantContext,
);
}
if missing_verified_answers {
messages.insert(
/*index*/ 0,
GuardianRootMessage::IncompleteVerifiedAnswers,
);
}
if missing_root_instructions {
messages.insert(
/*index*/ 0,
GuardianRootMessage::IncompleteRootInstructions,
);
}
messages.insert(/*index*/ 0, GuardianRootMessage::RetainedContextScope);
if !legacy_messages.is_empty() {
legacy_messages.push_front(GuardianRootMessage::LegacyContextScope);
legacy_messages.extend(messages);
messages = legacy_messages.into_iter().collect();
}
PreparedRootContext {
messages,
missing_root_instructions,
handoff_selection,
latest_legacy_turn_id,
}
}