From ce8df5c5204b17afd915b2f98defca1a1f0f1ec9 Mon Sep 17 00:00:00 2001
From: YOUR NAME <mle@openai.com>
Date: Thu, 26 Mar 2026 09:23:37 -0700
Subject: [PATCH] Stabilize multi-agent assign_task test

---
 .../src/tools/handlers/multi_agents_tests.rs  | 49 +++++++++++++++----
 1 file changed, 39 insertions(+), 10 deletions(-)

diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs
index ab40b20587..85336785ce 100644
--- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs
+++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs
@@ -109,6 +109,18 @@ fn history_contains_inter_agent_communication(
     })
 }
 
+fn response_input_items_contain_inter_agent_communication(
+    input_items: &[ResponseInputItem],
+    expected: &InterAgentCommunication,
+) -> bool {
+    let response_items: Vec<ResponseItem> = input_items
+        .iter()
+        .cloned()
+        .map(ResponseItem::from)
+        .collect();
+    history_contains_inter_agent_communication(&response_items, expected)
+}
+
 #[derive(Clone, Copy)]
 struct NeverEndingTask;
 
@@ -922,6 +934,13 @@ async fn multi_agent_v2_assign_task_interrupts_busy_child_without_losing_message
         .iter()
         .filter_map(|(id, op)| (*id == agent_id).then_some(op))
         .collect();
+    let expected_communication = InterAgentCommunication::new(
+        AgentPath::root(),
+        AgentPath::try_from("/root/worker").expect("agent path"),
+        Vec::new(),
+        "continue".to_string(),
+        true,
+    );
     assert!(ops_for_agent.iter().any(|op| matches!(op, Op::Interrupt)));
     assert!(ops_for_agent.iter().any(|op| {
         matches!(
@@ -931,9 +950,13 @@ async fn multi_agent_v2_assign_task_interrupts_busy_child_without_losing_message
                     && communication.recipient.as_str() == "/root/worker"
                     && communication.other_recipients.is_empty()
                     && communication.content == "continue"
+                    && communication.trigger_turn
         )
     }));
 
+    // `assign_task` wakes the interrupted child immediately, so the redirected
+    // envelope may still be buffered in turn state before it is materialized
+    // into history on slower runners.
     timeout(Duration::from_secs(5), async {
         loop {
             let history_items = thread
@@ -943,16 +966,22 @@ async fn multi_agent_v2_assign_task_interrupts_busy_child_without_losing_message
                 .await
                 .raw_items()
                 .to_vec();
-            let saw_envelope = history_contains_inter_agent_communication(
-                &history_items,
-                &InterAgentCommunication::new(
-                    AgentPath::root(),
-                    AgentPath::try_from("/root/worker").expect("agent path"),
-                    Vec::new(),
-                    "continue".to_string(),
-                    true,
-                ),
-            );
+            let pending_input_items = thread.codex.session.pending_input_snapshot().await;
+            let queued_input_items = thread
+                .codex
+                .session
+                .queued_response_items_for_next_turn_snapshot()
+                .await;
+            let saw_envelope =
+                history_contains_inter_agent_communication(&history_items, &expected_communication)
+                    || response_input_items_contain_inter_agent_communication(
+                        &pending_input_items,
+                        &expected_communication,
+                    )
+                    || response_input_items_contain_inter_agent_communication(
+                        &queued_input_items,
+                        &expected_communication,
+                    );
             let saw_user_message = history_items.iter().any(|item| {
                 matches!(
                     item,