From 527244910fb851cea6147334dbc08f8fbce4cb9d Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 23 Mar 2026 23:27:55 +0000 Subject: [PATCH] feat: custom watcher for multi-agent v2 (#15576) The new wait tool just returns `Wait timed out.` or `Wait completed.`. The actual content is done through the notification watcher --- .../src/tools/handlers/multi_agents_tests.rs | 71 ++++++++++++++++--- .../tools/handlers/multi_agents_v2/wait.rs | 38 +++++----- codex-rs/core/src/tools/spec.rs | 65 +++++++++++------ codex-rs/core/src/tools/spec_tests.rs | 6 +- 4 files changed, 123 insertions(+), 57 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index b780986d9..873b55d66 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -9,12 +9,14 @@ use crate::config::types::ShellEnvironmentPolicy; use crate::function_tool::FunctionCallError; use crate::protocol::AgentStatus; use crate::protocol::AskForApproval; +use crate::protocol::EventMsg; use crate::protocol::FileSystemSandboxPolicy; use crate::protocol::NetworkSandboxPolicy; use crate::protocol::Op; use crate::protocol::SandboxPolicy; use crate::protocol::SessionSource; use crate::protocol::SubAgentSource; +use crate::protocol::TurnCompleteEvent; use crate::state::TaskKind; use crate::tasks::SessionTask; use crate::tasks::SessionTaskContext; @@ -1414,7 +1416,7 @@ async fn multi_agent_v2_wait_agent_accepts_targets_argument() { assert_eq!( result, crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult { - status: HashMap::from([(target, AgentStatus::NotFound)]), + message: "Wait completed.".to_string(), timed_out: false, } ); @@ -1582,12 +1584,7 @@ async fn wait_agent_returns_final_status_without_timeout() { } #[tokio::test] -async fn multi_agent_v2_wait_agent_returns_statuses_keyed_by_path() { - #[derive(Debug, Deserialize)] - struct SpawnAgentResult { - task_name: String, - } - +async fn multi_agent_v2_wait_agent_returns_summary_for_named_targets() { let (mut session, mut turn) = make_session_and_context().await; let manager = thread_manager(); let root = manager @@ -1617,9 +1614,7 @@ async fn multi_agent_v2_wait_agent_returns_statuses_keyed_by_path() { )) .await .expect("spawn_agent should succeed"); - let (content, _) = expect_text_output(spawn_output); - let spawn_result: SpawnAgentResult = - serde_json::from_str(&content).expect("spawn result should parse"); + let _ = expect_text_output(spawn_output); let agent_id = session .services @@ -1667,13 +1662,67 @@ async fn multi_agent_v2_wait_agent_returns_statuses_keyed_by_path() { assert_eq!( result, crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult { - status: HashMap::from([(spawn_result.task_name, AgentStatus::Shutdown)]), + message: "Wait completed.".to_string(), timed_out: false, } ); assert_eq!(success, None); } +#[tokio::test] +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(); + session.services.agent_control = manager.agent_control(); + let mut config = (*turn.config).clone(); + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + turn.config = Arc::new(config.clone()); + + 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 + .handle(invocation( + Arc::new(session), + Arc::new(turn), + "wait_agent", + function_payload(json!({ + "targets": [agent_id.to_string()], + "timeout_ms": 1000 + })), + )) + .await + .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!( + result, + crate::tools::handlers::multi_agents_v2::wait::WaitAgentResult { + message: "Wait completed.".to_string(), + timed_out: false, + } + ); + assert!(!content.contains("sensitive child output")); + assert_eq!(success, None); +} + #[tokio::test] async fn close_agent_submits_shutdown_and_returns_previous_status() { let (mut session, turn) = make_session_and_context().await; 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 e4bc7d3f4..1b1dc90e4 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 @@ -35,21 +35,12 @@ impl ToolHandler for Handler { 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()); - let mut target_by_thread_id = HashMap::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(); - target_by_thread_id.insert( - *receiver_thread_id, - agent_metadata - .agent_path - .as_ref() - .map(ToString::to_string) - .unwrap_or_else(|| receiver_thread_id.to_string()), - ); receiver_agents.push(CollabAgentRef { thread_id: *receiver_thread_id, agent_nickname: agent_metadata.agent_nickname, @@ -152,18 +143,7 @@ impl ToolHandler for Handler { 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 result = WaitAgentResult { - status: statuses - .into_iter() - .filter_map(|(thread_id, status)| { - target_by_thread_id - .get(&thread_id) - .cloned() - .map(|target| (target, status)) - }) - .collect(), - timed_out, - }; + let result = WaitAgentResult::from_timed_out(timed_out); session .send_event( @@ -191,10 +171,24 @@ struct WaitArgs { #[derive(Debug, Deserialize, Serialize, PartialEq, Eq)] pub(crate) struct WaitAgentResult { - pub(crate) status: HashMap, + pub(crate) message: String, pub(crate) timed_out: bool, } +impl WaitAgentResult { + fn from_timed_out(timed_out: bool) -> Self { + let message = if timed_out { + "Wait timed out." + } else { + "Wait completed." + }; + Self { + message: message.to_string(), + timed_out, + } + } +} + impl ToolOutput for WaitAgentResult { fn log_preview(&self) -> String { tool_output_json_text(self, "wait_agent") diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index d8419e1f5..b115e572e 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -178,23 +178,41 @@ fn resume_agent_output_schema() -> JsonValue { }) } -fn wait_output_schema() -> JsonValue { - json!({ - "type": "object", - "properties": { - "status": { - "type": "object", - "description": "Final statuses keyed by canonical task name when available, otherwise by agent id.", - "additionalProperties": agent_status_output_schema() +fn wait_output_schema(multi_agent_v2: bool) -> JsonValue { + if multi_agent_v2 { + json!({ + "type": "object", + "properties": { + "message": { + "type": "string", + "description": "Brief wait summary without the agent's final content." + }, + "timed_out": { + "type": "boolean", + "description": "Whether the wait call returned due to timeout before any agent reached a final status." + } }, - "timed_out": { - "type": "boolean", - "description": "Whether the wait call returned due to timeout before any agent reached a final status." - } - }, - "required": ["status", "timed_out"], - "additionalProperties": false - }) + "required": ["message", "timed_out"], + "additionalProperties": false + }) + } else { + json!({ + "type": "object", + "properties": { + "status": { + "type": "object", + "description": "Final statuses keyed by canonical task name when available, otherwise by agent id.", + "additionalProperties": agent_status_output_schema() + }, + "timed_out": { + "type": "boolean", + "description": "Whether the wait call returned due to timeout before any agent reached a final status." + } + }, + "required": ["status", "timed_out"], + "additionalProperties": false + }) + } } fn close_agent_output_schema() -> JsonValue { @@ -1422,7 +1440,7 @@ fn create_resume_agent_tool() -> ToolSpec { }) } -fn create_wait_agent_tool() -> ToolSpec { +fn create_wait_agent_tool(multi_agent_v2: bool) -> ToolSpec { let mut properties = BTreeMap::new(); properties.insert( "targets".to_string(), @@ -1445,8 +1463,13 @@ fn create_wait_agent_tool() -> ToolSpec { ToolSpec::Function(ResponsesApiTool { name: "wait_agent".to_string(), - description: "Wait for agents to reach a final status. Completed statuses may include the agent's final message. Returns empty status when timed out. Once the agent reaches a final status, a notification message will be received containing the same completed status." - .to_string(), + description: if multi_agent_v2 { + "Wait for agents to reach a final status. Returns a brief wait summary instead of the agent's final content. Returns a timeout summary when no agent reaches a final status before the deadline." + .to_string() + } else { + "Wait for agents to reach a final status. Completed statuses may include the agent's final message. Returns empty status when timed out. Once the agent reaches a final status, a notification message will be received containing the same completed status." + .to_string() + }, strict: false, defer_loading: None, parameters: JsonSchema::Object { @@ -1454,7 +1477,7 @@ fn create_wait_agent_tool() -> ToolSpec { required: Some(vec!["targets".to_string()]), additional_properties: Some(false.into()), }, - output_schema: Some(wait_output_schema()), + output_schema: Some(wait_output_schema(multi_agent_v2)), }) } @@ -3006,7 +3029,7 @@ pub(crate) fn build_specs_with_discoverable_tools( } push_tool_spec( &mut builder, - create_wait_agent_tool(), + create_wait_agent_tool(config.multi_agent_v2), /*supports_parallel_tool_calls*/ false, config.code_mode_enabled, ); diff --git a/codex-rs/core/src/tools/spec_tests.rs b/codex-rs/core/src/tools/spec_tests.rs index 1cb6bb166..b33748dd3 100644 --- a/codex-rs/core/src/tools/spec_tests.rs +++ b/codex-rs/core/src/tools/spec_tests.rs @@ -469,7 +469,7 @@ fn test_full_toolset_specs_for_gpt5_codex_unified_exec_web_search() { create_view_image_tool(config.can_request_original_image_detail), create_spawn_agent_tool(&config), create_send_input_tool(), - create_wait_agent_tool(), + create_wait_agent_tool(config.multi_agent_v2), create_close_agent_tool(), ] { expected.insert(tool_name(&spec).to_string(), spec); @@ -607,8 +607,8 @@ fn test_build_specs_multi_agent_v2_uses_task_names_and_hides_resume() { .as_ref() .expect("wait_agent should define output schema"); assert_eq!( - output_schema["properties"]["status"]["description"], - json!("Final statuses keyed by canonical task name when available, otherwise by agent id.") + output_schema["properties"]["message"]["description"], + json!("Brief wait summary without the agent's final content.") ); assert_lacks_tool_name(&tools, "resume_agent"); }