From fed333a68b01ae2df265cc7d039c7bc58f28622f Mon Sep 17 00:00:00 2001
From: starr-openai <starr@openai.com>
Date: Mon, 18 May 2026 10:12:33 -0700
Subject: [PATCH] Extend rollback event waits in flaky core tests

---
 .../core/tests/suite/compact_resume_fork.rs     | 17 ++++++++++++-----
 codex-rs/core/tests/suite/model_switching.rs    |  9 ++++++---
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/codex-rs/core/tests/suite/compact_resume_fork.rs b/codex-rs/core/tests/suite/compact_resume_fork.rs
index b2b02f5e66..89087bcd87 100644
--- a/codex-rs/core/tests/suite/compact_resume_fork.rs
+++ b/codex-rs/core/tests/suite/compact_resume_fork.rs
@@ -41,6 +41,7 @@ use serde_json::Value;
 use serde_json::json;
 use std::sync::Arc;
 use tempfile::TempDir;
+use tokio::time::Duration;
 use wiremock::MockServer;
 
 const AFTER_SECOND_RESUME: &str = "AFTER_SECOND_RESUME";
@@ -456,8 +457,12 @@ async fn snapshot_rollback_past_compaction_replays_append_only_history() -> Resu
     base.submit(Op::ThreadRollback { num_turns: 1 })
         .await
         .expect("submit thread rollback");
-    let rollback_event =
-        wait_for_event(&base, |ev| matches!(ev, EventMsg::ThreadRolledBack(_))).await;
+    let rollback_event = core_test_support::wait_for_event_with_timeout(
+        &base,
+        |ev| matches!(ev, EventMsg::ThreadRolledBack(_)),
+        Duration::from_secs(20),
+    )
+    .await;
     let EventMsg::ThreadRolledBack(rollback_event) = rollback_event else {
         panic!("expected thread rolled back event");
     };
@@ -570,9 +575,11 @@ async fn snapshot_rollback_followup_turn_trims_context_updates() -> Result<()> {
     conversation
         .submit(Op::ThreadRollback { num_turns: 1 })
         .await?;
-    let rollback_event = wait_for_event(&conversation, |ev| {
-        matches!(ev, EventMsg::ThreadRolledBack(_))
-    })
+    let rollback_event = core_test_support::wait_for_event_with_timeout(
+        &conversation,
+        |ev| matches!(ev, EventMsg::ThreadRolledBack(_)),
+        Duration::from_secs(20),
+    )
     .await;
     let EventMsg::ThreadRolledBack(rollback_event) = rollback_event else {
         panic!("expected thread rolled back event");
diff --git a/codex-rs/core/tests/suite/model_switching.rs b/codex-rs/core/tests/suite/model_switching.rs
index 009dfda60f..7eb78751d2 100644
--- a/codex-rs/core/tests/suite/model_switching.rs
+++ b/codex-rs/core/tests/suite/model_switching.rs
@@ -37,6 +37,7 @@ use core_test_support::wait_for_event;
 use pretty_assertions::assert_eq;
 use std::path::Path;
 use std::path::PathBuf;
+use tokio::time::Duration;
 use wiremock::MockServer;
 
 fn read_only_user_turn(test: &TestCodex, items: Vec<UserInput>, model: String) -> Op {
@@ -785,9 +786,11 @@ async fn thread_rollback_after_generated_image_drops_entire_image_turn_history()
     test.codex
         .submit(Op::ThreadRollback { num_turns: 1 })
         .await?;
-    wait_for_event(&test.codex, |ev| {
-        matches!(ev, EventMsg::ThreadRolledBack(_))
-    })
+    core_test_support::wait_for_event_with_timeout(
+        &test.codex,
+        |ev| matches!(ev, EventMsg::ThreadRolledBack(_)),
+        Duration::from_secs(20),
+    )
     .await;
 
     test.codex