diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs b/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs index eb89ec23db..248490838f 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs @@ -13,14 +13,16 @@ use codex_protocol::models::plaintext_agent_message_content; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::TruncationPolicy; +use self::window::TranscriptWindow; use super::truncation::TruncationObservation; +mod window; + pub(crate) const MAX_MESSAGE_ENTRY_TOKENS: usize = 2_000; pub(crate) const MAX_TOOL_ENTRY_TOKENS: usize = 1_000; pub(crate) const MAX_MESSAGE_TRANSCRIPT_TOKENS: usize = 10_000; pub(crate) const MAX_TOOL_TRANSCRIPT_TOKENS: usize = 10_000; pub(crate) const MAX_RECENT_NON_USER_ENTRIES: usize = 40; -const MIN_RECENT_TOOL_ENTRIES: usize = 5; const MAX_TRANSCRIPT_IMAGES: usize = 4; const MAX_TRANSCRIPT_IMAGE_BYTES: usize = 8 * 1024 * 1024; const MANUAL_APPROVAL_DEVELOPER_PREFIX: &str = @@ -52,44 +54,6 @@ pub(crate) struct RenderedImages { pub(crate) omitted_bytes: usize, } -fn retained_tokens( - entries: &[TranscriptEntry], - retained: &VecDeque, - kind: TranscriptEntryKind, -) -> usize { - retained - .iter() - .filter_map(|&index| { - let entry = &entries[index]; - let same_token_pool = match kind { - TranscriptEntryKind::ProtectedMessage | TranscriptEntryKind::Message => matches!( - entry.kind, - TranscriptEntryKind::ProtectedMessage | TranscriptEntryKind::Message - ), - TranscriptEntryKind::User | TranscriptEntryKind::Tool => entry.kind == kind, - }; - same_token_pool.then_some(entry.tokens) - }) - .sum() -} - -fn evict_oldest_matching( - entries: &[TranscriptEntry], - retained: &mut VecDeque, - count: usize, - mut matches: impl FnMut(TranscriptEntryKind) -> bool, -) { - let mut remaining = count; - retained.retain(|&index| { - if remaining > 0 && matches(entries[index].kind) { - remaining -= 1; - false - } else { - true - } - }); -} - #[derive(Clone, Debug, Eq, PartialEq)] pub(crate) struct TranscriptConfig { pub(crate) sources: Vec, @@ -433,157 +397,12 @@ impl TranscriptConfig { let available_message_tokens = self .max_message_transcript_tokens .saturating_sub(message_tokens); - let mut retained_non_user: VecDeque = VecDeque::new(); - - // Replay non-user entries as an append-only bounded buffer. Evict ordinary evidence - // before protected conversation messages, removing half of whichever pool overflows so - // the retained prefix stays stable between overflow points without storing sampler state. - for (index, entry) in entries.iter().enumerate() { - if entry.kind == TranscriptEntryKind::User { - continue; - } - - let token_budget = match entry.kind { - TranscriptEntryKind::Tool => self.max_tool_transcript_tokens, - TranscriptEntryKind::ProtectedMessage | TranscriptEntryKind::Message => { - available_message_tokens - } - TranscriptEntryKind::User => unreachable!("user entries were selected separately"), - }; - - let protected_tokens = if entry.kind == TranscriptEntryKind::Message { - retained_non_user - .iter() - .filter_map(|&retained_index| { - let retained = &entries[retained_index]; - (retained.kind == TranscriptEntryKind::ProtectedMessage) - .then_some(retained.tokens) - }) - .sum() - } else { - 0 - }; - - // An entry that cannot fit beside its non-evictable context must not evict evidence. - if protected_tokens + entry.tokens > token_budget - || self.max_recent_non_user_entries == 0 - { - continue; - } - - let eviction_order: &[TranscriptEntryKind] = match entry.kind { - TranscriptEntryKind::ProtectedMessage => &[ - TranscriptEntryKind::Message, - TranscriptEntryKind::ProtectedMessage, - ], - TranscriptEntryKind::Message => &[TranscriptEntryKind::Message], - TranscriptEntryKind::Tool => &[TranscriptEntryKind::Tool], - TranscriptEntryKind::User => unreachable!("user entries were selected separately"), - }; - - for &eviction_kind in eviction_order { - if retained_tokens(&entries, &retained_non_user, entry.kind) - .saturating_add(entry.tokens) - <= token_budget - { - break; - } - - let retained_kind_count = retained_non_user - .iter() - .filter(|&&retained_index| entries[retained_index].kind == eviction_kind) - .count(); - if retained_kind_count == 0 { - continue; - } - - evict_oldest_matching( - &entries, - &mut retained_non_user, - retained_kind_count.div_ceil(2), - |kind| kind == eviction_kind, - ); - while retained_tokens(&entries, &retained_non_user, entry.kind) - .saturating_add(entry.tokens) - > token_budget - && retained_non_user - .iter() - .any(|&retained_index| entries[retained_index].kind == eviction_kind) - { - evict_oldest_matching( - &entries, - &mut retained_non_user, - /*count*/ 1, - |kind| kind == eviction_kind, - ); - } - } - - if retained_tokens(&entries, &retained_non_user, entry.kind) - .saturating_add(entry.tokens) - > token_budget - { - continue; - } - - if retained_non_user.len().saturating_add(1) > self.max_recent_non_user_entries { - let retained_tool_count = retained_non_user - .iter() - .filter(|&&retained_index| { - entries[retained_index].kind == TranscriptEntryKind::Tool - }) - .count(); - // Keep one slot available for protected context in smaller configured windows. - let minimum_existing_tool_entries = MIN_RECENT_TOOL_ENTRIES - .min(self.max_recent_non_user_entries.saturating_sub(1)) - .saturating_sub(usize::from(entry.kind == TranscriptEntryKind::Tool)); - let mut removable_tool_count = - retained_tool_count.saturating_sub(minimum_existing_tool_entries); - let expendable_message_count = retained_non_user - .iter() - .filter(|&&retained_index| { - entries[retained_index].kind == TranscriptEntryKind::Message - }) - .count(); - let ordinary_evidence_count = expendable_message_count + retained_tool_count; - let removable_count = expendable_message_count + removable_tool_count; - - if removable_count > 0 { - evict_oldest_matching( - &entries, - &mut retained_non_user, - ordinary_evidence_count.div_ceil(2).min(removable_count), - |kind| match kind { - TranscriptEntryKind::Message => true, - TranscriptEntryKind::Tool if removable_tool_count > 0 => { - removable_tool_count -= 1; - true - } - TranscriptEntryKind::User - | TranscriptEntryKind::ProtectedMessage - | TranscriptEntryKind::Tool => false, - }, - ); - } else if matches!( - entry.kind, - TranscriptEntryKind::ProtectedMessage | TranscriptEntryKind::Tool - ) { - let entries_to_evict = retained_non_user.len().div_ceil(2); - evict_oldest_matching( - &entries, - &mut retained_non_user, - entries_to_evict, - |kind| kind == TranscriptEntryKind::ProtectedMessage, - ); - } else { - continue; - } - } - - retained_non_user.push_back(index); + let mut window = TranscriptWindow::new(&entries, self, available_message_tokens); + for index in 0..entries.len() { + window.insert(index); } - for index in retained_non_user { + for index in window.into_indices() { included[index] = true; } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/transcript/window.rs b/codex-rs/ext/guardian-v2/src/async_scorer/transcript/window.rs new file mode 100644 index 0000000000..fbb6eaa31b --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/transcript/window.rs @@ -0,0 +1,223 @@ +use std::collections::VecDeque; + +use super::TranscriptConfig; +use super::TranscriptEntry; +use super::TranscriptEntryKind; + +const MIN_RECENT_TOOL_ENTRIES: usize = 5; + +#[derive(Default)] +struct EntryPool { + indices: VecDeque, + tokens: usize, +} + +impl EntryPool { + fn push(&mut self, index: usize, tokens: usize) { + self.indices.push_back(index); + self.tokens += tokens; + } + + fn pop_oldest(&mut self, entries: &[TranscriptEntry]) { + if let Some(index) = self.indices.pop_front() { + self.tokens -= entries[index].tokens; + } + } + + fn evict_oldest(&mut self, count: usize, entries: &[TranscriptEntry]) { + for _ in 0..count.min(self.indices.len()) { + self.pop_oldest(entries); + } + } +} + +pub(super) struct TranscriptWindow<'a> { + entries: &'a [TranscriptEntry], + message_token_budget: usize, + tool_token_budget: usize, + max_entries: usize, + protected_messages: EntryPool, + ordinary_messages: EntryPool, + tools: EntryPool, +} + +impl<'a> TranscriptWindow<'a> { + pub(super) fn new( + entries: &'a [TranscriptEntry], + config: &TranscriptConfig, + message_token_budget: usize, + ) -> Self { + Self { + entries, + message_token_budget, + tool_token_budget: config.max_tool_transcript_tokens, + max_entries: config.max_recent_non_user_entries, + protected_messages: EntryPool::default(), + ordinary_messages: EntryPool::default(), + tools: EntryPool::default(), + } + } + + pub(super) fn insert(&mut self, index: usize) { + let kind = self.entries[index].kind; + let tokens = self.entries[index].tokens; + + if kind == TranscriptEntryKind::User + || self.max_entries == 0 + || !self.make_room_for_token_budget(kind, tokens) + || !self.make_room_for_entry_limit(kind) + { + return; + } + + self.pool_mut(kind).push(index, tokens); + } + + pub(super) fn into_indices(self) -> impl Iterator { + self.protected_messages + .indices + .into_iter() + .chain(self.ordinary_messages.indices) + .chain(self.tools.indices) + } + + fn make_room_for_token_budget(&mut self, kind: TranscriptEntryKind, tokens: usize) -> bool { + let token_budget = match kind { + TranscriptEntryKind::ProtectedMessage | TranscriptEntryKind::Message => { + self.message_token_budget + } + TranscriptEntryKind::Tool => self.tool_token_budget, + TranscriptEntryKind::User => unreachable!("user entries were selected separately"), + }; + let protected_tokens = if kind == TranscriptEntryKind::Message { + self.protected_messages.tokens + } else { + 0 + }; + + // Rejected commentary must not evict the protected context it cannot fit beside. + if protected_tokens.saturating_add(tokens) > token_budget { + return false; + } + + let eviction_order: &[TranscriptEntryKind] = match kind { + TranscriptEntryKind::ProtectedMessage => &[ + TranscriptEntryKind::Message, + TranscriptEntryKind::ProtectedMessage, + ], + TranscriptEntryKind::Message => &[TranscriptEntryKind::Message], + TranscriptEntryKind::Tool => &[TranscriptEntryKind::Tool], + TranscriptEntryKind::User => unreachable!("user entries were selected separately"), + }; + let entries = self.entries; + + for &eviction_kind in eviction_order { + if self.has_token_capacity(kind, tokens, token_budget) { + break; + } + + let pool_len = self.pool(eviction_kind).indices.len(); + if pool_len == 0 { + continue; + } + + self.pool_mut(eviction_kind) + .evict_oldest(pool_len.div_ceil(2), entries); + while !self.has_token_capacity(kind, tokens, token_budget) + && !self.pool(eviction_kind).indices.is_empty() + { + self.pool_mut(eviction_kind).pop_oldest(entries); + } + } + + self.has_token_capacity(kind, tokens, token_budget) + } + + fn make_room_for_entry_limit(&mut self, kind: TranscriptEntryKind) -> bool { + let retained_count = self.protected_messages.indices.len() + + self.ordinary_messages.indices.len() + + self.tools.indices.len(); + if retained_count < self.max_entries { + return true; + } + + // Smaller windows still leave one slot available for protected context. + let minimum_existing_tools = MIN_RECENT_TOOL_ENTRIES + .min(self.max_entries.saturating_sub(1)) + .saturating_sub(usize::from(kind == TranscriptEntryKind::Tool)); + let removable_tools = self + .tools + .indices + .len() + .saturating_sub(minimum_existing_tools); + let removable_count = self.ordinary_messages.indices.len() + removable_tools; + + if removable_count > 0 { + let evidence_count = self.ordinary_messages.indices.len() + self.tools.indices.len(); + let eviction_count = evidence_count.div_ceil(2).min(removable_count); + + for _ in 0..eviction_count { + let oldest_tool = (self.tools.indices.len() > minimum_existing_tools) + .then(|| self.tools.indices.front()) + .flatten(); + let pool = match (self.ordinary_messages.indices.front(), oldest_tool) { + (Some(message), Some(tool)) if message < tool => &mut self.ordinary_messages, + (Some(_), Some(_)) | (None, Some(_)) => &mut self.tools, + (Some(_), None) => &mut self.ordinary_messages, + (None, None) => unreachable!("removable evidence was counted before eviction"), + }; + pool.pop_oldest(self.entries); + } + + return true; + } + + if matches!( + kind, + TranscriptEntryKind::ProtectedMessage | TranscriptEntryKind::Tool + ) { + let protected_count = self.protected_messages.indices.len(); + self.protected_messages + .evict_oldest(protected_count.div_ceil(2), self.entries); + return true; + } + + false + } + + fn has_token_capacity( + &self, + kind: TranscriptEntryKind, + tokens: usize, + token_budget: usize, + ) -> bool { + let retained_tokens = match kind { + TranscriptEntryKind::ProtectedMessage | TranscriptEntryKind::Message => self + .protected_messages + .tokens + .saturating_add(self.ordinary_messages.tokens), + TranscriptEntryKind::Tool => self.tools.tokens, + TranscriptEntryKind::User => unreachable!("user entries were selected separately"), + }; + + retained_tokens.saturating_add(tokens) <= token_budget + } + + fn pool(&self, kind: TranscriptEntryKind) -> &EntryPool { + match kind { + TranscriptEntryKind::ProtectedMessage => &self.protected_messages, + TranscriptEntryKind::Message => &self.ordinary_messages, + TranscriptEntryKind::Tool => &self.tools, + TranscriptEntryKind::User => unreachable!("user entries were selected separately"), + } + } + + fn pool_mut(&mut self, kind: TranscriptEntryKind) -> &mut EntryPool { + match kind { + TranscriptEntryKind::ProtectedMessage => &mut self.protected_messages, + TranscriptEntryKind::Message => &mut self.ordinary_messages, + TranscriptEntryKind::Tool => &mut self.tools, + TranscriptEntryKind::User => unreachable!("user entries were selected separately"), + } + } +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/transcript_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/transcript_tests.rs index c2e6ea3b7f..2901da9fe1 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/transcript_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/transcript_tests.rs @@ -432,6 +432,47 @@ fn transcript_evicts_protected_messages_in_cacheable_chunks() { } } +#[test] +fn transcript_preserves_latest_final_when_reserved_tools_fill_entry_window() { + let mut items = vec![ + user_message_item("Find a flight to New York."), + assistant_message("Earlier booking details.", MessagePhase::FinalAnswer), + assistant_message( + "I found a $450 flight. Should I book it?", + MessagePhase::FinalAnswer, + ), + user_message_item("Yes."), + ]; + items.extend((0..3).map(|index| ResponseItem::FunctionCall { + id: None, + name: "exec_command".to_string(), + namespace: None, + arguments: format!("recent evidence {index}"), + encrypted_function_args: None, + call_id: format!("call-{index}"), + internal_chat_message_metadata_passthrough: None, + })); + + let transcript = TranscriptConfig { + max_recent_non_user_entries: 4, + ..TranscriptConfig::default() + } + .build(&items) + .entries; + + assert_eq!( + transcript, + vec![ + "[1] user: Find a flight to New York.\n", + "[3] assistant: I found a $450 flight. Should I book it?\n", + "[4] user: Yes.\n", + "[5] tool exec_command call: recent evidence 0\n", + "[6] tool exec_command call: recent evidence 1\n", + "[7] tool exec_command call: recent evidence 2\n", + ] + ); +} + #[test] fn transcript_does_not_protect_legacy_inter_agent_instructions() { let communication = InterAgentCommunication::new(