mirror of
https://github.com/openai/codex.git
synced 2026-10-20 00:18:49 +00:00
Optimize Guardian transcript window retention (#40597)
## What changed - Track protected messages, ordinary messages, and tool entries in separate pools with cached token totals. - Evict the oldest eligible evidence across pools while preserving transcript token and entry limits and the reserved recent tool entries. ## Testing - Add regression coverage that a small, tool-heavy transcript window retains the latest final assistant reply. GitOrigin-RevId: 7aa78d0b0609d8e45cb82e4f7773c0877a4d8224
This commit is contained in:
felixxia-oai
@@ -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<usize>,
|
||||
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<usize>,
|
||||
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<TranscriptSource>,
|
||||
@@ -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<usize> = 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;
|
||||
}
|
||||
|
||||
|
||||
@@ -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<usize>,
|
||||
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<Item = usize> {
|
||||
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"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user