openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-31 19:28:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add elastic budgeting for Guardian context preparation (#58454)

Browse Source 
## What changed

Introduce `ContextBudgeting::Elastic` to prepare captured root conversation and Node REPL evidence without the fixed per-section request caps. Elastic preparation preserves message text and candidates for later allocation while retaining storage omission notices and legacy source deduplication.

Keep `ContextBudgeting::Fixed` as the default and use it at existing call sites.

## Testing

Add tests covering elastic root message preservation and legacy source deduplication, plus REPL text, images, and omission counts in text-only and multimodal modes.

GitOrigin-RevId: ce4fae9d725694456525a413fb511166df6f5e2c
This commit is contained in:
ioanbudea-OAI
2026-10-30 23:54:38 +00:00
committed by copyberry
parent a795f961be
commit 26cc286314
11 changed files with 227 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/agent/control/user_authorization.rs
+6 -2
View File
@@ -47,7 +47,10 @@ struct CapturedRootConversation {
}
impl RootConversationSource for CapturedRootConversation {
fn prepare(&self) -> codex_guardian_context::PreparedRootContext {
fn prepare(
&self,
budgeting: codex_guardian_context::ContextBudgeting,
) -> codex_guardian_context::PreparedRootContext {
let root_history = &self.history;
// Join calls to host-confirmed outputs in this snapshot. Older outputs without
// captured text cannot establish what was sent, including after a hook rewrite.
@@ -245,6 +248,7 @@ impl RootConversationSource for CapturedRootConversation {
codex_guardian_context::prepare_root_context(
&request_candidates,
root_history.original_legacy_user_messages(),
budgeting,
)
}
}
@@ -298,7 +302,7 @@ impl LocalAgentControl {
history: root_history,
handoff_boundaries,
});
let prepared = request_candidates.prepare();
let prepared = request_candidates.prepare(codex_guardian_context::ContextBudgeting::Fixed);
let root_evidence = root_thread
.session
.services
codex-rs/core/src/guardian/input_budget.rs
+6 -4
View File
@@ -31,7 +31,9 @@ use crate::session::turn::build_prompt;
use crate::session::turn_context::TurnContext;
#[derive(Clone)]
pub(crate) struct PendingReviewContext(pub ComposedContext);
pub(crate) struct PendingReviewContext {
pub context: ComposedContext,
}
pub(super) use codex_guardian_reviewer::CheckpointRecovery;
@@ -76,7 +78,7 @@ pub(crate) async fn check_pending(session: &Session, turn: &TurnContext) -> Code
codex_protocol::protocol::TruncationPolicy::Bytes(base.text.len()).token_budget();
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();
let mut context = pending.context.clone();
let history = session
.clone_history()
.await
@@ -129,7 +131,7 @@ pub(crate) async fn finalize(
"Guardian expects one review input".to_owned(),
));
};
let mut context = pending.0.clone();
let mut context = pending.context.clone();
let model = &step.settings.model_info;
let history = session.clone_history().await;
let recovery = session
@@ -361,7 +363,7 @@ pub(crate) async fn finalize(
}
}
step.turn.extension_data.insert(RetainedReviewContext {
context: pending.0.retained_instructions(),
context: pending.context.retained_instructions(),
history_version,
});
session
codex-rs/core/src/guardian/input_budget_tests.rs
+4 -4
View File
@@ -61,7 +61,7 @@ async fn finalization_exports_loss_totals_only_for_accepted_input(window: i64) {
session
.services
.thread_extension_data
.insert(PendingReviewContext(context));
.insert(PendingReviewContext { context });
assert_eq!(
finalize(&session, &step, &mut input, HistoryTruncation::Preserve)
.await
@@ -177,7 +177,7 @@ async fn cancelled_startup_does_not_record_unselected_review_evidence() {
session
.services
.thread_extension_data
.insert(PendingReviewContext(context));
.insert(PendingReviewContext { context });
session
.set_session_startup_prewarm(
crate::session::startup_prewarm::SessionStartupPrewarmHandle::new(
@@ -234,7 +234,7 @@ async fn feasibility_rejects_required_evidence_inside_the_reserved_margin() {
session
.services
.thread_extension_data
.insert(PendingReviewContext(context));
.insert(PendingReviewContext { context });
assert!(check_pending(&session, &turn).await.is_err());
}
@@ -257,7 +257,7 @@ async fn finalization_overflow_marks_the_reviewer_exhausted() {
session
.services
.thread_extension_data
.insert(PendingReviewContext(context));
.insert(PendingReviewContext { context });
assert!(
finalize(&session, &step, &mut input, HistoryTruncation::Preserve)
codex-rs/core/src/guardian/review_session.rs
+3 -3
View File
@@ -555,9 +555,9 @@ async fn run_review_on_session(
.session
.services
.thread_extension_data
.insert(super::input_budget::PendingReviewContext(
prompt_items.context,
));
.insert(super::input_budget::PendingReviewContext {
context: prompt_items.context,
});
review_session
.session
.services
codex-rs/ext/guardian-reviewer/src/node_repl_review_evidence.rs
+1
View File
@@ -264,6 +264,7 @@ pub struct NodeReplReviewEvidenceSnapshot {
impl NodeReplReviewEvidenceSnapshot {
pub fn context(&self, mode: NodeReplReviewEvidenceMode) -> NodeReplContext<'_> {
NodeReplContext {
budgeting: codex_guardian_context::ContextBudgeting::Fixed,
responses: &self.responses,
omitted_responses: self.omitted_responses,
mode,
codex-rs/guardian-context/src/lib.rs
+1
View File
@@ -104,6 +104,7 @@ pub use composition::ComposedContext;
pub use composition::ContextPresentation;
pub use composition::MAX_TEXT_PART_ESTIMATED_TOKENS;
pub use composition::PreparedTranscript;
pub use preparation::ContextBudgeting;
pub use preparation::ContextProfile;
pub use preparation::SyncImagePreparer;
pub use preparation::SyncPreparation;
codex-rs/guardian-context/src/preparation/mod.rs
+16 -1
View File
@@ -32,6 +32,20 @@ pub(crate) const MAX_INSTRUCTION_TOKENS: usize = 900;
const MAX_APPROVAL_REASON_TOKENS: usize = 512;
const MAX_DESCRIPTION_TOKENS: usize = 400;
/// Request selection policy; persistent capture and storage bounds are independent.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum ContextBudgeting {
/// Preserve existing per-section request selection.
#[default]
Fixed,
/// Keep captured candidates complete until optional entry capping and allocation.
Elastic,
}
#[cfg(test)]
#[path = "source_tests.rs"]
mod source_tests;
/// Prepares the host-selected transcript suffix and all synchronous context sections.
/// Cursor selection remains with the host; its offset addresses the collected entries.
/// No storage, reviewer history or cursor state is changed by this operation.
@@ -159,7 +173,7 @@ pub(crate) fn prepare_sync_sections(context: &mut CollectedContext) {
if let crate::ContextSection::RootConversation { content } = section
&& let crate::RootConversationContent::Collected(candidates) = content
{
let prepared = candidates.prepare();
let prepared = candidates.prepare(ContextBudgeting::Fixed);
*content = crate::RootConversationContent::Prepared(
crate::authorization::root_conversation_items(prepared.messages),
);
@@ -178,6 +192,7 @@ pub(crate) fn prepare_sync_sections(context: &mut CollectedContext) {
responses,
omitted_responses: *omitted_responses,
mode: *mode,
budgeting: ContextBudgeting::Fixed,
};
let items = snapshot.render_inputs();
evidence.content = crate::node_repl::NodeReplEvidenceContent::Prepared(items);
codex-rs/guardian-context/src/preparation/repl.rs
+14 -1
View File
@@ -51,6 +51,8 @@ pub struct NodeReplContext<'a> {
pub responses: &'a [Arc<NodeReplResponse>],
pub omitted_responses: u64,
pub mode: NodeReplReviewEvidenceMode,
/// Controls request selection only; host storage and delivery cursors stay unchanged.
pub budgeting: crate::ContextBudgeting,
}
/// Request-local snapshot, selected and rendered only during sync preparation.
@@ -153,6 +155,7 @@ impl NodeReplContext<'_> {
intro.len().saturating_add(closing.len()),
NodeReplLayout::Multimodal,
self.omitted_responses,
self.budgeting,
);
let omitted_responses = selected.omitted_responses;
@@ -245,6 +248,7 @@ impl ContextualUserFragment for NodeReplContext<'_> {
.saturating_add(body.len()),
NodeReplLayout::Text,
self.omitted_responses,
self.budgeting,
);
let omitted_responses = selected.omitted_responses;
@@ -274,7 +278,9 @@ impl ContextualUserFragment for NodeReplContext<'_> {
body.push('\n');
}
}
debug_assert!(body.len() <= max_body_bytes);
debug_assert!(
self.budgeting == crate::ContextBudgeting::Elastic || body.len() <= max_body_bytes
);
input_text(body)
}
}
@@ -298,7 +304,14 @@ fn select_node_repl_responses(
framing_bytes: usize,
layout: NodeReplLayout,
mut omitted_responses: u64,
budgeting: crate::ContextBudgeting,
) -> NodeReplSelection {
if budgeting == crate::ContextBudgeting::Elastic {
return NodeReplSelection {
first: 0,
omitted_responses,
};
}
let notice_bytes = match layout {
NodeReplLayout::Text => 64,
NodeReplLayout::Multimodal => 160,
codex-rs/guardian-context/src/preparation/root.rs
+64 -25
View File
@@ -1,7 +1,9 @@
//! Fixed root request selection, including legacy source matching and omission notices.
//! Root request preparation, including legacy source matching and omission notices.
//! Fixed mode keeps its existing caps; elastic mode leaves size selection to admission.
//! Collection and persistent storage keep their existing bounds and source ownership.
//! Assistant candidates borrow one immutable host snapshot until synchronous preparation ends.
use crate::ContextBudgeting;
use crate::GuardianRootMessage;
use codex_history::RetainedContextOrder;
use codex_history::RetainedUserMessage;
@@ -69,7 +71,8 @@ impl std::fmt::Debug for RootConversationCandidates<'_> {
/// An immutable host snapshot that lends sources only during synchronous preparation.
/// Implementations must preserve capture-time provenance and avoid reading live session state.
pub trait RootConversationSource: Send + Sync {
fn prepare(&self) -> PreparedRootContext;
/// Prepares one request view without changing captured sources or storage policy.
fn prepare(&self, budgeting: ContextBudgeting) -> PreparedRootContext;
}
/// Shared ownership keeps captured source payloads alive without copying them.
@@ -82,8 +85,8 @@ impl RootConversationInput {
Self(std::sync::Arc::new(source))
}
pub fn prepare(&self) -> PreparedRootContext {
self.0.prepare()
pub fn prepare(&self, budgeting: ContextBudgeting) -> PreparedRootContext {
self.0.prepare(budgeting)
}
}
@@ -174,11 +177,15 @@ fn prepare_root_assistant_messages(
retained_assistants: &[crate::RootRetainedAssistant<'_>],
delivered_messages: &[crate::RootDeliveredMessage<'_>],
mut missing_assistant_context: bool,
budgeting: ContextBudgeting,
) -> PreparedRootAssistants {
let delivered = delivered_messages
.iter()
.take(MAX_ROOT_MESSAGES)
.map(|message| (message.call_id, message.text))
.enumerate()
.take_while(|(index, _)| {
budgeting == ContextBudgeting::Elastic || *index < MAX_ROOT_MESSAGES
})
.map(|(_, message)| (message.call_id, message.text))
.collect::<HashMap<_, _>>();
let live = assistant_messages
.iter()
@@ -205,8 +212,11 @@ fn prepare_root_assistant_messages(
});
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));
.and_then(|retained| root_retained_assistant(retained.message, budgeting))
.unwrap_or_else(|| match budgeting {
ContextBudgeting::Fixed => crate::truncate_text(&text, MAX_ROOT_MESSAGE_TOKENS),
ContextBudgeting::Elastic => text,
});
let message = if order.is_some() {
GuardianRootMessage::Assistant(text)
} else {
@@ -229,7 +239,7 @@ fn prepare_root_assistant_messages(
{
return None;
}
let text = root_retained_assistant(source.message);
let text = root_retained_assistant(source.message, budgeting);
missing_assistant_context |= text.is_none();
text.map(|text| RootMessageCandidate {
order: Some(source.order),
@@ -244,8 +254,16 @@ fn prepare_root_assistant_messages(
}
}
fn root_retained_assistant(message: &codex_history::RetainedUserMessage) -> Option<String> {
crate::retained_assistant_message(message).map(|_| message.text.clone())
fn root_retained_assistant(
message: &codex_history::RetainedUserMessage,
budgeting: ContextBudgeting,
) -> Option<String> {
match budgeting {
ContextBudgeting::Fixed => {
crate::retained_assistant_message(message).map(|_| message.text.clone())
}
ContextBudgeting::Elastic => message.complete.then(|| message.text.clone()),
}
}
/// Bounded root view plus the existing cache and skill-scope selection metadata.
@@ -256,10 +274,11 @@ pub struct PreparedRootContext {
pub latest_user_turn_id: Option<String>,
}
/// Applies the fixed root request policy without discarding the captured sources.
/// Prepares the selected request policy without changing the captured sources.
pub fn prepare_root_context(
input: &RootConversationCandidates<'_>,
legacy_candidates: impl IntoIterator<Item = RetainedUserMessage>,
budgeting: ContextBudgeting,
) -> PreparedRootContext {
let mut missing_root_instructions = input.missing_root_instructions;
let mut missing_verified_answers = input.missing_verified_answers;
@@ -269,11 +288,14 @@ pub fn prepare_root_context(
.iter()
.filter_map(|candidate| {
if let GuardianRootMessage::UserInput(text) = &candidate.message
&& budgeting == ContextBudgeting::Fixed
&& !crate::verified_answers::answer_fits(text.len())
{
return None;
}
let message = if let GuardianRootMessage::User(text) = &candidate.message {
let message = if let GuardianRootMessage::User(text) = &candidate.message
&& budgeting == ContextBudgeting::Fixed
{
GuardianRootMessage::User(crate::truncate_text(text, MAX_ROOT_MESSAGE_TOKENS))
} else {
candidate.message.clone()
@@ -288,23 +310,29 @@ pub fn prepare_root_context(
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 {
if missing_root_instructions && (budgeting == ContextBudgeting::Elastic || 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.
// ID-less checkpoints identify sources by their stored excerpt, independently
// of whether request preparation will render the original text.
let retained_text = crate::truncate_text(&candidate.text, MAX_ROOT_MESSAGE_TOKENS);
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
candidate.turn_id == source.turn_id && retained_text == source.text
}
}) {
continue;
}
let text = match budgeting {
ContextBudgeting::Fixed => retained_text,
ContextBudgeting::Elastic => candidate.text,
};
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 {
if budgeting == ContextBudgeting::Fixed && legacy_messages.len() > legacy_capacity {
legacy_messages.pop_front();
}
}
@@ -314,11 +342,17 @@ pub fn prepare_root_context(
&input.retained_assistants,
&input.delivered_messages,
input.missing_assistant_context,
budgeting,
);
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);
let handoff_selection = match budgeting {
ContextBudgeting::Fixed => {
select_root_handoff_indices(&messages, &input.handoff_boundaries)
}
ContextBudgeting::Elastic => None,
};
if let Some(selected) = &handoff_selection {
messages = messages
.into_iter()
@@ -337,7 +371,10 @@ pub fn prepare_root_context(
})
.collect();
}
let removed = messages.len().saturating_sub(MAX_ROOT_MESSAGES);
let removed = match budgeting {
ContextBudgeting::Fixed => messages.len().saturating_sub(MAX_ROOT_MESSAGES),
ContextBudgeting::Elastic => 0,
};
for message in &messages[..removed] {
missing_root_instructions |= matches!(
message.message,
@@ -350,11 +387,13 @@ pub fn prepare_root_context(
}
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()),
);
if budgeting == ContextBudgeting::Fixed {
legacy_messages.drain(
..legacy_messages
.len()
.saturating_sub(MAX_ROOT_MESSAGES - messages.len()),
);
}
let mut messages = messages
.into_iter()
.map(|message| message.message)
codex-rs/guardian-context/src/preparation/source_tests.rs
+111
View File
@@ -0,0 +1,111 @@
#![cfg(test)]
//! Checks that elastic preparation keeps captured root and REPL evidence.
use super::*;
use crate::GuardianRootMessage;
use crate::NodeReplContext;
use crate::NodeReplResponse;
use crate::NodeReplReviewEvidenceMode;
use crate::RootConversationCandidates;
use crate::RootMessageCandidate;
use codex_history::RetainedContextOrder;
use codex_protocol::user_input::UserInput;
use pretty_assertions::assert_eq;
use std::sync::Arc;
#[test]
fn elastic_root_preparation_preserves_captured_messages() {
const ROOT_RECORDS: usize = 20;
const TEXT_REPEATS: usize = 1_000;
const LEGACY_SOURCE_TOKENS: usize = 900;
let text = "root evidence ".repeat(TEXT_REPEATS);
let legacy = codex_history::RetainedUserMessage {
turn_id: "legacy-turn".to_owned(),
message_id: None,
text: text.clone(),
complete: false,
origin: codex_history::UserInputOrigin::User,
phase: None,
};
let retained = codex_history::RetainedUserMessage {
text: crate::truncate_text(&text, LEGACY_SOURCE_TOKENS),
..legacy.clone()
};
let mut input = RootConversationCandidates {
user_messages: (0..ROOT_RECORDS)
.map(|index| RootMessageCandidate {
order: Some(RetainedContextOrder::Local(u64::try_from(index).unwrap())),
message: GuardianRootMessage::User(text.clone()),
})
.collect(),
assistant_messages: Vec::new(),
retained_assistants: Vec::new(),
delivered_messages: Vec::new(),
retained_user_sources: vec![retained.clone()],
handoff_boundaries: Vec::new(),
latest_user_turn_id: None,
missing_root_instructions: true,
missing_assistant_context: false,
missing_verified_answers: false,
};
input.user_messages[0].message = GuardianRootMessage::User(retained.text.clone());
let fixed = root::prepare_root_context(&input, vec![legacy.clone()], ContextBudgeting::Fixed);
let elastic = root::prepare_root_context(&input, vec![legacy], ContextBudgeting::Elastic);
let mut expected = vec![GuardianRootMessage::User(text); ROOT_RECORDS];
expected[0] = GuardianRootMessage::User(retained.text);
expected.insert(
/*index*/ 0,
GuardianRootMessage::IncompleteRootInstructions,
);
expected.insert(/*index*/ 0, GuardianRootMessage::RetainedContextScope);
assert_eq!(elastic.messages, expected);
assert!(fixed.messages.len() < elastic.messages.len());
}
#[test]
fn elastic_repl_rendering_keeps_stored_text_images_and_omissions() {
const TEXT_BYTES: usize = 32_001;
const STORAGE_OMISSIONS: u64 = 7;
let text = "x".repeat(TEXT_BYTES);
let image = UserInput::Image {
image: ImageReference::Inline {
image_url: "data:image/png;base64,original".to_owned(),
},
detail: None,
};
let responses = [Arc::new(NodeReplResponse {
sequence: 1,
provenance: "source".to_owned(),
items: vec![
UserInput::Text {
text: text.clone(),
text_elements: Vec::new(),
},
image.clone(),
],
})];
for mode in [
NodeReplReviewEvidenceMode::TextOnly,
NodeReplReviewEvidenceMode::Multimodal,
] {
let [fixed, elastic] =
[ContextBudgeting::Fixed, ContextBudgeting::Elastic].map(|budgeting| {
NodeReplContext {
responses: &responses,
omitted_responses: STORAGE_OMISSIONS,
mode,
budgeting,
}
.render_inputs()
});
let rendered = serde_json::to_string(&elastic).unwrap();
assert!(!serde_json::to_string(&fixed).unwrap().contains(&text));
assert!(rendered.contains(&text));
assert!(rendered.contains(&format!("node_repl_responses=\\\"{STORAGE_OMISSIONS}\\\"")));
assert_eq!(
elastic.contains(&image),
mode == NodeReplReviewEvidenceMode::Multimodal
);
}
}
codex-rs/guardian-context/src/registry_tests.rs
+1
View File
@@ -235,6 +235,7 @@ fn reused_registry_preserves_section_identity_and_source_roles() {
items: repl_items.to_vec(),
})];
let repl = super::NodeReplContext {
budgeting: crate::ContextBudgeting::Fixed,
responses: &repl_responses,
omitted_responses: 0,
mode: super::NodeReplReviewEvidenceMode::TextOnly,