diff --git a/codex-rs/core/src/agent/agent_resolver.rs b/codex-rs/core/src/agent/agent_resolver.rs index 3d1f75f57..14f21c731 100644 --- a/codex-rs/core/src/agent/agent_resolver.rs +++ b/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, - turn: &Arc, - targets: Vec, -) -> Result, 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, turn: &Arc) { session .services diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 2def29af9..e440f9838 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/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; diff --git a/codex-rs/core/src/agent/mailbox.rs b/codex-rs/core/src/agent/mailbox.rs new file mode 100644 index 000000000..c32823647 --- /dev/null +++ b/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, + next_seq: AtomicU64, + seq_tx: watch::Sender, +} + +pub(crate) struct MailboxReceiver { + rx: mpsc::UnboundedReceiver, + pending_mails: VecDeque, +} + +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 { + 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 { + 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()); + } +} diff --git a/codex-rs/core/src/agent/mod.rs b/codex-rs/core/src/agent/mod.rs index 350962dc0..a60fc3004 100644 --- a/codex-rs/core/src/agent/mod.rs +++ b/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; diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 4dca8dc68..9a8cc9cea 100644 --- a/codex-rs/core/src/codex.rs +++ b/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>, pub(crate) conversation: Arc, pub(crate) active_turn: Mutex>, - idle_pending_input: Mutex>, + mailbox: Mailbox, + mailbox_rx: Mutex, + idle_pending_input: Mutex>, // TODO (jif) merge with mailbox! pub(crate) guardian_review_session: GuardianReviewSessionManager, pub(crate) services: SessionServices, js_repl: Arc, @@ -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 { + 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) -> 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 { - 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::>() + }; + 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) { 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; } } diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index f1a665277..f5e5d4463 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/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, diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 0635718c9..033f07aa1 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/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) diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 601f5cb27..e1db8b78c 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/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) { - 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.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, + 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; } } 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 8006ee220..091b42157 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/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 = diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2.rs index c25f0918e..e572ec0af 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2.rs +++ b/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; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs index 1b1dc90e4..164bcac60 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs +++ b/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::>(); - 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, timeout_ms: Option, } @@ -207,24 +116,12 @@ impl ToolOutput for WaitAgentResult { } } -async fn wait_for_final_status( - session: std::sync::Arc, - thread_id: ThreadId, - mut status_rx: Receiver, -) -> 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, + deadline: Instant, +) -> bool { + match timeout_at(deadline, mailbox_seq_rx.changed()).await { + Ok(Ok(())) => true, + Ok(Err(_)) | Err(_) => false, } }