Order unified exec lifecycle events reliably

codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs

@@ -7,9 +7,8 @@ use crate::tools::registry::CoreToolRuntime;
 use crate::tools::registry::PostToolUsePayload;
 use crate::tools::registry::PreToolUsePayload;
 use crate::tools::registry::ToolExecutor;
+use crate::unified_exec::WriteStdinInteractionEvent;
 use crate::unified_exec::WriteStdinRequest;
-use codex_protocol::protocol::EventMsg;
-use codex_protocol::protocol::TerminalInteractionEvent;
 use codex_tools::ToolName;
 use codex_tools::ToolSpec;
 use serde::Deserialize;
@@ -80,28 +79,16 @@ impl WriteStdinHandler {
                 yield_time_ms: args.yield_time_ms,
                 max_output_tokens: args.max_output_tokens,
                 truncation_policy: turn.model_info.truncation_policy.into(),
+                interaction_event: Some(WriteStdinInteractionEvent {
+                    session: &session,
+                    turn: &turn,
+                }),
             })
             .await
             .map_err(|err| {
                 FunctionCallError::RespondToModel(format!("write_stdin failed: {err}"))
             })?;
 
-        // Empty stdin is a background poll, so emit it only while there is
-        // still a live process for the UI to wait on. Non-empty stdin is a real
-        // terminal interaction and should remain visible even if it completes
-        // the process before the response returns.
-        if !args.chars.is_empty() || response.process_id.is_some() {
-            let process_id = response.process_id.unwrap_or(args.session_id);
-            let interaction = TerminalInteractionEvent {
-                call_id: response.event_call_id.clone(),
-                process_id: process_id.to_string(),
-                stdin: args.chars.clone(),
-            };
-            session
-                .send_event(turn.as_ref(), EventMsg::TerminalInteraction(interaction))
-                .await;
-        }
-
         Ok(boxed_tool_output(response))
     }
 }

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

@@ -1,12 +1,14 @@
 use std::pin::Pin;
 use std::sync::Arc;
+use std::sync::atomic::Ordering;
 
 use tokio::sync::Mutex;
 use tokio::time::Duration;
 use tokio::time::Instant;
 use tokio::time::Sleep;
 
 use super::UnifiedExecContext;
+use super::process::OutputHandles;
 use super::process::UnifiedExecProcess;
 use crate::exec::MAX_EXEC_OUTPUT_DELTAS_PER_CALL;
 use crate::session::session::Session;
@@ -45,6 +47,11 @@ pub(crate) fn start_streaming_output(
     let mut receiver = process.output_receiver();
     let output_drained = process.output_drained_notify();
     let exit_token = process.cancellation_token();
+    let OutputHandles {
+        output_closed,
+        output_closed_notify,
+        ..
+    } = process.output_handles();
 
     let session_ref = Arc::clone(&context.session);
     let turn_ref = Arc::clone(&context.turn);
@@ -57,8 +64,19 @@ pub(crate) fn start_streaming_output(
         let mut emitted_deltas: usize = 0;
 
         let mut grace_sleep: Option<Pin<Box<Sleep>>> = None;
+        let output_closed_notified = output_closed_notify.notified();
+        tokio::pin!(output_closed_notified);
+        let mut output_complete = false;
 
         loop {
+            // Register before checking the atomic so a close between the check
+            // and the select cannot miss the notification.
+            output_closed_notified.as_mut().enable();
+            if grace_sleep.is_some() && output_closed.load(Ordering::Acquire) {
+                output_complete = true;
+                break;
+            }
+
             tokio::select! {
                 _ = exit_token.cancelled(), if grace_sleep.is_none() => {
                     let deadline = Instant::now() + TRAILING_OUTPUT_GRACE;
@@ -70,18 +88,21 @@ pub(crate) fn start_streaming_output(
                         sleep.as_mut().await;
                     }
                 }, if grace_sleep.is_some() => {
-                    output_drained.notify_one();
                     break;
                 }
 
+                _ = &mut output_closed_notified, if grace_sleep.is_some() => {
+                    output_closed_notified.set(output_closed_notify.notified());
+                }
+
                 received = receiver.recv() => {
                     let chunk = match received {
                         Ok(chunk) => chunk,
                         Err(RecvError::Lagged(_)) => {
                             continue;
                         },
                         Err(RecvError::Closed) => {
-                            output_drained.notify_one();
+                            output_complete = true;
                             break;
                         }
                     };
@@ -98,6 +119,35 @@ pub(crate) fn start_streaming_output(
                 }
             }
         }
+
+        output_complete |= output_closed.load(Ordering::Acquire);
+        if output_complete {
+            // Output producers publish all chunks before setting output_closed
+            // with Release ordering, so the Acquire above makes this a final
+            // safe drain.
+            loop {
+                let chunk = match receiver.try_recv() {
+                    Ok(chunk) => chunk,
+                    Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) => continue,
+                    Err(
+                        tokio::sync::broadcast::error::TryRecvError::Empty
+                        | tokio::sync::broadcast::error::TryRecvError::Closed,
+                    ) => break,
+                };
+
+                process_chunk(
+                    &mut pending,
+                    &transcript,
+                    &call_id,
+                    &session_ref,
+                    &turn_ref,
+                    &mut emitted_deltas,
+                    chunk,
+                )
+                .await;
+            }
+        }
+        output_drained.notify_one();
     });
 }
 
