From 9e29c221b9cc110db6480a16237f45723c1dd9b3 Mon Sep 17 00:00:00 2001 From: Hunter B Date: Thu, 4 Jun 2026 18:22:46 -0700 Subject: [PATCH] fix(mcp): preserve underscored MCP server names in tool routing parse_prefixed_name now matches the qualified mcp__ name against the set of registered server names (connections + configured servers) and prefers the longest matching server name, instead of naively splitting on the first underscore. Tools on servers whose names contain underscores (e.g. "my_db") are now reachable, and an overlapping pair like "my" and "my_db" routes to the correct server. Falls back to the legacy first-underscore split when no registered server matches, preserving backward compatibility. Harvested from PR #2747 by @cyq1017; supersedes the equivalent fix in PR #2746 by @puneetdixit200. Both contributors diagnosed and fixed issue #2744; #2747 landed for its longest-match tie-break test coverage. Fixes #2744. Co-authored-by: cyq1017 <61975706+cyq1017@users.noreply.github.com> Co-authored-by: puneetdixit200 <236133619+puneetdixit200@users.noreply.github.com> --- .github/AUTHOR_MAP | 1 + crates/tui/src/mcp.rs | 132 +++++++++++++++++++++++++++++++++++++++++- 2 files changed, 131 insertions(+), 2 deletions(-) diff --git a/.github/AUTHOR_MAP b/.github/AUTHOR_MAP index 26bca8cf..d1997277 100644 --- a/.github/AUTHOR_MAP +++ b/.github/AUTHOR_MAP @@ -88,3 +88,4 @@ hsdbeebou = hsdbeebou <284843096+hsdbeebou@users.noreply.github.com> tdccccc = tdccccc <79492752+tdccccc@users.noreply.github.com> greyfreedom = greyfreedom <11493871+greyfreedom@users.noreply.github.com> greyfreedom@163.com = greyfreedom <11493871+greyfreedom@users.noreply.github.com> +puneetdixit200 = puneetdixit200 <236133619+puneetdixit200@users.noreply.github.com> diff --git a/crates/tui/src/mcp.rs b/crates/tui/src/mcp.rs index 32a342b1..038bfe16 100644 --- a/crates/tui/src/mcp.rs +++ b/crates/tui/src/mcp.rs @@ -2241,10 +2241,30 @@ impl McpPool { /// Parse a prefixed name into (server_name, tool_name) fn parse_prefixed_name<'a>(&self, prefixed_name: &'a str) -> Result<(&'a str, &'a str)> { - if !prefixed_name.starts_with("mcp_") { + let Some(rest) = prefixed_name.strip_prefix("mcp_") else { anyhow::bail!("Invalid MCP tool name: {prefixed_name}"); + }; + + let mut best_match: Option<(&str, &str)> = None; + for server in self.connections.keys().chain(self.config.servers.keys()) { + let Some(tool) = rest + .strip_prefix(server) + .and_then(|tail| tail.strip_prefix('_')) + else { + continue; + }; + if tool.is_empty() { + continue; + } + if best_match.is_none_or(|(matched, _)| server.len() > matched.len()) { + best_match = Some((&rest[..server.len()], tool)); + } } - let rest = &prefixed_name[4..]; + + if let Some((server, tool)) = best_match { + return Ok((server, tool)); + } + let Some((server, tool)) = rest.split_once('_') else { anyhow::bail!("Invalid MCP tool name format: {prefixed_name}"); }; @@ -3643,6 +3663,114 @@ mod tests { ); } + #[tokio::test] + async fn mcp_pool_call_tool_preserves_server_names_with_underscores() { + let sent = Arc::new(Mutex::new(Vec::new())); + let transport = ScriptedValueTransport { + sent: Arc::clone(&sent), + responses: VecDeque::from([json_frame(serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "result": {"ok": true} + }))]), + }; + let mut conn = test_connection(Box::new(transport)); + conn.name = "my_db".to_string(); + conn.tools = vec![McpTool { + name: "execute_sql".to_string(), + description: None, + input_schema: serde_json::json!({}), + }]; + + let mut pool = McpPool::new(McpConfig { + timeouts: McpTimeouts::default(), + servers: HashMap::new(), + }); + pool.connections.insert("my_db".to_string(), conn); + + let result = pool + .call_tool( + "mcp_my_db_execute_sql", + serde_json::json!({"query": "select 1"}), + ) + .await + .unwrap(); + + assert_eq!(result, serde_json::json!({"ok": true})); + let sent = sent.lock().unwrap(); + assert_eq!(sent[0]["method"], "tools/call"); + assert_eq!(sent[0]["params"]["name"], "execute_sql"); + assert_eq!( + sent[0]["params"]["arguments"], + serde_json::json!({"query": "select 1"}) + ); + } + + #[tokio::test] + async fn mcp_pool_call_tool_prefers_longest_matching_server_name() { + let sent_short = Arc::new(Mutex::new(Vec::new())); + let short_transport = ScriptedValueTransport { + sent: Arc::clone(&sent_short), + responses: VecDeque::from([json_frame(serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "result": {"short": true} + }))]), + }; + let mut short_conn = test_connection(Box::new(short_transport)); + short_conn.name = "my".to_string(); + short_conn.tools = vec![McpTool { + name: "db_execute_sql".to_string(), + description: None, + input_schema: serde_json::json!({}), + }]; + + let sent_long = Arc::new(Mutex::new(Vec::new())); + let long_transport = ScriptedValueTransport { + sent: Arc::clone(&sent_long), + responses: VecDeque::from([json_frame(serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "result": {"long": true} + }))]), + }; + let mut long_conn = test_connection(Box::new(long_transport)); + long_conn.name = "my_db".to_string(); + long_conn.tools = vec![McpTool { + name: "execute_sql".to_string(), + description: None, + input_schema: serde_json::json!({}), + }]; + + let mut pool = McpPool::new(McpConfig { + timeouts: McpTimeouts::default(), + servers: HashMap::new(), + }); + pool.connections.insert("my".to_string(), short_conn); + pool.connections.insert("my_db".to_string(), long_conn); + + let result = pool + .call_tool( + "mcp_my_db_execute_sql", + serde_json::json!({"query": "select 1"}), + ) + .await + .unwrap(); + + assert_eq!(result, serde_json::json!({"long": true})); + assert!( + sent_short.lock().unwrap().is_empty(), + "the shorter server name must not receive the tool call" + ); + let sent_long = sent_long.lock().unwrap(); + assert_eq!(sent_long[0]["method"], "tools/call"); + assert_eq!(sent_long[0]["params"]["name"], "execute_sql"); + assert_eq!( + sent_long[0]["params"]["arguments"], + serde_json::json!({"query": "select 1"}) + ); + } + #[tokio::test] async fn json_rpc_session_error_is_marked_stale() { let sent = Arc::new(Mutex::new(Vec::new()));