diff --git a/codex-rs/core/src/unified_exec/async_watcher.rs b/codex-rs/core/src/unified_exec/async_watcher.rs index 1537fb974a..2b9b132d9c 100644 --- a/codex-rs/core/src/unified_exec/async_watcher.rs +++ b/codex-rs/core/src/unified_exec/async_watcher.rs @@ -116,7 +116,7 @@ pub(crate) fn spawn_exit_watcher( process_id: i32, transcript: Arc>, started_at: Instant, -) { +) -> tokio::task::JoinHandle<()> { let exit_token = process.cancellation_token(); let output_drained = process.output_drained_notify(); @@ -157,7 +157,7 @@ pub(crate) fn spawn_exit_watcher( ) .await; } - }); + }) } async fn process_chunk( diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 114c026b8d..0fd317de22 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -159,6 +159,7 @@ impl Default for UnifiedExecProcessManager { struct ProcessEntry { process: Arc, + exit_watcher: Option>, call_id: String, process_id: i32, cwd: AbsolutePathBuf, diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index ad39d62b96..284b02ddf3 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -126,6 +126,7 @@ async fn exec_command_with_tty( if process_started_alive { let entry = ProcessEntry { process: Arc::clone(&process), + exit_watcher: None, call_id: context.call_id.clone(), process_id, cwd: cwd.clone(), @@ -672,6 +673,7 @@ async fn terminating_initial_exec_command_rechecks_initial_response_state() -> a process_id, ProcessEntry { process, + exit_watcher: None, call_id: "call".to_string(), process_id, cwd, @@ -745,6 +747,7 @@ async fn terminating_during_stdin_poll_returns_exited_response() -> anyhow::Resu process_id, ProcessEntry { process: Arc::clone(&process), + exit_watcher: None, call_id: "call".to_string(), process_id, cwd, diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index f175b0ac04..cb7d690c74 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -580,7 +580,10 @@ impl UnifiedExecProcessManager { process_id, .. } => (Some(process_id), exit_code), - ProcessStatus::Exited { exit_code, entry } => { + ProcessStatus::Exited { + exit_code, + mut entry, + } => { if let Err(message) = finish_deferred_network_approval_after_process_exit_for_session( Some(&context.session), @@ -590,6 +593,9 @@ impl UnifiedExecProcessManager { { return Err(fail_process_with_message(entry.process.as_ref(), message)); } + if let Some(exit_watcher) = entry.exit_watcher.take() { + let _ = exit_watcher.await; + } process.check_for_sandbox_denial_with_text(&text).await?; (None, exit_code) } @@ -897,8 +903,21 @@ impl UnifiedExecProcessManager { transcript: Arc>, initial_exec_command_active: Arc, ) { + let exit_watcher = spawn_exit_watcher( + Arc::clone(&process), + Arc::clone(&context.session), + Arc::clone(&context.turn), + context.call_id.clone(), + command.to_vec(), + cwd_uri, + path_convention, + process_id, + transcript, + started_at, + ); let entry = ProcessEntry { process: Arc::clone(&process), + exit_watcher: Some(exit_watcher), call_id: context.call_id.clone(), process_id, cwd: cwd.clone(), @@ -921,19 +940,6 @@ impl UnifiedExecProcessManager { unregister_network_approval_for_entry(&pruned_entry).await; pruned_entry.process.terminate(); } - - spawn_exit_watcher( - Arc::clone(&process), - Arc::clone(&context.session), - Arc::clone(&context.turn), - context.call_id.clone(), - command.to_vec(), - cwd_uri, - path_convention, - process_id, - transcript, - started_at, - ); } pub(crate) async fn open_session_with_exec_env( diff --git a/codex-rs/core/tests/suite/unified_exec.rs b/codex-rs/core/tests/suite/unified_exec.rs index 946ef666f2..94d0a144f5 100644 --- a/codex-rs/core/tests/suite/unified_exec.rs +++ b/codex-rs/core/tests/suite/unified_exec.rs @@ -756,7 +756,6 @@ async fn unified_exec_full_lifecycle_with_background_end_event() -> Result<()> { let mut begin_event = None; let mut end_event = None; - let mut task_completed = false; loop { let msg = wait_for_event(&test.codex, |_| true).await; @@ -768,15 +767,13 @@ async fn unified_exec_full_lifecycle_with_background_end_event() -> Result<()> { "expected a single ExecCommandEnd event for this call id" ); end_event = Some(ev); - if task_completed && end_event.is_some() { - break; - } } EventMsg::TurnComplete(_) => { - task_completed = true; - if task_completed && end_event.is_some() { - break; - } + assert!( + end_event.is_some(), + "short-lived command should emit ExecCommandEnd before TurnComplete" + ); + break; } _ => {} }