From 76cf0f45ec44e488297cd27e1b1a951f02c31bfe Mon Sep 17 00:00:00 2001 From: ioanbudea-OAI Date: Fri, 13 Nov 2026 14:50:02 +0000 Subject: [PATCH] Reserve item ID overhead before elastic context allocation (#61347) ## Why Elastic request allocation did not account for item IDs added after evidence selection. Their token cost could push the completed request over its input limit instead of allowing optional evidence to be dropped. ## What changed Reserve IDs for the possible new request items before allocating evidence, including their JSON escaping and field overhead in the global request budget. Assign the same reserved IDs to the selected items while preserving existing IDs and committed history. ## Testing Add coverage for escaped IDs and native tool call outputs, and a regression test with long IDs that verifies optional images are omitted while committed history and new transcript text are retained within budget. Update budget boundary tests to distinguish fixed rejection from elastic evidence reduction. GitOrigin-RevId: 621d6d83e3cc95efb9eef4976dff9b7799a6e7c8 --- .../guardian-v2/src/async_scorer/request.rs | 24 +++-- .../src/async_scorer/request_tests.rs | 98 +++++++++++++++++-- codex-rs/guardian-context/src/allocation.rs | 90 +++++++++++++++++ .../guardian-context/src/allocation_tests.rs | 60 ++++++++++++ codex-rs/guardian-context/src/budget.rs | 2 +- codex-rs/guardian-context/src/lib.rs | 1 + 6 files changed, 257 insertions(+), 18 deletions(-) diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/request.rs b/codex-rs/ext/guardian-v2/src/async_scorer/request.rs index 7439107c5d..81f23e7047 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/request.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/request.rs @@ -7,6 +7,7 @@ use codex_context_fragments::ContextualUserFragment; use codex_core::context::GuardianBudgetOmission; use codex_guardian_context::RequestBudget; +use codex_guardian_context::ReservedElasticItemIds; use codex_guardian_context::TranscriptCursor; use codex_guardian_context::TruncationObservation; use codex_guardian_context::estimate_input_tokens; @@ -64,17 +65,20 @@ impl LunaSamplingRequest { .map(|item| estimate_input_tokens(&item.item)) .fold(0usize, usize::saturating_add); let maximum = max_input_tokens.saturating_sub(INPUT_TOKEN_MARGIN); - if let Some(policy) = &evidence.elastic { - context = context + let ids = if let Some(policy) = &evidence.elastic { + let prepared = context .prepare_elastic_entries(policy) .ok()? .prepare_delivery(history) - .ok()? + .ok()?; + let ids = ReservedElasticItemIds::reserve(&prepared, &self.item_id_generator); + context = prepared .enforce_budget( RequestBudget { max_input_tokens: maximum, existing_context_tokens: existing_context_tokens - .saturating_add(setup_tokens), + .saturating_add(setup_tokens) + .saturating_add(ids.tokens()), }, GuardianBudgetOmission .render() @@ -82,11 +86,17 @@ impl LunaSamplingRequest { .unwrap_or_default(), ) .ok()?; - } + Some(ids) + } else { + None + }; let truncations = std::mem::take(&mut context.truncations); let section_costs = context.section_costs().collect(); - let (mut input, new_input_tokens) = - LunaSampler::finish_input(self, setup, context.into_annotated_messages()); + let mut suffix = context.into_annotated_messages(); + if let Some(ids) = ids { + ids.assign(&mut suffix).ok()?; + } + let (mut input, new_input_tokens) = LunaSampler::finish_input(self, setup, suffix); let input_tokens = existing_context_tokens.saturating_add(new_input_tokens); if input_tokens > maximum { return None; 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 7a562b879b..bd3632cfa6 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 @@ -1,4 +1,5 @@ use super::*; +use crate::async_scorer::sampler::INPUT_TOKEN_MARGIN; use crate::async_scorer::transcript::CollectedTranscript; use crate::async_scorer::transcript::ContextInput; use crate::async_scorer::transcript::TranscriptConfig; @@ -235,6 +236,7 @@ fn preparation_preserves_committed_prefix_and_selects_only_new_entries_and_image #[test_case::test_case(ContextBudgeting::Fixed; "fixed")] #[test_case::test_case(ContextBudgeting::Elastic; "elastic")] fn rejected_preparation_preserves_committed_history_and_cursor(mode: ContextBudgeting) { + const FRESH_REQUEST_HEADROOM: usize = 128; let mut history = History { items: vec![responses::user_message_item("Inspect README.md.")], version: 1, @@ -244,7 +246,8 @@ fn rejected_preparation_preserves_committed_history_and_cursor(mode: ContextBudg let first = config .prepare_retained(&collect(&history, &[], mode), &mut state, INPUT_BUDGET) .unwrap(); - let limit = first.input_tokens + 256; + // Leave room for the fresh request's allocation estimate and ID reservation. + let limit = first.input_tokens + INPUT_TOKEN_MARGIN + FRESH_REQUEST_HEADROOM; let mut committed = first.input; committed.push(responses::user_message_item(&"output ".repeat(/*n*/ 400)).into()); state.complete_review(first.cursor); @@ -278,6 +281,62 @@ fn rejected_preparation_preserves_committed_history_and_cursor(mode: ContextBudg assert_eq!(state.snapshot().unwrap().history(), &committed); } +#[test] +fn elastic_fits_ids_without_resetting_committed_history() { + const ID_PADDING_BYTES: usize = 1_024; + let mode = ContextBudgeting::Elastic; + let mut config = sample_request("parent"); + let original_generator = std::sync::Arc::clone(&config.item_id_generator); + let padding = "x".repeat(ID_PADDING_BYTES); + config.item_id_generator = std::sync::Arc::new(move |prefix| { + ResponseItemId::from_server(format!("{}{padding}", original_generator(prefix))) + }); + let mut history = History { + items: vec![responses::user_message_item("Inspect README.md.")], + version: 1, + }; + let mut state = ConversationState::default(); + let first = config + .prepare_retained(&collect(&history, &[], mode), &mut state, INPUT_BUDGET) + .unwrap(); + state.complete_review(first.cursor); + state.commit_snapshot(first.input.clone()); + history + .items + .push(responses::user_message_item("Inspect LICENSE next.")); + let image = ContentItem::InputImage { + image: ImageReference::Inline { + image_url: "data:image/png;base64,b3B0aW9uYWw=".to_owned(), + }, + detail: None, + }; + let evidence = collect(&history, std::slice::from_ref(&image), mode); + let mut generous = ConversationState::from(state.snapshot().unwrap().clone()); + let full = config + .prepare_retained(&evidence, &mut generous, INPUT_BUDGET) + .unwrap(); + assert!(full.input.iter().any(|item| matches!( + &item.item, ResponseItem::Message { content, .. } if content.contains(&image) + ))); + let maximum = full.input_tokens - 1; + let bounded = config + .prepare_retained(&evidence, &mut state, maximum + INPUT_TOKEN_MARGIN) + .expect("allocate evidence with room for IDs"); + assert_eq!(&bounded.input[..first.input.len()], first.input); + assert!(bounded.input_tokens <= maximum); + assert!(bounded.input.iter().all(|item| item.id().is_some())); + let suffix = serde_json::to_string( + &bounded.input[first.input.len()..] + .iter() + .map(|item| &item.item) + .collect::>(), + ) + .unwrap(); + assert!(suffix.contains("Inspect LICENSE next.")); + assert!(!suffix.contains("data:image/png;base64,b3B0aW9uYWw=")); + assert!(suffix.contains("")); +} + #[test_case::test_case(ContextBudgeting::Fixed; "fixed")] #[test_case::test_case(ContextBudgeting::Elastic; "elastic")] fn complete_budget_counts_setup_checkpoint_images_and_trusted_evidence(mode: ContextBudgeting) { @@ -307,13 +366,32 @@ fn complete_budget_counts_setup_checkpoint_images_and_trusted_evidence(mode: Con config.parent_compaction.as_ref() ); assert!(request.input_tokens > 10_000); - assert!( - config - .prepare_retained(&evidence, &mut state, request.input_tokens + 255) - .is_none() - ); - // The complete input fits exactly at the reserved-output boundary. - config - .prepare_retained(&evidence, &mut state, request.input_tokens + 256) - .unwrap(); + let maximum = request.input_tokens - 1; + let smaller = config.prepare_retained(&evidence, &mut state, maximum + INPUT_TOKEN_MARGIN); + match mode { + ContextBudgeting::Fixed => { + assert!(smaller.is_none()); + let exact = config + .prepare_retained( + &evidence, + &mut state, + request.input_tokens + INPUT_TOKEN_MARGIN, + ) + .unwrap(); + assert_eq!(exact.input_tokens, request.input_tokens); + assert!(exact.input.iter().any(|item| matches!( + &item.item, ResponseItem::Message { content, .. } if content.contains(&images[0]) + ))); + } + ContextBudgeting::Elastic => { + let smaller = smaller.expect("Elastic can shed optional evidence"); + assert!(smaller.input_tokens <= maximum); + assert!( + smaller + .input + .iter() + .any(|item| Some(&item.item) == config.parent_compaction.as_ref()) + ); + } + } } diff --git a/codex-rs/guardian-context/src/allocation.rs b/codex-rs/guardian-context/src/allocation.rs index ba836d21e8..f2da198a06 100644 --- a/codex-rs/guardian-context/src/allocation.rs +++ b/codex-rs/guardian-context/src/allocation.rs @@ -1,6 +1,7 @@ //! Funds fixed section minimums and scales only the headroom up to section maxima. //! Required records precede Optional records within each section. //! Protected framing and exact actions use global space. Selection preserves prepared bytes and order. +//! Async callers can reserve wire space for IDs in the new suffix before selection. use std::collections::BTreeMap; use std::collections::HashSet; @@ -21,7 +22,10 @@ use crate::budget::section_tokens; use crate::composition::SectionDelivery; use crate::composition::SectionOutput; use codex_history::ResponseItemEnvelope; +use codex_protocol::ItemIdGenerator; +use codex_protocol::ResponseItemId; use codex_protocol::models::ContentItem; +use codex_protocol::protocol::TruncationPolicy; const REQUEST_BUDGET_EXCEEDED: SectionError = SectionError::EvidenceLimitExceeded { section: "request_budget", @@ -281,6 +285,92 @@ impl PreparedElasticContext<'_> { } } +const MESSAGE_ID_PREFIX: &str = "msg"; +const ID_FIELD_FRAMING: &str = r#","id":"#; + +/// Actual IDs reserved for possible new messages and their conservative wire token cost. +pub struct ReservedElasticItemIds { + by_prefix: BTreeMap<&'static str, Vec>, + tokens: usize, +} + +impl ReservedElasticItemIds { + /// Reserves IDs for the largest possible emitted suffix before dropping evidence. + /// Removing a native item can merge adjacent user messages, never split them. + pub fn reserve(prepared: &PreparedElasticContext<'_>, generate: &ItemIdGenerator) -> Self { + let mut reserved = Self { + by_prefix: BTreeMap::new(), + tokens: 0, + }; + let mut has_user_content = false; + let mut reserve = |prefix| { + let id = generate(prefix); + // Measure the wire string, including escaping and the optional field framing. + let wire_bytes = + ByteCount::measure(|counter| serde_json::to_writer(counter, id.as_str())) + .saturating_add(ID_FIELD_FRAMING.len()); + reserved.tokens = reserved + .tokens + .saturating_add(TruncationPolicy::Bytes(wire_bytes).token_budget()); + reserved.by_prefix.entry(prefix).or_default().push(id); + }; + for section in &prepared.context.sections { + match §ion.delivery { + SectionDelivery::UserContent(items) => { + // Omission-state changes may add framing to retained instructions. + has_user_content |= + !items.is_empty() || section.id == "retained_user_instructions"; + } + SectionDelivery::Message(message) => { + if std::mem::take(&mut has_user_content) { + reserve(MESSAGE_ID_PREFIX); + } + let item = &message.content.item; + if item.id().is_none() + && let Some(prefix) = item.id_prefix() + { + reserve(prefix); + } + if let Some(prefix) = crate::native_tool_call::evidence_output(item) + .and_then(|output| output.id_prefix()) + { + reserve(prefix); + } + } + } + } + // The final user message can consist solely of the allocation omission notice. + reserve(MESSAGE_ID_PREFIX); + reserved + } + /// Global request overhead, separate from section evidence quotas. + pub fn tokens(&self) -> usize { + self.tokens + } + + /// Uses the same IDs that were measured without touching the existing prefix. + pub fn assign(&self, suffix: &mut [ResponseItemEnvelope]) -> Result<(), SectionError> { + let mut ids = self + .by_prefix + .iter() + .map(|(&prefix, ids)| (prefix, ids.iter())) + .collect::>(); + for item in suffix { + if item.id().is_none() + && let Some(prefix) = item.id_prefix() + { + let id = ids.get_mut(prefix).and_then(Iterator::next).ok_or( + SectionError::UnsupportedDelivery { + section: "response_item_ids", + }, + )?; + item.set_id(Some(id.clone())); + } + } + Ok(()) + } +} + /// Protected global cost and priority-ordered demand for each section. struct ContextDemand { protected_tokens: usize, diff --git a/codex-rs/guardian-context/src/allocation_tests.rs b/codex-rs/guardian-context/src/allocation_tests.rs index 586811259a..72c7e9f71a 100644 --- a/codex-rs/guardian-context/src/allocation_tests.rs +++ b/codex-rs/guardian-context/src/allocation_tests.rs @@ -2,8 +2,68 @@ use super::*; use crate::ElasticSectionLimits; +use crate::estimate_input_tokens; use codex_protocol::models::ImageReference; +use codex_protocol::models::ResponseItem; +use codex_protocol::models::input_text; +use codex_protocol_test_support::uuid_v7_item_id_generator_for_test; use pretty_assertions::assert_eq; +use std::sync::Arc; + +#[test] +fn reserves_escaped_ids_for_user_messages_and_native_call_outputs() { + let call: ResponseItem = serde_json::from_value(serde_json::json!({ + "type": "function_call", + "id": "fc_committed", + "call_id": "call", + "name": "read_file", + "arguments": "{}", + })) + .unwrap(); + let call_id = call.id().cloned(); + let sections = [ + SectionOutput { + id: "conversation_transcript", + delivery: SectionDelivery::user_content(vec![Budgeted::required(input_text("before"))]), + }, + SectionOutput { + id: "conversation_transcript", + delivery: SectionDelivery::Message(Budgeted::required(Box::new(call.into()))), + }, + SectionOutput { + id: "budget_omission", + delivery: SectionDelivery::user_content(vec![Budgeted::required(input_text("after"))]), + }, + ]; + let prepared = PreparedElasticContext { + context: ComposedContext { + sections: sections.to_vec(), + truncations: Vec::new(), + }, + config: &ElasticBudgetConfig::default(), + delivered_omissions: None, + }; + let uuid = uuid_v7_item_id_generator_for_test(); + let generator: ItemIdGenerator = + Arc::new(move |prefix| ResponseItemId::from_server(format!("{}\"\\", uuid(prefix)))); + let reserved = ReservedElasticItemIds::reserve(&prepared, &generator); + let mut suffix = prepared.context.into_annotated_messages(); + let measure = |items: &[ResponseItemEnvelope]| { + items + .iter() + .map(|item| estimate_input_tokens(&item.item)) + .sum::() + }; + let before = measure(&suffix); + reserved.assign(&mut suffix).unwrap(); + assert!(suffix.iter().all(|item| item.id().is_some())); + let delivered_call = suffix + .iter() + .find(|item| matches!(item.item, ResponseItem::FunctionCall { .. })) + .unwrap(); + assert_eq!(delivered_call.id(), call_id.as_ref()); + assert!(measure(&suffix) - before <= reserved.tokens()); +} #[test] fn groups_select_recent_whole_entries_and_redistribute_packing_gaps() { diff --git a/codex-rs/guardian-context/src/budget.rs b/codex-rs/guardian-context/src/budget.rs index 4bd5d5332c..f21ed75fb0 100644 --- a/codex-rs/guardian-context/src/budget.rs +++ b/codex-rs/guardian-context/src/budget.rs @@ -571,7 +571,7 @@ impl ByteCount { Self::measure(|counter| serde_json::to_writer(counter, item)) } - fn measure(serialize: impl FnOnce(&mut Self) -> serde_json::Result<()>) -> usize { + pub(super) fn measure(serialize: impl FnOnce(&mut Self) -> serde_json::Result<()>) -> usize { let mut count = Self::default(); if serialize(&mut count).is_err() { return usize::MAX; diff --git a/codex-rs/guardian-context/src/lib.rs b/codex-rs/guardian-context/src/lib.rs index 977f6e93e1..0fa99cd659 100644 --- a/codex-rs/guardian-context/src/lib.rs +++ b/codex-rs/guardian-context/src/lib.rs @@ -71,6 +71,7 @@ pub use retained_instructions::retained_assistant_message; mod action; mod allocation; pub use allocation::PreparedElasticContext; +pub use allocation::ReservedElasticItemIds; mod enforcement; pub(crate) use enforcement::BudgetPriority; pub use enforcement::Budgeted;