diff --git a/src/config/request_context.rs b/src/config/request_context.rs index ae4215f..14f0a01 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -23,8 +23,8 @@ use crate::function::{ user_interaction::USER_FUNCTION_PREFIX, }; use crate::mcp::{ - MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX, - MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpAuthReason, McpAuthRequired, is_auth_required_error, + MCP_INVOKE_META_FUNCTION_NAME_PREFIX, McpAuthReason, McpAuthRequired, is_auth_required_error, + is_mcp_meta_function, mcp_meta_function_names, }; use crate::rag::Rag; use crate::supervisor::Supervisor; @@ -2011,11 +2011,7 @@ impl RequestContext { .functions .declarations() .iter() - .filter(|v| { - !v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - && !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - && !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }) + .filter(|v| !is_mcp_meta_function(&v.name)) .map(|v| v.name.to_string()) .collect(); @@ -2025,11 +2021,7 @@ impl RequestContext { .functions() .declarations() .iter() - .filter(|v| { - !v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - && !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - && !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }) + .filter(|v| !is_mcp_meta_function(&v.name)) .map(|v| v.name.to_string()), ); } @@ -2102,11 +2094,7 @@ impl RequestContext { .declarations() .to_vec() .into_iter() - .filter(|v| { - !v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - && !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - && !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }) + .filter(|v| !is_mcp_meta_function(&v.name)) .collect(); if let Some(ref tool_names) = role_filter { @@ -2155,11 +2143,7 @@ impl RequestContext { .functions .declarations() .iter() - .filter(|v| { - v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - || v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - || v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }) + .filter(|v| is_mcp_meta_function(&v.name)) .map(|v| v.name.to_string()) .collect(); if let Some(agent) = &self.agent { @@ -2168,12 +2152,7 @@ impl RequestContext { .functions() .declarations() .iter() - .filter(|v| { - v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - || v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - || v.name - .starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }) + .filter(|v| is_mcp_meta_function(&v.name)) .map(|v| v.name.to_string()), ); } @@ -2190,39 +2169,15 @@ impl RequestContext { let item_invoke_name = format!("{}_{item}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX); - let item_search_name = - format!("{}_{item}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX); - let item_describe_name = - format!("{}_{item}", MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX); if let Some(values) = app.mapping_mcp_servers.get(item) { server_names.extend( values .split(',') - .flat_map(|v| { - vec![ - format!( - "{}_{}", - MCP_INVOKE_META_FUNCTION_NAME_PREFIX, - v.to_string() - ), - format!( - "{}_{}", - MCP_SEARCH_META_FUNCTION_NAME_PREFIX, - v.to_string() - ), - format!( - "{}_{}", - MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, - v.to_string() - ), - ] - }) + .flat_map(mcp_meta_function_names) .filter(|v| mcp_declaration_names.contains(v)), ) } else if mcp_declaration_names.contains(&item_invoke_name) { - server_names.insert(item_invoke_name); - server_names.insert(item_search_name); - server_names.insert(item_describe_name); + server_names.extend(mcp_meta_function_names(item)); } } } @@ -2251,11 +2206,7 @@ impl RequestContext { .declarations() .to_vec() .into_iter() - .filter(|v| { - v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - || v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - || v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }) + .filter(|v| is_mcp_meta_function(&v.name)) .collect(); if let Some(ref server_names) = role_filter { diff --git a/src/config/tool_scope.rs b/src/config/tool_scope.rs index 9836ca5..a68ebc5 100644 --- a/src/config/tool_scope.rs +++ b/src/config/tool_scope.rs @@ -63,10 +63,10 @@ impl McpRuntime { .get(server) .cloned() .with_context(|| format!("{server} MCP server not found in runtime"))?; - let tools = server_handle.list_tools(None).await?; + let tools = server_handle.list_all_tools().await?; let mut items = HashMap::new(); - for tool in tools.tools { + for tool in tools { let item = CatalogItem { name: tool.name.to_string(), server: server.to_string(), @@ -110,9 +110,8 @@ impl McpRuntime { .with_context(|| format!("{server} MCP server not found in runtime"))?; let tool_schema = server_handle - .list_tools(None) + .list_all_tools() .await? - .tools .into_iter() .find(|item| item.name == tool) .ok_or_else(|| anyhow!("{tool} not found in {server} MCP server catalog"))? diff --git a/src/function/mod.rs b/src/function/mod.rs index d17f042..72f3ba6 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -16,7 +16,7 @@ use crate::config::ensure_parent_exists; use crate::config::paths; use crate::mcp::{ MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX, - MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpServersConfig, + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpServersConfig, is_mcp_meta_function, }; use crate::parsers::{bash, python, typescript}; use anyhow::{Context, Result, anyhow, bail}; @@ -284,14 +284,9 @@ pub async fn eval_tool_calls( } } - let (mcp_calls, sequential_calls): (Vec<_>, Vec<_>) = - to_execute.into_iter().partition(|(_, call)| { - call.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) - || call.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) - || call - .name - .starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) - }); + let (mcp_calls, sequential_calls): (Vec<_>, Vec<_>) = to_execute + .into_iter() + .partition(|(_, call)| is_mcp_meta_function(&call.name)); if !mcp_calls.is_empty() { let ctx_ref: &RequestContext = ctx; @@ -2128,6 +2123,11 @@ mod tests { assert_eq!(MCP_INVOKE_META_FUNCTION_NAME_PREFIX, "mcp_invoke"); assert_eq!(MCP_SEARCH_META_FUNCTION_NAME_PREFIX, "mcp_search"); assert_eq!(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, "mcp_describe"); + assert_eq!(crate::mcp::MCP_READ_META_FUNCTION_NAME_PREFIX, "mcp_read"); + assert_eq!( + crate::mcp::MCP_PROMPT_META_FUNCTION_NAME_PREFIX, + "mcp_prompt" + ); } #[test] diff --git a/src/mcp/mod.rs b/src/mcp/mod.rs index e6b6222..80c9e57 100644 --- a/src/mcp/mod.rs +++ b/src/mcp/mod.rs @@ -34,6 +34,29 @@ use tokio::process::Command; pub const MCP_INVOKE_META_FUNCTION_NAME_PREFIX: &str = "mcp_invoke"; pub const MCP_SEARCH_META_FUNCTION_NAME_PREFIX: &str = "mcp_search"; pub const MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX: &str = "mcp_describe"; +pub const MCP_READ_META_FUNCTION_NAME_PREFIX: &str = "mcp_read"; +pub const MCP_PROMPT_META_FUNCTION_NAME_PREFIX: &str = "mcp_prompt"; + +pub const MCP_META_FUNCTION_PREFIXES: [&str; 5] = [ + MCP_INVOKE_META_FUNCTION_NAME_PREFIX, + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + MCP_READ_META_FUNCTION_NAME_PREFIX, + MCP_PROMPT_META_FUNCTION_NAME_PREFIX, +]; + +pub fn is_mcp_meta_function(name: &str) -> bool { + MCP_META_FUNCTION_PREFIXES + .iter() + .any(|prefix| name.starts_with(prefix)) +} + +pub fn mcp_meta_function_names(server: &str) -> Vec { + MCP_META_FUNCTION_PREFIXES + .iter() + .map(|prefix| format!("{prefix}_{server}")) + .collect() +} pub type ConnectedServer = RunningService; @@ -347,11 +370,11 @@ impl McpRegistry { Err(e) => return Err(e), }; - let tools = service.list_tools(None).await?; + let tools = service.list_all_tools().await?; debug!("Available tools for MCP server {id}: {tools:?}"); let mut items_vec = Vec::new(); - for t in tools.tools { + for t in tools { let name = t.name.to_string(); let description = t.description.unwrap_or_default().to_string(); items_vec.push(CatalogItem { @@ -1185,6 +1208,51 @@ mod tests { assert_eq!(MCP_INVOKE_META_FUNCTION_NAME_PREFIX, "mcp_invoke"); assert_eq!(MCP_SEARCH_META_FUNCTION_NAME_PREFIX, "mcp_search"); assert_eq!(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, "mcp_describe"); + assert_eq!(MCP_READ_META_FUNCTION_NAME_PREFIX, "mcp_read"); + assert_eq!(MCP_PROMPT_META_FUNCTION_NAME_PREFIX, "mcp_prompt"); + } + + #[test] + fn is_mcp_meta_function_classifies_names() { + assert!(is_mcp_meta_function("mcp_invoke_github")); + assert!(is_mcp_meta_function("mcp_search_github")); + assert!(is_mcp_meta_function("mcp_describe_github")); + assert!(is_mcp_meta_function("mcp_read_github")); + assert!(is_mcp_meta_function("mcp_prompt_github")); + assert!(!is_mcp_meta_function("mcp_gateway_tool")); + assert!(!is_mcp_meta_function("fs_read")); + assert!(!is_mcp_meta_function("")); + assert!(!is_mcp_meta_function("mcp_")); + } + + #[test] + fn meta_function_prefixes_are_not_prefixes_of_each_other() { + for (i, a) in MCP_META_FUNCTION_PREFIXES.iter().enumerate() { + for (j, b) in MCP_META_FUNCTION_PREFIXES.iter().enumerate() { + if i != j { + assert!(!b.starts_with(a), "{a} is a prefix of {b}"); + } + } + } + } + + #[test] + fn is_mcp_meta_function_preserves_lax_prefix_matching() { + assert!(is_mcp_meta_function("mcp_invoker_x")); + } + + #[test] + fn mcp_meta_function_names_returns_all_prefixes_in_order() { + assert_eq!( + mcp_meta_function_names("github"), + vec![ + "mcp_invoke_github", + "mcp_search_github", + "mcp_describe_github", + "mcp_read_github", + "mcp_prompt_github", + ] + ); } #[test]