feat: add mailbox concept for wait (#16010)

Add a mailbox we can use for inter-agent communication
`wait` is now based on it and don't take target anymore

codex-rs/core/src/agent/agent_resolver.rs

@@ -28,25 +28,6 @@ pub(crate) async fn resolve_agent_target(
         })
 }
 
-/// Resolves multiple tool-facing agent targets to thread ids.
-pub(crate) async fn resolve_agent_targets(
-    session: &Arc<Session>,
-    turn: &Arc<TurnContext>,
-    targets: Vec<String>,
-) -> Result<Vec<ThreadId>, FunctionCallError> {
-    if targets.is_empty() {
-        return Err(FunctionCallError::RespondToModel(
-            "agent targets must be non-empty".to_string(),
-        ));
-    }
-
-    let mut resolved = Vec::with_capacity(targets.len());
-    for target in &targets {
-        resolved.push(resolve_agent_target(session, turn, target).await?);
-    }
-    Ok(resolved)
-}
-
 fn register_session_root(session: &Arc<Session>, turn: &Arc<TurnContext>) {
     session
         .services

codex-rs/core/src/agent/control_tests.rs

@@ -448,12 +448,7 @@ async fn send_inter_agent_communication_without_turn_queues_message_without_trig
 
     timeout(Duration::from_secs(5), async {
         loop {
-            if thread
-                .codex
-                .session
-                .has_queued_response_items_for_next_turn()
-                .await
-            {
+            if thread.codex.session.has_pending_input().await {
                 break;
             }
             sleep(Duration::from_millis(10)).await;

codex-rs/core/src/agent/mailbox.rs

@@ -0,0 +1,161 @@
+use codex_protocol::protocol::InterAgentCommunication;
+use std::collections::VecDeque;
+use std::sync::atomic::AtomicU64;
+use std::sync::atomic::Ordering;
+use tokio::sync::mpsc;
+use tokio::sync::watch;
+
+#[cfg(test)]
+use codex_protocol::AgentPath;
+
+pub(crate) struct Mailbox {
+    tx: mpsc::UnboundedSender<InterAgentCommunication>,
+    next_seq: AtomicU64,
+    seq_tx: watch::Sender<u64>,
+}
+
+pub(crate) struct MailboxReceiver {
+    rx: mpsc::UnboundedReceiver<InterAgentCommunication>,
+    pending_mails: VecDeque<InterAgentCommunication>,
+}
+
+impl Mailbox {
+    pub(crate) fn new() -> (Self, MailboxReceiver) {
+        let (tx, rx) = mpsc::unbounded_channel();
+        let (seq_tx, _) = watch::channel(0);
+        (
+            Self {
+                tx,
+                next_seq: AtomicU64::new(0),
+                seq_tx,
+            },
+            MailboxReceiver {
+                rx,
+                pending_mails: VecDeque::new(),
+            },
+        )
+    }
+
+    pub(crate) fn subscribe(&self) -> watch::Receiver<u64> {
+        self.seq_tx.subscribe()
+    }
+
+    pub(crate) fn send(&self, communication: InterAgentCommunication) -> u64 {
+        let seq = self.next_seq.fetch_add(1, Ordering::Relaxed) + 1;
+        let _ = self.tx.send(communication);
+        self.seq_tx.send_replace(seq);
+        seq
+    }
+}
+
+impl MailboxReceiver {
+    fn sync_pending_mails(&mut self) {
+        while let Ok(mail) = self.rx.try_recv() {
+            self.pending_mails.push_back(mail);
+        }
+    }
+
+    pub(crate) fn has_pending(&mut self) -> bool {
+        self.sync_pending_mails();
+        !self.pending_mails.is_empty()
+    }
+
+    pub(crate) fn has_pending_trigger_turn(&mut self) -> bool {
+        self.sync_pending_mails();
+        self.pending_mails.iter().any(|mail| mail.trigger_turn)
+    }
+
+    pub(crate) fn drain(&mut self) -> Vec<InterAgentCommunication> {
+        self.sync_pending_mails();
+        self.pending_mails.drain(..).collect()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use pretty_assertions::assert_eq;
+
+    fn make_mail(
+        author: AgentPath,
+        recipient: AgentPath,
+        content: &str,
+        trigger_turn: bool,
+    ) -> InterAgentCommunication {
+        InterAgentCommunication::new(
+            author,
+            recipient,
+            Vec::new(),
+            content.to_string(),
+            trigger_turn,
+        )
+    }
+
+    #[tokio::test]
+    async fn mailbox_assigns_monotonic_sequence_numbers() {
+        let (mailbox, _receiver) = Mailbox::new();
+        let mut seq_rx = mailbox.subscribe();
+
+        let seq_a = mailbox.send(make_mail(
+            AgentPath::root(),
+            AgentPath::try_from("/root/worker").expect("agent path"),
+            "one",
+            /*trigger_turn*/ false,
+        ));
+        let seq_b = mailbox.send(make_mail(
+            AgentPath::root(),
+            AgentPath::try_from("/root/worker").expect("agent path"),
+            "two",
+            /*trigger_turn*/ false,
+        ));
+
+        seq_rx.changed().await.expect("first seq update");
+        assert_eq!(*seq_rx.borrow(), seq_b);
+        assert_eq!(seq_a, 1);
+        assert_eq!(seq_b, 2);
+    }
+
+    #[tokio::test]
+    async fn mailbox_drains_in_delivery_order() {
+        let (mailbox, mut receiver) = Mailbox::new();
+        let mail_one = make_mail(
+            AgentPath::root(),
+            AgentPath::try_from("/root/worker").expect("agent path"),
+            "one",
+            /*trigger_turn*/ false,
+        );
+        let mail_two = make_mail(
+            AgentPath::try_from("/root/worker").expect("agent path"),
+            AgentPath::root(),
+            "two",
+            /*trigger_turn*/ false,
+        );
+
+        mailbox.send(mail_one.clone());
+        mailbox.send(mail_two.clone());
+
+        assert_eq!(receiver.drain(), vec![mail_one, mail_two]);
+        assert!(!receiver.has_pending());
+    }
+
+    #[tokio::test]
+    async fn mailbox_tracks_pending_trigger_turn_mail() {
+        let (mailbox, mut receiver) = Mailbox::new();
+
+        mailbox.send(make_mail(
+            AgentPath::root(),
+            AgentPath::try_from("/root/worker").expect("agent path"),
+            "queued",
+            /*trigger_turn*/ false,
+        ));
+        assert!(!receiver.has_pending_trigger_turn());
+
+        mailbox.send(make_mail(
+            AgentPath::root(),
+            AgentPath::try_from("/root/worker").expect("agent path"),
+            "wake",
+            /*trigger_turn*/ true,
+        ));
+        assert!(receiver.has_pending_trigger_turn());
+    }
+}

codex-rs/core/src/agent/mod.rs

@@ -1,11 +1,14 @@
 pub(crate) mod agent_resolver;
 pub(crate) mod control;
+pub(crate) mod mailbox;
 mod registry;
 pub(crate) mod role;
 pub(crate) mod status;
 
 pub(crate) use codex_protocol::protocol::AgentStatus;
 pub(crate) use control::AgentControl;
+pub(crate) use mailbox::Mailbox;
+pub(crate) use mailbox::MailboxReceiver;
 pub(crate) use registry::exceeds_thread_spawn_depth_limit;
 pub(crate) use registry::next_thread_spawn_depth;
 pub(crate) use status::agent_status_from_event;

codex-rs/core/src/codex.rs

@@ -11,6 +11,8 @@ use crate::CodexAuth;
 use crate::SandboxState;
 use crate::agent::AgentControl;
 use crate::agent::AgentStatus;
+use crate::agent::Mailbox;
+use crate::agent::MailboxReceiver;
 use crate::agent::agent_status_from_event;
 use crate::apps::render_apps_section;
 use crate::auth_env_telemetry::collect_auth_env_telemetry;
@@ -97,6 +99,7 @@ use codex_protocol::permissions::FileSystemSandboxPolicy;
 use codex_protocol::permissions::NetworkSandboxPolicy;
 use codex_protocol::protocol::FileChange;
 use codex_protocol::protocol::HasLegacyEvent;
+use codex_protocol::protocol::InterAgentCommunication;
 use codex_protocol::protocol::ItemCompletedEvent;
 use codex_protocol::protocol::ItemStartedEvent;
 use codex_protocol::protocol::RawResponseItemEvent;
@@ -806,7 +809,9 @@ pub(crate) struct Session {
     pending_mcp_server_refresh_config: Mutex<Option<McpServerRefreshConfig>>,
     pub(crate) conversation: Arc<RealtimeConversationManager>,
     pub(crate) active_turn: Mutex<Option<ActiveTurn>>,
-    idle_pending_input: Mutex<Vec<ResponseInputItem>>,
+    mailbox: Mailbox,
+    mailbox_rx: Mutex<MailboxReceiver>,
+    idle_pending_input: Mutex<Vec<ResponseInputItem>>, // TODO (jif) merge with mailbox!
     pub(crate) guardian_review_session: GuardianReviewSessionManager,
     pub(crate) services: SessionServices,
     js_repl: Arc<JsReplHandle>,
@@ -1907,6 +1912,7 @@ impl Session {
         let (out_of_band_elicitation_paused, _out_of_band_elicitation_paused_rx) =
             watch::channel(false);
 
+        let (mailbox, mailbox_rx) = Mailbox::new();
         let sess = Arc::new(Session {
             conversation_id,
             tx_event: tx_event.clone(),
@@ -1917,6 +1923,8 @@ impl Session {
             pending_mcp_server_refresh_config: Mutex::new(None),
             conversation: Arc::new(RealtimeConversationManager::new()),
             active_turn: Mutex::new(None),
+            mailbox,
+            mailbox_rx: Mutex::new(mailbox_rx),
             idle_pending_input: Mutex::new(Vec::new()),
             guardian_review_session: GuardianReviewSessionManager::default(),
             services,
@@ -3957,6 +3965,18 @@ impl Session {
         }
     }
 
+    pub(crate) fn subscribe_mailbox_seq(&self) -> watch::Receiver<u64> {
+        self.mailbox.subscribe()
+    }
+
+    pub(crate) fn enqueue_mailbox_communication(&self, communication: InterAgentCommunication) {
+        self.mailbox.send(communication);
+    }
+
+    pub(crate) async fn has_trigger_turn_mailbox_items(&self) -> bool {
+        self.mailbox_rx.lock().await.has_pending_trigger_turn()
+    }
+
     pub async fn prepend_pending_input(&self, input: Vec<ResponseInputItem>) -> Result<(), ()> {
         let mut active = self.active_turn.lock().await;
         match active.as_mut() {
@@ -3970,17 +3990,37 @@ impl Session {
     }
 
     pub async fn get_pending_input(&self) -> Vec<ResponseInputItem> {
-        let mut active = self.active_turn.lock().await;
-        match active.as_mut() {
-            Some(at) => {
-                let mut ts = at.turn_state.lock().await;
-                ts.take_pending_input()
+        let pending_input = {
+            let mut active = self.active_turn.lock().await;
+            match active.as_mut() {
+                Some(at) => {
+                    let mut ts = at.turn_state.lock().await;
+                    ts.take_pending_input()
+                }
+                None => Vec::new(),
             }
-            None => Vec::with_capacity(0),
+        };
+        let mailbox_items = {
+            let mut mailbox_rx = self.mailbox_rx.lock().await;
+            mailbox_rx
+                .drain()
+                .into_iter()
+                .map(|mail| mail.to_response_input_item())
+                .collect::<Vec<_>>()
+        };
+        if pending_input.is_empty() {
+            mailbox_items
+        } else if mailbox_items.is_empty() {
+            pending_input
+        } else {
+            let mut pending_input = pending_input;
+            pending_input.extend(mailbox_items);
+            pending_input
         }
     }
 
     /// Queue response items to be injected into the next active turn created for this session.
+    #[cfg(test)]
     pub(crate) async fn queue_response_items_for_next_turn(&self, items: Vec<ResponseInputItem>) {
         if items.is_empty() {
             return;
@@ -3999,6 +4039,9 @@ impl Session {
     }
 
     pub async fn has_pending_input(&self) -> bool {
+        if self.mailbox_rx.lock().await.has_pending() {
+            return true;
+        }
         let active = self.active_turn.lock().await;
         match active.as_ref() {
             Some(at) => {
@@ -4664,18 +4707,10 @@ mod handlers {
         sub_id: String,
         communication: InterAgentCommunication,
     ) {
-        let pending_item = communication.to_response_input_item();
-        if let Ok(()) = sess.inject_response_items(vec![pending_item.clone()]).await {
-            return;
-        }
-
-        sess.queue_response_items_for_next_turn(vec![pending_item])
-            .await;
-        if communication.trigger_turn {
-            let turn_context = sess.new_default_turn_with_sub_id(sub_id).await;
-            sess.maybe_emit_unknown_model_warning_for_turn(turn_context.as_ref())
-                .await;
-            sess.spawn_task(turn_context, Vec::new(), crate::tasks::RegularTask::new())
+        let trigger_turn = communication.trigger_turn;
+        sess.enqueue_mailbox_communication(communication);
+        if trigger_turn {
+            sess.ensure_task_for_pending_inputs_with_sub_id(sub_id)
                 .await;
         }
     }

codex-rs/core/src/codex_tests.rs

@@ -2737,6 +2737,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
         skills_outcome,
     );
 
+    let (mailbox, mailbox_rx) = crate::agent::Mailbox::new();
     let session = Session {
         conversation_id,
         tx_event,
@@ -2747,6 +2748,8 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
         pending_mcp_server_refresh_config: Mutex::new(None),
         conversation: Arc::new(RealtimeConversationManager::new()),
         active_turn: Mutex::new(None),
+        mailbox,
+        mailbox_rx: Mutex::new(mailbox_rx),
         idle_pending_input: Mutex::new(Vec::new()),
         guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(),
         services,
@@ -3577,6 +3580,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx(
         skills_outcome,
     ));
 
+    let (mailbox, mailbox_rx) = crate::agent::Mailbox::new();
     let session = Arc::new(Session {
         conversation_id,
         tx_event,
@@ -3587,6 +3591,8 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx(
         pending_mcp_server_refresh_config: Mutex::new(None),
         conversation: Arc::new(RealtimeConversationManager::new()),
         active_turn: Mutex::new(None),
+        mailbox,
+        mailbox_rx: Mutex::new(mailbox_rx),
         idle_pending_input: Mutex::new(Vec::new()),
         guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(),
         services,

codex-rs/core/src/codex_thread.rs

@@ -180,10 +180,7 @@ impl CodexThread {
                 .session
                 .queue_response_items_for_next_turn(items)
                 .await;
-            self.codex
-                .session
-                .ensure_task_for_queued_response_items()
-                .await;
+            self.codex.session.ensure_task_for_pending_inputs().await;
         }
 
         Ok(submission_id)

codex-rs/core/src/tasks/mod.rs

@@ -231,6 +231,7 @@ impl Session {
         };
 
         let queued_response_items = self.take_queued_response_items_for_next_turn().await;
+        let mailbox_items = self.get_pending_input().await;
         let mut active = self.active_turn.lock().await;
         let mut turn = ActiveTurn::default();
         let mut turn_state = turn.turn_state.lock().await;
@@ -238,6 +239,9 @@ impl Session {
         for item in queued_response_items {
             turn_state.push_pending_input(item);
         }
+        for item in mailbox_items {
+            turn_state.push_pending_input(item);
+        }
         drop(turn_state);
 
         let timer = turn_context
@@ -258,8 +262,27 @@ impl Session {
         *active = Some(turn);
     }
 
-    pub(crate) async fn ensure_task_for_queued_response_items(self: &Arc<Self>) {
-        if !self.has_queued_response_items_for_next_turn().await {
+    /// Starts a regular turn when queued next-turn items or trigger-turn mailbox mail are waiting.
+    ///
+    /// This helper generates a fresh sub-id for the synthetic turn before delegating to the
+    /// explicit-sub-id variant.
+    pub(crate) async fn ensure_task_for_pending_inputs(self: &Arc<Self>) {
+        self.ensure_task_for_pending_inputs_with_sub_id(uuid::Uuid::new_v4().to_string())
+            .await;
+    }
+
+    /// Starts a regular turn with the provided sub-id when pending input should wake an idle
+    /// session.
+    ///
+    /// The turn is created only when there are queued next-turn items or mailbox mail marked with
+    /// `trigger_turn`, and only if the session is currently idle.
+    pub(crate) async fn ensure_task_for_pending_inputs_with_sub_id(
+        self: &Arc<Self>,
+        sub_id: String,
+    ) {
+        if !self.has_queued_response_items_for_next_turn().await
+            && !self.has_trigger_turn_mailbox_items().await
+        {
             return;
         }
 
@@ -267,7 +290,7 @@ impl Session {
             return;
         }
 
-        let turn_context = self.new_default_turn().await;
+        let turn_context = self.new_default_turn_with_sub_id(sub_id).await;
         self.maybe_emit_unknown_model_warning_for_turn(turn_context.as_ref())
             .await;
         self.start_task(turn_context, Vec::new(), RegularTask::new())
@@ -284,7 +307,7 @@ impl Session {
             active_turn.clear_pending().await;
         }
         if reason == TurnAbortReason::Interrupted {
-            self.ensure_task_for_queued_response_items().await;
+            self.ensure_task_for_pending_inputs().await;
         }
     }
 

codex-rs/core/src/tools/handlers/multi_agents_tests.rs

@@ -940,12 +940,7 @@ async fn multi_agent_v2_send_message_interrupts_busy_child_without_triggering_tu
 
     timeout(Duration::from_secs(5), async {
         loop {
-            if !thread
-                .codex
-                .session
-                .has_queued_response_items_for_next_turn()
-                .await
-            {
+            if !thread.codex.session.has_pending_input().await {
                 tokio::time::sleep(Duration::from_millis(10)).await;
                 continue;
             }
@@ -1788,27 +1783,78 @@ async fn wait_agent_rejects_empty_targets() {
 }
 
 #[tokio::test]
-async fn multi_agent_v2_wait_agent_accepts_targets_argument() {
+async fn multi_agent_v2_wait_agent_accepts_timeout_only_argument() {
     let (mut session, mut turn) = make_session_and_context().await;
-    let target = ThreadId::new().to_string();
     let manager = thread_manager();
+    let root = manager
+        .start_thread((*turn.config).clone())
+        .await
+        .expect("root thread should start");
     session.services.agent_control = manager.agent_control();
+    session.conversation_id = root.thread_id;
     let mut config = (*turn.config).clone();
     config
         .features
         .enable(Feature::MultiAgentV2)
         .expect("test config should allow feature update");
     turn.config = Arc::new(config);
-    let invocation = invocation(
-        Arc::new(session),
-        Arc::new(turn),
-        "wait_agent",
-        function_payload(json!({"targets": [target.clone()]})),
-    );
-    let output = WaitAgentHandlerV2
-        .handle(invocation)
+    let session = Arc::new(session);
+    let turn = Arc::new(turn);
+
+    SpawnAgentHandlerV2
+        .handle(invocation(
+            session.clone(),
+            turn.clone(),
+            "spawn_agent",
+            function_payload(json!({
+                "message": "boot worker",
+                "task_name": "worker"
+            })),
+        ))
         .await
-        .expect("targets should be accepted in v2 mode");
+        .expect("spawn worker");
+    let agent_id = session
+        .services
+        .agent_control
+        .resolve_agent_reference(session.conversation_id, &turn.session_source, "worker")
+        .await
+        .expect("worker should resolve");
+    let worker_path = session
+        .services
+        .agent_control
+        .get_agent_metadata(agent_id)
+        .expect("worker metadata")
+        .agent_path
+        .expect("worker path");
+
+    let wait_task = tokio::spawn({
+        let session = session.clone();
+        let turn = turn.clone();
+        async move {
+            WaitAgentHandlerV2
+                .handle(invocation(
+                    session,
+                    turn,
+                    "wait_agent",
+                    function_payload(json!({"timeout_ms": 1000})),
+                ))
+                .await
+        }
+    });
+    tokio::task::yield_now().await;
+
+    session.enqueue_mailbox_communication(InterAgentCommunication::new(
+        worker_path,
+        AgentPath::root(),
+        Vec::new(),
+        "hello from worker".to_string(),
+        /*trigger_turn*/ false,
+    ));
+
+    let output = wait_task
+        .await
+        .expect("wait task should join")
+        .expect("timeout-only args should be accepted in v2 mode");
     let (content, success) = expect_text_output(output);
     let result: crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult =
         serde_json::from_str(&content).expect("wait_agent result should be json");
@@ -1983,7 +2029,7 @@ async fn wait_agent_returns_final_status_without_timeout() {
 }
 
 #[tokio::test]
-async fn multi_agent_v2_wait_agent_returns_summary_for_named_targets() {
+async fn multi_agent_v2_wait_agent_returns_summary_for_mailbox_activity() {
     let (mut session, mut turn) = make_session_and_context().await;
     let manager = thread_manager();
     let root = manager
@@ -2025,37 +2071,229 @@ async fn multi_agent_v2_wait_agent_returns_summary_for_named_targets() {
         )
         .await
         .expect("relative path should resolve");
-    let mut status_rx = manager
-        .agent_control()
-        .subscribe_status(agent_id)
-        .await
-        .expect("subscribe should succeed");
+    let worker_path = session
+        .services
+        .agent_control
+        .get_agent_metadata(agent_id)
+        .expect("worker metadata")
+        .agent_path
+        .expect("worker path");
+    let wait_task = tokio::spawn({
+        let session = session.clone();
+        let turn = turn.clone();
+        async move {
+            WaitAgentHandlerV2
+                .handle(invocation(
+                    session,
+                    turn,
+                    "wait_agent",
+                    function_payload(json!({"timeout_ms": 1000})),
+                ))
+                .await
+        }
+    });
+    tokio::task::yield_now().await;
 
-    let child_thread = manager
-        .get_thread(agent_id)
-        .await
-        .expect("child should exist");
-    let _ = child_thread
-        .submit(Op::Shutdown {})
-        .await
-        .expect("shutdown should submit");
-    let _ = timeout(Duration::from_secs(1), status_rx.changed())
-        .await
-        .expect("shutdown status should arrive");
+    session.enqueue_mailbox_communication(InterAgentCommunication::new(
+        worker_path,
+        AgentPath::root(),
+        Vec::new(),
+        "completed".to_string(),
+        /*trigger_turn*/ false,
+    ));
 
-    let wait_output = WaitAgentHandlerV2
+    let wait_output = wait_task
+        .await
+        .expect("wait task should join")
+        .expect("wait_agent should succeed");
+    let (content, success) = expect_text_output(wait_output);
+    let result: crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult =
+        serde_json::from_str(&content).expect("wait_agent result should be json");
+    assert_eq!(
+        result,
+        crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult {
+            message: "Wait completed.".to_string(),
+            timed_out: false,
+        }
+    );
+    assert_eq!(success, None);
+}
+
+#[tokio::test]
+async fn multi_agent_v2_wait_agent_waits_for_new_mail_after_start() {
+    let (mut session, mut turn) = make_session_and_context().await;
+    let manager = thread_manager();
+    let root = manager
+        .start_thread((*turn.config).clone())
+        .await
+        .expect("root thread should start");
+    session.services.agent_control = manager.agent_control();
+    session.conversation_id = root.thread_id;
+    let mut config = (*turn.config).clone();
+    config
+        .features
+        .enable(Feature::MultiAgentV2)
+        .expect("test config should allow feature update");
+    turn.config = Arc::new(config);
+    let session = Arc::new(session);
+    let turn = Arc::new(turn);
+
+    SpawnAgentHandlerV2
         .handle(invocation(
-            session,
-            turn,
-            "wait_agent",
+            session.clone(),
+            turn.clone(),
+            "spawn_agent",
             function_payload(json!({
-                "targets": ["test_process"],
-                "timeout_ms": 1000
+                "message": "boot worker",
+                "task_name": "worker"
             })),
         ))
         .await
+        .expect("spawn worker");
+    let agent_id = session
+        .services
+        .agent_control
+        .resolve_agent_reference(session.conversation_id, &turn.session_source, "worker")
+        .await
+        .expect("worker should resolve");
+    let worker_path = session
+        .services
+        .agent_control
+        .get_agent_metadata(agent_id)
+        .expect("worker metadata")
+        .agent_path
+        .expect("worker path");
+
+    session.enqueue_mailbox_communication(InterAgentCommunication::new(
+        worker_path.clone(),
+        AgentPath::root(),
+        Vec::new(),
+        "already queued".to_string(),
+        /*trigger_turn*/ false,
+    ));
+
+    let wait_task = tokio::spawn({
+        let session = session.clone();
+        let turn = turn.clone();
+        async move {
+            WaitAgentHandlerV2
+                .handle(invocation(
+                    session,
+                    turn,
+                    "wait_agent",
+                    function_payload(json!({"timeout_ms": 1000})),
+                ))
+                .await
+        }
+    });
+    tokio::task::yield_now().await;
+    tokio::time::sleep(Duration::from_millis(50)).await;
+    assert!(
+        !wait_task.is_finished(),
+        "mail already queued before wait should not wake wait_agent"
+    );
+
+    session.enqueue_mailbox_communication(InterAgentCommunication::new(
+        worker_path,
+        AgentPath::root(),
+        Vec::new(),
+        "new mail".to_string(),
+        /*trigger_turn*/ false,
+    ));
+
+    let output = wait_task
+        .await
+        .expect("wait task should join")
         .expect("wait_agent should succeed");
-    let (content, success) = expect_text_output(wait_output);
+    let (content, success) = expect_text_output(output);
+    let result: crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult =
+        serde_json::from_str(&content).expect("wait_agent result should be json");
+    assert_eq!(
+        result,
+        crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult {
+            message: "Wait completed.".to_string(),
+            timed_out: false,
+        }
+    );
+    assert_eq!(success, None);
+}
+
+#[tokio::test]
+async fn multi_agent_v2_wait_agent_wakes_on_any_mailbox_notification() {
+    let (mut session, mut turn) = make_session_and_context().await;
+    let manager = thread_manager();
+    let root = manager
+        .start_thread((*turn.config).clone())
+        .await
+        .expect("root thread should start");
+    session.services.agent_control = manager.agent_control();
+    session.conversation_id = root.thread_id;
+    let mut config = (*turn.config).clone();
+    config
+        .features
+        .enable(Feature::MultiAgentV2)
+        .expect("test config should allow feature update");
+    turn.config = Arc::new(config);
+    let session = Arc::new(session);
+    let turn = Arc::new(turn);
+
+    for task_name in ["worker_a", "worker_b"] {
+        SpawnAgentHandlerV2
+            .handle(invocation(
+                session.clone(),
+                turn.clone(),
+                "spawn_agent",
+                function_payload(json!({
+                    "message": format!("boot {task_name}"),
+                    "task_name": task_name
+                })),
+            ))
+            .await
+            .expect("spawn worker");
+    }
+    let worker_b_id = session
+        .services
+        .agent_control
+        .resolve_agent_reference(session.conversation_id, &turn.session_source, "worker_b")
+        .await
+        .expect("worker_b should resolve");
+    let worker_b_path = session
+        .services
+        .agent_control
+        .get_agent_metadata(worker_b_id)
+        .expect("worker_b metadata")
+        .agent_path
+        .expect("worker_b path");
+
+    let wait_task = tokio::spawn({
+        let session = session.clone();
+        let turn = turn.clone();
+        async move {
+            WaitAgentHandlerV2
+                .handle(invocation(
+                    session,
+                    turn,
+                    "wait_agent",
+                    function_payload(json!({"timeout_ms": 1000})),
+                ))
+                .await
+        }
+    });
+    tokio::task::yield_now().await;
+
+    session.enqueue_mailbox_communication(InterAgentCommunication::new(
+        worker_b_path,
+        AgentPath::root(),
+        Vec::new(),
+        "from worker b".to_string(),
+        /*trigger_turn*/ false,
+    ));
+
+    let output = wait_task
+        .await
+        .expect("wait task should join")
+        .expect("wait_agent should succeed");
+    let (content, success) = expect_text_output(output);
     let result: crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult =
         serde_json::from_str(&content).expect("wait_agent result should be json");
     assert_eq!(
@@ -2072,41 +2310,73 @@ async fn multi_agent_v2_wait_agent_returns_summary_for_named_targets() {
 async fn multi_agent_v2_wait_agent_does_not_return_completed_content() {
     let (mut session, mut turn) = make_session_and_context().await;
     let manager = thread_manager();
+    let root = manager
+        .start_thread((*turn.config).clone())
+        .await
+        .expect("root thread should start");
     session.services.agent_control = manager.agent_control();
+    session.conversation_id = root.thread_id;
     let mut config = (*turn.config).clone();
     config
         .features
         .enable(Feature::MultiAgentV2)
         .expect("test config should allow feature update");
-    turn.config = Arc::new(config.clone());
+    turn.config = Arc::new(config);
+    let session = Arc::new(session);
+    let turn = Arc::new(turn);
 
-    let thread = manager.start_thread(config).await.expect("start thread");
-    let agent_id = thread.thread_id;
-    let child_turn = thread.thread.codex.session.new_default_turn().await;
-    thread
-        .thread
-        .codex
-        .session
-        .send_event(
-            child_turn.as_ref(),
-            EventMsg::TurnComplete(TurnCompleteEvent {
-                turn_id: child_turn.sub_id.clone(),
-                last_agent_message: Some("sensitive child output".to_string()),
-            }),
-        )
-        .await;
-
-    let output = WaitAgentHandlerV2
+    SpawnAgentHandlerV2
         .handle(invocation(
-            Arc::new(session),
-            Arc::new(turn),
-            "wait_agent",
+            session.clone(),
+            turn.clone(),
+            "spawn_agent",
             function_payload(json!({
-                "targets": [agent_id.to_string()],
-                "timeout_ms": 1000
+                "message": "boot worker",
+                "task_name": "worker"
             })),
         ))
         .await
+        .expect("spawn worker");
+    let agent_id = session
+        .services
+        .agent_control
+        .resolve_agent_reference(session.conversation_id, &turn.session_source, "worker")
+        .await
+        .expect("worker should resolve");
+    let worker_path = session
+        .services
+        .agent_control
+        .get_agent_metadata(agent_id)
+        .expect("worker metadata")
+        .agent_path
+        .expect("worker path");
+    let wait_task = tokio::spawn({
+        let session = session.clone();
+        let turn = turn.clone();
+        async move {
+            WaitAgentHandlerV2
+                .handle(invocation(
+                    session,
+                    turn,
+                    "wait_agent",
+                    function_payload(json!({"timeout_ms": 1000})),
+                ))
+                .await
+        }
+    });
+    tokio::task::yield_now().await;
+
+    session.enqueue_mailbox_communication(InterAgentCommunication::new(
+        worker_path,
+        AgentPath::root(),
+        Vec::new(),
+        "sensitive child output".to_string(),
+        /*trigger_turn*/ false,
+    ));
+
+    let output = wait_task
+        .await
+        .expect("wait task should join")
         .expect("wait_agent should succeed");
     let (content, success) = expect_text_output(output);
     let result: crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult =

codex-rs/core/src/tools/handlers/multi_agents_v2.rs

@@ -2,9 +2,7 @@
 
 use crate::agent::AgentStatus;
 use crate::agent::agent_resolver::resolve_agent_target;
-use crate::agent::agent_resolver::resolve_agent_targets;
 use crate::agent::exceeds_thread_spawn_depth_limit;
-use crate::codex::Session;
 use crate::function_tool::FunctionCallError;
 use crate::tools::context::ToolInvocation;
 use crate::tools::context::ToolOutput;
@@ -15,12 +13,10 @@ use crate::tools::registry::ToolHandler;
 use crate::tools::registry::ToolKind;
 use async_trait::async_trait;
 use codex_protocol::AgentPath;
-use codex_protocol::ThreadId;
 use codex_protocol::models::ResponseInputItem;
 use codex_protocol::openai_models::ReasoningEffort;
 use codex_protocol::protocol::CollabAgentInteractionBeginEvent;
 use codex_protocol::protocol::CollabAgentInteractionEndEvent;
-use codex_protocol::protocol::CollabAgentRef;
 use codex_protocol::protocol::CollabAgentSpawnBeginEvent;
 use codex_protocol::protocol::CollabAgentSpawnEndEvent;
 use codex_protocol::protocol::CollabCloseBeginEvent;

codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs

@@ -1,11 +1,6 @@
 use super::*;
-use crate::agent::status::is_final;
-use futures::FutureExt;
-use futures::StreamExt;
-use futures::stream::FuturesUnordered;
 use std::collections::HashMap;
 use std::time::Duration;
-use tokio::sync::watch::Receiver;
 use tokio::time::Instant;
 use tokio::time::timeout_at;
 
@@ -33,21 +28,6 @@ impl ToolHandler for Handler {
         } = invocation;
         let arguments = function_arguments(payload)?;
         let args: WaitArgs = parse_arguments(&arguments)?;
-        let receiver_thread_ids = resolve_agent_targets(&session, &turn, args.targets).await?;
-        let mut receiver_agents = Vec::with_capacity(receiver_thread_ids.len());
-        for receiver_thread_id in &receiver_thread_ids {
-            let agent_metadata = session
-                .services
-                .agent_control
-                .get_agent_metadata(*receiver_thread_id)
-                .unwrap_or_default();
-            receiver_agents.push(CollabAgentRef {
-                thread_id: *receiver_thread_id,
-                agent_nickname: agent_metadata.agent_nickname,
-                agent_role: agent_metadata.agent_role,
-            });
-        }
-
         let timeout_ms = args.timeout_ms.unwrap_or(DEFAULT_WAIT_TIMEOUT_MS);
         let timeout_ms = match timeout_ms {
             ms if ms <= 0 => {
@@ -63,86 +43,17 @@ impl ToolHandler for Handler {
                 &turn,
                 CollabWaitingBeginEvent {
                     sender_thread_id: session.conversation_id,
-                    receiver_thread_ids: receiver_thread_ids.clone(),
-                    receiver_agents: receiver_agents.clone(),
+                    receiver_thread_ids: Vec::new(),
+                    receiver_agents: Vec::new(),
                     call_id: call_id.clone(),
                 }
                 .into(),
             )
             .await;
 
-        let mut status_rxs = Vec::with_capacity(receiver_thread_ids.len());
-        let mut initial_final_statuses = Vec::new();
-        for id in &receiver_thread_ids {
-            match session.services.agent_control.subscribe_status(*id).await {
-                Ok(rx) => {
-                    let status = rx.borrow().clone();
-                    if is_final(&status) {
-                        initial_final_statuses.push((*id, status));
-                    }
-                    status_rxs.push((*id, rx));
-                }
-                Err(crate::error::CodexErr::ThreadNotFound(_)) => {
-                    initial_final_statuses.push((*id, AgentStatus::NotFound));
-                }
-                Err(err) => {
-                    let mut statuses = HashMap::with_capacity(1);
-                    statuses.insert(*id, session.services.agent_control.get_status(*id).await);
-                    session
-                        .send_event(
-                            &turn,
-                            CollabWaitingEndEvent {
-                                sender_thread_id: session.conversation_id,
-                                call_id: call_id.clone(),
-                                agent_statuses: build_wait_agent_statuses(
-                                    &statuses,
-                                    &receiver_agents,
-                                ),
-                                statuses,
-                            }
-                            .into(),
-                        )
-                        .await;
-                    return Err(collab_agent_error(*id, err));
-                }
-            }
-        }
-
-        let statuses = if !initial_final_statuses.is_empty() {
-            initial_final_statuses
-        } else {
-            let mut futures = FuturesUnordered::new();
-            for (id, rx) in status_rxs {
-                let session = session.clone();
-                futures.push(wait_for_final_status(session, id, rx));
-            }
-            let mut results = Vec::new();
-            let deadline = Instant::now() + Duration::from_millis(timeout_ms as u64);
-            loop {
-                match timeout_at(deadline, futures.next()).await {
-                    Ok(Some(Some(result))) => {
-                        results.push(result);
-                        break;
-                    }
-                    Ok(Some(None)) => continue,
-                    Ok(None) | Err(_) => break,
-                }
-            }
-            if !results.is_empty() {
-                loop {
-                    match futures.next().now_or_never() {
-                        Some(Some(Some(result))) => results.push(result),
-                        Some(Some(None)) => continue,
-                        Some(None) | None => break,
-                    }
-                }
-            }
-            results
-        };
-
-        let timed_out = statuses.is_empty();
-        let statuses_by_id = statuses.clone().into_iter().collect::<HashMap<_, _>>();
-        let agent_statuses = build_wait_agent_statuses(&statuses_by_id, &receiver_agents);
+        let mut mailbox_seq_rx = session.subscribe_mailbox_seq();
+        let deadline = Instant::now() + Duration::from_millis(timeout_ms as u64);
+        let timed_out = !wait_for_mailbox_change(&mut mailbox_seq_rx, deadline).await;
         let result = WaitAgentResult::from_timed_out(timed_out);
 
         session
@@ -151,8 +62,8 @@ impl ToolHandler for Handler {
                 CollabWaitingEndEvent {
                     sender_thread_id: session.conversation_id,
                     call_id,
-                    agent_statuses,
-                    statuses: statuses_by_id,
+                    agent_statuses: Vec::new(),
+                    statuses: HashMap::new(),
                 }
                 .into(),
             )
@@ -164,8 +75,6 @@ impl ToolHandler for Handler {
 
 #[derive(Debug, Deserialize)]
 struct WaitArgs {
-    #[serde(default)]
-    targets: Vec<String>,
     timeout_ms: Option<i64>,
 }
 
@@ -207,24 +116,12 @@ impl ToolOutput for WaitAgentResult {
     }
 }
 
-async fn wait_for_final_status(
-    session: std::sync::Arc<Session>,
-    thread_id: ThreadId,
-    mut status_rx: Receiver<AgentStatus>,
-) -> Option<(ThreadId, AgentStatus)> {
-    let mut status = status_rx.borrow().clone();
-    if is_final(&status) {
-        return Some((thread_id, status));
-    }
-
-    loop {
-        if status_rx.changed().await.is_err() {
-            let latest = session.services.agent_control.get_status(thread_id).await;
-            return is_final(&latest).then_some((thread_id, latest));
-        }
-        status = status_rx.borrow().clone();
-        if is_final(&status) {
-            return Some((thread_id, status));
-        }
+async fn wait_for_mailbox_change(
+    mailbox_seq_rx: &mut tokio::sync::watch::Receiver<u64>,
+    deadline: Instant,
+) -> bool {
+    match timeout_at(deadline, mailbox_seq_rx.changed()).await {
+        Ok(Ok(())) => true,
+        Ok(Err(_)) | Err(_) => false,
     }
 }