mirror of
https://github.com/openai/codex.git
synced 2026-09-29 16:57:06 +08:00
Add aggregate budget enforcement for Guardian context (#44166)
## What changed Add `ComposedContext::enforce_budget` to fit evidence within the input token allowance after reserving existing context. Preserve required content and message boundaries, reserve an omission notice, and return an error if required evidence cannot fit. Carry retention policies through transcript rendering and composition. Protect user messages, protected messages, and the newest five tool entries. Remove oversized optional items first, then evict commentary, older tool evidence, and images in priority order. Add image admission support and record omitted content in truncation observations. ## Testing Add tests for existing-context reservations, required-message preservation, image omission, framing costs, and protection of the newest five tool entries in both context profiles. GitOrigin-RevId: 8a0f245bad82f562323b4533442385ddbcb6e541
This commit is contained in:
@@ -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) {
|
||||
(
|
||||
"<guardian_context_omission>",
|
||||
"</guardian_context_omission>",
|
||||
)
|
||||
}
|
||||
|
||||
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()
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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(
|
||||
"<no retained transcript entries>".to_owned(),
|
||||
));
|
||||
}
|
||||
(
|
||||
transcript
|
||||
.items
|
||||
.push("<no retained transcript entries>".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
|
||||
|
||||
@@ -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<Item = (&'static str, SectionCost)> + '_ {
|
||||
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 {
|
||||
|
||||
@@ -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
|
||||
]),
|
||||
}],
|
||||
|
||||
@@ -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<String>,
|
||||
pub items: Vec<Budgeted<String>>,
|
||||
pub omission_note: Option<String>,
|
||||
pub truncations: Vec<TruncationObservation>,
|
||||
}
|
||||
@@ -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<ContentItem>),
|
||||
UserContent(Vec<Budgeted<ContentItem>>),
|
||||
Message(Box<ResponseItem>),
|
||||
}
|
||||
|
||||
/// 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<SectionOutput>,
|
||||
pub truncations: Vec<TruncationObservation>,
|
||||
@@ -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<String>) -> 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)));
|
||||
|
||||
@@ -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<T> {
|
||||
pub content: T,
|
||||
pub(crate) retention: Retention,
|
||||
}
|
||||
|
||||
impl<T> Budgeted<T> {
|
||||
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<T> std::fmt::Debug for Budgeted<T> {
|
||||
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<ImageDetail>) -> 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<Self, SectionError> {
|
||||
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<TruncationObservation>,
|
||||
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;
|
||||
@@ -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<ContentItem>| 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")])]
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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::<Vec<_>>();
|
||||
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<_>>(),
|
||||
vec![
|
||||
Retention::Optional(BudgetPriority::Tool),
|
||||
Retention::Optional(BudgetPriority::Tool),
|
||||
Retention::Optional(BudgetPriority::Tool),
|
||||
Retention::Required,
|
||||
Retention::Required,
|
||||
Retention::Required,
|
||||
Retention::Required,
|
||||
Retention::Required
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user