diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 4c852d48fa..870047dcfe 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -3952,20 +3952,11 @@ impl Session { name: &str, namespace: &Option, ) -> Option<(String, String)> { - let tool_name = if let Some(namespace) = namespace { - if name.starts_with(namespace.as_str()) { - name - } else { - &format!("{namespace}{name}") - } - } else { - name - }; self.services .mcp_connection_manager .read() .await - .parse_tool_name(tool_name) + .parse_tool_name(name, namespace.as_deref()) .await } diff --git a/codex-rs/core/src/mcp_connection_manager.rs b/codex-rs/core/src/mcp_connection_manager.rs index 74422d0740..4147702067 100644 --- a/codex-rs/core/src/mcp_connection_manager.rs +++ b/codex-rs/core/src/mcp_connection_manager.rs @@ -198,6 +198,79 @@ where qualified_tools } +fn parsed_tool_name(tool: &ToolInfo) -> (String, String) { + (tool.server_name.clone(), tool.tool.name.to_string()) +} + +fn normalize_requested_tool_name(name: &str) -> String { + sanitize_name(name).replace('-', "_") +} + +fn trim_optional_leading_underscore(value: &str) -> &str { + value.strip_prefix('_').unwrap_or(value) +} + +fn tool_name_matches_requested_name(tool: &ToolInfo, requested_name: &str) -> bool { + if tool.tool_name == requested_name || tool.tool.name.as_ref() == requested_name { + return true; + } + + if normalize_requested_tool_name(tool.tool.name.as_ref()) == requested_name { + return true; + } + + tool.server_name == CODEX_APPS_MCP_SERVER_NAME + && trim_optional_leading_underscore(tool.tool_name.as_str()) + == trim_optional_leading_underscore(requested_name) +} + +fn find_unique_matching_tool<'a>( + tools: &'a HashMap, + requested_name: &str, + namespace: Option<&str>, +) -> Option<&'a ToolInfo> { + let mut matches = tools.values().filter(|tool| { + namespace.map_or_else( + || tool.server_name == CODEX_APPS_MCP_SERVER_NAME, + |namespace| tool.tool_namespace == namespace, + ) && tool_name_matches_requested_name(tool, requested_name) + }); + let first = matches.next()?; + if matches.next().is_some() { + return None; + } + Some(first) +} + +fn parse_tool_name_from_tools( + tools: &HashMap, + tool_name: &str, + namespace: Option<&str>, +) -> Option<(String, String)> { + if let Some(tool) = tools.get(tool_name) { + return Some(parsed_tool_name(tool)); + } + + if let Some(namespace) = namespace { + let qualified_name = if tool_name.starts_with(namespace) { + tool_name.to_string() + } else { + format!("{namespace}{tool_name}") + }; + if let Some(tool) = tools.get(qualified_name.as_str()) { + return Some(parsed_tool_name(tool)); + } + + if let Some(stripped_name) = tool_name.strip_prefix(namespace) + && let Some(tool) = find_unique_matching_tool(tools, stripped_name, Some(namespace)) + { + return Some(parsed_tool_name(tool)); + } + } + + find_unique_matching_tool(tools, tool_name, namespace).map(parsed_tool_name) +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct ToolInfo { pub(crate) server_name: String, @@ -1095,11 +1168,13 @@ impl McpConnectionManager { .with_context(|| format!("resources/read failed for `{server}` ({uri})")) } - pub async fn parse_tool_name(&self, tool_name: &str) -> Option<(String, String)> { - self.list_all_tools() - .await - .get(tool_name) - .map(|tool| (tool.server_name.clone(), tool.tool.name.to_string())) + pub async fn parse_tool_name( + &self, + tool_name: &str, + namespace: Option<&str>, + ) -> Option<(String, String)> { + let tools = self.list_all_tools().await; + parse_tool_name_from_tools(&tools, tool_name, namespace) } pub async fn notify_sandbox_state_change(&self, sandbox_state: &SandboxState) -> Result<()> { diff --git a/codex-rs/core/src/mcp_connection_manager_tests.rs b/codex-rs/core/src/mcp_connection_manager_tests.rs index c5f7fc4a40..efb54e965d 100644 --- a/codex-rs/core/src/mcp_connection_manager_tests.rs +++ b/codex-rs/core/src/mcp_connection_manager_tests.rs @@ -1,7 +1,9 @@ use super::*; use codex_protocol::protocol::GranularApprovalConfig; use codex_protocol::protocol::McpAuthStatus; +use pretty_assertions::assert_eq; use rmcp::model::JsonObject; +use std::collections::HashMap; use std::collections::HashSet; use std::sync::Arc; use tempfile::tempdir; @@ -45,6 +47,39 @@ fn create_test_tool_with_connector( tool } +fn create_namespaced_codex_app_tool( + qualified_name: &str, + tool_name: &str, + raw_tool_name: &str, + namespace: &str, + connector_id: &str, + connector_name: &str, +) -> (String, ToolInfo) { + ( + qualified_name.to_string(), + ToolInfo { + server_name: CODEX_APPS_MCP_SERVER_NAME.to_string(), + tool_name: tool_name.to_string(), + tool_namespace: namespace.to_string(), + tool: Tool { + name: raw_tool_name.to_string().into(), + title: None, + description: Some(format!("Test tool: {raw_tool_name}").into()), + input_schema: Arc::new(JsonObject::default()), + output_schema: None, + annotations: None, + execution: None, + icons: None, + meta: None, + }, + connector_id: Some(connector_id.to_string()), + connector_name: Some(connector_name.to_string()), + plugin_display_names: Vec::new(), + connector_description: None, + }, + ) +} + fn create_codex_apps_tools_cache_context( codex_home: PathBuf, account_id: Option<&str>, @@ -181,6 +216,77 @@ fn test_qualify_tools_sanitizes_invalid_characters() { ); } +#[test] +fn parse_tool_name_from_tools_matches_unique_codex_app_alias_without_namespace() { + let tools = HashMap::from([create_namespaced_codex_app_tool( + "mcp__codex_apps__google_docs_create_document", + "_create_document", + "google_docs_create_document", + "mcp__codex_apps__google_docs", + "google_docs", + "Google Docs", + )]); + + assert_eq!( + parse_tool_name_from_tools(&tools, "_create_document", None), + Some(( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + "google_docs_create_document".to_string(), + )) + ); +} + +#[test] +fn parse_tool_name_from_tools_rejects_ambiguous_codex_app_alias_without_namespace() { + let tools = HashMap::from([ + create_namespaced_codex_app_tool( + "mcp__codex_apps__google_docs_create_document", + "_create_document", + "google_docs_create_document", + "mcp__codex_apps__google_docs", + "google_docs", + "Google Docs", + ), + create_namespaced_codex_app_tool( + "mcp__codex_apps__notion_create_document", + "_create_document", + "notion_create_document", + "mcp__codex_apps__notion", + "notion", + "Notion", + ), + ]); + + assert_eq!( + parse_tool_name_from_tools(&tools, "_create_document", None), + None + ); +} + +#[test] +fn parse_tool_name_from_tools_matches_namespaced_codex_app_raw_name() { + let tools = HashMap::from([create_namespaced_codex_app_tool( + "mcp__codex_apps__google_docs_create_document", + "_create_document", + "google_docs_create_document", + "mcp__codex_apps__google_docs", + "google_docs", + "Google Docs", + )]); + + assert_eq!( + parse_tool_name_from_tools( + &tools, + "google_docs_create_document", + Some("mcp__codex_apps__google_docs"), + ), + Some(( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + "google_docs_create_document".to_string(), + )) + ); +} + #[test] fn tool_filter_allows_by_default() { let filter = ToolFilter::default();