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

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() {
+async fn multi_agent_v2_wait_agent_returns_summary_for_named_targets() {
-    #[derive(Debug, Deserialize)]
-    struct SpawnAgentResult {
-        task_name: String,
-    }
-
     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 _ = expect_text_output(spawn_output);
-    let spawn_result: SpawnAgentResult =
-        serde_json::from_str(&content).expect("spawn result should parse");
 
     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;

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::<HashMap<_, _>>();
         let agent_statuses = build_wait_agent_statuses(&statuses_by_id, &receiver_agents);
-        let result = WaitAgentResult {
+        let result = WaitAgentResult::from_timed_out(timed_out);
-            status: statuses
-                .into_iter()
-                .filter_map(|(thread_id, status)| {
-                    target_by_thread_id
-                        .get(&thread_id)
-                        .cloned()
-                        .map(|target| (target, status))
-                })
-                .collect(),
-            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<String, AgentStatus>,
+    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")

codex-rs/core/src/tools/spec.rs

@@ -178,23 +178,41 @@ fn resume_agent_output_schema() -> JsonValue {
     })
 }
 
-fn wait_output_schema() -> JsonValue {
+fn wait_output_schema(multi_agent_v2: bool) -> JsonValue {
-    json!({
+    if multi_agent_v2 {
-        "type": "object",
+        json!({
-        "properties": {
+            "type": "object",
-            "status": {
+            "properties": {
-                "type": "object",
+                "message": {
-                "description": "Final statuses keyed by canonical task name when available, otherwise by agent id.",
+                    "type": "string",
-                "additionalProperties": agent_status_output_schema()
+                    "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": {
+            "required": ["message", "timed_out"],
-                "type": "boolean",
+            "additionalProperties": false
-                "description": "Whether the wait call returned due to timeout before any agent reached a final status."
+        })
-            }
+    } else {
-        },
+        json!({
-        "required": ["status", "timed_out"],
+            "type": "object",
-        "additionalProperties": false
+            "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."
+        description: if multi_agent_v2 {
-            .to_string(),
+            "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,
         );

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"],
+        output_schema["properties"]["message"]["description"],
-        json!("Final statuses keyed by canonical task name when available, otherwise by agent id.")
+        json!("Brief wait summary without the agent's final content.")
     );
     assert_lacks_tool_name(&tools, "resume_agent");
 }