diff --git a/codex-rs/core/src/unified_exec/async_watcher.rs b/codex-rs/core/src/unified_exec/async_watcher.rs index 283846d67a..702b9442ac 100644 --- a/codex-rs/core/src/unified_exec/async_watcher.rs +++ b/codex-rs/core/src/unified_exec/async_watcher.rs @@ -50,18 +50,11 @@ struct Emitter { struct Buffer { pending: Vec, - transcript: Arc>, emitter: Emitter, } -/// Spawn a background task that continuously reads from the PTY, appends to the -/// shared transcript, and emits ExecCommandOutputDelta events on UTF‑8 -/// boundaries. -pub(crate) fn start_streaming_output( - process: &UnifiedExecProcess, - context: &UnifiedExecContext, - transcript: Arc>, -) { +/// Spawn a background task that emits ExecCommandOutputDelta events on UTF‑8 boundaries. +pub(crate) fn start_streaming_output(process: &UnifiedExecProcess, context: &UnifiedExecContext) { let mut receiver = process.output_receiver(); let output_drained = process.output_drained_notify(); let exit_token = process.cancellation_token(); @@ -83,7 +76,6 @@ pub(crate) fn start_streaming_output( let mut output: Buffer = Buffer { pending: Vec::new(), - transcript, emitter, }; @@ -258,13 +250,7 @@ impl Buffer { "a frame must fit one UTF-8 scalar" ) }; - let Self { - pending, - transcript, - emitter, - } = self; - - transcript.lock().await.push_chunk(&bytes); + let Self { pending, emitter } = self; // Reuse a producer chunk when it fits, retaining only an incomplete // UTF-8 suffix for the next push. @@ -299,7 +285,6 @@ impl Buffer { async fn finish(self) { let Self { pending, - transcript: _, mut emitter, } = self; debug_assert!( 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 c2019a70ab..0b6aac6ad4 100644 --- a/codex-rs/core/src/unified_exec/async_watcher_tests.rs +++ b/codex-rs/core/src/unified_exec/async_watcher_tests.rs @@ -8,6 +8,7 @@ use super::start_streaming_output; 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::UnifiedExecProcess; @@ -58,8 +59,7 @@ async fn streaming_output_harness() -> anyhow::Result { tokio_util::sync::CancellationToken::new(), "streaming-output-test".to_string(), ); - let transcript = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default())); - start_streaming_output(&process, &context, Arc::clone(&transcript)); + let transcript = process.transcript(); Ok(StreamingOutputHarness { process, @@ -71,6 +71,76 @@ async fn streaming_output_harness() -> anyhow::Result { }) } +#[test_case::test_case(b""; "no_late_output")] +#[test_case::test_case(b"late\n"; "with_late_output")] +#[tokio::test] +async fn completed_output_preserves_bytes_before_subscription( + late_output: &[u8], +) -> anyhow::Result<()> { + let StreamingOutputHarness { + process, + stdout_tx, + exit_tx, + transcript, + context, + rx_event, + } = streaming_output_harness().await?; + + let collected = process.output_handles().output_notify.notified(); + tokio::pin!(collected); + collected.as_mut().enable(); + stdout_tx.send(b"early\n".to_vec())?; + collected.await; + + let model_output = UnifiedExecProcessManager::collect_output_until_deadline( + process.output_handles(), + /*pause_state*/ None, + Instant::now(), + ) + .await; + assert_eq!(model_output.to_bytes(), b"early\n"); + + start_streaming_output(&process, &context); + #[allow(deprecated)] + let cwd = context.step_context.turn.cwd.clone().into(); + spawn_exit_watcher( + Arc::clone(&process), + &context, + vec!["proof".to_string()], + cwd, + /*process_id*/ 123, + /*plugin_attribution*/ None, + transcript, + Instant::now(), + /*network_denial_monitor*/ None, + /*plugin_metrics_sidecar*/ None, + ); + stdout_tx.send(late_output.to_vec())?; + drop(stdout_tx); + exit_tx.send(0).expect("send exit"); + + let item = tokio::time::timeout(Duration::from_secs(2), async { + loop { + if let EventMsg::ItemCompleted(completed) = rx_event.recv().await?.msg + && let TurnItem::CommandExecution(item) = completed.item + { + return Ok::<_, async_channel::RecvError>(item); + } + } + }) + .await??; + let expected_output = String::from_utf8([b"early\n".as_slice(), late_output].concat())?; + assert_eq!( + (item.status, item.exit_code, item.aggregated_output), + ( + CommandExecutionStatus::Completed, + Some(0), + Some(expected_output) + ) + ); + Ok(()) +} + #[tokio::test] async fn streaming_output_preserves_multibyte_characters_across_chunks() -> anyhow::Result<()> { let StreamingOutputHarness { @@ -79,8 +149,9 @@ async fn streaming_output_preserves_multibyte_characters_across_chunks() -> anyh exit_tx, transcript, rx_event, - .. + context, } = streaming_output_harness().await?; + start_streaming_output(&process, &context); let output_drained = process.output_drained_notify(); let drained = output_drained.notified(); tokio::pin!(drained); @@ -121,8 +192,10 @@ async fn streaming_output_finishes_on_close_without_waiting_for_grace() -> anyho stdout_tx, exit_tx, transcript, + context, .. } = streaming_output_harness().await?; + start_streaming_output(&process, &context); let output_drained = process.output_drained_notify(); let drained = output_drained.notified(); tokio::pin!(drained); @@ -161,8 +234,9 @@ async fn streaming_output_keeps_grace_as_fallback_without_close() -> anyhow::Res exit_tx, transcript, rx_event, - .. + context, } = streaming_output_harness().await?; + start_streaming_output(&process, &context); let output_drained = process.output_drained_notify(); let drained = output_drained.notified(); tokio::pin!(drained); @@ -211,6 +285,7 @@ async fn exit_watcher_waits_for_late_network_denial_before_classifying_end() -> mut context, rx_event, } = streaming_output_harness().await?; + start_streaming_output(&process, &context); tokio::time::pause(); let process_for_late_denial = Arc::clone(&process); @@ -311,12 +386,10 @@ fn utf8_boundary_batches_malformed_output() { } #[tokio::test] -async fn streaming_output_bounds_invalid_bytes_and_keeps_the_full_transcript() { +async fn streaming_output_bounds_invalid_bytes() { let (session, turn, rx_event) = make_session_and_context_with_rx().await; - let transcript = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default())); let mut output = Buffer::<8> { pending: Vec::new(), - transcript: Arc::clone(&transcript), emitter: Emitter { remaining_deltas: 2, session, @@ -345,11 +418,4 @@ async fn streaming_output_bounds_invalid_bytes_and_keeps_the_full_transcript() { b"\xf0\x9f\x98\x80\xff\xff\xff".to_vec(), ] ); - - let mut expected_transcript = bytes.to_vec(); - expected_transcript.extend([0xfe, 0xfe]); - assert_eq!( - transcript.lock().await.to_bytes_with_omission_marker(), - expected_transcript - ); } diff --git a/codex-rs/core/src/unified_exec/process.rs b/codex-rs/core/src/unified_exec/process.rs index ecdea4d68b..9f425b7a46 100644 --- a/codex-rs/core/src/unified_exec/process.rs +++ b/codex-rs/core/src/unified_exec/process.rs @@ -92,6 +92,8 @@ 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, @@ -135,6 +137,7 @@ impl UnifiedExecProcess { process_handle, output_tx, output, + transcript: Arc::new(Mutex::new(HeadTailBuffer::default())), output_drained, interaction_lock: Arc::new(Mutex::new(())), state_tx, @@ -176,6 +179,10 @@ 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() } @@ -360,6 +367,7 @@ impl UnifiedExecProcess { managed.output_task = Some(Self::spawn_local_output_task( output_rx, managed.output_handles().clone(), + managed.transcript(), managed.output_tx.clone(), )); @@ -409,6 +417,7 @@ 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(), )); @@ -437,6 +446,7 @@ impl UnifiedExecProcess { fn spawn_exec_server_output_task( started: StartedExecProcess, output_handles: OutputHandles, + transcript: Arc>, output_tx: broadcast::Sender>, state_tx: watch::Sender, ) -> JoinHandle<()> { @@ -517,6 +527,7 @@ 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); @@ -560,6 +571,7 @@ 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); @@ -604,6 +616,7 @@ 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 { @@ -621,6 +634,7 @@ 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 cb50158963..68efb3bf17 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -557,7 +557,7 @@ impl UnifiedExecProcessManager { ) }); - let transcript = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default())); + let transcript = process.transcript(); let model_context = context.step_context.model_context(); let mut event_ctx = ToolEventCtx::new( context.session.as_ref(), @@ -595,7 +595,7 @@ impl UnifiedExecProcessManager { ); emitter.emit(event_ctx, ToolEventStage::Begin).await; - start_streaming_output(&process, context, Arc::clone(&transcript)); + start_streaming_output(&process, context); let start = Instant::now(); // Persist live sessions before the initial yield wait so interrupting the // turn cannot drop the last Arc and terminate the background process. diff --git a/codex-rs/core/tests/suite/unified_exec.rs b/codex-rs/core/tests/suite/unified_exec.rs index 4ca281058f..b17be39003 100644 --- a/codex-rs/core/tests/suite/unified_exec.rs +++ b/codex-rs/core/tests/suite/unified_exec.rs @@ -1136,9 +1136,9 @@ async fn unified_exec_full_lifecycle_with_background_end_event() -> Result<()> { let test = builder.build_with_auto_env(&server).await?; let call_id = "uexec-full-lifecycle"; - // This timing force the long-standing PTY + // Print before the subscriber attaches, then keep the process alive. let args = json!({ - "cmd": "sleep 0.5; printf 'HELLO-FULL-LIFECYCLE'", + "cmd": "printf 'EARLY-OUTPUT'; sleep 0.5; printf 'HELLO-FULL-LIFECYCLE'", "yield_time_ms": 1000, }); @@ -1205,10 +1205,9 @@ async fn unified_exec_full_lifecycle_with_background_end_event() -> Result<()> { end_event.process_id.is_some(), "end event should include process_id emitted by background watcher" ); - assert!( - end_event.aggregated_output.contains("HELLO-FULL-LIFECYCLE"), - "aggregated_output should contain the full PTY transcript; got {:?}", - end_event.aggregated_output + assert_eq!( + end_event.aggregated_output, + "EARLY-OUTPUTHELLO-FULL-LIFECYCLE" ); Ok(()) }