diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index aa6923fd1..7ae1d1428 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -44,7 +44,6 @@ pub use js_repl::JsReplResetHandler; pub use list_dir::ListDirHandler; pub use mcp::McpHandler; pub use mcp_resource::McpResourceHandler; -pub use multi_agents::MultiAgentHandler; pub use plan::PlanHandler; pub use read_file::ReadFileHandler; pub use request_permissions::RequestPermissionsHandler; diff --git a/codex-rs/core/src/tools/handlers/multi_agents.rs b/codex-rs/core/src/tools/handlers/multi_agents.rs index 193a4e5e7..87b11916a 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents.rs @@ -43,9 +43,13 @@ use codex_protocol::user_input::UserInput; use serde::Deserialize; use serde::Serialize; use std::collections::HashMap; +use std::sync::Arc; -/// Function-tool handler for the multi-agent collaboration API. -pub struct MultiAgentHandler; +pub(crate) use close_agent::Handler as CloseAgentHandler; +pub(crate) use resume_agent::Handler as ResumeAgentHandler; +pub(crate) use send_input::Handler as SendInputHandler; +pub(crate) use spawn::Handler as SpawnAgentHandler; +pub(crate) use wait::Handler as WaitHandler; /// Minimum wait timeout to prevent tight polling loops from burning CPU. pub(crate) const MIN_WAIT_TIMEOUT_MS: i64 = 10_000; @@ -57,47 +61,12 @@ struct CloseAgentArgs { id: String, } -#[async_trait] -impl ToolHandler for MultiAgentHandler { - type Output = FunctionToolOutput; - - fn kind(&self) -> ToolKind { - ToolKind::Function - } - - fn matches_kind(&self, payload: &ToolPayload) -> bool { - matches!(payload, ToolPayload::Function { .. }) - } - - async fn handle(&self, invocation: ToolInvocation) -> Result { - let ToolInvocation { - session, - turn, - tool_name, - payload, - call_id, - .. - } = invocation; - - let arguments = match payload { - ToolPayload::Function { arguments } => arguments, - _ => { - return Err(FunctionCallError::RespondToModel( - "collab handler received unsupported payload".to_string(), - )); - } - }; - - match tool_name.as_str() { - "spawn_agent" => spawn::handle(session, turn, call_id, arguments).await, - "send_input" => send_input::handle(session, turn, call_id, arguments).await, - "resume_agent" => resume_agent::handle(session, turn, call_id, arguments).await, - "wait" => wait::handle(session, turn, call_id, arguments).await, - "close_agent" => close_agent::handle(session, turn, call_id, arguments).await, - other => Err(FunctionCallError::RespondToModel(format!( - "unsupported collab tool {other}" - ))), - } +fn function_arguments(payload: ToolPayload) -> Result { + match payload { + ToolPayload::Function { arguments } => Ok(arguments), + _ => Err(FunctionCallError::RespondToModel( + "collab handler received unsupported payload".to_string(), + )), } } @@ -109,7 +78,145 @@ mod spawn { use crate::agent::exceeds_thread_spawn_depth_limit; use crate::agent::next_thread_spawn_depth; - use std::sync::Arc; + + pub(crate) struct Handler; + + #[async_trait] + impl ToolHandler for Handler { + type Output = FunctionToolOutput; + + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + fn matches_kind(&self, payload: &ToolPayload) -> bool { + matches!(payload, ToolPayload::Function { .. }) + } + + async fn handle( + &self, + invocation: ToolInvocation, + ) -> Result { + let ToolInvocation { + session, + turn, + payload, + call_id, + .. + } = invocation; + let arguments = function_arguments(payload)?; + let args: SpawnAgentArgs = parse_arguments(&arguments)?; + let role_name = args + .agent_type + .as_deref() + .map(str::trim) + .filter(|role| !role.is_empty()); + let input_items = parse_collab_input(args.message, args.items)?; + let prompt = input_preview(&input_items); + let session_source = turn.session_source.clone(); + let child_depth = next_thread_spawn_depth(&session_source); + let max_depth = turn.config.agent_max_depth; + if exceeds_thread_spawn_depth_limit(child_depth, max_depth) { + return Err(FunctionCallError::RespondToModel( + "Agent depth limit reached. Solve the task yourself.".to_string(), + )); + } + session + .send_event( + &turn, + CollabAgentSpawnBeginEvent { + call_id: call_id.clone(), + sender_thread_id: session.conversation_id, + prompt: prompt.clone(), + model: args.model.clone().unwrap_or_default(), + reasoning_effort: args.reasoning_effort.unwrap_or_default(), + } + .into(), + ) + .await; + let mut config = + build_agent_spawn_config(&session.get_base_instructions().await, turn.as_ref())?; + apply_requested_spawn_agent_model_overrides( + &session, + turn.as_ref(), + &mut config, + args.model.as_deref(), + args.reasoning_effort, + ) + .await?; + apply_role_to_config(&mut config, role_name) + .await + .map_err(FunctionCallError::RespondToModel)?; + apply_spawn_agent_runtime_overrides(&mut config, turn.as_ref())?; + apply_spawn_agent_overrides(&mut config, child_depth); + + let result = session + .services + .agent_control + .spawn_agent_with_options( + config, + input_items, + Some(thread_spawn_source( + session.conversation_id, + child_depth, + role_name, + )), + SpawnAgentOptions { + fork_parent_spawn_call_id: args.fork_context.then(|| call_id.clone()), + }, + ) + .await + .map_err(collab_spawn_error); + let (new_thread_id, status) = match &result { + Ok(thread_id) => ( + Some(*thread_id), + session.services.agent_control.get_status(*thread_id).await, + ), + Err(_) => (None, AgentStatus::NotFound), + }; + let (new_agent_nickname, new_agent_role) = match new_thread_id { + Some(thread_id) => session + .services + .agent_control + .get_agent_nickname_and_role(thread_id) + .await + .unwrap_or((None, None)), + None => (None, None), + }; + let nickname = new_agent_nickname.clone(); + session + .send_event( + &turn, + CollabAgentSpawnEndEvent { + call_id, + sender_thread_id: session.conversation_id, + new_thread_id, + new_agent_nickname, + new_agent_role, + prompt, + model: args.model.clone().unwrap_or_default(), + reasoning_effort: args.reasoning_effort.unwrap_or_default(), + status, + } + .into(), + ) + .await; + let new_thread_id = result?; + let role_tag = role_name.unwrap_or(DEFAULT_ROLE_NAME); + turn.session_telemetry + .counter("codex.multi_agent.spawn", 1, &[("role", role_tag)]); + + let content = serde_json::to_string(&SpawnAgentResult { + agent_id: new_thread_id.to_string(), + nickname, + }) + .map_err(|err| { + FunctionCallError::Fatal(format!("failed to serialize spawn_agent result: {err}")) + })?; + + Ok(FunctionToolOutput::from_text(content, Some(true))) + } + } #[derive(Debug, Deserialize)] struct SpawnAgentArgs { @@ -127,129 +234,105 @@ mod spawn { agent_id: String, nickname: Option, } - - pub async fn handle( - session: Arc, - turn: Arc, - call_id: String, - arguments: String, - ) -> Result { - let args: SpawnAgentArgs = parse_arguments(&arguments)?; - let role_name = args - .agent_type - .as_deref() - .map(str::trim) - .filter(|role| !role.is_empty()); - let input_items = parse_collab_input(args.message, args.items)?; - let prompt = input_preview(&input_items); - let session_source = turn.session_source.clone(); - let child_depth = next_thread_spawn_depth(&session_source); - let max_depth = turn.config.agent_max_depth; - if exceeds_thread_spawn_depth_limit(child_depth, max_depth) { - return Err(FunctionCallError::RespondToModel( - "Agent depth limit reached. Solve the task yourself.".to_string(), - )); - } - session - .send_event( - &turn, - CollabAgentSpawnBeginEvent { - call_id: call_id.clone(), - sender_thread_id: session.conversation_id, - prompt: prompt.clone(), - model: args.model.clone().unwrap_or_default(), - reasoning_effort: args.reasoning_effort.unwrap_or_default(), - } - .into(), - ) - .await; - let mut config = - build_agent_spawn_config(&session.get_base_instructions().await, turn.as_ref())?; - apply_requested_spawn_agent_model_overrides( - &session, - turn.as_ref(), - &mut config, - args.model.as_deref(), - args.reasoning_effort, - ) - .await?; - apply_role_to_config(&mut config, role_name) - .await - .map_err(FunctionCallError::RespondToModel)?; - apply_spawn_agent_runtime_overrides(&mut config, turn.as_ref())?; - apply_spawn_agent_overrides(&mut config, child_depth); - - let result = session - .services - .agent_control - .spawn_agent_with_options( - config, - input_items, - Some(thread_spawn_source( - session.conversation_id, - child_depth, - role_name, - )), - SpawnAgentOptions { - fork_parent_spawn_call_id: args.fork_context.then(|| call_id.clone()), - }, - ) - .await - .map_err(collab_spawn_error); - let (new_thread_id, status) = match &result { - Ok(thread_id) => ( - Some(*thread_id), - session.services.agent_control.get_status(*thread_id).await, - ), - Err(_) => (None, AgentStatus::NotFound), - }; - let (new_agent_nickname, new_agent_role) = match new_thread_id { - Some(thread_id) => session - .services - .agent_control - .get_agent_nickname_and_role(thread_id) - .await - .unwrap_or((None, None)), - None => (None, None), - }; - let nickname = new_agent_nickname.clone(); - session - .send_event( - &turn, - CollabAgentSpawnEndEvent { - call_id, - sender_thread_id: session.conversation_id, - new_thread_id, - new_agent_nickname, - new_agent_role, - prompt, - model: args.model.clone().unwrap_or_default(), - reasoning_effort: args.reasoning_effort.unwrap_or_default(), - status, - } - .into(), - ) - .await; - let new_thread_id = result?; - let role_tag = role_name.unwrap_or(DEFAULT_ROLE_NAME); - turn.session_telemetry - .counter("codex.multi_agent.spawn", 1, &[("role", role_tag)]); - - let content = serde_json::to_string(&SpawnAgentResult { - agent_id: new_thread_id.to_string(), - nickname, - }) - .map_err(|err| { - FunctionCallError::Fatal(format!("failed to serialize spawn_agent result: {err}")) - })?; - - Ok(FunctionToolOutput::from_text(content, Some(true))) - } } mod send_input { use super::*; - use std::sync::Arc; + + pub(crate) struct Handler; + + #[async_trait] + impl ToolHandler for Handler { + type Output = FunctionToolOutput; + + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + fn matches_kind(&self, payload: &ToolPayload) -> bool { + matches!(payload, ToolPayload::Function { .. }) + } + + async fn handle( + &self, + invocation: ToolInvocation, + ) -> Result { + let ToolInvocation { + session, + turn, + payload, + call_id, + .. + } = invocation; + let arguments = function_arguments(payload)?; + let args: SendInputArgs = parse_arguments(&arguments)?; + let receiver_thread_id = agent_id(&args.id)?; + let input_items = parse_collab_input(args.message, args.items)?; + let prompt = input_preview(&input_items); + let (receiver_agent_nickname, receiver_agent_role) = session + .services + .agent_control + .get_agent_nickname_and_role(receiver_thread_id) + .await + .unwrap_or((None, None)); + if args.interrupt { + session + .services + .agent_control + .interrupt_agent(receiver_thread_id) + .await + .map_err(|err| collab_agent_error(receiver_thread_id, err))?; + } + session + .send_event( + &turn, + CollabAgentInteractionBeginEvent { + call_id: call_id.clone(), + sender_thread_id: session.conversation_id, + receiver_thread_id, + prompt: prompt.clone(), + } + .into(), + ) + .await; + let result = session + .services + .agent_control + .send_input(receiver_thread_id, input_items) + .await + .map_err(|err| collab_agent_error(receiver_thread_id, err)); + let status = session + .services + .agent_control + .get_status(receiver_thread_id) + .await; + session + .send_event( + &turn, + CollabAgentInteractionEndEvent { + call_id, + sender_thread_id: session.conversation_id, + receiver_thread_id, + receiver_agent_nickname, + receiver_agent_role, + prompt, + status, + } + .into(), + ) + .await; + let submission_id = result?; + + let content = + serde_json::to_string(&SendInputResult { submission_id }).map_err(|err| { + FunctionCallError::Fatal(format!( + "failed to serialize send_input result: {err}" + )) + })?; + + Ok(FunctionToolOutput::from_text(content, Some(true))) + } + } #[derive(Debug, Deserialize)] struct SendInputArgs { @@ -264,83 +347,128 @@ mod send_input { struct SendInputResult { submission_id: String, } - - pub async fn handle( - session: Arc, - turn: Arc, - call_id: String, - arguments: String, - ) -> Result { - let args: SendInputArgs = parse_arguments(&arguments)?; - let receiver_thread_id = agent_id(&args.id)?; - let input_items = parse_collab_input(args.message, args.items)?; - let prompt = input_preview(&input_items); - let (receiver_agent_nickname, receiver_agent_role) = session - .services - .agent_control - .get_agent_nickname_and_role(receiver_thread_id) - .await - .unwrap_or((None, None)); - if args.interrupt { - session - .services - .agent_control - .interrupt_agent(receiver_thread_id) - .await - .map_err(|err| collab_agent_error(receiver_thread_id, err))?; - } - session - .send_event( - &turn, - CollabAgentInteractionBeginEvent { - call_id: call_id.clone(), - sender_thread_id: session.conversation_id, - receiver_thread_id, - prompt: prompt.clone(), - } - .into(), - ) - .await; - let result = session - .services - .agent_control - .send_input(receiver_thread_id, input_items) - .await - .map_err(|err| collab_agent_error(receiver_thread_id, err)); - let status = session - .services - .agent_control - .get_status(receiver_thread_id) - .await; - session - .send_event( - &turn, - CollabAgentInteractionEndEvent { - call_id, - sender_thread_id: session.conversation_id, - receiver_thread_id, - receiver_agent_nickname, - receiver_agent_role, - prompt, - status, - } - .into(), - ) - .await; - let submission_id = result?; - - let content = serde_json::to_string(&SendInputResult { submission_id }).map_err(|err| { - FunctionCallError::Fatal(format!("failed to serialize send_input result: {err}")) - })?; - - Ok(FunctionToolOutput::from_text(content, Some(true))) - } } mod resume_agent { use super::*; use crate::agent::next_thread_spawn_depth; - use std::sync::Arc; + + pub(crate) struct Handler; + + #[async_trait] + impl ToolHandler for Handler { + type Output = FunctionToolOutput; + + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + fn matches_kind(&self, payload: &ToolPayload) -> bool { + matches!(payload, ToolPayload::Function { .. }) + } + + async fn handle( + &self, + invocation: ToolInvocation, + ) -> Result { + let ToolInvocation { + session, + turn, + payload, + call_id, + .. + } = invocation; + let arguments = function_arguments(payload)?; + let args: ResumeAgentArgs = parse_arguments(&arguments)?; + let receiver_thread_id = agent_id(&args.id)?; + let (receiver_agent_nickname, receiver_agent_role) = session + .services + .agent_control + .get_agent_nickname_and_role(receiver_thread_id) + .await + .unwrap_or((None, None)); + let child_depth = next_thread_spawn_depth(&turn.session_source); + let max_depth = turn.config.agent_max_depth; + if exceeds_thread_spawn_depth_limit(child_depth, max_depth) { + return Err(FunctionCallError::RespondToModel( + "Agent depth limit reached. Solve the task yourself.".to_string(), + )); + } + + session + .send_event( + &turn, + CollabResumeBeginEvent { + call_id: call_id.clone(), + sender_thread_id: session.conversation_id, + receiver_thread_id, + receiver_agent_nickname: receiver_agent_nickname.clone(), + receiver_agent_role: receiver_agent_role.clone(), + } + .into(), + ) + .await; + + let mut status = session + .services + .agent_control + .get_status(receiver_thread_id) + .await; + let error = if matches!(status, AgentStatus::NotFound) { + match try_resume_closed_agent(&session, &turn, receiver_thread_id, child_depth) + .await + { + Ok(resumed_status) => { + status = resumed_status; + None + } + Err(err) => { + status = session + .services + .agent_control + .get_status(receiver_thread_id) + .await; + Some(err) + } + } + } else { + None + }; + + let (receiver_agent_nickname, receiver_agent_role) = session + .services + .agent_control + .get_agent_nickname_and_role(receiver_thread_id) + .await + .unwrap_or((receiver_agent_nickname, receiver_agent_role)); + session + .send_event( + &turn, + CollabResumeEndEvent { + call_id, + sender_thread_id: session.conversation_id, + receiver_thread_id, + receiver_agent_nickname, + receiver_agent_role, + status: status.clone(), + } + .into(), + ) + .await; + + if let Some(err) = error { + return Err(err); + } + turn.session_telemetry + .counter("codex.multi_agent.resume", 1, &[]); + + let content = serde_json::to_string(&ResumeAgentResult { status }).map_err(|err| { + FunctionCallError::Fatal(format!("failed to serialize resume_agent result: {err}")) + })?; + + Ok(FunctionToolOutput::from_text(content, Some(true))) + } + } #[derive(Debug, Deserialize)] struct ResumeAgentArgs { @@ -352,101 +480,6 @@ mod resume_agent { pub(super) status: AgentStatus, } - pub async fn handle( - session: Arc, - turn: Arc, - call_id: String, - arguments: String, - ) -> Result { - let args: ResumeAgentArgs = parse_arguments(&arguments)?; - let receiver_thread_id = agent_id(&args.id)?; - let (receiver_agent_nickname, receiver_agent_role) = session - .services - .agent_control - .get_agent_nickname_and_role(receiver_thread_id) - .await - .unwrap_or((None, None)); - let child_depth = next_thread_spawn_depth(&turn.session_source); - let max_depth = turn.config.agent_max_depth; - if exceeds_thread_spawn_depth_limit(child_depth, max_depth) { - return Err(FunctionCallError::RespondToModel( - "Agent depth limit reached. Solve the task yourself.".to_string(), - )); - } - - session - .send_event( - &turn, - CollabResumeBeginEvent { - call_id: call_id.clone(), - sender_thread_id: session.conversation_id, - receiver_thread_id, - receiver_agent_nickname: receiver_agent_nickname.clone(), - receiver_agent_role: receiver_agent_role.clone(), - } - .into(), - ) - .await; - - let mut status = session - .services - .agent_control - .get_status(receiver_thread_id) - .await; - let error = if matches!(status, AgentStatus::NotFound) { - // If the thread is no longer active, attempt to restore it from rollout. - match try_resume_closed_agent(&session, &turn, receiver_thread_id, child_depth).await { - Ok(resumed_status) => { - status = resumed_status; - None - } - Err(err) => { - status = session - .services - .agent_control - .get_status(receiver_thread_id) - .await; - Some(err) - } - } - } else { - None - }; - - let (receiver_agent_nickname, receiver_agent_role) = session - .services - .agent_control - .get_agent_nickname_and_role(receiver_thread_id) - .await - .unwrap_or((receiver_agent_nickname, receiver_agent_role)); - session - .send_event( - &turn, - CollabResumeEndEvent { - call_id, - sender_thread_id: session.conversation_id, - receiver_thread_id, - receiver_agent_nickname, - receiver_agent_role, - status: status.clone(), - } - .into(), - ) - .await; - - if let Some(err) = error { - return Err(err); - } - turn.session_telemetry - .counter("codex.multi_agent.resume", 1, &[]); - - let content = serde_json::to_string(&ResumeAgentResult { status }).map_err(|err| { - FunctionCallError::Fatal(format!("failed to serialize resume_agent result: {err}")) - })?; - - Ok(FunctionToolOutput::from_text(content, Some(true))) - } - async fn try_resume_closed_agent( session: &Arc, turn: &Arc, @@ -480,13 +513,184 @@ pub(crate) mod wait { use futures::StreamExt; use futures::stream::FuturesUnordered; use std::collections::HashMap; - use std::sync::Arc; use std::time::Duration; use tokio::sync::watch::Receiver; use tokio::time::Instant; use tokio::time::timeout_at; + pub(crate) struct Handler; + + #[async_trait] + impl ToolHandler for Handler { + type Output = FunctionToolOutput; + + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + fn matches_kind(&self, payload: &ToolPayload) -> bool { + matches!(payload, ToolPayload::Function { .. }) + } + + async fn handle( + &self, + invocation: ToolInvocation, + ) -> Result { + let ToolInvocation { + session, + turn, + payload, + call_id, + .. + } = invocation; + let arguments = function_arguments(payload)?; + let args: WaitArgs = parse_arguments(&arguments)?; + if args.ids.is_empty() { + return Err(FunctionCallError::RespondToModel( + "ids must be non-empty".to_owned(), + )); + } + let receiver_thread_ids = args + .ids + .iter() + .map(|id| agent_id(id)) + .collect::, _>>()?; + let mut receiver_agents = Vec::with_capacity(receiver_thread_ids.len()); + for receiver_thread_id in &receiver_thread_ids { + let (agent_nickname, agent_role) = session + .services + .agent_control + .get_agent_nickname_and_role(*receiver_thread_id) + .await + .unwrap_or((None, None)); + receiver_agents.push(CollabAgentRef { + thread_id: *receiver_thread_id, + agent_nickname, + agent_role, + }); + } + + let timeout_ms = args.timeout_ms.unwrap_or(DEFAULT_WAIT_TIMEOUT_MS); + let timeout_ms = match timeout_ms { + ms if ms <= 0 => { + return Err(FunctionCallError::RespondToModel( + "timeout_ms must be greater than zero".to_owned(), + )); + } + ms => ms.clamp(MIN_WAIT_TIMEOUT_MS, MAX_WAIT_TIMEOUT_MS), + }; + + session + .send_event( + &turn, + CollabWaitingBeginEvent { + sender_thread_id: session.conversation_id, + receiver_thread_ids: receiver_thread_ids.clone(), + receiver_agents: receiver_agents.clone(), + 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(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.into_iter() { + 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 statuses_map = statuses.clone().into_iter().collect::>(); + let agent_statuses = build_wait_agent_statuses(&statuses_map, &receiver_agents); + let result = WaitResult { + status: statuses_map.clone(), + timed_out: statuses.is_empty(), + }; + + session + .send_event( + &turn, + CollabWaitingEndEvent { + sender_thread_id: session.conversation_id, + call_id, + agent_statuses, + statuses: statuses_map, + } + .into(), + ) + .await; + + let content = serde_json::to_string(&result).map_err(|err| { + FunctionCallError::Fatal(format!("failed to serialize wait result: {err}")) + })?; + + Ok(FunctionToolOutput::from_text(content, None)) + } + } + #[derive(Debug, Deserialize)] struct WaitArgs { ids: Vec, @@ -499,164 +703,6 @@ pub(crate) mod wait { pub(crate) timed_out: bool, } - pub async fn handle( - session: Arc, - turn: Arc, - call_id: String, - arguments: String, - ) -> Result { - let args: WaitArgs = parse_arguments(&arguments)?; - if args.ids.is_empty() { - return Err(FunctionCallError::RespondToModel( - "ids must be non-empty".to_owned(), - )); - } - let receiver_thread_ids = args - .ids - .iter() - .map(|id| agent_id(id)) - .collect::, _>>()?; - let mut receiver_agents = Vec::with_capacity(receiver_thread_ids.len()); - for receiver_thread_id in &receiver_thread_ids { - let (agent_nickname, agent_role) = session - .services - .agent_control - .get_agent_nickname_and_role(*receiver_thread_id) - .await - .unwrap_or((None, None)); - receiver_agents.push(CollabAgentRef { - thread_id: *receiver_thread_id, - agent_nickname, - agent_role, - }); - } - - // Validate timeout. - // Very short timeouts encourage busy-polling loops in the orchestrator prompt and can - // cause high CPU usage even with a single active worker, so clamp to a minimum. - let timeout_ms = args.timeout_ms.unwrap_or(DEFAULT_WAIT_TIMEOUT_MS); - let timeout_ms = match timeout_ms { - ms if ms <= 0 => { - return Err(FunctionCallError::RespondToModel( - "timeout_ms must be greater than zero".to_owned(), - )); - } - ms => ms.clamp(MIN_WAIT_TIMEOUT_MS, MAX_WAIT_TIMEOUT_MS), - }; - - session - .send_event( - &turn, - CollabWaitingBeginEvent { - sender_thread_id: session.conversation_id, - receiver_thread_ids: receiver_thread_ids.clone(), - receiver_agents: receiver_agents.clone(), - 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(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 { - // Wait for the first agent to reach a final status. - let mut futures = FuturesUnordered::new(); - for (id, rx) in status_rxs.into_iter() { - 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() { - // Drain the unlikely last elements to prevent race. - loop { - match futures.next().now_or_never() { - Some(Some(Some(result))) => results.push(result), - Some(Some(None)) => continue, - Some(None) | None => break, - } - } - } - results - }; - - // Convert payload. - let statuses_map = statuses.clone().into_iter().collect::>(); - let agent_statuses = build_wait_agent_statuses(&statuses_map, &receiver_agents); - let result = WaitResult { - status: statuses_map.clone(), - timed_out: statuses.is_empty(), - }; - - // Final event emission. - session - .send_event( - &turn, - CollabWaitingEndEvent { - sender_thread_id: session.conversation_id, - call_id, - agent_statuses, - statuses: statuses_map, - } - .into(), - ) - .await; - - let content = serde_json::to_string(&result).map_err(|err| { - FunctionCallError::Fatal(format!("failed to serialize wait result: {err}")) - })?; - - Ok(FunctionToolOutput::from_text(content, None)) - } - async fn wait_for_final_status( session: Arc, thread_id: ThreadId, @@ -682,97 +728,117 @@ pub(crate) mod wait { pub mod close_agent { use super::*; - use std::sync::Arc; + + pub(crate) struct Handler; + + #[async_trait] + impl ToolHandler for Handler { + type Output = FunctionToolOutput; + + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + fn matches_kind(&self, payload: &ToolPayload) -> bool { + matches!(payload, ToolPayload::Function { .. }) + } + + async fn handle( + &self, + invocation: ToolInvocation, + ) -> Result { + let ToolInvocation { + session, + turn, + payload, + call_id, + .. + } = invocation; + let arguments = function_arguments(payload)?; + let args: CloseAgentArgs = parse_arguments(&arguments)?; + let agent_id = agent_id(&args.id)?; + let (receiver_agent_nickname, receiver_agent_role) = session + .services + .agent_control + .get_agent_nickname_and_role(agent_id) + .await + .unwrap_or((None, None)); + session + .send_event( + &turn, + CollabCloseBeginEvent { + call_id: call_id.clone(), + sender_thread_id: session.conversation_id, + receiver_thread_id: agent_id, + } + .into(), + ) + .await; + let status = match session + .services + .agent_control + .subscribe_status(agent_id) + .await + { + Ok(mut status_rx) => status_rx.borrow_and_update().clone(), + Err(err) => { + let status = session.services.agent_control.get_status(agent_id).await; + session + .send_event( + &turn, + CollabCloseEndEvent { + call_id: call_id.clone(), + sender_thread_id: session.conversation_id, + receiver_thread_id: agent_id, + receiver_agent_nickname: receiver_agent_nickname.clone(), + receiver_agent_role: receiver_agent_role.clone(), + status, + } + .into(), + ) + .await; + return Err(collab_agent_error(agent_id, err)); + } + }; + let result = if !matches!(status, AgentStatus::Shutdown) { + session + .services + .agent_control + .shutdown_agent(agent_id) + .await + .map_err(|err| collab_agent_error(agent_id, err)) + .map(|_| ()) + } else { + Ok(()) + }; + session + .send_event( + &turn, + CollabCloseEndEvent { + call_id, + sender_thread_id: session.conversation_id, + receiver_thread_id: agent_id, + receiver_agent_nickname, + receiver_agent_role, + status: status.clone(), + } + .into(), + ) + .await; + result?; + + let content = serde_json::to_string(&CloseAgentResult { status }).map_err(|err| { + FunctionCallError::Fatal(format!("failed to serialize close_agent result: {err}")) + })?; + + Ok(FunctionToolOutput::from_text(content, Some(true))) + } + } #[derive(Debug, Deserialize, Serialize)] pub(super) struct CloseAgentResult { pub(super) status: AgentStatus, } - - pub async fn handle( - session: Arc, - turn: Arc, - call_id: String, - arguments: String, - ) -> Result { - let args: CloseAgentArgs = parse_arguments(&arguments)?; - let agent_id = agent_id(&args.id)?; - let (receiver_agent_nickname, receiver_agent_role) = session - .services - .agent_control - .get_agent_nickname_and_role(agent_id) - .await - .unwrap_or((None, None)); - session - .send_event( - &turn, - CollabCloseBeginEvent { - call_id: call_id.clone(), - sender_thread_id: session.conversation_id, - receiver_thread_id: agent_id, - } - .into(), - ) - .await; - let status = match session - .services - .agent_control - .subscribe_status(agent_id) - .await - { - Ok(mut status_rx) => status_rx.borrow_and_update().clone(), - Err(err) => { - let status = session.services.agent_control.get_status(agent_id).await; - session - .send_event( - &turn, - CollabCloseEndEvent { - call_id: call_id.clone(), - sender_thread_id: session.conversation_id, - receiver_thread_id: agent_id, - receiver_agent_nickname: receiver_agent_nickname.clone(), - receiver_agent_role: receiver_agent_role.clone(), - status, - } - .into(), - ) - .await; - return Err(collab_agent_error(agent_id, err)); - } - }; - let result = if !matches!(status, AgentStatus::Shutdown) { - session - .services - .agent_control - .shutdown_agent(agent_id) - .await - .map_err(|err| collab_agent_error(agent_id, err)) - .map(|_| ()) - } else { - Ok(()) - }; - session - .send_event( - &turn, - CollabCloseEndEvent { - call_id, - sender_thread_id: session.conversation_id, - receiver_thread_id: agent_id, - receiver_agent_nickname, - receiver_agent_role, - status: status.clone(), - } - .into(), - ) - .await; - result?; - - let content = serde_json::to_string(&CloseAgentResult { status }).map_err(|err| { - FunctionCallError::Fatal(format!("failed to serialize close_agent result: {err}")) - })?; - - Ok(FunctionToolOutput::from_text(content, Some(true))) - } } fn agent_id(id: &str) -> Result { 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 1aee1fbf1..0f3674de8 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -78,7 +78,7 @@ async fn handler_rejects_non_function_payloads() { input: "hello".to_string(), }, ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = SpawnAgentHandler.handle(invocation).await else { panic!("payload should be rejected"); }; assert_eq!( @@ -89,24 +89,6 @@ async fn handler_rejects_non_function_payloads() { ); } -#[tokio::test] -async fn handler_rejects_unknown_tool() { - let (session, turn) = make_session_and_context().await; - let invocation = invocation( - Arc::new(session), - Arc::new(turn), - "unknown_tool", - function_payload(json!({})), - ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { - panic!("tool should be rejected"); - }; - assert_eq!( - err, - FunctionCallError::RespondToModel("unsupported collab tool unknown_tool".to_string()) - ); -} - #[tokio::test] async fn spawn_agent_rejects_empty_message() { let (session, turn) = make_session_and_context().await; @@ -116,7 +98,7 @@ async fn spawn_agent_rejects_empty_message() { "spawn_agent", function_payload(json!({"message": " "})), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = SpawnAgentHandler.handle(invocation).await else { panic!("empty message should be rejected"); }; assert_eq!( @@ -137,7 +119,7 @@ async fn spawn_agent_rejects_when_message_and_items_are_both_set() { "items": [{"type": "mention", "name": "drive", "path": "app://drive"}] })), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = SpawnAgentHandler.handle(invocation).await else { panic!("message+items should be rejected"); }; assert_eq!( @@ -183,7 +165,7 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() { "agent_type": "explorer" })), ); - let output = MultiAgentHandler + let output = SpawnAgentHandler .handle(invocation) .await .expect("spawn_agent should succeed"); @@ -216,7 +198,7 @@ async fn spawn_agent_errors_when_manager_dropped() { "spawn_agent", function_payload(json!({"message": "hello"})), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = SpawnAgentHandler.handle(invocation).await else { panic!("spawn should fail without a manager"); }; assert_eq!( @@ -276,7 +258,7 @@ async fn spawn_agent_reapplies_runtime_sandbox_after_role_config() { "agent_type": "explorer" })), ); - let output = MultiAgentHandler + let output = SpawnAgentHandler .handle(invocation) .await .expect("spawn_agent should succeed"); @@ -321,7 +303,7 @@ async fn spawn_agent_rejects_when_depth_limit_exceeded() { "spawn_agent", function_payload(json!({"message": "hello"})), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = SpawnAgentHandler.handle(invocation).await else { panic!("spawn should fail when depth limit exceeded"); }; assert_eq!( @@ -360,7 +342,7 @@ async fn spawn_agent_allows_depth_up_to_configured_max_depth() { "spawn_agent", function_payload(json!({"message": "hello"})), ); - let output = MultiAgentHandler + let output = SpawnAgentHandler .handle(invocation) .await .expect("spawn should succeed within configured depth"); @@ -386,7 +368,7 @@ async fn send_input_rejects_empty_message() { "send_input", function_payload(json!({"id": ThreadId::new().to_string(), "message": ""})), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = SendInputHandler.handle(invocation).await else { panic!("empty message should be rejected"); }; assert_eq!( @@ -408,7 +390,7 @@ async fn send_input_rejects_when_message_and_items_are_both_set() { "items": [{"type": "mention", "name": "drive", "path": "app://drive"}] })), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = SendInputHandler.handle(invocation).await else { panic!("message+items should be rejected"); }; assert_eq!( @@ -428,7 +410,7 @@ async fn send_input_rejects_invalid_id() { "send_input", function_payload(json!({"id": "not-a-uuid", "message": "hi"})), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = SendInputHandler.handle(invocation).await else { panic!("invalid id should be rejected"); }; let FunctionCallError::RespondToModel(msg) = err else { @@ -449,7 +431,7 @@ async fn send_input_reports_missing_agent() { "send_input", function_payload(json!({"id": agent_id.to_string(), "message": "hi"})), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = SendInputHandler.handle(invocation).await else { panic!("missing agent should be reported"); }; assert_eq!( @@ -476,7 +458,7 @@ async fn send_input_interrupts_before_prompt() { "interrupt": true })), ); - MultiAgentHandler + SendInputHandler .handle(invocation) .await .expect("send_input should succeed"); @@ -517,7 +499,7 @@ async fn send_input_accepts_structured_items() { ] })), ); - MultiAgentHandler + SendInputHandler .handle(invocation) .await .expect("send_input should succeed"); @@ -557,7 +539,7 @@ async fn resume_agent_rejects_invalid_id() { "resume_agent", function_payload(json!({"id": "not-a-uuid"})), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = ResumeAgentHandler.handle(invocation).await else { panic!("invalid id should be rejected"); }; let FunctionCallError::RespondToModel(msg) = err else { @@ -578,7 +560,7 @@ async fn resume_agent_reports_missing_agent() { "resume_agent", function_payload(json!({"id": agent_id.to_string()})), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = ResumeAgentHandler.handle(invocation).await else { panic!("missing agent should be reported"); }; assert_eq!( @@ -603,7 +585,7 @@ async fn resume_agent_noops_for_active_agent() { function_payload(json!({"id": agent_id.to_string()})), ); - let output = MultiAgentHandler + let output = ResumeAgentHandler .handle(invocation) .await .expect("resume_agent should succeed"); @@ -666,7 +648,7 @@ async fn resume_agent_restores_closed_agent_and_accepts_send_input() { "resume_agent", function_payload(json!({"id": agent_id.to_string()})), ); - let output = MultiAgentHandler + let output = ResumeAgentHandler .handle(resume_invocation) .await .expect("resume_agent should succeed"); @@ -682,7 +664,7 @@ async fn resume_agent_restores_closed_agent_and_accepts_send_input() { "send_input", function_payload(json!({"id": agent_id.to_string(), "message": "hello"})), ); - let output = MultiAgentHandler + let output = SendInputHandler .handle(send_invocation) .await .expect("send_input should succeed after resume"); @@ -723,7 +705,7 @@ async fn resume_agent_rejects_when_depth_limit_exceeded() { "resume_agent", function_payload(json!({"id": ThreadId::new().to_string()})), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = ResumeAgentHandler.handle(invocation).await else { panic!("resume should fail when depth limit exceeded"); }; assert_eq!( @@ -746,7 +728,7 @@ async fn wait_rejects_non_positive_timeout() { "timeout_ms": 0 })), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = WaitHandler.handle(invocation).await else { panic!("non-positive timeout should be rejected"); }; assert_eq!( @@ -764,7 +746,7 @@ async fn wait_rejects_invalid_id() { "wait", function_payload(json!({"ids": ["invalid"]})), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = WaitHandler.handle(invocation).await else { panic!("invalid id should be rejected"); }; let FunctionCallError::RespondToModel(msg) = err else { @@ -782,7 +764,7 @@ async fn wait_rejects_empty_ids() { "wait", function_payload(json!({"ids": []})), ); - let Err(err) = MultiAgentHandler.handle(invocation).await else { + let Err(err) = WaitHandler.handle(invocation).await else { panic!("empty ids should be rejected"); }; assert_eq!( @@ -807,7 +789,7 @@ async fn wait_returns_not_found_for_missing_agents() { "timeout_ms": 1000 })), ); - let output = MultiAgentHandler + let output = WaitHandler .handle(invocation) .await .expect("wait should succeed"); @@ -841,7 +823,7 @@ async fn wait_times_out_when_status_is_not_final() { "timeout_ms": MIN_WAIT_TIMEOUT_MS })), ); - let output = MultiAgentHandler + let output = WaitHandler .handle(invocation) .await .expect("wait should succeed"); @@ -882,11 +864,7 @@ async fn wait_clamps_short_timeouts_to_minimum() { })), ); - let early = timeout( - Duration::from_millis(50), - MultiAgentHandler.handle(invocation), - ) - .await; + let early = timeout(Duration::from_millis(50), WaitHandler.handle(invocation)).await; assert!( early.is_err(), "wait should not return before the minimum timeout clamp" @@ -931,7 +909,7 @@ async fn wait_returns_final_status_without_timeout() { "timeout_ms": 1000 })), ); - let output = MultiAgentHandler + let output = WaitHandler .handle(invocation) .await .expect("wait should succeed"); @@ -964,7 +942,7 @@ async fn close_agent_submits_shutdown_and_returns_status() { "close_agent", function_payload(json!({"id": agent_id.to_string()})), ); - let output = MultiAgentHandler + let output = CloseAgentHandler .handle(invocation) .await .expect("close_agent should succeed"); diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index b61856275..47647084c 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -2335,7 +2335,6 @@ pub(crate) fn build_specs_with_discoverable_tools( use crate::tools::handlers::ListDirHandler; use crate::tools::handlers::McpHandler; use crate::tools::handlers::McpResourceHandler; - use crate::tools::handlers::MultiAgentHandler; use crate::tools::handlers::PlanHandler; use crate::tools::handlers::ReadFileHandler; use crate::tools::handlers::RequestPermissionsHandler; @@ -2347,6 +2346,11 @@ pub(crate) fn build_specs_with_discoverable_tools( use crate::tools::handlers::ToolSuggestHandler; use crate::tools::handlers::UnifiedExecHandler; use crate::tools::handlers::ViewImageHandler; + use crate::tools::handlers::multi_agents::CloseAgentHandler; + use crate::tools::handlers::multi_agents::ResumeAgentHandler; + use crate::tools::handlers::multi_agents::SendInputHandler; + use crate::tools::handlers::multi_agents::SpawnAgentHandler; + use crate::tools::handlers::multi_agents::WaitHandler; use std::sync::Arc; let mut builder = ToolRegistryBuilder::new(); @@ -2715,7 +2719,6 @@ pub(crate) fn build_specs_with_discoverable_tools( } if config.collab_tools { - let multi_agent_handler = Arc::new(MultiAgentHandler); push_tool_spec( &mut builder, create_spawn_agent_tool(config), @@ -2746,11 +2749,11 @@ pub(crate) fn build_specs_with_discoverable_tools( false, config.code_mode_enabled, ); - builder.register_handler("spawn_agent", multi_agent_handler.clone()); - builder.register_handler("send_input", multi_agent_handler.clone()); - builder.register_handler("resume_agent", multi_agent_handler.clone()); - builder.register_handler("wait", multi_agent_handler.clone()); - builder.register_handler("close_agent", multi_agent_handler); + builder.register_handler("spawn_agent", Arc::new(SpawnAgentHandler)); + builder.register_handler("send_input", Arc::new(SendInputHandler)); + builder.register_handler("resume_agent", Arc::new(ResumeAgentHandler)); + builder.register_handler("wait", Arc::new(WaitHandler)); + builder.register_handler("close_agent", Arc::new(CloseAgentHandler)); } if config.agent_jobs_tools {