From ed5944ba1db98f0ef13572bfee1452d4e52f80a1 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Mon, 11 May 2026 17:11:31 -0700 Subject: [PATCH] Simplify MCP tool handler plumbing (#21595) ## Why The MCP tool path had accumulated a few core-owned special cases: a dedicated payload variant, resolver plumbing, a legacy `AfterToolUse` translation path, and a side channel for parallel-call metadata. That made `ToolRegistry` and the spec builder know more about MCP than they needed to. This change moves MCP-specific execution details back onto `ToolInfo` and `McpHandler` so `codex-core` can treat MCP calls like normal function calls while still preserving MCP-specific dispatch and telemetry behavior where it belongs. ## What changed - removed `resolve_mcp_tool_info`, `ToolPayload::Mcp`, `ToolKind`, and the remaining registry-side MCP resolver path - stored MCP routing metadata directly on `McpHandler` and `ToolInfo`, including `supports_parallel_tool_calls` - deleted the legacy `AfterToolUse` consumer in `core`, which removes the need for handler-specific `after_tool_use_payload` implementations - switched tool-result telemetry to handler-provided tags and kept MCP-specific dispatch payload construction inside the handler - simplified tool spec planning/building by passing `ToolInfo` directly and dropping the direct/deferred MCP wrapper structs and the parallel-server side table ## Testing - `cargo check -p codex-core -p codex-mcp -p codex-otel` - `cargo test -p codex-core mcp_parallel_support_uses_exact_payload_server` - `cargo test -p codex-core direct_mcp_tools_register_namespaced_handlers` - `cargo test -p codex-core search_tool_description_lists_each_mcp_source_once` - `cargo test -p codex-mcp list_all_tools_uses_startup_snapshot_while_client_is_pending` - `just fix -p codex-core -p codex-mcp -p codex-otel` --- codex-rs/codex-mcp/src/connection_manager.rs | 43 ++++---- .../codex-mcp/src/connection_manager_tests.rs | 81 +++++++------- codex-rs/codex-mcp/src/rmcp_client.rs | 2 + codex-rs/codex-mcp/src/tools.rs | 6 + codex-rs/core/src/connectors_tests.rs | 6 + codex-rs/core/src/mcp_tool_exposure_test.rs | 2 + codex-rs/core/src/session/mcp.rs | 13 --- codex-rs/core/src/session/mod.rs | 2 - codex-rs/core/src/session/tests.rs | 5 +- codex-rs/core/src/session/turn.rs | 2 - codex-rs/core/src/stream_events_utils.rs | 2 +- .../src/tools/code_mode/execute_handler.rs | 5 - codex-rs/core/src/tools/code_mode/mod.rs | 29 +---- .../core/src/tools/code_mode/wait_handler.rs | 5 - codex-rs/core/src/tools/context.rs | 32 +----- codex-rs/core/src/tools/context_tests.rs | 30 ++--- .../agent_jobs/report_agent_job_result.rs | 5 - .../agent_jobs/spawn_agents_on_csv.rs | 5 - .../core/src/tools/handlers/apply_patch.rs | 5 - codex-rs/core/src/tools/handlers/dynamic.rs | 5 - .../src/tools/handlers/extension_tools.rs | 5 - .../src/tools/handlers/goal/create_goal.rs | 5 - .../core/src/tools/handlers/goal/get_goal.rs | 5 - .../src/tools/handlers/goal/update_goal.rs | 5 - codex-rs/core/src/tools/handlers/mcp.rs | 103 ++++++++++-------- .../list_mcp_resource_templates.rs | 5 - .../mcp_resource/list_mcp_resources.rs | 5 - .../mcp_resource/read_mcp_resource.rs | 5 - .../core/src/tools/handlers/multi_agents.rs | 1 - .../handlers/multi_agents/close_agent.rs | 4 - .../handlers/multi_agents/resume_agent.rs | 4 - .../tools/handlers/multi_agents/send_input.rs | 4 - .../src/tools/handlers/multi_agents/spawn.rs | 4 - .../src/tools/handlers/multi_agents/wait.rs | 4 - .../src/tools/handlers/multi_agents_v2.rs | 1 - .../handlers/multi_agents_v2/close_agent.rs | 4 - .../handlers/multi_agents_v2/followup_task.rs | 4 - .../handlers/multi_agents_v2/list_agents.rs | 4 - .../handlers/multi_agents_v2/send_message.rs | 4 - .../tools/handlers/multi_agents_v2/spawn.rs | 4 - .../tools/handlers/multi_agents_v2/wait.rs | 4 - codex-rs/core/src/tools/handlers/plan.rs | 5 - .../src/tools/handlers/request_permissions.rs | 5 - .../tools/handlers/request_plugin_install.rs | 5 - .../src/tools/handlers/request_user_input.rs | 5 - .../tools/handlers/shell/container_exec.rs | 5 - .../src/tools/handlers/shell/local_shell.rs | 5 - .../src/tools/handlers/shell/shell_command.rs | 5 - .../src/tools/handlers/shell/shell_handler.rs | 5 - codex-rs/core/src/tools/handlers/test_sync.rs | 5 - .../core/src/tools/handlers/tool_search.rs | 7 +- .../src/tools/handlers/unavailable_tool.rs | 5 - .../handlers/unified_exec/exec_command.rs | 5 - .../handlers/unified_exec/write_stdin.rs | 5 - .../core/src/tools/handlers/view_image.rs | 5 - codex-rs/core/src/tools/registry.rs | 91 +++++++++------- codex-rs/core/src/tools/registry_tests.rs | 4 - codex-rs/core/src/tools/router.rs | 41 ++----- codex-rs/core/src/tools/router_tests.rs | 102 ++++++++++------- codex-rs/core/src/tools/spec.rs | 32 ++---- codex-rs/core/src/tools/spec_plan.rs | 36 +++--- codex-rs/core/src/tools/spec_plan_tests.rs | 81 ++++++++++---- codex-rs/core/src/tools/spec_plan_types.rs | 23 +--- codex-rs/core/src/tools/spec_tests.rs | 13 ++- .../core/src/tools/tool_dispatch_trace.rs | 9 -- .../src/tools/tool_dispatch_trace_tests.rs | 5 - codex-rs/otel/src/events/session_telemetry.rs | 20 ++-- .../tests/suite/otel_export_routing_policy.rs | 18 ++- codex-rs/otel/tests/suite/runtime_summary.rs | 3 +- codex-rs/rollout-trace/src/tool_dispatch.rs | 31 +----- 70 files changed, 412 insertions(+), 638 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 1b09ac5da..c5593a559 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -7,7 +7,6 @@ //! `codex-core`. use std::collections::HashMap; -use std::collections::HashSet; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -44,7 +43,6 @@ use codex_config::Constrained; use codex_config::McpServerTransportConfig; use codex_config::types::OAuthCredentialsStoreMode; use codex_login::CodexAuth; -use codex_protocol::ToolName; use codex_protocol::mcp::CallToolResult; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; @@ -131,17 +129,6 @@ impl McpConnectionManager { .is_none_or(|metadata| metadata.pollutes_memory) } - pub fn parallel_tool_call_server_names(&self) -> HashSet { - self.server_metadata - .iter() - .filter_map(|(name, metadata)| { - metadata - .supports_parallel_tool_calls - .then_some(name.clone()) - }) - .collect() - } - pub fn is_host_owned_codex_apps_server(&self, server_name: &str) -> bool { self.host_owned_codex_apps_enabled && server_name == CODEX_APPS_MCP_SERVER_NAME } @@ -373,7 +360,11 @@ impl McpConnectionManager { let Some(server_tools) = managed_client.listed_tools().await else { continue; }; - tools.extend(server_tools); + tools.extend( + server_tools + .into_iter() + .map(|tool| self.with_server_metadata(tool)), + ); } normalize_tools_for_model(tools) } @@ -424,11 +415,26 @@ impl McpConnectionManager { .into_iter() .map(|mut tool| { tool.tool = tool_with_model_visible_input_schema(&tool.tool); - tool + self.with_server_metadata(tool) }); Ok(normalize_tools_for_model(tools)) } + fn with_server_metadata(&self, mut tool: ToolInfo) -> ToolInfo { + let Some(metadata) = self.server_metadata.get(&tool.server_name) else { + tool.supports_parallel_tool_calls = false; + tool.server_origin = None; + return tool; + }; + + tool.supports_parallel_tool_calls = metadata.supports_parallel_tool_calls; + tool.server_origin = metadata + .origin + .as_ref() + .map(|origin| origin.as_str().to_string()); + tool + } + /// Returns a single map that contains all resources. Each key is the /// server name and the value is a vector of resources. pub async fn list_all_resources(&self) -> HashMap> { @@ -662,13 +668,6 @@ impl McpConnectionManager { .with_context(|| format!("resources/read failed for `{server}` ({uri})")) } - pub async fn resolve_tool_info(&self, tool_name: &ToolName) -> Option { - let all_tools = self.list_all_tools().await; - all_tools - .into_iter() - .find(|tool| tool.canonical_tool_name() == *tool_name) - } - async fn client_by_name(&self, name: &str) -> Result { self.clients .get(name) diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index b2cf5f94f..5d04dc2a3 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -10,6 +10,7 @@ use crate::elicitation::elicitation_is_rejected_by_policy; use crate::rmcp_client::AsyncManagedClient; use crate::rmcp_client::ManagedClient; use crate::rmcp_client::StartupOutcomeError; +use crate::server::McpServerOrigin; use crate::tools::ToolFilter; use crate::tools::ToolInfo; use crate::tools::filter_tools; @@ -38,6 +39,8 @@ fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { let tool_namespace = format!("mcp__{server_name}__"); ToolInfo { server_name: server_name.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, callable_name: tool_name.to_string(), callable_namespace: tool_namespace, namespace_description: None, @@ -712,44 +715,6 @@ async fn list_all_tools_uses_startup_snapshot_while_client_is_pending() { assert_eq!(tool.callable_name, "calendar_create_event"); } -#[tokio::test] -async fn resolve_tool_info_accepts_canonical_namespaced_tool_names() { - let startup_tools = vec![create_test_tool("rmcp", "echo")]; - let pending_client = futures::future::pending::>() - .boxed() - .shared(); - let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); - let permission_profile = Constrained::allow_any(PermissionProfile::default()); - let mut manager = - McpConnectionManager::new_uninitialized(&approval_policy, &permission_profile); - manager.clients.insert( - "rmcp".to_string(), - AsyncManagedClient { - client: pending_client, - startup_snapshot: Some(startup_tools), - startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), - cancel_token: CancellationToken::new(), - }, - ); - - let tool = manager - .resolve_tool_info(&ToolName::namespaced("mcp__rmcp__", "echo")) - .await - .expect("split MCP tool namespace and name should resolve"); - - let expected = ("rmcp", "mcp__rmcp__", "echo", "echo"); - assert_eq!( - ( - tool.server_name.as_str(), - tool.callable_namespace.as_str(), - tool.callable_name.as_str(), - tool.tool.name.as_ref(), - ), - expected - ); -} - #[tokio::test] async fn list_all_tools_blocks_while_client_is_pending_without_startup_snapshot() { let pending_client = futures::future::pending::>() @@ -842,6 +807,46 @@ async fn list_all_tools_uses_startup_snapshot_when_client_startup_fails() { assert_eq!(tool.callable_name, "calendar_create_event"); } +#[tokio::test] +async fn list_all_tools_adds_server_metadata_to_cached_tools() { + let server_name = "docs"; + let startup_tools = vec![create_test_tool(server_name, "search")]; + let pending_client = futures::future::pending::>() + .boxed() + .shared(); + let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = + McpConnectionManager::new_uninitialized(&approval_policy, &permission_profile); + manager.server_metadata.insert( + server_name.to_string(), + McpServerMetadata { + pollutes_memory: true, + origin: Some(McpServerOrigin::StreamableHttp( + "https://docs.example".to_string(), + )), + supports_parallel_tool_calls: true, + }, + ); + manager.clients.insert( + server_name.to_string(), + AsyncManagedClient { + client: pending_client, + startup_snapshot: Some(startup_tools), + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + }, + ); + + let tools = manager.list_all_tools().await; + assert_eq!(tools.len(), 1); + let tool = &tools[0]; + assert_eq!(tool.server_name, server_name); + assert!(tool.supports_parallel_tool_calls); + assert_eq!(tool.server_origin.as_deref(), Some("https://docs.example")); +} + #[test] fn elicitation_capability_uses_2025_06_18_shape_for_form_only_support() { let capability = Some(ElicitationCapability::default()); diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index bca85f00d..fcc9d1c86 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -375,6 +375,8 @@ pub(crate) async fn list_tools_for_client_uncached( }; ToolInfo { server_name: server_name.to_owned(), + supports_parallel_tool_calls: false, + server_origin: None, callable_name, callable_namespace, namespace_description, diff --git a/codex-rs/codex-mcp/src/tools.rs b/codex-rs/codex-mcp/src/tools.rs index cb3d8baba..5b2fba2ef 100644 --- a/codex-rs/codex-mcp/src/tools.rs +++ b/codex-rs/codex-mcp/src/tools.rs @@ -29,6 +29,12 @@ pub(crate) const MCP_TOOLS_CACHE_WRITE_DURATION_METRIC: &str = pub struct ToolInfo { /// Raw MCP server name used for routing the tool call. pub server_name: String, + /// Whether calls routed to this server may run in parallel. + #[serde(default)] + pub supports_parallel_tool_calls: bool, + /// MCP server origin used for telemetry and diagnostics, when known. + #[serde(default)] + pub server_origin: Option, /// Model-visible tool name used in Responses API tool declarations. #[serde(rename = "tool_name", alias = "callable_name")] pub callable_name: String, diff --git a/codex-rs/core/src/connectors_tests.rs b/codex-rs/core/src/connectors_tests.rs index 1f6f0e1ea..c5798d7c1 100644 --- a/codex-rs/core/src/connectors_tests.rs +++ b/codex-rs/core/src/connectors_tests.rs @@ -120,6 +120,8 @@ fn codex_app_tool( ToolInfo { server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, callable_name: tool_name.to_string(), callable_namespace: tool_namespace, namespace_description: None, @@ -190,6 +192,8 @@ fn accessible_connectors_from_mcp_tools_carries_plugin_display_names() { ), ToolInfo { server_name: "sample".to_string(), + supports_parallel_tool_calls: false, + server_origin: None, callable_name: "echo".to_string(), callable_namespace: "sample".to_string(), namespace_description: None, @@ -305,6 +309,8 @@ fn merge_connectors_unions_and_dedupes_plugin_display_names() { fn accessible_connectors_from_mcp_tools_preserves_description() { let mcp_tools = vec![ToolInfo { server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, callable_name: "calendar_create_event".to_string(), callable_namespace: "mcp__codex_apps__calendar".to_string(), namespace_description: Some("Plan events".to_string()), diff --git a/codex-rs/core/src/mcp_tool_exposure_test.rs b/codex-rs/core/src/mcp_tool_exposure_test.rs index c3f467800..f9468d898 100644 --- a/codex-rs/core/src/mcp_tool_exposure_test.rs +++ b/codex-rs/core/src/mcp_tool_exposure_test.rs @@ -49,6 +49,8 @@ fn make_mcp_tool( ) -> ToolInfo { ToolInfo { server_name: server_name.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, callable_name: callable_name.to_string(), callable_namespace: callable_namespace.to_string(), namespace_description: None, diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 102a9f211..d32ecf2b1 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -263,19 +263,6 @@ impl Session { .await } - #[expect( - clippy::await_holding_invalid_type, - reason = "MCP tool metadata reads through the session-owned manager guard" - )] - pub(crate) async fn resolve_mcp_tool_info(&self, tool_name: &ToolName) -> Option { - self.services - .mcp_connection_manager - .read() - .await - .resolve_tool_info(tool_name) - .await - } - async fn refresh_mcp_servers_inner( &self, turn_context: &TurnContext, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 5f7f3fd16..127f03496 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -64,7 +64,6 @@ use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_login::default_client::originator; use codex_mcp::McpConnectionManager; use codex_mcp::McpRuntimeEnvironment; -use codex_mcp::ToolInfo; use codex_mcp::codex_apps_tools_cache_key; use codex_models_manager::manager::RefreshStrategy; use codex_models_manager::manager::SharedModelsManager; @@ -75,7 +74,6 @@ use codex_otel::current_span_trace_id; use codex_otel::current_span_w3c_trace_context; use codex_otel::set_parent_from_w3c_trace_context; use codex_protocol::ThreadId; -use codex_protocol::ToolName; use codex_protocol::account::PlanType as AccountPlanType; use codex_protocol::approvals::ElicitationRequestEvent; use codex_protocol::approvals::ExecPolicyAmendment; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 9574cae30..94570785f 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -543,7 +543,6 @@ fn test_tool_runtime(session: Arc, turn_context: Arc) -> T mcp_tools: None, deferred_mcp_tools: None, unavailable_called_tools: Vec::new(), - parallel_mcp_server_names: HashSet::new(), discoverable_tools: None, extension_tool_bundles: Vec::new(), dynamic_tools: turn_context.dynamic_tools.as_slice(), @@ -8641,7 +8640,6 @@ async fn fatal_tool_error_stops_turn_and_reports_error() { deferred_mcp_tools, mcp_tools: Some(tools), unavailable_called_tools: Vec::new(), - parallel_mcp_server_names: HashSet::new(), discoverable_tools: None, extension_tool_bundles: Vec::new(), dynamic_tools: turn_context.dynamic_tools.as_slice(), @@ -8655,8 +8653,7 @@ async fn fatal_tool_error_stops_turn_and_reports_error() { input: "{}".to_string(), }; - let call = ToolRouter::build_tool_call(session.as_ref(), item.clone()) - .await + let call = ToolRouter::build_tool_call(item.clone()) .expect("build tool call") .expect("tool call present"); let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 4a75ed063..71ea34da8 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1161,7 +1161,6 @@ pub(crate) async fn built_tools( .list_all_tools() .or_cancel(cancellation_token) .await?; - let parallel_mcp_server_names = mcp_connection_manager.parallel_tool_call_server_names(); drop(mcp_connection_manager); let loaded_plugins = sess .services @@ -1267,7 +1266,6 @@ pub(crate) async fn built_tools( mcp_tools, deferred_mcp_tools, unavailable_called_tools, - parallel_mcp_server_names, discoverable_tools, extension_tool_bundles: extension_tool_bundles(sess), dynamic_tools: turn_context.dynamic_tools.as_slice(), diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 14a4e975a..b4f00ee86 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -228,7 +228,7 @@ pub(crate) async fn handle_output_item_done( let mut output = OutputItemResult::default(); let plan_mode = ctx.turn_context.collaboration_mode.mode == ModeKind::Plan; - match ToolRouter::build_tool_call(ctx.sess.as_ref(), item.clone()).await { + match ToolRouter::build_tool_call(item.clone()) { // The model emitted a tool call; log it, persist the item immediately, and queue the tool execution. Ok(Some(call)) => { ctx.sess diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 0e11cd166..80a72a88e 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -3,7 +3,6 @@ use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -94,10 +93,6 @@ impl ToolHandler for CodeModeExecuteHandler { Some(self.spec.clone()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Custom { .. }) } diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 0cf7ed898..9ee8e0352 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -276,7 +276,6 @@ async fn build_nested_router(exec: &ExecContext) -> ToolRouter { let nested_tools_config = exec.turn.tools_config.for_code_mode_nested_tools(); let mcp_connection_manager = exec.session.services.mcp_connection_manager.read().await; let listed_mcp_tools = mcp_connection_manager.list_all_tools().await; - let parallel_mcp_server_names = mcp_connection_manager.parallel_tool_call_server_names(); ToolRouter::from_config( &nested_tools_config, @@ -284,7 +283,6 @@ async fn build_nested_router(exec: &ExecContext) -> ToolRouter { deferred_mcp_tools: None, mcp_tools: Some(listed_mcp_tools), unavailable_called_tools: Vec::new(), - parallel_mcp_server_names, discoverable_tools: None, extension_tool_bundles: extension_tool_bundles(exec.session.as_ref()), dynamic_tools: exec.turn.dynamic_tools.as_slice(), @@ -293,7 +291,7 @@ async fn build_nested_router(exec: &ExecContext) -> ToolRouter { } async fn call_nested_tool( - exec: ExecContext, + _exec: ExecContext, tool_runtime: ToolCallRuntime, invocation: CodeModeNestedToolCall, cancellation_token: CancellationToken, @@ -310,29 +308,14 @@ async fn call_nested_tool( ))); } - let (tool_call_name, payload) = - if let Some(tool_info) = exec.session.resolve_mcp_tool_info(&tool_name).await { - let raw_arguments = match serialize_function_tool_arguments(&tool_name, input) { - Ok(raw_arguments) => raw_arguments, - Err(error) => return Err(FunctionCallError::RespondToModel(error)), - }; - ( - tool_info.canonical_tool_name(), - ToolPayload::Mcp { - server: tool_info.server_name, - tool: tool_info.tool.name.to_string(), - raw_arguments, - }, - ) - } else { - match build_nested_tool_payload(tool_runtime.find_spec(&tool_name), &tool_name, input) { - Ok(payload) => (tool_name, payload), - Err(error) => return Err(FunctionCallError::RespondToModel(error)), - } + let payload = + match build_nested_tool_payload(tool_runtime.find_spec(&tool_name), &tool_name, input) { + Ok(payload) => payload, + Err(error) => return Err(FunctionCallError::RespondToModel(error)), }; let call = ToolCall { - tool_name: tool_call_name, + tool_name, call_id: format!("{PUBLIC_TOOL_NAME}-{}", uuid::Uuid::new_v4()), payload, }; diff --git a/codex-rs/core/src/tools/code_mode/wait_handler.rs b/codex-rs/core/src/tools/code_mode/wait_handler.rs index 21191be02..40ca82939 100644 --- a/codex-rs/core/src/tools/code_mode/wait_handler.rs +++ b/codex-rs/core/src/tools/code_mode/wait_handler.rs @@ -5,7 +5,6 @@ use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -52,10 +51,6 @@ impl ToolHandler for CodeModeWaitHandler { Some(create_wait_tool()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index cf6b3fac9..76a99fe8f 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -58,23 +58,10 @@ pub struct ToolInvocation { #[derive(Clone, Debug)] pub enum ToolPayload { - Function { - arguments: String, - }, - ToolSearch { - arguments: SearchToolCallParams, - }, - Custom { - input: String, - }, - LocalShell { - params: ShellToolCallParams, - }, - Mcp { - server: String, - tool: String, - raw_arguments: String, - }, + Function { arguments: String }, + ToolSearch { arguments: SearchToolCallParams }, + Custom { input: String }, + LocalShell { params: ShellToolCallParams }, } impl ToolPayload { @@ -84,7 +71,6 @@ impl ToolPayload { ToolPayload::ToolSearch { arguments } => Cow::Owned(arguments.query.clone()), ToolPayload::Custom { input } => Cow::Borrowed(input), ToolPayload::LocalShell { params } => Cow::Owned(params.command.join(" ")), - ToolPayload::Mcp { raw_arguments, .. } => Cow::Borrowed(raw_arguments), } } } @@ -363,10 +349,6 @@ impl ToolOutput for AbortedToolOutput { execution: "client".to_string(), tools: Vec::new(), }, - ToolPayload::Mcp { .. } => ResponseInputItem::McpToolCallOutput { - call_id: call_id.to_string(), - output: CallToolResult::from_error_text(self.message.clone()), - }, _ => function_tool_response( call_id, payload, @@ -515,10 +497,8 @@ pub(crate) fn response_input_to_code_mode_result(response: ResponseInputItem) -> }, ResponseInputItem::ToolSearchOutput { tools, .. } => JsonValue::Array(tools), ResponseInputItem::McpToolCallOutput { output, .. } => { - output.code_mode_result(&ToolPayload::Mcp { - server: String::new(), - tool: String::new(), - raw_arguments: String::new(), + output.code_mode_result(&ToolPayload::Function { + arguments: String::new(), }) } } diff --git a/codex-rs/core/src/tools/context_tests.rs b/codex-rs/core/src/tools/context_tests.rs index 29adec5bd..6dc424831 100644 --- a/codex-rs/core/src/tools/context_tests.rs +++ b/codex-rs/core/src/tools/context_tests.rs @@ -61,10 +61,8 @@ fn mcp_code_mode_result_serializes_full_call_tool_result() { })), }; - let result = output.code_mode_result(&ToolPayload::Mcp { - server: "server".to_string(), - tool: "tool".to_string(), - raw_arguments: "{}".to_string(), + let result = output.code_mode_result(&ToolPayload::Function { + arguments: "{}".to_string(), }); assert_eq!( @@ -106,10 +104,8 @@ fn mcp_tool_output_response_item_includes_wall_time() { let response = output.to_response_item( "mcp-call-1", - &ToolPayload::Mcp { - server: "server".to_string(), - tool: "tool".to_string(), - raw_arguments: "{}".to_string(), + &ToolPayload::Function { + arguments: "{}".to_string(), }, ); @@ -160,10 +156,8 @@ fn mcp_tool_output_response_item_truncates_large_structured_content() { let response = output.to_response_item( "mcp-call-large", - &ToolPayload::Mcp { - server: "server".to_string(), - tool: "tool".to_string(), - raw_arguments: "{}".to_string(), + &ToolPayload::Function { + arguments: "{}".to_string(), }, ); @@ -205,10 +199,8 @@ fn mcp_tool_output_response_item_preserves_content_items() { let response = output.to_response_item( "mcp-call-2", - &ToolPayload::Mcp { - server: "server".to_string(), - tool: "tool".to_string(), - raw_arguments: "{}".to_string(), + &ToolPayload::Function { + arguments: "{}".to_string(), }, ); @@ -259,10 +251,8 @@ fn mcp_tool_output_code_mode_result_stays_raw_call_tool_result() { truncation_policy: TruncationPolicy::Bytes(64), }; - let result = output.code_mode_result(&ToolPayload::Mcp { - server: "server".to_string(), - tool: "tool".to_string(), - raw_arguments: "{}".to_string(), + let result = output.code_mode_result(&ToolPayload::Function { + arguments: "{}".to_string(), }); assert_eq!( diff --git a/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs b/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs index a7a36a49d..ad6815582 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs @@ -4,7 +4,6 @@ use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::handlers::agent_jobs_spec::create_report_agent_job_result_tool; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -23,10 +22,6 @@ impl ToolHandler for ReportAgentJobResultHandler { Some(create_report_agent_job_result_tool()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs b/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs index a1d0b5e7a..183387abf 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs/spawn_agents_on_csv.rs @@ -4,7 +4,6 @@ use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::handlers::agent_jobs_spec::create_spawn_agents_on_csv_tool; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -23,10 +22,6 @@ impl ToolHandler for SpawnAgentsOnCsvHandler { Some(create_spawn_agents_on_csv_tool()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/apply_patch.rs b/codex-rs/core/src/tools/handlers/apply_patch.rs index 34e466304..211862aa2 100644 --- a/codex-rs/core/src/tools/handlers/apply_patch.rs +++ b/codex-rs/core/src/tools/handlers/apply_patch.rs @@ -30,7 +30,6 @@ use crate::tools::registry::PostToolUsePayload; use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolArgumentDiffConsumer; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use crate::tools::runtimes::apply_patch::ApplyPatchRequest; use crate::tools::runtimes::apply_patch::ApplyPatchRuntime; use crate::tools::sandboxing::ToolCtx; @@ -307,10 +306,6 @@ impl ToolHandler for ApplyPatchHandler { Some(create_apply_patch_freeform_tool(self.multi_environment)) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Custom { .. }) } diff --git a/codex-rs/core/src/tools/handlers/dynamic.rs b/codex-rs/core/src/tools/handlers/dynamic.rs index 549edd514..59450a837 100644 --- a/codex-rs/core/src/tools/handlers/dynamic.rs +++ b/codex-rs/core/src/tools/handlers/dynamic.rs @@ -6,7 +6,6 @@ use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use crate::turn_timing::now_unix_timestamp_ms; use codex_protocol::dynamic_tools::DynamicToolCallRequest; use codex_protocol::dynamic_tools::DynamicToolResponse; @@ -36,10 +35,6 @@ impl ToolHandler for DynamicToolHandler { self.tool_name.clone() } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn is_mutating(&self, _invocation: &ToolInvocation) -> bool { true } diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index d41efe895..f76708478 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -17,7 +17,6 @@ use crate::tools::hook_names::HookToolName; use crate::tools::registry::PostToolUsePayload; use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; pub(crate) struct BundledToolOutput { value: Value, @@ -80,10 +79,6 @@ impl ToolHandler for BundledToolHandler { Some(self.spec.clone()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { self.arguments_from_payload(payload).is_some() } diff --git a/codex-rs/core/src/tools/handlers/goal/create_goal.rs b/codex-rs/core/src/tools/handlers/goal/create_goal.rs index 37ca319ad..db33445a2 100644 --- a/codex-rs/core/src/tools/handlers/goal/create_goal.rs +++ b/codex-rs/core/src/tools/handlers/goal/create_goal.rs @@ -7,7 +7,6 @@ use crate::tools::handlers::goal_spec::CREATE_GOAL_TOOL_NAME; use crate::tools::handlers::goal_spec::create_create_goal_tool; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -29,10 +28,6 @@ impl ToolHandler for CreateGoalHandler { Some(create_create_goal_tool()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, diff --git a/codex-rs/core/src/tools/handlers/goal/get_goal.rs b/codex-rs/core/src/tools/handlers/goal/get_goal.rs index f1af6dc5b..d3c587b3a 100644 --- a/codex-rs/core/src/tools/handlers/goal/get_goal.rs +++ b/codex-rs/core/src/tools/handlers/goal/get_goal.rs @@ -5,7 +5,6 @@ use crate::tools::context::ToolPayload; use crate::tools::handlers::goal_spec::GET_GOAL_TOOL_NAME; use crate::tools::handlers::goal_spec::create_get_goal_tool; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -26,10 +25,6 @@ impl ToolHandler for GetGoalHandler { Some(create_get_goal_tool()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, payload, .. diff --git a/codex-rs/core/src/tools/handlers/goal/update_goal.rs b/codex-rs/core/src/tools/handlers/goal/update_goal.rs index bdb2315a6..8e80a9013 100644 --- a/codex-rs/core/src/tools/handlers/goal/update_goal.rs +++ b/codex-rs/core/src/tools/handlers/goal/update_goal.rs @@ -8,7 +8,6 @@ use crate::tools::handlers::goal_spec::UPDATE_GOAL_TOOL_NAME; use crate::tools::handlers::goal_spec::create_update_goal_tool; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_protocol::protocol::ThreadGoalStatus; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -31,10 +30,6 @@ impl ToolHandler for UpdateGoalHandler { Some(create_update_goal_tool()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 35dbc3bc0..c9907f06c 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -8,22 +8,22 @@ use crate::tools::context::McpToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; -use crate::tools::flat_tool_name; use crate::tools::hook_names::HookToolName; use crate::tools::registry::PostToolUsePayload; use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; +use crate::tools::registry::ToolTelemetryTags; +use codex_mcp::ToolInfo; use codex_tools::ToolName; use serde_json::Value; pub struct McpHandler { - tool_name: ToolName, + tool_info: ToolInfo, } impl McpHandler { - pub fn new(tool_name: ToolName) -> Self { - Self { tool_name } + pub fn new(tool_info: ToolInfo) -> Self { + Self { tool_info } } } @@ -31,22 +31,29 @@ impl ToolHandler for McpHandler { type Output = McpToolOutput; fn tool_name(&self) -> ToolName { - self.tool_name.clone() + self.tool_info.canonical_tool_name() } - fn kind(&self) -> ToolKind { - ToolKind::Mcp + fn supports_parallel_tool_calls(&self) -> bool { + self.tool_info.supports_parallel_tool_calls + } + + async fn telemetry_tags(&self, _invocation: &ToolInvocation) -> ToolTelemetryTags { + let mut tags = vec![("mcp_server", self.tool_info.server_name.clone())]; + if let Some(origin) = &self.tool_info.server_origin { + tags.push(("mcp_server_origin", origin.clone())); + } + tags } fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option { - let ToolPayload::Mcp { raw_arguments, .. } = &invocation.payload else { + let ToolPayload::Function { arguments } = &invocation.payload else { return None; }; - let tool_name = &self.tool_name; Some(PreToolUsePayload { - tool_name: HookToolName::new(flat_tool_name(tool_name).into_owned()), - tool_input: mcp_hook_tool_input(raw_arguments), + tool_name: HookToolName::new(self.tool_name().to_string()), + tool_input: mcp_hook_tool_input(arguments), }) } @@ -55,15 +62,14 @@ impl ToolHandler for McpHandler { invocation: &ToolInvocation, result: &Self::Output, ) -> Option { - let ToolPayload::Mcp { .. } = &invocation.payload else { + let ToolPayload::Function { .. } = &invocation.payload else { return None; }; let tool_response = result.post_tool_use_response(&invocation.call_id, &invocation.payload)?; - let tool_name = &self.tool_name; Some(PostToolUsePayload { - tool_name: HookToolName::new(flat_tool_name(tool_name).into_owned()), + tool_name: HookToolName::new(self.tool_name().to_string()), tool_use_id: invocation.call_id.clone(), tool_input: result.tool_input.clone(), tool_response, @@ -80,11 +86,7 @@ impl ToolHandler for McpHandler { } = invocation; let payload = match payload { - ToolPayload::Mcp { - server, - tool, - raw_arguments, - } => (server, tool, raw_arguments), + ToolPayload::Function { arguments } => arguments, _ => { return Err(FunctionCallError::RespondToModel( "mcp handler received unsupported payload".to_string(), @@ -92,19 +94,15 @@ impl ToolHandler for McpHandler { } }; - let (server, tool, raw_arguments) = payload; - let arguments_str = raw_arguments; - let started = Instant::now(); - let hook_tool_name = flat_tool_name(&self.tool_name); let result = handle_mcp_tool_call( Arc::clone(&session), &turn, call_id.clone(), - server, - tool, - hook_tool_name.into_owned(), - arguments_str, + self.tool_info.server_name.clone(), + self.tool_info.tool.name.to_string(), + self.tool_name().to_string(), + payload, ) .await; @@ -139,10 +137,8 @@ mod tests { #[tokio::test] async fn mcp_pre_tool_use_payload_uses_model_tool_name_and_raw_args() { - let payload = ToolPayload::Mcp { - server: "memory".to_string(), - tool: "create_entities".to_string(), - raw_arguments: json!({ + let payload = ToolPayload::Function { + arguments: json!({ "entities": [{ "name": "Ada", "entityType": "person" @@ -151,10 +147,7 @@ mod tests { .to_string(), }; let (session, turn) = make_session_and_context().await; - let handler = McpHandler::new(codex_tools::ToolName::namespaced( - "mcp__memory__", - "create_entities", - )); + let handler = McpHandler::new(tool_info("memory", "mcp__memory__", "create_entities")); assert_eq!( handler.pre_tool_use_payload(&ToolInvocation { @@ -181,10 +174,8 @@ mod tests { #[tokio::test] async fn mcp_post_tool_use_payload_uses_model_tool_name_args_and_result() { - let payload = ToolPayload::Mcp { - server: "filesystem".to_string(), - tool: "read_file".to_string(), - raw_arguments: json!({ "path": "/tmp/notes.txt" }).to_string(), + let payload = ToolPayload::Function { + arguments: json!({ "path": "/tmp/notes.txt" }).to_string(), }; let output = McpToolOutput { result: codex_protocol::mcp::CallToolResult { @@ -206,10 +197,7 @@ mod tests { truncation_policy: codex_utils_output_truncation::TruncationPolicy::Bytes(1024), }; let (session, turn) = make_session_and_context().await; - let handler = McpHandler::new(codex_tools::ToolName::namespaced( - "mcp__filesystem__", - "read_file", - )); + let handler = McpHandler::new(tool_info("filesystem", "mcp__filesystem__", "read_file")); let invocation = ToolInvocation { session: session.into(), turn: turn.into(), @@ -245,4 +233,31 @@ mod tests { fn mcp_hook_tool_input_defaults_empty_args_to_object() { assert_eq!(mcp_hook_tool_input(" "), json!({})); } + + fn tool_info(server_name: &str, callable_namespace: &str, tool_name: &str) -> ToolInfo { + ToolInfo { + server_name: server_name.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, + callable_name: tool_name.to_string(), + callable_namespace: callable_namespace.to_string(), + namespace_description: None, + tool: rmcp::model::Tool { + name: tool_name.to_string().into(), + title: None, + description: None, + input_schema: Arc::new(rmcp::model::object(serde_json::json!({ + "type": "object", + }))), + output_schema: None, + annotations: None, + execution: None, + icons: None, + meta: None, + }, + connector_id: None, + connector_name: None, + plugin_display_names: Vec::new(), + } + } } diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs index 5e42bc3c0..14a5a9d2b 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs @@ -6,7 +6,6 @@ use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::handlers::mcp_resource_spec::create_list_mcp_resource_templates_tool; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_protocol::models::function_call_output_content_items_to_text; use codex_protocol::protocol::McpInvocation; use codex_tools::ToolName; @@ -41,10 +40,6 @@ impl ToolHandler for ListMcpResourceTemplatesHandler { true } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - #[expect( clippy::await_holding_invalid_type, reason = "MCP resource template listing reads through the session-owned manager guard" diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs index 77b224cd1..eaa5a54e0 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs @@ -6,7 +6,6 @@ use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::handlers::mcp_resource_spec::create_list_mcp_resources_tool; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_protocol::models::function_call_output_content_items_to_text; use codex_protocol::protocol::McpInvocation; use codex_tools::ToolName; @@ -41,10 +40,6 @@ impl ToolHandler for ListMcpResourcesHandler { true } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - #[expect( clippy::await_holding_invalid_type, reason = "MCP resource listing reads through the session-owned manager guard" diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs index 50944aefb..51bb6fa67 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs @@ -6,7 +6,6 @@ use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::handlers::mcp_resource_spec::create_read_mcp_resource_tool; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_protocol::models::function_call_output_content_items_to_text; use codex_protocol::protocol::McpInvocation; use codex_tools::ToolName; @@ -41,10 +40,6 @@ impl ToolHandler for ReadMcpResourceHandler { true } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, diff --git a/codex-rs/core/src/tools/handlers/multi_agents.rs b/codex-rs/core/src/tools/handlers/multi_agents.rs index 71ef84fd1..126263396 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents.rs @@ -16,7 +16,6 @@ use crate::tools::context::ToolPayload; pub(crate) use crate::tools::handlers::multi_agents_common::*; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_protocol::ThreadId; use codex_protocol::models::ResponseInputItem; use codex_protocol::openai_models::ReasoningEffort; diff --git a/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs index 7d47290c1..e4788dd9b 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs @@ -16,10 +16,6 @@ impl ToolHandler for Handler { Some(create_close_agent_tool_v1()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs index 0b86c9abd..fa77bc37b 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/resume_agent.rs @@ -18,10 +18,6 @@ impl ToolHandler for Handler { Some(create_resume_agent_tool()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs b/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs index cc65a00f4..cdd46a90d 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs @@ -17,10 +17,6 @@ impl ToolHandler for Handler { Some(create_send_input_tool_v1()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs index 40e9cc5d3..031765895 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -33,10 +33,6 @@ impl ToolHandler for Handler { Some(create_spawn_agent_tool_v1(self.options.clone())) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/multi_agents/wait.rs b/codex-rs/core/src/tools/handlers/multi_agents/wait.rs index 9e63450a6..17739606f 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/wait.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/wait.rs @@ -38,10 +38,6 @@ impl ToolHandler for Handler { Some(create_wait_agent_tool_v1(self.options)) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2.rs index a477c25ca..b40d85523 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2.rs @@ -9,7 +9,6 @@ use crate::tools::context::ToolPayload; use crate::tools::handlers::multi_agents_common::*; use crate::tools::handlers::parse_arguments; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_protocol::AgentPath; use codex_protocol::models::ResponseInputItem; use codex_protocol::openai_models::ReasoningEffort; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs index d3a290d36..6f2984c33 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/close_agent.rs @@ -16,10 +16,6 @@ impl ToolHandler for Handler { Some(create_close_agent_tool_v2()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/followup_task.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/followup_task.rs index 147d238b7..2077d6693 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/followup_task.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/followup_task.rs @@ -19,10 +19,6 @@ impl ToolHandler for Handler { Some(create_followup_task_tool()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/list_agents.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/list_agents.rs index 37365f2a5..b75f28675 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/list_agents.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/list_agents.rs @@ -16,10 +16,6 @@ impl ToolHandler for Handler { Some(create_list_agents_tool()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/send_message.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/send_message.rs index 38a0cc3ab..c9e206070 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/send_message.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/send_message.rs @@ -19,10 +19,6 @@ impl ToolHandler for Handler { Some(create_send_message_tool()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs index fc44d3df4..6642c5e4b 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs @@ -35,10 +35,6 @@ impl ToolHandler for Handler { Some(create_spawn_agent_tool_v2(self.options.clone())) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs index d5fbe49c7..15790e4ce 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs @@ -30,10 +30,6 @@ impl ToolHandler for Handler { Some(create_wait_agent_tool_v2(self.options)) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/plan.rs b/codex-rs/core/src/tools/handlers/plan.rs index 2995fbbc2..a3c8aa1be 100644 --- a/codex-rs/core/src/tools/handlers/plan.rs +++ b/codex-rs/core/src/tools/handlers/plan.rs @@ -4,7 +4,6 @@ use crate::tools::context::ToolOutput; use crate::tools::context::ToolPayload; use crate::tools::handlers::plan_spec::create_update_plan_tool; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_protocol::config_types::ModeKind; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; @@ -55,10 +54,6 @@ impl ToolHandler for PlanHandler { Some(create_update_plan_tool()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, diff --git a/codex-rs/core/src/tools/handlers/request_permissions.rs b/codex-rs/core/src/tools/handlers/request_permissions.rs index 01d182057..59c117905 100644 --- a/codex-rs/core/src/tools/handlers/request_permissions.rs +++ b/codex-rs/core/src/tools/handlers/request_permissions.rs @@ -9,7 +9,6 @@ use crate::tools::handlers::parse_arguments_with_base_path; use crate::tools::handlers::shell_spec::create_request_permissions_tool; use crate::tools::handlers::shell_spec::request_permissions_tool_description; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -28,10 +27,6 @@ impl ToolHandler for RequestPermissionsHandler { )) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install.rs b/codex-rs/core/src/tools/handlers/request_plugin_install.rs index 7a9ace384..c1dbdee31 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install.rs @@ -35,7 +35,6 @@ use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::handlers::request_plugin_install_spec::create_request_plugin_install_tool; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; #[derive(Default)] pub struct RequestPluginInstallHandler { @@ -65,10 +64,6 @@ impl ToolHandler for RequestPluginInstallHandler { true } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - #[expect( clippy::await_holding_invalid_type, reason = "plugin install discovery reads through the session-owned manager guard" 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 6d2623485..583f62c0c 100644 --- a/codex-rs/core/src/tools/handlers/request_user_input.rs +++ b/codex-rs/core/src/tools/handlers/request_user_input.rs @@ -9,7 +9,6 @@ use crate::tools::handlers::request_user_input_spec::normalize_request_user_inpu use crate::tools::handlers::request_user_input_spec::request_user_input_tool_description; use crate::tools::handlers::request_user_input_spec::request_user_input_unavailable_message; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_protocol::config_types::ModeKind; use codex_protocol::request_user_input::RequestUserInputArgs; use codex_tools::ToolName; @@ -32,10 +31,6 @@ impl ToolHandler for RequestUserInputHandler { )) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { session, diff --git a/codex-rs/core/src/tools/handlers/shell/container_exec.rs b/codex-rs/core/src/tools/handlers/shell/container_exec.rs index 0ada4a16f..969b34af1 100644 --- a/codex-rs/core/src/tools/handlers/shell/container_exec.rs +++ b/codex-rs/core/src/tools/handlers/shell/container_exec.rs @@ -11,7 +11,6 @@ use crate::tools::handlers::resolve_workdir_base_path; use crate::tools::registry::PostToolUsePayload; use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use crate::tools::runtimes::shell::ShellRuntimeBackend; use super::RunExecLikeArgs; @@ -29,10 +28,6 @@ impl ToolHandler for ContainerExecHandler { ToolName::plain("container.exec") } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/shell/local_shell.rs b/codex-rs/core/src/tools/handlers/shell/local_shell.rs index 1bb765e6d..c8b44951f 100644 --- a/codex-rs/core/src/tools/handlers/shell/local_shell.rs +++ b/codex-rs/core/src/tools/handlers/shell/local_shell.rs @@ -10,7 +10,6 @@ use crate::tools::hook_names::HookToolName; use crate::tools::registry::PostToolUsePayload; use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use crate::tools::runtimes::shell::ShellRuntimeBackend; use codex_tools::ToolSpec; @@ -46,10 +45,6 @@ impl ToolHandler for LocalShellHandler { self.include_spec } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::LocalShell { .. }) } diff --git a/codex-rs/core/src/tools/handlers/shell/shell_command.rs b/codex-rs/core/src/tools/handlers/shell/shell_command.rs index e64c4666c..91238b40d 100644 --- a/codex-rs/core/src/tools/handlers/shell/shell_command.rs +++ b/codex-rs/core/src/tools/handlers/shell/shell_command.rs @@ -21,7 +21,6 @@ use crate::tools::hook_names::HookToolName; use crate::tools::registry::PostToolUsePayload; use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use crate::tools::runtimes::shell::ShellRuntimeBackend; use codex_tools::ToolSpec; @@ -144,10 +143,6 @@ impl ToolHandler for ShellCommandHandler { self.options.is_some() } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/shell/shell_handler.rs b/codex-rs/core/src/tools/handlers/shell/shell_handler.rs index 5a5101957..112fd22c3 100644 --- a/codex-rs/core/src/tools/handlers/shell/shell_handler.rs +++ b/codex-rs/core/src/tools/handlers/shell/shell_handler.rs @@ -16,7 +16,6 @@ use crate::tools::handlers::resolve_workdir_base_path; use crate::tools::registry::PostToolUsePayload; use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use crate::tools::runtimes::shell::ShellRuntimeBackend; use codex_tools::ToolSpec; @@ -78,10 +77,6 @@ impl ToolHandler for ShellHandler { self.options.is_some() } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/test_sync.rs b/codex-rs/core/src/tools/handlers/test_sync.rs index 6254b9482..6ef39ffd1 100644 --- a/codex-rs/core/src/tools/handlers/test_sync.rs +++ b/codex-rs/core/src/tools/handlers/test_sync.rs @@ -15,7 +15,6 @@ use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::handlers::test_sync_spec::create_test_sync_tool; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -71,10 +70,6 @@ impl ToolHandler for TestSyncHandler { true } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { payload, .. } = invocation; diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index 70410db3a..2883f92f9 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -4,7 +4,6 @@ use crate::tools::context::ToolPayload; use crate::tools::context::ToolSearchOutput; use crate::tools::handlers::tool_search_spec::create_tool_search_tool; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use crate::tools::tool_search_entry::ToolSearchEntry; use bm25::Document; use bm25::Language; @@ -68,10 +67,6 @@ impl ToolHandler for ToolSearchHandler { true } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle( &self, invocation: ToolInvocation, @@ -399,6 +394,8 @@ mod tests { fn tool_info(server_name: &str, tool_name: &str, description_prefix: &str) -> ToolInfo { ToolInfo { server_name: server_name.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, callable_name: tool_name.to_string(), callable_namespace: format!("mcp__{server_name}__"), namespace_description: None, diff --git a/codex-rs/core/src/tools/handlers/unavailable_tool.rs b/codex-rs/core/src/tools/handlers/unavailable_tool.rs index 74e74305f..e297374a7 100644 --- a/codex-rs/core/src/tools/handlers/unavailable_tool.rs +++ b/codex-rs/core/src/tools/handlers/unavailable_tool.rs @@ -3,7 +3,6 @@ use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -48,10 +47,6 @@ impl ToolHandler for UnavailableToolHandler { self.spec.clone() } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle(&self, invocation: ToolInvocation) -> Result { let ToolInvocation { payload, .. } = invocation; diff --git a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs index 959ea1f6f..d18764a93 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs @@ -16,7 +16,6 @@ use crate::tools::hook_names::HookToolName; use crate::tools::registry::PostToolUsePayload; use crate::tools::registry::PreToolUsePayload; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use crate::unified_exec::ExecCommandRequest; use crate::unified_exec::UnifiedExecContext; use crate::unified_exec::UnifiedExecError; @@ -88,10 +87,6 @@ impl ToolHandler for ExecCommandHandler { true } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs b/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs index b7be04f2a..3250a9d0a 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec/write_stdin.rs @@ -5,7 +5,6 @@ use crate::tools::context::ToolPayload; use crate::tools::handlers::parse_arguments; use crate::tools::registry::PostToolUsePayload; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use crate::unified_exec::WriteStdinRequest; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::TerminalInteractionEvent; @@ -42,10 +41,6 @@ impl ToolHandler for WriteStdinHandler { Some(create_write_stdin_tool()) } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!(payload, ToolPayload::Function { .. }) } diff --git a/codex-rs/core/src/tools/handlers/view_image.rs b/codex-rs/core/src/tools/handlers/view_image.rs index 6e4a25c0a..4ec377d91 100644 --- a/codex-rs/core/src/tools/handlers/view_image.rs +++ b/codex-rs/core/src/tools/handlers/view_image.rs @@ -21,7 +21,6 @@ use crate::tools::handlers::resolve_tool_environment; use crate::tools::handlers::view_image_spec::ViewImageToolOptions; use crate::tools::handlers::view_image_spec::create_view_image_tool; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use codex_tools::ToolName; use codex_tools::ToolSpec; @@ -77,10 +76,6 @@ impl ToolHandler for ViewImageHandler { true } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle(&self, invocation: ToolInvocation) -> Result { if !invocation .turn diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 5193d1591..f1f9db0dd 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -32,11 +32,7 @@ use futures::future::BoxFuture; use serde_json::Value; use tracing::warn; -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -pub enum ToolKind { - Function, - Mcp, -} +pub(crate) type ToolTelemetryTags = Vec<(&'static str, String)>; pub trait ToolHandler: Send + Sync { type Output: ToolOutput + 'static; @@ -52,17 +48,20 @@ pub trait ToolHandler: Send + Sync { false } - fn kind(&self) -> ToolKind; - fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!( - (self.kind(), payload), - (ToolKind::Function, ToolPayload::Function { .. }) - | (ToolKind::Function, ToolPayload::ToolSearch { .. }) - | (ToolKind::Mcp, ToolPayload::Mcp { .. }) + payload, + ToolPayload::Function { .. } | ToolPayload::ToolSearch { .. } ) } + fn telemetry_tags( + &self, + _invocation: &ToolInvocation, + ) -> impl std::future::Future + Send { + async { Vec::new() } + } + /// Returns `true` if the [ToolInvocation] *might* mutate the environment of the /// user (through file system, OS operations, ...). /// This function must remains defensive and return `true` if a doubt exist on the @@ -168,12 +167,19 @@ pub(crate) struct PostToolUsePayload { } trait AnyToolHandler: Send + Sync { + fn supports_parallel_tool_calls(&self) -> bool; + fn matches_kind(&self, payload: &ToolPayload) -> bool; fn is_mutating<'a>(&'a self, invocation: &'a ToolInvocation) -> BoxFuture<'a, bool>; fn pre_tool_use_payload(&self, invocation: &ToolInvocation) -> Option; + fn telemetry_tags<'a>( + &'a self, + invocation: &'a ToolInvocation, + ) -> BoxFuture<'a, ToolTelemetryTags>; + fn create_diff_consumer(&self) -> Option>; fn handle_any<'a>( &'a self, @@ -185,6 +191,10 @@ impl AnyToolHandler for T where T: ToolHandler, { + fn supports_parallel_tool_calls(&self) -> bool { + ToolHandler::supports_parallel_tool_calls(self) + } + fn matches_kind(&self, payload: &ToolPayload) -> bool { ToolHandler::matches_kind(self, payload) } @@ -197,6 +207,13 @@ where ToolHandler::pre_tool_use_payload(self, invocation) } + fn telemetry_tags<'a>( + &'a self, + invocation: &'a ToolInvocation, + ) -> BoxFuture<'a, ToolTelemetryTags> { + Box::pin(ToolHandler::telemetry_tags(self, invocation)) + } + fn create_diff_consumer(&self) -> Option> { ToolHandler::create_diff_consumer(self) } @@ -259,6 +276,10 @@ impl ToolRegistry { self.handler(name)?.create_diff_consumer() } + pub(crate) fn supports_parallel_tool_calls(&self, name: &ToolName) -> Option { + Some(self.handler(name)?.supports_parallel_tool_calls()) + } + #[expect( clippy::await_holding_invalid_type, reason = "tool dispatch must keep active-turn accounting atomic" @@ -273,7 +294,7 @@ impl ToolRegistry { let otel = invocation.turn.session_telemetry.clone(); let payload_for_response = invocation.payload.clone(); let log_payload = payload_for_response.log_payload(); - let metric_tags = [ + let base_tool_result_tags = [ ( "sandbox", permission_profile_sandbox_tag( @@ -290,21 +311,6 @@ impl ToolRegistry { ), ), ]; - let (mcp_server, mcp_server_origin) = match &invocation.payload { - ToolPayload::Mcp { server, .. } => { - let manager = invocation - .session - .services - .mcp_connection_manager - .read() - .await; - let origin = manager.server_origin(server).map(str::to_owned); - (Some(server.clone()), origin) - } - _ => (None, None), - }; - let mcp_server_ref = mcp_server.as_deref(); - let mcp_server_origin_ref = mcp_server_origin.as_deref(); { let mut active = invocation.session.active_turn.lock().await; @@ -315,7 +321,6 @@ impl ToolRegistry { } let dispatch_trace = ToolDispatchTrace::start(&invocation); - let handler = match self.handler(&tool_name) { Some(handler) => handler, None => { @@ -327,9 +332,8 @@ impl ToolRegistry { Duration::ZERO, /*success*/ false, &message, - &metric_tags, - mcp_server_ref, - mcp_server_origin_ref, + &base_tool_result_tags, + /*extra_trace_fields*/ &[], ); let err = FunctionCallError::RespondToModel(message); dispatch_trace.record_failed(&err); @@ -337,6 +341,19 @@ impl ToolRegistry { } }; + let telemetry_tags = handler.telemetry_tags(&invocation).await; + let mut tool_result_tags = + Vec::with_capacity(base_tool_result_tags.len() + telemetry_tags.len()); + let mut extra_trace_fields = Vec::new(); + tool_result_tags.extend_from_slice(&base_tool_result_tags); + for (key, value) in &telemetry_tags { + if matches!(*key, "mcp_server" | "mcp_server_origin") { + extra_trace_fields.push((*key, value.as_str())); + } else { + tool_result_tags.push((*key, value.as_str())); + } + } + if !handler.matches_kind(&invocation.payload) { let message = format!("tool {tool_name} invoked with incompatible payload"); otel.tool_result_with_tags( @@ -346,9 +363,8 @@ impl ToolRegistry { Duration::ZERO, /*success*/ false, &message, - &metric_tags, - mcp_server_ref, - mcp_server_origin_ref, + &tool_result_tags, + &extra_trace_fields, ); let err = FunctionCallError::Fatal(message); dispatch_trace.record_failed(&err); @@ -379,9 +395,8 @@ impl ToolRegistry { tool_name_flat.as_ref(), &call_id_owned, log_payload.as_ref(), - &metric_tags, - mcp_server_ref, - mcp_server_origin_ref, + &tool_result_tags, + &extra_trace_fields, || { let handler = handler.clone(); let response_cell = &response_cell; @@ -406,7 +421,7 @@ impl ToolRegistry { ) .await; let success = match &result { - Ok((_preview, success)) => *success, + Ok((_, success)) => *success, Err(_) => false, }; emit_metric_for_tool_read(&invocation, success).await; diff --git a/codex-rs/core/src/tools/registry_tests.rs b/codex-rs/core/src/tools/registry_tests.rs index d445b196a..be5223a11 100644 --- a/codex-rs/core/src/tools/registry_tests.rs +++ b/codex-rs/core/src/tools/registry_tests.rs @@ -15,10 +15,6 @@ impl ToolHandler for TestHandler { self.tool_name.clone() } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle(&self, _invocation: ToolInvocation) -> Result { Ok(crate::tools::context::FunctionToolOutput::from_text( "ok".to_string(), diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 4649caa04..51fd40a26 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -40,14 +40,12 @@ pub struct ToolRouter { registry: ToolRegistry, specs: Vec, model_visible_specs: Vec, - parallel_mcp_server_names: HashSet, } pub(crate) struct ToolRouterParams<'a> { pub(crate) mcp_tools: Option>, pub(crate) deferred_mcp_tools: Option>, pub(crate) unavailable_called_tools: Vec, - pub(crate) parallel_mcp_server_names: HashSet, pub(crate) discoverable_tools: Option>, pub(crate) extension_tool_bundles: Vec, pub(crate) dynamic_tools: &'a [DynamicToolSpec], @@ -59,7 +57,6 @@ impl ToolRouter { mcp_tools, deferred_mcp_tools, unavailable_called_tools, - parallel_mcp_server_names, discoverable_tools, extension_tool_bundles, dynamic_tools, @@ -99,7 +96,6 @@ impl ToolRouter { registry, specs, model_visible_specs, - parallel_mcp_server_names, } } @@ -166,20 +162,13 @@ impl ToolRouter { } pub fn tool_supports_parallel(&self, call: &ToolCall) -> bool { - match &call.payload { - // MCP parallel support is configured per server, including for deferred - // tools that may not have a matching spec entry. Use the parsed payload - // server so similarly named servers/tools cannot collide. - ToolPayload::Mcp { server, .. } => self.parallel_mcp_server_names.contains(server), - _ => self.configured_tool_supports_parallel(&call.tool_name), - } + self.registry + .supports_parallel_tool_calls(&call.tool_name) + .unwrap_or_else(|| self.configured_tool_supports_parallel(&call.tool_name)) } #[instrument(level = "trace", skip_all, err)] - pub async fn build_tool_call( - session: &Session, - item: ResponseItem, - ) -> Result, FunctionCallError> { + pub fn build_tool_call(item: ResponseItem) -> Result, FunctionCallError> { match item { ResponseItem::FunctionCall { name, @@ -189,23 +178,11 @@ impl ToolRouter { .. } => { let tool_name = ToolName::new(namespace, name); - if let Some(tool_info) = session.resolve_mcp_tool_info(&tool_name).await { - Ok(Some(ToolCall { - tool_name: tool_info.canonical_tool_name(), - call_id, - payload: ToolPayload::Mcp { - server: tool_info.server_name, - tool: tool_info.tool.name.to_string(), - raw_arguments: arguments, - }, - })) - } else { - Ok(Some(ToolCall { - tool_name, - call_id, - payload: ToolPayload::Function { arguments }, - })) - } + Ok(Some(ToolCall { + tool_name, + call_id, + payload: ToolPayload::Function { arguments }, + })) } ResponseItem::ToolSearchCall { call_id: Some(call_id), diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index 1eacbcf6c..dc06f7ed0 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -1,4 +1,3 @@ -use std::collections::HashSet; use std::sync::Arc; use crate::config::Config; @@ -99,7 +98,6 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow deferred_mcp_tools: None, mcp_tools: Some(mcp_tools), unavailable_called_tools: Vec::new(), - parallel_mcp_server_names: HashSet::new(), discoverable_tools: None, extension_tool_bundles: Vec::new(), dynamic_tools: turn.dynamic_tools.as_slice(), @@ -132,21 +130,15 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow #[tokio::test] async fn build_tool_call_uses_namespace_for_registry_name() -> anyhow::Result<()> { - let (session, _) = make_session_and_context().await; - let session = Arc::new(session); let tool_name = "create_event".to_string(); - let call = ToolRouter::build_tool_call( - &session, - ResponseItem::FunctionCall { - id: None, - name: tool_name.clone(), - namespace: Some("mcp__codex_apps__calendar".to_string()), - arguments: "{}".to_string(), - call_id: "call-namespace".to_string(), - }, - ) - .await? + let call = ToolRouter::build_tool_call(ResponseItem::FunctionCall { + id: None, + name: tool_name.clone(), + namespace: Some("mcp__codex_apps__calendar".to_string()), + arguments: "{}".to_string(), + call_id: "call-namespace".to_string(), + })? .expect("function_call should produce a tool call"); assert_eq!( @@ -171,9 +163,21 @@ async fn mcp_parallel_support_uses_exact_payload_server() -> anyhow::Result<()> &turn.tools_config, ToolRouterParams { deferred_mcp_tools: None, - mcp_tools: None, + mcp_tools: Some(vec![ + mcp_tool_info( + "echo", + /*supports_parallel_tool_calls*/ true, + "mcp__echo__", + "query_with_delay", + ), + mcp_tool_info( + "hello_echo", + /*supports_parallel_tool_calls*/ false, + "mcp__hello_echo__", + "query_with_delay", + ), + ]), unavailable_called_tools: Vec::new(), - parallel_mcp_server_names: HashSet::from(["echo".to_string()]), discoverable_tools: None, extension_tool_bundles: Vec::new(), dynamic_tools: turn.dynamic_tools.as_slice(), @@ -183,10 +187,8 @@ async fn mcp_parallel_support_uses_exact_payload_server() -> anyhow::Result<()> let deferred_call = ToolCall { tool_name: ToolName::namespaced("mcp__echo__", "query_with_delay"), call_id: "call-deferred".to_string(), - payload: ToolPayload::Mcp { - server: "echo".to_string(), - tool: "query_with_delay".to_string(), - raw_arguments: "{}".to_string(), + payload: ToolPayload::Function { + arguments: "{}".to_string(), }, }; assert!(router.tool_supports_parallel(&deferred_call)); @@ -194,10 +196,8 @@ async fn mcp_parallel_support_uses_exact_payload_server() -> anyhow::Result<()> let different_server_call = ToolCall { tool_name: ToolName::namespaced("mcp__hello_echo__", "query_with_delay"), call_id: "call-other-server".to_string(), - payload: ToolPayload::Mcp { - server: "hello_echo".to_string(), - tool: "query_with_delay".to_string(), - raw_arguments: "{}".to_string(), + payload: ToolPayload::Function { + arguments: "{}".to_string(), }, }; assert!(!router.tool_supports_parallel(&different_server_call)); @@ -241,7 +241,6 @@ async fn model_visible_specs_filter_deferred_dynamic_tools() -> anyhow::Result<( deferred_mcp_tools: None, mcp_tools: None, unavailable_called_tools: Vec::new(), - parallel_mcp_server_names: HashSet::new(), discoverable_tools: None, extension_tool_bundles: Vec::new(), dynamic_tools: &dynamic_tools, @@ -265,6 +264,38 @@ async fn model_visible_specs_filter_deferred_dynamic_tools() -> anyhow::Result<( Ok(()) } +fn mcp_tool_info( + server_name: &str, + supports_parallel_tool_calls: bool, + callable_namespace: &str, + tool_name: &str, +) -> codex_mcp::ToolInfo { + codex_mcp::ToolInfo { + server_name: server_name.to_string(), + supports_parallel_tool_calls, + server_origin: None, + callable_name: tool_name.to_string(), + callable_namespace: callable_namespace.to_string(), + namespace_description: None, + tool: rmcp::model::Tool { + name: tool_name.to_string().into(), + title: None, + description: Some("Test MCP tool".to_string().into()), + input_schema: Arc::new(rmcp::model::object(json!({ + "type": "object", + }))), + output_schema: None, + annotations: None, + execution: None, + icons: None, + meta: None, + }, + connector_id: None, + connector_name: None, + plugin_display_names: Vec::new(), + } +} + #[tokio::test] async fn extension_tool_bundles_are_model_visible_and_dispatchable() -> anyhow::Result<()> { let (mut session, turn) = make_session_and_context().await; @@ -276,7 +307,6 @@ async fn extension_tool_bundles_are_model_visible_and_dispatchable() -> anyhow:: deferred_mcp_tools: None, mcp_tools: None, unavailable_called_tools: Vec::new(), - parallel_mcp_server_names: HashSet::new(), discoverable_tools: None, extension_tool_bundles: extension_tool_bundles(&session), dynamic_tools: turn.dynamic_tools.as_slice(), @@ -297,17 +327,13 @@ async fn extension_tool_bundles_are_model_visible_and_dispatchable() -> anyhow:: "expected extension-provided tool to be visible to the model" ); - let call = ToolRouter::build_tool_call( - &session, - ResponseItem::FunctionCall { - id: None, - name: "extension_echo".to_string(), - namespace: None, - arguments: json!({ "message": "hello" }).to_string(), - call_id: "call-extension".to_string(), - }, - ) - .await? + let call = ToolRouter::build_tool_call(ResponseItem::FunctionCall { + id: None, + name: "extension_echo".to_string(), + namespace: None, + arguments: json!({ "message": "hello" }).to_string(), + call_id: "call-extension".to_string(), + })? .expect("function_call should produce a tool call"); let result = router diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 2faac645a..fc21ff337 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -8,8 +8,6 @@ use crate::tools::handlers::multi_agents_spec::WaitAgentTimeoutOptions; use crate::tools::registry::ToolRegistryBuilder; use crate::tools::spec_plan::build_tool_registry_builder; use crate::tools::spec_plan_types::ToolNamespace; -use crate::tools::spec_plan_types::ToolRegistryBuildDeferredTool; -use crate::tools::spec_plan_types::ToolRegistryBuildMcpTool; use crate::tools::spec_plan_types::ToolRegistryBuildParams; use codex_mcp::ToolInfo; use codex_protocol::dynamic_tools::DynamicToolSpec; @@ -35,20 +33,14 @@ pub(crate) fn tool_user_shell_type(user_shell: &Shell) -> ToolUserShellType { } } -struct McpToolPlanInputs<'a> { - mcp_tools: Vec>, +struct McpToolPlanInputs { + mcp_tools: Vec, tool_namespaces: HashMap, } -fn map_mcp_tools_for_plan(mcp_tools: &[ToolInfo]) -> McpToolPlanInputs<'_> { +fn map_mcp_tools_for_plan(mcp_tools: &[ToolInfo]) -> McpToolPlanInputs { McpToolPlanInputs { - mcp_tools: mcp_tools - .iter() - .map(|tool| ToolRegistryBuildMcpTool { - name: tool.canonical_tool_name(), - tool: &tool.tool, - }) - .collect(), + mcp_tools: mcp_tools.to_vec(), tool_namespaces: mcp_tools .iter() .map(|tool| { @@ -78,17 +70,7 @@ pub(crate) fn build_specs_with_discoverable_tools( use crate::tools::tool_search_entry::build_tool_search_entries_for_config; let mcp_tool_plan_inputs = mcp_tools.as_deref().map(map_mcp_tools_for_plan); - let deferred_mcp_tool_sources = deferred_mcp_tools.as_ref().map(|tools| { - tools - .iter() - .map(|tool| ToolRegistryBuildDeferredTool { - name: tool.canonical_tool_name(), - server_name: tool.server_name.as_str(), - connector_name: tool.connector_name.as_deref(), - description: tool.namespace_description.as_deref(), - }) - .collect::>() - }); + let deferred_mcp_tool_sources = deferred_mcp_tools.as_deref(); let default_agent_type_description = crate::agent::role::spawn_tool_spec::build(&std::collections::BTreeMap::new()); let min_wait_timeout_ms = if config.multi_agent_v2 { @@ -108,7 +90,7 @@ pub(crate) fn build_specs_with_discoverable_tools( .collect::>(); let tool_search_entries = build_tool_search_entries_for_config( config, - deferred_mcp_tools.as_deref(), + deferred_mcp_tool_sources, &deferred_dynamic_tools, ); let mut builder = build_tool_registry_builder( @@ -117,7 +99,7 @@ pub(crate) fn build_specs_with_discoverable_tools( mcp_tools: mcp_tool_plan_inputs .as_ref() .map(|inputs| inputs.mcp_tools.as_slice()), - deferred_mcp_tools: deferred_mcp_tool_sources.as_deref(), + deferred_mcp_tools: deferred_mcp_tool_sources, tool_namespaces: mcp_tool_plan_inputs .as_ref() .map(|inputs| &inputs.tool_namespaces), diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 7fded6140..b07a572ef 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -47,6 +47,7 @@ use crate::tools::hosted_spec::create_web_search_tool; use crate::tools::registry::ToolRegistryBuilder; use crate::tools::spec_plan_types::ToolRegistryBuildParams; use crate::tools::spec_plan_types::agent_type_description; +use codex_mcp::ToolInfo; use codex_protocol::openai_models::ConfigShellToolType; use codex_tools::ResponsesApiNamespace; use codex_tools::ResponsesApiNamespaceTool; @@ -227,9 +228,9 @@ pub fn build_tool_registry_builder( .map(|deferred_mcp_tools| { collect_tool_search_source_infos(deferred_mcp_tools.iter().map(|tool| { ToolSearchSource { - server_name: tool.server_name, - connector_name: tool.connector_name, - description: tool.description, + server_name: tool.server_name.as_str(), + connector_name: tool.connector_name.as_deref(), + description: tool.namespace_description.as_deref(), } })) }) @@ -340,24 +341,26 @@ pub fn build_tool_registry_builder( } if let Some(mcp_tools) = params.mcp_tools { - let mut entries = mcp_tools.to_vec(); - entries.sort_by(|a, b| a.name.cmp(&b.name)); + let mut entries = mcp_tools + .iter() + .map(|tool| (tool.canonical_tool_name(), tool)) + .collect::>(); + entries.sort_by(|(left, _), (right, _)| left.cmp(right)); let mut namespace_entries = BTreeMap::new(); - for tool in entries { - let Some(namespace) = tool.name.namespace.as_ref() else { - let tool_name = &tool.name; + for (tool_name, tool) in entries { + let Some(namespace) = tool_name.namespace.as_ref() else { tracing::error!("Skipping MCP tool `{tool_name}`: MCP tools must be namespaced"); continue; }; namespace_entries .entry(namespace.clone()) .or_insert_with(Vec::new) - .push(tool); + .push((tool_name, tool)); } for (namespace, mut entries) in namespace_entries { - entries.sort_by_key(|tool| tool.name.name.clone()); + entries.sort_by_key(|(tool_name, _)| tool_name.name.clone()); let tool_namespace = params .tool_namespaces .and_then(|namespaces| namespaces.get(&namespace)); @@ -373,14 +376,13 @@ pub fn build_tool_registry_builder( default_namespace_description(namespace_name) }); let mut tools = Vec::new(); - for tool in entries { - match mcp_tool_to_responses_api_tool(&tool.name, tool.tool) { + for (tool_name, tool) in entries { + match mcp_tool_to_responses_api_tool(&tool_name, &tool.tool) { Ok(converted_tool) => { tools.push(ResponsesApiNamespaceTool::Function(converted_tool)); - builder.register_handler(Arc::new(McpHandler::new(tool.name))); + builder.register_handler(Arc::new(McpHandler::new(tool.clone()))); } Err(error) => { - let tool_name = &tool.name; tracing::error!( "Failed to convert `{tool_name}` MCP tool to OpenAI tool: {error:?}" ); @@ -429,11 +431,11 @@ pub fn build_tool_registry_builder( .mcp_tools .into_iter() .flatten() - .map(|direct| direct.name.clone()) + .map(ToolInfo::canonical_tool_name) .collect::>(); for tool in deferred_mcp_tools { - if !directly_registered_mcp_tools.contains(&tool.name) { - builder.register_handler(Arc::new(McpHandler::new(tool.name.clone()))); + if !directly_registered_mcp_tools.contains(&tool.canonical_tool_name()) { + builder.register_handler(Arc::new(McpHandler::new(tool.clone()))); } } } diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index fdae0452e..8079f91da 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -26,11 +26,10 @@ use crate::tools::handlers::view_image_spec::ViewImageToolOptions; use crate::tools::handlers::view_image_spec::create_view_image_tool; use crate::tools::registry::ToolRegistry; use crate::tools::spec_plan_types::ToolNamespace; -use crate::tools::spec_plan_types::ToolRegistryBuildDeferredTool; -use crate::tools::spec_plan_types::ToolRegistryBuildMcpTool; use codex_app_server_protocol::AppInfo; use codex_features::Feature; use codex_features::Features; +use codex_mcp::ToolInfo; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::WebSearchConfig; use codex_protocol::config_types::WebSearchMode; @@ -2413,10 +2412,10 @@ fn search_capable_model_info() -> ModelInfo { } } -fn build_specs<'a>( +fn build_specs( config: &ToolsConfig, mcp_tools: Option>, - deferred_mcp_tools: Option>>, + deferred_mcp_tools: Option>, dynamic_tools: &[DynamicToolSpec], ) -> (Vec, ToolRegistry) { build_specs_with_discoverable_tools( @@ -2429,10 +2428,10 @@ fn build_specs<'a>( ) } -fn build_specs_with_discoverable_tools<'a>( +fn build_specs_with_discoverable_tools( config: &ToolsConfig, mcp_tools: Option>, - deferred_mcp_tools: Option>>, + deferred_mcp_tools: Option>, discoverable_tools: Option>, extension_tool_bundles: &[codex_tool_api::ToolBundle], dynamic_tools: &[DynamicToolSpec], @@ -2448,10 +2447,10 @@ fn build_specs_with_discoverable_tools<'a>( ) } -fn build_specs_with_optional_tool_namespaces<'a>( +fn build_specs_with_optional_tool_namespaces( config: &ToolsConfig, mcp_tools: Option>, - deferred_mcp_tools: Option>>, + deferred_mcp_tools: Option>, tool_namespaces: Option>, discoverable_tools: Option>, extension_tool_bundles: &[codex_tool_api::ToolBundle], @@ -2460,10 +2459,7 @@ fn build_specs_with_optional_tool_namespaces<'a>( let mcp_tool_inputs = mcp_tools.as_ref().map(|mcp_tools| { mcp_tools .iter() - .map(|(name, tool)| ToolRegistryBuildMcpTool { - name: name.clone(), - tool, - }) + .map(|(name, tool)| tool_info_from_parts(name, tool.clone())) .collect::>() }); let builder = build_tool_registry_builder( @@ -2497,6 +2493,33 @@ fn mcp_tool(name: &str, description: &str, input_schema: serde_json::Value) -> r } } +fn tool_info_from_parts(name: &ToolName, tool: rmcp::model::Tool) -> ToolInfo { + ToolInfo { + server_name: server_name_from_tool_name(name), + supports_parallel_tool_calls: false, + server_origin: None, + callable_name: name.name.clone(), + callable_namespace: name.namespace.clone().unwrap_or_default(), + namespace_description: None, + tool, + connector_id: None, + connector_name: None, + plugin_display_names: Vec::new(), + } +} + +fn server_name_from_tool_name(name: &ToolName) -> String { + name.namespace + .as_deref() + .and_then(|namespace| { + namespace + .strip_prefix("mcp__") + .and_then(|suffix| suffix.strip_suffix("__")) + }) + .unwrap_or_else(|| name.namespace.as_deref().unwrap_or("test_server")) + .to_string() +} + #[test] fn code_mode_augments_mcp_tool_descriptions_with_structured_output_sample() { let model_info = model_info(); @@ -2588,18 +2611,28 @@ fn discoverable_connector(id: &str, name: &str, description: &str) -> Discoverab })) } -fn deferred_mcp_tool<'a>( - tool_name: &'a str, - tool_namespace: &'a str, - server_name: &'a str, - connector_name: Option<&'a str>, - description: Option<&'a str>, -) -> ToolRegistryBuildDeferredTool<'a> { - ToolRegistryBuildDeferredTool { - name: ToolName::namespaced(tool_namespace, tool_name), - server_name, - connector_name, - description, +fn deferred_mcp_tool( + tool_name: &str, + tool_namespace: &str, + server_name: &str, + connector_name: Option<&str>, + description: Option<&str>, +) -> ToolInfo { + ToolInfo { + server_name: server_name.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, + callable_name: tool_name.to_string(), + callable_namespace: tool_namespace.to_string(), + namespace_description: description.map(str::to_string), + tool: mcp_tool( + tool_name, + description.unwrap_or("Deferred MCP tool"), + json!({}), + ), + connector_id: None, + connector_name: connector_name.map(str::to_string), + plugin_display_names: Vec::new(), } } diff --git a/codex-rs/core/src/tools/spec_plan_types.rs b/codex-rs/core/src/tools/spec_plan_types.rs index 7ac5c0f52..9e9ff555c 100644 --- a/codex-rs/core/src/tools/spec_plan_types.rs +++ b/codex-rs/core/src/tools/spec_plan_types.rs @@ -1,15 +1,15 @@ use crate::tools::handlers::multi_agents_spec::WaitAgentTimeoutOptions; +use codex_mcp::ToolInfo; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_tool_api::ToolBundle as ExtensionToolBundle; use codex_tools::DiscoverableTool; -use codex_tools::ToolName; use codex_tools::ToolsConfig; use std::collections::HashMap; #[derive(Clone, Copy)] pub struct ToolRegistryBuildParams<'a> { - pub mcp_tools: Option<&'a [ToolRegistryBuildMcpTool<'a>]>, - pub deferred_mcp_tools: Option<&'a [ToolRegistryBuildDeferredTool<'a>]>, + pub mcp_tools: Option<&'a [ToolInfo]>, + pub deferred_mcp_tools: Option<&'a [ToolInfo]>, pub tool_namespaces: Option<&'a HashMap>, pub discoverable_tools: Option<&'a [DiscoverableTool]>, pub extension_tool_bundles: &'a [ExtensionToolBundle], @@ -25,23 +25,6 @@ pub struct ToolNamespace { pub description: Option, } -/// Direct MCP tool metadata needed to expose the Responses API namespace tool -/// while registering its runtime handler with the canonical namespace/name -/// identity. -#[derive(Debug, Clone)] -pub struct ToolRegistryBuildMcpTool<'a> { - pub name: ToolName, - pub tool: &'a rmcp::model::Tool, -} - -#[derive(Debug, Clone)] -pub struct ToolRegistryBuildDeferredTool<'a> { - pub name: ToolName, - pub server_name: &'a str, - pub connector_name: Option<&'a str>, - pub description: Option<&'a str>, -} - pub(crate) fn agent_type_description( config: &ToolsConfig, default_agent_type_description: &str, diff --git a/codex-rs/core/src/tools/spec_tests.rs b/codex-rs/core/src/tools/spec_tests.rs index 06343fce0..c97d3ccde 100644 --- a/codex-rs/core/src/tools/spec_tests.rs +++ b/codex-rs/core/src/tools/spec_tests.rs @@ -60,6 +60,8 @@ fn mcp_tool(name: &str, description: &str, input_schema: serde_json::Value) -> r fn mcp_tool_info(tool: rmcp::model::Tool) -> ToolInfo { ToolInfo { server_name: "test_server".to_string(), + supports_parallel_tool_calls: false, + server_origin: None, callable_name: tool.name.to_string(), callable_namespace: "mcp__test_server__".to_string(), namespace_description: None, @@ -78,6 +80,8 @@ fn mcp_tool_info_with_display_name(display_name: &str, tool: rmcp::model::Tool) ToolInfo { server_name: "test_server".to_string(), + supports_parallel_tool_calls: false, + server_origin: None, callable_name, callable_namespace, namespace_description: None, @@ -352,7 +356,6 @@ async fn assert_model_tools( mcp_tools: None, deferred_mcp_tools: None, unavailable_called_tools: Vec::new(), - parallel_mcp_server_names: std::collections::HashSet::new(), discoverable_tools: None, extension_tool_bundles: Vec::new(), dynamic_tools: &[], @@ -895,6 +898,8 @@ async fn search_tool_description_falls_back_to_connector_name_without_descriptio /*mcp_tools*/ None, Some(vec![ToolInfo { server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, callable_name: "_create_event".to_string(), callable_namespace: "mcp__codex_apps__calendar".to_string(), namespace_description: None, @@ -943,6 +948,8 @@ async fn search_tool_registers_namespaced_mcp_tool_aliases() { Some(vec![ ToolInfo { server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, callable_name: "_create_event".to_string(), callable_namespace: "mcp__codex_apps__calendar".to_string(), namespace_description: None, @@ -957,6 +964,8 @@ async fn search_tool_registers_namespaced_mcp_tool_aliases() { }, ToolInfo { server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), + supports_parallel_tool_calls: false, + server_origin: None, callable_name: "_list_events".to_string(), callable_namespace: "mcp__codex_apps__calendar".to_string(), namespace_description: None, @@ -971,6 +980,8 @@ async fn search_tool_registers_namespaced_mcp_tool_aliases() { }, ToolInfo { server_name: "rmcp".to_string(), + supports_parallel_tool_calls: false, + server_origin: None, callable_name: "echo".to_string(), callable_namespace: "mcp__rmcp__".to_string(), namespace_description: None, diff --git a/codex-rs/core/src/tools/tool_dispatch_trace.rs b/codex-rs/core/src/tools/tool_dispatch_trace.rs index 344b68634..6a0e63579 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace.rs @@ -120,15 +120,6 @@ fn tool_dispatch_payload(payload: &ToolPayload) -> ToolDispatchPayload { additional_permissions: params.additional_permissions.clone(), justification: params.justification.clone(), }, - ToolPayload::Mcp { - server, - tool, - raw_arguments, - } => ToolDispatchPayload::Mcp { - server: server.clone(), - tool: tool.clone(), - raw_arguments: raw_arguments.clone(), - }, } } diff --git a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs index 99d90c345..f0c808b30 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace_tests.rs @@ -22,7 +22,6 @@ use crate::tools::context::ToolCallSource; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolPayload; use crate::tools::registry::ToolHandler; -use crate::tools::registry::ToolKind; use crate::tools::registry::ToolRegistry; use crate::turn_diff_tracker::TurnDiffTracker; @@ -37,10 +36,6 @@ impl ToolHandler for TestHandler { self.tool_name.clone() } - fn kind(&self) -> ToolKind { - ToolKind::Function - } - async fn handle(&self, _invocation: ToolInvocation) -> Result { Ok(FunctionToolOutput::from_text("ok".to_string(), Some(true))) } diff --git a/codex-rs/otel/src/events/session_telemetry.rs b/codex-rs/otel/src/events/session_telemetry.rs index e377419a6..49d8b8eb7 100644 --- a/codex-rs/otel/src/events/session_telemetry.rs +++ b/codex-rs/otel/src/events/session_telemetry.rs @@ -65,6 +65,12 @@ const RESPONSES_API_ENGINE_SERVICE_TTFT_FIELD: &str = "engine_service_ttft_total const RESPONSES_API_ENGINE_IAPI_TBT_FIELD: &str = "engine_iapi_tbt_across_engine_calls_ms"; const RESPONSES_API_ENGINE_SERVICE_TBT_FIELD: &str = "engine_service_tbt_across_engine_calls_ms"; +fn trace_field_value<'a>(fields: &'a [(&str, &str)], key: &str) -> Option<&'a str> { + fields + .iter() + .find_map(|(field_key, value)| (*field_key == key).then_some(*value)) +} + #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct AuthEnvTelemetryMetadata { pub openai_api_key_env_present: bool, @@ -943,8 +949,7 @@ impl SessionTelemetry { call_id: &str, arguments: &str, extra_tags: &[(&str, &str)], - mcp_server: Option<&str>, - mcp_server_origin: Option<&str>, + extra_trace_fields: &[(&str, &str)], f: F, ) -> Result<(String, bool), E> where @@ -969,8 +974,7 @@ impl SessionTelemetry { success, output.as_ref(), extra_tags, - mcp_server, - mcp_server_origin, + extra_trace_fields, ); result @@ -1010,8 +1014,7 @@ impl SessionTelemetry { success: bool, output: &str, extra_tags: &[(&str, &str)], - mcp_server: Option<&str>, - mcp_server_origin: Option<&str>, + extra_trace_fields: &[(&str, &str)], ) { let success_str = if success { "true" } else { "false" }; let mut tags = Vec::with_capacity(2 + extra_tags.len()); @@ -1020,8 +1023,9 @@ impl SessionTelemetry { tags.extend_from_slice(extra_tags); self.counter(TOOL_CALL_COUNT_METRIC, /*inc*/ 1, &tags); self.record_duration(TOOL_CALL_DURATION_METRIC, duration, &tags); - let mcp_server = mcp_server.unwrap_or(""); - let mcp_server_origin = mcp_server_origin.unwrap_or(""); + let mcp_server = trace_field_value(extra_trace_fields, "mcp_server").unwrap_or(""); + let mcp_server_origin = + trace_field_value(extra_trace_fields, "mcp_server_origin").unwrap_or(""); log_event!( self, event.name = "codex.tool_result", diff --git a/codex-rs/otel/tests/suite/otel_export_routing_policy.rs b/codex-rs/otel/tests/suite/otel_export_routing_policy.rs index fa02cc7f3..ef30fbcd8 100644 --- a/codex-rs/otel/tests/suite/otel_export_routing_policy.rs +++ b/codex-rs/otel/tests/suite/otel_export_routing_policy.rs @@ -249,8 +249,10 @@ fn otel_export_routing_policy_routes_tool_result_log_and_trace_events() { /*success*/ true, "secret output\nsecond line", &[], - Some("internal-mcp"), - Some("stdio"), + &[ + ("mcp_server", "internal-mcp"), + ("mcp_server_origin", "stdio"), + ], ); }); @@ -277,6 +279,10 @@ fn otel_export_routing_policy_routes_tool_result_log_and_trace_events() { tool_log_attrs.get("mcp_server").map(String::as_str), Some("internal-mcp") ); + assert_eq!( + tool_log_attrs.get("mcp_server_origin").map(String::as_str), + Some("stdio") + ); let spans = span_exporter.get_finished_spans().expect("span export"); assert_eq!(spans.len(), 1); @@ -299,14 +305,6 @@ fn otel_export_routing_policy_routes_tool_result_log_and_trace_events() { .map(String::as_str), Some("2") ); - assert_eq!( - tool_trace_attrs.get("tool_origin").map(String::as_str), - Some("mcp") - ); - assert_eq!( - tool_trace_attrs.get("mcp_tool").map(String::as_str), - Some("true") - ); assert!(!tool_trace_attrs.contains_key("arguments")); assert!(!tool_trace_attrs.contains_key("output")); assert!(!tool_trace_attrs.contains_key("mcp_server")); diff --git a/codex-rs/otel/tests/suite/runtime_summary.rs b/codex-rs/otel/tests/suite/runtime_summary.rs index 21fc7bf76..7c869482c 100644 --- a/codex-rs/otel/tests/suite/runtime_summary.rs +++ b/codex-rs/otel/tests/suite/runtime_summary.rs @@ -44,8 +44,7 @@ fn runtime_metrics_summary_collects_tool_api_and_streaming_metrics() -> Result<( /*success*/ true, "ok", &[], - /*mcp_server*/ None, - /*mcp_server_origin*/ None, + /*extra_trace_fields*/ &[], ); manager.record_api_request( /*attempt*/ 1, diff --git a/codex-rs/rollout-trace/src/tool_dispatch.rs b/codex-rs/rollout-trace/src/tool_dispatch.rs index 1c444a87d..12d058c43 100644 --- a/codex-rs/rollout-trace/src/tool_dispatch.rs +++ b/codex-rs/rollout-trace/src/tool_dispatch.rs @@ -93,11 +93,6 @@ pub enum ToolDispatchPayload { additional_permissions: Option, justification: Option, }, - Mcp { - server: String, - tool: String, - raw_arguments: String, - }, } /// Result data returned from a dispatch-level tool call. @@ -263,14 +258,7 @@ fn requester_fields( } } -fn dispatched_tool_kind(tool_name: &str, payload: &ToolDispatchPayload) -> ToolCallKind { - if let ToolDispatchPayload::Mcp { server, tool, .. } = payload { - return ToolCallKind::Mcp { - server: server.clone(), - tool: tool.clone(), - }; - } - +fn dispatched_tool_kind(tool_name: &str, _payload: &ToolDispatchPayload) -> ToolCallKind { match tool_name { "exec_command" | "local_shell" | "shell" | "shell_command" => ToolCallKind::ExecCommand, "write_stdin" => ToolCallKind::WriteStdin, @@ -291,12 +279,8 @@ fn dispatched_tool_kind(tool_name: &str, payload: &ToolDispatchPayload) -> ToolC fn dispatched_tool_label( tool_name: &str, tool_namespace: Option<&str>, - payload: &ToolDispatchPayload, + _payload: &ToolDispatchPayload, ) -> String { - if let ToolDispatchPayload::Mcp { server, tool, .. } = payload { - return format!("mcp:{server}:{tool}"); - } - match tool_namespace { Some(namespace) => format!("{namespace}.{tool_name}"), None => tool_name.to_string(), @@ -310,7 +294,6 @@ impl ToolDispatchPayload { ToolDispatchPayload::ToolSearch { arguments } => truncate_preview(&arguments.query), ToolDispatchPayload::Custom { input } => truncate_preview(input), ToolDispatchPayload::LocalShell { command, .. } => truncate_preview(&command.join(" ")), - ToolDispatchPayload::Mcp { raw_arguments, .. } => truncate_preview(raw_arguments), } } @@ -346,16 +329,6 @@ impl ToolDispatchPayload { "additional_permissions": additional_permissions, "justification": justification, }), - ToolDispatchPayload::Mcp { - server, - tool, - raw_arguments, - } => json!({ - "type": "mcp", - "server": server, - "tool": tool, - "raw_arguments": raw_arguments, - }), } } }