diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 6a58a9c373..474fa61918 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -4,6 +4,7 @@ use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::session::turn_context::TurnEnvironment; use crate::turn_diff_tracker::TurnDiffTracker; +use crate::unified_exec::output::CapturedStream; use crate::unified_exec::resolve_max_tokens; use codex_protocol::ResponseItemId; use codex_protocol::mcp::CallToolResult; @@ -16,16 +17,11 @@ use codex_protocol::models::function_call_output_content_items_to_text; use codex_tools::LoadableToolSpec; use codex_tools::ToolName; use codex_utils_audio::estimate_audio_token_count; -use codex_utils_head_tail_buffer::format_output_omission_marker; use codex_utils_output_truncation::TruncationPolicy; -use codex_utils_output_truncation::approx_token_count; -use codex_utils_output_truncation::formatted_truncate_text; use codex_utils_output_truncation::truncate_function_output_payload; -use codex_utils_output_truncation::truncate_text; use codex_utils_output_truncation::with_serialization_allowance; use serde::Serialize; use serde_json::Value as JsonValue; -use std::num::NonZeroUsize; use std::sync::Arc; use std::sync::OnceLock; use std::sync::atomic::AtomicBool; @@ -425,28 +421,22 @@ pub struct ExecCommandToolOutput { pub event_call_id: String, pub chunk_id: String, pub wall_time: Duration, - /// Raw bytes returned for this unified exec call before any truncation. - pub raw_output: Vec, + /// Output retained by collection, before model-facing truncation. + pub(crate) raw_output: CapturedStream, pub truncation_policy: TruncationPolicy, pub max_output_tokens: Option, pub process_id: Option, pub exit_code: Option, pub original_token_count: Option, - /// Bytes omitted by the output collection cap before model-facing truncation. - pub output_omitted_bytes: Option, pub hook_command: Option, } impl ToolOutput for ExecCommandToolOutput { fn log_output(&self) -> String { // The telemetry budget must not inherit the model's output-token limit. - let mut output = String::from_utf8_lossy(&self.raw_output).into_owned(); - if let Some(omitted_bytes) = self.output_omitted_bytes { - let marker = format_output_omission_marker(omitted_bytes.get()); - if !output.contains(&marker) { - output = format!("{marker}\n{output}"); - } - } + let output = self + .raw_output + .render(/*policy*/ None, self.original_token_count); format!("{}\n{output}", self.response_header()) } @@ -512,7 +502,7 @@ impl ToolOutput for ExecCommandToolOutput { original_token_count: self.original_token_count, output: match self.max_output_tokens { Some(max_tokens) => self.truncated_output(max_tokens), - None => String::from_utf8_lossy(&self.raw_output).to_string(), + None => self.raw_output.text.clone(), }, }; @@ -537,32 +527,8 @@ impl ExecCommandToolOutput { } fn truncated_output_with_policy(&self, policy: TruncationPolicy) -> String { - let text = String::from_utf8_lossy(&self.raw_output).to_string(); - let Some(omitted_bytes) = self.output_omitted_bytes else { - return formatted_truncate_text(&text, policy); - }; - - let marker = format_output_omission_marker(omitted_bytes.get()); - if text.len() <= policy.byte_budget() { - return if text.contains(&marker) { - text - } else { - format!("{marker}\n{text}") - }; - } - - let original_token_count = self - .original_token_count - .unwrap_or_else(|| approx_token_count(&text)); - let truncated = truncate_text(&text, policy); - let omission_notice = if truncated.contains(&marker) { - String::new() - } else { - format!("{marker}\n") - }; - format!( - "Warning: truncated output (original token count: {original_token_count})\n{omission_notice}\n{truncated}" - ) + self.raw_output + .render(Some(policy), self.original_token_count) } fn response_header(&self) -> String { diff --git a/codex-rs/core/src/tools/context_tests.rs b/codex-rs/core/src/tools/context_tests.rs index bd0cb0c5b5..7781ddccd9 100644 --- a/codex-rs/core/src/tools/context_tests.rs +++ b/codex-rs/core/src/tools/context_tests.rs @@ -1,7 +1,9 @@ use super::*; +use crate::unified_exec::output::CapturedStream; use codex_protocol::models::DEFAULT_IMAGE_DETAIL; use codex_protocol::models::ImageReference; use codex_protocol::models::SearchToolCallParams; +use codex_utils_head_tail_buffer::format_output_omission_marker; use core_test_support::assert_regex_match; use pretty_assertions::assert_eq; use serde_json::json; @@ -416,13 +418,12 @@ fn exec_command_tool_output_formats_truncated_response() { event_call_id: "call-42".to_string(), chunk_id: "abc123".to_string(), wall_time: std::time::Duration::from_millis(1250), - raw_output: b"token one token two token three token four token five".to_vec(), + raw_output: "token one token two token three token four token five".into(), truncation_policy: TruncationPolicy::Tokens(10_000), max_output_tokens: Some(4), process_id: None, exit_code: Some(0), original_token_count: Some(10), - output_omitted_bytes: None, hook_command: None, }; assert_eq!( @@ -462,8 +463,7 @@ fn exec_command_tool_output_reserves_metadata_budget_and_preserves_policy_units( }; let raw_output = (1..=150) .map(|line| format!("{line}\n")) - .collect::() - .into_bytes(); + .collect::(); for (policy, marker) in [ (TruncationPolicy::Bytes(200), "chars truncated"), @@ -473,13 +473,12 @@ fn exec_command_tool_output_reserves_metadata_budget_and_preserves_policy_units( event_call_id: "call-42".to_string(), chunk_id: "abc123".to_string(), wall_time: std::time::Duration::from_millis(/*millis*/ 1250), - raw_output: raw_output.clone(), + raw_output: raw_output.as_str().into(), truncation_policy: policy, max_output_tokens: None, process_id: None, exit_code: Some(0), original_token_count: Some(123), - output_omitted_bytes: None, hook_command: None, } .to_response_item("call-42", &payload); @@ -511,33 +510,31 @@ fn exec_command_tool_output_preserves_omission_metadata_when_truncated() { "HEAD-{}\n{marker}\nTAIL-{}", "a".repeat(/*n*/ 100), "z".repeat(/*n*/ 100) - ) - .into_bytes(); + ); let mut output = ExecCommandToolOutput { event_call_id: "call-omitted".to_string(), chunk_id: "abc123".to_string(), wall_time: std::time::Duration::from_millis(/*millis*/ 1250), - raw_output, + raw_output: CapturedStream { + text: raw_output.clone(), + omitted_bytes: 123_456, + }, truncation_policy: TruncationPolicy::Tokens(10_000), max_output_tokens: Some(4), process_id: None, exit_code: Some(0), original_token_count: Some(42_000), - output_omitted_bytes: NonZeroUsize::new(/*n*/ 123_456), hook_command: None, }; let expected_header = "Chunk ID: abc123\nWall time: 1.2500 seconds\nProcess exited with code 0\nOriginal token count: 42000\nOutput:\n"; assert_eq!( output.log_output(), - format!( - "{expected_header}{}", - String::from_utf8_lossy(&output.raw_output) - ) + format!("{expected_header}{raw_output}") ); let response = output.to_response_item("call-omitted", &payload); // Collection may report omitted bytes without including the marker in its text. - output.raw_output = b"remaining output".to_vec(); + output.raw_output.text = "remaining output".to_string(); assert_eq!( output.log_output(), format!("{expected_header}{marker}\nremaining output") diff --git a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs index ca15be3f73..1a188c90a1 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs @@ -34,6 +34,7 @@ use crate::unified_exec::UnifiedExecContext; use crate::unified_exec::UnifiedExecError; use crate::unified_exec::UnifiedExecProcessManager; use crate::unified_exec::generate_chunk_id; +use crate::unified_exec::output::CapturedStream; use codex_features::Feature; use codex_otel::SessionTelemetry; use codex_otel::TOOL_CALL_UNIFIED_EXEC_METRIC; @@ -374,13 +375,15 @@ impl ExecCommandHandler { event_call_id: String::new(), chunk_id: String::new(), wall_time: std::time::Duration::ZERO, - raw_output: output.into_text().into_bytes(), + raw_output: CapturedStream { + text: output.into_text(), + omitted_bytes: 0, + }, truncation_policy: step_context.settings.model_info.truncation_policy.into(), max_output_tokens, process_id: None, exit_code: None, original_token_count: None, - output_omitted_bytes: None, hook_command: None, })); } @@ -438,7 +441,10 @@ impl ExecCommandHandler { event_call_id: context.call_id.clone(), chunk_id: generate_chunk_id(), wall_time: output.duration, - raw_output: output_text.into_bytes(), + raw_output: CapturedStream { + text: output_text, + omitted_bytes: output_omitted_bytes.map_or(0, std::num::NonZeroUsize::get), + }, truncation_policy: step_context.settings.model_info.truncation_policy.into(), max_output_tokens, // Sandbox denial is terminal, so there is no live @@ -446,7 +452,6 @@ impl ExecCommandHandler { process_id: None, exit_code: Some(output.exit_code), original_token_count: Some(original_token_count), - output_omitted_bytes, hook_command: Some(hook_command), })) } diff --git a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs index 0881f7eaff..ee8b481bba 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs @@ -524,13 +524,12 @@ async fn exec_command_post_tool_use_payload_uses_output_for_noninteractive_one_s event_call_id: "call-43".to_string(), chunk_id: "chunk-1".to_string(), wall_time: std::time::Duration::from_millis(498), - raw_output: b"three".to_vec(), + raw_output: "three".into(), truncation_policy: TEST_TRUNCATION_POLICY, max_output_tokens: None, process_id: None, exit_code: Some(0), original_token_count: None, - output_omitted_bytes: None, hook_command: Some("echo three".to_string()), }; let invocation = invocation_for_payload("exec_command", "call-43", payload).await; @@ -555,13 +554,12 @@ async fn exec_command_post_tool_use_payload_uses_output_for_interactive_completi event_call_id: "call-44".to_string(), chunk_id: "chunk-1".to_string(), wall_time: std::time::Duration::from_millis(498), - raw_output: b"three".to_vec(), + raw_output: "three".into(), truncation_policy: TEST_TRUNCATION_POLICY, max_output_tokens: None, process_id: None, exit_code: Some(0), original_token_count: None, - output_omitted_bytes: None, hook_command: Some("echo three".to_string()), }; let invocation = invocation_for_payload("exec_command", "call-44", payload).await; @@ -587,13 +585,12 @@ async fn exec_command_post_tool_use_payload_skips_running_sessions() { event_call_id: "event-45".to_string(), chunk_id: "chunk-1".to_string(), wall_time: std::time::Duration::from_millis(498), - raw_output: b"three".to_vec(), + raw_output: "three".into(), truncation_policy: TEST_TRUNCATION_POLICY, max_output_tokens: None, process_id: Some(45), exit_code: None, original_token_count: None, - output_omitted_bytes: None, hook_command: Some("echo three".to_string()), }; let invocation = invocation_for_payload("exec_command", "call-45", payload).await; @@ -614,13 +611,12 @@ async fn write_stdin_post_tool_use_payload_uses_original_exec_call_id_and_comman event_call_id: "exec-call-45".to_string(), chunk_id: "chunk-2".to_string(), wall_time: std::time::Duration::from_millis(498), - raw_output: b"finished\n".to_vec(), + raw_output: "finished\n".into(), truncation_policy: TEST_TRUNCATION_POLICY, max_output_tokens: None, process_id: None, exit_code: Some(0), original_token_count: None, - output_omitted_bytes: None, hook_command: Some("sleep 1; echo finished".to_string()), }; let invocation = invocation_for_payload("write_stdin", "write-stdin-call", payload).await; @@ -646,26 +642,24 @@ async fn write_stdin_post_tool_use_payload_keeps_parallel_session_metadata_separ event_call_id: "exec-call-a".to_string(), chunk_id: "chunk-a".to_string(), wall_time: std::time::Duration::from_millis(498), - raw_output: b"alpha\n".to_vec(), + raw_output: "alpha\n".into(), truncation_policy: TEST_TRUNCATION_POLICY, max_output_tokens: None, process_id: None, exit_code: Some(0), original_token_count: None, - output_omitted_bytes: None, hook_command: Some("sleep 2; echo alpha".to_string()), }; let output_b = ExecCommandToolOutput { event_call_id: "exec-call-b".to_string(), chunk_id: "chunk-b".to_string(), wall_time: std::time::Duration::from_millis(498), - raw_output: b"beta\n".to_vec(), + raw_output: "beta\n".into(), truncation_policy: TEST_TRUNCATION_POLICY, max_output_tokens: None, process_id: None, exit_code: Some(0), original_token_count: None, - output_omitted_bytes: None, hook_command: Some("sleep 1; echo beta".to_string()), }; let invocation_b = invocation_for_payload("write_stdin", "write-call-b", payload.clone()).await; diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index a7a8899c5e..26a0cc5258 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -49,6 +49,7 @@ use codex_core_plugins::PluginMetricsSidecar; mod async_watcher; mod errors; mod oneshot; +pub(crate) mod output; mod process; mod process_manager; pub(crate) use process_manager::exec_env_policy_from_shell_policy; diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index 5fb49dbd25..6be0a6ec6a 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -9,6 +9,7 @@ use crate::session::tests::make_session_and_context; use crate::session::turn_context::TurnContext; use crate::tools::context::ExecCommandToolOutput; use crate::unified_exec::WriteStdinRequest; +use crate::unified_exec::output::CapturedStream; use codex_exec_server::ExecProcess; use codex_exec_server::ExecProcessEventReceiver; use codex_exec_server::ExecProcessFuture; @@ -27,7 +28,6 @@ use core_test_support::skip_if_sandbox; use core_test_support::test_codex::test_env as remote_test_env; use pretty_assertions::assert_eq; use std::collections::HashMap; -use std::num::NonZeroUsize; use std::path::PathBuf; use std::sync::Arc; use tokio::sync::Notify; @@ -179,8 +179,7 @@ async fn exec_command_with_tty( collected_output.total_bytes(), )) .unwrap_or(usize::MAX); - let output_omitted_bytes = NonZeroUsize::new(collected_output.omitted_bytes()); - let collected = collected_output.to_bytes_with_omission_marker(); + let collected = CapturedStream::from(&collected_output); let has_exited = process.has_exited(); let exit_code = process.exit_code(); let response_process_id = if process_started_alive && !has_exited { @@ -212,7 +211,6 @@ async fn exec_command_with_tty( process_id: response_process_id, exit_code, original_token_count: Some(original_token_count), - output_omitted_bytes, hook_command: Some(cmd.to_string()), }) } diff --git a/codex-rs/core/src/unified_exec/output.rs b/codex-rs/core/src/unified_exec/output.rs new file mode 100644 index 0000000000..d035ca0ba9 --- /dev/null +++ b/codex-rs/core/src/unified_exec/output.rs @@ -0,0 +1,63 @@ +//! Text retained for one unified exec output stream and the bytes omitted by its collection cap. + +use codex_utils_head_tail_buffer::HeadTailBuffer; +use codex_utils_head_tail_buffer::format_output_omission_marker; +use codex_utils_output_truncation::TruncationPolicy; +use codex_utils_output_truncation::approx_token_count; +use codex_utils_output_truncation::formatted_truncate_text; +use codex_utils_output_truncation::truncate_text; + +#[derive(Clone, Debug, Default, PartialEq)] +pub(crate) struct CapturedStream { + pub(crate) text: String, + pub(crate) omitted_bytes: usize, +} + +impl From<&str> for CapturedStream { + fn from(text: &str) -> Self { + Self { + text: text.to_string(), + omitted_bytes: 0, + } + } +} + +impl From<&HeadTailBuffer> for CapturedStream { + fn from(buffer: &HeadTailBuffer) -> Self { + Self { + text: String::from_utf8_lossy(&buffer.to_bytes_with_omission_marker()).into_owned(), + omitted_bytes: buffer.omitted_bytes(), + } + } +} + +impl CapturedStream { + pub(crate) fn render( + &self, + policy: Option, + original_tokens: Option, + ) -> String { + let truncated = match policy { + Some(policy) if self.omitted_bytes == 0 => { + return formatted_truncate_text(&self.text, policy); + } + Some(policy) => truncate_text(&self.text, policy), + None if self.omitted_bytes == 0 => return self.text.clone(), + None => self.text.clone(), + }; + let marker = format_output_omission_marker(self.omitted_bytes); + let notice = if truncated.contains(&marker) { + String::new() + } else { + format!("{marker}\n") + }; + if policy.is_none_or(|policy| self.text.len() <= policy.byte_budget()) { + if notice.is_empty() { + return truncated; + } + return format!("{notice}{truncated}"); + } + let tokens = original_tokens.unwrap_or_else(|| approx_token_count(&self.text)); + format!("Warning: truncated output (original token count: {tokens})\n{notice}\n{truncated}") + } +} diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index b07b410be6..b55848d9b1 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -1,3 +1,4 @@ +use crate::unified_exec::output::CapturedStream; use rand::Rng; use std::cmp::Reverse; use std::collections::HashMap; @@ -686,8 +687,8 @@ impl UnifiedExecProcessManager { )) .unwrap_or(usize::MAX); let output_omitted_bytes = NonZeroUsize::new(collected_output.omitted_bytes()); - let collected = collected_output.to_bytes_with_omission_marker(); - let text = String::from_utf8_lossy(&collected).to_string(); + let collected = CapturedStream::from(&collected_output); + let text = collected.text.clone(); let chunk_id = generate_chunk_id(); if deferred_network_approval .as_ref() @@ -872,7 +873,6 @@ impl UnifiedExecProcessManager { process_id: response_process_id, exit_code, original_token_count: Some(original_token_count), - output_omitted_bytes, hook_command: Some(request.hook_command.clone()), }; @@ -1061,8 +1061,7 @@ impl UnifiedExecProcessManager { collected_output.total_bytes(), )) .unwrap_or(usize::MAX); - let output_omitted_bytes = NonZeroUsize::new(collected_output.omitted_bytes()); - let collected = collected_output.to_bytes_with_omission_marker(); + let collected = CapturedStream::from(&collected_output); let chunk_id = generate_chunk_id(); if network_approval .as_ref() @@ -1126,7 +1125,6 @@ impl UnifiedExecProcessManager { process_id, exit_code, original_token_count: Some(original_token_count), - output_omitted_bytes, hook_command: Some(hook_command), };