Add whole-entry elastic allocation for Guardian context (#58388)

## What changed

Add `PreparedElasticContext::enforce_budget` to reserve existing context, required evidence, history, and an omission notice before allocating remaining tokens across optional evidence groups. Scale minimum reservations to fit, favor recent whole entries, and redistribute unused capacity while preserving retained content and order.

Account for shared section framing and keep images whole. Record omitted entries and reject requests when required evidence cannot fit. Supply defaults for unconfigured groups and rename allocation groups to match their source sections, retaining previous serialized names as aliases.

## Testing

Add unit tests for recent-entry selection, redistribution of packing gaps, whole-image retention, shared framing, default limits, and preservation of required instructions and historical approvals under tight budgets.

GitOrigin-RevId: a3c789a5936ec3f22893a31a0014c0902197a281

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>]) -> 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>]) -> 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<Candidate>,
+    total_cost: usize,
+    quota: usize,
+    used: usize,
+}
+
+pub(crate) fn allocation_section(
+    section: &'static str,
+    retention: Retention,
+) -> Result<ElasticSection, SectionError> {
+    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<ComposedContext, SectionError> {
+        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(&notice);
+        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<ElasticSection, GroupDemand>,
+    /// User-content counts for shared framing; native messages have no wrapper.
+    content_counts: Vec<Option<usize>>,
+}
+
+fn collect_demand(context: &ComposedContext) -> Result<ContextDemand, SectionError> {
+    let mut required_tokens = 0usize;
+    let mut demands = BTreeMap::<ElasticSection, GroupDemand>::new();
+    let mut content_counts = Vec::with_capacity(context.sections.len());
+    for (section_index, section) in context.sections.iter().enumerate() {
+        let required_count = match &section.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 &section.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::<usize>();
+    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 &section.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::<Vec<_>>();
+                crate::enforcement::retain_content(
+                    section,
+                    &mut context.truncations,
+                    |index, _| keep[index],
+                );
+            }
+            SectionDelivery::Message(_) => {}
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "allocation_tests.rs"]
+mod tests;

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"));
+    }
+}

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<ElasticSection, ElasticSectionLimits>,
 }
 
+impl ElasticBudgetConfig {
+    pub(crate) fn section_limits(&self, section: ElasticSection) -> ElasticSectionLimits {
+        self.sections.get(&section).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 {

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<TruncationObservation>,
     mut retain: impl FnMut(usize, &mut SectionContent) -> bool,

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;