From aeaf7ee3adcaea0dcf88057c049208e2c091336a Mon Sep 17 00:00:00 2001
From: starr-openai <starr@openai.com>
Date: Fri, 8 May 2026 18:13:51 -0700
Subject: [PATCH] Drain events while polling remote tool output

Keep remote routing tests from stalling before the first model request by consuming Codex events while waiting for the captured tool output.

Co-authored-by: Codex <noreply@openai.com>
---
 codex-rs/core/tests/suite/remote_env.rs | 14 ++++++++++----
 codex-rs/core/tests/suite/view_image.rs | 14 ++++++++++----
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs
index 9b8ad68bbf..b380389919 100644
--- a/codex-rs/core/tests/suite/remote_env.rs
+++ b/codex-rs/core/tests/suite/remote_env.rs
@@ -41,19 +41,25 @@ use tempfile::TempDir;
 use tokio::time::Duration;
 
 async fn wait_for_function_call_output(
+    test: &TestCodex,
     response_mock: &core_test_support::responses::ResponseMock,
     call_id: &str,
 ) -> Result<String> {
     tokio::time::timeout(Duration::from_secs(60), async {
         loop {
             if let Some(output) = response_mock.function_call_output_text(call_id) {
-                return output;
+                return Ok(output);
+            }
+            tokio::select! {
+                event = test.codex.next_event() => {
+                    let _ = event.context("codex event stream ended while waiting for function_call_output")?;
+                }
+                _ = tokio::time::sleep(Duration::from_millis(50)) => {}
             }
-            tokio::time::sleep(Duration::from_millis(50)).await;
         }
     })
     .await
-    .with_context(|| format!("timed out waiting for function_call_output for {call_id}"))
+    .with_context(|| format!("timed out waiting for function_call_output for {call_id}"))?
 }
 async fn unified_exec_test(server: &wiremock::MockServer) -> Result<TestCodex> {
     let mut builder = test_codex().with_config(|config| {
@@ -194,7 +200,7 @@ async fn exec_command_routing_output(
     test.submit_turn_with_environments_no_wait("route exec command", environments)
         .await?;
 
-    wait_for_function_call_output(&response_mock, call_id).await
+    wait_for_function_call_output(test, &response_mock, call_id).await
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
diff --git a/codex-rs/core/tests/suite/view_image.rs b/codex-rs/core/tests/suite/view_image.rs
index 222f8fab76..e5361521ed 100644
--- a/codex-rs/core/tests/suite/view_image.rs
+++ b/codex-rs/core/tests/suite/view_image.rs
@@ -65,19 +65,25 @@ use wiremock::matchers::body_string_contains;
 const VIEW_IMAGE_TURN_COMPLETE_TIMEOUT: Duration = Duration::from_secs(30);
 
 async fn wait_for_function_call_output(
+    test: &TestCodex,
     response_mock: &responses::ResponseMock,
     call_id: &str,
 ) -> anyhow::Result<Value> {
     tokio::time::timeout(Duration::from_secs(60), async {
         loop {
             if let Some(output) = response_mock.function_call_output(call_id) {
-                return output;
+                return Ok(output);
+            }
+            tokio::select! {
+                event = test.codex.next_event() => {
+                    let _ = event.context("codex event stream ended while waiting for function_call_output")?;
+                }
+                _ = tokio::time::sleep(Duration::from_millis(50)) => {}
             }
-            tokio::time::sleep(Duration::from_millis(50)).await;
         }
     })
     .await
-    .with_context(|| format!("timed out waiting for function_call_output for {call_id}"))
+    .with_context(|| format!("timed out waiting for function_call_output for {call_id}"))?
 }
 
 fn disabled_user_turn(test: &TestCodex, items: Vec<UserInput>, model: String) -> Op {
@@ -559,7 +565,7 @@ async fn view_image_routes_to_selected_remote_environment() -> anyhow::Result<()
     )
     .await?;
 
-    let output = wait_for_function_call_output(&response_mock, call_id).await?;
+    let output = wait_for_function_call_output(&test, &response_mock, call_id).await?;
     let output_items = output
         .get("output")
         .and_then(Value::as_array)