@@ -114,13 +164,22 @@ pub(crate) fn spawn_exit_watcher(
     process_id: i32,
     transcript: Arc<Mutex<HeadTailBuffer>>,
     started_at: Instant,
+    network_denial_monitor: Option<tokio::task::JoinHandle<()>>,
 ) {
     let exit_token = process.cancellation_token();
     let output_drained = process.output_drained_notify();
+    let interaction_lock = process.interaction_lock();
 
     tokio::spawn(async move {
         exit_token.cancelled().await;
         output_drained.notified().await;
+        // Deferred network denial deliberately remains observable for a short
+        // window after process exit. Do not classify the terminal event until
+        // that monitor has settled, even when output closes immediately.
+        if let Some(network_denial_monitor) = network_denial_monitor {
+            let _ = network_denial_monitor.await;
+        }
+        let _interaction_guard = interaction_lock.lock_owned().await;
 
         let duration = Instant::now().saturating_duration_since(started_at);
         if let Some(message) = process.failure_message() {

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

@@ -1,6 +1,202 @@
+use std::sync::Arc;
+
+use super::TRAILING_OUTPUT_GRACE;
+use super::spawn_exit_watcher;
 use super::split_valid_utf8_prefix_with_max;
+use super::start_streaming_output;
+use crate::session::tests::make_session_and_context_with_rx;
+use crate::unified_exec::UnifiedExecContext;
+use crate::unified_exec::head_tail_buffer::HeadTailBuffer;
+use crate::unified_exec::process::NoopSpawnLifecycle;
+use crate::unified_exec::process::UnifiedExecProcess;
+use codex_protocol::items::CommandExecutionStatus;
+use codex_protocol::items::TurnItem;
+use codex_protocol::protocol::Event;
+use codex_protocol::protocol::EventMsg;
+use codex_sandboxing::SandboxType;
 
 use pretty_assertions::assert_eq;
+use tokio::time::Duration;
+use tokio::time::Instant;
+
+struct StreamingOutputHarness {
+    process: Arc<UnifiedExecProcess>,
+    stdout_tx: tokio::sync::broadcast::Sender<Vec<u8>>,
+    exit_tx: tokio::sync::oneshot::Sender<i32>,
+    transcript: Arc<tokio::sync::Mutex<HeadTailBuffer>>,
+    context: UnifiedExecContext,
+    rx_event: async_channel::Receiver<Event>,
+}
+
+async fn streaming_output_harness() -> anyhow::Result<StreamingOutputHarness> {
+    let (writer_tx, _writer_rx) = tokio::sync::mpsc::channel::<Vec<u8>>(1);
+    let (stdout_tx, stdout_rx) = tokio::sync::broadcast::channel::<Vec<u8>>(8);
+    let (exit_tx, exit_rx) = tokio::sync::oneshot::channel::<i32>();
+    let spawned = codex_utils_pty::spawn_from_driver(codex_utils_pty::ProcessDriver {
+        writer_tx,
+        stdout_rx,
+        stderr_rx: None,
+        exit_rx,
+        terminator: None,
+        writer_handle: None,
+        resizer: None,
+    });
+    let process = Arc::new(
+        UnifiedExecProcess::from_spawned(spawned, SandboxType::None, Box::new(NoopSpawnLifecycle))
+            .await?,
+    );
+    let (session, turn, rx_event) = make_session_and_context_with_rx().await;
+    let context = UnifiedExecContext::new(session, turn, "streaming-output-test".to_string());
+    let transcript = Arc::new(tokio::sync::Mutex::new(HeadTailBuffer::default()));
+    start_streaming_output(&process, &context, Arc::clone(&transcript));
+
+    Ok(StreamingOutputHarness {
+        process,
+        stdout_tx,
+        exit_tx,
+        transcript,
+        context,
+        rx_event,
+    })
+}
+
+#[tokio::test]
+async fn streaming_output_finishes_on_close_without_waiting_for_grace() -> anyhow::Result<()> {
+    let StreamingOutputHarness {
+        process,
+        stdout_tx,
+        exit_tx,
+        transcript,
+        ..
+    } = streaming_output_harness().await?;
+    let output_drained = process.output_drained_notify();
+    let drained = output_drained.notified();
+    tokio::pin!(drained);
+
+    tokio::time::pause();
+    let exited_at = Instant::now();
+    exit_tx.send(0).expect("send exit");
+    tokio::spawn(async move {
+        tokio::time::sleep(Duration::from_millis(50)).await;
+        stdout_tx
+            .send(b"LATE-OUTPUT-MARKER".to_vec())
+            .expect("send late output");
+    });
+
+    (&mut drained).await;
+    let elapsed = Instant::now().saturating_duration_since(exited_at);
+    tokio::time::resume();
+
+    assert!(
+        elapsed >= Duration::from_millis(50) && elapsed < TRAILING_OUTPUT_GRACE,
+        "output close should finish before the grace fallback: {elapsed:?}"
+    );
+    assert_eq!(
+        transcript.lock().await.to_bytes_with_omission_marker(),
+        b"LATE-OUTPUT-MARKER"
+    );
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn streaming_output_keeps_grace_as_fallback_without_close() -> anyhow::Result<()> {
+    let StreamingOutputHarness {
+        process,
+        stdout_tx: _stdout_tx,
+        exit_tx,
+        ..
+    } = streaming_output_harness().await?;
+    let output_drained = process.output_drained_notify();
+    let drained = output_drained.notified();
+    tokio::pin!(drained);
+
+    tokio::time::pause();
+    let exited_at = Instant::now();
+    exit_tx.send(0).expect("send exit");
+    (&mut drained).await;
+    let elapsed = Instant::now().saturating_duration_since(exited_at);
+    tokio::time::resume();
+
+    assert!(
+        elapsed >= TRAILING_OUTPUT_GRACE
+            && elapsed <= TRAILING_OUTPUT_GRACE + Duration::from_millis(10),
+        "missing output close should use the grace fallback: {elapsed:?}"
+    );
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn exit_watcher_waits_for_late_network_denial_before_classifying_end() -> anyhow::Result<()> {
+    let StreamingOutputHarness {
+        process,
+        stdout_tx,
+        exit_tx,
+        transcript,
+        context,
+        rx_event,
+    } = streaming_output_harness().await?;
+
+    tokio::time::pause();
+    let process_for_late_denial = Arc::clone(&process);
+    let (late_denial_armed_tx, late_denial_armed_rx) = tokio::sync::oneshot::channel();
+    let network_denial_monitor = tokio::spawn(async move {
+        let sleep = tokio::time::sleep(Duration::from_millis(10));
+        tokio::pin!(sleep);
+        late_denial_armed_tx.send(()).expect("arm late denial");
+        sleep.await;
+        process_for_late_denial.fail_and_terminate("LATE_DENIAL".to_string());
+    });
+    late_denial_armed_rx.await.expect("late denial armed");
+
+    #[allow(deprecated)]
+    let cwd = context.turn.cwd.clone().into();
+    spawn_exit_watcher(
+        Arc::clone(&process),
+        Arc::clone(&context.session),
+        Arc::clone(&context.turn),
+        context.call_id,
+        vec!["proof".to_string()],
+        cwd,
+        /*process_id*/ 123,
+        transcript,
+        Instant::now(),
+        Some(network_denial_monitor),
+    );
+
+    let exited_at = Instant::now();
+    exit_tx.send(0).expect("send exit");
+    drop(stdout_tx);
+
+    let event = rx_event.recv().await.expect("command end event");
+    let elapsed = Instant::now().saturating_duration_since(exited_at);
+    tokio::time::resume();
+    let EventMsg::ItemCompleted(completed) = event.msg else {
+        panic!("expected ItemCompleted");
+    };
+    let TurnItem::CommandExecution(item) = completed.item else {
+        panic!("expected CommandExecution");
+    };
+    assert_eq!(
+        (
+            item.status,
+            item.exit_code,
+            item.aggregated_output.as_deref()
+        ),
+        (
+            CommandExecutionStatus::Failed,
+            Some(-1),
+            Some("LATE_DENIAL")
+        )
+    );
+    assert!(
+        elapsed >= Duration::from_millis(10) && elapsed < TRAILING_OUTPUT_GRACE,
+        "completion should wait for denial without falling back to the output grace: {elapsed:?}"
+    );
+
+    Ok(())
+}
 
 #[test]
 fn split_valid_utf8_prefix_respects_max_bytes_for_ascii() {

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

@@ -116,6 +116,18 @@ pub(crate) struct WriteStdinRequest<'a> {
     pub yield_time_ms: u64,
     pub max_output_tokens: Option<usize>,
     pub truncation_policy: TruncationPolicy,
+    pub interaction_event: Option<WriteStdinInteractionEvent<'a>>,
+}
+
+pub(crate) struct WriteStdinInteractionEvent<'a> {
+    pub session: &'a Arc<Session>,
+    pub turn: &'a Arc<TurnContext>,
+}
+
+impl std::fmt::Debug for WriteStdinInteractionEvent<'_> {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter.write_str("WriteStdinInteractionEvent")
+    }
 }
 
 #[derive(Default)]