From 52a0ab1201c2c3fda775f9b45d4c3326a4d9b806 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Thu, 16 Apr 2026 21:05:27 +0800 Subject: [PATCH] key mcp tool listings by ToolName --- codex-rs/codex-mcp/src/mcp/mod.rs | 5 +- .../codex-mcp/src/mcp_connection_manager.rs | 21 +- .../src/mcp_connection_manager_tests.rs | 40 ++- codex-rs/codex-mcp/src/mcp_tool_names.rs | 23 +- codex-rs/core/src/codex/turn.rs | 2 +- codex-rs/core/src/connectors.rs | 5 +- codex-rs/core/src/connectors_tests.rs | 132 ++++----- codex-rs/core/src/mcp_tool_exposure.rs | 11 +- codex-rs/core/src/mcp_tool_exposure_test.rs | 110 ++++--- codex-rs/core/src/plugins/injection.rs | 3 +- .../core/src/tools/handlers/tool_search.rs | 23 +- codex-rs/core/src/tools/router.rs | 4 +- codex-rs/core/src/tools/spec.rs | 18 +- codex-rs/core/src/tools/spec_tests.rs | 269 ++++++++---------- codex-rs/core/src/unavailable_tool.rs | 7 +- 15 files changed, 342 insertions(+), 331 deletions(-) diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 1dc9db0789..1054619bd6 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -597,7 +597,7 @@ async fn collect_mcp_server_status_snapshot_from_manager( ); let mut tools_by_server = HashMap::>::new(); - for (_qualified_name, tool_info) in tools { + for (_tool_name, tool_info) in tools { let raw_tool_name = tool_info.tool.name.to_string(); let Some(tool) = protocol_tool_from_rmcp_tool(&raw_tool_name, &tool_info.tool) else { continue; @@ -655,7 +655,8 @@ pub async fn collect_mcp_snapshot_from_manager_with_detail( let tools = tools .into_iter() .filter_map(|(name, tool)| { - protocol_tool_from_rmcp_tool(&name, &tool.tool).map(|tool| (name, tool)) + let display_name = name.display(); + protocol_tool_from_rmcp_tool(&display_name, &tool.tool).map(|tool| (display_name, tool)) }) .collect::>(); diff --git a/codex-rs/codex-mcp/src/mcp_connection_manager.rs b/codex-rs/codex-mcp/src/mcp_connection_manager.rs index d1696b5931..e803880daa 100644 --- a/codex-rs/codex-mcp/src/mcp_connection_manager.rs +++ b/codex-rs/codex-mcp/src/mcp_connection_manager.rs @@ -3,8 +3,8 @@ //! The [`McpConnectionManager`] owns one [`codex_rmcp_client::RmcpClient`] per //! configured server (keyed by the *server name*). It offers convenience //! helpers to query the available tools across *all* servers and returns them -//! in a single aggregated map using the model-visible fully-qualified tool name -//! as the key. +//! in a single aggregated map using the model-visible callable tool name as the +//! key. use std::borrow::Cow; use std::collections::HashMap; @@ -874,10 +874,11 @@ impl McpConnectionManager { failures } - /// Returns a single map that contains all tools. Each key is the - /// fully-qualified name for the tool. + /// Returns a single map that contains all tools keyed by model-visible + /// callable name. Each key's `display()` is unique and <= 64 bytes so it + /// can be used at flat protocol/display boundaries. #[instrument(level = "trace", skip_all)] - pub async fn list_all_tools(&self) -> HashMap { + pub async fn list_all_tools(&self) -> HashMap { let mut tools = Vec::new(); for managed_client in self.clients.values() { let Some(server_tools) = managed_client.listed_tools().await else { @@ -893,7 +894,7 @@ impl McpConnectionManager { /// On success, the refreshed tools replace the cache contents and the /// latest filtered tool map is returned directly to the caller. On /// failure, the existing cache remains unchanged. - pub async fn hard_refresh_codex_apps_tools_cache(&self) -> Result> { + pub async fn hard_refresh_codex_apps_tools_cache(&self) -> Result> { let managed_client = self .clients .get(CODEX_APPS_MCP_SERVER_NAME) @@ -1174,9 +1175,7 @@ impl McpConnectionManager { pub async fn resolve_tool_info(&self, tool_name: &ToolName) -> Option { let all_tools = self.list_all_tools().await; - all_tools - .into_values() - .find(|tool| tool.canonical_tool_name() == *tool_name) + all_tools.get(tool_name).cloned() } } @@ -1236,8 +1235,8 @@ fn filter_tools(tools: Vec, filter: &ToolFilter) -> Vec { } pub fn filter_non_codex_apps_mcp_tools_only( - mcp_tools: &HashMap, -) -> HashMap { + mcp_tools: &HashMap, +) -> HashMap { mcp_tools .iter() .filter(|(_, tool)| tool.server_name != CODEX_APPS_MCP_SERVER_NAME) 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 cf2889ccde..fabf6fbb2d 100644 --- a/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs @@ -252,8 +252,8 @@ fn test_qualify_tools_short_non_duplicated_names() { let qualified_tools = qualify_tools(tools); assert_eq!(qualified_tools.len(), 2); - assert!(qualified_tools.contains_key("mcp__server1__tool1")); - assert!(qualified_tools.contains_key("mcp__server1__tool2")); + assert!(qualified_tools.contains_key(&ToolName::namespaced("mcp__server1__", "tool1"))); + assert!(qualified_tools.contains_key(&ToolName::namespaced("mcp__server1__", "tool2"))); } #[test] @@ -267,7 +267,9 @@ fn test_qualify_tools_duplicated_names_skipped() { // Only the first tool should remain, the second is skipped assert_eq!(qualified_tools.len(), 1); - assert!(qualified_tools.contains_key("mcp__server1__duplicate_tool")); + assert!( + qualified_tools.contains_key(&ToolName::namespaced("mcp__server1__", "duplicate_tool")) + ); } #[test] @@ -289,7 +291,7 @@ fn test_qualify_tools_long_names_same_server() { assert_eq!(qualified_tools.len(), 2); - let mut keys: Vec<_> = qualified_tools.keys().cloned().collect(); + let mut keys: Vec<_> = qualified_tools.keys().map(ToolName::display).collect(); keys.sort(); assert!(keys.iter().all(|key| key.len() == 64)); @@ -309,10 +311,13 @@ fn test_qualify_tools_sanitizes_invalid_characters() { assert_eq!(qualified_tools.len(), 1); let (qualified_name, tool) = qualified_tools.into_iter().next().expect("one tool"); - assert_eq!(qualified_name, "mcp__server_one__tool_two_three"); + assert_eq!( + qualified_name, + ToolName::namespaced("mcp__server_one__", "tool_two_three") + ); assert_eq!( format!("{}{}", tool.callable_namespace, tool.callable_name), - qualified_name + qualified_name.display() ); // The key and callable parts are sanitized for model-visible tool calls, but @@ -324,6 +329,7 @@ fn test_qualify_tools_sanitizes_invalid_characters() { assert!( qualified_name + .display() .chars() .all(|c| c.is_ascii_alphanumeric() || c == '_'), "qualified name must be code-mode compatible: {qualified_name:?}" @@ -338,7 +344,10 @@ fn test_qualify_tools_keeps_hyphenated_mcp_tools_callable() { assert_eq!(qualified_tools.len(), 1); let (qualified_name, tool) = qualified_tools.into_iter().next().expect("one tool"); - assert_eq!(qualified_name, "mcp__music_studio__get_strudel_guide"); + assert_eq!( + qualified_name, + ToolName::namespaced("mcp__music_studio__", "get_strudel_guide") + ); assert_eq!(tool.callable_namespace, "mcp__music_studio__"); assert_eq!(tool.callable_name, "get_strudel_guide"); assert_eq!(tool.tool.name, "get-strudel-guide"); @@ -368,9 +377,10 @@ fn test_qualify_tools_disambiguates_sanitized_namespace_collisions() { .collect::>(); assert_eq!(raw_servers, HashSet::from(["basic-server", "basic_server"])); assert!( - qualified_tools - .keys() - .all(|key| key.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')), + qualified_tools.keys().all(|key| key + .display() + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_')), "qualified names must be code-mode compatible: {qualified_tools:?}" ); } @@ -641,7 +651,10 @@ async fn list_all_tools_uses_startup_snapshot_while_client_is_pending() { let tools = manager.list_all_tools().await; let tool = tools - .get("mcp__codex_apps__calendar_create_event") + .get(&ToolName::namespaced( + "mcp__codex_apps__", + "calendar_create_event", + )) .expect("tool from startup cache"); assert_eq!(tool.server_name, CODEX_APPS_MCP_SERVER_NAME); assert_eq!(tool.callable_name, "calendar_create_event"); @@ -759,7 +772,10 @@ async fn list_all_tools_uses_startup_snapshot_when_client_startup_fails() { let tools = manager.list_all_tools().await; let tool = tools - .get("mcp__codex_apps__calendar_create_event") + .get(&ToolName::namespaced( + "mcp__codex_apps__", + "calendar_create_event", + )) .expect("tool from startup cache"); assert_eq!(tool.server_name, CODEX_APPS_MCP_SERVER_NAME); assert_eq!(tool.callable_name, "calendar_create_event"); diff --git a/codex-rs/codex-mcp/src/mcp_tool_names.rs b/codex-rs/codex-mcp/src/mcp_tool_names.rs index 5a323dfe71..0f01c1c0fb 100644 --- a/codex-rs/codex-mcp/src/mcp_tool_names.rs +++ b/codex-rs/codex-mcp/src/mcp_tool_names.rs @@ -9,6 +9,7 @@ use tracing::warn; use crate::mcp::sanitize_responses_api_tool_name; use crate::mcp_connection_manager::ToolInfo; +use codex_protocol::ToolName; const MCP_TOOL_NAME_DELIMITER: &str = "__"; const MAX_TOOL_NAME_LENGTH: usize = 64; @@ -71,10 +72,10 @@ fn unique_callable_parts( tool_name: &str, raw_identity: &str, used_names: &mut HashSet, -) -> (String, String, String) { - let qualified_name = format!("{namespace}{tool_name}"); - if qualified_name.len() <= MAX_TOOL_NAME_LENGTH && used_names.insert(qualified_name.clone()) { - return (namespace.to_string(), tool_name.to_string(), qualified_name); +) -> (String, String) { + let display_name = ToolName::namespaced(namespace, tool_name).display(); + if display_name.len() <= MAX_TOOL_NAME_LENGTH && used_names.insert(display_name) { + return (namespace.to_string(), tool_name.to_string()); } let mut attempt = 0_u32; @@ -86,9 +87,9 @@ fn unique_callable_parts( }; let (namespace, tool_name) = fit_callable_parts_with_hash(namespace, tool_name, &hash_input); - let qualified_name = format!("{namespace}{tool_name}"); - if used_names.insert(qualified_name.clone()) { - return (namespace, tool_name, qualified_name); + let display_name = ToolName::namespaced(&namespace, &tool_name).display(); + if used_names.insert(display_name) { + return (namespace, tool_name); } attempt = attempt.saturating_add(1); } @@ -103,12 +104,12 @@ struct CallableToolCandidate { callable_name: String, } -/// Returns a qualified-name lookup for MCP tools. +/// Returns a callable-name lookup for MCP tools. /// /// Raw MCP server/tool names are kept on each [`ToolInfo`] for protocol calls, while /// `callable_namespace` / `callable_name` are sanitized and, when necessary, hashed so /// every model-visible `mcp__namespace__tool` name is unique and <= 64 bytes. -pub(crate) fn qualify_tools(tools: I) -> HashMap +pub(crate) fn qualify_tools(tools: I) -> HashMap where I: IntoIterator, { @@ -188,7 +189,7 @@ where let mut used_names = HashSet::new(); let mut qualified_tools = HashMap::new(); for mut candidate in candidates { - let (callable_namespace, callable_name, qualified_name) = unique_callable_parts( + let (callable_namespace, callable_name) = unique_callable_parts( &candidate.callable_namespace, &candidate.callable_name, &candidate.raw_tool_identity, @@ -196,7 +197,7 @@ where ); candidate.tool.callable_namespace = callable_namespace; candidate.tool.callable_name = callable_name; - qualified_tools.insert(qualified_name, candidate.tool); + qualified_tools.insert(candidate.tool.canonical_tool_name(), candidate.tool); } qualified_tools } diff --git a/codex-rs/core/src/codex/turn.rs b/codex-rs/core/src/codex/turn.rs index 5749bda84f..0f52f90e3c 100644 --- a/codex-rs/core/src/codex/turn.rs +++ b/codex-rs/core/src/codex/turn.rs @@ -1232,7 +1232,7 @@ pub(crate) async fn built_tools( let exposed_tool_names = mcp_tools .iter() .chain(deferred_mcp_tools.iter()) - .flat_map(|tools| tools.keys().map(String::as_str)) + .flat_map(|tools| tools.keys().map(ToolName::display)) .collect::>(); collect_unavailable_called_tools(input, &exposed_tool_names) } else { diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 103ce251fc..debc067735 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -16,6 +16,7 @@ use codex_connectors::DirectoryListResponse; use codex_login::token_data::TokenData; use codex_protocol::protocol::SandboxPolicy; use codex_tools::DiscoverableTool; +use codex_tools::ToolName; use rmcp::model::ToolAnnotations; use serde::Deserialize; use serde::de::DeserializeOwned; @@ -160,7 +161,7 @@ pub async fn list_cached_accessible_connectors_from_mcp_tools( pub(crate) fn refresh_accessible_connectors_cache_from_mcp_tools( config: &Config, auth: Option<&CodexAuth>, - mcp_tools: &HashMap, + mcp_tools: &HashMap, ) { if !config.features.enabled(Feature::Apps) { return; @@ -480,7 +481,7 @@ async fn chatgpt_get_request_with_token( } pub(crate) fn accessible_connectors_from_mcp_tools( - mcp_tools: &HashMap, + mcp_tools: &HashMap, ) -> Vec { // ToolInfo already carries plugin provenance, so app-level plugin sources // can be derived here instead of requiring a separate enrichment pass. diff --git a/codex-rs/core/src/connectors_tests.rs b/codex-rs/core/src/connectors_tests.rs index e462c7939e..b67d85035b 100644 --- a/codex-rs/core/src/connectors_tests.rs +++ b/codex-rs/core/src/connectors_tests.rs @@ -21,6 +21,7 @@ use codex_connectors::metadata::sanitize_name; use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::ToolInfo; +use codex_tools::ToolName; use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; use rmcp::model::JsonObject; @@ -128,6 +129,13 @@ fn codex_app_tool( } } +fn mcp_tool_map(tools: [ToolInfo; N]) -> HashMap { + tools + .into_iter() + .map(|tool| (tool.canonical_tool_name(), tool)) + .collect() +} + fn with_accessible_connectors_cache_cleared(f: impl FnOnce() -> R) -> R { let previous = { let mut cache_guard = ACCESSIBLE_CONNECTORS_CACHE @@ -173,39 +181,30 @@ fn merge_connectors_replaces_plugin_placeholder_name_with_accessible_name() { #[test] fn accessible_connectors_from_mcp_tools_carries_plugin_display_names() { - let tools = HashMap::from([ - ( - "mcp__codex_apps__calendar_list_events".to_string(), - codex_app_tool( - "calendar_list_events", - "calendar", - /*connector_name*/ None, - &["sample", "sample"], - ), + let tools = mcp_tool_map([ + codex_app_tool( + "calendar_list_events", + "calendar", + /*connector_name*/ None, + &["sample", "sample"], ), - ( - "mcp__codex_apps__calendar_create_event".to_string(), - codex_app_tool( - "calendar_create_event", - "calendar", - Some("Google Calendar"), - &["beta", "sample"], - ), - ), - ( - "mcp__sample__echo".to_string(), - ToolInfo { - server_name: "sample".to_string(), - callable_name: "echo".to_string(), - callable_namespace: "sample".to_string(), - server_instructions: None, - tool: test_tool_definition("echo"), - connector_id: None, - connector_name: None, - connector_description: None, - plugin_display_names: plugin_names(&["ignored"]), - }, + codex_app_tool( + "calendar_create_event", + "calendar", + Some("Google Calendar"), + &["beta", "sample"], ), + ToolInfo { + server_name: "sample".to_string(), + callable_name: "echo".to_string(), + callable_namespace: "sample".to_string(), + server_instructions: None, + tool: test_tool_definition("echo"), + connector_id: None, + connector_name: None, + connector_description: None, + plugin_display_names: plugin_names(&["ignored"]), + }, ]); let connectors = accessible_connectors_from_mcp_tools(&tools); @@ -240,24 +239,18 @@ async fn refresh_accessible_connectors_cache_from_mcp_tools_writes_latest_instal .expect("config should load"); let _ = config.features.set_enabled(Feature::Apps, /*enabled*/ true); let cache_key = accessible_connectors_cache_key(&config, /*auth*/ None); - let tools = HashMap::from([ - ( - "mcp__codex_apps__calendar_list_events".to_string(), - codex_app_tool( - "calendar_list_events", - "calendar", - Some("Google Calendar"), - &["calendar-plugin"], - ), + let tools = mcp_tool_map([ + codex_app_tool( + "calendar_list_events", + "calendar", + Some("Google Calendar"), + &["calendar-plugin"], ), - ( - "mcp__codex_apps__openai_hidden".to_string(), - codex_app_tool( - "openai_hidden", - "connector_openai_hidden", - Some("Hidden"), - &[], - ), + codex_app_tool( + "openai_hidden", + "connector_openai_hidden", + Some("Hidden"), + &[], ), ]); @@ -317,30 +310,27 @@ fn merge_connectors_unions_and_dedupes_plugin_display_names() { #[test] fn accessible_connectors_from_mcp_tools_preserves_description() { - let mcp_tools = HashMap::from([( - "mcp__codex_apps__calendar_create_event".to_string(), - ToolInfo { - server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), - callable_name: "calendar_create_event".to_string(), - callable_namespace: "mcp__codex_apps__calendar".to_string(), - server_instructions: None, - tool: Tool { - name: "calendar_create_event".to_string().into(), - title: None, - description: Some("Create a calendar event".into()), - input_schema: Arc::new(JsonObject::default()), - output_schema: None, - annotations: None, - execution: None, - icons: None, - meta: None, - }, - connector_id: Some("calendar".to_string()), - connector_name: Some("Calendar".to_string()), - connector_description: Some("Plan events".to_string()), - plugin_display_names: Vec::new(), + let mcp_tools = mcp_tool_map([ToolInfo { + server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), + callable_name: "calendar_create_event".to_string(), + callable_namespace: "mcp__codex_apps__calendar".to_string(), + server_instructions: None, + tool: Tool { + name: "calendar_create_event".to_string().into(), + title: None, + description: Some("Create a calendar event".into()), + input_schema: Arc::new(JsonObject::default()), + output_schema: None, + annotations: None, + execution: None, + icons: None, + meta: None, }, - )]); + connector_id: Some("calendar".to_string()), + connector_name: Some("Calendar".to_string()), + connector_description: Some("Plan events".to_string()), + plugin_display_names: Vec::new(), + }]); assert_eq!( accessible_connectors_from_mcp_tools(&mcp_tools), diff --git a/codex-rs/core/src/mcp_tool_exposure.rs b/codex-rs/core/src/mcp_tool_exposure.rs index 3917a1d5e7..e73d477d00 100644 --- a/codex-rs/core/src/mcp_tool_exposure.rs +++ b/codex-rs/core/src/mcp_tool_exposure.rs @@ -5,6 +5,7 @@ use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::ToolInfo as McpToolInfo; use codex_mcp::filter_non_codex_apps_mcp_tools_only; +use codex_tools::ToolName; use codex_tools::ToolsConfig; use crate::config::Config; @@ -13,12 +14,12 @@ use crate::connectors; pub(crate) const DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD: usize = 100; pub(crate) struct McpToolExposure { - pub(crate) direct_tools: HashMap, - pub(crate) deferred_tools: Option>, + pub(crate) direct_tools: HashMap, + pub(crate) deferred_tools: Option>, } pub(crate) fn build_mcp_tool_exposure( - all_mcp_tools: &HashMap, + all_mcp_tools: &HashMap, connectors: Option<&[connectors::AppInfo]>, explicitly_enabled_connectors: &[connectors::AppInfo], config: &Config, @@ -59,10 +60,10 @@ pub(crate) fn build_mcp_tool_exposure( } fn filter_codex_apps_mcp_tools( - mcp_tools: &HashMap, + mcp_tools: &HashMap, connectors: &[connectors::AppInfo], config: &Config, -) -> HashMap { +) -> HashMap { let allowed: HashSet<&str> = connectors .iter() .map(|connector| connector.id.as_str()) diff --git a/codex-rs/core/src/mcp_tool_exposure_test.rs b/codex-rs/core/src/mcp_tool_exposure_test.rs index 3cf4634970..ad9ef69719 100644 --- a/codex-rs/core/src/mcp_tool_exposure_test.rs +++ b/codex-rs/core/src/mcp_tool_exposure_test.rs @@ -11,6 +11,7 @@ use codex_protocol::config_types::WebSearchMode; use codex_protocol::config_types::WindowsSandboxLevel; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; +use codex_tools::ToolName; use codex_tools::ToolsConfig; use codex_tools::ToolsConfigParams; use pretty_assertions::assert_eq; @@ -54,9 +55,18 @@ fn make_mcp_tool( format!("mcp__{server_name}__") }; + let callable_name = if server_name == CODEX_APPS_MCP_SERVER_NAME { + connector_id + .and_then(|connector_id| tool_name.strip_prefix(connector_id)) + .unwrap_or(tool_name) + .to_string() + } else { + tool_name.to_string() + }; + ToolInfo { server_name: server_name.to_string(), - callable_name: tool_name.to_string(), + callable_name, callable_namespace: tool_namespace, server_instructions: None, tool: Tool { @@ -77,16 +87,14 @@ fn make_mcp_tool( } } -fn numbered_mcp_tools(count: usize) -> HashMap { +fn numbered_mcp_tools(count: usize) -> HashMap { (0..count) .map(|index| { let tool_name = format!("tool_{index}"); - ( - format!("mcp__rmcp__{tool_name}"), - make_mcp_tool( - "rmcp", &tool_name, /*connector_id*/ None, /*connector_name*/ None, - ), - ) + let tool = make_mcp_tool( + "rmcp", &tool_name, /*connector_id*/ None, /*connector_name*/ None, + ); + (tool.canonical_tool_name(), tool) }) .collect() } @@ -127,9 +135,13 @@ async fn directly_exposes_small_effective_tool_sets() { &tools_config, ); - let mut direct_tool_names: Vec<_> = exposure.direct_tools.keys().cloned().collect(); + let mut direct_tool_names: Vec<_> = exposure + .direct_tools + .keys() + .map(ToolName::display) + .collect(); direct_tool_names.sort(); - let mut expected_tool_names: Vec<_> = mcp_tools.keys().cloned().collect(); + let mut expected_tool_names: Vec<_> = mcp_tools.keys().map(ToolName::display).collect(); expected_tool_names.sort(); assert_eq!(direct_tool_names, expected_tool_names); assert!(exposure.deferred_tools.is_none()); @@ -154,9 +166,9 @@ async fn searches_large_effective_tool_sets() { .deferred_tools .as_ref() .expect("large tool sets should be discoverable through tool_search"); - let mut deferred_tool_names: Vec<_> = deferred_tools.keys().cloned().collect(); + let mut deferred_tool_names: Vec<_> = deferred_tools.keys().map(ToolName::display).collect(); deferred_tool_names.sort(); - let mut expected_tool_names: Vec<_> = mcp_tools.keys().cloned().collect(); + let mut expected_tool_names: Vec<_> = mcp_tools.keys().map(ToolName::display).collect(); expected_tool_names.sort(); assert_eq!(deferred_tool_names, expected_tool_names); } @@ -166,15 +178,13 @@ async fn directly_exposes_explicit_apps_without_deferred_overlap() { let config = test_config().await; let tools_config = tools_config_for_mcp_tool_exposure(/*search_tool*/ true).await; let mut mcp_tools = numbered_mcp_tools(DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD - 1); - mcp_tools.extend([( - "mcp__codex_apps__calendar_create_event".to_string(), - make_mcp_tool( - CODEX_APPS_MCP_SERVER_NAME, - "calendar_create_event", - Some("calendar"), - Some("Calendar"), - ), - )]); + let calendar_tool = make_mcp_tool( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_create_event", + Some("calendar"), + Some("Calendar"), + ); + mcp_tools.insert(calendar_tool.canonical_tool_name(), calendar_tool); let connectors = vec![make_connector("calendar", "Calendar")]; let exposure = build_mcp_tool_exposure( @@ -185,7 +195,11 @@ async fn directly_exposes_explicit_apps_without_deferred_overlap() { &tools_config, ); - let mut tool_names: Vec = exposure.direct_tools.into_keys().collect(); + let mut tool_names: Vec = exposure + .direct_tools + .into_keys() + .map(|name| name.display()) + .collect(); tool_names.sort(); assert_eq!( tool_names, @@ -200,13 +214,16 @@ async fn directly_exposes_explicit_apps_without_deferred_overlap() { .as_ref() .expect("large tool sets should be discoverable through tool_search"); assert!( - tool_names - .iter() - .all(|direct_tool_name| !deferred_tools.contains_key(direct_tool_name)), + tool_names.iter().all(|direct_tool_name| !deferred_tools + .keys() + .any(|deferred_tool_name| deferred_tool_name.display() == *direct_tool_name)), "direct tools should not also be deferred: {tool_names:?}" ); - assert!(!deferred_tools.contains_key("mcp__codex_apps__calendar_create_event")); - assert!(deferred_tools.contains_key("mcp__rmcp__tool_0")); + assert!(!deferred_tools.contains_key(&ToolName::namespaced( + "mcp__codex_apps__calendar", + "_create_event" + ))); + assert!(deferred_tools.contains_key(&ToolName::namespaced("mcp__rmcp__", "tool_0"))); } #[tokio::test] @@ -217,22 +234,18 @@ async fn always_defer_feature_preserves_explicit_apps() { .enable(Feature::ToolSearchAlwaysDeferMcpTools) .expect("test config should allow feature update"); let tools_config = tools_config_for_mcp_tool_exposure(/*search_tool*/ true).await; + let rmcp_tool = make_mcp_tool( + "rmcp", "tool", /*connector_id*/ None, /*connector_name*/ None, + ); + let calendar_tool = make_mcp_tool( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_create_event", + Some("calendar"), + Some("Calendar"), + ); let mcp_tools = HashMap::from([ - ( - "mcp__rmcp__tool".to_string(), - make_mcp_tool( - "rmcp", "tool", /*connector_id*/ None, /*connector_name*/ None, - ), - ), - ( - "mcp__codex_apps__calendar_create_event".to_string(), - make_mcp_tool( - CODEX_APPS_MCP_SERVER_NAME, - "calendar_create_event", - Some("calendar"), - Some("Calendar"), - ), - ), + (rmcp_tool.canonical_tool_name(), rmcp_tool), + (calendar_tool.canonical_tool_name(), calendar_tool), ]); let connectors = vec![make_connector("calendar", "Calendar")]; @@ -244,7 +257,11 @@ async fn always_defer_feature_preserves_explicit_apps() { &tools_config, ); - let mut direct_tool_names: Vec = exposure.direct_tools.into_keys().collect(); + let mut direct_tool_names: Vec = exposure + .direct_tools + .into_keys() + .map(|name| name.display()) + .collect(); direct_tool_names.sort(); assert_eq!( direct_tool_names, @@ -254,6 +271,9 @@ async fn always_defer_feature_preserves_explicit_apps() { .deferred_tools .as_ref() .expect("MCP tools should be discoverable through tool_search"); - assert!(deferred_tools.contains_key("mcp__rmcp__tool")); - assert!(!deferred_tools.contains_key("mcp__codex_apps__calendar_create_event")); + assert!(deferred_tools.contains_key(&ToolName::namespaced("mcp__rmcp__", "tool"))); + assert!(!deferred_tools.contains_key(&ToolName::namespaced( + "mcp__codex_apps__calendar", + "_create_event" + ))); } diff --git a/codex-rs/core/src/plugins/injection.rs b/codex-rs/core/src/plugins/injection.rs index 2f6ec26260..7938fcfb37 100644 --- a/codex-rs/core/src/plugins/injection.rs +++ b/codex-rs/core/src/plugins/injection.rs @@ -10,10 +10,11 @@ use crate::plugins::PluginCapabilitySummary; use crate::plugins::render_explicit_plugin_instructions; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::ToolInfo; +use codex_tools::ToolName; pub(crate) fn build_plugin_injections( mentioned_plugins: &[PluginCapabilitySummary], - mcp_tools: &HashMap, + mcp_tools: &HashMap, available_connectors: &[connectors::AppInfo], ) -> Vec { if mentioned_plugins.is_empty() { diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index 4bab746ac2..4692139ba2 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -11,6 +11,7 @@ use bm25::SearchEngineBuilder; use codex_mcp::ToolInfo; use codex_tools::TOOL_SEARCH_DEFAULT_LIMIT; use codex_tools::TOOL_SEARCH_TOOL_NAME; +use codex_tools::ToolName; use codex_tools::ToolSearchResultSource; use codex_tools::collect_tool_search_output_tools; @@ -18,14 +19,14 @@ const COMPUTER_USE_MCP_SERVER_NAME: &str = "computer-use"; const COMPUTER_USE_TOOL_SEARCH_LIMIT: usize = 20; pub struct ToolSearchHandler { - entries: Vec<(String, ToolInfo)>, + entries: Vec<(ToolName, ToolInfo)>, search_engine: SearchEngine, } impl ToolSearchHandler { - pub fn new(tools: std::collections::HashMap) -> Self { - let mut entries: Vec<(String, ToolInfo)> = tools.into_iter().collect(); - entries.sort_by(|a, b| a.0.cmp(&b.0)); + pub fn new(tools: std::collections::HashMap) -> Self { + let mut entries: Vec<(ToolName, ToolInfo)> = tools.into_iter().collect(); + entries.sort_by_key(|entry| entry.0.display()); let documents: Vec> = entries .iter() @@ -111,7 +112,7 @@ impl ToolSearchHandler { query: &str, limit: usize, use_default_limit: bool, - ) -> Vec<&(String, ToolInfo)> { + ) -> Vec<&(ToolName, ToolInfo)> { let mut results = self .search_engine .search(query, limit) @@ -137,7 +138,7 @@ impl ToolSearchHandler { } } -fn limit_results_per_server(results: Vec<&(String, ToolInfo)>) -> Vec<&(String, ToolInfo)> { +fn limit_results_per_server(results: Vec<&(ToolName, ToolInfo)>) -> Vec<&(ToolName, ToolInfo)> { results .into_iter() .scan( @@ -165,9 +166,9 @@ fn default_limit_for_server(server_name: &str) -> usize { } } -fn build_search_text(name: &str, info: &ToolInfo) -> String { +fn build_search_text(name: &ToolName, info: &ToolInfo) -> String { let mut parts = vec![ - name.to_string(), + name.display(), info.callable_name.clone(), info.tool.name.to_string(), info.server_name.clone(), @@ -320,12 +321,12 @@ mod tests { server_name: &str, description_prefix: &str, count: usize, - ) -> std::collections::HashMap { + ) -> std::collections::HashMap { (0..count) .map(|index| { let tool_name = format!("tool_{index:03}"); ( - format!("mcp__{server_name}__{tool_name}"), + ToolName::namespaced(format!("mcp__{server_name}__"), tool_name.clone()), tool_info(server_name, &tool_name, description_prefix), ) }) @@ -360,7 +361,7 @@ mod tests { } } - fn count_results_for_server(results: &[&(String, ToolInfo)], server_name: &str) -> usize { + fn count_results_for_server(results: &[&(ToolName, ToolInfo)], server_name: &str) -> usize { results .iter() .filter(|(_, tool)| tool.server_name == server_name) diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index d57ca5473a..95fb85adca 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -43,8 +43,8 @@ pub struct ToolRouter { } pub(crate) struct ToolRouterParams<'a> { - pub(crate) mcp_tools: Option>, - pub(crate) deferred_mcp_tools: Option>, + 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>, diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 1a1a9f2e5f..eb50b8e3bc 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -41,12 +41,12 @@ struct McpToolPlanInputs<'a> { 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 - .values() - .map(|tool| ToolRegistryPlanMcpTool { - name: tool.canonical_tool_name(), + .iter() + .map(|(name, tool)| ToolRegistryPlanMcpTool { + name: name.clone(), tool: &tool.tool, }) .collect(), @@ -70,8 +70,8 @@ fn map_mcp_tools_for_plan(mcp_tools: &HashMap) -> McpToolPlanI pub(crate) fn build_specs_with_discoverable_tools( config: &ToolsConfig, - mcp_tools: Option>, - deferred_mcp_tools: Option>, + mcp_tools: Option>, + deferred_mcp_tools: Option>, unavailable_called_tools: Vec, discoverable_tools: Option>, dynamic_tools: &[DynamicToolSpec], @@ -113,9 +113,9 @@ pub(crate) fn build_specs_with_discoverable_tools( let mcp_tool_plan_inputs = mcp_tools.as_ref().map(map_mcp_tools_for_plan); let deferred_mcp_tool_sources = deferred_mcp_tools.as_ref().map(|tools| { tools - .values() - .map(|tool| ToolRegistryPlanDeferredTool { - name: tool.canonical_tool_name(), + .iter() + .map(|(name, tool)| ToolRegistryPlanDeferredTool { + name: name.clone(), server_name: tool.server_name.as_str(), connector_name: tool.connector_name.as_deref(), connector_description: tool.connector_description.as_deref(), diff --git a/codex-rs/core/src/tools/spec_tests.rs b/codex-rs/core/src/tools/spec_tests.rs index 9de9c81a92..f6c30ce429 100644 --- a/codex-rs/core/src/tools/spec_tests.rs +++ b/codex-rs/core/src/tools/spec_tests.rs @@ -88,6 +88,13 @@ fn mcp_tool_info_with_display_name(display_name: &str, tool: rmcp::model::Tool) } } +fn mcp_tool_map(tools: [ToolInfo; N]) -> HashMap { + tools + .into_iter() + .map(|tool| (tool.canonical_tool_name(), tool)) + .collect() +} + fn discoverable_connector(id: &str, name: &str, description: &str) -> DiscoverableTool { let slug = name.replace(' ', "-").to_lowercase(); DiscoverableTool::Connector(Box::new(AppInfo { @@ -265,8 +272,8 @@ async fn model_info_from_models_json(slug: &str) -> ModelInfo { /// Builds the tool registry builder while collecting tool specs for later serialization. fn build_specs( config: &ToolsConfig, - mcp_tools: Option>, - deferred_mcp_tools: Option>, + mcp_tools: Option>, + deferred_mcp_tools: Option>, dynamic_tools: &[DynamicToolSpec], ) -> ToolRegistryBuilder { build_specs_with_unavailable_tools( @@ -280,8 +287,8 @@ fn build_specs( fn build_specs_with_unavailable_tools( config: &ToolsConfig, - mcp_tools: Option>, - deferred_mcp_tools: Option>, + mcp_tools: Option>, + deferred_mcp_tools: Option>, unavailable_called_tools: Vec, dynamic_tools: &[DynamicToolSpec], ) -> ToolRegistryBuilder { @@ -880,24 +887,21 @@ async fn search_tool_description_falls_back_to_connector_name_without_descriptio let (tools, _) = build_specs( &tools_config, /*mcp_tools*/ None, - Some(HashMap::from([( - "mcp__codex_apps__calendar_create_event".to_string(), - ToolInfo { - server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), - callable_name: "_create_event".to_string(), - callable_namespace: "mcp__codex_apps__calendar".to_string(), - server_instructions: None, - tool: mcp_tool( - "calendar_create_event", - "Create calendar event", - serde_json::json!({"type": "object"}), - ), - connector_id: Some("calendar".to_string()), - connector_name: Some("Calendar".to_string()), - plugin_display_names: Vec::new(), - connector_description: None, - }, - )])), + Some(mcp_tool_map([ToolInfo { + server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), + callable_name: "_create_event".to_string(), + callable_namespace: "mcp__codex_apps__calendar".to_string(), + server_instructions: None, + tool: mcp_tool( + "calendar_create_event", + "Create calendar event", + serde_json::json!({"type": "object"}), + ), + connector_id: Some("calendar".to_string()), + connector_name: Some("Calendar".to_string()), + plugin_display_names: Vec::new(), + connector_description: None, + }])), &[], ) .build(); @@ -931,57 +935,48 @@ async fn search_tool_registers_namespaced_mcp_tool_aliases() { let (_, registry) = build_specs( &tools_config, /*mcp_tools*/ None, - Some(HashMap::from([ - ( - "mcp__codex_apps__calendar_create_event".to_string(), - ToolInfo { - server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), - callable_name: "_create_event".to_string(), - callable_namespace: "mcp__codex_apps__calendar".to_string(), - server_instructions: None, - tool: mcp_tool( - "calendar-create-event", - "Create calendar event", - serde_json::json!({"type": "object"}), - ), - connector_id: Some("calendar".to_string()), - connector_name: Some("Calendar".to_string()), - connector_description: None, - plugin_display_names: Vec::new(), - }, - ), - ( - "mcp__codex_apps__calendar_list_events".to_string(), - ToolInfo { - server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), - callable_name: "_list_events".to_string(), - callable_namespace: "mcp__codex_apps__calendar".to_string(), - server_instructions: None, - tool: mcp_tool( - "calendar-list-events", - "List calendar events", - serde_json::json!({"type": "object"}), - ), - connector_id: Some("calendar".to_string()), - connector_name: Some("Calendar".to_string()), - connector_description: None, - plugin_display_names: Vec::new(), - }, - ), - ( - "mcp__rmcp__echo".to_string(), - ToolInfo { - server_name: "rmcp".to_string(), - callable_name: "echo".to_string(), - callable_namespace: "mcp__rmcp__".to_string(), - server_instructions: None, - tool: mcp_tool("echo", "Echo", serde_json::json!({"type": "object"})), - connector_id: None, - connector_name: None, - connector_description: None, - plugin_display_names: Vec::new(), - }, - ), + Some(mcp_tool_map([ + ToolInfo { + server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), + callable_name: "_create_event".to_string(), + callable_namespace: "mcp__codex_apps__calendar".to_string(), + server_instructions: None, + tool: mcp_tool( + "calendar-create-event", + "Create calendar event", + serde_json::json!({"type": "object"}), + ), + connector_id: Some("calendar".to_string()), + connector_name: Some("Calendar".to_string()), + connector_description: None, + plugin_display_names: Vec::new(), + }, + ToolInfo { + server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), + callable_name: "_list_events".to_string(), + callable_namespace: "mcp__codex_apps__calendar".to_string(), + server_instructions: None, + tool: mcp_tool( + "calendar-list-events", + "List calendar events", + serde_json::json!({"type": "object"}), + ), + connector_id: Some("calendar".to_string()), + connector_name: Some("Calendar".to_string()), + connector_description: None, + plugin_display_names: Vec::new(), + }, + ToolInfo { + server_name: "rmcp".to_string(), + callable_name: "echo".to_string(), + callable_namespace: "mcp__rmcp__".to_string(), + server_instructions: None, + tool: mcp_tool("echo", "Echo", serde_json::json!({"type": "object"})), + connector_id: None, + connector_name: None, + connector_description: None, + plugin_display_names: Vec::new(), + }, ])), &[], ) @@ -1015,14 +1010,11 @@ async fn direct_mcp_tools_register_namespaced_handlers() { 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"}), - )), - )])), + Some(mcp_tool_map([mcp_tool_info(mcp_tool( + "echo", + "Echo", + serde_json::json!({"type": "object"}), + ))])), /*deferred_mcp_tools*/ None, &[], ) @@ -1101,20 +1093,17 @@ async fn test_mcp_tool_property_missing_type_defaults_to_string() { let (tools, _) = build_specs( &tools_config, - Some(HashMap::from([( - "dash/search".to_string(), - mcp_tool_info_with_display_name( - "dash/search", - mcp_tool( - "search", - "Search docs", - serde_json::json!({ - "type": "object", - "properties": { - "query": {"description": "search query"} - } - }), - ), + Some(mcp_tool_map([mcp_tool_info_with_display_name( + "dash/search", + mcp_tool( + "search", + "Search docs", + serde_json::json!({ + "type": "object", + "properties": { + "query": {"description": "search query"} + } + }), ), )])), /*deferred_mcp_tools*/ None, @@ -1164,18 +1153,15 @@ async fn test_mcp_tool_preserves_integer_schema() { let (tools, _) = build_specs( &tools_config, - Some(HashMap::from([( - "dash/paginate".to_string(), - mcp_tool_info_with_display_name( - "dash/paginate", - mcp_tool( - "paginate", - "Pagination", - serde_json::json!({ - "type": "object", - "properties": {"page": {"type": "integer"}} - }), - ), + Some(mcp_tool_map([mcp_tool_info_with_display_name( + "dash/paginate", + mcp_tool( + "paginate", + "Pagination", + serde_json::json!({ + "type": "object", + "properties": {"page": {"type": "integer"}} + }), ), )])), /*deferred_mcp_tools*/ None, @@ -1226,18 +1212,15 @@ async fn test_mcp_tool_array_without_items_gets_default_string_items() { let (tools, _) = build_specs( &tools_config, - Some(HashMap::from([( - "dash/tags".to_string(), - mcp_tool_info_with_display_name( - "dash/tags", - mcp_tool( - "tags", - "Tags", - serde_json::json!({ - "type": "object", - "properties": {"tags": {"type": "array"}} - }), - ), + Some(mcp_tool_map([mcp_tool_info_with_display_name( + "dash/tags", + mcp_tool( + "tags", + "Tags", + serde_json::json!({ + "type": "object", + "properties": {"tags": {"type": "array"}} + }), ), )])), /*deferred_mcp_tools*/ None, @@ -1290,20 +1273,17 @@ async fn test_mcp_tool_anyof_defaults_to_string() { let (tools, _) = build_specs( &tools_config, - Some(HashMap::from([( - "dash/value".to_string(), - mcp_tool_info_with_display_name( - "dash/value", - mcp_tool( - "value", - "AnyOf Value", - serde_json::json!({ - "type": "object", - "properties": { - "value": {"anyOf": [{"type": "string"}, {"type": "number"}]} - } - }), - ), + Some(mcp_tool_map([mcp_tool_info_with_display_name( + "dash/value", + mcp_tool( + "value", + "AnyOf Value", + serde_json::json!({ + "type": "object", + "properties": { + "value": {"anyOf": [{"type": "string"}, {"type": "number"}]} + } + }), ), )])), /*deferred_mcp_tools*/ None, @@ -1358,16 +1338,14 @@ async fn test_get_openai_tools_mcp_tools_with_additional_properties_schema() { }); let (tools, _) = build_specs( &tools_config, - Some(HashMap::from([( - "test_server/do_something_cool".to_string(), - mcp_tool_info_with_display_name( - "test_server/do_something_cool", - mcp_tool( - "do_something_cool", - "Do something cool", - serde_json::json!({ - "type": "object", - "properties": { + Some(mcp_tool_map([mcp_tool_info_with_display_name( + "test_server/do_something_cool", + mcp_tool( + "do_something_cool", + "Do something cool", + serde_json::json!({ + "type": "object", + "properties": { "string_argument": {"type": "string"}, "number_argument": {"type": "number"}, "object_argument": { @@ -1384,11 +1362,10 @@ async fn test_get_openai_tools_mcp_tools_with_additional_properties_schema() { }, "required": ["addtl_prop"], "additionalProperties": false - } } } - }), - ), + } + }), ), )])), /*deferred_mcp_tools*/ None, diff --git a/codex-rs/core/src/unavailable_tool.rs b/codex-rs/core/src/unavailable_tool.rs index 39ba21f9b8..9884ec5835 100644 --- a/codex-rs/core/src/unavailable_tool.rs +++ b/codex-rs/core/src/unavailable_tool.rs @@ -6,7 +6,7 @@ use codex_tools::ToolName; pub(crate) fn collect_unavailable_called_tools( input: &[ResponseItem], - exposed_tool_names: &HashSet<&str>, + exposed_tool_names: &HashSet, ) -> Vec { let mut unavailable_tools = BTreeMap::new(); @@ -78,7 +78,10 @@ mod tests { #[test] fn collect_unavailable_called_tools_skips_currently_available_tools() { - let exposed_tool_names = HashSet::from(["mcp__server__lookup", "mcp__server__search"]); + let exposed_tool_names = HashSet::from([ + "mcp__server__lookup".to_string(), + "mcp__server__search".to_string(), + ]); let input = vec![ function_call("mcp__server__lookup", /*namespace*/ None), function_call("mcp__server__search", /*namespace*/ None),