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

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<PluginCommandAttribution>,
-    transcript: Arc<Mutex<HeadTailBuffer>>,
+    output_buffer: Arc<Mutex<OutputBuffers>>,
     started_at: Instant,
     network_denial_monitor: Option<tokio::task::JoinHandle<()>>,
     plugin_metrics_sidecar: Option<SharedPluginMetricsSidecar>,
@@ -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<String>,
     plugin_attribution: Option<PluginCommandAttribution>,
-    transcript: Arc<Mutex<HeadTailBuffer>>,
+    output_buffer: Arc<Mutex<OutputBuffers>>,
     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<String>,
     plugin_attribution: Option<PluginCommandAttribution>,
-    transcript: Arc<Mutex<HeadTailBuffer>>,
+    output_buffer: Arc<Mutex<OutputBuffers>>,
     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<Mutex<HeadTailBuffer>>,
+    output_buffer: &Arc<Mutex<OutputBuffers>>,
     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)]

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<UnifiedExecProcess>,
     stdout_tx: tokio::sync::broadcast::Sender<Vec<u8>>,
     exit_tx: tokio::sync::oneshot::Sender<i32>,
-    transcript: Arc<tokio::sync::Mutex<HeadTailBuffer>>,
+    output_buffer: Arc<tokio::sync::Mutex<OutputBuffers>>,
     context: UnifiedExecContext,
     rx_event: async_channel::Receiver<Event>,
 }
@@ -59,13 +59,13 @@ async fn streaming_output_harness() -> anyhow::Result<StreamingOutputHarness> {
         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,

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<const MAX_BYTES: usize = UNIFIED_EXEC_OUTPUT_MAX_BYTES> {
+    pub(super) pending: HeadTailBuffer<MAX_BYTES>,
+    pub(super) transcript: HeadTailBuffer<MAX_BYTES>,
+}
+
+impl<const MAX_BYTES: usize> OutputBuffers<MAX_BYTES> {
+    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<const MAX_BYTES: usize = UNIFIED_EXEC_OUTPUT_MAX_BYTES> {
-    pub(crate) output_buffer: Arc<Mutex<HeadTailBuffer<MAX_BYTES>>>,
+    pub(crate) output_buffer: Arc<Mutex<OutputBuffers<MAX_BYTES>>>,
     pub(crate) output_notify: Arc<Notify>,
     pub(crate) output_closed: Arc<AtomicBool>,
     pub(crate) output_closed_notify: Arc<Notify>,
@@ -92,8 +107,6 @@ pub(crate) struct UnifiedExecProcess {
     process_handle: ProcessHandle,
     output_tx: broadcast::Sender<Vec<u8>>,
     output: OutputHandles,
-    // Keep completion output independent of polling and lossy delta delivery.
-    transcript: Arc<Mutex<HeadTailBuffer>>,
     output_drained: Arc<Notify>,
     interaction_lock: Arc<Mutex<()>>,
     state_tx: watch::Sender<ProcessState>,
@@ -123,7 +136,7 @@ impl UnifiedExecProcess {
         spawn_lifecycle: Option<SpawnLifecycleHandle>,
     ) -> 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<Mutex<HeadTailBuffer>> {
-        Arc::clone(&self.transcript)
-    }
-
     pub(super) fn output_receiver(&self) -> tokio::sync::broadcast::Receiver<Vec<u8>> {
         self.output_tx.subscribe()
     }
@@ -285,7 +293,7 @@ impl UnifiedExecProcess {
 
     async fn snapshot_output(&self) -> Vec<u8> {
         let guard = self.output.output_buffer.lock().await;
-        guard.to_bytes()
+        guard.pending.to_bytes()
     }
 
     pub(crate) fn sandbox_type(&self) -> Option<SandboxType> {
@@ -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<Mutex<HeadTailBuffer>>,
         output_tx: broadcast::Sender<Vec<u8>>,
         state_tx: watch::Sender<ProcessState>,
     ) -> 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<Vec<u8>>,
         output_handles: OutputHandles,
-        transcript: Arc<Mutex<HeadTailBuffer>>,
         output_tx: broadcast::Sender<Vec<u8>>,
     ) -> 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);

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<PluginCommandAttribution>,
-    transcript: Arc<tokio::sync::Mutex<HeadTailBuffer>>,
+    output_buffer: Arc<tokio::sync::Mutex<OutputBuffers>>,
     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<DeferredNetworkApproval>,
         network_denial_monitor: Option<tokio::task::JoinHandle<()>>,
         metrics_sidecar: Option<PluginMetricsSidecar>,
-        transcript: Arc<tokio::sync::Mutex<HeadTailBuffer>>,
+        output_buffer: Arc<tokio::sync::Mutex<OutputBuffers>>,
         initial_exec_command_active: Arc<AtomicBool>,
     ) {
         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 {

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),