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::(&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",