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