diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 16d7dd114a..e9811e7a57 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3345,6 +3345,7 @@ dependencies = [ "codex-core", "codex-extension-api", "codex-features", + "codex-guardian-context", "codex-history", "codex-http-client", "codex-login", diff --git a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs index e4051a8b2c..5e6298237d 100644 --- a/codex-rs/app-server/tests/suite/v2/guardian_v2.rs +++ b/codex-rs/app-server/tests/suite/v2/guardian_v2.rs @@ -61,6 +61,9 @@ use super::mcp_tool::TEST_TOOL_NAME; use super::mcp_tool::start_mcp_server; use super::mcp_tool::start_mcp_server_with_tools; +#[path = "guardian_v2_history_tests.rs"] +mod history; + const TIMEOUT: Duration = Duration::from_secs(30); const MODEL: &str = "mock-model"; const REQUIRED_MODEL: &str = "protected-model"; diff --git a/codex-rs/app-server/tests/suite/v2/guardian_v2_history_tests.rs b/codex-rs/app-server/tests/suite/v2/guardian_v2_history_tests.rs new file mode 100644 index 0000000000..19f8453031 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/guardian_v2_history_tests.rs @@ -0,0 +1,273 @@ +//! Exercises async review retention across app-server tool execution, compaction, and rollback. + +use std::sync::Arc; +use std::sync::Mutex; + +use anyhow::Result; +use app_test_support::ChatGptAuthFixture; +use app_test_support::MockResponsesConfig; +use app_test_support::TestAppServer; +use app_test_support::write_chatgpt_auth; +use axum::Json; +use axum::Router; +use axum::http::header; +use axum::routing::get; +use axum::routing::post; +use codex_app_server_protocol::ApprovalsReviewer; +use codex_app_server_protocol::AskForApproval; +use codex_app_server_protocol::ThreadCompactStartParams; +use codex_app_server_protocol::ThreadCompactStartResponse; +use codex_app_server_protocol::ThreadHistoryMode; +use codex_app_server_protocol::ThreadRollbackParams; +use codex_app_server_protocol::ThreadRollbackResponse; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::TurnCompletedNotification; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStatus; +use codex_app_server_protocol::UserInput; +use codex_config::types::AuthCredentialsStoreMode; +use codex_features::Feature; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use tempfile::TempDir; +use tokio::net::TcpListener; +use tokio::time::timeout; + +use super::MockResponsesState; +use super::TEST_SERVER_NAME; +use super::TEST_TOOL_NAME; +use super::TIMEOUT; +use super::luna_websocket; +use super::start_mcp_server; +use super::wait_for_luna_request; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn async_guardian_retains_evidence_after_compaction_and_discards_it_after_rollback() +-> Result<()> { + skip_if_no_network!(Ok(())); + + const RESTRICTION: &str = "Only publish to a private repository."; + const EVIDENCE: &str = "repository is private"; + const SUMMARY: &str = "Repository inspection was summarized."; + let expected_output = format!("\"echoed\":\"{EVIDENCE}\""); + let parent_requests = Arc::new(Mutex::new(Vec::::new())); + let compact_requests = Arc::new(Mutex::new(Vec::::new())); + let classifier = Arc::new(MockResponsesState { + luna_score: 1.0, + ..Default::default() + }); + let router = Router::new() + .route( + "/v1/responses", + get(luna_websocket).post({ + let parent_requests = Arc::clone(&parent_requests); + move |Json(request): Json| { + let parent_requests = Arc::clone(&parent_requests); + async move { + let events = if request["client_metadata"]["x-openai-subagent"] + == "guardian" + { + vec![ + responses::ev_assistant_message("review", r#"{"outcome":"allow"}"#), + responses::ev_completed("review"), + ] + } else { + let mut requests = parent_requests.lock().unwrap(); + let step = requests.len(); + requests.push(request); + if step.is_multiple_of(/*rhs*/ 2) { + let message = if step == 0 { + EVIDENCE + } else { + "current inspection" + }; + vec![ + responses::ev_function_call_with_namespace( + &format!("inspect-{step}"), + &format!("mcp__{TEST_SERVER_NAME}"), + TEST_TOOL_NAME, + &json!({"message": message}).to_string(), + ), + responses::ev_completed(&format!("inspect-{step}")), + ] + } else { + vec![responses::ev_completed(&format!("done-{step}"))] + } + }; + ( + [(header::CONTENT_TYPE, "text/event-stream")], + responses::sse(events), + ) + } + } + }), + ) + .route( + "/v1/responses/compact", + post({ + let compact_requests = Arc::clone(&compact_requests); + move |Json(request): Json| { + let compact_requests = Arc::clone(&compact_requests); + async move { + compact_requests.lock().unwrap().push(request); + Json(json!({"output": [ + {"type": "message", "role": "assistant", "content": [ + {"type": "output_text", "text": SUMMARY} + ]}, + {"type": "compaction", "encrypted_content": "encrypted summary"} + ]})) + } + } + }), + ) + .with_state(Arc::clone(&classifier)); + let listener = TcpListener::bind("127.0.0.1:0").await?; + let responses_url = format!("http://{}", listener.local_addr()?); + let responses_server = tokio::spawn(async move { + axum::serve(listener, router).await.unwrap(); + }); + let (mcp_url, mcp_server) = start_mcp_server(/*sensitive_action*/ None).await?; + let codex_home = TempDir::new()?; + MockResponsesConfig::new(&responses_url) + .with_provider_name("OpenAI") + .with_provider_config("requires_openai_auth = true\nsupports_websockets = false") + .with_root_config("approvals_reviewer = \"auto_review\"\nmodel_auto_compact_token_limit = 1000000") + .enable_feature(Feature::GuardianApproval) + .disable_feature(Feature::EnableRequestCompression) + .disable_feature(Feature::RemoteCompactionV2) + .disable_feature(Feature::TokenBudget) + .with_extra_config(&format!( + "[mcp_servers.{TEST_SERVER_NAME}]\nurl = \"{mcp_url}/mcp\"\ndefault_tools_approval_mode = \"prompt\"\n\n[features.guardianv2]\nenabled = true\n\n[features.guardianv2.review_scope]\ncomputer_use_only = false" + )) + .write(codex_home.path())?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("access-chatgpt").plan_type("pro"), + AuthCredentialsStoreMode::File, + )?; + let mut app_server = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .with_env_overrides(&[("OPENAI_API_KEY", None)]) + .build_initialized_with_timeout(TIMEOUT) + .await?; + let thread_id = app_server + .start_thread(ThreadStartParams { + approval_policy: Some(AskForApproval::OnRequest), + approvals_reviewer: Some(ApprovalsReviewer::AutoReview), + history_mode: Some(ThreadHistoryMode::Legacy), + ..Default::default() + }) + .await? + .thread + .id; + + for (index, prompt) in [ + RESTRICTION, + "Recheck the repository.", + "Inspect a different repository.", + ] + .into_iter() + .enumerate() + { + if index == 1 { + let id = app_server + .send_thread_compact_start_request(ThreadCompactStartParams { + thread_id: thread_id.clone(), + }) + .await?; + let _: ThreadCompactStartResponse = + timeout(TIMEOUT, app_server.read_response(id)).await??; + let completed: TurnCompletedNotification = + timeout(TIMEOUT, app_server.read_notification("turn/completed")).await??; + assert_eq!(completed.turn.status, TurnStatus::Completed); + assert_eq!(compact_requests.lock().unwrap().len(), 1); + } else if index == 2 { + let id = app_server + .send_thread_rollback_request(ThreadRollbackParams { + thread_id: thread_id.clone(), + num_turns: 2, + }) + .await?; + let _: ThreadRollbackResponse = + timeout(TIMEOUT, app_server.read_response(id)).await??; + } + + let completed = timeout( + TIMEOUT, + app_server.start_turn_and_wait_for_completion(TurnStartParams { + thread_id: thread_id.clone(), + input: vec![UserInput::Text { + text: prompt.to_owned(), + text_elements: Vec::new(), + }], + ..Default::default() + }), + ) + .await??; + assert_eq!(completed.turn.status, TurnStatus::Completed); + let request = wait_for_luna_request(&classifier, index).await?; + let content = request["input"] + .as_array() + .unwrap() + .iter() + .filter(|item| item["role"] == "user") + .filter_map(|item| item["content"].as_array()) + .flatten() + .filter_map(|part| part["text"].as_str()) + .collect::(); + let transcript = content + .split_once(">>> TRANSCRIPT START\n") + .unwrap() + .1 + .split_once(">>> TRANSCRIPT END") + .unwrap() + .0; + assert!( + transcript.contains(prompt), + "current user input missing: {transcript}" + ); + + let parent = parent_requests.lock().unwrap(); + assert_eq!(parent.len(), (index + 1) * 2); + if index == 0 { + let input = parent[1]["input"].as_array().unwrap(); + let output = input + .iter() + .find(|item| item["type"] == "function_call_output") + .expect("the MCP tool must actually execute before compaction"); + let output = output["output"].as_str().unwrap(); + assert!(output.contains(&expected_output), "{output}"); + } else { + let parent_input = serde_json::to_string(&parent[index * 2]["input"])?; + assert!(!parent_input.contains(RESTRICTION)); + assert!(!parent_input.contains(EVIDENCE)); + if index == 1 { + assert!(parent_input.contains(SUMMARY)); + assert!( + transcript.contains(RESTRICTION), + "retained user restriction missing: {transcript}" + ); + assert!(transcript.contains(&format!("tool {TEST_TOOL_NAME} call:"))); + assert!(transcript.contains(&format!("tool {TEST_TOOL_NAME} result:"))); + assert!( + transcript.contains(&expected_output), + "retained MCP output missing: {transcript}" + ); + } else { + assert!(!transcript.contains(RESTRICTION)); + assert!(!transcript.contains(EVIDENCE)); + assert!(!transcript.contains("Recheck the repository.")); + assert!(!transcript.contains("\"echoed\":\"current inspection\"")); + } + } + classifier.allow_luna.notify_one(); + } + + app_server.shutdown_gracefully().await?; + mcp_server.abort(); + responses_server.abort(); + Ok(()) +} diff --git a/codex-rs/ext/guardian-v2/Cargo.toml b/codex-rs/ext/guardian-v2/Cargo.toml index 3e28970a36..242102a99c 100644 --- a/codex-rs/ext/guardian-v2/Cargo.toml +++ b/codex-rs/ext/guardian-v2/Cargo.toml @@ -17,6 +17,7 @@ codex-api = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-features = { workspace = true } +codex-guardian-context = { workspace = true } codex-history = { workspace = true } codex-http-client = { workspace = true } codex-login = { workspace = true } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs index 10d2d8c678..6eb7c7f6dd 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs @@ -557,7 +557,7 @@ impl GuardianV2Extension { guardian_evidence.user_input_fragments(input.conversation_history.as_ref()); let transcript = guardian_config .transcript - .build(input.conversation_history.review_items()); + .build_snapshot(input.conversation_history.as_ref()); let local_trusted_skill_paths = guardian_evidence.trusted_skill_paths(input.turn_id); let node_repl_images = if guardian_config.transcript.include_images { input diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs index 719aff31c1..9df19ecda7 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs @@ -39,6 +39,9 @@ use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ImageDetail; use codex_protocol::models::InternalChatMessageMetadataPassthrough; +use codex_protocol::models::LocalShellAction; +use codex_protocol::models::LocalShellExecAction; +use codex_protocol::models::LocalShellStatus; use codex_protocol::models::MessagePhase; use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::openai_models::GuardianV2ModelConfig; @@ -81,6 +84,8 @@ use crate::async_scorer::sampler::CLASSIFICATION_TOKEN_USAGE_METRIC; use crate::async_scorer::sampler::INITIAL_WEBSOCKET_CONNECTIONS; use crate::async_scorer::sampler::LunaSampler; use crate::async_scorer::sampler::MODEL; +use crate::async_scorer::transcript::MAX_MESSAGE_ENTRY_TOKENS; +use crate::async_scorer::transcript::MAX_TOOL_ENTRY_TOKENS; use crate::async_scorer::transcript::truncate_entry; use crate::async_scorer::truncation::CLASSIFICATION_TRUNCATION_BYTES_METRIC; use crate::async_scorer::truncation::CLASSIFICATION_TRUNCATION_METRIC; @@ -2666,6 +2671,56 @@ async fn contributor_sends_compacted_conversation_history_to_luna() -> Result<() }, ] })); + history.extend( + [ + (None, None, "unattributed anonymous output"), + (Some("missing-call"), None, "unattributed orphaned output"), + ( + Some("missing-named-call"), + Some("notifications"), + "named orphaned output", + ), + (None, Some("notifications"), "attributed notification"), + ] + .map(|(call_id, name, text)| ResponseItem::FunctionCallOutput { + id: None, + call_id: call_id.map(str::to_string), + name: name.map(str::to_string), + namespace: Some("slack".to_string()), + output: FunctionCallOutputPayload::from_text(text.to_string()), + internal_chat_message_metadata_passthrough: None, + }), + ); + history.push(ResponseItem::CustomToolCallOutput { + id: None, + call_id: "missing-custom-call".to_string(), + name: None, + output: FunctionCallOutputPayload::from_text("unattributed custom output".to_string()), + internal_chat_message_metadata_passthrough: None, + }); + history.extend([ + ResponseItem::LocalShellCall { + id: None, + call_id: Some("shell-1".to_string()), + status: LocalShellStatus::Completed, + action: LocalShellAction::Exec(LocalShellExecAction { + command: vec!["echo".to_string(), "hello".to_string()], + timeout_ms: None, + working_directory: None, + env: None, + user: None, + }), + internal_chat_message_metadata_passthrough: None, + }, + ResponseItem::FunctionCallOutput { + id: None, + call_id: Some("shell-1".to_string()), + name: None, + namespace: None, + output: FunctionCallOutputPayload::from_text("local shell evidence".to_string()), + internal_chat_message_metadata_passthrough: None, + }, + ]); let (request, _test, _registry) = sample_conversation_history( history, @@ -2705,6 +2760,42 @@ async fn contributor_sends_compacted_conversation_history_to_luna() -> Result<() .any(|entry| entry.contains("result evidence 0:")) ); assert!(entries.iter().any(|entry| entry.contains("( + pub(crate) fn build_snapshot( &self, - items: impl IntoIterator, + history: &dyn ConversationHistorySnapshot, + ) -> RenderedTranscript { + let history = SnapshotHistory(history); + let retention = TranscriptRetentionConfig { + max_message_transcript_tokens: self.max_message_transcript_tokens, + max_tool_transcript_tokens: self.max_tool_transcript_tokens, + max_recent_non_user_entries: self.max_recent_non_user_entries, + }; + let transcript = ConversationTranscriptConfig { + options: ConversationTranscriptOptions { + include_tool_calls: self.sources.contains(&TranscriptSource::ToolCalls), + include_tool_outputs: self.sources.contains(&TranscriptSource::ToolOutputs), + include_reasoning: self.sources.contains(&TranscriptSource::Reasoning), + }, + entry_limits: TranscriptEntryLimits { + message_tokens: self.max_message_entry_tokens, + tool_tokens: self.max_tool_entry_tokens, + node_repl_output_tokens: self.max_tool_entry_tokens, + }, + }; + let entries = collect_transcript(&history, &transcript); + Self::render(entries, &retention) + } + + fn render( + transcript_entries: impl IntoIterator, + retention: &TranscriptRetentionConfig, ) -> RenderedTranscript { let mut entries = Vec::new(); - let mut tool_names_by_call_id = HashMap::new(); - for item in items { - let (role, text) = match item { - ResponseItem::Message { role, content, .. } => { - let text = content - .iter() - .filter_map(|item| match item { - ContentItem::InputText { text } | ContentItem::OutputText { text } - if !text.is_empty() => - { - Some(text.as_str()) - } - ContentItem::InputText { .. } - | ContentItem::OutputText { .. } - | ContentItem::InputImage { .. } - | ContentItem::InputAudio { .. } => None, - }) - .collect::>() - .join("\n"); - if text.trim().is_empty() { - continue; - } - let include_message = match role.as_str() { - "user" | "assistant" => true, - "developer" => text.starts_with(MANUAL_APPROVAL_DEVELOPER_PREFIX), - _ => false, - }; - if !include_message { - continue; - } - (role.clone(), text) - } - ResponseItem::AgentMessage { - author, content, .. - } => { - let Some(text) = plaintext_agent_message_content(content) else { - continue; - }; - ( - "assistant".to_owned(), - format!("Agent message from {author}:\n{text}"), - ) - } - ResponseItem::FunctionCall { - name, - arguments, - call_id, - .. - } - | ResponseItem::CustomToolCall { - name, - input: arguments, - call_id, - .. - } => { - tool_names_by_call_id.insert(call_id.as_str(), name.as_str()); - if !self.sources.contains(&TranscriptSource::ToolCalls) { - continue; - } - (format!("tool {name} call"), arguments.clone()) - } - ResponseItem::FunctionCallOutput { - call_id: Some(call_id), - output, - .. - } - | ResponseItem::CustomToolCallOutput { - call_id, output, .. - } => { - if !self.sources.contains(&TranscriptSource::ToolOutputs) { - continue; - } - let Some(output) = output.body.to_text() else { - continue; - }; - let role = tool_names_by_call_id.get(call_id.as_str()).map_or_else( - || "tool result".to_owned(), - |name| format!("tool {name} result"), - ); - (role, output) - } - ResponseItem::FunctionCallOutput { - call_id: None, - name, - namespace, - output, - .. - } => { - if !self.sources.contains(&TranscriptSource::ToolOutputs) { - continue; - } - let output = output - .body - .to_text() - .unwrap_or_else(|| "[non-text output]".into()); - let role = match (name.as_deref(), namespace.as_deref()) { - (Some(name), Some(namespace)) => { - format!("tool {namespace}.{name} result") - } - (Some(name), None) => format!("tool {name} result"), - (None, _) => "tool result".to_owned(), - }; - (role, output) - } - ResponseItem::Reasoning { - summary, content, .. - } => { - if !self.sources.contains(&TranscriptSource::Reasoning) { - continue; - } - let text = summary - .iter() - .map(|item| match item { - ReasoningItemReasoningSummary::SummaryText { text } => text.as_str(), - }) - .chain(content.iter().flatten().map(|item| match item { - ReasoningItemContent::ReasoningText { text } - | ReasoningItemContent::Text { text } => text.as_str(), - })) - .filter(|text| !text.trim().is_empty()) - .collect::>() - .join("\n"); - ("reasoning".to_owned(), text) - } - ResponseItem::LocalShellCall { action, .. } => { - if !self.sources.contains(&TranscriptSource::ToolCalls) { - continue; - } - let Ok(text) = serde_json::to_string(action) else { - continue; - }; - ("tool shell call".to_owned(), text) - } - ResponseItem::WebSearchCall { action, .. } => { - if !self.sources.contains(&TranscriptSource::ToolCalls) { - continue; - } - let Some(action) = action else { - continue; - }; - let Ok(text) = serde_json::to_string(action) else { - continue; - }; - ("tool web_search call".to_owned(), text) - } - ResponseItem::AdditionalTools { .. } - | ResponseItem::ImageGenerationCall { .. } - | ResponseItem::ToolSearchCall { .. } - | ResponseItem::ToolSearchOutput { .. } - | ResponseItem::Compaction { .. } - | ResponseItem::CompactionTrigger { .. } - | ResponseItem::ContextCompaction { .. } - | ResponseItem::Other => continue, - }; - if text.trim().is_empty() { - continue; - } - - let kind = match (role.as_str(), item) { - ("user", _) => TranscriptEntryKind::User, - ("developer", ResponseItem::Message { .. }) => { + for entry in transcript_entries { + let role = entry.kind.role(); + let kind = match &entry.kind { + ConversationTranscriptEntryKind::User => TranscriptEntryKind::User, + ConversationTranscriptEntryKind::Developer + | ConversationTranscriptEntryKind::ProtectedAssistant => { TranscriptEntryKind::ProtectedMessage } - ( - "assistant", - ResponseItem::Message { - phase: None | Some(MessagePhase::FinalAnswer), - content, - .. - }, - ) if !InterAgentCommunication::is_message_content(content) => { - TranscriptEntryKind::ProtectedMessage + ConversationTranscriptEntryKind::Assistant + | ConversationTranscriptEntryKind::Reasoning => TranscriptEntryKind::Message, + ConversationTranscriptEntryKind::ToolCall(_) + | ConversationTranscriptEntryKind::ToolOutput(_) + | ConversationTranscriptEntryKind::NodeReplToolOutput(_) => { + TranscriptEntryKind::Tool } - (role, _) if role.starts_with("tool ") => TranscriptEntryKind::Tool, - _ => TranscriptEntryKind::Message, }; - let token_cap = match kind { - TranscriptEntryKind::Tool => self.max_tool_entry_tokens, - TranscriptEntryKind::User - | TranscriptEntryKind::ProtectedMessage - | TranscriptEntryKind::Message => self.max_message_entry_tokens, - }; - let original_bytes = text.len(); - let text = truncate_entry(&text, token_cap); + let original_bytes = entry.original_bytes; + let text = entry.text; let retained_bytes = text.len(); let entry_number = entries.len() + 1; let text = format!("[{entry_number}] {role}: {text}\n"); @@ -377,7 +240,7 @@ impl TranscriptConfig { if let Some(&latest_user_index) = user_indices.last() && !included[latest_user_index] && message_tokens + entries[latest_user_index].tokens - <= self.max_message_transcript_tokens + <= retention.max_message_transcript_tokens { included[latest_user_index] = true; message_tokens += entries[latest_user_index].tokens; @@ -385,7 +248,7 @@ impl TranscriptConfig { for &index in user_indices.iter().rev() { if included[index] - || message_tokens + entries[index].tokens > self.max_message_transcript_tokens + || message_tokens + entries[index].tokens > retention.max_message_transcript_tokens { continue; } @@ -394,10 +257,10 @@ impl TranscriptConfig { message_tokens += entries[index].tokens; } - let available_message_tokens = self + let available_message_tokens = retention .max_message_transcript_tokens .saturating_sub(message_tokens); - let mut window = TranscriptWindow::new(&entries, self, available_message_tokens); + let mut window = TranscriptWindow::new(&entries, retention, available_message_tokens); for index in 0..entries.len() { window.insert(index); } @@ -441,34 +304,12 @@ impl TranscriptConfig { } } -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(); +struct SnapshotHistory<'a>(&'a dyn ConversationHistorySnapshot); + +impl SectionHistory for SnapshotHistory<'_> { + fn items(&self) -> Box + Send + '_> { + self.0.review_items() } - - 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)] 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 index fbb6eaa31b..ccc72649ac 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/transcript/window.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/transcript/window.rs @@ -1,6 +1,7 @@ use std::collections::VecDeque; -use super::TranscriptConfig; +use codex_guardian_context::TranscriptRetentionConfig; + use super::TranscriptEntry; use super::TranscriptEntryKind; @@ -44,7 +45,7 @@ pub(super) struct TranscriptWindow<'a> { impl<'a> TranscriptWindow<'a> { pub(super) fn new( entries: &'a [TranscriptEntry], - config: &TranscriptConfig, + config: &TranscriptRetentionConfig, message_token_budget: usize, ) -> Self { Self { 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 2901da9fe1..4cb2d5dfb2 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 @@ -1,3 +1,4 @@ +use codex_extension_api::ConversationHistorySnapshot; use codex_extension_api::ResponseItem; use codex_protocol::AgentPath; use codex_protocol::models::AgentMessageInputContent; @@ -21,6 +22,22 @@ use super::TranscriptConfig; use super::TranscriptSource; use super::truncate_entry; +struct TestConversationHistory<'a>(&'a [ResponseItem]); + +impl ConversationHistorySnapshot for TestConversationHistory<'_> { + fn history_version(&self) -> u64 { + 0 + } + + fn user_message_revision(&self) -> u64 { + self.0.iter().filter(|item| item.is_user_message()).count() as u64 + } + + fn items(&self) -> Box + Send + '_> { + Box::new(self.0.iter()) + } +} + fn assistant_message(text: impl Into, phase: MessagePhase) -> ResponseItem { ResponseItem::Message { id: None, @@ -80,7 +97,9 @@ fn transcript_keeps_conversation_and_configured_sources() { }, ]; - let transcript = TranscriptConfig::default().build(&items).entries; + let transcript = TranscriptConfig::default() + .build_snapshot(&TestConversationHistory(&items)) + .entries; assert_eq!( transcript, vec![ @@ -95,7 +114,9 @@ fn transcript_keeps_conversation_and_configured_sources() { ..TranscriptConfig::default() }; - let transcript = output_and_reasoning.build(&items).entries; + let transcript = output_and_reasoning + .build_snapshot(&TestConversationHistory(&items)) + .entries; assert_eq!( transcript, vec![ @@ -110,7 +131,9 @@ fn transcript_keeps_conversation_and_configured_sources() { ..TranscriptConfig::default() }; - let transcript = calls_only.build(&items).entries; + let transcript = calls_only + .build_snapshot(&TestConversationHistory(&items)) + .entries; assert_eq!( transcript, vec![ @@ -128,6 +151,7 @@ fn transcript_truncates_oversized_entries_without_splitting_characters() { "{prefix}{}{suffix}", "é".repeat(TruncationPolicy::Tokens(MAX_MESSAGE_ENTRY_TOKENS).byte_budget()) ); + let original_bytes = oversized_message.len(); let items = vec![ ResponseItem::Message { id: None, @@ -149,7 +173,8 @@ fn transcript_truncates_oversized_entries_without_splitting_characters() { }, ]; - let transcript = TranscriptConfig::default().build(&items).entries; + let rendered = TranscriptConfig::default().build_snapshot(&TestConversationHistory(&items)); + let transcript = rendered.entries; assert_eq!(transcript.len(), 2); let user_entry = &transcript[0]; @@ -162,6 +187,22 @@ fn transcript_truncates_oversized_entries_without_splitting_characters() { + "[1] user: \n".len() ); assert_eq!(transcript[1], "[2] assistant: latest response\n"); + assert_eq!( + rendered + .truncations + .iter() + .map(|observation| ( + observation.component, + observation.original_bytes, + observation.retained_bytes, + )) + .collect::>(), + vec![( + "transcript_user", + original_bytes, + user_entry.len() - "[1] user: \n".len() + )] + ); } #[test] @@ -188,7 +229,9 @@ fn transcript_preserves_first_and_latest_user_messages_and_recent_history() { internal_chat_message_metadata_passthrough: None, }); - let transcript = TranscriptConfig::default().build(&items).entries; + let transcript = TranscriptConfig::default() + .build_snapshot(&TestConversationHistory(&items)) + .entries; assert!(transcript[0].starts_with("[1] user: user turn 0:")); assert!( @@ -242,7 +285,7 @@ fn transcript_preserves_user_restrictions_before_final_assistant_messages() { max_message_transcript_tokens: message_budget, ..TranscriptConfig::default() } - .build(&items) + .build_snapshot(&TestConversationHistory(&items)) .entries; assert_eq!( @@ -271,7 +314,7 @@ fn transcript_preserves_recent_tool_evidence_when_protected_messages_fill_entry_ max_recent_non_user_entries: 4, ..TranscriptConfig::default() } - .build(&items); + .build_snapshot(&TestConversationHistory(&items)); assert_eq!( transcript.entries, @@ -325,7 +368,9 @@ fn transcript_reserves_five_recent_tool_entries_from_protected_messages() { ) })); - let transcript = TranscriptConfig::default().build(&items).entries; + let transcript = TranscriptConfig::default() + .build_snapshot(&TestConversationHistory(&items)) + .entries; let tool_entries = transcript .iter() .filter(|entry| entry.contains("tool exec_command ")) @@ -369,7 +414,7 @@ fn rejected_commentary_does_not_evict_retained_message_evidence() { max_message_transcript_tokens: message_budget, ..TranscriptConfig::default() } - .build(&items) + .build_snapshot(&TestConversationHistory(&items)) .entries; assert_eq!(transcript, vec![protected_entry, retained_entry]); @@ -398,7 +443,9 @@ fn transcript_evicts_protected_messages_in_cacheable_chunks() { items.extend((0..message_count).map(|index| { assistant_message(format!("final answer {index}"), MessagePhase::FinalAnswer) })); - config.build(&items).entries + config + .build_snapshot(&TestConversationHistory(&items)) + .entries }; assert_eq!( @@ -457,7 +504,7 @@ fn transcript_preserves_latest_final_when_reserved_tools_fill_entry_window() { max_recent_non_user_entries: 4, ..TranscriptConfig::default() } - .build(&items) + .build_snapshot(&TestConversationHistory(&items)) .entries; assert_eq!( @@ -519,7 +566,7 @@ fn transcript_does_not_protect_legacy_inter_agent_instructions() { max_recent_non_user_entries: 2, ..TranscriptConfig::default() } - .build(&items) + .build_snapshot(&TestConversationHistory(&items)) .entries; assert_eq!( @@ -554,7 +601,9 @@ fn transcript_reserves_separate_budget_for_recent_tool_evidence() { internal_chat_message_metadata_passthrough: None, })); - let transcript = TranscriptConfig::default().build(&items).entries; + let transcript = TranscriptConfig::default() + .build_snapshot(&TestConversationHistory(&items)) + .entries; assert!(transcript[0].contains("user turn 0:")); assert!( @@ -603,7 +652,9 @@ fn transcript_reserves_separate_budget_for_recent_tool_evidence() { call_id: "call-12".to_string(), internal_chat_message_metadata_passthrough: None, }); - let next_transcript = TranscriptConfig::default().build(&items).entries; + let next_transcript = TranscriptConfig::default() + .build_snapshot(&TestConversationHistory(&items)) + .entries; let next_first_retained_tool = next_transcript .iter() .find(|entry| entry.contains("tool exec_command call:")) @@ -655,7 +706,7 @@ fn transcript_preserves_newest_manual_approval_when_message_budget_overflows() { max_message_transcript_tokens: message_budget, ..TranscriptConfig::default() } - .build(&items) + .build_snapshot(&TestConversationHistory(&items)) .entries; assert_eq!(transcript, vec![approval_entry]); @@ -704,7 +755,7 @@ fn rejected_message_does_not_evict_retained_tool_entries() { max_message_transcript_tokens: TruncationPolicy::Bytes(user_entry.len()).token_budget(), ..TranscriptConfig::default() } - .build(&items) + .build_snapshot(&TestConversationHistory(&items)) .entries; assert_eq!(transcript, expected); @@ -737,7 +788,9 @@ fn transcript_evicts_non_user_entries_in_cacheable_chunks() { internal_chat_message_metadata_passthrough: None, }), ); - config.build(&items).entries + config + .build_snapshot(&TestConversationHistory(&items)) + .entries }; let first_overflow = build_transcript(5); @@ -797,7 +850,9 @@ fn transcript_truncates_tool_results_using_standard_budget() { }, ]; - let transcript = TranscriptConfig::default().build(&items).entries; + let transcript = TranscriptConfig::default() + .build_snapshot(&TestConversationHistory(&items)) + .entries; let result = transcript .iter() .find(|entry| entry.contains(" result: ")) @@ -814,7 +869,7 @@ fn transcript_truncates_tool_results_using_standard_budget() { } #[test] -fn transcript_preserves_named_unpaired_tool_source() { +fn transcript_preserves_outputs_with_call_ids_or_explicit_names() { let mut items = vec![ResponseItem::FunctionCallOutput { id: None, call_id: None, @@ -823,10 +878,29 @@ fn transcript_preserves_named_unpaired_tool_source() { output: FunctionCallOutputPayload::from_text("new message".to_owned()), internal_chat_message_metadata_passthrough: None, }]; + items.extend( + [ + (None, "anonymous output"), + (Some("missing-call"), "orphaned function output"), + ] + .map(|(call_id, text)| ResponseItem::FunctionCallOutput { + id: None, + call_id: call_id.map(str::to_string), + name: None, + namespace: None, + output: FunctionCallOutputPayload::from_text(text.to_string()), + internal_chat_message_metadata_passthrough: None, + }), + ); assert_eq!( - TranscriptConfig::default().build(&items).entries, - vec!["[1] tool slack.notifications result: new message\n"] + TranscriptConfig::default() + .build_snapshot(&TestConversationHistory(&items)) + .entries, + vec![ + "[1] tool slack.notifications result: new message\n", + "[2] tool result: orphaned function output\n", + ] ); if let ResponseItem::FunctionCallOutput { output, .. } = &mut items[0] { @@ -838,8 +912,13 @@ fn transcript_preserves_named_unpaired_tool_source() { ]); } assert_eq!( - TranscriptConfig::default().build(&items).entries, - vec!["[1] tool slack.notifications result: [non-text output]\n"] + TranscriptConfig::default() + .build_snapshot(&TestConversationHistory(&items)) + .entries, + vec![ + "[1] tool slack.notifications result: [non-text output]\n", + "[2] tool result: orphaned function output\n", + ] ); } @@ -870,7 +949,7 @@ fn configured_reasoning_counts_against_message_budget() { sources: vec![TranscriptSource::Reasoning], ..TranscriptConfig::default() } - .build(&items) + .build_snapshot(&TestConversationHistory(&items)) .entries; assert!(transcript[0].contains("user turn 0:")); @@ -927,7 +1006,9 @@ fn transcript_keeps_only_manual_approval_developer_messages() { }, ]; - let transcript = TranscriptConfig::default().build(&items).entries; + let transcript = TranscriptConfig::default() + .build_snapshot(&TestConversationHistory(&items)) + .entries; assert_eq!( transcript, vec![format!("[1] developer: {approval_text}\n")] @@ -1001,7 +1082,9 @@ fn transcript_omits_media_payloads_and_keeps_readable_content() { }, ]; - let transcript = TranscriptConfig::default().build(&items).entries; + let transcript = TranscriptConfig::default() + .build_snapshot(&TestConversationHistory(&items)) + .entries; assert_eq!( transcript, vec![ @@ -1057,7 +1140,9 @@ fn transcript_omits_encrypted_messages_arguments_and_tool_outputs() { }, ]; - let transcript = TranscriptConfig::default().build(&items).entries; + let transcript = TranscriptConfig::default() + .build_snapshot(&TestConversationHistory(&items)) + .entries; assert_eq!( transcript, vec![ diff --git a/codex-rs/ext/guardian-v2/src/sync_reviewer/prompt.rs b/codex-rs/ext/guardian-v2/src/sync_reviewer/prompt.rs index f50c863cac..dca170bbf7 100644 --- a/codex-rs/ext/guardian-v2/src/sync_reviewer/prompt.rs +++ b/codex-rs/ext/guardian-v2/src/sync_reviewer/prompt.rs @@ -94,7 +94,7 @@ pub(super) fn build( _ => None, }) .collect(); - let transcript = transcript_config.build(input.conversation_history.items()); + let transcript = transcript_config.build_snapshot(input.conversation_history.as_ref()); let images = render_images( input.conversation_history.as_ref(), transcript_config,