diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index a6dd9e1a741c..72018904dd2d 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -7,6 +7,7 @@ use crate::tools::TELEMETRY_PREVIEW_MAX_BYTES; use crate::tools::TELEMETRY_PREVIEW_MAX_LINES; use crate::tools::TELEMETRY_PREVIEW_TRUNCATION_NOTICE; use crate::turn_diff_tracker::TurnDiffTracker; +use crate::unified_exec::format_output_omission_marker; use crate::unified_exec::resolve_max_tokens; use codex_protocol::mcp::CallToolResult; use codex_protocol::models::FunctionCallOutputBody; @@ -17,10 +18,13 @@ use codex_protocol::models::function_call_output_content_items_to_text; use codex_tools::LoadableToolSpec; use codex_tools::ToolName; 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; use codex_utils_string::take_bytes_at_char_boundary; use serde::Serialize; use serde_json::Value as JsonValue; +use std::num::NonZeroUsize; use std::sync::Arc; use std::time::Duration; use tokio::sync::Mutex; @@ -319,6 +323,8 @@ pub struct ExecCommandToolOutput { 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, } @@ -406,7 +412,32 @@ impl ExecCommandToolOutput { pub(crate) fn truncated_output(&self, max_tokens: usize) -> String { let text = String::from_utf8_lossy(&self.raw_output).to_string(); - formatted_truncate_text(&text, TruncationPolicy::Tokens(max_tokens)) + let policy = TruncationPolicy::Tokens(max_tokens); + 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}" + ) } fn response_text(&self) -> String { diff --git a/codex-rs/core/src/tools/context_tests.rs b/codex-rs/core/src/tools/context_tests.rs index 6ac5b6978f25..caed87dfa5bb 100644 --- a/codex-rs/core/src/tools/context_tests.rs +++ b/codex-rs/core/src/tools/context_tests.rs @@ -434,6 +434,7 @@ fn exec_command_tool_output_formats_truncated_response() { process_id: None, exit_code: Some(0), original_token_count: Some(10), + output_omitted_bytes: None, hook_command: None, } .to_response_item("call-42", &payload); @@ -461,3 +462,42 @@ fn exec_command_tool_output_formats_truncated_response() { other => panic!("expected FunctionCallOutput, got {other:?}"), } } + +#[test] +fn exec_command_tool_output_preserves_omission_metadata_when_truncated() { + let payload = ToolPayload::Function { + arguments: "{}".to_string(), + }; + let marker = format_output_omission_marker(/*omitted_bytes*/ 123_456); + let raw_output = format!( + "HEAD-{}\n{marker}\nTAIL-{}", + "a".repeat(/*n*/ 100), + "z".repeat(/*n*/ 100) + ) + .into_bytes(); + let response = ExecCommandToolOutput { + event_call_id: "call-omitted".to_string(), + chunk_id: "abc123".to_string(), + wall_time: std::time::Duration::from_millis(/*millis*/ 1250), + raw_output, + 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, + } + .to_response_item("call-omitted", &payload); + + let ResponseInputItem::FunctionCallOutput { output, .. } = response else { + panic!("expected FunctionCallOutput"); + }; + let text = output + .body + .to_text() + .expect("exec output should serialize as text"); + assert!(text.contains("Original token count: 42000")); + assert!(text.contains("Warning: truncated output (original token count: 42000)")); + assert_eq!(text.matches(&marker).count(), 1); +} 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 546ad459e1c7..fc48f1ed81a8 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 @@ -335,6 +335,7 @@ impl ExecCommandHandler { process_id: None, exit_code: None, original_token_count: None, + output_omitted_bytes: None, hook_command: None, })); } @@ -367,9 +368,15 @@ impl ExecCommandHandler { .await { Ok(response) => Ok(boxed_tool_output(response)), - Err(UnifiedExecError::SandboxDenied { output, .. }) => { + Err(UnifiedExecError::SandboxDenied { + output, + original_token_count, + output_omitted_bytes, + .. + }) => { let output_text = output.aggregated_output.text; - let original_token_count = approx_token_count(&output_text); + let original_token_count = + original_token_count.unwrap_or_else(|| approx_token_count(&output_text)); Ok(boxed_tool_output(ExecCommandToolOutput { event_call_id: context.call_id.clone(), chunk_id: generate_chunk_id(), @@ -382,6 +389,7 @@ 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 6f859a8cc4fb..9b107dd83d00 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec_tests.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec_tests.rs @@ -301,6 +301,7 @@ async fn exec_command_post_tool_use_payload_uses_output_for_noninteractive_one_s 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; @@ -331,6 +332,7 @@ async fn exec_command_post_tool_use_payload_uses_output_for_interactive_completi 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; @@ -362,6 +364,7 @@ async fn exec_command_post_tool_use_payload_skips_running_sessions() { 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; @@ -388,6 +391,7 @@ async fn write_stdin_post_tool_use_payload_uses_original_exec_call_id_and_comman 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; @@ -419,6 +423,7 @@ async fn write_stdin_post_tool_use_payload_keeps_parallel_session_metadata_separ 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 { @@ -431,6 +436,7 @@ async fn write_stdin_post_tool_use_payload_keeps_parallel_session_metadata_separ 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/async_watcher.rs b/codex-rs/core/src/unified_exec/async_watcher.rs index 38010754e96e..c1d61ebdcee6 100644 --- a/codex-rs/core/src/unified_exec/async_watcher.rs +++ b/codex-rs/core/src/unified_exec/async_watcher.rs @@ -326,7 +326,7 @@ async fn resolve_aggregated_output( return fallback; } - String::from_utf8_lossy(&guard.to_bytes()).to_string() + String::from_utf8_lossy(&guard.to_bytes_with_omission_marker()).to_string() } #[cfg(test)] diff --git a/codex-rs/core/src/unified_exec/errors.rs b/codex-rs/core/src/unified_exec/errors.rs index b4e2882c9bbd..dd401046ae67 100644 --- a/codex-rs/core/src/unified_exec/errors.rs +++ b/codex-rs/core/src/unified_exec/errors.rs @@ -1,5 +1,6 @@ use codex_protocol::exec_output::ExecToolCallOutput; use codex_utils_path_uri::PathUri; +use std::num::NonZeroUsize; use thiserror::Error; #[derive(Debug, Error)] @@ -23,6 +24,8 @@ pub(crate) enum UnifiedExecError { SandboxDenied { message: String, output: ExecToolCallOutput, + original_token_count: Option, + output_omitted_bytes: Option, }, #[error("{path} is not valid on {}", std::env::consts::OS)] ForeignPath { path: PathUri }, @@ -38,6 +41,29 @@ impl UnifiedExecError { } pub(crate) fn sandbox_denied(message: String, output: ExecToolCallOutput) -> Self { - Self::SandboxDenied { message, output } + Self::SandboxDenied { + message, + output, + original_token_count: None, + output_omitted_bytes: None, + } + } + + pub(crate) fn with_output_collection_metadata( + self, + original_token_count: usize, + output_omitted_bytes: Option, + ) -> Self { + match self { + Self::SandboxDenied { + message, output, .. + } => Self::SandboxDenied { + message, + output, + original_token_count: Some(original_token_count), + output_omitted_bytes, + }, + other => other, + } } } diff --git a/codex-rs/core/src/unified_exec/head_tail_buffer.rs b/codex-rs/core/src/unified_exec/head_tail_buffer.rs index 52039e149597..fb04f34f0bb7 100644 --- a/codex-rs/core/src/unified_exec/head_tail_buffer.rs +++ b/codex-rs/core/src/unified_exec/head_tail_buffer.rs @@ -1,4 +1,5 @@ use crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES; +use crate::unified_exec::format_output_omission_marker; use std::collections::VecDeque; /// A capped buffer that preserves a stable prefix ("head") and suffix ("tail"), @@ -6,14 +7,13 @@ use std::collections::VecDeque; /// symmetric meaning 50% of the capacity is allocated to the head and 50% is /// allocated to the tail. #[derive(Debug)] +#[cfg_attr(test, derive(Eq, PartialEq))] pub(crate) struct HeadTailBuffer { max_bytes: usize, head_budget: usize, tail_budget: usize, - head: VecDeque>, - tail: VecDeque>, - head_bytes: usize, - tail_bytes: usize, + head: Vec, + tail: VecDeque, omitted_bytes: usize, } @@ -35,10 +35,8 @@ impl HeadTailBuffer { max_bytes, head_budget, tail_budget, - head: VecDeque::new(), + head: Vec::new(), tail: VecDeque::new(), - head_bytes: 0, - tail_bytes: 0, omitted_bytes: 0, } } @@ -47,7 +45,7 @@ impl HeadTailBuffer { #[allow(dead_code)] /// Total bytes currently retained by the buffer (head + tail). pub(crate) fn retained_bytes(&self) -> usize { - self.head_bytes.saturating_add(self.tail_bytes) + self.head.len().saturating_add(self.tail.len()) } // Used for tests. @@ -57,37 +55,32 @@ impl HeadTailBuffer { self.omitted_bytes } + /// Total bytes observed by the buffer, including bytes omitted by the cap. + pub(crate) fn total_bytes(&self) -> usize { + self.retained_bytes().saturating_add(self.omitted_bytes) + } + /// Append a chunk of bytes to the buffer. /// /// Bytes are first added to the head until the head budget is full; any /// remaining bytes are added to the tail, with older tail bytes being /// dropped to preserve the tail budget. pub(crate) fn push_chunk(&mut self, chunk: Vec) { + if chunk.is_empty() { + return; + } if self.max_bytes == 0 { self.omitted_bytes = self.omitted_bytes.saturating_add(chunk.len()); return; } // Fill the head budget first, then keep a capped tail. - if self.head_bytes < self.head_budget { - let remaining_head = self.head_budget.saturating_sub(self.head_bytes); - if chunk.len() <= remaining_head { - self.head_bytes = self.head_bytes.saturating_add(chunk.len()); - self.head.push_back(chunk); - return; - } - - // Split the chunk: part goes to head, remainder goes to tail. - let (head_part, tail_part) = chunk.split_at(remaining_head); - if !head_part.is_empty() { - self.head_bytes = self.head_bytes.saturating_add(head_part.len()); - self.head.push_back(head_part.to_vec()); - } - self.push_to_tail(tail_part.to_vec()); - return; + let remaining_head = self.head_budget.saturating_sub(self.head.len()); + let head_len = remaining_head.min(chunk.len()); + if head_len > 0 { + self.head.extend_from_slice(&chunk[..head_len]); } - - self.push_to_tail(chunk); + self.push_to_tail(&chunk[head_len..]); } /// Snapshot the retained output as a list of chunks. @@ -95,9 +88,13 @@ impl HeadTailBuffer { /// The returned chunks are ordered as: head chunks first, then tail chunks. /// Omitted bytes are not represented in the snapshot. pub(crate) fn snapshot_chunks(&self) -> Vec> { - let mut out = Vec::new(); - out.extend(self.head.iter().cloned()); - out.extend(self.tail.iter().cloned()); + let mut out = Vec::with_capacity(2); + if !self.head.is_empty() { + out.push(self.head.clone()); + } + if !self.tail.is_empty() { + out.push(self.tail.iter().copied().collect()); + } out } @@ -107,29 +104,58 @@ impl HeadTailBuffer { /// Omitted bytes are not represented in the returned value. pub(crate) fn to_bytes(&self) -> Vec { let mut out = Vec::with_capacity(self.retained_bytes()); - for chunk in self.head.iter() { - out.extend_from_slice(chunk); - } - for chunk in self.tail.iter() { - out.extend_from_slice(chunk); - } + out.extend_from_slice(&self.head); + out.extend(self.tail.iter().copied()); out } - /// Drain all retained chunks from the buffer and reset its state. - /// - /// The drained chunks are returned in head-then-tail order. Omitted bytes - /// are discarded along with the retained content. - pub(crate) fn drain_chunks(&mut self) -> Vec> { - let mut out: Vec> = self.head.drain(..).collect(); - out.extend(self.tail.drain(..)); - self.head_bytes = 0; - self.tail_bytes = 0; - self.omitted_bytes = 0; + /// Return the retained output with an explicit marker between the head and + /// tail when bytes were omitted. + pub(crate) fn to_bytes_with_omission_marker(&self) -> Vec { + if self.omitted_bytes == 0 { + return self.to_bytes(); + } + + let marker = format_output_omission_marker(self.omitted_bytes); + let marker_delimiter_bytes = 2; + let mut out = Vec::with_capacity( + self.retained_bytes() + .saturating_add(marker.len()) + .saturating_add(marker_delimiter_bytes), + ); + out.extend_from_slice(&self.head); + out.push(b'\n'); + out.extend_from_slice(marker.as_bytes()); + out.push(b'\n'); + out.extend(self.tail.iter().copied()); out } - fn push_to_tail(&mut self, chunk: Vec) { + /// Drain the retained output and omission metadata, resetting this buffer's + /// contents while preserving its configured capacity. + pub(crate) fn drain(&mut self) -> Self { + Self { + max_bytes: self.max_bytes, + head_budget: self.head_budget, + tail_budget: self.tail_budget, + head: std::mem::take(&mut self.head), + tail: std::mem::take(&mut self.tail), + omitted_bytes: std::mem::take(&mut self.omitted_bytes), + } + } + + /// Append retained output from another buffer and preserve any omissions it + /// already recorded. + pub(crate) fn push_buffer(&mut self, mut buffer: Self) { + self.push_chunk(std::mem::take(&mut buffer.head)); + self.push_chunk(buffer.tail.drain(..).collect()); + self.omitted_bytes = self.omitted_bytes.saturating_add(buffer.omitted_bytes); + } + + fn push_to_tail(&mut self, chunk: &[u8]) { + if chunk.is_empty() { + return; + } if self.tail_budget == 0 { self.omitted_bytes = self.omitted_bytes.saturating_add(chunk.len()); return; @@ -139,41 +165,26 @@ impl HeadTailBuffer { // This single chunk is larger than the whole tail budget. Keep only the last // tail_budget bytes and drop everything else. let start = chunk.len().saturating_sub(self.tail_budget); - let kept = chunk[start..].to_vec(); + let kept = &chunk[start..]; let dropped = chunk.len().saturating_sub(kept.len()); self.omitted_bytes = self .omitted_bytes - .saturating_add(self.tail_bytes) + .saturating_add(self.tail.len()) .saturating_add(dropped); self.tail.clear(); - self.tail_bytes = kept.len(); - self.tail.push_back(kept); + self.tail.extend(kept); return; } - self.tail_bytes = self.tail_bytes.saturating_add(chunk.len()); - self.tail.push_back(chunk); + self.tail.extend(chunk); self.trim_tail_to_budget(); } fn trim_tail_to_budget(&mut self) { - let mut excess = self.tail_bytes.saturating_sub(self.tail_budget); - while excess > 0 { - match self.tail.front_mut() { - Some(front) if excess >= front.len() => { - excess -= front.len(); - self.tail_bytes = self.tail_bytes.saturating_sub(front.len()); - self.omitted_bytes = self.omitted_bytes.saturating_add(front.len()); - self.tail.pop_front(); - } - Some(front) => { - front.drain(..excess); - self.tail_bytes = self.tail_bytes.saturating_sub(excess); - self.omitted_bytes = self.omitted_bytes.saturating_add(excess); - break; - } - None => break, - } + let excess = self.tail.len().saturating_sub(self.tail_budget); + if excess > 0 { + drop(self.tail.drain(..excess)); + self.omitted_bytes = self.omitted_bytes.saturating_add(excess); } } } diff --git a/codex-rs/core/src/unified_exec/head_tail_buffer_tests.rs b/codex-rs/core/src/unified_exec/head_tail_buffer_tests.rs index 34da5be5f9de..c086984d930e 100644 --- a/codex-rs/core/src/unified_exec/head_tail_buffer_tests.rs +++ b/codex-rs/core/src/unified_exec/head_tail_buffer_tests.rs @@ -16,6 +16,10 @@ fn keeps_prefix_and_suffix_when_over_budget() { let rendered = String::from_utf8_lossy(&buf.to_bytes()).to_string(); assert!(rendered.starts_with("01234")); assert!(rendered.ends_with("89ab")); + assert_eq!( + String::from_utf8_lossy(&buf.to_bytes_with_omission_marker()), + "01234\n... 2 bytes omitted ...\n789ab" + ); } #[test] @@ -40,17 +44,21 @@ fn head_budget_zero_keeps_only_last_byte_in_tail() { } #[test] -fn draining_resets_state() { +fn draining_resets_state_and_push_buffer_preserves_omissions() { let mut buf = HeadTailBuffer::new(/*max_bytes*/ 10); buf.push_chunk(b"0123456789".to_vec()); buf.push_chunk(b"ab".to_vec()); - let drained = buf.drain_chunks(); - assert!(!drained.is_empty()); + let drained = buf.drain(); + let mut collected = HeadTailBuffer::new(/*max_bytes*/ 10); + collected.push_buffer(drained); assert_eq!(buf.retained_bytes(), 0); assert_eq!(buf.omitted_bytes(), 0); assert_eq!(buf.to_bytes(), b"".to_vec()); + assert_eq!(collected.to_bytes(), b"01234789ab".to_vec()); + assert_eq!(collected.omitted_bytes(), 2); + assert_eq!(collected.total_bytes(), 12); } #[test] @@ -87,3 +95,20 @@ fn fills_head_then_tail_across_multiple_chunks() { assert_eq!(buf.to_bytes(), b"012346789a".to_vec()); assert_eq!(buf.omitted_bytes(), 1); } + +#[test] +fn empty_and_tiny_chunks_have_bounded_metadata() { + let mut buf = HeadTailBuffer::new(/*max_bytes*/ 10); + + for byte in b"0123456789ab" { + buf.push_chunk(Vec::new()); + buf.push_chunk(vec![*byte]); + } + + assert_eq!( + buf.snapshot_chunks(), + vec![b"01234".to_vec(), b"789ab".to_vec()] + ); + assert_eq!(buf.retained_bytes(), 10); + assert_eq!(buf.omitted_bytes(), 2); +} diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 34d88c814919..21c1a17f0965 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -178,6 +178,10 @@ pub(crate) fn resolve_max_tokens(max_tokens: Option) -> usize { max_tokens.unwrap_or(DEFAULT_MAX_OUTPUT_TOKENS) } +pub(crate) fn format_output_omission_marker(omitted_bytes: usize) -> String { + format!("... {omitted_bytes} bytes omitted ...") +} + pub(crate) fn generate_chunk_id() -> String { let mut rng = rng(); (0..6) diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index 81403c8e52a0..7e3cf8083aa8 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -22,12 +22,13 @@ use codex_exec_server::WriteStatus; use codex_sandboxing::SandboxType; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_output_truncation::TruncationPolicy; -use codex_utils_output_truncation::approx_token_count; +use codex_utils_output_truncation::approx_tokens_from_byte_count; use core_test_support::skip_if_no_remote_env; 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; @@ -154,7 +155,7 @@ async fn exec_command_with_tty( cancellation_token, } = process.output_handles(); let deadline = started_at + Duration::from_millis(yield_time_ms); - let collected = UnifiedExecProcessManager::collect_output_until_deadline( + let collected_output = UnifiedExecProcessManager::collect_output_until_deadline( &output_buffer, &output_notify, &output_closed, @@ -165,7 +166,12 @@ async fn exec_command_with_tty( ) .await; let wall_time = Instant::now().saturating_duration_since(started_at); - let text = String::from_utf8_lossy(&collected).to_string(); + let original_token_count = usize::try_from(approx_tokens_from_byte_count( + 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 has_exited = process.has_exited(); let exit_code = process.exit_code(); let response_process_id = if process_started_alive && !has_exited { @@ -196,7 +202,8 @@ async fn exec_command_with_tty( max_output_tokens: None, process_id: response_process_id, exit_code, - original_token_count: Some(approx_token_count(&text)), + original_token_count: Some(original_token_count), + output_omitted_bytes, hook_command: Some(cmd.to_string()), }) } @@ -327,17 +334,11 @@ fn push_chunk_preserves_prefix_and_suffix() { assert_eq!(buffer.retained_bytes(), UNIFIED_EXEC_OUTPUT_MAX_BYTES); let snapshot = buffer.snapshot_chunks(); - - let first = snapshot.first().expect("expected at least one chunk"); - assert_eq!(first.first(), Some(&b'a')); - assert!(snapshot.iter().any(|chunk| chunk.as_slice() == b"b")); - assert_eq!( - snapshot - .last() - .expect("expected at least one chunk") - .as_slice(), - b"c" - ); + let head_bytes = UNIFIED_EXEC_OUTPUT_MAX_BYTES / 2; + let tail_bytes = UNIFIED_EXEC_OUTPUT_MAX_BYTES - head_bytes; + let mut expected_tail = vec![b'a'; tail_bytes - 2]; + expected_tail.extend_from_slice(b"bc"); + assert_eq!(snapshot, vec![vec![b'a'; head_bytes], expected_tail]); } #[test] @@ -876,7 +877,8 @@ async fn unified_exec_uses_remote_exec_server_when_configured() -> anyhow::Resul /*pause_state*/ None, Instant::now() + Duration::from_millis(2_500), ) - .await; + .await + .to_bytes_with_omission_marker(); assert!(String::from_utf8_lossy(&collected).contains("remote-unified-exec")); Ok(()) diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index fe3ec45759bb..db233397b054 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -2,6 +2,7 @@ use rand::Rng; use std::cmp::Reverse; use std::collections::HashMap; use std::collections::HashSet; +use std::num::NonZeroUsize; use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; @@ -63,7 +64,7 @@ use codex_protocol::error::SandboxErr; use codex_protocol::protocol::ExecCommandSource; use codex_sandboxing::SandboxCommand; use codex_tools::ToolName; -use codex_utils_output_truncation::approx_token_count; +use codex_utils_output_truncation::approx_tokens_from_byte_count; use codex_utils_path_uri::PathUri; const UNIFIED_EXEC_ENV: [(&str, &str); 10] = [ @@ -487,7 +488,7 @@ impl UnifiedExecProcessManager { cancellation_token, } = process.output_handles(); let deadline = start + Duration::from_millis(yield_time_ms); - let collected = Self::collect_output_until_deadline( + let collected_output = Self::collect_output_until_deadline( &output_buffer, &output_notify, &output_closed, @@ -499,6 +500,12 @@ impl UnifiedExecProcessManager { .await; let wall_time = Instant::now().saturating_duration_since(start); + let original_token_count = usize::try_from(approx_tokens_from_byte_count( + 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 text = String::from_utf8_lossy(&collected).to_string(); let chunk_id = generate_chunk_id(); if deferred_network_approval @@ -565,7 +572,15 @@ impl UnifiedExecProcessManager { { return Err(fail_process_with_message(entry.process.as_ref(), message)); } - process.check_for_sandbox_denial_with_text(&text).await?; + process + .check_for_sandbox_denial_with_text(&text) + .await + .map_err(|err| { + err.with_output_collection_metadata( + original_token_count, + output_omitted_bytes, + ) + })?; (None, exit_code) } ProcessStatus::Unknown => { @@ -612,11 +627,15 @@ impl UnifiedExecProcessManager { .await; self.release_process_id(request.process_id).await; - process.check_for_sandbox_denial_with_text(&text).await?; + process + .check_for_sandbox_denial_with_text(&text) + .await + .map_err(|err| { + err.with_output_collection_metadata(original_token_count, output_omitted_bytes) + })?; (None, exit_code) }; - let original_token_count = approx_token_count(&text); let response = ExecCommandToolOutput { event_call_id: context.call_id.clone(), chunk_id, @@ -627,6 +646,7 @@ 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()), }; @@ -699,7 +719,7 @@ impl UnifiedExecProcessManager { }; let start = Instant::now(); let deadline = start + Duration::from_millis(yield_time_ms); - let collected = Self::collect_output_until_deadline( + let collected_output = Self::collect_output_until_deadline( &output_buffer, &output_notify, &output_closed, @@ -711,8 +731,12 @@ impl UnifiedExecProcessManager { .await; let wall_time = Instant::now().saturating_duration_since(start); - let text = String::from_utf8_lossy(&collected).to_string(); - let original_token_count = approx_token_count(&text); + let original_token_count = usize::try_from(approx_tokens_from_byte_count( + 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 chunk_id = generate_chunk_id(); if network_approval .as_ref() @@ -782,6 +806,7 @@ impl UnifiedExecProcessManager { process_id, exit_code, original_token_count: Some(original_token_count), + output_omitted_bytes, hook_command: Some(hook_command), }; @@ -1207,10 +1232,10 @@ impl UnifiedExecProcessManager { cancellation_token: &CancellationToken, mut pause_state: Option>, mut deadline: Instant, - ) -> Vec { + ) -> HeadTailBuffer { const POST_EXIT_CLOSE_WAIT_CAP: Duration = Duration::from_millis(50); - let mut collected: Vec = Vec::with_capacity(4096); + let mut collected = HeadTailBuffer::default(); let mut exit_signal_received = cancellation_token.is_cancelled(); let mut post_exit_deadline: Option = None; loop { @@ -1220,17 +1245,20 @@ impl UnifiedExecProcessManager { &mut post_exit_deadline, ) .await; - let drained_chunks: Vec>; + let drained_output: HeadTailBuffer; + let has_drained_output: bool; let mut wait_for_output = None; { let mut guard = output_buffer.lock().await; - drained_chunks = guard.drain_chunks(); - if drained_chunks.is_empty() { + drained_output = guard.drain(); + has_drained_output = + drained_output.retained_bytes() > 0 || drained_output.omitted_bytes() > 0; + if !has_drained_output { wait_for_output = Some(output_notify.notified()); } } - if drained_chunks.is_empty() { + if !has_drained_output { exit_signal_received |= cancellation_token.is_cancelled(); if exit_signal_received && output_closed.load(std::sync::atomic::Ordering::Acquire) { @@ -1275,9 +1303,7 @@ impl UnifiedExecProcessManager { continue; } - for chunk in drained_chunks { - collected.extend_from_slice(&chunk); - } + collected.push_buffer(drained_output); exit_signal_received |= cancellation_token.is_cancelled(); if Instant::now() >= deadline { diff --git a/codex-rs/core/src/unified_exec/process_manager_tests.rs b/codex-rs/core/src/unified_exec/process_manager_tests.rs index 20a930b5525c..9163862ae7a2 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -234,6 +234,93 @@ fn initial_exec_yield_time_has_no_platform_floor() { ); } +#[tokio::test] +async fn output_collection_stays_bounded_across_repeated_drains() { + let output_buffer = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default())); + let output_notify = Arc::new(Notify::new()); + let output_closed = Arc::new(AtomicBool::new(false)); + let output_closed_notify = Arc::new(Notify::new()); + let cancellation_token = CancellationToken::new(); + + let collect = UnifiedExecProcessManager::collect_output_until_deadline( + &output_buffer, + &output_notify, + &output_closed, + &output_closed_notify, + &cancellation_token, + /*pause_state*/ None, + Instant::now() + Duration::from_secs(5), + ); + let produce = async { + for byte in [b'a', b'b', b'c'] { + output_buffer.lock().await.push_chunk( + vec![byte; crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES], + ); + output_notify.notify_one(); + tokio::time::timeout(Duration::from_secs(1), async { + loop { + if output_buffer.lock().await.retained_bytes() == 0 { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("collector should drain each chunk"); + } + + output_closed.store(true, Ordering::Release); + cancellation_token.cancel(); + output_closed_notify.notify_waiters(); + output_notify.notify_waiters(); + }; + + let (collected, ()) = tokio::join!(collect, produce); + let mut expected = HeadTailBuffer::default(); + for byte in [b'a', b'b', b'c'] { + expected.push_chunk(vec![ + byte; + crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES + ]); + } + assert_eq!(collected, expected); +} + +#[tokio::test] +async fn output_collection_preserves_omissions_from_drained_buffer() { + let mut buffered_output = HeadTailBuffer::default(); + buffered_output.push_chunk(vec![ + b'a'; + crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES + ]); + buffered_output.push_chunk(b"overflow".to_vec()); + let mut expected = HeadTailBuffer::default(); + expected.push_chunk(vec![ + b'a'; + crate::unified_exec::UNIFIED_EXEC_OUTPUT_MAX_BYTES + ]); + expected.push_chunk(b"overflow".to_vec()); + let output_buffer = Arc::new(tokio::sync::Mutex::new(buffered_output)); + let output_notify = Arc::new(Notify::new()); + let output_closed = Arc::new(AtomicBool::new(true)); + let output_closed_notify = Arc::new(Notify::new()); + let cancellation_token = CancellationToken::new(); + cancellation_token.cancel(); + + let collected = UnifiedExecProcessManager::collect_output_until_deadline( + &output_buffer, + &output_notify, + &output_closed, + &output_closed_notify, + &cancellation_token, + /*pause_state*/ None, + Instant::now() + Duration::from_secs(1), + ) + .await; + + assert_eq!(collected, expected); +} + #[tokio::test] async fn network_denial_fallback_message_names_sandbox_network_proxy() { let message = network_denial_message_for_session(/*session*/ None, /*deferred*/ None).await; diff --git a/codex-rs/core/tests/suite/unified_exec.rs b/codex-rs/core/tests/suite/unified_exec.rs index 42e0214fa12e..46b361b73ea7 100644 --- a/codex-rs/core/tests/suite/unified_exec.rs +++ b/codex-rs/core/tests/suite/unified_exec.rs @@ -17,6 +17,7 @@ use codex_protocol::protocol::ExecCommandSource; use codex_protocol::protocol::ExecCommandStatus; use codex_protocol::protocol::Op; use codex_protocol::user_input::UserInput; +use codex_utils_output_truncation::approx_tokens_from_byte_count; use codex_utils_path_uri::PathUri; use core_test_support::TempDirExt; use core_test_support::assert_regex_match; @@ -2898,17 +2899,27 @@ async fn unified_exec_formats_large_output_summary() -> Result<()> { }); let test = builder.build_with_auto_env(&server).await?; - let script = r#"python3 - <<'PY' + let output_line = "token token \n"; + let output_repetitions = 100_000; + let original_output_bytes = + b"HEAD\n".len() + output_line.len() * output_repetitions + b"TAIL\n".len(); + let expected_original_token_count = + usize::try_from(approx_tokens_from_byte_count(original_output_bytes)).unwrap_or(usize::MAX); + let script = format!( + r#"python3 - <<'PY' import sys -sys.stdout.write("token token \n" * 5000) +sys.stdout.write("HEAD\n") +sys.stdout.write("token token \n" * {output_repetitions}) +sys.stdout.write("TAIL\n") PY -"#; +"# + ); let call_id = "uexec-large-output"; let args = serde_json::json!({ "cmd": script, "max_output_tokens": 100, - "yield_time_ms": 500, + "yield_time_ms": 3_000, }); let responses = vec![ @@ -2926,6 +2937,18 @@ PY submit_unified_exec_turn(&test, "summarize large output", PermissionProfile::Disabled).await?; + let end_event = wait_for_event_match(&test.codex, |event| match event { + EventMsg::ExecCommandEnd(event) if event.call_id == call_id => Some(event.clone()), + _ => None, + }) + .await; + assert!(end_event.aggregated_output.contains("HEAD\n")); + assert!(end_event.aggregated_output.contains("TAIL\n")); + assert_regex_match( + r"\.\.\. \d+ bytes omitted \.\.\.", + &end_event.aggregated_output, + ); + wait_for_event(&test.codex, |event| { matches!(event, EventMsg::TurnComplete(_)) }) @@ -2942,13 +2965,16 @@ PY let large_output = outputs.get(call_id).expect("missing large output summary"); let output_text = large_output.output.replace("\r\n", "\n"); - let truncated_pattern = r"(?s)^Warning: truncated output \(original token count: \d+\)\nTotal output lines: \d+\n\n(token token \n){5,}.*…\d+ tokens truncated….*(token token \n){5,}$"; - assert_regex_match(truncated_pattern, &output_text); - - let original_tokens = large_output - .original_token_count - .expect("missing original_token_count for large output summary"); - assert!(original_tokens > 0); + assert!(output_text.starts_with(&format!( + "Warning: truncated output (original token count: {expected_original_token_count})\n" + ))); + assert_regex_match(r"\.\.\. \d+ bytes omitted \.\.\.", &output_text); + assert!(output_text.contains("HEAD\n")); + assert!(output_text.contains("TAIL\n")); + assert_eq!( + large_output.original_token_count, + Some(expected_original_token_count) + ); Ok(()) }