From b4e66fb8ec33cf337d9560d24b93a2c482e0effb Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Fri, 10 Apr 2026 18:30:53 -0700 Subject: [PATCH] register all mcps w/ namespace --- .../codex-mcp/src/mcp_connection_manager.rs | 26 +++++--- .../src/mcp_connection_manager_tests.rs | 62 +++++++++++++++++++ codex-rs/core/src/tools/code_mode/mod.rs | 18 +++--- codex-rs/core/src/tools/js_repl/mod.rs | 39 +++++++----- codex-rs/core/src/tools/router.rs | 19 +++--- codex-rs/core/src/tools/spec.rs | 16 +++-- codex-rs/core/src/tools/spec_tests.rs | 37 +++++++++++ codex-rs/tools/src/lib.rs | 1 + codex-rs/tools/src/tool_registry_plan.rs | 20 +++--- .../tools/src/tool_registry_plan_tests.rs | 22 ++++++- .../tools/src/tool_registry_plan_types.rs | 13 +++- 11 files changed, 219 insertions(+), 54 deletions(-) diff --git a/codex-rs/codex-mcp/src/mcp_connection_manager.rs b/codex-rs/codex-mcp/src/mcp_connection_manager.rs index 9021f7d7a9..ce9d955f1d 100644 --- a/codex-rs/codex-mcp/src/mcp_connection_manager.rs +++ b/codex-rs/codex-mcp/src/mcp_connection_manager.rs @@ -1192,13 +1192,25 @@ impl McpConnectionManager { } pub async fn resolve_tool_info(&self, name: &str, namespace: Option<&str>) -> Option { - let qualified_name = match namespace { - Some(namespace) if name.starts_with(namespace) => name.to_string(), - Some(namespace) => format!("{namespace}{name}"), - None => name.to_string(), - }; - - self.list_all_tools().await.get(&qualified_name).cloned() + let all_tools = self.list_all_tools().await; + if let Some(namespace) = namespace { + let qualified_name = format!("{namespace}{name}"); + if let Some(tool) = all_tools.get(&qualified_name) + && tool.callable_namespace == namespace + && tool.callable_name == name + { + return Some(tool.clone()); + } + if let Some(tool) = all_tools.get(name) + && tool.callable_namespace == namespace + && format!("{namespace}{}", tool.callable_name) == name + { + return Some(tool.clone()); + } + None + } else { + all_tools.get(name).cloned() + } } pub async fn notify_sandbox_state_change(&self, sandbox_state: &SandboxState) -> Result<()> { diff --git a/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs b/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs index 2c0d6fda68..f7cd64880c 100644 --- a/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs @@ -646,6 +646,68 @@ 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_flat_and_split_callable_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 sandbox_policy = Constrained::allow_any(SandboxPolicy::new_read_only_policy()); + let mut manager = McpConnectionManager::new_uninitialized(&approval_policy, &sandbox_policy); + 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()), + }, + ); + + let flat = manager + .resolve_tool_info("mcp__rmcp__echo", None) + .await + .expect("flat qualified MCP tool name should resolve"); + let split = manager + .resolve_tool_info("echo", Some("mcp__rmcp__")) + .await + .expect("split MCP tool namespace and name should resolve"); + let split_with_flat_name = manager + .resolve_tool_info("mcp__rmcp__echo", Some("mcp__rmcp__")) + .await + .expect("flat MCP tool name with namespace should resolve"); + + let expected = ("rmcp", "mcp__rmcp__", "echo", "echo"); + assert_eq!( + ( + flat.server_name.as_str(), + flat.callable_namespace.as_str(), + flat.callable_name.as_str(), + flat.tool.name.as_ref(), + ), + expected + ); + assert_eq!( + ( + split.server_name.as_str(), + split.callable_namespace.as_str(), + split.callable_name.as_str(), + split.tool.name.as_ref(), + ), + expected + ); + assert_eq!( + ( + split_with_flat_name.server_name.as_str(), + split_with_flat_name.callable_namespace.as_str(), + split_with_flat_name.callable_name.as_str(), + split_with_flat_name.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::>() diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 99c8ec0e0d..ecdda91768 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -284,28 +284,32 @@ async fn call_nested_tool( ))); } - let payload = if let Some(tool_info) = exec + let (tool_call_name, payload) = if let Some(tool_info) = exec .session .resolve_mcp_tool_info(&tool_name, /*namespace*/ None) .await { - match serialize_function_tool_arguments(&tool_name, input) { - Ok(raw_arguments) => ToolPayload::Mcp { + let raw_arguments = match serialize_function_tool_arguments(&tool_name, input) { + Ok(raw_arguments) => raw_arguments, + Err(error) => return Err(FunctionCallError::RespondToModel(error)), + }; + ( + ToolName::namespaced(tool_info.callable_namespace, tool_info.callable_name), + ToolPayload::Mcp { server: tool_info.server_name, tool: tool_info.tool.name.to_string(), raw_arguments, }, - Err(error) => return Err(FunctionCallError::RespondToModel(error)), - } + ) } else { match build_nested_tool_payload(tool_runtime.find_spec(&tool_name), &tool_name, input) { - Ok(payload) => payload, + Ok(payload) => (ToolName::plain(tool_name.clone()), payload), Err(error) => return Err(FunctionCallError::RespondToModel(error)), } }; let call = ToolCall { - tool_name: ToolName::plain(tool_name.clone()), + tool_name: tool_call_name, call_id: format!("{PUBLIC_TOOL_NAME}-{}", uuid::Uuid::new_v4()), payload, }; diff --git a/codex-rs/core/src/tools/js_repl/mod.rs b/codex-rs/core/src/tools/js_repl/mod.rs index e307f9718c..a4e4b23055 100644 --- a/codex-rs/core/src/tools/js_repl/mod.rs +++ b/codex-rs/core/src/tools/js_repl/mod.rs @@ -1572,29 +1572,40 @@ impl JsReplManager { }, ); - let payload = if let Some(tool_info) = exec + let (tool_call_name, payload) = if let Some(tool_info) = exec .session .resolve_mcp_tool_info(&req.tool_name, /*namespace*/ None) .await { - crate::tools::context::ToolPayload::Mcp { - server: tool_info.server_name, - tool: tool_info.tool.name.to_string(), - raw_arguments: req.arguments.clone(), - } + ( + codex_tools::ToolName::namespaced( + tool_info.callable_namespace, + tool_info.callable_name, + ), + crate::tools::context::ToolPayload::Mcp { + server: tool_info.server_name, + tool: tool_info.tool.name.to_string(), + raw_arguments: req.arguments.clone(), + }, + ) } else if is_freeform_tool(&router.specs(), &req.tool_name) { - crate::tools::context::ToolPayload::Custom { - input: req.arguments.clone(), - } + ( + codex_tools::ToolName::plain(req.tool_name.clone()), + crate::tools::context::ToolPayload::Custom { + input: req.arguments.clone(), + }, + ) } else { - crate::tools::context::ToolPayload::Function { - arguments: req.arguments.clone(), - } + ( + codex_tools::ToolName::plain(req.tool_name.clone()), + crate::tools::context::ToolPayload::Function { + arguments: req.arguments.clone(), + }, + ) }; - let tool_name = req.tool_name.clone(); let call = crate::tools::router::ToolCall { - tool_name: codex_tools::ToolName::plain(tool_name.clone()), + tool_name: tool_call_name, call_id: req.id.clone(), payload, }; diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index af9d278585..94ad92188c 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -126,16 +126,15 @@ impl ToolRouter { call_id, .. } => { - let mcp_tool = session + if let Some(tool_info) = session .resolve_mcp_tool_info(&name, namespace.as_deref()) - .await; - let tool_name = match namespace { - Some(namespace) => ToolName::namespaced(namespace, name), - None => ToolName::plain(name), - }; - if let Some(tool_info) = mcp_tool { + .await + { Ok(Some(ToolCall { - tool_name, + tool_name: ToolName::namespaced( + tool_info.callable_namespace, + tool_info.callable_name, + ), call_id, payload: ToolPayload::Mcp { server: tool_info.server_name, @@ -144,6 +143,10 @@ impl ToolRouter { }, })) } else { + let tool_name = match namespace { + Some(namespace) => ToolName::namespaced(namespace, name), + None => ToolName::plain(name), + }; Ok(Some(ToolCall { tool_name, call_id, diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index bfad3cdb35..cb0eb08193 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -11,6 +11,7 @@ use codex_tools::DiscoverableTool; use codex_tools::ToolHandlerKind; use codex_tools::ToolNamespace; use codex_tools::ToolRegistryPlanDeferredTool; +use codex_tools::ToolRegistryPlanMcpTool; use codex_tools::ToolRegistryPlanParams; use codex_tools::ToolUserShellType; use codex_tools::ToolsConfig; @@ -29,16 +30,21 @@ pub(crate) fn tool_user_shell_type(user_shell: &Shell) -> ToolUserShellType { } } -struct McpToolPlanInputs { - mcp_tools: HashMap, +struct McpToolPlanInputs<'a> { + mcp_tools: Vec>, tool_namespaces: HashMap, } -fn map_mcp_tools_for_plan(mcp_tools: &HashMap) -> McpToolPlanInputs { +fn map_mcp_tools_for_plan(mcp_tools: &HashMap) -> McpToolPlanInputs<'_> { McpToolPlanInputs { mcp_tools: mcp_tools .iter() - .map(|(name, tool)| (name.clone(), tool.tool.clone())) + .map(|(qualified_name, tool)| ToolRegistryPlanMcpTool { + qualified_name: qualified_name.clone(), + callable_name: tool.callable_name.clone(), + callable_namespace: tool.callable_namespace.clone(), + tool: &tool.tool, + }) .collect(), tool_namespaces: mcp_tools .iter() @@ -114,7 +120,7 @@ pub(crate) fn build_specs_with_discoverable_tools( ToolRegistryPlanParams { mcp_tools: mcp_tool_plan_inputs .as_ref() - .map(|inputs| &inputs.mcp_tools), + .map(|inputs| inputs.mcp_tools.as_slice()), deferred_mcp_tools: deferred_mcp_tool_sources.as_deref(), tool_namespaces: mcp_tool_plan_inputs .as_ref() diff --git a/codex-rs/core/src/tools/spec_tests.rs b/codex-rs/core/src/tools/spec_tests.rs index da76a1bb28..4e96214aa2 100644 --- a/codex-rs/core/src/tools/spec_tests.rs +++ b/codex-rs/core/src/tools/spec_tests.rs @@ -911,6 +911,43 @@ fn search_tool_registers_namespaced_mcp_tool_aliases() { assert!(registry.has_handler(&mcp_alias)); } +#[test] +fn direct_mcp_tools_register_namespaced_handlers() { + let config = test_config(); + let model_info = construct_model_info_offline("gpt-5-codex", &config); + let mut features = Features::with_defaults(); + features.enable(Feature::UnifiedExec); + let available_models = Vec::new(); + let tools_config = ToolsConfig::new(&ToolsConfigParams { + model_info: &model_info, + available_models: &available_models, + features: &features, + image_generation_tool_auth_allowed: true, + web_search_mode: Some(WebSearchMode::Cached), + session_source: SessionSource::Cli, + sandbox_policy: &SandboxPolicy::DangerFullAccess, + windows_sandbox_level: WindowsSandboxLevel::Disabled, + }); + + let (_, registry) = build_specs( + &tools_config, + Some(HashMap::from([( + "mcp__test_server__echo".to_string(), + mcp_tool_info(mcp_tool( + "echo", + "Echo", + serde_json::json!({"type": "object"}), + )), + )])), + /*deferred_mcp_tools*/ None, + &[], + ) + .build(); + + assert!(registry.has_handler(&ToolName::namespaced("mcp__test_server__", "echo"))); + assert!(!registry.has_handler(&ToolName::plain("mcp__test_server__echo"))); +} + #[test] fn test_mcp_tool_property_missing_type_defaults_to_string() { let config = test_config(); diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index c310222a0f..55b12329cb 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -120,6 +120,7 @@ pub use tool_registry_plan_types::ToolHandlerSpec; pub use tool_registry_plan_types::ToolNamespace; pub use tool_registry_plan_types::ToolRegistryPlan; pub use tool_registry_plan_types::ToolRegistryPlanDeferredTool; +pub use tool_registry_plan_types::ToolRegistryPlanMcpTool; pub use tool_registry_plan_types::ToolRegistryPlanParams; pub use tool_spec::ConfiguredToolSpec; pub use tool_spec::ResponsesApiWebSearchFilters; diff --git a/codex-rs/tools/src/tool_registry_plan.rs b/codex-rs/tools/src/tool_registry_plan.rs index 82b4c3a142..050d05c7cd 100644 --- a/codex-rs/tools/src/tool_registry_plan.rs +++ b/codex-rs/tools/src/tool_registry_plan.rs @@ -61,7 +61,6 @@ use crate::request_user_input_tool_description; use crate::tool_registry_plan_types::agent_type_description; use codex_protocol::openai_models::ApplyPatchToolType; use codex_protocol::openai_models::ConfigShellToolType; -use rmcp::model::Tool as McpTool; use std::collections::BTreeMap; pub fn build_tool_registry_plan( @@ -471,25 +470,26 @@ pub fn build_tool_registry_plan( } if let Some(mcp_tools) = params.mcp_tools { - let mut entries: Vec<(String, &McpTool)> = mcp_tools - .iter() - .map(|(name, tool)| (name.clone(), tool)) - .collect(); - entries.sort_by(|left, right| left.0.cmp(&right.0)); + let mut entries = mcp_tools.to_vec(); + entries.sort_by(|left, right| left.qualified_name.cmp(&right.qualified_name)); - for (name, tool) in entries { - match mcp_tool_to_responses_api_tool(name.clone(), tool) { + for tool in entries { + let qualified_name = tool.qualified_name.clone(); + match mcp_tool_to_responses_api_tool(qualified_name.clone(), tool.tool) { Ok(converted_tool) => { plan.push_spec( ToolSpec::Function(converted_tool), /*supports_parallel_tool_calls*/ false, config.code_mode_enabled, ); - plan.register_handler(name, ToolHandlerKind::Mcp); + plan.register_handler( + ToolName::namespaced(tool.callable_namespace, tool.callable_name), + ToolHandlerKind::Mcp, + ); } Err(error) => { tracing::error!( - "Failed to convert {name:?} MCP tool to OpenAI tool: {error:?}" + "Failed to convert {qualified_name:?} MCP tool to OpenAI tool: {error:?}" ); } } diff --git a/codex-rs/tools/src/tool_registry_plan_tests.rs b/codex-rs/tools/src/tool_registry_plan_tests.rs index 9da58849c1..c042e3d750 100644 --- a/codex-rs/tools/src/tool_registry_plan_tests.rs +++ b/codex-rs/tools/src/tool_registry_plan_tests.rs @@ -13,6 +13,7 @@ use crate::ResponsesApiWebSearchUserLocation; use crate::ToolHandlerSpec; use crate::ToolNamespace; use crate::ToolRegistryPlanDeferredTool; +use crate::ToolRegistryPlanMcpTool; use crate::ToolsConfigParams; use crate::WaitAgentTimeoutOptions; use crate::mcp_call_tool_result_output_schema; @@ -1890,10 +1891,29 @@ fn build_specs_with_optional_tool_namespaces<'a>( discoverable_tools: Option>, dynamic_tools: &[DynamicToolSpec], ) -> (Vec, Vec) { + let mcp_tool_inputs = mcp_tools.as_ref().map(|mcp_tools| { + mcp_tools + .iter() + .map(|(qualified_name, tool)| { + let raw_tool_name = tool.name.as_ref(); + let callable_namespace = qualified_name + .strip_suffix(raw_tool_name) + .filter(|namespace| !namespace.is_empty()) + .unwrap_or("mcp__test_server__"); + + ToolRegistryPlanMcpTool { + qualified_name: qualified_name.clone(), + callable_name: raw_tool_name.to_string(), + callable_namespace: callable_namespace.to_string(), + tool, + } + }) + .collect::>() + }); let plan = build_tool_registry_plan( config, ToolRegistryPlanParams { - mcp_tools: mcp_tools.as_ref(), + mcp_tools: mcp_tool_inputs.as_deref(), deferred_mcp_tools: deferred_mcp_tools.as_deref(), tool_namespaces: tool_namespaces.as_ref(), discoverable_tools: discoverable_tools.as_deref(), diff --git a/codex-rs/tools/src/tool_registry_plan_types.rs b/codex-rs/tools/src/tool_registry_plan_types.rs index 7459954dcd..1c0d6ed9b1 100644 --- a/codex-rs/tools/src/tool_registry_plan_types.rs +++ b/codex-rs/tools/src/tool_registry_plan_types.rs @@ -6,7 +6,6 @@ use crate::ToolsConfig; use crate::WaitAgentTimeoutOptions; use crate::augment_tool_spec_for_code_mode; use codex_protocol::dynamic_tools::DynamicToolSpec; -use rmcp::model::Tool as McpTool; use std::collections::HashMap; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -58,7 +57,7 @@ pub struct ToolRegistryPlan { #[derive(Debug, Clone, Copy)] pub struct ToolRegistryPlanParams<'a> { - pub mcp_tools: Option<&'a HashMap>, + pub mcp_tools: Option<&'a [ToolRegistryPlanMcpTool<'a>]>, pub deferred_mcp_tools: Option<&'a [ToolRegistryPlanDeferredTool<'a>]>, pub tool_namespaces: Option<&'a HashMap>, pub discoverable_tools: Option<&'a [DiscoverableTool]>, @@ -73,6 +72,16 @@ pub struct ToolNamespace { pub description: Option, } +/// Direct MCP tool metadata needed to expose the flat Responses API tool while +/// registering its runtime handler with the canonical namespace/name split. +#[derive(Debug, Clone)] +pub struct ToolRegistryPlanMcpTool<'a> { + pub qualified_name: String, + pub callable_name: String, + pub callable_namespace: String, + pub tool: &'a rmcp::model::Tool, +} + #[derive(Debug, Clone, Copy)] pub struct ToolRegistryPlanDeferredTool<'a> { pub tool_name: &'a str,