diff --git a/codex-rs/ext/guardian-v2/src/extension_tests.rs b/codex-rs/ext/guardian-v2/src/extension_tests.rs index 72d029a2b4..352a67ce69 100644 --- a/codex-rs/ext/guardian-v2/src/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/extension_tests.rs @@ -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, +) -> 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::>(); + 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::>(); + + 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(" 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, ) -> Vec { - 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::>(); + + 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!(""); + 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; diff --git a/codex-rs/ext/guardian-v2/src/transcript_tests.rs b/codex-rs/ext/guardian-v2/src/transcript_tests.rs index 301bc9ab4e..7c49b65106 100644 --- a/codex-rs/ext/guardian-v2/src/transcript_tests.rs +++ b/codex-rs/ext/guardian-v2/src/transcript_tests.rs @@ -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::() <= 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(">(); + 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::() + <= 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::>(); + 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::>(); + 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::() + <= 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(">(); + 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::() + <= 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("