Add typed multi-agent tool outputs (#14536)

## Summary
- return typed `ToolOutput` values from the multi-agent handlers instead
of plain `FunctionToolOutput`
- keep the regular function-call response shape as JSON text while
exposing structured values to code mode
- add output schemas for `spawn_agent`, `send_input`, `resume_agent`,
`wait`, and `close_agent`

## Verification
- `just fmt`
- focused multi-agent and integration tests passed earlier in this
branch during iteration
- after the final edit, I only reran formatting before opening this PR

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

@@ -16,6 +16,7 @@ use crate::function_tool::FunctionCallError;
 use crate::models_manager::manager::RefreshStrategy;
 use crate::tools::context::FunctionToolOutput;
 use crate::tools::context::ToolInvocation;
+use crate::tools::context::ToolOutput;
 use crate::tools::context::ToolPayload;
 use crate::tools::handlers::parse_arguments;
 use crate::tools::registry::ToolHandler;
@@ -23,6 +24,7 @@ use crate::tools::registry::ToolKind;
 use async_trait::async_trait;
 use codex_protocol::ThreadId;
 use codex_protocol::models::BaseInstructions;
+use codex_protocol::models::ResponseInputItem;
 use codex_protocol::openai_models::ReasoningEffort;
 use codex_protocol::openai_models::ReasoningEffortPreset;
 use codex_protocol::protocol::CollabAgentInteractionBeginEvent;
@@ -42,6 +44,7 @@ use codex_protocol::protocol::SubAgentSource;
 use codex_protocol::user_input::UserInput;
 use serde::Deserialize;
 use serde::Serialize;
+use serde_json::Value as JsonValue;
 use std::collections::HashMap;
 use std::sync::Arc;
 
@@ -70,6 +73,38 @@ fn function_arguments(payload: ToolPayload) -> Result<String, FunctionCallError>
     }
 }
 
