From 4891c4e35fe8b306433f6c0f04dd74bd9e4257e9 Mon Sep 17 00:00:00 2001 From: Steve Coffey Date: Thu, 24 Sep 2026 01:33:41 +0000 Subject: [PATCH] Update unified exec output buffers atomically (#47712) ## Why Appending output to the completion transcript and pending polling buffer under separate locks allows cancellation to interrupt a chunk update between the two buffers. ## What changed Store both buffers in `OutputBuffers` behind a shared mutex and append each chunk to both while holding that lock. Polling drains only the pending buffer, while completion events read the retained transcript. Apply this to both local and exec-server output handling. ## Testing Extend the repeated-drain test to verify that the completion transcript preserves the expected bounded output after polling drains the pending buffer. GitOrigin-RevId: ef8519582937b0d8c42e46375226501338df4a01 --- .../core/src/unified_exec/async_watcher.rs | 24 +++++------ .../src/unified_exec/async_watcher_tests.rs | 40 ++++++++++++------- codex-rs/core/src/unified_exec/process.rs | 35 ++++++++-------- .../core/src/unified_exec/process_manager.rs | 23 ++++++----- .../src/unified_exec/process_manager_tests.rs | 13 +++--- 5 files changed, 75 insertions(+), 60 deletions(-) diff --git a/codex-rs/core/src/unified_exec/async_watcher.rs b/codex-rs/core/src/unified_exec/async_watcher.rs index 702b9442ac..f30c750cbb 100644 --- a/codex-rs/core/src/unified_exec/async_watcher.rs +++ b/codex-rs/core/src/unified_exec/async_watcher.rs @@ -9,6 +9,7 @@ use tokio::time::Sleep; use super::SharedPluginMetricsSidecar; use super::UnifiedExecContext; +use super::process::OutputBuffers; use super::process::OutputHandles; use super::process::UnifiedExecProcess; use super::take_plugin_metrics_sidecar; @@ -20,7 +21,6 @@ use crate::tools::events::ToolEmitter; use crate::tools::events::ToolEventCtx; use crate::tools::events::ToolEventFailure; use crate::tools::events::ToolEventStage; -use crate::unified_exec::head_tail_buffer::HeadTailBuffer; use codex_core_plugins::PluginCommandAttribution; use codex_protocol::exec_output::ExecToolCallOutput; use codex_protocol::exec_output::StreamOutput; @@ -162,7 +162,7 @@ pub(crate) fn spawn_exit_watcher( cwd: PathUri, process_id: i32, plugin_attribution: Option, - transcript: Arc>, + output_buffer: Arc>, started_at: Instant, network_denial_monitor: Option>, plugin_metrics_sidecar: Option, @@ -203,7 +203,7 @@ pub(crate) fn spawn_exit_watcher( cwd, Some(process_id.to_string()), plugin_attribution, - transcript, + output_buffer, String::new(), message, duration, @@ -231,7 +231,7 @@ pub(crate) fn spawn_exit_watcher( cwd, Some(process_id.to_string()), plugin_attribution, - transcript, + output_buffer, String::new(), exit_code, duration, @@ -338,13 +338,13 @@ pub(crate) async fn emit_exec_end_for_unified_exec( cwd: PathUri, process_id: Option, plugin_attribution: Option, - transcript: Arc>, + output_buffer: Arc>, fallback_output: String, exit_code: i32, duration: Duration, timed_out: bool, ) { - let aggregated_output = resolve_aggregated_output(&transcript, fallback_output).await; + let aggregated_output = resolve_aggregated_output(&output_buffer, fallback_output).await; let output = ExecToolCallOutput { exit_code, stdout: StreamOutput::new(aggregated_output.clone()), @@ -390,13 +390,13 @@ pub(crate) async fn emit_failed_exec_end_for_unified_exec( cwd: PathUri, process_id: Option, plugin_attribution: Option, - transcript: Arc>, + output_buffer: Arc>, fallback_output: String, message: String, duration: Duration, ) { let stdout = if fallback_output.is_empty() { - resolve_aggregated_output(&transcript, fallback_output).await + resolve_aggregated_output(&output_buffer, fallback_output).await } else { fallback_output }; @@ -457,15 +457,15 @@ fn utf8_boundary(bytes: &[u8]) -> usize { } async fn resolve_aggregated_output( - transcript: &Arc>, + output_buffer: &Arc>, fallback: String, ) -> String { - let guard = transcript.lock().await; - if guard.retained_bytes() == 0 { + let guard = output_buffer.lock().await; + if guard.transcript.retained_bytes() == 0 { return fallback; } - String::from_utf8_lossy(&guard.to_bytes_with_omission_marker()).to_string() + String::from_utf8_lossy(&guard.transcript.to_bytes_with_omission_marker()).to_string() } #[cfg(test)] diff --git a/codex-rs/core/src/unified_exec/async_watcher_tests.rs b/codex-rs/core/src/unified_exec/async_watcher_tests.rs index 0b6aac6ad4..d5e1b9ba3e 100644 --- a/codex-rs/core/src/unified_exec/async_watcher_tests.rs +++ b/codex-rs/core/src/unified_exec/async_watcher_tests.rs @@ -9,8 +9,8 @@ use super::utf8_boundary; use crate::session::tests::make_session_and_context_with_rx; use crate::unified_exec::UnifiedExecContext; use crate::unified_exec::UnifiedExecProcessManager; -use crate::unified_exec::head_tail_buffer::HeadTailBuffer; use crate::unified_exec::process::NoopSpawnLifecycle; +use crate::unified_exec::process::OutputBuffers; use crate::unified_exec::process::UnifiedExecProcess; use codex_protocol::items::CommandExecutionStatus; use codex_protocol::items::TurnItem; @@ -28,7 +28,7 @@ struct StreamingOutputHarness { process: Arc, stdout_tx: tokio::sync::broadcast::Sender>, exit_tx: tokio::sync::oneshot::Sender, - transcript: Arc>, + output_buffer: Arc>, context: UnifiedExecContext, rx_event: async_channel::Receiver, } @@ -59,13 +59,13 @@ async fn streaming_output_harness() -> anyhow::Result { tokio_util::sync::CancellationToken::new(), "streaming-output-test".to_string(), ); - let transcript = process.transcript(); + let output_buffer = Arc::clone(&process.output_handles().output_buffer); Ok(StreamingOutputHarness { process, stdout_tx, exit_tx, - transcript, + output_buffer, context, rx_event, }) @@ -81,7 +81,7 @@ async fn completed_output_preserves_bytes_before_subscription( process, stdout_tx, exit_tx, - transcript, + output_buffer, context, rx_event, } = streaming_output_harness().await?; @@ -110,7 +110,7 @@ async fn completed_output_preserves_bytes_before_subscription( cwd, /*process_id*/ 123, /*plugin_attribution*/ None, - transcript, + output_buffer, Instant::now(), /*network_denial_monitor*/ None, /*plugin_metrics_sidecar*/ None, @@ -147,7 +147,7 @@ async fn streaming_output_preserves_multibyte_characters_across_chunks() -> anyh process, stdout_tx, exit_tx, - transcript, + output_buffer, rx_event, context, } = streaming_output_harness().await?; @@ -177,7 +177,11 @@ async fn streaming_output_preserves_multibyte_characters_across_chunks() -> anyh } ); assert_eq!( - transcript.lock().await.to_bytes_with_omission_marker(), + output_buffer + .lock() + .await + .transcript + .to_bytes_with_omission_marker(), "é".as_bytes() ); assert!(rx_event.try_recv().is_err()); @@ -191,7 +195,7 @@ async fn streaming_output_finishes_on_close_without_waiting_for_grace() -> anyho process, stdout_tx, exit_tx, - transcript, + output_buffer, context, .. } = streaming_output_harness().await?; @@ -219,7 +223,11 @@ async fn streaming_output_finishes_on_close_without_waiting_for_grace() -> anyho "output close should finish before the grace fallback: {elapsed:?}" ); assert_eq!( - transcript.lock().await.to_bytes_with_omission_marker(), + output_buffer + .lock() + .await + .transcript + .to_bytes_with_omission_marker(), b"LATE-OUTPUT-MARKER\xc3" ); @@ -232,7 +240,7 @@ async fn streaming_output_keeps_grace_as_fallback_without_close() -> anyhow::Res process, stdout_tx, exit_tx, - transcript, + output_buffer, rx_event, context, } = streaming_output_harness().await?; @@ -255,7 +263,11 @@ async fn streaming_output_keeps_grace_as_fallback_without_close() -> anyhow::Res "missing output close should use the grace fallback: {elapsed:?}" ); assert_eq!( - transcript.lock().await.to_bytes_with_omission_marker(), + output_buffer + .lock() + .await + .transcript + .to_bytes_with_omission_marker(), vec![0xc3] ); let event = rx_event.try_recv().expect("receive final output delta"); @@ -281,7 +293,7 @@ async fn exit_watcher_waits_for_late_network_denial_before_classifying_end() -> process, stdout_tx, exit_tx, - transcript, + output_buffer, mut context, rx_event, } = streaming_output_harness().await?; @@ -314,7 +326,7 @@ async fn exit_watcher_waits_for_late_network_denial_before_classifying_end() -> cwd, /*process_id*/ 123, /*plugin_attribution*/ None, - transcript, + output_buffer, Instant::now(), Some(network_denial_monitor), /*plugin_metrics_sidecar*/ None, diff --git a/codex-rs/core/src/unified_exec/process.rs b/codex-rs/core/src/unified_exec/process.rs index 9f425b7a46..8305c02894 100644 --- a/codex-rs/core/src/unified_exec/process.rs +++ b/codex-rs/core/src/unified_exec/process.rs @@ -58,10 +58,25 @@ pub(crate) struct NoopSpawnLifecycle; impl SpawnLifecycle for NoopSpawnLifecycle {} +/// Output pending model polling and the retained completion transcript. +/// Append both under the same lock so cancellation cannot split an update. +#[derive(Default)] +pub(crate) struct OutputBuffers { + pub(super) pending: HeadTailBuffer, + pub(super) transcript: HeadTailBuffer, +} + +impl OutputBuffers { + pub(super) fn push_chunk(&mut self, chunk: &[u8]) { + self.pending.push_chunk(chunk); + self.transcript.push_chunk(chunk); + } +} + /// Shared output state exposed to polling and streaming consumers. #[derive(Clone)] pub(crate) struct OutputHandles { - pub(crate) output_buffer: Arc>>, + pub(crate) output_buffer: Arc>>, pub(crate) output_notify: Arc, pub(crate) output_closed: Arc, pub(crate) output_closed_notify: Arc, @@ -92,8 +107,6 @@ pub(crate) struct UnifiedExecProcess { process_handle: ProcessHandle, output_tx: broadcast::Sender>, output: OutputHandles, - // Keep completion output independent of polling and lossy delta delivery. - transcript: Arc>, output_drained: Arc, interaction_lock: Arc>, state_tx: watch::Sender, @@ -123,7 +136,7 @@ impl UnifiedExecProcess { spawn_lifecycle: Option, ) -> Self { let output = OutputHandles { - output_buffer: Arc::new(Mutex::new(HeadTailBuffer::default())), + output_buffer: Arc::new(Mutex::new(OutputBuffers::default())), output_notify: Arc::new(Notify::new()), output_closed: Arc::new(AtomicBool::new(false)), output_closed_notify: Arc::new(Notify::new()), @@ -137,7 +150,6 @@ impl UnifiedExecProcess { process_handle, output_tx, output, - transcript: Arc::new(Mutex::new(HeadTailBuffer::default())), output_drained, interaction_lock: Arc::new(Mutex::new(())), state_tx, @@ -179,10 +191,6 @@ impl UnifiedExecProcess { &self.output } - pub(super) fn transcript(&self) -> Arc> { - Arc::clone(&self.transcript) - } - pub(super) fn output_receiver(&self) -> tokio::sync::broadcast::Receiver> { self.output_tx.subscribe() } @@ -285,7 +293,7 @@ impl UnifiedExecProcess { async fn snapshot_output(&self) -> Vec { let guard = self.output.output_buffer.lock().await; - guard.to_bytes() + guard.pending.to_bytes() } pub(crate) fn sandbox_type(&self) -> Option { @@ -367,7 +375,6 @@ impl UnifiedExecProcess { managed.output_task = Some(Self::spawn_local_output_task( output_rx, managed.output_handles().clone(), - managed.transcript(), managed.output_tx.clone(), )); @@ -417,7 +424,6 @@ impl UnifiedExecProcess { managed.output_task = Some(Self::spawn_exec_server_output_task( started, output_handles, - managed.transcript(), managed.output_tx.clone(), managed.state_tx.clone(), )); @@ -446,7 +452,6 @@ impl UnifiedExecProcess { fn spawn_exec_server_output_task( started: StartedExecProcess, output_handles: OutputHandles, - transcript: Arc>, output_tx: broadcast::Sender>, state_tx: watch::Sender, ) -> JoinHandle<()> { @@ -527,7 +532,6 @@ impl UnifiedExecProcess { } = response; for chunk in chunks.into_iter().filter(|chunk| chunk.seq > last_seq) { let bytes = chunk.chunk.into_inner(); - transcript.lock().await.push_chunk(&bytes); let mut guard = output_buffer.lock().await; guard.push_chunk(&bytes); drop(guard); @@ -571,7 +575,6 @@ impl UnifiedExecProcess { } last_seq = chunk.seq; let bytes = chunk.chunk.into_inner(); - transcript.lock().await.push_chunk(&bytes); let mut guard = output_buffer.lock().await; guard.push_chunk(&bytes); drop(guard); @@ -616,7 +619,6 @@ impl UnifiedExecProcess { fn spawn_local_output_task( mut receiver: tokio::sync::broadcast::Receiver>, output_handles: OutputHandles, - transcript: Arc>, output_tx: broadcast::Sender>, ) -> JoinHandle<()> { let OutputHandles { @@ -634,7 +636,6 @@ impl UnifiedExecProcess { loop { match receiver.recv().await { Ok(chunk) => { - transcript.lock().await.push_chunk(&chunk); let mut guard = output_buffer.lock().await; guard.push_chunk(&chunk); drop(guard); diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 68efb3bf17..4f31284c11 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -62,6 +62,7 @@ use crate::unified_exec::async_watcher::start_streaming_output; use crate::unified_exec::clamp_yield_time; use crate::unified_exec::generate_chunk_id; use crate::unified_exec::head_tail_buffer::HeadTailBuffer; +use crate::unified_exec::process::OutputBuffers; use crate::unified_exec::process::OutputHandles; use crate::unified_exec::process::SpawnLifecycleHandle; use crate::unified_exec::process::UnifiedExecProcess; @@ -399,7 +400,7 @@ async fn emit_failed_initial_exec_end_if_unstored( request: &ExecCommandRequest, cwd: PathUri, plugin_attribution: Option, - transcript: Arc>, + output_buffer: Arc>, fallback_output: String, message: String, wall_time: Duration, @@ -418,7 +419,7 @@ async fn emit_failed_initial_exec_end_if_unstored( cwd, Some(request.process_id.to_string()), plugin_attribution, - transcript, + output_buffer, fallback_output, message, wall_time, @@ -557,7 +558,7 @@ impl UnifiedExecProcessManager { ) }); - let transcript = process.transcript(); + let output_buffer = Arc::clone(&process.output_handles().output_buffer); let model_context = context.step_context.model_context(); let mut event_ctx = ToolEventCtx::new( context.session.as_ref(), @@ -617,7 +618,7 @@ impl UnifiedExecProcessManager { deferred_network_approval.clone(), network_denial_monitor, metrics_sidecar, - Arc::clone(&transcript), + Arc::clone(&output_buffer), Arc::clone(&initial_exec_command_active), ) .await; @@ -684,7 +685,7 @@ impl UnifiedExecProcessManager { &request, cwd.clone(), plugin_attribution.clone(), - Arc::clone(&transcript), + Arc::clone(&output_buffer), text.clone(), message.clone(), wall_time, @@ -706,7 +707,7 @@ impl UnifiedExecProcessManager { &request, cwd.clone(), plugin_attribution.clone(), - Arc::clone(&transcript), + Arc::clone(&output_buffer), text.clone(), message.clone(), wall_time, @@ -785,7 +786,7 @@ impl UnifiedExecProcessManager { &request, cwd.clone(), plugin_attribution.clone(), - Arc::clone(&transcript), + Arc::clone(&output_buffer), text.clone(), message.clone(), wall_time, @@ -809,7 +810,7 @@ impl UnifiedExecProcessManager { cwd.clone(), Some(process_id.to_string()), plugin_attribution.clone(), - Arc::clone(&transcript), + Arc::clone(&output_buffer), text.clone(), exit, wall_time, @@ -1206,7 +1207,7 @@ impl UnifiedExecProcessManager { network_approval: Option, network_denial_monitor: Option>, metrics_sidecar: Option, - transcript: Arc>, + output_buffer: Arc>, initial_exec_command_active: Arc, ) { let plugin_metrics_sidecar = @@ -1246,7 +1247,7 @@ impl UnifiedExecProcessManager { cwd, process_id, plugin_attribution, - transcript, + output_buffer, started_at, network_denial_monitor, plugin_metrics_sidecar, @@ -1610,7 +1611,7 @@ impl UnifiedExecProcessManager { let mut wait_for_output = None; { let mut guard = output_buffer.lock().await; - drained_output = std::mem::take(&mut *guard); + drained_output = std::mem::take(&mut guard.pending); has_drained_output = drained_output.retained_bytes() > 0 || drained_output.omitted_bytes() > 0; if !has_drained_output { 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 74cfeb4592..5ce38c4930 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -336,7 +336,7 @@ fn initial_exec_yield_time_has_no_platform_floor() { #[tokio::test] async fn output_collection_stays_bounded_across_repeated_drains() { let chunks: [&[u8]; 4] = [b"01234567", b"89ABCDEF", b"ghijklmnopq", b"rs"]; - let output_buffer = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::<10>::default())); + let output_buffer = Arc::new(tokio::sync::Mutex::new(OutputBuffers::<10>::default())); let output_notify = Arc::new(Notify::new()); let output_closed = Arc::new(AtomicBool::new(false)); let output_closed_notify = Arc::new(Notify::new()); @@ -360,7 +360,7 @@ async fn output_collection_stays_bounded_across_repeated_drains() { output_notify.notify_one(); tokio::time::timeout(Duration::from_secs(1), async { loop { - if output_buffer.lock().await.retained_bytes() == 0 { + if output_buffer.lock().await.pending.retained_bytes() == 0 { break; } tokio::task::yield_now().await; @@ -382,11 +382,12 @@ async fn output_collection_stays_bounded_across_repeated_drains() { expected.push_chunk(chunk); } assert_eq!(collected, expected); + assert_eq!(output_buffer.lock().await.transcript, expected); } #[tokio::test] async fn output_collection_preserves_omissions_from_drained_buffer() { - let mut buffered_output = HeadTailBuffer::<10>::default(); + let mut buffered_output = OutputBuffers::<10>::default(); buffered_output.push_chunk(&[b'a'; 10]); buffered_output.push_chunk(b"overflow"); let mut expected = HeadTailBuffer::<10>::default(); @@ -477,8 +478,8 @@ async fn failed_initial_end_for_unstored_process_uses_fallback_output() { prefix_rule: None, }; - let transcript = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default())); - transcript.lock().await.push_chunk(b"PARTIAL_TRANSCRIPT"); + let output_buffer = Arc::new(tokio::sync::Mutex::new(OutputBuffers::default())); + output_buffer.lock().await.push_chunk(b"PARTIAL_TRANSCRIPT"); emit_failed_initial_exec_end_if_unstored( /*process_started_alive*/ false, @@ -488,7 +489,7 @@ async fn failed_initial_end_for_unstored_process_uses_fallback_output() { #[allow(deprecated)] turn.cwd.clone().into(), /*plugin_attribution*/ None, - transcript, + output_buffer, "PRE_DENIAL_MARKER".to_string(), "Network access denied".to_string(), Duration::from_millis(7),