fix: flaky test 2 (#7818)

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

@@ -1,9 +1,11 @@
 use std::path::PathBuf;
+use std::pin::Pin;
 use std::sync::Arc;
 
 use tokio::sync::Mutex;
 use tokio::time::Duration;
 use tokio::time::Instant;
+use tokio::time::Sleep;
 
 use crate::codex::Session;
 use crate::codex::TurnContext;
@@ -21,6 +23,8 @@ use super::CommandTranscript;
 use super::UnifiedExecContext;
 use super::session::UnifiedExecSession;
 
+pub(crate) const TRAILING_OUTPUT_GRACE: Duration = Duration::from_millis(100);
+
 /// Spawn a background task that continuously reads from the PTY, appends to the
 /// shared transcript, and emits ExecCommandOutputDelta events on UTF‑8
 /// boundaries.
@@ -30,39 +34,58 @@ pub(crate) fn start_streaming_output(
     transcript: Arc<Mutex<CommandTranscript>>,
 ) {
     let mut receiver = session.output_receiver();
+    let output_drained = session.output_drained_notify();
+    let exit_token = session.cancellation_token();
+
     let session_ref = Arc::clone(&context.session);
     let turn_ref = Arc::clone(&context.turn);
     let call_id = context.call_id.clone();
-    let cancellation_token = session.cancellation_token();
 
     tokio::spawn(async move {
-        let mut pending: Vec<u8> = Vec::new();
+        use tokio::sync::broadcast::error::RecvError;
+
+        let mut pending = Vec::<u8>::new();
+
+        let mut grace_sleep: Option<Pin<Box<Sleep>>> = None;
+
         loop {
             tokio::select! {
-                _ = cancellation_token.cancelled() => break,
-                result = receiver.recv() => match result {
-                    Ok(chunk) => {
-                        pending.extend_from_slice(&chunk);
-                        while let Some(prefix) = split_valid_utf8_prefix(&mut pending) {
-                            {
-                                let mut guard = transcript.lock().await;
-                                guard.append(&prefix);
-                            }
-
-                            let event = ExecCommandOutputDeltaEvent {
-                                call_id: call_id.clone(),
-                                stream: ExecOutputStream::Stdout,
-                                chunk: prefix,
-                            };
-                            session_ref
-                                .send_event(turn_ref.as_ref(), EventMsg::ExecCommandOutputDelta(event))
-                                .await;
-                        }
-                    }
-                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
-                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
+                _ = exit_token.cancelled(), if grace_sleep.is_none() => {
+                    let deadline = Instant::now() + TRAILING_OUTPUT_GRACE;
+                    grace_sleep.replace(Box::pin(tokio::time::sleep_until(deadline)));
                 }
-            };
+
+                _ = async {
+                    if let Some(sleep) = grace_sleep.as_mut() {
+                        sleep.as_mut().await;
+                    }
+                }, if grace_sleep.is_some() => {
+                    output_drained.notify_one();
+                    break;
+                }
+
+                received = receiver.recv() => {
+                    let chunk = match received {
+                        Ok(chunk) => chunk,
+                        Err(RecvError::Lagged(_)) => {
+                            continue;
+                        },
+                        Err(RecvError::Closed) => {
+                            output_drained.notify_one();
+                            break;
+                        }
+                    };
+
+                    process_chunk(
+                        &mut pending,
+                        &transcript,
+                        &call_id,
+                        &session_ref,
+                        &turn_ref,
+                        chunk,
+                    ).await;
+                }
+            }
         }
     });
 }
@@ -82,9 +105,11 @@ pub(crate) fn spawn_exit_watcher(
     started_at: Instant,
 ) {
     let exit_token = session.cancellation_token();
+    let output_drained = session.output_drained_notify();
 
     tokio::spawn(async move {
         exit_token.cancelled().await;
+        output_drained.notified().await;
 
         let exit_code = session.exit_code().unwrap_or(-1);
         let duration = Instant::now().saturating_duration_since(started_at);
@@ -104,6 +129,32 @@ pub(crate) fn spawn_exit_watcher(
     });
 }
 
+async fn process_chunk(
+    pending: &mut Vec<u8>,
+    transcript: &Arc<Mutex<CommandTranscript>>,
+    call_id: &str,
+    session_ref: &Arc<Session>,
+    turn_ref: &Arc<TurnContext>,
+    chunk: Vec<u8>,
+) {
+    pending.extend_from_slice(&chunk);
+    while let Some(prefix) = split_valid_utf8_prefix(pending) {
+        {
+            let mut guard = transcript.lock().await;
+            guard.append(&prefix);
+        }
+
+        let event = ExecCommandOutputDeltaEvent {
+            call_id: call_id.to_string(),
+            stream: ExecOutputStream::Stdout,
+            chunk: prefix,
+        };
+        session_ref
+            .send_event(turn_ref.as_ref(), EventMsg::ExecCommandOutputDelta(event))
+            .await;
+    }
+}
+
 /// Emit an ExecCommandEnd event for a unified exec session, using the transcript
 /// as the primary source of aggregated_output and falling back to the provided
 /// text when the transcript is empty.

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

@@ -79,6 +79,7 @@ pub(crate) struct UnifiedExecSession {
     output_buffer: OutputBuffer,
     output_notify: Arc<Notify>,
     cancellation_token: CancellationToken,
+    output_drained: Arc<Notify>,
     output_task: JoinHandle<()>,
     sandbox_type: SandboxType,
 }
@@ -92,27 +93,21 @@ impl UnifiedExecSession {
         let output_buffer = Arc::new(Mutex::new(OutputBufferState::default()));
         let output_notify = Arc::new(Notify::new());
         let cancellation_token = CancellationToken::new();
+        let output_drained = Arc::new(Notify::new());
         let mut receiver = initial_output_rx;
         let buffer_clone = Arc::clone(&output_buffer);
         let notify_clone = Arc::clone(&output_notify);
-        let cancellation_token_clone = cancellation_token.clone();
         let output_task = tokio::spawn(async move {
             loop {
-                tokio::select! {
-                    _ = cancellation_token_clone.cancelled() => break,
-                    result = receiver.recv() => match result {
-                        Ok(chunk) => {
-                            let mut guard = buffer_clone.lock().await;
-                            guard.push_chunk(chunk);
-                            drop(guard);
-                            notify_clone.notify_waiters();
-                        }
-                        Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
-                        Err(tokio::sync::broadcast::error::RecvError::Closed) => {
-                            cancellation_token_clone.cancel();
-                            break;
-                        }
+                match receiver.recv().await {
+                    Ok(chunk) => {
+                        let mut guard = buffer_clone.lock().await;
+                        guard.push_chunk(chunk);
+                        drop(guard);
+                        notify_clone.notify_waiters();
                     }
+                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
+                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
                 };
             }
         });
@@ -122,6 +117,7 @@ impl UnifiedExecSession {
             output_buffer,
             output_notify,
             cancellation_token,
+            output_drained,
             output_task,
             sandbox_type,
         }
@@ -147,6 +143,10 @@ impl UnifiedExecSession {
         self.cancellation_token.clone()
     }
 
+    pub(super) fn output_drained_notify(&self) -> Arc<Notify> {
+        Arc::clone(&self.output_drained)
+    }
+
     pub(super) fn has_exited(&self) -> bool {
         self.session.has_exited()
     }