diff --git a/codex-rs/core/src/unified_exec/process.rs b/codex-rs/core/src/unified_exec/process.rs index 09e0cd1824..b3ddb02952 100644 --- a/codex-rs/core/src/unified_exec/process.rs +++ b/codex-rs/core/src/unified_exec/process.rs @@ -37,6 +37,8 @@ use super::process_state::ProcessState; use crate::shell_snapshot::ShellSnapshotFile; const EARLY_EXIT_GRACE_PERIOD: Duration = Duration::from_millis(150); +const REMOTE_OUTPUT_DRAIN_GRACE: Duration = Duration::from_secs(5); + /// Output pending model polling and the retained completion transcript. /// Append both under the same lock so cancellation cannot split an update. #[derive(Default)] @@ -60,6 +62,9 @@ pub(crate) struct OutputHandles, pub(crate) output_closed_notify: Arc, pub(crate) cancellation_token: CancellationToken, + pub(super) drain_policy: OutputDrainPolicy, + /// The initial exec call returned or was cancelled, so background draining may begin. + pub(crate) collection_finished: CancellationToken, } struct OutputTaskGuard { @@ -102,6 +107,8 @@ pub(crate) struct UnifiedExecProcess { output_task: Option>, sandbox_type: Option, timed_out: AtomicBool, + /// Explicit termination overrides retention of remote output after exit. + pub(super) termination_requested: AtomicBool, // The shell may still need to replay this file after process startup returns. pub(crate) _shell_snapshot: Option>, } @@ -117,13 +124,19 @@ impl std::fmt::Debug for UnifiedExecProcess { } impl UnifiedExecProcess { - fn new(process_handle: ProcessHandle, sandbox_type: Option) -> Self { + fn new( + process_handle: ProcessHandle, + sandbox_type: Option, + drain_policy: OutputDrainPolicy, + ) -> Self { let output = OutputHandles { 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()), cancellation_token: CancellationToken::new(), + drain_policy, + collection_finished: CancellationToken::new(), }; let output_drained = Arc::new(Notify::new()); let (output_tx, _) = broadcast::channel(64); @@ -140,6 +153,7 @@ impl UnifiedExecProcess { output_task: None, sandbox_type, timed_out: AtomicBool::new(false), + termination_requested: AtomicBool::new(false), _shell_snapshot: None, } } @@ -219,6 +233,7 @@ impl UnifiedExecProcess { } fn finish_termination(&self) { + self.termination_requested.store(true, Ordering::Release); self.output.cancellation_token.cancel(); if let Some(output_task) = &self.output_task { output_task.abort(); @@ -351,6 +366,7 @@ impl UnifiedExecProcess { let mut managed = Self::new( ProcessHandle::Local(Box::new(process_handle)), Some(sandbox_type), + OutputDrainPolicy::BoundedAfterExit, ); managed.output_task = Some(Self::spawn_local_output_task( output_rx, @@ -400,7 +416,7 @@ impl UnifiedExecProcess { // Older peers do not report this field. In that case, skip local // classification rather than attributing a violation to a guessed backend. let sandbox_type = started.sandbox_type; - let mut managed = Self::new(process_handle, sandbox_type); + let mut managed = Self::new(process_handle, sandbox_type, output_drain_policy); let output_handles = managed.output_handles().clone(); managed.output_task = Some(Self::spawn_exec_server_output_task( started, @@ -408,6 +424,7 @@ impl UnifiedExecProcess { output_handles, managed.output_tx.clone(), managed.state_tx.clone(), + managed.interaction_lock(), )); let mut state_rx = managed.state_rx.clone(); @@ -437,6 +454,7 @@ impl UnifiedExecProcess { output_handles: OutputHandles, output_tx: broadcast::Sender>, state_tx: watch::Sender, + interaction_lock: Arc>, ) -> JoinHandle<()> { let OutputHandles { output_buffer, @@ -444,20 +462,42 @@ impl UnifiedExecProcess { output_closed, output_closed_notify, cancellation_token, + collection_finished, + .. } = output_handles; let process = started.process; let mut events = process.subscribe_events(); + let mut state_rx = state_tx.subscribe(); tokio::spawn(async move { let _output_task_guard = OutputTaskGuard { output_closed: Arc::clone(&output_closed), output_closed_notify: Arc::clone(&output_closed_notify), }; + let drain = async { + collection_finished.cancelled().await; + let _ = state_rx.wait_for(|state| state.has_exited).await; + tokio::time::sleep(REMOTE_OUTPUT_DRAIN_GRACE).await; + // An active poll keeps its own deadline, including transport recovery. + interaction_lock.lock_owned().await + }; + tokio::pin!(drain); + let mut final_drain = false; let mut last_seq: u64 = 0; loop { - let event = match events.recv().await { - Ok(event) => Some(event), - Err(broadcast::error::RecvError::Lagged(_)) => None, - Err(broadcast::error::RecvError::Closed) => { + let received = tokio::select! { + received = events.recv() => Some(received), + guard = &mut drain, if output_drain_policy == OutputDrainPolicy::WaitForOutputClosure + && !final_drain => { + // Later polls must remain able to collect output while this read recovers. + drop(guard); + final_drain = true; + None + } + }; + let event = match received { + Some(Ok(event)) => Some(event), + Some(Err(broadcast::error::RecvError::Lagged(_))) | None => None, + Some(Err(broadcast::error::RecvError::Closed)) => { let state = state_tx.borrow().clone(); let _ = state_tx.send_replace( state.failed("exec-server process event stream closed".to_string()), @@ -486,6 +526,8 @@ impl UnifiedExecProcess { || event_seq.is_some_and(|seq| seq > last_seq.saturating_add(1)) || missing_sandbox_denial { + // The final drain read waits for transport recovery and collects retained + // output before closing an idle session with inherited pipes still open. let response = match process .read( Some(last_seq), @@ -542,7 +584,7 @@ impl UnifiedExecProcess { if exited && output_drain_policy == OutputDrainPolicy::BoundedAfterExit { cancellation_token.cancel(); } - if closed { + if closed || final_drain { output_closed.store(true, Ordering::Release); output_closed_notify.notify_waiters(); cancellation_token.cancel(); @@ -602,6 +644,12 @@ impl UnifiedExecProcess { } } } + if final_drain { + tokio::spawn(async move { + let _ = process.terminate().await; + }); + } + cancellation_token.cancel(); }) } diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 25df52c2ca..cda8c59dd8 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -271,6 +271,7 @@ struct PreparedProcessHandles { struct InitialExecCommandGuard { active: Option>, metrics_sidecar: Option, + collection_finished: CancellationToken, } impl InitialExecCommandGuard { @@ -293,6 +294,7 @@ impl Drop for InitialExecCommandGuard { if let Some(active) = self.active.as_ref() { active.store(false, Ordering::Release); } + self.collection_finished.cancel(); } } @@ -390,7 +392,7 @@ fn fail_process_with_message(process: &UnifiedExecProcess, message: String) -> U #[allow(clippy::too_many_arguments)] async fn emit_failed_initial_exec_end_if_unstored( - process_started_alive: bool, + process_stored: bool, sandbox_type: Option, context: &UnifiedExecContext, request: &ExecCommandRequest, @@ -401,7 +403,7 @@ async fn emit_failed_initial_exec_end_if_unstored( message: String, wall_time: Duration, ) { - if process_started_alive { + if process_stored { return; } @@ -592,13 +594,30 @@ impl UnifiedExecProcessManager { ); emitter.emit(event_ctx, ToolEventStage::Begin).await; + // Serialize remote collection with later polls and final transcript publication. + let _interaction_guard = + if process.output_handles().drain_policy == OutputDrainPolicy::WaitForOutputClosure { + Some(process.interaction_lock().lock_owned().await) + } else { + None + }; 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. - let process_started_alive = !process.has_exited() && process.exit_code().is_none(); - let mut initial_exec_command_guard = if process_started_alive { + // Persist live sessions and remote output that remains pollable after an early exit. + let process_stored = (!process.has_exited() && process.exit_code().is_none()) + || (process.output_handles().drain_policy == OutputDrainPolicy::WaitForOutputClosure + && !process + .output_handles() + .output_closed + .load(Ordering::Acquire)); + let mut initial_exec_command_guard = InitialExecCommandGuard { + active: None, + metrics_sidecar, + collection_finished: process.output_handles().collection_finished.clone(), + }; + if process_stored { let initial_exec_command_active = Arc::new(AtomicBool::new(true)); + initial_exec_command_guard.active = Some(Arc::clone(&initial_exec_command_active)); self.store_process( Arc::clone(&process), context, @@ -613,21 +632,12 @@ impl UnifiedExecProcessManager { request.tty, deferred_network_approval.clone(), network_denial_monitor, - metrics_sidecar, + initial_exec_command_guard.metrics_sidecar.take(), Arc::clone(&output_buffer), Arc::clone(&initial_exec_command_active), ) .await; - InitialExecCommandGuard { - active: Some(initial_exec_command_active), - metrics_sidecar: None, - } - } else { - InitialExecCommandGuard { - active: None, - metrics_sidecar, - } - }; + } let yield_time_ms = clamp_yield_time(request.yield_time_ms); // For the initial exec_command call, we both stream output to events @@ -675,7 +685,7 @@ impl UnifiedExecProcessManager { ) .await; emit_failed_initial_exec_end_if_unstored( - process_started_alive, + process_stored, process.sandbox_type(), context, &request, @@ -697,7 +707,7 @@ impl UnifiedExecProcessManager { ) .await; emit_failed_initial_exec_end_if_unstored( - process_started_alive, + process_stored, process.sandbox_type(), context, &request, @@ -716,15 +726,19 @@ impl UnifiedExecProcessManager { return Err(UnifiedExecError::process_failed(message)); } let process_id = request.process_id; - let (response_process_id, exit_code) = if process_started_alive { - match self.refresh_process_state(process_id).await { + let (response_process_id, exit_code) = if process_stored { + let retention = if completion.is_some() { + OutputRetention::UntilExit + } else { + OutputRetention::UntilCollected + }; + match self.refresh_process_state(process_id, retention).await { ProcessStatus::Alive { exit_code, process_id, .. } => (Some(process_id), exit_code), ProcessStatus::Exited { exit_code, entry } => { - // Collection is finished even if the remote output stream remains open. entry.process.cancellation_token().cancel(); if let Err(message) = finish_deferred_network_approval_after_process_exit_for_session( @@ -778,7 +792,7 @@ impl UnifiedExecProcessManager { .await; if let Err(message) = finish_result { emit_failed_initial_exec_end_if_unstored( - process_started_alive, + process_stored, process.sandbox_type(), context, &request, @@ -981,9 +995,10 @@ impl UnifiedExecProcessManager { } = self .prepare_process_handles(process_id, &locked_process) .await?; - let mut status_after_write = None; - - if !request.input.is_empty() { + let remote_exited = + process.has_exited() && output.drain_policy == OutputDrainPolicy::WaitForOutputClosure; + // A write racing with remote exit still collects the final output. + if !request.input.is_empty() && !remote_exited { if !tty { if request.input == INTERRUPT { process.interrupt().await?; @@ -998,14 +1013,11 @@ impl UnifiedExecProcessManager { tokio::time::sleep(Duration::from_millis(100)).await; } Err(err) => { - let status = self.refresh_process_state(process_id).await; - if matches!(status, ProcessStatus::Exited { .. }) { - status_after_write = Some(status); - } else if matches!(err, UnifiedExecError::ProcessFailed { .. }) { - process.terminate(); - self.release_process_id(process_id).await; - return Err(err); - } else { + if !process.has_exited() { + if matches!(err, UnifiedExecError::ProcessFailed { .. }) { + process.terminate(); + self.release_process_id(process_id).await; + } return Err(err); } } @@ -1059,15 +1071,10 @@ impl UnifiedExecProcessManager { return Err(UnifiedExecError::process_failed(message)); } - // After polling, refresh_process_state tells us whether the PTY is - // still alive or has exited and been removed from the store; we thread - // that through so the handler can tag or suppress TerminalInteraction - // with an appropriate process_id and exit_code. - let status = if let Some(status) = status_after_write { - status - } else { - self.refresh_process_state(process_id).await - }; + // Only finalize remote sessions after both process exit and output collection. + let status = self + .refresh_process_state(process_id, OutputRetention::UntilCollected) + .await; let (process_id, exit_code, event_call_id) = match status { ProcessStatus::Alive { exit_code, @@ -1075,7 +1082,6 @@ impl UnifiedExecProcessManager { process_id, } => (Some(process_id), exit_code, call_id), ProcessStatus::Exited { exit_code, entry } => { - // Release the exit watcher after the final poll, not on remote process exit. entry.process.cancellation_token().cancel(); let call_id = entry.call_id.clone(); if let Err(message) = @@ -1130,7 +1136,11 @@ impl UnifiedExecProcessManager { Ok(response) } - async fn refresh_process_state(&self, process_id: i32) -> ProcessStatus { + async fn refresh_process_state( + &self, + process_id: i32, + retention: OutputRetention, + ) -> ProcessStatus { let mut store = self.process_store.lock().await; let Some(entry) = store.processes.get_mut(&process_id) else { return ProcessStatus::Unknown; @@ -1140,6 +1150,24 @@ impl UnifiedExecProcessManager { let process_id = entry.process_id; if entry.process.has_exited() { + let output = entry.process.output_handles(); + // Never wait for output under the shared process-store lock. Retain busy output + // for another poll; closed, empty output cannot grow again. + if retention == OutputRetention::UntilCollected + && output.drain_policy == OutputDrainPolicy::WaitForOutputClosure + && !entry.process.termination_requested.load(Ordering::Acquire) + && (!output.output_closed.load(Ordering::Acquire) + || output + .output_buffer + .try_lock() + .map_or(true, |buffer| buffer.pending.total_bytes() > 0)) + { + return ProcessStatus::Alive { + exit_code: None, + call_id: entry.call_id.clone(), + process_id, + }; + } let Some(entry) = store.remove(process_id) else { return ProcessStatus::Unknown; }; @@ -1233,13 +1261,6 @@ impl UnifiedExecProcessManager { store.processes.insert(process_id, entry); pruned_entry }; - // prune_processes_if_needed runs while holding process_store; do async - // network-approval cleanup only after dropping that lock. - if let Some(pruned_entry) = pruned_entry { - unregister_network_approval_for_entry(&pruned_entry).await; - pruned_entry.process.terminate(); - } - spawn_exit_watcher( Arc::clone(&process), context, @@ -1252,6 +1273,11 @@ impl UnifiedExecProcessManager { network_denial_monitor, plugin_metrics_sidecar, ); + // Install the watcher before cancellable cleanup, and outside the process-store lock. + if let Some(pruned_entry) = pruned_entry { + unregister_network_approval_for_entry(&pruned_entry).await; + pruned_entry.process.terminate(); + } } #[allow(clippy::too_many_arguments)] @@ -1589,6 +1615,7 @@ impl UnifiedExecProcessManager { output_closed, output_closed_notify, cancellation_token, + .. } = output; let mut collected = HeadTailBuffer::default(); let mut exit_signal_received = cancellation_token.is_cancelled(); @@ -1848,10 +1875,12 @@ impl UnifiedExecProcessManager { (Arc::clone(&entry.process), entry.process.has_exited()) }; - if !already_exited && process.terminate_confirmed().await.is_err() { + if already_exited { + // Stop the output producer even when only its inherited pipes remain open. + process.terminate(); + } else if process.terminate_confirmed().await.is_err() { return false; } - process.cancellation_token().cancel(); let entry = { let mut store = self.process_store.lock().await; @@ -1875,6 +1904,12 @@ impl UnifiedExecProcessManager { } } +#[derive(PartialEq, Eq)] +enum OutputRetention { + UntilExit, + UntilCollected, +} + enum ProcessStatus { Alive { exit_code: Option, 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 081ab93599..cc560d3cba 100644 --- a/codex-rs/core/src/unified_exec/process_manager_tests.rs +++ b/codex-rs/core/src/unified_exec/process_manager_tests.rs @@ -355,6 +355,8 @@ async fn output_collection_stays_bounded_across_repeated_drains() { output_closed: Arc::clone(&output_closed), output_closed_notify: Arc::clone(&output_closed_notify), cancellation_token: cancellation_token.clone(), + collection_finished: CancellationToken::new(), + drain_policy: OutputDrainPolicy::BoundedAfterExit, }; let collect = UnifiedExecProcessManager::collect_output_until_deadline( @@ -413,6 +415,8 @@ async fn output_collection_preserves_omissions_from_drained_buffer() { output_closed, output_closed_notify, cancellation_token, + collection_finished: CancellationToken::new(), + drain_policy: OutputDrainPolicy::BoundedAfterExit, }; let collected = UnifiedExecProcessManager::collect_output_until_deadline( @@ -495,7 +499,7 @@ async fn failed_initial_end_for_unstored_process_uses_fallback_output() { output_buffer.lock().await.push_chunk(b"PARTIAL_TRANSCRIPT"); emit_failed_initial_exec_end_if_unstored( - /*process_started_alive*/ false, + /*process_stored*/ false, Some(codex_protocol::sandbox::SandboxType::WindowsMxc), &context, &request, diff --git a/codex-rs/core/tests/suite/unified_exec_process_events.rs b/codex-rs/core/tests/suite/unified_exec_process_events.rs index 49ae63deb7..e9b98e4b55 100644 --- a/codex-rs/core/tests/suite/unified_exec_process_events.rs +++ b/codex-rs/core/tests/suite/unified_exec_process_events.rs @@ -64,6 +64,7 @@ use tokio_tungstenite::tungstenite::Message; const CALL_ID: &str = "pushed-remote-process-events"; const POLL_CALL_ID: &str = "poll-remote-process-events"; +const DRAIN_POLL_CALL_ID: &str = "poll-during-remote-drain"; const COMPLETE_OUTPUT: &str = "pushed remote output\n"; const RECOVERED_OUTPUT: &str = "recovered missing output\n"; const RETAINED_OUTPUT: &str = "retained output\n"; @@ -79,7 +80,9 @@ enum PushedExecScenario { OutputOpenAtDeadline, OutputOpenAtFinalPoll, OutputOpenWhenCancelled, + OutputOpenWhenResumableCancelled, OutputOpenWhenTerminated, + OutputOpenWhenTerminatedBeforeYield, CodexPath, CodexPathDefaultFeature, CodexPathDisabledFeature, @@ -240,6 +243,7 @@ async fn serve_exec_with_pushed_events( scenario: PushedExecScenario, process_started: Arc, output_ready: Arc, + drain_read_ready: Arc, ) -> PushedExecServerResult { let unrestricted_patch = matches!( scenario, @@ -434,9 +438,11 @@ async fn serve_exec_with_pushed_events( PushedExecScenario::OutputOpenAtDeadline | PushedExecScenario::OutputOpenAtFinalPoll | PushedExecScenario::OutputOpenWhenCancelled + | PushedExecScenario::OutputOpenWhenResumableCancelled | PushedExecScenario::OutputOpenWhenTerminated + | PushedExecScenario::OutputOpenWhenTerminatedBeforeYield ); - if output_stays_open { + if output_stays_open || matches!(scenario, PushedExecScenario::OutputAfterExit) { output_ready.notified().await; } let mut process_read_requests = 0; @@ -446,7 +452,9 @@ async fn serve_exec_with_pushed_events( | PushedExecScenario::OutputOpenAtDeadline | PushedExecScenario::OutputOpenAtFinalPoll | PushedExecScenario::OutputOpenWhenCancelled - | PushedExecScenario::OutputOpenWhenTerminated => { + | PushedExecScenario::OutputOpenWhenResumableCancelled + | PushedExecScenario::OutputOpenWhenTerminated + | PushedExecScenario::OutputOpenWhenTerminatedBeforeYield => { send_exec_server_json( &mut websocket, json!({ @@ -493,22 +501,36 @@ async fn serve_exec_with_pushed_events( ) .await; } + if matches!(scenario, PushedExecScenario::OutputAfterExit) { + // Keep the tail in transit until the initial collection yields. + output_ready.notified().await; + } // Model foreground output delayed in transit beyond the local post-exit grace. - tokio::time::sleep(Duration::from_millis(250)).await; - send_exec_server_json( - &mut websocket, - json!({ - "method": "process/output", - "params": { - "processId": &process_id, - "seq": 3, - "stream": "stdout", - "chunk": BASE64_STANDARD.encode("remote output\n"), - } - }), + // An active final poll must also outlive the five-second background drain. + tokio::time::sleep( + if matches!(scenario, PushedExecScenario::OutputOpenAtFinalPoll) { + Duration::from_secs(6) + } else { + Duration::from_millis(250) + }, ) .await; - if !output_stays_open { + if !matches!(scenario, PushedExecScenario::OutputOpenAtDeadline) { + send_exec_server_json( + &mut websocket, + json!({ + "method": "process/output", + "params": { + "processId": &process_id, + "seq": 3, + "stream": "stdout", + "chunk": BASE64_STANDARD.encode("remote output\n"), + } + }), + ) + .await; + } + if !output_stays_open || matches!(scenario, PushedExecScenario::OutputOpenAtFinalPoll) { send_exec_server_json( &mut websocket, json!({ @@ -604,23 +626,72 @@ async fn serve_exec_with_pushed_events( PushedExecScenario::ReplayGap => {} } + let mut pending_drain_response = None; loop { let wait = if output_stays_open { STARTUP_TIMEOUT } else { Duration::from_secs(/*secs*/ 5) }; - let request = read_exec_server_json(&mut websocket, wait).await; + let request = tokio::select! { + request = read_exec_server_json(&mut websocket, wait) => request, + _ = output_ready.notified(), if pending_drain_response.is_some() => { + send_exec_server_json( + &mut websocket, + pending_drain_response.take().expect("pending final read"), + ).await; + continue; + } + }; match request["method"].as_str() { + Some("fs/getMetadata") => { + // Starting the polling turn still needs filesystem metadata while the read waits. + send_exec_server_json( + &mut websocket, + json!({ + "id": request["id"], + "error": { "code": -32004, "message": "not found" } + }), + ) + .await; + } Some("process/read") => { process_read_requests += 1; let result = match scenario { + PushedExecScenario::OutputOpenAtDeadline => { + assert_eq!(request["params"]["afterSeq"], 2); + // Hold the final read until a later poll yields its buffered output. + // The tail remains recoverable even after that poll's deadline. + drain_read_ready.notify_one(); + json!({ + "chunks": [{ + "seq": 3, + "stream": "stdout", + "chunk": BASE64_STANDARD.encode("remote output\n"), + }], + "nextSeq": 4, + "exited": true, + "exitCode": 0, + "closed": false, + "failure": null, + "sandboxDenied": false, + }) + } + PushedExecScenario::OutputOpenWhenResumableCancelled => json!({ + "chunks": [], + "nextSeq": 4, + "exited": true, + "exitCode": 0, + "closed": false, + "failure": null, + "sandboxDenied": false, + }), PushedExecScenario::OutputAfterExit | PushedExecScenario::OutputAfterRecoveredExit - | PushedExecScenario::OutputOpenAtDeadline | PushedExecScenario::OutputOpenAtFinalPoll | PushedExecScenario::OutputOpenWhenCancelled - | PushedExecScenario::OutputOpenWhenTerminated => { + | PushedExecScenario::OutputOpenWhenTerminated + | PushedExecScenario::OutputOpenWhenTerminatedBeforeYield => { panic!("unexpected read after the delayed-output sequence") } PushedExecScenario::Complete @@ -701,14 +772,12 @@ async fn serve_exec_with_pushed_events( }) } }; - send_exec_server_json( - &mut websocket, - json!({ - "id": request["id"], - "result": result, - }), - ) - .await; + let response = json!({ "id": request["id"], "result": result }); + if matches!(scenario, PushedExecScenario::OutputOpenAtDeadline) { + pending_drain_response = Some(response); + } else { + send_exec_server_json(&mut websocket, response).await; + } if matches!(scenario, PushedExecScenario::ReplayGap) && process_read_requests == 1 { send_exec_server_json( &mut websocket, @@ -748,7 +817,9 @@ async fn serve_exec_with_pushed_events( #[test_case(PushedExecScenario::OutputOpenAtDeadline, ManagedNetworkScenario::None, /*foreign_cwd*/ false ; "output_open_at_deadline")] #[test_case(PushedExecScenario::OutputOpenAtFinalPoll, ManagedNetworkScenario::None, /*foreign_cwd*/ false ; "output_open_at_final_poll")] #[test_case(PushedExecScenario::OutputOpenWhenCancelled, ManagedNetworkScenario::None, /*foreign_cwd*/ false ; "output_open_when_cancelled")] +#[test_case(PushedExecScenario::OutputOpenWhenResumableCancelled, ManagedNetworkScenario::None, /*foreign_cwd*/ false ; "output_open_when_resumable_cancelled")] #[test_case(PushedExecScenario::OutputOpenWhenTerminated, ManagedNetworkScenario::None, /*foreign_cwd*/ false ; "output_open_when_terminated")] +#[test_case(PushedExecScenario::OutputOpenWhenTerminatedBeforeYield, ManagedNetworkScenario::None, /*foreign_cwd*/ false ; "output_open_when_terminated_before_yield")] #[test_case(PushedExecScenario::CodexPath, ManagedNetworkScenario::None, false ; "codex_path_is_exported_before_nested_command")] #[test_case(PushedExecScenario::CodexPathDefaultFeature, ManagedNetworkScenario::None, false ; "codex_path_is_not_exported_by_default")] #[test_case(PushedExecScenario::CodexPathDisabledFeature, ManagedNetworkScenario::None, false ; "disabled_codex_path_preserves_original_login_command")] @@ -784,8 +855,19 @@ async fn exec_command_consumes_pushed_remote_process_events( ) -> Result<()> { let finish_at_deadline = matches!(scenario, PushedExecScenario::OutputOpenAtDeadline); let finish_at_final_poll = matches!(scenario, PushedExecScenario::OutputOpenAtFinalPoll); - let cancel_after_exit = matches!(scenario, PushedExecScenario::OutputOpenWhenCancelled); - let terminate_after_exit = matches!(scenario, PushedExecScenario::OutputOpenWhenTerminated); + let resumable_cancel = matches!( + scenario, + PushedExecScenario::OutputOpenWhenResumableCancelled + ); + let cancel_after_exit = + resumable_cancel || matches!(scenario, PushedExecScenario::OutputOpenWhenCancelled); + let output_after_yield = matches!(scenario, PushedExecScenario::OutputAfterExit); + let terminate_initial_collection = matches!( + scenario, + PushedExecScenario::OutputOpenWhenTerminatedBeforeYield + ); + let terminate_after_exit = terminate_initial_collection + || matches!(scenario, PushedExecScenario::OutputOpenWhenTerminated); let explicit_path = "/user/configured/bin:/user/other/bin"; let managed_network_configured = !matches!(managed_network, ManagedNetworkScenario::None); let managed_network_enabled = matches!(managed_network, ManagedNetworkScenario::Enabled { .. }); @@ -856,7 +938,8 @@ async fn exec_command_consumes_pushed_remote_process_events( | PushedExecScenario::CodexPathDisabledFeature => "sh -c 'command -v rg'".to_string(), _ => "pwd".to_string(), }, - "yield_time_ms": 1_000, + "tty": output_after_yield, + "yield_time_ms": if resumable_cancel || terminate_initial_collection { 30_000 } else { 1_000 }, }) .to_string(), ), @@ -876,7 +959,7 @@ async fn exec_command_consumes_pushed_remote_process_events( if cancel_after_exit { responses.truncate(/*len*/ 1); } - if finish_at_final_poll { + if output_after_yield || finish_at_final_poll { responses.insert( /*index*/ 1, sse(vec![ @@ -884,21 +967,48 @@ async fn exec_command_consumes_pushed_remote_process_events( ev_function_call( POLL_CALL_ID, "write_stdin", - &json!({ "session_id": 1000, "chars": "", "yield_time_ms": 1_000 }).to_string(), + &json!({ + "session_id": 1000, + // Exit is already observed; skip stdin but still collect the tail. + "chars": if output_after_yield { "ignored\n" } else { "" }, + // The final poll receives output after the five-second background drain. + "yield_time_ms": if finish_at_final_poll { 10_000 } else { 1_000 }, + }) + .to_string(), ), ev_completed("resp-poll"), ]), ); } + if finish_at_deadline { + for call_id in [DRAIN_POLL_CALL_ID, POLL_CALL_ID] { + responses.push(sse(vec![ + ev_response_created(call_id), + ev_function_call( + call_id, + "write_stdin", + &json!({ "session_id": 1000, "chars": "", "yield_time_ms": 1_000 }).to_string(), + ), + ev_completed(call_id), + ])); + } + responses.push(sse(vec![ + ev_response_created("resp-drained"), + ev_assistant_message("msg-drained", "done"), + ev_completed("resp-drained"), + ])); + } let response_mock = mount_sse_sequence(&server, responses).await; let exec_server_url = format!("ws://{}", listener.local_addr()?); let process_started = Arc::new(AtomicBool::new(false)); let output_ready = Arc::new(tokio::sync::Notify::new()); + let drain_read_ready = Arc::new(tokio::sync::Notify::new()); let exec_server = tokio::spawn(serve_exec_with_pushed_events( listener, scenario, Arc::clone(&process_started), Arc::clone(&output_ready), + Arc::clone(&drain_read_ready), )); let mut builder = test_codex().with_exec_server_url(exec_server_url); if managed_network_configured { @@ -960,7 +1070,7 @@ timeout = 900 .expect("write managed-network test config"); }); } - if cancel_after_exit { + if matches!(scenario, PushedExecScenario::OutputOpenWhenCancelled) { builder = builder.with_cloud_config_bundle( CloudConfigBundleFixture::loader_with_enterprise_requirement( "[features]\nunified_exec = false\nshell_tool = true\n", @@ -1145,9 +1255,15 @@ timeout = 900 let mut command_end = None; let mut stopped = false; let mut turn_finished = false; + let mut drain_poll_started = false; let wait_for_end = finish_at_deadline || finish_at_final_poll || cancel_after_exit || terminate_after_exit; - let event_timeout = if finish_at_deadline || finish_at_final_poll || terminate_after_exit { + let event_timeout = if finish_at_deadline + || finish_at_final_poll + || terminate_after_exit + || resumable_cancel + || output_after_yield + { STARTUP_TIMEOUT } else { Duration::from_secs(/*secs*/ 5) @@ -1162,9 +1278,9 @@ timeout = 900 EventMsg::ExecCommandBegin(event) if event.call_id == CALL_ID => { saw_exec_command_begin = true; exec_command_begin_command = Some(event.command); - if finish_at_deadline || cancel_after_exit { - // Storage follows the streaming subscription. Wait for it before - // exiting so both deadline cleanup and cancellation see a live session. + if cancel_after_exit || output_after_yield || terminate_initial_collection { + // Storage follows the streaming subscription, so the fixture's + // exit and tail cannot race ahead of the cancellation trigger. timeout(STARTUP_TIMEOUT, async { while test.codex.list_background_terminals().await.is_empty() { tokio::task::yield_now().await; @@ -1227,10 +1343,19 @@ timeout = 900 stopped = true; } EventMsg::RawResponseItem(raw) - if (finish_at_final_poll || terminate_after_exit) + if finish_at_deadline + && matches!(&raw.item, ResponseItem::FunctionCallOutput { call_id, .. } if call_id.as_deref() == Some(DRAIN_POLL_CALL_ID)) => + { + output_ready.notify_one(); + } + EventMsg::RawResponseItem(raw) + if ((terminate_after_exit && !terminate_initial_collection) + || finish_at_deadline + || finish_at_final_poll + || output_after_yield) && matches!(&raw.item, ResponseItem::FunctionCallOutput { call_id, .. } if call_id.as_deref() == Some(CALL_ID)) => { - // Polling and explicit termination must observe exit after the initial call yielded. + // Release exit/output only after the initial exec_command yielded. output_ready.notify_one(); } EventMsg::ExecCommandEnd(event) if event.call_id == CALL_ID => { @@ -1239,6 +1364,16 @@ timeout = 900 EventMsg::TurnAborted(_) if cancel_after_exit => { turn_finished = true; } + EventMsg::TurnComplete(_) if finish_at_deadline && !drain_poll_started => { + timeout(STARTUP_TIMEOUT, drain_read_ready.notified()).await?; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Poll the command.".into(), + text_elements: Vec::new(), + }])) + .await?; + drain_poll_started = true; + } EventMsg::TurnComplete(_) if !cancel_after_exit => turn_finished = true, _ => {} } @@ -1421,10 +1556,24 @@ timeout = 900 "toolCallId": CALL_ID, }), ); + if resumable_cancel { + assert_eq!(exec_server_result.process_read_requests, 1); + // The producer is closed without another poll; its unread output stays collectable. + assert!( + test.codex + .terminate_background_terminal(/*process_id*/ 1000) + .await + ); + } if cancel_after_exit || terminate_after_exit { assert!(stopped); - assert_eq!(exec_server_result.process_read_requests, 0); - return Ok(()); + assert_eq!( + exec_server_result.process_read_requests, + usize::from(resumable_cancel) + ); + if !terminate_initial_collection { + return Ok(()); + } } if foreign_cwd && !managed_network_configured && !deny_read { let permissions = response_mock.requests()[0] @@ -1575,15 +1724,15 @@ timeout = 900 let request = response_mock .last_request() .context("model should receive the exec_command output")?; - let output_call_id = if finish_at_final_poll { - let (initial_output, _) = request + let output_call_id = if output_after_yield || finish_at_final_poll || finish_at_deadline { + let (initial, _) = request .function_call_output_content_and_success(CALL_ID) - .context("initial exec_command should yield a session")?; - assert!( - initial_output - .context("initial exec_command output should contain text")? - .contains("Process running with session ID 1000") - ); + .context("initial collection")?; + let initial = initial.context("initial output")?; + assert!(initial.contains("Process running with session ID 1000")); + if output_after_yield { + assert!(initial.contains("pushed ")); + } assert!( !test .codex @@ -1594,6 +1743,14 @@ timeout = 900 } else { CALL_ID }; + if finish_at_deadline { + let (pending, _) = request + .function_call_output_content_and_success(DRAIN_POLL_CALL_ID) + .context("poll during the final read should yield")?; + let pending = pending.context("poll should contain buffered output")?; + assert!(pending.contains("Process running with session ID 1000")); + assert!(pending.contains("pushed ")); + } let (output, success) = request .function_call_output_content_and_success(output_call_id) .context("exec_command output should be model visible")?; @@ -1605,12 +1762,18 @@ timeout = 900 | PushedExecScenario::OutputOpenAtDeadline | PushedExecScenario::OutputOpenAtFinalPoll | PushedExecScenario::OutputOpenWhenCancelled - | PushedExecScenario::OutputOpenWhenTerminated => { + | PushedExecScenario::OutputOpenWhenResumableCancelled + | PushedExecScenario::OutputOpenWhenTerminated + | PushedExecScenario::OutputOpenWhenTerminatedBeforeYield => { assert_ne!(success, Some(false)); assert!(saw_exec_command_begin); assert!(output.contains("Process exited with code 0")); assert!( - output.contains(COMPLETE_OUTPUT), + output.contains(if output_after_yield || finish_at_deadline { + "remote output\n" + } else { + COMPLETE_OUTPUT + }), "missing remote output tail: {output}" ); assert_eq!( @@ -1618,6 +1781,7 @@ timeout = 900 usize::from(matches!( scenario, PushedExecScenario::OutputAfterRecoveredExit + | PushedExecScenario::OutputOpenAtDeadline )) ); }