From d71e0426940b75f7dea0c149f2129f0b86c17f20 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Mon, 9 Mar 2026 22:49:44 -0600 Subject: [PATCH] Enforce single tool output type in codex handlers (#14157) We'll need to associate output schema with each tool. Each tool can only have on output type. --- .../tests/suite/v2/dynamic_tools.rs | 6 +- codex-rs/core/src/codex_tests.rs | 13 +- codex-rs/core/src/codex_tests_guardian.rs | 12 +- codex-rs/core/src/tools/context.rs | 120 ++++++++---------- .../core/src/tools/handlers/agent_jobs.rs | 21 ++- .../core/src/tools/handlers/apply_patch.rs | 29 ++--- codex-rs/core/src/tools/handlers/artifacts.rs | 15 ++- codex-rs/core/src/tools/handlers/code_mode.rs | 12 +- codex-rs/core/src/tools/handlers/dynamic.rs | 12 +- .../core/src/tools/handlers/grep_files.rs | 23 ++-- codex-rs/core/src/tools/handlers/js_repl.rs | 30 ++--- codex-rs/core/src/tools/handlers/list_dir.rs | 12 +- codex-rs/core/src/tools/handlers/mcp.rs | 48 +++++-- .../core/src/tools/handlers/mcp_resource.rs | 51 +++----- .../core/src/tools/handlers/multi_agents.rs | 56 +++----- codex-rs/core/src/tools/handlers/plan.rs | 12 +- codex-rs/core/src/tools/handlers/read_file.rs | 15 ++- .../src/tools/handlers/request_permissions.rs | 12 +- .../src/tools/handlers/request_user_input.rs | 12 +- .../src/tools/handlers/search_tool_bm25.rs | 17 +-- codex-rs/core/src/tools/handlers/shell.rs | 18 +-- codex-rs/core/src/tools/handlers/test_sync.rs | 12 +- .../core/src/tools/handlers/unified_exec.rs | 12 +- .../core/src/tools/handlers/view_image.rs | 12 +- codex-rs/core/src/tools/registry.rs | 94 +++++++++++--- codex-rs/core/tests/suite/view_image.rs | 24 +--- 26 files changed, 345 insertions(+), 355 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs b/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs index f3945c2b4..338593c46 100644 --- a/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs +++ b/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs @@ -281,11 +281,7 @@ async fn dynamic_tool_call_round_trip_sends_text_content_items_to_model() -> Res .iter() .find_map(|body| function_call_output_payload(body, call_id)) .context("expected function_call_output in follow-up request")?; - let expected_payload = FunctionCallOutputPayload::from_content_items(vec![ - FunctionCallOutputContentItem::InputText { - text: "dynamic-ok".to_string(), - }, - ]); + let expected_payload = FunctionCallOutputPayload::from_text("dynamic-ok".to_string()); assert_eq!(payload, expected_payload); Ok(()) diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index bf8d970a8..d9a3a508c 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -46,7 +46,7 @@ use crate::state::TaskKind; use crate::tasks::SessionTask; use crate::tasks::SessionTaskContext; use crate::tools::ToolRouter; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::handlers::ShellHandler; @@ -89,11 +89,10 @@ use std::time::Duration as StdDuration; #[path = "codex_tests_guardian.rs"] mod guardian_tests; -fn expect_text_tool_output(output: &dyn std::any::Any) -> String { - let Some(output) = output.downcast_ref::() else { - panic!("unexpected tool output"); - }; - output.text.clone() +use codex_protocol::models::function_call_output_content_items_to_text; + +fn expect_text_tool_output(output: &FunctionToolOutput) -> String { + function_call_output_content_items_to_text(&output.body).unwrap_or_default() } struct InstructionsTestCase { @@ -4144,7 +4143,7 @@ async fn rejects_escalated_permissions_when_policy_not_on_request() { }) .await; - let output = expect_text_tool_output(&*resp2.expect("expected Ok result")); + let output = expect_text_tool_output(&resp2.expect("expected Ok result")); #[derive(Deserialize, PartialEq, Eq, Debug)] struct ResponseExecMetadata { diff --git a/codex-rs/core/src/codex_tests_guardian.rs b/codex-rs/core/src/codex_tests_guardian.rs index 9e3bea6b6..0479f355a 100644 --- a/codex-rs/core/src/codex_tests_guardian.rs +++ b/codex-rs/core/src/codex_tests_guardian.rs @@ -8,7 +8,7 @@ use crate::features::Feature; use crate::guardian::GUARDIAN_SUBAGENT_NAME; use crate::protocol::AskForApproval; use crate::sandboxing::SandboxPermissions; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::turn_diff_tracker::TurnDiffTracker; use codex_app_server_protocol::ConfigLayerSource; use codex_execpolicy::Decision; @@ -16,6 +16,7 @@ use codex_execpolicy::Evaluation; use codex_execpolicy::RuleMatch; use codex_protocol::models::NetworkPermissions; use codex_protocol::models::PermissionProfile; +use codex_protocol::models::function_call_output_content_items_to_text; use codex_protocol::permissions::FileSystemSandboxPolicy; use codex_protocol::permissions::NetworkSandboxPolicy; use codex_utils_absolute_path::AbsolutePathBuf; @@ -33,11 +34,8 @@ use std::fs; use std::sync::Arc; use tempfile::tempdir; -fn expect_text_output(output: &dyn std::any::Any) -> String { - let Some(output) = output.downcast_ref::() else { - panic!("unexpected tool output"); - }; - output.text.clone() +fn expect_text_output(output: &FunctionToolOutput) -> String { + function_call_output_content_items_to_text(&output.body).unwrap_or_default() } #[tokio::test] @@ -159,7 +157,7 @@ async fn guardian_allows_shell_additional_permissions_requests_past_policy_valid }) .await; - let output = expect_text_output(&*resp.expect("expected Ok result")); + let output = expect_text_output(&resp.expect("expected Ok result")); #[derive(Deserialize, PartialEq, Eq, Debug)] struct ResponseExecMetadata { diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 0bfebfa64..41d11e6b3 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -10,8 +10,8 @@ use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ShellToolCallParams; +use codex_protocol::models::function_call_output_content_items_to_text; use codex_utils_string::take_bytes_at_char_boundary; -use std::any::Any; use std::borrow::Cow; use std::sync::Arc; use tokio::sync::Mutex; @@ -64,16 +64,14 @@ impl ToolPayload { } } -pub trait ToolOutput: Any + Send { +pub trait ToolOutput: Send { fn log_preview(&self) -> String; fn success_for_logging(&self) -> bool; - fn into_response(self: Box, call_id: &str, payload: &ToolPayload) -> ResponseInputItem; + fn into_response(self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem; } -pub type ToolOutputBox = Box; - pub struct McpToolOutput { pub result: Result, } @@ -87,8 +85,8 @@ impl ToolOutput for McpToolOutput { self.result.is_ok() } - fn into_response(self: Box, call_id: &str, _payload: &ToolPayload) -> ResponseInputItem { - let Self { result } = *self; + fn into_response(self, call_id: &str, _payload: &ToolPayload) -> ResponseInputItem { + let Self { result } = self; ResponseInputItem::McpToolCallOutput { call_id: call_id.to_string(), result, @@ -96,42 +94,34 @@ impl ToolOutput for McpToolOutput { } } -pub struct TextToolOutput { - pub text: String, +pub struct FunctionToolOutput { + pub body: Vec, pub success: Option, } -impl ToolOutput for TextToolOutput { - fn log_preview(&self) -> String { - telemetry_preview(&self.text) - } - - fn success_for_logging(&self) -> bool { - self.success.unwrap_or(true) - } - - fn into_response(self: Box, call_id: &str, payload: &ToolPayload) -> ResponseInputItem { - let Self { text, success } = *self; - function_tool_response( - call_id, - payload, - FunctionCallOutputBody::Text(text), +impl FunctionToolOutput { + pub fn from_text(text: String, success: Option) -> Self { + Self { + body: vec![FunctionCallOutputContentItem::InputText { text }], success, - ) + } + } + + pub fn from_content( + content: Vec, + success: Option, + ) -> Self { + Self { + body: content, + success, + } } } -pub struct ContentToolOutput { - pub content: Vec, - pub success: Option, -} - -impl ToolOutput for ContentToolOutput { +impl ToolOutput for FunctionToolOutput { fn log_preview(&self) -> String { telemetry_preview( - &FunctionCallOutputBody::ContentItems(self.content.clone()) - .to_text() - .unwrap_or_default(), + &function_call_output_content_items_to_text(&self.body).unwrap_or_default(), ) } @@ -139,23 +129,25 @@ impl ToolOutput for ContentToolOutput { self.success.unwrap_or(true) } - fn into_response(self: Box, call_id: &str, payload: &ToolPayload) -> ResponseInputItem { - let Self { content, success } = *self; - function_tool_response( - call_id, - payload, - FunctionCallOutputBody::ContentItems(content), - success, - ) + fn into_response(self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem { + let Self { body, success } = self; + function_tool_response(call_id, payload, body, success) } } fn function_tool_response( call_id: &str, payload: &ToolPayload, - body: FunctionCallOutputBody, + body: Vec, success: Option, ) -> ResponseInputItem { + let body = match body.as_slice() { + [FunctionCallOutputContentItem::InputText { text }] => { + FunctionCallOutputBody::Text(text.clone()) + } + _ => FunctionCallOutputBody::ContentItems(body), + }; + if matches!(payload, ToolPayload::Custom { .. }) { return ResponseInputItem::CustomToolCallOutput { call_id: call_id.to_string(), @@ -219,17 +211,14 @@ mod tests { let payload = ToolPayload::Custom { input: "patch".to_string(), }; - let response = Box::new(TextToolOutput { - text: "patched".to_string(), - success: Some(true), - }) - .into_response("call-42", &payload); + let response = FunctionToolOutput::from_text("patched".to_string(), Some(true)) + .into_response("call-42", &payload); match response { ResponseInputItem::CustomToolCallOutput { call_id, output } => { assert_eq!(call_id, "call-42"); - assert_eq!(output.text_content(), Some("patched")); - assert!(output.content_items().is_none()); + assert_eq!(output.content_items(), None); + assert_eq!(output.body.to_text().as_deref(), Some("patched")); assert_eq!(output.success, Some(true)); } other => panic!("expected CustomToolCallOutput, got {other:?}"), @@ -241,17 +230,14 @@ mod tests { let payload = ToolPayload::Function { arguments: "{}".to_string(), }; - let response = Box::new(TextToolOutput { - text: "ok".to_string(), - success: Some(true), - }) - .into_response("fn-1", &payload); + let response = FunctionToolOutput::from_text("ok".to_string(), Some(true)) + .into_response("fn-1", &payload); match response { ResponseInputItem::FunctionCallOutput { call_id, output } => { assert_eq!(call_id, "fn-1"); - assert_eq!(output.text_content(), Some("ok")); - assert!(output.content_items().is_none()); + assert_eq!(output.content_items(), None); + assert_eq!(output.body.to_text().as_deref(), Some("ok")); assert_eq!(output.success, Some(true)); } other => panic!("expected FunctionCallOutput, got {other:?}"), @@ -263,8 +249,8 @@ mod tests { let payload = ToolPayload::Custom { input: "patch".to_string(), }; - let response = Box::new(ContentToolOutput { - content: vec![ + let response = FunctionToolOutput::from_content( + vec![ FunctionCallOutputContentItem::InputText { text: "line 1".to_string(), }, @@ -276,8 +262,8 @@ mod tests { text: "line 2".to_string(), }, ], - success: Some(true), - }) + Some(true), + ) .into_response("call-99", &payload); match response { @@ -305,14 +291,18 @@ mod tests { #[test] fn log_preview_uses_content_items_when_plain_text_is_missing() { - let output = ContentToolOutput { - content: vec![FunctionCallOutputContentItem::InputText { + let output = FunctionToolOutput::from_content( + vec![FunctionCallOutputContentItem::InputText { text: "preview".to_string(), }], - success: Some(true), - }; + Some(true), + ); assert_eq!(output.log_preview(), "preview"); + assert_eq!( + function_call_output_content_items_to_text(&output.body), + Some("preview".to_string()) + ); } #[test] diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index a3b747d50..ff02a3fbd 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -6,9 +6,8 @@ use crate::codex::TurnContext; use crate::config::Config; use crate::error::CodexErr; use crate::function_tool::FunctionCallError; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::handlers::multi_agents::build_agent_spawn_config; use crate::tools::handlers::parse_arguments; @@ -175,6 +174,8 @@ impl JobProgressEmitter { #[async_trait] impl ToolHandler for BatchJobHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } @@ -183,7 +184,7 @@ impl ToolHandler for BatchJobHandler { matches!(payload, ToolPayload::Function { .. }) } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -223,7 +224,7 @@ mod spawn_agents_on_csv { session: Arc, turn: Arc, arguments: String, - ) -> Result { + ) -> Result { let args: SpawnAgentsOnCsvArgs = parse_arguments(arguments.as_str())?; if args.instruction.trim().is_empty() { return Err(FunctionCallError::RespondToModel( @@ -456,10 +457,7 @@ mod spawn_agents_on_csv { "failed to serialize spawn_agents_on_csv result: {err}" )) })?; - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } } @@ -469,7 +467,7 @@ mod report_agent_job_result { pub async fn handle( session: Arc, arguments: String, - ) -> Result { + ) -> Result { let args: ReportAgentJobResultArgs = parse_arguments(arguments.as_str())?; if !args.result.is_object() { return Err(FunctionCallError::RespondToModel( @@ -505,10 +503,7 @@ mod report_agent_job_result { "failed to serialize report_agent_job_result result: {err}" )) })?; - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } } diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index aeca890c7..e31a47ab1 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -11,10 +11,9 @@ use crate::client_common::tools::ToolSpec; use crate::codex::Session; use crate::codex::TurnContext; use crate::function_tool::FunctionCallError; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::SharedTurnDiffTracker; -use crate::tools::context::TextToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::events::ToolEmitter; use crate::tools::events::ToolEventCtx; @@ -92,6 +91,8 @@ fn write_permissions_for_paths(file_paths: &[AbsolutePathBuf]) -> Option ToolKind { ToolKind::Function } @@ -107,7 +108,7 @@ impl ToolHandler for ApplyPatchHandler { true } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -140,10 +141,7 @@ impl ToolHandler for ApplyPatchHandler { match apply_patch::apply_patch(turn.as_ref(), changes).await { InternalApplyPatchInvocation::Output(item) => { let content = item?; - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } InternalApplyPatchInvocation::DelegateToExec(apply) => { let changes = convert_apply_patch_to_protocol(&apply.action); @@ -204,10 +202,7 @@ impl ToolHandler for ApplyPatchHandler { Some(&tracker), ); let content = emitter.finish(event_ctx, out).await?; - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } } } @@ -241,7 +236,7 @@ pub(crate) async fn intercept_apply_patch( tracker: Option<&SharedTurnDiffTracker>, call_id: &str, tool_name: &str, -) -> Result, FunctionCallError> { +) -> Result, FunctionCallError> { match codex_apply_patch::maybe_parse_apply_patch_verified(command, cwd) { codex_apply_patch::MaybeApplyPatchVerified::Body(changes) => { session @@ -255,10 +250,7 @@ pub(crate) async fn intercept_apply_patch( match apply_patch::apply_patch(turn.as_ref(), changes).await { InternalApplyPatchInvocation::Output(item) => { let content = item?; - Ok(Some(Box::new(TextToolOutput { - text: content, - success: Some(true), - }))) + Ok(Some(FunctionToolOutput::from_text(content, Some(true)))) } InternalApplyPatchInvocation::DelegateToExec(apply) => { let changes = convert_apply_patch_to_protocol(&apply.action); @@ -317,10 +309,7 @@ pub(crate) async fn intercept_apply_patch( tracker.as_ref().copied(), ); let content = emitter.finish(event_ctx, out).await?; - Ok(Some(Box::new(TextToolOutput { - text: content, - success: Some(true), - }))) + Ok(Some(FunctionToolOutput::from_text(content, Some(true)))) } } } diff --git a/codex-rs/core/src/tools/handlers/artifacts.rs b/codex-rs/core/src/tools/handlers/artifacts.rs index 38f228291..df239495e 100644 --- a/codex-rs/core/src/tools/handlers/artifacts.rs +++ b/codex-rs/core/src/tools/handlers/artifacts.rs @@ -16,9 +16,8 @@ use crate::exec::StreamOutput; use crate::features::Feature; use crate::function_tool::FunctionCallError; use crate::protocol::ExecCommandSource; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::events::ToolEmitter; use crate::tools::events::ToolEventCtx; @@ -42,6 +41,8 @@ struct ArtifactsToolArgs { #[async_trait] impl ToolHandler for ArtifactsHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } @@ -54,7 +55,7 @@ impl ToolHandler for ArtifactsHandler { true } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -112,10 +113,10 @@ impl ToolHandler for ArtifactsHandler { ) .await; - Ok(Box::new(TextToolOutput { - text: format_artifact_output(&output), - success: Some(success), - })) + Ok(FunctionToolOutput::from_text( + format_artifact_output(&output), + Some(success), + )) } } diff --git a/codex-rs/core/src/tools/handlers/code_mode.rs b/codex-rs/core/src/tools/handlers/code_mode.rs index c4994d059..025e85004 100644 --- a/codex-rs/core/src/tools/handlers/code_mode.rs +++ b/codex-rs/core/src/tools/handlers/code_mode.rs @@ -3,9 +3,8 @@ use async_trait::async_trait; use crate::features::Feature; use crate::function_tool::FunctionCallError; use crate::tools::code_mode; -use crate::tools::context::ContentToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; @@ -14,6 +13,8 @@ pub struct CodeModeHandler; #[async_trait] impl ToolHandler for CodeModeHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } @@ -22,7 +23,7 @@ impl ToolHandler for CodeModeHandler { matches!(payload, ToolPayload::Custom { .. }) } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -47,9 +48,6 @@ impl ToolHandler for CodeModeHandler { }; let content_items = code_mode::execute(session, turn, tracker, code).await?; - Ok(Box::new(ContentToolOutput { - content: content_items, - success: Some(true), - })) + Ok(FunctionToolOutput::from_content(content_items, Some(true))) } } diff --git a/codex-rs/core/src/tools/handlers/dynamic.rs b/codex-rs/core/src/tools/handlers/dynamic.rs index e06ffc7d6..23285bc90 100644 --- a/codex-rs/core/src/tools/handlers/dynamic.rs +++ b/codex-rs/core/src/tools/handlers/dynamic.rs @@ -1,9 +1,8 @@ use crate::codex::Session; use crate::codex::TurnContext; use crate::function_tool::FunctionCallError; -use crate::tools::context::ContentToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; @@ -23,6 +22,8 @@ pub struct DynamicToolHandler; #[async_trait] impl ToolHandler for DynamicToolHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } @@ -31,7 +32,7 @@ impl ToolHandler for DynamicToolHandler { true } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -67,10 +68,7 @@ impl ToolHandler for DynamicToolHandler { .into_iter() .map(FunctionCallOutputContentItem::from) .collect::>(); - Ok(Box::new(ContentToolOutput { - content: body, - success: Some(success), - })) + Ok(FunctionToolOutput::from_content(body, Some(success))) } } diff --git a/codex-rs/core/src/tools/handlers/grep_files.rs b/codex-rs/core/src/tools/handlers/grep_files.rs index 9d9434d85..071ecec70 100644 --- a/codex-rs/core/src/tools/handlers/grep_files.rs +++ b/codex-rs/core/src/tools/handlers/grep_files.rs @@ -7,9 +7,8 @@ use tokio::process::Command; use tokio::time::timeout; use crate::function_tool::FunctionCallError; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; @@ -38,11 +37,13 @@ struct GrepFilesArgs { #[async_trait] impl ToolHandler for GrepFilesHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { payload, turn, .. } = invocation; let arguments = match payload { @@ -86,15 +87,15 @@ impl ToolHandler for GrepFilesHandler { run_rg_search(pattern, include.as_deref(), &search_path, limit, &turn.cwd).await?; if search_results.is_empty() { - Ok(Box::new(TextToolOutput { - text: "No matches found.".to_string(), - success: Some(false), - })) + Ok(FunctionToolOutput::from_text( + "No matches found.".to_string(), + Some(false), + )) } else { - Ok(Box::new(TextToolOutput { - text: search_results.join("\n"), - success: Some(true), - })) + Ok(FunctionToolOutput::from_text( + search_results.join("\n"), + Some(true), + )) } } } diff --git a/codex-rs/core/src/tools/handlers/js_repl.rs b/codex-rs/core/src/tools/handlers/js_repl.rs index dec739942..d0404a9e2 100644 --- a/codex-rs/core/src/tools/handlers/js_repl.rs +++ b/codex-rs/core/src/tools/handlers/js_repl.rs @@ -9,10 +9,8 @@ use crate::exec::StreamOutput; use crate::features::Feature; use crate::function_tool::FunctionCallError; use crate::protocol::ExecCommandSource; -use crate::tools::context::ContentToolOutput; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::events::ToolEmitter; use crate::tools::events::ToolEventCtx; @@ -96,6 +94,8 @@ async fn emit_js_repl_exec_end( } #[async_trait] impl ToolHandler for JsReplHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } @@ -107,7 +107,7 @@ impl ToolHandler for JsReplHandler { ) } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -175,26 +175,22 @@ impl ToolHandler for JsReplHandler { .await; if items.is_empty() { - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } else { - Ok(Box::new(ContentToolOutput { - content: items, - success: Some(true), - })) + Ok(FunctionToolOutput::from_content(items, Some(true))) } } } #[async_trait] impl ToolHandler for JsReplResetHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { if !invocation.session.features().enabled(Feature::JsRepl) { return Err(FunctionCallError::RespondToModel( "js_repl is disabled by feature flag".to_string(), @@ -202,10 +198,10 @@ impl ToolHandler for JsReplResetHandler { } let manager = invocation.turn.js_repl.manager().await?; manager.reset().await?; - Ok(Box::new(TextToolOutput { - text: "js_repl kernel reset".to_string(), - success: Some(true), - })) + Ok(FunctionToolOutput::from_text( + "js_repl kernel reset".to_string(), + Some(true), + )) } } diff --git a/codex-rs/core/src/tools/handlers/list_dir.rs b/codex-rs/core/src/tools/handlers/list_dir.rs index fd0062102..fb65b3282 100644 --- a/codex-rs/core/src/tools/handlers/list_dir.rs +++ b/codex-rs/core/src/tools/handlers/list_dir.rs @@ -10,9 +10,8 @@ use serde::Deserialize; use tokio::fs; use crate::function_tool::FunctionCallError; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; @@ -48,11 +47,13 @@ struct ListDirArgs { #[async_trait] impl ToolHandler for ListDirHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { payload, .. } = invocation; let arguments = match payload { @@ -102,10 +103,7 @@ impl ToolHandler for ListDirHandler { let mut output = Vec::with_capacity(entries.len() + 1); output.push(format!("Absolute path: {}", path.display())); output.extend(entries); - Ok(Box::new(TextToolOutput { - text: output.join("\n"), - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(output.join("\n"), Some(true))) } } diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 932cdce89..1564206be 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -3,11 +3,9 @@ use std::sync::Arc; use crate::function_tool::FunctionCallError; use crate::mcp_tool_call::handle_mcp_tool_call; -use crate::tools::context::ContentToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::McpToolOutput; -use crate::tools::context::TextToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; @@ -15,13 +13,43 @@ use codex_protocol::models::ResponseInputItem; pub struct McpHandler; +pub enum McpHandlerOutput { + Mcp(McpToolOutput), + Function(FunctionToolOutput), +} + +impl crate::tools::context::ToolOutput for McpHandlerOutput { + fn log_preview(&self) -> String { + match self { + Self::Mcp(output) => output.log_preview(), + Self::Function(output) => output.log_preview(), + } + } + + fn success_for_logging(&self) -> bool { + match self { + Self::Mcp(output) => output.success_for_logging(), + Self::Function(output) => output.success_for_logging(), + } + } + + fn into_response(self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem { + match self { + Self::Mcp(output) => output.into_response(call_id, payload), + Self::Function(output) => output.into_response(call_id, payload), + } + } +} + #[async_trait] impl ToolHandler for McpHandler { + type Output = McpHandlerOutput; + fn kind(&self) -> ToolKind { ToolKind::Mcp } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -58,16 +86,18 @@ impl ToolHandler for McpHandler { match response { ResponseInputItem::McpToolCallOutput { result, .. } => { - Ok(Box::new(McpToolOutput { result })) + Ok(McpHandlerOutput::Mcp(McpToolOutput { result })) } ResponseInputItem::FunctionCallOutput { output, .. } => { let success = output.success; match output.body { - codex_protocol::models::FunctionCallOutputBody::Text(text) => { - Ok(Box::new(TextToolOutput { text, success })) - } + codex_protocol::models::FunctionCallOutputBody::Text(text) => Ok( + McpHandlerOutput::Function(FunctionToolOutput::from_text(text, success)), + ), codex_protocol::models::FunctionCallOutputBody::ContentItems(content) => { - Ok(Box::new(ContentToolOutput { content, success })) + Ok(McpHandlerOutput::Function( + FunctionToolOutput::from_content(content, success), + )) } } } diff --git a/codex-rs/core/src/tools/handlers/mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource.rs index c0f0691e2..d1480063b 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource.rs @@ -5,6 +5,7 @@ use std::time::Instant; use async_trait::async_trait; use codex_protocol::mcp::CallToolResult; +use codex_protocol::models::function_call_output_content_items_to_text; use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; use rmcp::model::PaginatedRequestParams; @@ -24,9 +25,8 @@ use crate::protocol::EventMsg; use crate::protocol::McpInvocation; use crate::protocol::McpToolCallBeginEvent; use crate::protocol::McpToolCallEndEvent; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; @@ -180,11 +180,13 @@ struct ReadResourcePayload { #[async_trait] impl ToolHandler for McpResourceHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -245,7 +247,7 @@ async fn handle_list_resources( turn: Arc, call_id: String, arguments: Option, -) -> Result { +) -> Result { let args: ListResourcesArgs = parse_args_with_default(arguments.clone())?; let ListResourcesArgs { server, cursor } = args; let server = normalize_optional_string(server); @@ -298,12 +300,8 @@ async fn handle_list_resources( match payload_result { Ok(payload) => match serialize_function_output(payload) { Ok(output) => { - let Some(output_text) = - (&*output as &dyn std::any::Any).downcast_ref::() - else { - unreachable!("MCP resource handler should return text output"); - }; - let content = output_text.text.clone(); + let content = + function_call_output_content_items_to_text(&output.body).unwrap_or_default(); let duration = start.elapsed(); emit_tool_call_end( &session, @@ -311,7 +309,7 @@ async fn handle_list_resources( &call_id, invocation, duration, - Ok(call_tool_result_from_content(&content, output_text.success)), + Ok(call_tool_result_from_content(&content, output.success)), ) .await; Ok(output) @@ -353,7 +351,7 @@ async fn handle_list_resource_templates( turn: Arc, call_id: String, arguments: Option, -) -> Result { +) -> Result { let args: ListResourceTemplatesArgs = parse_args_with_default(arguments.clone())?; let ListResourceTemplatesArgs { server, cursor } = args; let server = normalize_optional_string(server); @@ -408,12 +406,8 @@ async fn handle_list_resource_templates( match payload_result { Ok(payload) => match serialize_function_output(payload) { Ok(output) => { - let Some(output_text) = - (&*output as &dyn std::any::Any).downcast_ref::() - else { - unreachable!("MCP resource handler should return text output"); - }; - let content = output_text.text.clone(); + let content = + function_call_output_content_items_to_text(&output.body).unwrap_or_default(); let duration = start.elapsed(); emit_tool_call_end( &session, @@ -421,7 +415,7 @@ async fn handle_list_resource_templates( &call_id, invocation, duration, - Ok(call_tool_result_from_content(&content, output_text.success)), + Ok(call_tool_result_from_content(&content, output.success)), ) .await; Ok(output) @@ -463,7 +457,7 @@ async fn handle_read_resource( turn: Arc, call_id: String, arguments: Option, -) -> Result { +) -> Result { let args: ReadResourceArgs = parse_args(arguments.clone())?; let ReadResourceArgs { server, uri } = args; let server = normalize_required_string("server", server)?; @@ -503,12 +497,8 @@ async fn handle_read_resource( match payload_result { Ok(payload) => match serialize_function_output(payload) { Ok(output) => { - let Some(output_text) = - (&*output as &dyn std::any::Any).downcast_ref::() - else { - unreachable!("MCP resource handler should return text output"); - }; - let content = output_text.text.clone(); + let content = + function_call_output_content_items_to_text(&output.body).unwrap_or_default(); let duration = start.elapsed(); emit_tool_call_end( &session, @@ -516,7 +506,7 @@ async fn handle_read_resource( &call_id, invocation, duration, - Ok(call_tool_result_from_content(&content, output_text.success)), + Ok(call_tool_result_from_content(&content, output.success)), ) .await; Ok(output) @@ -620,7 +610,7 @@ fn normalize_required_string(field: &str, value: String) -> Result(payload: T) -> Result +fn serialize_function_output(payload: T) -> Result where T: Serialize, { @@ -630,10 +620,7 @@ where )) })?; - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } fn parse_arguments(raw_args: &str) -> Result, FunctionCallError> { diff --git a/codex-rs/core/src/tools/handlers/multi_agents.rs b/codex-rs/core/src/tools/handlers/multi_agents.rs index 58cca0094..aa121301c 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents.rs @@ -13,9 +13,8 @@ use crate::config::Config; use crate::error::CodexErr; use crate::features::Feature; use crate::function_tool::FunctionCallError; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; @@ -57,6 +56,8 @@ struct CloseAgentArgs { #[async_trait] impl ToolHandler for MultiAgentHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } @@ -65,7 +66,7 @@ impl ToolHandler for MultiAgentHandler { matches!(payload, ToolPayload::Function { .. }) } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -127,7 +128,7 @@ mod spawn { turn: Arc, call_id: String, arguments: String, - ) -> Result { + ) -> Result { let args: SpawnAgentArgs = parse_arguments(&arguments)?; let role_name = args .agent_type @@ -225,10 +226,7 @@ mod spawn { FunctionCallError::Fatal(format!("failed to serialize spawn_agent result: {err}")) })?; - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } } @@ -255,7 +253,7 @@ mod send_input { turn: Arc, call_id: String, arguments: String, - ) -> Result { + ) -> 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)?; @@ -318,10 +316,7 @@ mod send_input { FunctionCallError::Fatal(format!("failed to serialize send_input result: {err}")) })?; - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } } @@ -345,7 +340,7 @@ mod resume_agent { turn: Arc, call_id: String, arguments: String, - ) -> Result { + ) -> Result { let args: ResumeAgentArgs = parse_arguments(&arguments)?; let receiver_thread_id = agent_id(&args.id)?; let (receiver_agent_nickname, receiver_agent_role) = session @@ -432,10 +427,7 @@ mod resume_agent { FunctionCallError::Fatal(format!("failed to serialize resume_agent result: {err}")) })?; - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } async fn try_resume_closed_agent( @@ -495,7 +487,7 @@ pub(crate) mod wait { turn: Arc, call_id: String, arguments: String, - ) -> Result { + ) -> Result { let args: WaitArgs = parse_arguments(&arguments)?; if args.ids.is_empty() { return Err(FunctionCallError::RespondToModel( @@ -645,10 +637,7 @@ pub(crate) mod wait { FunctionCallError::Fatal(format!("failed to serialize wait result: {err}")) })?; - Ok(Box::new(TextToolOutput { - text: content, - success: None, - })) + Ok(FunctionToolOutput::from_text(content, None)) } async fn wait_for_final_status( @@ -688,7 +677,7 @@ pub mod close_agent { turn: Arc, call_id: String, arguments: String, - ) -> Result { + ) -> Result { let args: CloseAgentArgs = parse_arguments(&arguments)?; let agent_id = agent_id(&args.id)?; let (receiver_agent_nickname, receiver_agent_role) = session @@ -765,10 +754,7 @@ pub mod close_agent { FunctionCallError::Fatal(format!("failed to serialize close_agent result: {err}")) })?; - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } } @@ -993,8 +979,7 @@ mod tests { use crate::protocol::SandboxPolicy; use crate::protocol::SessionSource; use crate::protocol::SubAgentSource; - use crate::tools::context::TextToolOutput; - use crate::tools::context::ToolOutputBox; + use crate::tools::context::FunctionToolOutput; use crate::turn_diff_tracker::TurnDiffTracker; use codex_protocol::ThreadId; use codex_protocol::models::ContentItem; @@ -1040,11 +1025,12 @@ mod tests { ) } - fn expect_text_output(output: ToolOutputBox) -> (String, Option) { - let output = (&*output as &dyn std::any::Any) - .downcast_ref::() - .expect("expected text output"); - (output.text.clone(), output.success) + 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, + ) } #[tokio::test] diff --git a/codex-rs/core/src/tools/handlers/plan.rs b/codex-rs/core/src/tools/handlers/plan.rs index 638eb6379..6b810e0a6 100644 --- a/codex-rs/core/src/tools/handlers/plan.rs +++ b/codex-rs/core/src/tools/handlers/plan.rs @@ -3,9 +3,8 @@ use crate::client_common::tools::ToolSpec; use crate::codex::Session; use crate::codex::TurnContext; use crate::function_tool::FunctionCallError; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; @@ -63,11 +62,13 @@ At most one step can be in_progress at a time. #[async_trait] impl ToolHandler for PlanHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -88,10 +89,7 @@ impl ToolHandler for PlanHandler { let content = handle_update_plan(session.as_ref(), turn.as_ref(), arguments, call_id).await?; - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } } diff --git a/codex-rs/core/src/tools/handlers/read_file.rs b/codex-rs/core/src/tools/handlers/read_file.rs index 25cab314d..e88bf9baa 100644 --- a/codex-rs/core/src/tools/handlers/read_file.rs +++ b/codex-rs/core/src/tools/handlers/read_file.rs @@ -6,9 +6,8 @@ use codex_utils_string::take_bytes_at_char_boundary; use serde::Deserialize; use crate::function_tool::FunctionCallError; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; @@ -94,11 +93,13 @@ impl LineRecord { #[async_trait] impl ToolHandler for ReadFileHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { payload, .. } = invocation; let arguments = match payload { @@ -146,10 +147,10 @@ impl ToolHandler for ReadFileHandler { indentation::read_block(&path, offset, limit, indentation).await? } }; - Ok(Box::new(TextToolOutput { - text: collected.join("\n"), - success: Some(true), - })) + Ok(FunctionToolOutput::from_text( + collected.join("\n"), + Some(true), + )) } } diff --git a/codex-rs/core/src/tools/handlers/request_permissions.rs b/codex-rs/core/src/tools/handlers/request_permissions.rs index 3650a339e..f14a8bd71 100644 --- a/codex-rs/core/src/tools/handlers/request_permissions.rs +++ b/codex-rs/core/src/tools/handlers/request_permissions.rs @@ -3,9 +3,8 @@ use codex_protocol::request_permissions::RequestPermissionsArgs; use crate::function_tool::FunctionCallError; use crate::sandboxing::normalize_additional_permissions; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments_with_base_path; use crate::tools::registry::ToolHandler; @@ -20,11 +19,13 @@ pub struct RequestPermissionsHandler; #[async_trait] impl ToolHandler for RequestPermissionsHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -67,9 +68,6 @@ impl ToolHandler for RequestPermissionsHandler { )) })?; - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } } diff --git a/codex-rs/core/src/tools/handlers/request_user_input.rs b/codex-rs/core/src/tools/handlers/request_user_input.rs index 1a5ed6559..77def0b68 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input.rs @@ -1,7 +1,6 @@ use crate::function_tool::FunctionCallError; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; @@ -58,11 +57,13 @@ pub struct RequestUserInputHandler { #[async_trait] impl ToolHandler for RequestUserInputHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -115,10 +116,7 @@ impl ToolHandler for RequestUserInputHandler { )) })?; - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } } diff --git a/codex-rs/core/src/tools/handlers/search_tool_bm25.rs b/codex-rs/core/src/tools/handlers/search_tool_bm25.rs index 556883cbf..a038fd5ed 100644 --- a/codex-rs/core/src/tools/handlers/search_tool_bm25.rs +++ b/codex-rs/core/src/tools/handlers/search_tool_bm25.rs @@ -12,9 +12,8 @@ use crate::connectors; use crate::function_tool::FunctionCallError; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp_connection_manager::ToolInfo; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; @@ -74,11 +73,13 @@ impl ToolEntry { #[async_trait] impl ToolHandler for SearchToolBm25Handler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { payload, session, @@ -141,10 +142,7 @@ impl ToolHandler for SearchToolBm25Handler { "tools": [], }) .to_string(); - return Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })); + return Ok(FunctionToolOutput::from_text(content, Some(true))); } let documents: Vec> = entries @@ -184,10 +182,7 @@ impl ToolHandler for SearchToolBm25Handler { }) .to_string(); - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } } diff --git a/codex-rs/core/src/tools/handlers/shell.rs b/codex-rs/core/src/tools/handlers/shell.rs index 542f3a279..8c8d668f8 100644 --- a/codex-rs/core/src/tools/handlers/shell.rs +++ b/codex-rs/core/src/tools/handlers/shell.rs @@ -14,9 +14,8 @@ use crate::is_safe_command::is_known_safe_command; use crate::protocol::ExecCommandSource; use crate::shell::Shell; use crate::skills::maybe_emit_implicit_skill_invocation; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::events::ToolEmitter; use crate::tools::events::ToolEventCtx; @@ -142,6 +141,8 @@ impl From for ShellCommandHandler { #[async_trait] impl ToolHandler for ShellHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } @@ -165,7 +166,7 @@ impl ToolHandler for ShellHandler { } } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -224,6 +225,8 @@ impl ToolHandler for ShellHandler { #[async_trait] impl ToolHandler for ShellCommandHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } @@ -253,7 +256,7 @@ impl ToolHandler for ShellCommandHandler { .unwrap_or(true) } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -305,7 +308,7 @@ impl ToolHandler for ShellCommandHandler { } impl ShellHandler { - async fn run_exec_like(args: RunExecLikeArgs) -> Result { + async fn run_exec_like(args: RunExecLikeArgs) -> Result { let RunExecLikeArgs { tool_name, exec_params, @@ -449,10 +452,7 @@ impl ShellHandler { .map(|result| result.output); let event_ctx = ToolEventCtx::new(session.as_ref(), turn.as_ref(), &call_id, None); let content = emitter.finish(event_ctx, out).await?; - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } } diff --git a/codex-rs/core/src/tools/handlers/test_sync.rs b/codex-rs/core/src/tools/handlers/test_sync.rs index 210feda76..2d6e35148 100644 --- a/codex-rs/core/src/tools/handlers/test_sync.rs +++ b/codex-rs/core/src/tools/handlers/test_sync.rs @@ -10,9 +10,8 @@ use tokio::sync::Barrier; use tokio::time::sleep; use crate::function_tool::FunctionCallError; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; @@ -57,11 +56,13 @@ fn barrier_map() -> &'static tokio::sync::Mutex> { #[async_trait] impl ToolHandler for TestSyncHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { payload, .. } = invocation; let arguments = match payload { @@ -91,10 +92,7 @@ impl ToolHandler for TestSyncHandler { sleep(Duration::from_millis(delay)).await; } - Ok(Box::new(TextToolOutput { - text: "ok".to_string(), - success: Some(true), - })) + Ok(FunctionToolOutput::from_text("ok".to_string(), Some(true))) } } diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index 649a3ebab..721fd30d1 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -7,9 +7,8 @@ use crate::sandboxing::SandboxPermissions; use crate::shell::Shell; use crate::shell::get_shell_by_model_provided_path; use crate::skills::maybe_emit_implicit_skill_invocation; -use crate::tools::context::TextToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::handlers::apply_granted_turn_permissions; use crate::tools::handlers::apply_patch::intercept_apply_patch; @@ -83,6 +82,8 @@ fn default_tty() -> bool { #[async_trait] impl ToolHandler for UnifiedExecHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } @@ -114,7 +115,7 @@ impl ToolHandler for UnifiedExecHandler { !is_known_safe_command(&command) } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, turn, @@ -291,10 +292,7 @@ impl ToolHandler for UnifiedExecHandler { let content = format_response(&response); - Ok(Box::new(TextToolOutput { - text: content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_text(content, Some(true))) } } diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs index ffcf0f5dd..959e073d7 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -12,9 +12,8 @@ use crate::features::Feature; use crate::function_tool::FunctionCallError; use crate::protocol::EventMsg; use crate::protocol::ViewImageToolCallEvent; -use crate::tools::context::ContentToolOutput; +use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; @@ -32,11 +31,13 @@ struct ViewImageArgs { #[async_trait] impl ToolHandler for ViewImageHandler { + type Output = FunctionToolOutput; + fn kind(&self) -> ToolKind { ToolKind::Function } - async fn handle(&self, invocation: ToolInvocation) -> Result { + async fn handle(&self, invocation: ToolInvocation) -> Result { if !invocation .turn .model_info @@ -121,9 +122,6 @@ impl ToolHandler for ViewImageHandler { ) .await; - Ok(Box::new(ContentToolOutput { - content, - success: Some(true), - })) + Ok(FunctionToolOutput::from_content(content, Some(true))) } } diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 6e66ba289..eb74f90db 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -10,7 +10,7 @@ use crate::memories::usage::emit_metric_for_tool_read; use crate::protocol::SandboxPolicy; use crate::sandbox_tags::sandbox_tag; use crate::tools::context::ToolInvocation; -use crate::tools::context::ToolOutputBox; +use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; use async_trait::async_trait; use codex_hooks::HookEvent; @@ -32,6 +32,8 @@ pub enum ToolKind { #[async_trait] pub trait ToolHandler: Send + Sync { + type Output: ToolOutput + 'static; + fn kind(&self) -> ToolKind; fn matches_kind(&self, payload: &ToolPayload) -> bool { @@ -52,19 +54,68 @@ pub trait ToolHandler: Send + Sync { /// Perform the actual [ToolInvocation] and returns a [ToolOutput] containing /// the final output to return to the model. - async fn handle(&self, invocation: ToolInvocation) -> Result; + async fn handle(&self, invocation: ToolInvocation) -> Result; +} + +struct AnyToolResult { + preview: String, + success: bool, + response: ResponseInputItem, +} + +#[async_trait] +trait AnyToolHandler: Send + Sync { + fn matches_kind(&self, payload: &ToolPayload) -> bool; + + async fn is_mutating(&self, invocation: &ToolInvocation) -> bool; + + async fn handle_any( + &self, + invocation: ToolInvocation, + ) -> Result; +} + +#[async_trait] +impl AnyToolHandler for T +where + T: ToolHandler, +{ + fn matches_kind(&self, payload: &ToolPayload) -> bool { + ToolHandler::matches_kind(self, payload) + } + + async fn is_mutating(&self, invocation: &ToolInvocation) -> bool { + ToolHandler::is_mutating(self, invocation).await + } + + async fn handle_any( + &self, + invocation: ToolInvocation, + ) -> Result { + let call_id = invocation.call_id.clone(); + let payload = invocation.payload.clone(); + let output = self.handle(invocation).await?; + let preview = output.log_preview(); + let success = output.success_for_logging(); + let response = output.into_response(&call_id, &payload); + Ok(AnyToolResult { + preview, + success, + response, + }) + } } pub struct ToolRegistry { - handlers: HashMap>, + handlers: HashMap>, } impl ToolRegistry { - pub fn new(handlers: HashMap>) -> Self { + fn new(handlers: HashMap>) -> Self { Self { handlers } } - pub fn handler(&self, name: &str) -> Option> { + fn handler(&self, name: &str) -> Option> { self.handlers.get(name).map(Arc::clone) } @@ -163,7 +214,7 @@ impl ToolRegistry { } let is_mutating = handler.is_mutating(&invocation).await; - let output_cell = tokio::sync::Mutex::new(None); + let response_cell = tokio::sync::Mutex::new(None); let invocation_for_tool = invocation.clone(); let started = Instant::now(); @@ -177,19 +228,22 @@ impl ToolRegistry { mcp_server_origin_ref, || { let handler = handler.clone(); - let output_cell = &output_cell; + let response_cell = &response_cell; async move { if is_mutating { tracing::trace!("waiting for tool gate"); invocation_for_tool.turn.tool_call_gate.wait_ready().await; tracing::trace!("tool gate released"); } - match handler.handle(invocation_for_tool).await { - Ok(output) => { - let preview = output.log_preview(); - let success = output.success_for_logging(); - let mut guard = output_cell.lock().await; - *guard = Some(output); + match handler.handle_any(invocation_for_tool).await { + Ok(result) => { + let AnyToolResult { + preview, + success, + response, + } = result; + let mut guard = response_cell.lock().await; + *guard = Some(response); Ok((preview, success)) } Err(err) => Err(err), @@ -220,11 +274,11 @@ impl ToolRegistry { match result { Ok(_) => { - let mut guard = output_cell.lock().await; - let output = guard.take().ok_or_else(|| { + let mut guard = response_cell.lock().await; + let response = guard.take().ok_or_else(|| { FunctionCallError::Fatal("tool produced no output".to_string()) })?; - Ok(output.into_response(&call_id_owned, &payload_for_response)) + Ok(response) } Err(err) => Err(err), } @@ -247,7 +301,7 @@ impl ConfiguredToolSpec { } pub struct ToolRegistryBuilder { - handlers: HashMap>, + handlers: HashMap>, specs: Vec, } @@ -272,8 +326,12 @@ impl ToolRegistryBuilder { .push(ConfiguredToolSpec::new(spec, supports_parallel_tool_calls)); } - pub fn register_handler(&mut self, name: impl Into, handler: Arc) { + pub fn register_handler(&mut self, name: impl Into, handler: Arc) + where + H: ToolHandler + 'static, + { let name = name.into(); + let handler: Arc = handler; if self .handlers .insert(name.clone(), handler.clone()) diff --git a/codex-rs/core/tests/suite/view_image.rs b/codex-rs/core/tests/suite/view_image.rs index 3a3c89ac7..cf06b872e 100644 --- a/codex-rs/core/tests/suite/view_image.rs +++ b/codex-rs/core/tests/suite/view_image.rs @@ -857,25 +857,11 @@ async fn view_image_tool_placeholder_for_non_image_files() -> anyhow::Result<()> request.inputs_of_type("input_image").is_empty(), "non-image file should not produce an input_image message" ); - let function_output = request.function_call_output(call_id); - let output_items = function_output - .get("output") - .and_then(Value::as_array) - .expect("function_call_output should be a content item array"); - assert_eq!( - output_items.len(), - 1, - "non-image placeholder should be returned as a single content item" - ); - assert_eq!( - output_items[0].get("type").and_then(Value::as_str), - Some("input_text"), - "non-image placeholder should be returned as input_text" - ); - let placeholder = output_items[0] - .get("text") - .and_then(Value::as_str) - .expect("placeholder text present"); + let (placeholder, success) = request + .function_call_output_content_and_success(call_id) + .expect("function_call_output should be present"); + assert_eq!(success, None); + let placeholder = placeholder.expect("placeholder text present"); assert!( placeholder.contains("Codex could not read the local image at")