From fc70c09055f3b2e98430e4d5e7d7b3e243eda128 Mon Sep 17 00:00:00 2001
From: starr-openai <starr@openai.com>
Date: Fri, 8 May 2026 17:31:22 -0700
Subject: [PATCH] Fix view image output polling helper

Add an optional captured-response helper for function call output values so remote-routing tests can poll without panicking on earlier requests.

Co-authored-by: Codex <noreply@openai.com>
---
 codex-rs/core/tests/common/responses.rs | 21 +++++++++++++++++++--
 codex-rs/core/tests/suite/view_image.rs |  6 +-----
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/codex-rs/core/tests/common/responses.rs b/codex-rs/core/tests/common/responses.rs
index eab90c5aba..e4b884841b 100644
--- a/codex-rs/core/tests/common/responses.rs
+++ b/codex-rs/core/tests/common/responses.rs
@@ -80,6 +80,14 @@ impl ResponseMock {
             .iter()
             .find_map(|req| req.function_call_output_text(call_id))
     }
+
+    /// Returns the full matching `function_call_output` item for the provided
+    /// `call_id`, searching across all captured requests.
+    pub fn function_call_output(&self, call_id: &str) -> Option<Value> {
+        self.requests()
+            .iter()
+            .find_map(|req| req.maybe_function_call_output(call_id))
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -209,6 +217,10 @@ impl ResponsesRequest {
         self.call_output(call_id, "function_call_output")
     }
 
+    pub fn maybe_function_call_output(&self, call_id: &str) -> Option<Value> {
+        self.maybe_call_output(call_id, "function_call_output")
+    }
+
     pub fn custom_tool_call_output(&self, call_id: &str) -> Value {
         self.call_output(call_id, "custom_tool_call_output")
     }
@@ -218,13 +230,18 @@ impl ResponsesRequest {
     }
 
     pub fn call_output(&self, call_id: &str, call_type: &str) -> Value {
+        self.maybe_call_output(call_id, call_type)
+            .unwrap_or_else(|| panic!("function call output {call_id} item not found in request"))
+    }
+
+    pub fn maybe_call_output(&self, call_id: &str, call_type: &str) -> Option<Value> {
         self.input()
             .iter()
             .find(|item| {
-                item.get("type").unwrap() == call_type && item.get("call_id").unwrap() == call_id
+                item.get("type").and_then(Value::as_str) == Some(call_type)
+                    && item.get("call_id").and_then(Value::as_str) == Some(call_id)
             })
             .cloned()
-            .unwrap_or_else(|| panic!("function call output {call_id} item not found in request"))
     }
 
     /// Returns true if this request's `input` contains a `function_call` with
diff --git a/codex-rs/core/tests/suite/view_image.rs b/codex-rs/core/tests/suite/view_image.rs
index b3377f7d26..07e383d1a3 100644
--- a/codex-rs/core/tests/suite/view_image.rs
+++ b/codex-rs/core/tests/suite/view_image.rs
@@ -70,11 +70,7 @@ async fn wait_for_function_call_output(
 ) -> anyhow::Result<Value> {
     tokio::time::timeout(Duration::from_secs(25), async {
         loop {
-            if let Some(output) = response_mock
-                .requests()
-                .iter()
-                .find_map(|request| request.function_call_output(call_id))
-            {
+            if let Some(output) = response_mock.function_call_output(call_id) {
                 return output;
             }
             tokio::time::sleep(Duration::from_millis(50)).await;