diff --git a/codex-rs/guardian-context/src/allocation.rs b/codex-rs/guardian-context/src/allocation.rs new file mode 100644 index 0000000000..589414d1f9 --- /dev/null +++ b/codex-rs/guardian-context/src/allocation.rs @@ -0,0 +1,327 @@ +//! Allocates whole optional entries with shared section framing, preserving content and order. + +use std::collections::BTreeMap; +use std::collections::HashSet; + +use crate::BudgetPriority; +use crate::Budgeted; +use crate::ComposedContext; +use crate::ElasticBudgetConfig; +use crate::ElasticSection; +use crate::RequestBudget; +use crate::Retention; +use crate::SectionError; +use crate::TruncationObservation; +use crate::budget::ByteCount; +use crate::budget::content_framing_tokens; +use crate::budget::section_content_tokens; +use crate::budget::section_tokens; +use crate::composition::SectionDelivery; +use crate::composition::SectionOutput; +use codex_protocol::models::ContentItem; + +#[derive(Clone, Copy)] +struct Candidate { + section_index: usize, + item_index: usize, + cost: usize, +} + +impl Candidate { + fn cost(self, content_counts: &[Option]) -> usize { + self.cost + .saturating_add(content_counts[self.section_index].map_or(0, |count| { + content_framing_tokens(count + 1).saturating_sub(content_framing_tokens(count)) + })) + } + + fn select(self, content_counts: &mut [Option]) -> usize { + let cost = self.cost(content_counts); + if let Some(count) = &mut content_counts[self.section_index] { + *count += 1; + } + cost + } +} + +#[derive(Default)] +struct GroupDemand { + pending_items: Vec, + total_cost: usize, + quota: usize, + used: usize, +} + +pub(crate) fn allocation_section( + section: &'static str, + retention: Retention, +) -> Result { + Ok(match section { + "conversation_transcript" if retention == Retention::Optional(BudgetPriority::Tool) => { + ElasticSection::TranscriptTools + } + "conversation_transcript" => ElasticSection::TranscriptMessages, + "root_conversation" => ElasticSection::RootConversationMessages, + "retained_user_instructions" => ElasticSection::RetainedUserInstructions, + "trusted_user_answers" => ElasticSection::TrustedUserAnswers, + "sender_user_messages" => ElasticSection::SenderUserMessages, + "planned_action" => ElasticSection::PlannedAction, + "previous_reviews" => ElasticSection::PreviousReviews, + "trusted_tool" => ElasticSection::TrustedTools, + "trusted_skills" => ElasticSection::TrustedSkills, + "transcript_images" => ElasticSection::TranscriptImages, + "node_repl_evidence" => ElasticSection::NodeReplEvidence, + "permissions" => ElasticSection::PermissionContext, + _ => return Err(SectionError::UnsupportedDelivery { section }), + }) +} + +/// Prepared section text with a borrowed immutable cap policy. +#[derive(Clone)] +pub struct PreparedElasticContext<'a> { + pub(crate) context: ComposedContext, + pub(crate) config: &'a ElasticBudgetConfig, +} + +impl PreparedElasticContext<'_> { + /// Reserves required evidence and history, then selects whole optional entries. + pub fn enforce_budget( + self, + budget: RequestBudget, + omission_notice: String, + ) -> Result { + let mut context = self.context; + context + .sections + .retain(|section| section.id != "budget_omission"); + let remaining = budget + .max_input_tokens + .checked_sub(budget.existing_context_tokens) + .ok_or(SectionError::EvidenceLimitExceeded { + section: "request_budget", + })?; + let demand = collect_demand(&context)?; + let context_tokens = context.estimated_tokens(); + if context_tokens <= remaining && context.truncations.is_empty() { + return Ok(context); + } + let notice = SectionOutput { + id: "budget_omission", + delivery: SectionDelivery::user_content(vec![Budgeted::required( + ContentItem::InputText { + text: omission_notice, + }, + )]), + }; + let notice_tokens = section_tokens(¬ice); + if context_tokens.saturating_add(notice_tokens) <= remaining { + context.sections.push(notice); + return Ok(context); + } + let allowance = remaining + .checked_sub(demand.required_tokens.saturating_add(notice_tokens)) + .ok_or(SectionError::EvidenceLimitExceeded { + section: "request_budget", + })?; + let selected = allocate_items(demand, self.config, allowance); + retain_selected(&mut context, &selected); + context.sections.push(notice); + if context.estimated_tokens() > remaining { + return Err(SectionError::EvidenceLimitExceeded { + section: "request_budget", + }); + } + Ok(context) + } +} + +/// Required reservation and optional group demand for one pending request. +struct ContextDemand { + required_tokens: usize, + groups: BTreeMap, + /// User-content counts for shared framing; native messages have no wrapper. + content_counts: Vec>, +} + +fn collect_demand(context: &ComposedContext) -> Result { + let mut required_tokens = 0usize; + let mut demands = BTreeMap::::new(); + let mut content_counts = Vec::with_capacity(context.sections.len()); + for (section_index, section) in context.sections.iter().enumerate() { + let required_count = match §ion.delivery { + SectionDelivery::UserContent(items) => Some( + items + .iter() + .filter(|item| !matches!(item.retention, Retention::Optional(_))) + .count(), + ), + SectionDelivery::Message(_) => None, + }; + content_counts.push(required_count); + required_tokens = + required_tokens.saturating_add(required_count.map_or(0, content_framing_tokens)); + let mut count = required_count.unwrap_or(0); + let mut add = |item_index, priority, cost: usize| { + let group = allocation_section(section.id, Retention::Optional(priority))?; + let framing = required_count.map_or(0, |_| { + content_framing_tokens(count + 1).saturating_sub(content_framing_tokens(count)) + }); + count += 1; + let demand = demands.entry(group).or_default(); + demand.total_cost = demand + .total_cost + .saturating_add(cost.saturating_add(framing)); + demand.pending_items.push(Candidate { + section_index, + item_index, + cost, + }); + Ok(()) + }; + match §ion.delivery { + SectionDelivery::Message(message) => { + let cost = section_tokens(section); + if let Retention::Optional(priority) = message.retention { + add(/*item_index*/ 0, priority, cost)?; + } else { + required_tokens = required_tokens.saturating_add(cost); + } + } + SectionDelivery::UserContent(items) => { + for (index, item) in items.iter().enumerate() { + let cost = section_content_tokens(&item.content); + if let Retention::Optional(priority) = item.retention { + add(index, priority, cost)?; + } else { + required_tokens = required_tokens.saturating_add(cost); + } + } + } + } + } + Ok(ContextDemand { + required_tokens, + groups: demands, + content_counts, + }) +} + +fn allocate_items( + demand: ContextDemand, + config: &ElasticBudgetConfig, + allowance: usize, +) -> HashSet<(usize, usize)> { + let mut demands = demand.groups; + let mut content_counts = demand.content_counts; + for (group, demand) in demands.iter_mut() { + let limits = config.section_limits(*group); + // Prefer one whole entry per group when the remaining space permits it. + let smallest = demand + .pending_items + .iter() + .map(|item| item.cost(&content_counts)) + .min() + .unwrap_or(0); + demand.quota = limits + .min_total_section_tokens() + .max(smallest) + .min(demand.total_cost); + } + let minimum: u128 = demands.values().map(|demand| demand.quota as u128).sum(); + if minimum > allowance as u128 { + for demand in demands.values_mut() { + demand.quota = (demand.quota as u128 * allowance as u128 / minimum) as usize; + } + } + let mut available = allowance - demands.values().map(|demand| demand.quota).sum::(); + loop { + let active = demands + .values() + .filter(|demand| demand.total_cost > demand.quota) + .count(); + if active == 0 || available / active == 0 { + break; + } + let share = available / active; + for demand in demands.values_mut() { + let granted = share.min(demand.total_cost - demand.quota); + demand.quota += granted; + available -= granted; + } + } + let mut selected = HashSet::new(); + let mut remaining = allowance; + for demand in demands.values_mut() { + let mut deferred = Vec::new(); + for candidate in demand.pending_items.drain(..).rev() { + if candidate.cost(&content_counts) <= demand.quota - demand.used { + let cost = candidate.select(&mut content_counts); + selected.insert((candidate.section_index, candidate.item_index)); + demand.used += cost; + remaining -= cost; + } else { + deferred.push(candidate); + } + } + deferred.reverse(); + demand.pending_items = deferred; + } + // Return packing gaps to the least-served group that can still use them. + while remaining > 0 { + let Some(demand) = demands + .values_mut() + .filter(|demand| !demand.pending_items.is_empty()) + .min_by_key(|demand| demand.used) + else { + break; + }; + let Some(candidate) = demand.pending_items.pop() else { + continue; + }; + if candidate.cost(&content_counts) <= remaining { + let cost = candidate.select(&mut content_counts); + selected.insert((candidate.section_index, candidate.item_index)); + demand.used += cost; + remaining -= cost; + } + } + selected +} + +fn retain_selected(context: &mut ComposedContext, selected: &HashSet<(usize, usize)>) { + for (section_index, section) in context.sections.iter_mut().enumerate() { + match §ion.delivery { + SectionDelivery::Message(message) + if matches!(message.retention, Retention::Optional(_)) + && !selected.contains(&(section_index, 0)) => + { + context.truncations.push(TruncationObservation { + component: section.id, + original_bytes: ByteCount::item(&message.content.item), + retained_bytes: 0, + }); + section.delivery = SectionDelivery::UserContent(Vec::new()); + } + SectionDelivery::UserContent(items) => { + let keep = items + .iter() + .enumerate() + .map(|(index, item)| { + !matches!(item.retention, Retention::Optional(_)) + || selected.contains(&(section_index, index)) + }) + .collect::>(); + crate::enforcement::retain_content( + section, + &mut context.truncations, + |index, _| keep[index], + ); + } + SectionDelivery::Message(_) => {} + } + } +} + +#[cfg(test)] +#[path = "allocation_tests.rs"] +mod tests; diff --git a/codex-rs/guardian-context/src/allocation_tests.rs b/codex-rs/guardian-context/src/allocation_tests.rs new file mode 100644 index 0000000000..12c39e5dd6 --- /dev/null +++ b/codex-rs/guardian-context/src/allocation_tests.rs @@ -0,0 +1,118 @@ +#![cfg(test)] + +use super::*; +use crate::ElasticSectionLimits; +use codex_protocol::models::ImageReference; +use pretty_assertions::assert_eq; + +#[test] +fn groups_select_recent_whole_entries_and_redistribute_packing_gaps() { + const ENTRY_TOKENS: usize = 256; + const ITEM_COSTS: [usize; 3] = [50, 200, 180]; + const ALLOWANCE: usize = 660; + const NEWEST: usize = 2; + let groups = [ + ElasticSection::TranscriptMessages, + ElasticSection::TranscriptTools, + ]; + let limits = ElasticSectionLimits::new(ENTRY_TOKENS, ENTRY_TOKENS).unwrap(); + let config = ElasticBudgetConfig { + sections: groups.map(|group| (group, limits)).into(), + }; + let mut demand = ContextDemand { + required_tokens: 0, + content_counts: vec![None, None], + groups: BTreeMap::new(), + }; + for (section_index, group) in groups.into_iter().enumerate() { + let mut entries = GroupDemand::default(); + for (item_index, cost) in ITEM_COSTS.into_iter().enumerate() { + entries.pending_items.push(Candidate { + section_index, + item_index, + cost, + }); + entries.total_cost += cost; + } + demand.groups.insert(group, entries); + } + assert_eq!( + allocate_items(demand, &config, ALLOWANCE), + HashSet::from([(0, 0), (0, 1), (0, NEWEST), (1, 0), (1, NEWEST)]) + ); +} + +#[test] +fn whole_images_use_defaults_and_charge_shared_framing_once() { + const ENTRY_TOKENS: usize = 128; + const IMAGE_COUNT: usize = 3; + const TOOL_COUNT: usize = 40; + const TOOL_REPEATS: usize = 24; + const LIMITED_TOKENS: usize = 1_000; + const HISTORY_TOKENS: usize = 12_000; + const NOTICE: &str = "Evidence omitted"; + let mut images = Vec::new(); + for index in 0..IMAGE_COUNT { + let image = ContentItem::InputImage { + image: ImageReference::Inline { + image_url: format!("data:image/png;base64,original-{index}"), + }, + detail: None, + }; + images.push(Budgeted::optional(image.into(), BudgetPriority::Image)); + } + let tool = Budgeted::optional( + "tool evidence ".repeat(TOOL_REPEATS).into(), + BudgetPriority::Tool, + ); + let mut tools = vec![tool; TOOL_COUNT]; + tools.extend([ + Budgeted::required("exact user instruction".to_owned().into()), + Budgeted::historical("manual approval".to_owned().into()), + ]); + let context = ComposedContext { + sections: [ + ("transcript_images", images), + ("conversation_transcript", tools), + ] + .map(|(id, items)| SectionOutput { + id, + delivery: SectionDelivery::UserContent(items), + }) + .into(), + truncations: Vec::new(), + }; + let mut expected = context.clone(); + if let SectionDelivery::UserContent(images) = &mut expected.sections[0].delivery { + images.remove(/*index*/ 0); + } + expected.sections.push(SectionOutput { + id: "budget_omission", + delivery: SectionDelivery::UserContent(vec![Budgeted::required(NOTICE.to_owned().into())]), + }); + let limits = ElasticSectionLimits::new(ENTRY_TOKENS, ENTRY_TOKENS).unwrap(); + for sections in [ + BTreeMap::new(), + BTreeMap::from([(ElasticSection::TranscriptTools, limits)]), + ] { + let config = ElasticBudgetConfig { sections }; + let input = PreparedElasticContext { + context: context.clone(), + config: &config, + }; + let mut budget = RequestBudget { + max_input_tokens: context.estimated_tokens() + HISTORY_TOKENS, + existing_context_tokens: HISTORY_TOKENS, + }; + let untouched = input.clone().enforce_budget(budget, NOTICE.into()).unwrap(); + assert_eq!(untouched.into_messages(), context.clone().into_messages()); + budget.max_input_tokens = expected.estimated_tokens() + HISTORY_TOKENS; + let admitted = input.clone().enforce_budget(budget, NOTICE.into()).unwrap(); + assert_eq!(admitted.into_messages(), expected.clone().into_messages()); + budget.max_input_tokens = LIMITED_TOKENS + HISTORY_TOKENS; + let admitted = input.enforce_budget(budget, NOTICE.into()).unwrap(); + assert!(admitted.estimated_tokens() <= LIMITED_TOKENS); + let text = serde_json::to_string(&admitted.into_messages()).unwrap(); + assert!(text.contains("exact user instruction") && text.contains("manual approval")); + } +} diff --git a/codex-rs/guardian-context/src/budget.rs b/codex-rs/guardian-context/src/budget.rs index a3ab6e69a6..eee0c3e8d5 100644 --- a/codex-rs/guardian-context/src/budget.rs +++ b/codex-rs/guardian-context/src/budget.rs @@ -112,7 +112,7 @@ pub struct RequestBudget { pub existing_context_tokens: usize, } -/// Named optional allocation groups, independent of prompt placement. +/// Optional allocation groups named for the source section that owns the evidence. /// The transcript has separate message and tool groups despite sharing one section. /// Required evidence, including the action and manual approvals, bypasses these limits. #[derive( @@ -122,23 +122,33 @@ pub struct RequestBudget { pub enum ElasticSection { TranscriptMessages, TranscriptTools, - RootMessages, - RetainedInstructions, - UserAnswers, - SenderMessages, - ToolDescriptions, + #[serde(alias = "root_messages")] + RootConversationMessages, + #[serde(alias = "retained_instructions")] + RetainedUserInstructions, + #[serde(alias = "user_answers")] + TrustedUserAnswers, + #[serde(alias = "sender_messages")] + SenderUserMessages, + #[serde( + alias = "tool_descriptions", + alias = "planned_action_tool_descriptions" + )] + PlannedAction, PreviousReviews, TrustedTools, TrustedSkills, - Images, - ReplEvidence, + #[serde(alias = "images")] + TranscriptImages, + #[serde(alias = "repl_evidence")] + NodeReplEvidence, PermissionContext, } /// Per-group elastic limits, validated when decoded; not an activation switch. /// Entries configure only optional evidence in the named group. They cannot -/// downgrade required evidence. Omitted groups have no override; the allocator -/// must supply its defaults when it is introduced. +/// downgrade required evidence. Omitted groups use the synchronous entry caps, +/// reserving one capped entry as their minimum. #[derive(Clone, Debug, Default, Eq, PartialEq, serde::Serialize, serde::Deserialize)] #[serde(deny_unknown_fields)] pub struct ElasticBudgetConfig { @@ -146,6 +156,25 @@ pub struct ElasticBudgetConfig { pub sections: std::collections::BTreeMap, } +impl ElasticBudgetConfig { + pub(crate) fn section_limits(&self, section: ElasticSection) -> ElasticSectionLimits { + self.sections.get(§ion).copied().unwrap_or_else(|| { + let entries = crate::ContextProfile::synchronous().transcript.entry_limits; + let tokens = if section == ElasticSection::TranscriptMessages { + entries.message_tokens + } else if section == ElasticSection::TranscriptTools { + entries.tool_tokens + } else { + crate::preparation::MAX_INSTRUCTION_TOKENS + }; + ElasticSectionLimits { + min_total_section_tokens: tokens, + max_section_entry_tokens: tokens, + } + }) + } +} + /// Demand-capped minimum allocation and per-entry text allowance for optional evidence. /// The minimum is a reservation, not a maximum: a section may receive spare capacity. /// Preparation must cap optional entries once, including framing and omission markers, @@ -153,7 +182,7 @@ pub struct ElasticBudgetConfig { /// Entry caps cannot exceed the existing text-part delivery allowance. /// Opaque text messages must be kept or omitted whole, never shortened. /// Images stay whole and use section/global space rather than the text entry cap. -/// Request-level feasibility depends on the active groups and remaining global budget. +/// Minimum reservations scale down together to fit the available optional budget. /// These settings do not change storage bounds or activate elastic allocation. #[derive(Clone, Copy, Debug, Eq, PartialEq, serde::Serialize)] pub struct ElasticSectionLimits { diff --git a/codex-rs/guardian-context/src/enforcement.rs b/codex-rs/guardian-context/src/enforcement.rs index bf82ddc79a..1a81e21ca7 100644 --- a/codex-rs/guardian-context/src/enforcement.rs +++ b/codex-rs/guardian-context/src/enforcement.rs @@ -396,7 +396,7 @@ impl ComposedContext { } } -fn retain_content( +pub(crate) fn retain_content( section: &mut SectionOutput, truncations: &mut Vec, mut retain: impl FnMut(usize, &mut SectionContent) -> bool, diff --git a/codex-rs/guardian-context/src/lib.rs b/codex-rs/guardian-context/src/lib.rs index fb4f04f5c9..da45587d22 100644 --- a/codex-rs/guardian-context/src/lib.rs +++ b/codex-rs/guardian-context/src/lib.rs @@ -69,6 +69,8 @@ pub use retained_instructions::RetainedInstructionItem; pub use retained_instructions::retained_assistant_message; mod action; +mod allocation; +pub use allocation::PreparedElasticContext; mod enforcement; pub(crate) use enforcement::BudgetPriority; pub use enforcement::Budgeted;