Improve Guardian transcript context selection (#38558)

## What changed

- Replace the single byte limit with token-derived per-entry and separate message and tool transcript budgets.
- Preserve the first and latest user messages, then fill the remaining budgets with recent user, assistant, and tool context while limiting retained non-user entries.
- Truncate oversized entries in the middle without splitting UTF-8 characters, and report the approximate omitted token count.
- Exclude reasoning from the default transcript sources while retaining support when it is explicitly configured.

## Testing

- Cover oversized message and tool-result truncation, user-message retention, separate tool budgeting, configured reasoning, and the resulting Guardian sampling request.

GitOrigin-RevId: aaeaa5fb5d902b5a1581a8a819bf34724bde794b
This commit is contained in:
felixxia-oai
2026-08-14 11:12:33 +00:00
committed by copyberry
parent db5b4a4cfb
commit 486df09a00
3 changed files with 481 additions and 69 deletions

View File

@@ -25,6 +25,7 @@ use core_test_support::responses;
use core_test_support::responses::ev_assistant_message;
use core_test_support::responses::ev_completed;
use core_test_support::skip_if_no_network;
use core_test_support::test_codex::TestCodex;
use core_test_support::test_codex::test_codex;
use pretty_assertions::assert_eq;
use serde_json::json;
@@ -37,13 +38,11 @@ impl ConversationHistorySnapshot for TestConversationHistory {
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result<()> {
skip_if_no_network!(Ok(()));
async fn sample_conversation_history(
conversation_history: Vec<ResponseItem>,
) -> Result<(serde_json::Value, TestCodex)> {
let thread_server = responses::start_mock_server().await;
let test = test_codex().build_with_auto_env(&thread_server).await?;
let thread_id = test.session_configured.thread_id;
let events = vec![
ev_assistant_message("sample", r#"{"scores":{"action_risk":0.25}}"#),
ev_completed("response-1"),
@@ -83,7 +82,35 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result<
let tool_payload = ToolPayload::Function {
arguments: r#"{"path":"README.md"}"#.to_owned(),
};
let conversation_history = TestConversationHistory(vec![
let conversation_history = TestConversationHistory(conversation_history);
registry.tool_lifecycle_contributors()[0]
.on_tool_start(ToolStartInput {
session_store: &session_store,
thread_store,
turn_store: &turn_store,
turn_id: "turn-1",
call_id: "call-1",
tool_name: &tool_name,
payload: &tool_payload,
conversation_history: Arc::new(conversation_history),
source: ToolCallSource::Direct,
})
.await;
let request = tokio::time::timeout(
Duration::from_secs(5),
server.wait_for_request(/*connection_index*/ 1, /*request_index*/ 0),
)
.await?;
Ok((request.body_json(), test))
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result<()> {
skip_if_no_network!(Ok(()));
let (request, test) = sample_conversation_history(vec![
ResponseItem::Message {
id: None,
role: "user".to_owned(),
@@ -126,28 +153,10 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result<
call_id: "call-1".to_owned(),
internal_chat_message_metadata_passthrough: None,
},
]);
registry.tool_lifecycle_contributors()[0]
.on_tool_start(ToolStartInput {
session_store: &session_store,
thread_store,
turn_store: &turn_store,
turn_id: "turn-1",
call_id: "call-1",
tool_name: &tool_name,
payload: &tool_payload,
conversation_history: Arc::new(conversation_history),
source: ToolCallSource::Direct,
})
.await;
let request = tokio::time::timeout(
Duration::from_secs(5),
server.wait_for_request(/*connection_index*/ 1, /*request_index*/ 0),
)
])
.await?;
let request = request.body_json();
let thread_id = test.session_configured.thread_id;
let thread_store = test.codex.thread_extension_data();
assert_eq!(request["model"], "gpt-5.6-luna");
assert_eq!(
request["client_metadata"]["thread_id"],
@@ -166,10 +175,9 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result<
json!([
{"type": "input_text", "text": ">>> TRANSCRIPT START\n"},
{"type": "input_text", "text": "[1] user: Inspect the repository guidelines.\n"},
{"type": "input_text", "text": "[2] reasoning: Find the repository documentation.\n"},
{"type": "input_text", "text": "[3] tool list_dir call: {\"path\":\".\"}\n"},
{"type": "input_text", "text": "[4] tool list_dir result: README.md\n"},
{"type": "input_text", "text": "[5] tool read_file call: {\"path\":\"README.md\"}\n"},
{"type": "input_text", "text": "[2] tool list_dir call: {\"path\":\".\"}\n"},
{"type": "input_text", "text": "[3] tool list_dir result: README.md\n"},
{"type": "input_text", "text": "[4] tool read_file call: {\"path\":\"README.md\"}\n"},
{"type": "input_text", "text": ">>> TRANSCRIPT END\n\n"},
{
"type": "input_text",
@@ -217,3 +225,79 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result<
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn contributor_sends_compacted_conversation_history_to_luna() -> Result<()> {
skip_if_no_network!(Ok(()));
let mut history = (0..8)
.map(|index| ResponseItem::Message {
id: None,
role: "user".to_owned(),
content: vec![ContentItem::InputText {
text: format!("user turn {index}: {}", "authorization ".repeat(1_000)),
}],
phase: None,
internal_chat_message_metadata_passthrough: None,
})
.collect::<Vec<_>>();
history.extend((0..12).flat_map(|index| {
let call_id = format!("call-{index}");
[
ResponseItem::FunctionCall {
id: None,
name: "exec_command".to_owned(),
namespace: None,
arguments: format!("tool evidence {index}: {}", "signal ".repeat(1_000)),
encrypted_function_args: None,
call_id: call_id.clone(),
internal_chat_message_metadata_passthrough: None,
},
ResponseItem::FunctionCallOutput {
id: None,
call_id,
output: FunctionCallOutputPayload::from_text(format!(
"result evidence {index}: {}",
"signal ".repeat(1_000)
)),
internal_chat_message_metadata_passthrough: None,
},
]
}));
let (request, _test) = sample_conversation_history(history).await?;
let content = request["input"][2]["content"]
.as_array()
.expect("Luna request should contain separate transcript text items");
let entries = content
.iter()
.filter_map(|entry| entry["text"].as_str())
.collect::<Vec<_>>();
assert!(entries.iter().any(|entry| entry.contains("user turn 0:")));
assert!(entries.iter().any(|entry| entry.contains("user turn 7:")));
assert!(!entries.iter().any(|entry| entry.contains("user turn 1:")));
assert!(
entries
.iter()
.any(|entry| entry.contains("tool exec_command call: tool evidence 11:"))
);
assert!(
entries
.iter()
.any(|entry| entry.contains("tool exec_command result: result evidence 11:"))
);
assert!(
!entries
.iter()
.any(|entry| entry.contains("tool evidence 0:"))
);
assert!(
!entries
.iter()
.any(|entry| entry.contains("result evidence 0:"))
);
assert!(entries.iter().any(|entry| entry.contains("<truncated")));
Ok(())
}

View File

@@ -1,17 +1,33 @@
use std::collections::HashMap;
use std::collections::VecDeque;
use codex_extension_api::ResponseItem;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ReasoningItemContent;
use codex_protocol::models::ReasoningItemReasoningSummary;
use codex_protocol::models::plaintext_agent_message_content;
use codex_protocol::protocol::TruncationPolicy;
// Provisional approximation of an 80k-token transcript budget.
const MAX_TRANSCRIPT_BYTES: usize = 320 * 1024;
pub(crate) const MAX_MESSAGE_ENTRY_TOKENS: usize = 2_000;
const MAX_TOOL_ENTRY_TOKENS: usize = 1_000;
const MAX_MESSAGE_TRANSCRIPT_TOKENS: usize = 10_000;
const MAX_TOOL_TRANSCRIPT_TOKENS: usize = 10_000;
const MAX_RECENT_NON_USER_ENTRIES: usize = 40;
const MANUAL_APPROVAL_DEVELOPER_PREFIX: &str =
"The user has manually approved a specific action that was previously `Rejected`.";
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum TranscriptEntryKind {
User,
Message,
Tool,
}
struct TranscriptEntry {
kind: TranscriptEntryKind,
text: String,
tokens: usize,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum TranscriptSource {
ToolCalls,
@@ -27,11 +43,7 @@ pub(crate) struct TranscriptConfig {
impl Default for TranscriptConfig {
fn default() -> Self {
Self {
sources: vec![
TranscriptSource::ToolCalls,
TranscriptSource::ToolOutputs,
TranscriptSource::Reasoning,
],
sources: vec![TranscriptSource::ToolCalls, TranscriptSource::ToolOutputs],
}
}
}
@@ -41,10 +53,8 @@ impl TranscriptConfig {
&self,
items: impl IntoIterator<Item = &'a ResponseItem>,
) -> Vec<String> {
let mut transcript = VecDeque::new();
let mut transcript_bytes = 0;
let mut entries = Vec::new();
let mut tool_names_by_call_id = HashMap::new();
let mut entry_number = 0;
for item in items {
let (role, text) = match item {
@@ -177,34 +187,126 @@ impl TranscriptConfig {
if text.trim().is_empty() {
continue;
}
entry_number += 1;
let entry = format!("[{entry_number}] {role}: {text}\n");
transcript_bytes += entry.len();
transcript.push_back(entry);
while transcript_bytes > MAX_TRANSCRIPT_BYTES {
let Some(first_entry) = transcript.front_mut() else {
break;
};
let bytes_to_remove = transcript_bytes - MAX_TRANSCRIPT_BYTES;
if first_entry.len() <= bytes_to_remove {
transcript_bytes -= first_entry.len();
transcript.pop_front();
} else {
let mut first_retained_byte = bytes_to_remove;
while !first_entry.is_char_boundary(first_retained_byte) {
first_retained_byte += 1;
}
first_entry.drain(..first_retained_byte);
transcript_bytes -= first_retained_byte;
let kind = match role.as_str() {
"user" => TranscriptEntryKind::User,
role if role.starts_with("tool ") => TranscriptEntryKind::Tool,
_ => TranscriptEntryKind::Message,
};
let token_cap = match kind {
TranscriptEntryKind::Tool => MAX_TOOL_ENTRY_TOKENS,
TranscriptEntryKind::User | TranscriptEntryKind::Message => {
MAX_MESSAGE_ENTRY_TOKENS
}
};
let text = truncate_entry(&text, token_cap);
let entry_number = entries.len() + 1;
let text = format!("[{entry_number}] {role}: {text}\n");
let tokens = TruncationPolicy::Bytes(text.len()).token_budget();
entries.push(TranscriptEntry { kind, text, tokens });
}
let mut included = vec![false; entries.len()];
let mut message_tokens = 0;
let mut tool_tokens = 0;
let user_indices = entries
.iter()
.enumerate()
.filter_map(|(index, entry)| (entry.kind == TranscriptEntryKind::User).then_some(index))
.collect::<Vec<_>>();
if let Some(&first_user_index) = user_indices.first() {
included[first_user_index] = true;
message_tokens += entries[first_user_index].tokens;
}
if let Some(&latest_user_index) = user_indices.last()
&& !included[latest_user_index]
&& message_tokens + entries[latest_user_index].tokens <= MAX_MESSAGE_TRANSCRIPT_TOKENS
{
included[latest_user_index] = true;
message_tokens += entries[latest_user_index].tokens;
}
for &index in user_indices.iter().rev() {
if included[index]
|| message_tokens + entries[index].tokens > MAX_MESSAGE_TRANSCRIPT_TOKENS
{
continue;
}
included[index] = true;
message_tokens += entries[index].tokens;
}
let mut retained_non_user_entries = 0;
for (index, entry) in entries.iter().enumerate().rev() {
if entry.kind == TranscriptEntryKind::User
|| retained_non_user_entries >= MAX_RECENT_NON_USER_ENTRIES
{
continue;
}
let fits_budget = match entry.kind {
TranscriptEntryKind::Tool => {
tool_tokens + entry.tokens <= MAX_TOOL_TRANSCRIPT_TOKENS
}
TranscriptEntryKind::Message => {
message_tokens + entry.tokens <= MAX_MESSAGE_TRANSCRIPT_TOKENS
}
TranscriptEntryKind::User => unreachable!("user entries were selected separately"),
};
if !fits_budget {
continue;
}
included[index] = true;
retained_non_user_entries += 1;
match entry.kind {
TranscriptEntryKind::Tool => tool_tokens += entry.tokens,
TranscriptEntryKind::Message => message_tokens += entry.tokens,
TranscriptEntryKind::User => unreachable!("user entries were selected separately"),
}
}
transcript.into()
entries
.into_iter()
.enumerate()
.filter_map(|(index, entry)| included[index].then_some(entry.text))
.collect()
}
}
pub(crate) fn truncate_entry(text: &str, max_tokens: usize) -> String {
let max_bytes = TruncationPolicy::Tokens(max_tokens).byte_budget();
if text.len() <= max_bytes {
return text.to_owned();
}
let omitted_tokens =
TruncationPolicy::Bytes(text.len().saturating_sub(max_bytes)).token_budget();
let marker = format!("<truncated omitted_approx_tokens=\"{omitted_tokens}\" />");
if max_bytes <= marker.len() {
return marker;
}
let available_bytes = max_bytes - marker.len();
let prefix_bytes = available_bytes / 2;
let suffix_bytes = available_bytes - prefix_bytes;
let mut prefix_end = prefix_bytes;
while !text.is_char_boundary(prefix_end) {
prefix_end -= 1;
}
let mut suffix_start = text.len() - suffix_bytes;
while !text.is_char_boundary(suffix_start) {
suffix_start += 1;
}
format!("{}{marker}{}", &text[..prefix_end], &text[suffix_start..])
}
#[cfg(test)]
#[path = "transcript_tests.rs"]
mod tests;

View File

@@ -5,12 +5,17 @@ use codex_protocol::models::FunctionCallOutputContentItem;
use codex_protocol::models::FunctionCallOutputPayload;
use codex_protocol::models::ReasoningItemContent;
use codex_protocol::models::ReasoningItemReasoningSummary;
use codex_protocol::protocol::TruncationPolicy;
use pretty_assertions::assert_eq;
use super::MANUAL_APPROVAL_DEVELOPER_PREFIX;
use super::MAX_TRANSCRIPT_BYTES;
use super::MAX_MESSAGE_ENTRY_TOKENS;
use super::MAX_MESSAGE_TRANSCRIPT_TOKENS;
use super::MAX_TOOL_ENTRY_TOKENS;
use super::MAX_TOOL_TRANSCRIPT_TOKENS;
use super::TranscriptConfig;
use super::TranscriptSource;
use super::truncate_entry;
#[test]
fn transcript_keeps_conversation_and_configured_sources() {
@@ -66,7 +71,6 @@ fn transcript_keeps_conversation_and_configured_sources() {
"[1] user: Inspect the workspace.\n",
"[2] tool exec_command call: {}\n",
"[3] tool exec_command result: Workspace inspected.\n",
"[4] reasoning: Review the current files.\nPlaintext reasoning.\n",
]
);
@@ -99,13 +103,19 @@ fn transcript_keeps_conversation_and_configured_sources() {
}
#[test]
fn transcript_retains_the_most_recent_bounded_content() {
fn transcript_truncates_oversized_entries_without_splitting_characters() {
let prefix = "start é";
let suffix = "é end";
let oversized_message = format!(
"{prefix}{}{suffix}",
"é".repeat(TruncationPolicy::Tokens(MAX_MESSAGE_ENTRY_TOKENS).byte_budget())
);
let items = vec![
ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: "é".repeat(MAX_TRANSCRIPT_BYTES),
text: oversized_message,
}],
phase: None,
internal_chat_message_metadata_passthrough: None,
@@ -123,12 +133,227 @@ fn transcript_retains_the_most_recent_bounded_content() {
let transcript = TranscriptConfig::default().build(&items);
assert!(transcript.iter().map(String::len).sum::<usize>() <= MAX_TRANSCRIPT_BYTES);
assert_eq!(transcript.len(), 2);
let user_entry = &transcript[0];
assert!(user_entry.starts_with(&format!("[1] user: {prefix}")));
assert!(user_entry.contains("<truncated omitted_approx_tokens=\""));
assert!(user_entry.ends_with(&format!("{suffix}\n")));
assert!(
user_entry.len()
<= TruncationPolicy::Tokens(MAX_MESSAGE_ENTRY_TOKENS).byte_budget()
+ "[1] user: \n".len()
);
assert_eq!(transcript[1], "[2] assistant: latest response\n");
}
#[test]
fn transcript_preserves_first_and_latest_user_messages_and_recent_history() {
let oversized_user_message = "authorization ".repeat(1_000);
let mut items = (0..8)
.map(|index| ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: format!("user turn {index}: {oversized_user_message}"),
}],
phase: None,
internal_chat_message_metadata_passthrough: None,
})
.collect::<Vec<_>>();
items.push(ResponseItem::Message {
id: None,
role: "assistant".to_string(),
content: vec![ContentItem::OutputText {
text: "Most recent assistant context.".to_string(),
}],
phase: None,
internal_chat_message_metadata_passthrough: None,
});
let transcript = TranscriptConfig::default().build(&items);
assert!(transcript[0].starts_with("[1] user: user turn 0:"));
assert!(
transcript
.iter()
.any(|entry| entry.contains("latest response"))
.any(|entry| entry.starts_with("[8] user: user turn 7:"))
);
assert!(
!transcript
.iter()
.any(|entry| entry.contains("user turn 1:"))
);
assert!(
transcript
.iter()
.any(|entry| entry.contains("user turn 6:"))
);
assert_eq!(
transcript.last().map(String::as_str),
Some("[9] assistant: Most recent assistant context.\n")
);
assert!(
transcript
.iter()
.map(|entry| TruncationPolicy::Bytes(entry.len()).token_budget())
.sum::<usize>()
<= MAX_MESSAGE_TRANSCRIPT_TOKENS
);
}
#[test]
fn transcript_reserves_separate_budget_for_recent_tool_evidence() {
let mut items = (0..8)
.map(|index| ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: format!("user turn {index}: {}", "authorization ".repeat(1_000)),
}],
phase: None,
internal_chat_message_metadata_passthrough: None,
})
.collect::<Vec<_>>();
items.extend((0..12).map(|index| ResponseItem::FunctionCall {
id: None,
name: "exec_command".to_string(),
namespace: None,
arguments: format!("tool evidence {index}: {}", "signal ".repeat(1_000)),
encrypted_function_args: None,
call_id: format!("call-{index}"),
internal_chat_message_metadata_passthrough: None,
}));
let transcript = TranscriptConfig::default().build(&items);
assert!(transcript[0].contains("user turn 0:"));
assert!(
transcript
.iter()
.any(|entry| entry.contains("user turn 7:"))
);
assert!(
transcript.last().is_some_and(
|entry| entry.starts_with("[20] tool exec_command call: tool evidence 11:")
)
);
assert!(
!transcript
.iter()
.any(|entry| entry.contains("tool evidence 0:"))
);
let tool_entries = transcript
.iter()
.filter(|entry| entry.contains("tool exec_command call:"))
.collect::<Vec<_>>();
assert!(tool_entries.len() < 12);
assert!(tool_entries.iter().all(|entry| {
entry.len()
<= TruncationPolicy::Tokens(MAX_TOOL_ENTRY_TOKENS).byte_budget()
+ "[20] tool exec_command call: \n".len()
}));
assert!(
tool_entries
.iter()
.map(|entry| TruncationPolicy::Bytes(entry.len()).token_budget())
.sum::<usize>()
<= MAX_TOOL_TRANSCRIPT_TOKENS
);
}
#[test]
fn transcript_truncates_tool_results_using_standard_budget() {
let output = format!("first {} last", "evidence ".repeat(5_000));
let items = vec![
ResponseItem::FunctionCall {
id: None,
name: "node_repl__run".to_owned(),
namespace: None,
arguments: "{}".to_owned(),
encrypted_function_args: None,
call_id: "call-1".to_owned(),
internal_chat_message_metadata_passthrough: None,
},
ResponseItem::FunctionCallOutput {
id: None,
call_id: "call-1".to_owned(),
output: FunctionCallOutputPayload::from_text(output),
internal_chat_message_metadata_passthrough: None,
},
];
let transcript = TranscriptConfig::default().build(&items);
let result = transcript
.iter()
.find(|entry| entry.contains(" result: "))
.expect("the tool result should be retained");
let prefix = "[2] tool node_repl__run result: ";
assert!(result.starts_with(&format!("{prefix}first ")));
assert!(result.contains("<truncated omitted_approx_tokens=\""));
assert!(result.ends_with(" last\n"));
assert!(
result.len()
<= TruncationPolicy::Tokens(MAX_TOOL_ENTRY_TOKENS).byte_budget() + prefix.len() + 1
);
}
#[test]
fn configured_reasoning_counts_against_message_budget() {
let mut items = (0..8)
.map(|index| ResponseItem::Message {
id: None,
role: "user".to_string(),
content: vec![ContentItem::InputText {
text: format!("user turn {index}: {}", "authorization ".repeat(1_000)),
}],
phase: None,
internal_chat_message_metadata_passthrough: None,
})
.collect::<Vec<_>>();
items.push(ResponseItem::Reasoning {
id: None,
summary: vec![ReasoningItemReasoningSummary::SummaryText {
text: "Recent reasoning evidence.".to_string(),
}],
content: None,
encrypted_content: None,
internal_chat_message_metadata_passthrough: None,
});
let transcript = TranscriptConfig {
sources: vec![TranscriptSource::Reasoning],
}
.build(&items);
assert!(transcript[0].contains("user turn 0:"));
assert!(
transcript
.iter()
.any(|entry| entry.contains("user turn 7:"))
);
assert_eq!(
transcript.last().map(String::as_str),
Some("[9] reasoning: Recent reasoning evidence.\n")
);
assert!(
transcript
.iter()
.map(|entry| TruncationPolicy::Bytes(entry.len()).token_budget())
.sum::<usize>()
<= MAX_MESSAGE_TRANSCRIPT_TOKENS
);
}
#[test]
fn truncate_entry_preserves_prefix_suffix_and_utf8_boundaries() {
let text = format!("prefix é{}é suffix", "🦀".repeat(2_000));
let truncated = truncate_entry(&text, /*max_tokens*/ 200);
assert!(truncated.starts_with("prefix é"));
assert!(truncated.contains("<truncated omitted_approx_tokens=\""));
assert!(truncated.ends_with("é suffix"));
assert!(truncated.len() <= TruncationPolicy::Tokens(200).byte_budget());
}
#[test]
@@ -164,7 +389,8 @@ fn transcript_keeps_only_manual_approval_developer_messages() {
#[test]
fn transcript_omits_media_payloads_and_keeps_readable_content() {
let oversized_image = "A".repeat(MAX_TRANSCRIPT_BYTES + 1);
let oversized_image =
"A".repeat(TruncationPolicy::Tokens(MAX_MESSAGE_TRANSCRIPT_TOKENS).byte_budget() + 1);
let items = vec![
ResponseItem::Message {
id: None,