unified-exec: finish short-lived lifecycle before turn completion

codex-rs/core/src/unified_exec/async_watcher.rs

@@ -116,7 +116,7 @@ pub(crate) fn spawn_exit_watcher(
     process_id: i32,
     transcript: Arc<Mutex<HeadTailBuffer>>,
     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(

codex-rs/core/src/unified_exec/mod.rs

@@ -159,6 +159,7 @@ impl Default for UnifiedExecProcessManager {
 
 struct ProcessEntry {
     process: Arc<UnifiedExecProcess>,
+    exit_watcher: Option<tokio::task::JoinHandle<()>>,
     call_id: String,
     process_id: i32,
     cwd: AbsolutePathBuf,

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,

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<tokio::sync::Mutex<HeadTailBuffer>>,
         initial_exec_command_active: Arc<AtomicBool>,
     ) {
+        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(

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;
             }
             _ => {}
         }