From 7626f612748515d6d79e149c2ae37d7d783cf989 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Thu, 12 Mar 2026 18:10:10 -0700 Subject: [PATCH] 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 --- .../core/src/tools/handlers/multi_agents.rs | 181 ++++++++++++++---- .../src/tools/handlers/multi_agents_tests.rs | 34 +++- codex-rs/core/src/tools/spec.rs | 115 ++++++++++- 3 files changed, 278 insertions(+), 52 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/multi_agents.rs b/codex-rs/core/src/tools/handlers/multi_agents.rs index 87b11916a..6f699b20d 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents.rs +++ b/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 } } +fn tool_output_json_text(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( + call_id: &str, + payload: &ToolPayload, + value: &T, + success: Option, + 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(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, } + + 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, 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") + } } } 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 0f3674de8..548a2bc84 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/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) { - ( - codex_protocol::models::function_call_output_content_items_to_text(&output.body) - .unwrap_or_default(), - output.success, - ) +fn expect_text_output(output: T) -> (String, Option) +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] diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 47647084c..28f522847 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/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()), }) }