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;