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
This commit is contained in:
jif
2026-09-18 23:47:19 +00:00
committed by copyberry
parent 8933a816eb
commit d374a93b41
+32 -10
View File
@@ -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",