From 2c8b8194696d515a798e39f56efbbf2c04f16cf8 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 15 Jun 2026 14:50:53 +0200 Subject: [PATCH] bound unified exec output collection --- codex-rs/core/src/unified_exec/mod_tests.rs | 63 +++++++++++++++++-- .../core/src/unified_exec/process_manager.rs | 41 +++++++++--- 2 files changed, 89 insertions(+), 15 deletions(-) diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index ad39d62b96..7d45e85157 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -21,7 +21,6 @@ use codex_exec_server::WriteResponse; use codex_exec_server::WriteStatus; use codex_sandboxing::SandboxType; use codex_utils_output_truncation::TruncationPolicy; -use codex_utils_output_truncation::approx_token_count; use core_test_support::get_remote_test_env; use core_test_support::skip_if_sandbox; use core_test_support::test_codex::test_env as remote_test_env; @@ -163,7 +162,6 @@ 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 has_exited = process.has_exited(); let exit_code = process.exit_code(); let response_process_id = if process_started_alive && !has_exited { @@ -189,12 +187,12 @@ async fn exec_command_with_tty( event_call_id: context.call_id, chunk_id: generate_chunk_id(), wall_time, - raw_output: collected, + raw_output: collected.raw_output, truncation_policy: turn.truncation_policy, max_output_tokens: None, process_id: response_process_id, exit_code, - original_token_count: Some(approx_token_count(&text)), + original_token_count: Some(collected.original_token_count), hook_command: Some(cmd.to_string()), }) } @@ -351,6 +349,61 @@ fn head_tail_buffer_default_preserves_prefix_and_suffix() { assert!(rendered.ends_with(b"bc")); } +#[tokio::test] +async fn collect_output_bounds_multiple_drains_and_reports_original_size() { + let output_buffer = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default())); + let output_notify = Arc::new(Notify::new()); + let output_closed = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let output_closed_notify = Arc::new(Notify::new()); + let cancellation_token = tokio_util::sync::CancellationToken::new(); + let chunk_len = UNIFIED_EXEC_OUTPUT_MAX_BYTES * 3 / 4; + + output_buffer.lock().await.push_chunk(vec![b'a'; chunk_len]); + + let collector = tokio::spawn({ + let output_buffer = Arc::clone(&output_buffer); + let output_notify = Arc::clone(&output_notify); + let output_closed = Arc::clone(&output_closed); + let output_closed_notify = Arc::clone(&output_closed_notify); + let cancellation_token = cancellation_token.clone(); + async move { + 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 + } + }); + + 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 the first chunk"); + + output_buffer.lock().await.push_chunk(vec![b'b'; chunk_len]); + output_closed.store(true, std::sync::atomic::Ordering::Release); + output_closed_notify.notify_waiters(); + output_notify.notify_waiters(); + cancellation_token.cancel(); + + let collected = collector.await.expect("collector task should succeed"); + let half = UNIFIED_EXEC_OUTPUT_MAX_BYTES / 2; + let expected_output = [vec![b'a'; half], vec![b'b'; half]].concat(); + assert_eq!(collected.raw_output, expected_output); + assert_eq!(collected.original_token_count, chunk_len * 2 / 4); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn unified_exec_persists_across_requests() -> anyhow::Result<()> { skip_if_sandbox!(Ok(())); @@ -881,7 +934,7 @@ async fn unified_exec_uses_remote_exec_server_when_configured() -> anyhow::Resul ) .await; - assert!(String::from_utf8_lossy(&collected).contains("remote-unified-exec")); + assert!(String::from_utf8_lossy(&collected.raw_output).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 a44cc5253d..86347fc4b7 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -56,7 +56,7 @@ use codex_protocol::error::SandboxErr; use codex_protocol::protocol::ExecCommandSource; use codex_tools::ToolName; use codex_utils_absolute_path::AbsolutePathBuf; -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] = [ @@ -76,6 +76,11 @@ const NETWORK_ACCESS_DENIED_MESSAGE: &str = const LATE_NETWORK_DENIAL_GRACE_PERIOD: Duration = Duration::from_millis(100); const INTERRUPT: &str = "\u{3}"; +pub(super) struct CollectedOutput { + pub(super) raw_output: Vec, + pub(super) original_token_count: usize, +} + /// Test-only override for deterministic unified exec process IDs. /// /// In production builds this value should remain at its default (`false`) and @@ -462,7 +467,10 @@ impl UnifiedExecProcessManager { cancellation_token, } = process.output_handles(); let deadline = start + Duration::from_millis(yield_time_ms); - let collected = Self::collect_output_until_deadline( + let CollectedOutput { + raw_output: collected, + original_token_count, + } = Self::collect_output_until_deadline( &output_buffer, &output_notify, &output_closed, @@ -596,7 +604,6 @@ impl UnifiedExecProcessManager { (None, exit_code) }; - let original_token_count = approx_token_count(&text); let response = ExecCommandToolOutput { event_call_id: context.call_id.clone(), chunk_id, @@ -679,7 +686,10 @@ impl UnifiedExecProcessManager { }; let start = Instant::now(); let deadline = start + Duration::from_millis(yield_time_ms); - let collected = Self::collect_output_until_deadline( + let CollectedOutput { + raw_output: collected, + original_token_count, + } = Self::collect_output_until_deadline( &output_buffer, &output_notify, &output_closed, @@ -691,8 +701,6 @@ 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 chunk_id = generate_chunk_id(); if network_approval .as_ref() @@ -1120,10 +1128,11 @@ impl UnifiedExecProcessManager { cancellation_token: &CancellationToken, mut pause_state: Option>, mut deadline: Instant, - ) -> Vec { + ) -> CollectedOutput { 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 original_output_bytes = 0usize; let mut exit_signal_received = cancellation_token.is_cancelled(); let mut post_exit_deadline: Option = None; loop { @@ -1134,14 +1143,20 @@ impl UnifiedExecProcessManager { ) .await; let drained_chunks: Vec>; + let omitted_bytes: usize; let mut wait_for_output = None; { let mut guard = output_buffer.lock().await; + omitted_bytes = guard.omitted_bytes(); drained_chunks = guard.drain_chunks(); if drained_chunks.is_empty() { wait_for_output = Some(output_notify.notified()); } } + original_output_bytes = drained_chunks.iter().fold( + original_output_bytes.saturating_add(omitted_bytes), + |total, chunk| total.saturating_add(chunk.len()), + ); if drained_chunks.is_empty() { exit_signal_received |= cancellation_token.is_cancelled(); @@ -1189,7 +1204,7 @@ impl UnifiedExecProcessManager { } for chunk in drained_chunks { - collected.extend_from_slice(&chunk); + collected.push_chunk(chunk); } exit_signal_received |= cancellation_token.is_cancelled(); @@ -1198,7 +1213,13 @@ impl UnifiedExecProcessManager { } } - collected + CollectedOutput { + raw_output: collected.to_bytes(), + original_token_count: usize::try_from(approx_tokens_from_byte_count( + original_output_bytes, + )) + .unwrap_or(usize::MAX), + } } async fn extend_deadlines_while_paused(