Flaky CI fix (#1647)

Flushing before sending `TaskCompleteEvent` and ending the submission
loop to avoid race conditions.

codex-rs/core/src/codex.rs

@@ -812,6 +812,37 @@ async fn submission_loop(
                     }
                 });
             }
+            Op::Shutdown => {
+                info!("Shutting down Codex instance");
+
+                // Gracefully flush and shutdown rollout recorder on session end so tests
+                // that inspect the rollout file do not race with the background writer.
+                if let Some(sess_arc) = sess {
+                    let recorder_opt = sess_arc.rollout.lock().unwrap().take();
+                    if let Some(rec) = recorder_opt {
+                        if let Err(e) = rec.shutdown().await {
+                            warn!("failed to shutdown rollout recorder: {e}");
+                            let event = Event {
+                                id: sub.id.clone(),
+                                msg: EventMsg::Error(ErrorEvent {
+                                    message: "Failed to shutdown rollout recorder".to_string(),
+                                }),
+                            };
+                            if let Err(e) = tx_event.send(event).await {
+                                warn!("failed to send error message: {e:?}");
+                            }
+                        }
+                    }
+                }
+                let event = Event {
+                    id: sub.id.clone(),
+                    msg: EventMsg::ShutdownComplete,
+                };
+                if let Err(e) = tx_event.send(event).await {
+                    warn!("failed to send Shutdown event: {e}");
+                }
+                break;
+            }
         }
     }
     debug!("Agent loop exited");

codex-rs/core/src/protocol.rs

@@ -116,6 +116,9 @@ pub enum Op {
 
     /// Request a single history entry identified by `log_id` + `offset`.
     GetHistoryEntryRequest { offset: usize, log_id: u64 },
+
+    /// Request to shut down codex instance.
+    Shutdown,
 }
 
 /// Determines the conditions under which the user is consulted to approve
@@ -326,6 +329,9 @@ pub enum EventMsg {
 
     /// Response to GetHistoryEntryRequest.
     GetHistoryEntryResponse(GetHistoryEntryResponseEvent),
+
+    /// Notification that the agent is shutting down.
+    ShutdownComplete,
 }
 
 // Individual event payload types matching each `EventMsg` variant.

codex-rs/core/src/rollout.rs

@@ -14,6 +14,7 @@ use time::macros::format_description;
 use tokio::io::AsyncWriteExt;
 use tokio::sync::mpsc::Sender;
 use tokio::sync::mpsc::{self};
+use tokio::sync::oneshot;
 use tracing::info;
 use tracing::warn;
 use uuid::Uuid;
@@ -57,10 +58,10 @@ pub(crate) struct RolloutRecorder {
     tx: Sender<RolloutCmd>,
 }
 
-#[derive(Clone)]
 enum RolloutCmd {
     AddItems(Vec<ResponseItem>),
     UpdateState(SessionStateSnapshot),
+    Shutdown { ack: oneshot::Sender<()> },
 }
 
 impl RolloutRecorder {
@@ -204,6 +205,21 @@ impl RolloutRecorder {
         info!("Resumed rollout successfully from {path:?}");
         Ok((Self { tx }, saved))
     }
+
+    pub async fn shutdown(&self) -> std::io::Result<()> {
+        let (tx_done, rx_done) = oneshot::channel();
+        match self.tx.send(RolloutCmd::Shutdown { ack: tx_done }).await {
+            Ok(_) => rx_done
+                .await
+                .map_err(|e| IoError::other(format!("failed waiting for rollout shutdown: {e}"))),
+            Err(e) => {
+                warn!("failed to send rollout shutdown command: {e}");
+                Err(IoError::other(format!(
+                    "failed to send rollout shutdown command: {e}"
+                )))
+            }
+        }
+    }
 }
 
 struct LogFileInfo {
@@ -299,6 +315,9 @@ async fn rollout_writer(
                     let _ = file.flush().await;
                 }
             }
+            RolloutCmd::Shutdown { ack } => {
+                let _ = ack.send(());
+            }
         }
     }
 }