+fn tool_output_json_text<T>(value: &T, tool_name: &str) -> String
+where
+    T: Serialize,
+{
+    serde_json::to_string(value).unwrap_or_else(|err| {
+        JsonValue::String(format!("failed to serialize {tool_name} result: {err}")).to_string()
+    })
+}
+
+fn tool_output_response_item<T>(
+    call_id: &str,
+    payload: &ToolPayload,
+    value: &T,
+    success: Option<bool>,
+    tool_name: &str,
+) -> ResponseInputItem
+where
+    T: Serialize,
+{
+    FunctionToolOutput::from_text(tool_output_json_text(value, tool_name), success)
+        .to_response_item(call_id, payload)
+}
+
+fn tool_output_code_mode_result<T>(value: &T, tool_name: &str) -> JsonValue
+where
+    T: Serialize,
+{
+    serde_json::to_value(value).unwrap_or_else(|err| {
+        JsonValue::String(format!("failed to serialize {tool_name} result: {err}"))
+    })
+}
+
 mod spawn {
     use super::*;
     use crate::agent::control::SpawnAgentOptions;
@@ -83,7 +118,7 @@ mod spawn {
 
     #[async_trait]
     impl ToolHandler for Handler {
-        type Output = FunctionToolOutput;
+        type Output = SpawnAgentResult;
 
         fn kind(&self) -> ToolKind {
             ToolKind::Function
@@ -206,15 +241,10 @@ mod spawn {
             turn.session_telemetry
                 .counter("codex.multi_agent.spawn", 1, &[("role", role_tag)]);
 
-            let content = serde_json::to_string(&SpawnAgentResult {
+            Ok(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)))
         }
     }
 
@@ -230,10 +260,28 @@ mod spawn {
     }
 
     #[derive(Debug, Serialize)]
-    struct SpawnAgentResult {
+    pub(crate) struct SpawnAgentResult {
         agent_id: String,
         nickname: Option<String>,
     }
+
+    impl ToolOutput for SpawnAgentResult {
+        fn log_preview(&self) -> String {
+            tool_output_json_text(self, "spawn_agent")
+        }
+
+        fn success_for_logging(&self) -> bool {
+            true
+        }
+
+        fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem {
+            tool_output_response_item(call_id, payload, self, Some(true), "spawn_agent")
+        }
+
+        fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue {
+            tool_output_code_mode_result(self, "spawn_agent")
+        }
+    }
 }
 
 mod send_input {
@@ -243,7 +291,7 @@ mod send_input {
 
     #[async_trait]
     impl ToolHandler for Handler {
-        type Output = FunctionToolOutput;
+        type Output = SendInputResult;
 
         fn kind(&self) -> ToolKind {
             ToolKind::Function
@@ -323,14 +371,7 @@ mod send_input {
                 .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)))
+            Ok(SendInputResult { submission_id })
         }
     }
 
@@ -344,9 +385,27 @@ mod send_input {
     }
 
     #[derive(Debug, Serialize)]
-    struct SendInputResult {
+    pub(crate) struct SendInputResult {
         submission_id: String,
     }
+
+    impl ToolOutput for SendInputResult {
+        fn log_preview(&self) -> String {
+            tool_output_json_text(self, "send_input")
+        }
+
+        fn success_for_logging(&self) -> bool {
+            true
+        }
+
+        fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem {
+            tool_output_response_item(call_id, payload, self, Some(true), "send_input")
+        }
+
+        fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue {
+            tool_output_code_mode_result(self, "send_input")
+        }
+    }
 }
 
 mod resume_agent {
@@ -357,7 +416,7 @@ mod resume_agent {
 
     #[async_trait]
     impl ToolHandler for Handler {
-        type Output = FunctionToolOutput;
+        type Output = ResumeAgentResult;
 
         fn kind(&self) -> ToolKind {
             ToolKind::Function
@@ -462,11 +521,7 @@ mod resume_agent {
             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)))
+            Ok(ResumeAgentResult { status })
         }
     }
 
@@ -476,8 +531,26 @@ mod resume_agent {
     }
 
     #[derive(Debug, Deserialize, Serialize, PartialEq, Eq)]
-    pub(super) struct ResumeAgentResult {
-        pub(super) status: AgentStatus,
+    pub(crate) struct ResumeAgentResult {
+        pub(crate) status: AgentStatus,
+    }
+
+    impl ToolOutput for ResumeAgentResult {
+        fn log_preview(&self) -> String {
+            tool_output_json_text(self, "resume_agent")
+        }
+
+        fn success_for_logging(&self) -> bool {
+            true
+        }
+
+        fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem {
+            tool_output_response_item(call_id, payload, self, Some(true), "resume_agent")
+        }
+
+        fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue {
+            tool_output_code_mode_result(self, "resume_agent")
+        }
     }
 
     async fn try_resume_closed_agent(
@@ -523,7 +596,7 @@ pub(crate) mod wait {
 
     #[async_trait]
     impl ToolHandler for Handler {
-        type Output = FunctionToolOutput;
+        type Output = WaitResult;
 
         fn kind(&self) -> ToolKind {
             ToolKind::Function
@@ -683,11 +756,7 @@ pub(crate) mod wait {
                 )
                 .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))
+            Ok(result)
         }
     }
 
@@ -703,6 +772,24 @@ pub(crate) mod wait {
         pub(crate) timed_out: bool,
     }
 
+    impl ToolOutput for WaitResult {
+        fn log_preview(&self) -> String {
+            tool_output_json_text(self, "wait")
+        }
+
+        fn success_for_logging(&self) -> bool {
+            true
+        }
+
+        fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem {
+            tool_output_response_item(call_id, payload, self, None, "wait")
+        }
+
+        fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue {
+            tool_output_code_mode_result(self, "wait")
+        }
+    }
+
     async fn wait_for_final_status(
         session: Arc<Session>,
         thread_id: ThreadId,
@@ -733,7 +820,7 @@ pub mod close_agent {
 
     #[async_trait]
     impl ToolHandler for Handler {
-        type Output = FunctionToolOutput;
+        type Output = CloseAgentResult;
 
         fn kind(&self) -> ToolKind {
             ToolKind::Function
@@ -827,17 +914,31 @@ pub mod close_agent {
                 .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)))
+            Ok(CloseAgentResult { status })
         }
     }
 
     #[derive(Debug, Deserialize, Serialize)]
-    pub(super) struct CloseAgentResult {
-        pub(super) status: AgentStatus,
+    pub(crate) struct CloseAgentResult {
+        pub(crate) status: AgentStatus,
+    }
+
+    impl ToolOutput for CloseAgentResult {
+        fn log_preview(&self) -> String {
+            tool_output_json_text(self, "close_agent")
+        }
+
+        fn success_for_logging(&self) -> bool {
+            true
+        }
+
+        fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem {
+            tool_output_response_item(call_id, payload, self, Some(true), "close_agent")
+        }
+
+        fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue {
+            tool_output_code_mode_result(self, "close_agent")
+        }
     }
 }
 

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

@@ -12,10 +12,12 @@ use crate::protocol::Op;
 use crate::protocol::SandboxPolicy;
 use crate::protocol::SessionSource;
 use crate::protocol::SubAgentSource;
-use crate::tools::context::FunctionToolOutput;
+use crate::tools::context::ToolOutput;
 use crate::turn_diff_tracker::TurnDiffTracker;
 use codex_protocol::ThreadId;
 use codex_protocol::models::ContentItem;
+use codex_protocol::models::FunctionCallOutputBody;
+use codex_protocol::models::ResponseInputItem;
 use codex_protocol::models::ResponseItem;
 use codex_protocol::protocol::InitialHistory;
 use codex_protocol::protocol::RolloutItem;
@@ -59,12 +61,30 @@ fn thread_manager() -> ThreadManager {
     )
 }
 
-fn expect_text_output(output: FunctionToolOutput) -> (String, Option<bool>) {
-    (
-        codex_protocol::models::function_call_output_content_items_to_text(&output.body)
-            .unwrap_or_default(),
-        output.success,
-    )
+fn expect_text_output<T>(output: T) -> (String, Option<bool>)
+where
+    T: ToolOutput,
+{
+    let response = output.to_response_item(
+        "call-1",
+        &ToolPayload::Function {
+            arguments: "{}".to_string(),
+        },
+    );
+    match response {
+        ResponseInputItem::FunctionCallOutput { output, .. }
+        | ResponseInputItem::CustomToolCallOutput { output, .. } => {
+            let content = match output.body {
+                FunctionCallOutputBody::Text(text) => text,
+                FunctionCallOutputBody::ContentItems(items) => {
+                    codex_protocol::models::function_call_output_content_items_to_text(&items)
+                        .unwrap_or_default()
+                }
+            };
+            (content, output.success)
+        }
+        other => panic!("expected function output, got {other:?}"),
+    }
 }
 
 #[tokio::test]

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

@@ -91,6 +91,111 @@ fn unified_exec_output_schema() -> JsonValue {
         "additionalProperties": false
     })
 }
+
+fn agent_status_output_schema() -> JsonValue {
+    json!({
+        "oneOf": [
+            {
+                "type": "string",
+                "enum": ["pending_init", "running", "shutdown", "not_found"]
+            },
+            {
+                "type": "object",
+                "properties": {
+                    "completed": {
+                        "type": ["string", "null"]
+                    }
+                },
+                "required": ["completed"],
+                "additionalProperties": false
+            },
+            {
+                "type": "object",
+                "properties": {
+                    "errored": {
+                        "type": "string"
+                    }
+                },
+                "required": ["errored"],
+                "additionalProperties": false
+            }
+        ]
+    })
+}
+
+fn spawn_agent_output_schema() -> JsonValue {
+    json!({
+        "type": "object",
+        "properties": {
+            "agent_id": {
+                "type": "string",
+                "description": "Thread identifier for the spawned agent."
+            },
+            "nickname": {
+                "type": ["string", "null"],
+                "description": "User-facing nickname for the spawned agent when available."
+            }
+        },
+        "required": ["agent_id", "nickname"],
+        "additionalProperties": false
+    })
+}
+
+fn send_input_output_schema() -> JsonValue {
+    json!({
+        "type": "object",
+        "properties": {
+            "submission_id": {
+                "type": "string",
+                "description": "Identifier for the queued input submission."
+            }
+        },
+        "required": ["submission_id"],
+        "additionalProperties": false
+    })
+}
+
+fn resume_agent_output_schema() -> JsonValue {
+    json!({
+        "type": "object",
+        "properties": {
+            "status": agent_status_output_schema()
+        },
+        "required": ["status"],
+        "additionalProperties": false
+    })
+}
+
+fn wait_output_schema() -> JsonValue {
+    json!({
+        "type": "object",
+        "properties": {
+            "status": {
+                "type": "object",
+                "description": "Final statuses keyed by agent id for agents that finished before the timeout.",
+                "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 {
+    json!({
+        "type": "object",
+        "properties": {
+            "status": agent_status_output_schema()
+        },
+        "required": ["status"],
+        "additionalProperties": false
+    })
+}
+
 #[derive(Debug, Clone, Copy, Eq, PartialEq)]
 pub enum ShellCommandBackendConfig {
     Classic,
@@ -986,7 +1091,7 @@ fn create_spawn_agent_tool(config: &ToolsConfig) -> ToolSpec {
             required: None,
             additional_properties: Some(false.into()),
         },
-        output_schema: None,
+        output_schema: Some(spawn_agent_output_schema()),
     })
 }
 
@@ -1188,7 +1293,7 @@ fn create_send_input_tool() -> ToolSpec {
             required: Some(vec!["id".to_string()]),
             additional_properties: Some(false.into()),
         },
-        output_schema: None,
+        output_schema: Some(send_input_output_schema()),
     })
 }
 
@@ -1213,7 +1318,7 @@ fn create_resume_agent_tool() -> ToolSpec {
             required: Some(vec!["id".to_string()]),
             additional_properties: Some(false.into()),
         },
-        output_schema: None,
+        output_schema: Some(resume_agent_output_schema()),
     })
 }
 
@@ -1249,7 +1354,7 @@ fn create_wait_tool() -> ToolSpec {
             required: Some(vec!["ids".to_string()]),
             additional_properties: Some(false.into()),
         },
-        output_schema: None,
+        output_schema: Some(wait_output_schema()),
     })
 }
 
@@ -1385,7 +1490,7 @@ fn create_close_agent_tool() -> ToolSpec {
             required: Some(vec!["id".to_string()]),
             additional_properties: Some(false.into()),
         },
-        output_schema: None,
+        output_schema: Some(close_agent_output_schema()),
     })
 }