diff --git a/codex-rs/core/src/context/guardian_budget_omission.rs b/codex-rs/core/src/context/guardian_budget_omission.rs new file mode 100644 index 0000000000..3b7072fcc7 --- /dev/null +++ b/codex-rs/core/src/context/guardian_budget_omission.rs @@ -0,0 +1,32 @@ +//! Bounded reviewer notice for evidence omitted by the aggregate input budget. + +use super::ContextualUserFragment; +use codex_protocol::models::ContentItemKind; + +/// Identifies incomplete optional evidence without implying additional authority. +pub struct GuardianBudgetOmission; + +impl ContextualUserFragment for GuardianBudgetOmission { + fn content_kind(&self) -> ContentItemKind { + ContentItemKind("guardian.context_omission".to_owned()) + } + + fn role(&self) -> &'static str { + "user" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + ( + "", + "", + ) + } + + fn body(&self) -> String { + "Optional conversation evidence or images were omitted to fit the review input budget. Treat the remaining evidence as incomplete; omissions do not authorize actions.".to_owned() + } +} diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index 832cbce3b3..a203ff8661 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -11,7 +11,9 @@ mod developer_instructions; mod environment_context; mod environments_instructions; mod guardian_approved_action; +mod guardian_budget_omission; mod guardian_context_mode; +pub use guardian_budget_omission::GuardianBudgetOmission; mod guardian_followup_review_reminder; mod guardian_node_repl_policy; mod guardian_policy; diff --git a/codex-rs/core/src/guardian/prompt.rs b/codex-rs/core/src/guardian/prompt.rs index 9f5aea8daf..43f5b83214 100644 --- a/codex-rs/core/src/guardian/prompt.rs +++ b/codex-rs/core/src/guardian/prompt.rs @@ -1,4 +1,5 @@ use codex_extension_api::ConversationHistorySnapshot; +use codex_guardian_context::Budgeted; use codex_guardian_context::CollectedContext; use codex_guardian_context::ContextPresentation; use codex_guardian_context::ContextProfile; @@ -208,7 +209,9 @@ pub(crate) async fn build_guardian_prompt_items_with_parent_turn( let profile = ContextProfile::synchronous(); let mut transcript = profile.render_transcript(transcript_entries, offset); if transcript_entries.is_empty() { - transcript.items.push(placeholder.to_owned()); + transcript + .items + .push(Budgeted::required(placeholder.to_owned())); } let context = sections.compose(presentation, transcript)?; for (section, cost) in context.section_costs() { @@ -267,11 +270,18 @@ pub(crate) fn render_guardian_transcript_entries( let mut transcript = ContextProfile::synchronous().render_transcript(entries, /*entry_number_offset*/ 0); if entries.is_empty() { + transcript.items.push(Budgeted::required( + "".to_owned(), + )); + } + ( transcript .items - .push("".to_owned()); - } - (transcript.items, transcript.omission_note) + .into_iter() + .map(|item| item.content) + .collect(), + transcript.omission_note, + ) } /// Retains the human-readable conversation plus recent tool call / result diff --git a/codex-rs/guardian-context/src/budget.rs b/codex-rs/guardian-context/src/budget.rs index 6131d9bfd9..a400c1967b 100644 --- a/codex-rs/guardian-context/src/budget.rs +++ b/codex-rs/guardian-context/src/budget.rs @@ -57,6 +57,24 @@ pub struct SectionCost { } impl SectionCost { + fn add_content(mut self, item: &ContentItem) -> Self { + match item { + ContentItem::InputText { text } | ContentItem::OutputText { text } => { + self.text_bytes = self.text_bytes.saturating_add(text.len()); + } + ContentItem::InputImage { image_url, .. } => { + self.image_bytes = self.image_bytes.saturating_add(image_url.len()); + self.image_count = self.image_count.saturating_add(1); + } + ContentItem::InputAudio { audio_url } => { + // Guardian currently has no audio contributor. Count a future opaque + // payload conservatively until its consumer supplies modality costs. + self.text_bytes = self.text_bytes.saturating_add(audio_url.len()); + } + } + self + } + pub fn measurements(self) -> [(&'static str, usize); 4] { [ ("text_bytes", self.text_bytes), @@ -83,38 +101,21 @@ impl ComposedContext { /// Stable section names and numeric costs; never exposes evidence in diagnostics. pub fn section_costs(&self) -> impl Iterator + '_ { self.sections.iter().map(|section| { - let content = match §ion.delivery { - SectionDelivery::UserContent(content) => content, + let cost = match §ion.delivery { + SectionDelivery::UserContent(content) => content + .iter() + .map(|item| &item.content) + .fold(SectionCost::default(), SectionCost::add_content), SectionDelivery::Message(message) => match message.as_ref() { - ResponseItem::Message { content, .. } => content, - item => { - return ( - section.id, - SectionCost { - text_bytes: ByteCount::item(item), - ..SectionCost::default() - }, - ); - } + ResponseItem::Message { content, .. } => content + .iter() + .fold(SectionCost::default(), SectionCost::add_content), + item => SectionCost { + text_bytes: ByteCount::item(item), + ..SectionCost::default() + }, }, }; - let mut cost = SectionCost::default(); - for item in content { - match item { - ContentItem::InputText { text } | ContentItem::OutputText { text } => { - cost.text_bytes = cost.text_bytes.saturating_add(text.len()); - } - ContentItem::InputImage { image_url, .. } => { - cost.image_bytes = cost.image_bytes.saturating_add(image_url.len()); - cost.image_count = cost.image_count.saturating_add(1); - } - ContentItem::InputAudio { audio_url } => { - // Guardian currently has no audio contributor. Count a future opaque - // payload conservatively until its consumer supplies modality costs. - cost.text_bytes = cost.text_bytes.saturating_add(audio_url.len()); - } - } - } (section.id, cost) }) } @@ -138,20 +139,22 @@ pub(super) fn content_tokens(item: &ContentItem) -> usize { pub(super) fn section_tokens(section: &SectionOutput) -> usize { match §ion.delivery { - SectionDelivery::UserContent(content) if content.is_empty() => 0, - SectionDelivery::UserContent(content) => { - let separators = - TruncationPolicy::Bytes(content.len().saturating_sub(1)).token_budget(); - content.iter().map(content_tokens).fold( - estimate_input_tokens(&crate::composition::user_message(Vec::new())) - .saturating_add(separators), - usize::saturating_add, - ) - } + SectionDelivery::UserContent(content) => content + .iter() + .map(|item| content_tokens(&item.content)) + .fold(content_framing_tokens(content.len()), usize::saturating_add), SectionDelivery::Message(message) => estimate_input_tokens(message), } } +pub(super) fn content_framing_tokens(item_count: usize) -> usize { + if item_count == 0 { + return 0; + } + estimate_input_tokens(&crate::composition::user_message(Vec::new())) + .saturating_add(TruncationPolicy::Bytes(item_count - 1).token_budget()) +} + fn adjusted_tokens(mut bytes: usize, content: &[ContentItem]) -> usize { for item in content { if let ContentItem::InputImage { image_url, .. } = item { diff --git a/codex-rs/guardian-context/src/budget_tests.rs b/codex-rs/guardian-context/src/budget_tests.rs index b66ec2ea37..5af8e7460a 100644 --- a/codex-rs/guardian-context/src/budget_tests.rs +++ b/codex-rs/guardian-context/src/budget_tests.rs @@ -1,6 +1,7 @@ //! Costs distinguish evidence payloads from complete model request estimates. use super::*; +use crate::Budgeted; use crate::composition::SectionOutput; use crate::composition::user_message as message; use pretty_assertions::assert_eq; @@ -12,13 +13,13 @@ fn section_costs_keep_multimodal_payloads_separate() { SectionOutput { id: "transcript", delivery: SectionDelivery::UserContent(vec![ - ContentItem::InputText { + Budgeted::required(ContentItem::InputText { text: "évidence".to_owned(), - }, - ContentItem::InputImage { + }), + Budgeted::required(ContentItem::InputImage { image_url: "data:image/png;base64,AAAA".to_owned(), detail: None, - }, + }), ]), }, SectionOutput { @@ -76,9 +77,11 @@ fn section_estimate_bounds_the_delivered_message() { sections: vec![SectionOutput { id: "transcript", delivery: SectionDelivery::UserContent(vec![ - ContentItem::InputText { - text: text.to_owned(), - }; + Budgeted::required( + ContentItem::InputText { + text: text.to_owned(), + } + ); count ]), }], diff --git a/codex-rs/guardian-context/src/composition.rs b/codex-rs/guardian-context/src/composition.rs index f6db22a21b..e88d43f495 100644 --- a/codex-rs/guardian-context/src/composition.rs +++ b/codex-rs/guardian-context/src/composition.rs @@ -1,6 +1,7 @@ //! Composes collected evidence into ordered sections with explicit delivery. //! Profiles retain the host-selected transcript slice; composition owns //! framing, message boundaries and section placement, without retaining history. +//! Each content item keeps its selection policy until transport conversion. use codex_context_fragments::ContextualUserFragment; use codex_protocol::models::ContentItem; @@ -8,6 +9,8 @@ use codex_protocol::models::ResponseItem; use codex_protocol::user_input::UserInput; use crate::ActionPresentation; +use crate::BudgetPriority; +use crate::Budgeted; use crate::ContextSection; use crate::ConversationTranscriptEntry; use crate::SectionError; @@ -22,7 +25,7 @@ pub enum ContextPresentation<'a> { /// Host-selected, already bounded transcript entries and omission notice. pub struct RenderedTranscript { - pub items: Vec, + pub items: Vec>, pub omission_note: Option, pub truncations: Vec, } @@ -33,20 +36,21 @@ pub struct CollectedContext { } /// One section's delivery; separate messages retain their roles and annotations. -#[derive(PartialEq)] +#[derive(Clone, PartialEq)] pub(crate) enum SectionDelivery { - UserContent(Vec), + UserContent(Vec>), Message(Box), } /// Rendered evidence with a stable identity, independent of its source type. -#[derive(PartialEq)] +#[derive(Clone, PartialEq)] pub(crate) struct SectionOutput { pub id: &'static str, pub delivery: SectionDelivery, } /// Ordered sections ready for a consumer's transport adapter. +#[derive(Clone)] pub struct ComposedContext { pub(crate) sections: Vec, pub truncations: Vec, @@ -140,23 +144,41 @@ impl CollectedContext { transcript.take().ok_or(SectionError::UnsupportedDelivery { section: "conversation_transcript", })?; - let mut items = vec![start.to_owned()]; + let mut items = vec![Budgeted::required(start.to_owned())]; for (index, entry) in transcript.items.into_iter().enumerate() { - items.push(if session_id.is_some() { + let text = if session_id.is_some() { let prefix = if index == 0 { "" } else { "\n" }; - format!("{prefix}{entry}\n") + format!("{prefix}{}\n", entry.content) } else { - entry + entry.content + }; + items.push(Budgeted { + content: text, + retention: entry.retention, }); } - items.push(end.to_owned()); + items.push(Budgeted::required(end.to_owned())); if let Some(session_id) = session_id { - items.push(format!("Reviewed Codex session id: {session_id}\n")); + items.push(Budgeted::required(format!( + "Reviewed Codex session id: {session_id}\n" + ))); } if let Some(note) = transcript.omission_note { - items.push(format!("\n{note}\n")); + items.push(Budgeted::required(format!("\n{note}\n"))); } - (7, "conversation_transcript", text_content(items)) + ( + 7, + "conversation_transcript", + SectionDelivery::UserContent( + items + .into_iter() + .map(|item| Budgeted { + content: ContentItem::InputText { text: item.content }, + retention: item.retention, + }) + .collect(), + ), + ) } ContextSection::PermissionContext { items } => { (8, "permissions", text_content(items)) @@ -172,7 +194,13 @@ impl CollectedContext { ( if session_id.is_some() { 9 } else { 12 }, "transcript_images", - SectionDelivery::UserContent(images.images), + SectionDelivery::UserContent( + images + .images + .into_iter() + .map(|image| Budgeted::optional(image, BudgetPriority::Image)) + .collect(), + ), ) } ContextSection::NodeReplEvidence(evidence) => { @@ -180,10 +208,13 @@ impl CollectedContext { .items .into_iter() .map(|item| match item { - UserInput::Text { text, .. } => Ok(ContentItem::InputText { text }), - UserInput::Image { image_url, detail } => { - Ok(ContentItem::InputImage { image_url, detail }) + UserInput::Text { text, .. } => { + Ok(Budgeted::required(ContentItem::InputText { text })) } + UserInput::Image { image_url, detail } => Ok(Budgeted::optional( + ContentItem::InputImage { image_url, detail }, + BudgetPriority::Image, + )), _ => Err(SectionError::UnsupportedDelivery { section: "node_repl_evidence", }), @@ -213,7 +244,7 @@ fn text_content(items: Vec) -> SectionDelivery { SectionDelivery::UserContent( items .into_iter() - .map(|text| ContentItem::InputText { text }) + .map(|text| Budgeted::required(ContentItem::InputText { text })) .collect(), ) } @@ -229,7 +260,7 @@ impl ComposedContext { }); }; for item in content { - inputs.push(match item { + inputs.push(match item.content { ContentItem::InputText { text } => UserInput::Text { text, text_elements: Vec::new(), @@ -254,7 +285,9 @@ impl ComposedContext { let mut user_content = Vec::new(); for section in self.sections { match section.delivery { - SectionDelivery::UserContent(content) => user_content.extend(content), + SectionDelivery::UserContent(content) => { + user_content.extend(content.into_iter().map(|item| item.content)) + } SectionDelivery::Message(message) => { if !user_content.is_empty() { messages.push(user_message(std::mem::take(&mut user_content))); diff --git a/codex-rs/guardian-context/src/enforcement.rs b/codex-rs/guardian-context/src/enforcement.rs new file mode 100644 index 0000000000..a996875b45 --- /dev/null +++ b/codex-rs/guardian-context/src/enforcement.rs @@ -0,0 +1,214 @@ +//! Fits newly composed evidence into the remaining complete-request allowance. +//! Required evidence is never truncated. Optional content is removed in a stable +//! order, and a host-owned omission fragment is reserved before any removal. + +use std::collections::HashSet; + +use codex_protocol::models::ContentItem; +use codex_protocol::models::ImageDetail; + +use crate::ComposedContext; +use crate::RequestBudget; +use crate::SectionError; +use crate::TruncationObservation; +use crate::budget::content_framing_tokens; +use crate::budget::content_tokens; +use crate::budget::section_tokens; +use crate::composition::SectionDelivery; +use crate::composition::SectionOutput; + +/// Eviction priority within a profile; older items at the same priority go first. +#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)] +pub(crate) enum BudgetPriority { + Commentary, + Tool, + Image, +} + +/// Content and its selection policy move together through rendering and admission. +/// Consumers may add required content; optional priorities stay crate-owned. +#[derive(Clone, PartialEq)] +pub struct Budgeted { + pub content: T, + pub(crate) retention: Retention, +} + +impl Budgeted { + pub fn required(content: T) -> Self { + Self { + content, + retention: Retention::Required, + } + } + + pub(crate) fn optional(content: T, priority: BudgetPriority) -> Self { + Self { + content, + retention: Retention::Optional(priority), + } + } +} + +impl std::fmt::Debug for Budgeted { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("Budgeted") + .field("retention", &self.retention) + .finish_non_exhaustive() + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum Retention { + Required, + Optional(BudgetPriority), +} + +impl ComposedContext { + /// Applies host image admission before aggregate selection, preserving section + /// identity and each retained item's selection policy. + pub fn retain_images(&mut self, mut admit: impl FnMut(&str, &mut Option) -> bool) { + for section in &mut self.sections { + retain_content(section, &mut self.truncations, |_, item| match item { + ContentItem::InputImage { image_url, detail } => admit(image_url, detail), + _ => true, + }); + } + } + + /// Returns complete evidence that fits, or no context at all. The host owns + /// the bounded omission fragment and the existing reviewer history/checkpoint. + pub fn enforce_budget( + mut self, + budget: RequestBudget, + omission_notice: String, + ) -> Result { + let remaining = budget + .max_input_tokens + .checked_sub(budget.existing_context_tokens) + .ok_or(SectionError::EvidenceLimitExceeded { + section: "request_budget", + })?; + let current = self.estimated_tokens(); + if current <= remaining { + return Ok(self); + } + let notice = SectionOutput { + id: "budget_omission", + delivery: SectionDelivery::UserContent(vec![Budgeted::required( + ContentItem::InputText { + text: omission_notice, + }, + )]), + }; + let mut required_tokens = current.saturating_add(section_tokens(¬ice)); + let mut needed = required_tokens.saturating_sub(remaining); + let mut candidates = Vec::new(); + let mut remaining_items = vec![0; self.sections.len()]; + let mut candidate_framing = vec![0; self.sections.len()]; + for (section_index, section) in self.sections.iter().enumerate() { + if let SectionDelivery::UserContent(content) = §ion.delivery { + let mut required_items = content.len(); + for (index, item) in content.iter().enumerate() { + let Retention::Optional(priority) = item.retention else { + continue; + }; + required_items -= 1; + let tokens = content_tokens(&item.content); + required_tokens = required_tokens.saturating_sub(tokens); + candidates.push((priority, section_index, index, tokens)); + } + // Recompute framing with only required items, then charge each + // candidate for the framing needed to add it back on its own. + let required_framing = content_framing_tokens(required_items); + required_tokens = required_tokens.saturating_sub( + content_framing_tokens(content.len()).saturating_sub(required_framing), + ); + candidate_framing[section_index] = + content_framing_tokens(required_items + 1).saturating_sub(required_framing); + remaining_items[section_index] = content.len(); + } + } + // Evidence that cannot fit beside the required content and notice must + // leave first, without evicting useful smaller entries on its behalf. + let optional_allowance = + remaining + .checked_sub(required_tokens) + .ok_or(SectionError::EvidenceLimitExceeded { + section: "request_budget", + })?; + let mut removed = HashSet::new(); + let mut remove = |section_index: usize, index: usize, tokens: usize| { + if !removed.insert((section_index, index)) { + return 0; + } + let count = &mut remaining_items[section_index]; + let framing = content_framing_tokens(*count); + *count -= 1; + tokens.saturating_add(framing.saturating_sub(content_framing_tokens(*count))) + }; + candidates.retain(|&(_, section_index, index, tokens)| { + if tokens.saturating_add(candidate_framing[section_index]) <= optional_allowance { + return true; + } + needed = needed.saturating_sub(remove(section_index, index, tokens)); + false + }); + candidates.sort_unstable(); + for (_, section_index, index, tokens) in candidates { + if needed == 0 { + break; + } + needed = needed.saturating_sub(remove(section_index, index, tokens)); + } + if removed.is_empty() { + return Err(SectionError::EvidenceLimitExceeded { + section: "request_budget", + }); + } + for (section_index, section) in self.sections.iter_mut().enumerate() { + retain_content(section, &mut self.truncations, |index, _| { + !removed.contains(&(section_index, index)) + }); + } + self.sections.push(notice); + if self.estimated_tokens() > remaining { + return Err(SectionError::EvidenceLimitExceeded { + section: "request_budget", + }); + } + Ok(self) + } +} + +fn retain_content( + section: &mut SectionOutput, + truncations: &mut Vec, + mut retain: impl FnMut(usize, &mut ContentItem) -> bool, +) { + let SectionDelivery::UserContent(content) = &mut section.delivery else { + return; + }; + let mut index = 0; + content.retain_mut(|item| { + let keep = retain(index, &mut item.content); + index += 1; + if !keep { + let original_bytes = match &item.content { + ContentItem::InputText { text } | ContentItem::OutputText { text } => text.len(), + ContentItem::InputImage { image_url, .. } => image_url.len(), + ContentItem::InputAudio { audio_url } => audio_url.len(), + }; + truncations.push(TruncationObservation { + component: section.id, + original_bytes, + retained_bytes: 0, + }); + } + keep + }); +} + +#[cfg(test)] +#[path = "enforcement_tests.rs"] +mod tests; diff --git a/codex-rs/guardian-context/src/enforcement_tests.rs b/codex-rs/guardian-context/src/enforcement_tests.rs new file mode 100644 index 0000000000..609ed392d8 --- /dev/null +++ b/codex-rs/guardian-context/src/enforcement_tests.rs @@ -0,0 +1,223 @@ +//! Aggregate budgets preserve required evidence and explicit message boundaries. + +use super::*; +use crate::budget::section_tokens; +use crate::composition::user_message; +use pretty_assertions::assert_eq; + +fn text(value: &str) -> ContentItem { + ContentItem::InputText { + text: value.to_owned(), + } +} + +#[test] +fn budget_reserves_existing_context_and_preserves_required_messages() { + let trusted = crate::PreviousReviews::try_from_fragments(vec!["verified review".to_owned()]) + .unwrap() + .into_message(); + let image = ContentItem::InputImage { + image_url: "data:image/png;base64,AAAA".to_owned(), + detail: None, + }; + let make_context = || ComposedContext { + sections: vec![ + SectionOutput { + id: "conversation_transcript", + delivery: SectionDelivery::UserContent(vec![ + Budgeted::required(text("user restriction")), + Budgeted::optional( + text(&"old commentary".repeat(/*n*/ 200)), + BudgetPriority::Commentary, + ), + Budgeted::optional( + text(&"old tool output".repeat(/*n*/ 100)), + BudgetPriority::Tool, + ), + Budgeted::required(text("latest tool evidence")), + Budgeted::optional(image.clone(), BudgetPriority::Image), + ]), + }, + SectionOutput { + id: "previous_reviews", + delivery: SectionDelivery::Message(Box::new(trusted.clone())), + }, + SectionOutput { + id: "planned_action", + delivery: SectionDelivery::UserContent(vec![Budgeted::required(text( + "exact action", + ))]), + }, + ], + truncations: Vec::new(), + }; + let notice = SectionOutput { + id: "budget_omission", + delivery: SectionDelivery::UserContent(vec![Budgeted::required(text("evidence omitted"))]), + }; + let full = make_context(); + let available = full.estimated_tokens() + - content_tokens(&text(&"old commentary".repeat(/*n*/ 200))) + - content_tokens(&text(&"old tool output".repeat(/*n*/ 100))) + + section_tokens(¬ice); + let context = full + .enforce_budget( + RequestBudget { + max_input_tokens: available + 2_000, + existing_context_tokens: 2_000, + }, + "evidence omitted".to_owned(), + ) + .unwrap(); + assert!(context.estimated_tokens() <= available); + assert_eq!( + context.into_messages(), + vec![ + user_message(vec![ + text("user restriction"), + text("latest tool evidence"), + image.clone() + ]), + trusted.clone(), + user_message(vec![text("exact action"), text("evidence omitted")]), + ] + ); + let without_image = make_context() + .enforce_budget( + RequestBudget { + max_input_tokens: available - content_tokens(&image), + existing_context_tokens: 0, + }, + "evidence omitted".to_owned(), + ) + .unwrap(); + assert_eq!( + without_image.into_messages(), + vec![ + user_message(vec![text("user restriction"), text("latest tool evidence")]), + trusted.clone(), + user_message(vec![text("exact action"), text("evidence omitted")]), + ] + ); +} + +#[test] +fn image_omission_preserves_text_and_later_eviction_policy() { + let evidence = text(&"optional commentary ".repeat(/*n*/ 100)); + let mut context = ComposedContext { + sections: vec![SectionOutput { + id: "evidence", + delivery: SectionDelivery::UserContent(vec![ + Budgeted::optional( + ContentItem::InputImage { + image_url: "rejected-image".to_owned(), + detail: None, + }, + BudgetPriority::Image, + ), + Budgeted::optional(evidence.clone(), BudgetPriority::Commentary), + Budgeted::required(text("user restriction")), + ]), + }], + truncations: Vec::new(), + }; + let without_oversized_image = context + .clone() + .enforce_budget( + RequestBudget { + max_input_tokens: 1_000, + existing_context_tokens: 0, + }, + "evidence omitted".to_owned(), + ) + .unwrap(); + assert_eq!( + without_oversized_image.into_messages(), + vec![user_message(vec![ + evidence.clone(), + text("user restriction"), + text("evidence omitted") + ])] + ); + context.retain_images(|_, _| false); + let available = context.estimated_tokens(); + let retained = context + .clone() + .enforce_budget( + RequestBudget { + max_input_tokens: available, + existing_context_tokens: 0, + }, + "evidence omitted".to_owned(), + ) + .unwrap(); + assert_eq!( + retained.into_messages(), + vec![user_message(vec![evidence, text("user restriction")])] + ); + let smaller = context + .enforce_budget( + RequestBudget { + max_input_tokens: 100, + existing_context_tokens: 0, + }, + "evidence omitted".to_owned(), + ) + .unwrap(); + assert_eq!( + smaller.into_messages(), + vec![user_message(vec![ + text("user restriction"), + text("evidence omitted") + ])] + ); + + let older = ContentItem::InputImage { + image_url: "older-image".to_owned(), + detail: None, + }; + let newer = ContentItem::InputImage { + image_url: "newer-image".to_owned(), + detail: None, + }; + let image_section = |images: Vec| SectionOutput { + id: "transcript_images", + delivery: SectionDelivery::UserContent( + images + .into_iter() + .map(|image| Budgeted::optional(image, BudgetPriority::Image)) + .collect(), + ), + }; + let notice = SectionOutput { + id: "budget_omission", + delivery: SectionDelivery::UserContent(vec![Budgeted::required(text("evidence omitted"))]), + }; + let available = section_tokens(&image_section(vec![newer.clone()])) + section_tokens(¬ice); + // Removing the older image frees a separator in one section, or an entire + // wrapper in separate sections. Either way, the newer image fits exactly. + for sections in [ + vec![image_section(vec![older.clone(), newer.clone()])], + vec![ + image_section(vec![older]), + image_section(vec![newer.clone()]), + ], + ] { + let context = ComposedContext { + sections, + truncations: Vec::new(), + } + .enforce_budget( + RequestBudget { + max_input_tokens: available, + existing_context_tokens: 0, + }, + "evidence omitted".to_owned(), + ) + .unwrap(); + assert_eq!( + context.into_messages(), + vec![user_message(vec![newer.clone(), text("evidence omitted")])] + ); + } +} diff --git a/codex-rs/guardian-context/src/lib.rs b/codex-rs/guardian-context/src/lib.rs index e33f6562b3..df76d79600 100644 --- a/codex-rs/guardian-context/src/lib.rs +++ b/codex-rs/guardian-context/src/lib.rs @@ -48,6 +48,10 @@ pub use verified_answers::render_verified_answers; mod retained_instructions; mod action; +mod enforcement; +pub(crate) use enforcement::BudgetPriority; +pub use enforcement::Budgeted; +pub(crate) use enforcement::Retention; mod budget; mod composition; pub use budget::DEFAULT_MAX_INPUT_TOKENS; diff --git a/codex-rs/guardian-context/src/profile.rs b/codex-rs/guardian-context/src/profile.rs index 5e700f9bee..7a2efdfefb 100644 --- a/codex-rs/guardian-context/src/profile.rs +++ b/codex-rs/guardian-context/src/profile.rs @@ -4,12 +4,15 @@ use codex_protocol::protocol::TruncationPolicy; +use crate::BudgetPriority; +use crate::Budgeted; use crate::ContextTarget; use crate::ConversationTranscriptConfig; use crate::ConversationTranscriptEntry; use crate::ConversationTranscriptEntryKind; use crate::ConversationTranscriptOptions; use crate::RenderedTranscript; +use crate::Retention; use crate::TranscriptEntryLimits; use crate::TranscriptRetentionConfig; use crate::TruncationObservation; @@ -19,6 +22,8 @@ use crate::select_user_messages; use self::window::TranscriptWindow; mod window; +const MIN_RECENT_TOOL_ENTRIES: usize = 5; + /// Request-local policy resolved from the consumer's model and configuration. /// Registry scope follows `target`; source flags and caps apply before retention. #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -198,7 +203,7 @@ impl ContextProfile { && included.iter().any(|included| !included)) .then(|| "Some conversation entries were omitted.".to_owned()); let mut truncations = Vec::new(); - let items = entries + let mut items: Vec<_> = entries .into_iter() .enumerate() .filter_map(|(index, entry)| { @@ -221,9 +226,32 @@ impl ContextProfile { retained_bytes, }); } - included[index].then_some(entry.text) + if included[index] { + Some(match entry.kind { + TranscriptEntryKind::User | TranscriptEntryKind::ProtectedMessage => { + Budgeted::required(entry.text) + } + TranscriptEntryKind::Message => { + Budgeted::optional(entry.text, BudgetPriority::Commentary) + } + TranscriptEntryKind::Tool => { + Budgeted::optional(entry.text, BudgetPriority::Tool) + } + }) + } else { + None + } }) .collect(); + // Keep the newest tool evidence even when the aggregate allowance is tight. + for item in items + .iter_mut() + .rev() + .filter(|item| item.retention == Retention::Optional(BudgetPriority::Tool)) + .take(MIN_RECENT_TOOL_ENTRIES) + { + item.retention = Retention::Required; + } RenderedTranscript { items, omission_note, diff --git a/codex-rs/guardian-context/src/profile/window.rs b/codex-rs/guardian-context/src/profile/window.rs index fba26bcfe0..f82173b56e 100644 --- a/codex-rs/guardian-context/src/profile/window.rs +++ b/codex-rs/guardian-context/src/profile/window.rs @@ -8,7 +8,7 @@ use crate::TranscriptRetentionConfig; use super::TranscriptEntry; use super::TranscriptEntryKind; -const MIN_RECENT_TOOL_ENTRIES: usize = 5; +use super::MIN_RECENT_TOOL_ENTRIES; #[derive(Default)] struct EntryPool { diff --git a/codex-rs/guardian-context/src/profile_tests.rs b/codex-rs/guardian-context/src/profile_tests.rs index ea94cfc2cb..bd89168067 100644 --- a/codex-rs/guardian-context/src/profile_tests.rs +++ b/codex-rs/guardian-context/src/profile_tests.rs @@ -30,8 +30,11 @@ fn profiles_preserve_distinct_retention_and_original_numbering() { (sync.items, sync.omission_note), ( vec![ - "[8] user: inspect only".to_owned(), - "[10] assistant: working".to_owned() + Budgeted::required("[8] user: inspect only".to_owned()), + Budgeted::optional( + "[10] assistant: working".to_owned(), + BudgetPriority::Commentary + ) ], Some("Some conversation entries were omitted.".to_owned()), ), @@ -40,8 +43,8 @@ fn profiles_preserve_distinct_retention_and_original_numbering() { (asynchronous.items, asynchronous.omission_note), ( vec![ - "[1] user: inspect only\n".to_owned(), - "[2] developer: approved action\n".to_owned() + Budgeted::required("[1] user: inspect only\n".to_owned()), + Budgeted::required("[2] developer: approved action\n".to_owned()) ], None, ), @@ -59,3 +62,37 @@ fn profiles_preserve_distinct_retention_and_original_numbering() { vec![("transcript_message", "working".len(), 0)], ); } + +#[test] +fn profiles_reserve_the_newest_five_tool_entries_for_aggregate_enforcement() { + let entries = (0..8) + .map(|index| ConversationTranscriptEntry { + kind: ConversationTranscriptEntryKind::ToolOutput("tool result".to_owned()), + text: format!("result {index}"), + original_bytes: 8, + }) + .collect::>(); + for profile in [ + ContextProfile::synchronous(), + ContextProfile::asynchronous(), + ] { + let transcript = profile.render_transcript(&entries, /*entry_number_offset*/ 0); + assert_eq!( + transcript + .items + .iter() + .map(|item| item.retention) + .collect::>(), + vec![ + Retention::Optional(BudgetPriority::Tool), + Retention::Optional(BudgetPriority::Tool), + Retention::Optional(BudgetPriority::Tool), + Retention::Required, + Retention::Required, + Retention::Required, + Retention::Required, + Retention::Required + ] + ); + } +} diff --git a/codex-rs/guardian-context/src/truncation.rs b/codex-rs/guardian-context/src/truncation.rs index 896db0dd24..b38a0a7e74 100644 --- a/codex-rs/guardian-context/src/truncation.rs +++ b/codex-rs/guardian-context/src/truncation.rs @@ -42,6 +42,7 @@ pub fn truncate_text(text: &str, max_tokens: usize) -> String { mod tests; /// Actual evidence reduction, reported by consumers using their existing metrics. +#[derive(Clone)] pub struct TruncationObservation { pub component: &'static str, pub original_bytes: usize,