From d374a93b41426cfe7fef1d069d6f327e7195e814 Mon Sep 17 00:00:00 2001
From: jif <jif@openai.com>
Date: Fri, 18 Sep 2026 12:46:20 +0000
Subject: [PATCH] Handle completion timing in the multi-agent resume test
 (#46516)

## Why

A grandchild's completion can arrive while the worker is completing its response, triggering an extra sampling request to drain the message. The cold root resume test previously allowed only one worker completion request.

## What changed

Match worker completion requests by the nested spawn's `function_call_output` and allow one or two requests. Scope the mock to the initial worker's lifetime and drop it before the follow-up phase.

GitOrigin-RevId: 10e0a429e6e1f51b5f04a86e3351c5fa4df9cb66
---
 .../core/tests/suite/multi_agent_resume.rs    | 42 ++++++++++++++-----
 1 file changed, 32 insertions(+), 10 deletions(-)

diff --git a/codex-rs/core/tests/suite/multi_agent_resume.rs b/codex-rs/core/tests/suite/multi_agent_resume.rs
index 47bf151006..4e1ab9a56f 100644
--- a/codex-rs/core/tests/suite/multi_agent_resume.rs
+++ b/codex-rs/core/tests/suite/multi_agent_resume.rs
@@ -17,6 +17,7 @@ use core_test_support::responses::ev_response_created;
 use core_test_support::responses::mount_sse_once;
 use core_test_support::responses::mount_sse_once_match;
 use core_test_support::responses::sse;
+use core_test_support::responses::sse_response;
 use core_test_support::responses::start_mock_server;
 use core_test_support::test_codex::test_codex;
 use core_test_support::wait_for_event;
@@ -217,17 +218,37 @@ async fn cold_root_resume_restores_agent_identity_and_role_on_followup() -> Resu
         sse(vec![ev_completed("resp-parent-turn-assistant")]),
     )
     .await;
-    for (text, is_subagent) in [(NESTED_CALL_ID, true), (QUEUE_CALL_ID, false)] {
-        mount_sse_once_match(
-            &server,
-            move |request: &wiremock::Request| {
-                body_contains(request, text)
-                    && request_has_input_type(request, "agent_message") == is_subagent
-            },
-            sse(vec![ev_completed("resp-parent-turn-assistant")]),
-        )
+    // The grandchild's completion can arrive during the worker's completion response,
+    // causing one more sampling request to drain that message before the turn ends.
+    let worker_completion = wiremock::Mock::given(wiremock::matchers::method("POST"))
+        .and(wiremock::matchers::path("/v1/responses"))
+        .and(|request: &wiremock::Request| {
+            decoded_body(request)
+                .and_then(|body| serde_json::from_slice::<Value>(&body).ok())
+                .is_some_and(|body| {
+                    body["input"].as_array().is_some_and(|items| {
+                        items.iter().any(|item| {
+                            item["type"] == "function_call_output"
+                                && item["call_id"] == NESTED_CALL_ID
+                        })
+                    })
+                })
+        })
+        .respond_with(sse_response(sse(vec![ev_completed(
+            "resp-worker-complete",
+        )])))
+        .expect(1..=2)
+        .mount_as_scoped(&server)
         .await;
-    }
+    mount_sse_once_match(
+        &server,
+        |request: &wiremock::Request| {
+            body_contains(request, QUEUE_CALL_ID)
+                && !request_has_input_type(request, "agent_message")
+        },
+        sse(vec![ev_completed("resp-parent-turn-assistant")]),
+    )
+    .await;
     mount_sse_once_match(
         &server,
         |request: &wiremock::Request| {
@@ -371,6 +392,7 @@ async fn cold_root_resume_restores_agent_identity_and_role_on_followup() -> Resu
     worker_thread.shutdown_and_wait().await?;
     drop(sibling_thread);
     drop(worker_thread);
+    drop(worker_completion);
 
     let followup_args = serde_json::to_string(&json!({
         "target": "worker",