diff --git a/src/config/agent.rs b/src/config/agent.rs index 9a6b83f..40771fa 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -15,6 +15,7 @@ use crate::config::prompts::{ }; use crate::graph::types::RagNode; use crate::graph::{Graph, GraphParser, NodeType}; +use crate::mcp::McpServerFeatures; use crate::rag::RagInitConfig; use crate::vault::SECRET_RE; use anyhow::{Context, Result}; @@ -380,7 +381,7 @@ impl Agent { self.graph_rags.get(node_id).cloned() } - pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec) { + pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec) { self.functions.append_mcp_meta_functions(mcp_servers); } diff --git a/src/config/app_state.rs b/src/config/app_state.rs index c1f4769..7a6cb0e 100644 --- a/src/config/app_state.rs +++ b/src/config/app_state.rs @@ -70,7 +70,7 @@ impl AppState { let mut functions = Functions::init(config.visible_tools.as_ref().unwrap_or(&Vec::new()))?; if !mcp_registry.is_empty() && config.mcp_server_support { - functions.append_mcp_meta_functions(mcp_registry.list_started_servers()); + functions.append_mcp_meta_functions(mcp_registry.server_features()); } let mcp_registry = if mcp_registry.is_empty() { diff --git a/src/config/mod.rs b/src/config/mod.rs index 426ed3f..78cbd73 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -22,6 +22,8 @@ pub(crate) mod todo; mod tool_scope; mod update; +#[cfg(test)] +pub(crate) use self::agent::AgentConfig; pub use self::agent::{ Agent, AgentVariable, AgentVariables, complete_agent_variables, list_agents, list_agents_with_descriptions, diff --git a/src/config/request_context.rs b/src/config/request_context.rs index 14f0a01..965df3f 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -23,7 +23,7 @@ use crate::function::{ user_interaction::USER_FUNCTION_PREFIX, }; use crate::mcp::{ - MCP_INVOKE_META_FUNCTION_NAME_PREFIX, McpAuthReason, McpAuthRequired, is_auth_required_error, + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpAuthReason, McpAuthRequired, is_auth_required_error, is_mcp_meta_function, mcp_meta_function_names, }; use crate::rag::Rag; @@ -2167,8 +2167,8 @@ impl RequestContext { continue; } - let item_invoke_name = - format!("{}_{item}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX); + let item_search_name = + format!("{}_{item}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX); if let Some(values) = app.mapping_mcp_servers.get(item) { server_names.extend( values @@ -2176,7 +2176,7 @@ impl RequestContext { .flat_map(mcp_meta_function_names) .filter(|v| mcp_declaration_names.contains(v)), ) - } else if mcp_declaration_names.contains(&item_invoke_name) { + } else if mcp_declaration_names.contains(&item_search_name) { server_names.extend(mcp_meta_function_names(item)); } } @@ -3779,7 +3779,7 @@ impl RequestContext { functions.append_todo_functions(); } if !mcp_runtime.is_empty() { - functions.append_mcp_meta_functions(mcp_runtime.server_names()); + functions.append_mcp_meta_functions(mcp_runtime.server_features()); } if app.function_calling_support && policy.skills_enabled { functions.append_skill_functions(); @@ -4628,7 +4628,7 @@ mod tests { use crate::config::AppState; use crate::config::agent::AgentConfig; use crate::function::{ToolCall, skill}; - use crate::mcp::{McpServer, McpServersConfig, McpTransportType}; + use crate::mcp::{McpServer, McpServerFeatures, McpServersConfig, McpTransportType}; use crate::utils; use crate::utils::get_env_name; use crate::vault::Vault; @@ -4689,6 +4689,15 @@ mod tests { RequestContext::new(default_app_state(), WorkingMode::Cmd) } + fn tools_only_features(name: &str) -> McpServerFeatures { + McpServerFeatures { + name: name.to_string(), + tools: true, + resources: false, + prompts: false, + } + } + #[test] fn new_creates_clean_state() { let ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); @@ -5696,9 +5705,10 @@ mod tests { #[test] fn select_enabled_mcp_servers_all_returns_all_mcp_functions() { let mut ctx = create_test_ctx(); - ctx.tool_scope - .functions - .append_mcp_meta_functions(vec!["github".into(), "slack".into()]); + ctx.tool_scope.functions.append_mcp_meta_functions(vec![ + tools_only_features("github"), + tools_only_features("slack"), + ]); let mut role = Role::new("r", "p"); role.set_enabled_mcp_servers(Some(vec!["all".to_string()])); @@ -5713,9 +5723,10 @@ mod tests { #[test] fn select_enabled_mcp_servers_comma_filters() { let mut ctx = create_test_ctx(); - ctx.tool_scope - .functions - .append_mcp_meta_functions(vec!["github".into(), "slack".into()]); + ctx.tool_scope.functions.append_mcp_meta_functions(vec![ + tools_only_features("github"), + tools_only_features("slack"), + ]); let mut role = Role::new("r", "p"); role.set_enabled_mcp_servers(Some(vec!["github".to_string()])); @@ -5726,6 +5737,28 @@ mod tests { assert!(!names.contains(&"mcp_invoke_slack")); } + #[test] + fn select_enabled_mcp_servers_keeps_resources_only_server() { + let mut ctx = create_test_ctx(); + ctx.tool_scope + .functions + .append_mcp_meta_functions(vec![McpServerFeatures { + name: "res".to_string(), + tools: false, + resources: true, + prompts: false, + }]); + + let mut role = Role::new("r", "p"); + role.set_enabled_mcp_servers(Some(vec!["res".to_string()])); + + let fns = ctx.select_enabled_mcp_servers(&role); + let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect(); + assert!(names.contains(&"mcp_search_res")); + assert!(names.contains(&"mcp_describe_res")); + assert!(!names.contains(&"mcp_invoke_res")); + } + #[test] fn state_empty_context_has_no_context_flags() { let ctx = create_test_ctx(); diff --git a/src/config/tool_scope.rs b/src/config/tool_scope.rs index eac4e8b..4ca9d29 100644 --- a/src/config/tool_scope.rs +++ b/src/config/tool_scope.rs @@ -1,5 +1,5 @@ use crate::function::{Functions, ToolCallTracker}; -use crate::mcp::{CatalogItem, CatalogItemKind, ConnectedServer, McpRegistry}; +use crate::mcp::{CatalogItem, CatalogItemKind, ConnectedServer, McpRegistry, McpServerFeatures}; use anyhow::{Context, Result, anyhow}; use bm25::{Document, Language, SearchEngineBuilder}; @@ -49,8 +49,20 @@ impl McpRuntime { self.servers.get(name) } - pub fn server_names(&self) -> Vec { - self.servers.keys().cloned().collect() + pub fn server_features(&self) -> Vec { + let mut features: Vec = self + .servers + .iter() + .map(|(name, handle)| { + let info = handle.peer_info(); + McpServerFeatures::from_capabilities( + name.as_str(), + info.as_ref().map(|info| &info.capabilities), + ) + }) + .collect(); + features.sort_by(|a, b| a.name.cmp(&b.name)); + features } pub fn sync_from_registry(&mut self, registry: &McpRegistry) { @@ -65,12 +77,14 @@ impl McpRuntime { .get(server) .cloned() .with_context(|| format!("{server} MCP server not found in runtime"))?; - let capabilities = server_handle - .peer_info() - .map(|info| info.capabilities.clone()); + let info = server_handle.peer_info(); + let features = McpServerFeatures::from_capabilities( + server, + info.as_ref().map(|info| &info.capabilities), + ); let mut items = HashMap::new(); - if capabilities.as_ref().is_none_or(|c| c.tools.is_some()) { + if features.tools { match server_handle.list_all_tools().await { Ok(tools) => merge_catalog_items( &mut items, @@ -82,7 +96,7 @@ impl McpRuntime { } } - if capabilities.as_ref().is_some_and(|c| c.resources.is_some()) { + if features.resources { match server_handle.list_all_resources().await { Ok(resources) => merge_catalog_items( &mut items, @@ -105,7 +119,7 @@ impl McpRuntime { } } - if capabilities.as_ref().is_some_and(|c| c.prompts.is_some()) { + if features.prompts { match server_handle.list_all_prompts().await { Ok(prompts) => merge_catalog_items( &mut items, @@ -349,8 +363,9 @@ pub(crate) mod test_fixtures { use rmcp::{RoleServer, ServerHandler, ServiceExt}; use std::sync::atomic::{AtomicUsize, Ordering}; - #[derive(Clone, Default)] + #[derive(Clone)] pub(crate) struct FixtureServer { + pub(crate) tools_capability: bool, pub(crate) resources_capability: bool, pub(crate) prompts_capability: bool, pub(crate) fail_resource_listings: bool, @@ -358,9 +373,26 @@ pub(crate) mod test_fixtures { pub(crate) list_prompts_calls: Arc, } + impl Default for FixtureServer { + fn default() -> Self { + Self { + tools_capability: true, + resources_capability: false, + prompts_capability: false, + fail_resource_listings: false, + list_resources_calls: Arc::default(), + list_prompts_calls: Arc::default(), + } + } + } + impl ServerHandler for FixtureServer { fn get_info(&self) -> ServerInfo { - let mut capabilities = ServerCapabilities::builder().enable_tools().build(); + let mut capabilities = if self.tools_capability { + ServerCapabilities::builder().enable_tools().build() + } else { + ServerCapabilities::builder().build() + }; capabilities.resources = self.resources_capability.then(ResourcesCapability::default); capabilities.prompts = self.prompts_capability.then(PromptsCapability::default); ServerInfo::new(capabilities) @@ -493,7 +525,7 @@ mod tests { fn mcp_runtime_new_is_empty() { let runtime = McpRuntime::new(); assert!(runtime.is_empty()); - assert!(runtime.server_names().is_empty()); + assert!(runtime.server_features().is_empty()); } #[test] @@ -508,6 +540,61 @@ mod tests { assert!(runtime.get("nonexistent").is_none()); } + #[tokio::test] + async fn server_features_reports_fixture_capabilities() { + let (runtime, _server) = fixture_runtime(FixtureServer { + resources_capability: true, + ..Default::default() + }) + .await; + + let features = runtime.server_features(); + assert_eq!( + features, + vec![McpServerFeatures { + name: "fixture".to_string(), + tools: true, + resources: true, + prompts: false, + }] + ); + + let mut functions = Functions::default(); + functions.append_mcp_meta_functions(features); + assert_eq!(functions.declarations().len(), 3); + assert!(functions.contains("mcp_invoke_fixture")); + assert!(functions.contains("mcp_search_fixture")); + assert!(functions.contains("mcp_describe_fixture")); + } + + #[tokio::test] + async fn server_features_without_tools_capability_gates_invoke() { + let (runtime, _server) = fixture_runtime(FixtureServer { + tools_capability: false, + resources_capability: true, + ..Default::default() + }) + .await; + + let features = runtime.server_features(); + assert_eq!( + features, + vec![McpServerFeatures { + name: "fixture".to_string(), + tools: false, + resources: true, + prompts: false, + }] + ); + + let mut functions = Functions::default(); + functions.append_mcp_meta_functions(features); + assert_eq!(functions.declarations().len(), 2); + assert!(!functions.contains("mcp_invoke_fixture")); + assert!(functions.contains("mcp_search_fixture")); + assert!(functions.contains("mcp_describe_fixture")); + } + #[test] fn tool_scope_default_has_empty_mcp_runtime() { let scope = ToolScope::default(); diff --git a/src/function/mod.rs b/src/function/mod.rs index 7d00567..5ea626b 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -16,7 +16,9 @@ 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, is_mcp_meta_function, + MCP_META_FUNCTION_PREFIXES, MCP_PROMPT_META_FUNCTION_NAME_PREFIX, + MCP_READ_META_FUNCTION_NAME_PREFIX, MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpServerFeatures, + McpServersConfig, is_mcp_meta_function, }; use crate::parsers::{bash, python, typescript}; use anyhow::{Context, Result, anyhow, bail}; @@ -409,6 +411,18 @@ impl ToolResult { } } +fn gated_meta_function_prefixes(features: &McpServerFeatures) -> Vec<&'static str> { + MCP_META_FUNCTION_PREFIXES + .into_iter() + .filter(|&prefix| match prefix { + MCP_INVOKE_META_FUNCTION_NAME_PREFIX => features.tools, + MCP_READ_META_FUNCTION_NAME_PREFIX => features.resources, + MCP_PROMPT_META_FUNCTION_NAME_PREFIX => features.prompts, + _ => true, + }) + .collect() +} + #[derive(Debug, Clone, Default)] pub struct Functions { declarations: Vec, @@ -624,7 +638,7 @@ impl Functions { .retain(|f| !f.name.starts_with(RAG_FUNCTION_PREFIX)); } - pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec) { + pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec) { let mut invoke_function_properties = IndexMap::new(); invoke_function_properties.insert( "tool".to_string(), @@ -682,59 +696,72 @@ impl Functions { }, ); - for server in mcp_servers { + for features in mcp_servers { + let server = &features.name; let search_function_name = format!("{}_{server}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX); let describe_function_name = format!("{}_{server}", MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX); let invoke_function_name = format!("{}_{server}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX); - let invoke_function_declaration = FunctionDeclaration { - name: invoke_function_name.clone(), - description: formatdoc!( - r#" + for prefix in gated_meta_function_prefixes(&features) { + match prefix { + MCP_INVOKE_META_FUNCTION_NAME_PREFIX => { + self.declarations.push(FunctionDeclaration { + name: invoke_function_name.clone(), + description: formatdoc!( + r#" Invoke the specified tool on the {server} MCP server. Always call {describe_function_name} first to find the correct invocation schema for the given tool. "# - ), - parameters: JsonSchema { - type_value: Some("object".to_string()), - properties: Some(invoke_function_properties.clone()), - required: Some(vec!["tool".to_string()]), - ..Default::default() - }, - agent: false, - }; - let search_functions_declaration = FunctionDeclaration { - name: search_function_name.clone(), - description: formatdoc!( - r#" - Find candidate tools by keywords for the {server} MCP server. Returns small suggestions; fetch - schemas with {describe_function_name}. - "# - ), - parameters: JsonSchema { - type_value: Some("object".to_string()), - properties: Some(search_function_properties.clone()), - required: Some(vec!["query".to_string()]), - ..Default::default() - }, - agent: false, - }; - let describe_functions_declaration = FunctionDeclaration { - name: describe_function_name.clone(), - description: "Get the full schema or metadata for exactly one MCP catalog item: \ - a tool, resource, resource template, or prompt." - .to_string(), - parameters: JsonSchema { - type_value: Some("object".to_string()), - properties: Some(describe_function_properties.clone()), - required: Some(vec!["tool".to_string()]), - ..Default::default() - }, - agent: false, - }; - self.declarations.push(invoke_function_declaration); - self.declarations.push(search_functions_declaration); - self.declarations.push(describe_functions_declaration); + ), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(invoke_function_properties.clone()), + required: Some(vec!["tool".to_string()]), + ..Default::default() + }, + agent: false, + }); + } + MCP_SEARCH_META_FUNCTION_NAME_PREFIX => { + self.declarations.push(FunctionDeclaration { + name: search_function_name.clone(), + description: formatdoc!( + r#" + Find candidate tools by keywords for the {server} MCP server. Returns small suggestions; fetch + schemas with {describe_function_name}. + "# + ), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(search_function_properties.clone()), + required: Some(vec!["query".to_string()]), + ..Default::default() + }, + agent: false, + }); + } + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX => { + self.declarations.push(FunctionDeclaration { + name: describe_function_name.clone(), + description: "Get the full schema or metadata for exactly one MCP \ + catalog item: a tool, resource, resource template, or \ + prompt." + .to_string(), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(describe_function_properties.clone()), + required: Some(vec!["tool".to_string()]), + ..Default::default() + }, + agent: false, + }); + } + // The declaration is added alongside its handler. + MCP_READ_META_FUNCTION_NAME_PREFIX => {} + MCP_PROMPT_META_FUNCTION_NAME_PREFIX => {} + _ => debug_assert!(false, "unhandled MCP meta-function prefix: {prefix}"), + } + } } } @@ -1859,6 +1886,19 @@ mod tests { ToolCall::new(name.to_string(), args, Some("id1".to_string())) } + fn mcp_features(name: &str, tools: bool, resources: bool, prompts: bool) -> McpServerFeatures { + McpServerFeatures { + name: name.to_string(), + tools, + resources, + prompts, + } + } + + fn tools_only(name: &str) -> McpServerFeatures { + mcp_features(name, true, false, false) + } + fn run_async(f: F) -> F::Output { tokio::runtime::Builder::new_current_thread() .enable_all() @@ -2231,7 +2271,7 @@ mod tests { #[test] fn functions_append_mcp_meta_creates_three_per_server() { let mut f = Functions::default(); - f.append_mcp_meta_functions(vec!["github".to_string()]); + f.append_mcp_meta_functions(vec![tools_only("github")]); assert_eq!(f.declarations().len(), 3); assert!(f.contains("mcp_invoke_github")); assert!(f.contains("mcp_search_github")); @@ -2241,7 +2281,7 @@ mod tests { #[test] fn functions_append_mcp_meta_multiple_servers() { let mut f = Functions::default(); - f.append_mcp_meta_functions(vec!["github".into(), "slack".into()]); + f.append_mcp_meta_functions(vec![tools_only("github"), tools_only("slack")]); assert_eq!(f.declarations().len(), 6); assert!(f.contains("mcp_invoke_github")); assert!(f.contains("mcp_invoke_slack")); @@ -2254,6 +2294,132 @@ mod tests { assert!(f.is_empty()); } + #[test] + fn functions_append_mcp_meta_resources_only_omits_invoke() { + let mut f = Functions::default(); + f.append_mcp_meta_functions(vec![mcp_features("res", false, true, false)]); + assert_eq!(f.declarations().len(), 2); + assert!(!f.contains("mcp_invoke_res")); + assert!(f.contains("mcp_search_res")); + assert!(f.contains("mcp_describe_res")); + } + + #[test] + fn functions_append_mcp_meta_all_capabilities_emits_three() { + let mut f = Functions::default(); + f.append_mcp_meta_functions(vec![mcp_features("srv", true, true, true)]); + assert_eq!(f.declarations().len(), 3); + assert!(f.contains("mcp_invoke_srv")); + assert!(f.contains("mcp_search_srv")); + assert!(f.contains("mcp_describe_srv")); + } + + #[test] + fn features_from_missing_capabilities_fail_open_for_tools() { + let features = McpServerFeatures::from_capabilities("srv", None); + assert!(features.tools); + assert!(!features.resources); + assert!(!features.prompts); + + let mut f = Functions::default(); + f.append_mcp_meta_functions(vec![features]); + assert!(f.contains("mcp_invoke_srv")); + } + + #[test] + fn gated_prefixes_tools_only() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", true, false, false)), + vec![ + MCP_INVOKE_META_FUNCTION_NAME_PREFIX, + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + ] + ); + } + + #[test] + fn gated_prefixes_tools_and_resources_include_read() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", true, true, false)), + vec![ + MCP_INVOKE_META_FUNCTION_NAME_PREFIX, + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + MCP_READ_META_FUNCTION_NAME_PREFIX, + ] + ); + } + + #[test] + fn gated_prefixes_tools_and_prompts_include_prompt() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", true, false, true)), + vec![ + MCP_INVOKE_META_FUNCTION_NAME_PREFIX, + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + MCP_PROMPT_META_FUNCTION_NAME_PREFIX, + ] + ); + } + + #[test] + fn gated_prefixes_all_capabilities_include_all() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", true, true, true)), + MCP_META_FUNCTION_PREFIXES.to_vec() + ); + } + + #[test] + fn gated_prefixes_resources_only_omit_invoke() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", false, true, false)), + vec![ + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + MCP_READ_META_FUNCTION_NAME_PREFIX, + ] + ); + } + + #[test] + fn gated_prefixes_prompts_only_omit_invoke_and_read() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", false, false, true)), + vec![ + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + MCP_PROMPT_META_FUNCTION_NAME_PREFIX, + ] + ); + } + + #[test] + fn gated_prefixes_resources_and_prompts_omit_invoke() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", false, true, true)), + vec![ + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + MCP_READ_META_FUNCTION_NAME_PREFIX, + MCP_PROMPT_META_FUNCTION_NAME_PREFIX, + ] + ); + } + + #[test] + fn gated_prefixes_no_capabilities_keep_search_and_describe() { + assert_eq!( + gated_meta_function_prefixes(&mcp_features("srv", false, false, false)), + vec![ + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + ] + ); + } + #[test] fn functions_find_returns_declaration() { let mut f = Functions::default(); @@ -2285,7 +2451,7 @@ mod tests { #[test] fn functions_mcp_invoke_declaration_has_tool_and_arguments_params() { let mut f = Functions::default(); - f.append_mcp_meta_functions(vec!["srv".to_string()]); + f.append_mcp_meta_functions(vec![tools_only("srv")]); let decl = f.find("mcp_invoke_srv").unwrap(); let props = decl.parameters.properties.as_ref().unwrap(); assert!(props.contains_key("tool")); @@ -2297,7 +2463,7 @@ mod tests { #[test] fn functions_mcp_search_declaration_has_query_and_top_k_params() { let mut f = Functions::default(); - f.append_mcp_meta_functions(vec!["srv".to_string()]); + f.append_mcp_meta_functions(vec![tools_only("srv")]); let decl = f.find("mcp_search_srv").unwrap(); let props = decl.parameters.properties.as_ref().unwrap(); assert!(props.contains_key("query")); @@ -2307,7 +2473,7 @@ mod tests { #[test] fn functions_mcp_describe_declaration_has_tool_param() { let mut f = Functions::default(); - f.append_mcp_meta_functions(vec!["srv".to_string()]); + f.append_mcp_meta_functions(vec![tools_only("srv")]); let decl = f.find("mcp_describe_srv").unwrap(); let props = decl.parameters.properties.as_ref().unwrap(); assert!(props.contains_key("tool")); @@ -2316,7 +2482,7 @@ mod tests { #[test] fn functions_mcp_describe_declaration_has_optional_kind_param() { let mut f = Functions::default(); - f.append_mcp_meta_functions(vec!["srv".to_string()]); + f.append_mcp_meta_functions(vec![tools_only("srv")]); let decl = f.find("mcp_describe_srv").unwrap(); let props = decl.parameters.properties.as_ref().unwrap(); let kind = props.get("kind").unwrap(); diff --git a/src/function/supervisor.rs b/src/function/supervisor.rs index 05cdf65..ab1e7c1 100644 --- a/src/function/supervisor.rs +++ b/src/function/supervisor.rs @@ -630,14 +630,14 @@ async fn populate_agent_mcp_runtime(ctx: &mut RequestContext, server_ids: &[Stri } fn sync_agent_functions_to_ctx(ctx: &mut RequestContext) -> Result<()> { - let server_names = ctx.tool_scope.mcp_runtime.server_names(); + let server_features = ctx.tool_scope.mcp_runtime.server_features(); let functions = { let agent = ctx .agent .as_mut() .with_context(|| "Agent should be initialized")?; - if !server_names.is_empty() { - agent.append_mcp_meta_functions(server_names); + if !server_features.is_empty() { + agent.append_mcp_meta_functions(server_features); } agent.functions().clone() }; @@ -1453,7 +1453,8 @@ async fn summarize_output(ctx: &RequestContext, agent_name: &str, output: &str) #[cfg(test)] mod tests { use super::*; - use crate::config::{AppState, WorkingMode}; + use crate::config::test_fixtures::{FixtureServer, fixture_runtime}; + use crate::config::{AgentConfig, AppState, WorkingMode}; use crate::supervisor::escalation::{EscalationQueue, EscalationRequest}; use serde_json::json; use serial_test::serial; @@ -1510,6 +1511,27 @@ mod tests { .block_on(f) } + #[tokio::test] + async fn sync_agent_functions_gates_meta_functions_on_live_capabilities() { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + ctx.agent = Some(Agent::test_new(AgentConfig::default())); + let (runtime, _server) = fixture_runtime(FixtureServer { + tools_capability: false, + resources_capability: true, + ..FixtureServer::default() + }) + .await; + ctx.tool_scope.mcp_runtime = runtime; + + sync_agent_functions_to_ctx(&mut ctx).unwrap(); + + let functions = &ctx.tool_scope.functions; + assert_eq!(functions.declarations().len(), 2); + assert!(functions.contains("mcp_search_fixture")); + assert!(functions.contains("mcp_describe_fixture")); + assert!(!functions.contains("mcp_invoke_fixture")); + } + #[test] fn handle_list_running_empty_supervisor() { let mut ctx = ctx_with_supervisor(4, 3); diff --git a/src/mcp/mod.rs b/src/mcp/mod.rs index 64f6e88..d426ff4 100644 --- a/src/mcp/mod.rs +++ b/src/mcp/mod.rs @@ -16,6 +16,7 @@ use futures_util::{StreamExt, TryStreamExt, stream}; use http::{HeaderName, HeaderValue}; use indexmap::IndexMap; use indoc::formatdoc; +use rmcp::model::ServerCapabilities; use rmcp::service::RunningService; use rmcp::transport::StreamableHttpClientTransport; use rmcp::transport::TokioChildProcess; @@ -61,6 +62,28 @@ pub fn mcp_meta_function_names(server: &str) -> Vec { pub type ConnectedServer = RunningService; +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct McpServerFeatures { + pub name: String, + pub tools: bool, + pub resources: bool, + pub prompts: bool, +} + +impl McpServerFeatures { + pub fn from_capabilities( + name: impl Into, + capabilities: Option<&ServerCapabilities>, + ) -> Self { + Self { + name: name.into(), + tools: capabilities.is_none_or(|c| c.tools.is_some()), + resources: capabilities.is_some_and(|c| c.resources.is_some()), + prompts: capabilities.is_some_and(|c| c.prompts.is_some()), + } + } +} + #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize)] #[serde(rename_all = "snake_case")] pub enum CatalogItemKind { @@ -417,8 +440,20 @@ impl McpRegistry { &self.servers } - pub fn list_started_servers(&self) -> Vec { - self.servers.keys().cloned().collect() + pub fn server_features(&self) -> Vec { + let mut features: Vec = self + .servers + .iter() + .map(|(name, handle)| { + let info = handle.peer_info(); + McpServerFeatures::from_capabilities( + name.as_str(), + info.as_ref().map(|info| &info.capabilities), + ) + }) + .collect(); + features.sort_by(|a, b| a.name.cmp(&b.name)); + features } pub fn is_empty(&self) -> bool { @@ -1180,7 +1215,7 @@ mod tests { let registry = McpRegistry::default(); assert!(registry.is_empty()); - assert!(registry.list_started_servers().is_empty()); + assert!(registry.server_features().is_empty()); assert!(registry.mcp_config().is_none()); assert!(registry.log_path().is_none()); }