From 92d9464b387d595bc4edc00b7e98b9ad3d08994d Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 27 Aug 2026 14:43:18 -0600 Subject: [PATCH] feat(mcp): enforce per-server tool allowlists across runtime, jobs, agents, and graph nodes --- src/config/mcp_tool_policy.rs | 125 ++++++++++++--- src/config/mod.rs | 3 + src/config/request_context.rs | 278 +++++++++++++++++++++++++++++++++- src/config/skill_registry.rs | 4 + src/config/tool_scope.rs | 192 ++++++++++++++++++++++- src/function/agents.rs | 58 +++++++ src/function/jobs.rs | 44 ++++++ src/graph/executor.rs | 2 +- src/graph/llm.rs | 13 +- src/graph/map.rs | 8 +- src/graph/validator.rs | 10 +- 11 files changed, 695 insertions(+), 42 deletions(-) diff --git a/src/config/mcp_tool_policy.rs b/src/config/mcp_tool_policy.rs index 48581b1..db48c5e 100644 --- a/src/config/mcp_tool_policy.rs +++ b/src/config/mcp_tool_policy.rs @@ -8,7 +8,6 @@ use std::fmt; /// The configuration level that contributed a layer of tool patterns for an /// MCP server, as rendered in diagnostics. -#[allow(dead_code)] #[derive(Debug, Clone, PartialEq, Eq)] pub enum LayerSource { Global, @@ -34,21 +33,18 @@ impl fmt::Display for LayerSource { } } -#[allow(dead_code)] -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct CompiledPatterns { source: LayerSource, raw: Vec, regexes: Vec, } -#[allow(dead_code)] -#[derive(Debug, Default)] +#[derive(Debug, Clone, Default)] pub struct ToolFilter { layers: Vec, } -#[allow(dead_code)] impl ToolFilter { pub fn push_layer(&mut self, source: LayerSource, patterns: &[String]) { self.layers.push(CompiledPatterns { @@ -70,6 +66,7 @@ impl ToolFilter { /// The first matching raw pattern per layer, in layer order, or the /// source of the first layer with no match. + #[allow(dead_code)] pub fn allows_explain(&self, tool: &str) -> Result, &LayerSource> { let mut matched = Vec::with_capacity(self.layers.len()); for layer in &self.layers { @@ -86,11 +83,48 @@ impl ToolFilter { } Ok(matched) } + + /// Context-layer patterns that match none of the advertised tools + /// surviving the global layer — dead weight, usually a typo. + pub fn dead_context_patterns<'a>( + &'a self, + advertised: &[String], + ) -> Vec<(&'a LayerSource, &'a str)> { + let surviving: Vec<&String> = advertised + .iter() + .filter(|name| { + self.layers + .iter() + .filter(|layer| layer.source == LayerSource::Global) + .all(|layer| { + layer + .regexes + .iter() + .any(|regex| regex.is_match(name).unwrap_or(false)) + }) + }) + .collect(); + let mut dead = Vec::new(); + for layer in self + .layers + .iter() + .filter(|l| l.source != LayerSource::Global) + { + for (raw, regex) in layer.raw.iter().zip(&layer.regexes) { + if !surviving + .iter() + .any(|name| regex.is_match(name).unwrap_or(false)) + { + dead.push((&layer.source, raw.as_str())); + } + } + } + dead + } } /// Translates a glob pattern (`*` = any run of characters, `?` = exactly one) /// into an anchored regex. Patterns that fail to compile match nothing. -#[allow(dead_code)] fn compile_glob(pattern: &str) -> Regex { let translated = format!( "^{}$", @@ -104,22 +138,18 @@ fn compile_glob(pattern: &str) -> Regex { }) } -#[allow(dead_code)] fn never_matching_regex() -> Regex { Regex::new("(?!)").expect("'(?!)' is a valid never-matching regex") } -#[allow(dead_code)] pub struct SkillMcpLayer { pub name: String, pub enabled_servers: Vec, pub mcp_tools: IndexMap>, } -#[allow(dead_code)] pub struct McpToolPolicy; -#[allow(dead_code)] impl McpToolPolicy { #[allow(clippy::too_many_arguments)] pub fn effective( @@ -208,7 +238,6 @@ impl McpToolPolicy { } } -#[allow(dead_code)] fn push_level( filters: &mut HashMap, mcp_config: &McpServersConfig, @@ -234,7 +263,6 @@ fn push_level( /// a key that is an alias expands to every configured id in its /// comma-separated value; anything else is dropped. Keys expanding to the /// same server merge their pattern lists. -#[allow(dead_code)] fn expand_server_keys( mcp_config: &McpServersConfig, aliases: &IndexMap, @@ -248,11 +276,11 @@ fn expand_server_keys( .entry(key.to_string()) .or_default() .extend(patterns.iter().cloned()); - } else if let Some(mapped) = aliases.get(key) { - for mapped_id in mapped.split(',').map(str::trim) { - if mcp_config.mcp_servers.contains_key(mapped_id) { + } else { + for mapped_id in expand_mcp_server_alias(aliases, key) { + if mcp_config.mcp_servers.contains_key(&mapped_id) { expanded - .entry(mapped_id.to_string()) + .entry(mapped_id) .or_default() .extend(patterns.iter().cloned()); } @@ -262,6 +290,25 @@ fn expand_server_keys( expanded } +/// Expands a `mapping_mcp_servers` alias key into its comma-separated server +/// ids. A key with no alias entry expands to nothing. +pub(crate) fn expand_mcp_server_alias( + aliases: &IndexMap, + key: &str, +) -> Vec { + aliases + .get(key) + .map(|mapped| { + mapped + .split(',') + .map(str::trim) + .filter(|id| !id.is_empty()) + .map(str::to_string) + .collect() + }) + .unwrap_or_default() +} + #[cfg(test)] mod tests { use super::*; @@ -756,4 +803,48 @@ mod tests { ); assert_eq!(LayerSource::Node("n1".into()).to_string(), "node (n1)"); } + + #[test] + fn dead_context_patterns_flags_patterns_matching_nothing() { + let filter = layered(&[ + (LayerSource::Global, &["get_*"]), + (LayerSource::Role("dev".into()), &["get_issue", "set_*"]), + ]); + + let advertised = vec!["get_issue".to_string(), "set_topic".to_string()]; + let dead = filter.dead_context_patterns(&advertised); + + // set_* only matches set_topic, which the global layer hides. + assert_eq!(dead, vec![(&LayerSource::Role("dev".into()), "set_*")]); + } + + #[test] + fn dead_context_patterns_is_empty_when_every_pattern_is_live() { + let filter = layered(&[ + (LayerSource::Global, &["get_*"]), + (LayerSource::Session, &["get_issue"]), + ]); + + let advertised = vec!["get_issue".to_string()]; + assert!(filter.dead_context_patterns(&advertised).is_empty()); + } + + #[test] + fn dead_context_patterns_ignores_the_global_layer_itself() { + let filter = layered(&[(LayerSource::Global, &["zzz_*"])]); + + let advertised = vec!["get_issue".to_string()]; + assert!(filter.dead_context_patterns(&advertised).is_empty()); + } + + #[test] + fn expand_mcp_server_alias_splits_and_trims() { + let aliases = aliases(&[("gh", "github, gitlab,")]); + + assert_eq!( + expand_mcp_server_alias(&aliases, "gh"), + vec!["github".to_string(), "gitlab".to_string()] + ); + assert!(expand_mcp_server_alias(&aliases, "nope").is_empty()); + } } diff --git a/src/config/mod.rs b/src/config/mod.rs index d04db86..4d1a4f5 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -42,6 +42,9 @@ pub use self::install_remote::{ pub use self::macro_policy::{ MacroAllowlistLevel, MacroPolicy, MacroSource, MacroState, RESERVED_MACRO_NAMES, ResolvedMacro, }; +pub(crate) use self::mcp_tool_policy::expand_mcp_server_alias; +#[cfg(test)] +pub(crate) use self::mcp_tool_policy::{LayerSource, ToolFilter}; #[allow(unused_imports)] pub use self::request_context::{ RenderMode, RequestContext, effective_max_concurrent_jobs, jobs_enabled, diff --git a/src/config/request_context.rs b/src/config/request_context.rs index 94092c8..8018bab 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -1,4 +1,5 @@ use super::bundles::BundleStore; +use super::mcp_tool_policy::{McpToolPolicy, SkillMcpLayer, ToolFilter, expand_mcp_server_alias}; use super::rag_cache::{RagCache, RagKey}; use super::session::{INTERRUPTED_RESPONSE_TEXT, Session}; use super::skill::{SKILL_SCAFFOLD, Skill}; @@ -51,7 +52,6 @@ use crate::graph; use anyhow::{Context, Error, Result, bail}; use colored::Colorize; use gman::providers::SupportedProvider; -#[cfg(test)] use indexmap::IndexMap; use indoc::formatdoc; use inquire::{Confirm, MultiSelect, Text, list_option::ListOption, validator::Validation}; @@ -95,10 +95,10 @@ pub(crate) fn expand_enabled_mcp_server_ids( for item in enabled_mcp_servers.iter().map(|s| s.trim()) { if mcp_config.mcp_servers.contains_key(item) { ids.push(item.to_string()); - } else if let Some(mapped) = app.mapping_mcp_servers.get(item) { - for mapped_id in mapped.split(',').map(|s| s.trim()) { - if mcp_config.mcp_servers.contains_key(mapped_id) { - ids.push(mapped_id.to_string()); + } else { + for mapped_id in expand_mcp_server_alias(&app.mapping_mcp_servers, item) { + if mcp_config.mcp_servers.contains_key(&mapped_id) { + ids.push(mapped_id); } } } @@ -333,6 +333,10 @@ pub struct RequestContext { /// context owns every job in its supervisor. pub node_job_scope: Option>, + /// Set while a graph LLM node with `mcp_tools` is executing; re-applied as + /// the last filter layer by every `refresh_mcp_tool_filters` recompute. + pub active_node_mcp_tools: Option<(String, IndexMap>)>, + pub supervisor: Option>>, pub parent_supervisor: Option>>, pub self_agent_id: Option, @@ -369,6 +373,7 @@ impl RequestContext { tool_scope: ToolScope::default(), declared_function_names: Default::default(), node_job_scope: None, + active_node_mcp_tools: None, supervisor: None, parent_supervisor: None, self_agent_id: None, @@ -410,7 +415,7 @@ impl RequestContext { mcp_runtime.sync_from_registry(registry); } - Ok(Self { + let mut ctx = Self { app, macro_flag: false, macro_non_isolated: false, @@ -431,6 +436,7 @@ impl RequestContext { }, declared_function_names: Default::default(), node_job_scope: None, + active_node_mcp_tools: None, supervisor: None, parent_supervisor: None, self_agent_id: None, @@ -445,7 +451,9 @@ impl RequestContext { last_continuation_response: None, pending_prefill: None, render_mode: RenderMode::default(), - }) + }; + ctx.refresh_mcp_tool_filters(); + Ok(ctx) } /// Forks the context for one parallel branch of a graph super-step. @@ -480,6 +488,7 @@ impl RequestContext { tool_scope: self.tool_scope.clone(), declared_function_names: self.declared_function_names.clone(), node_job_scope: None, + active_node_mcp_tools: self.active_node_mcp_tools.clone(), supervisor: self.supervisor.clone(), parent_supervisor: self.parent_supervisor.clone(), self_agent_id: self.self_agent_id.clone(), @@ -527,6 +536,7 @@ impl RequestContext { }, declared_function_names: Default::default(), node_job_scope: None, + active_node_mcp_tools: None, supervisor: None, parent_supervisor: parent.supervisor.clone(), self_agent_id: Some(self_agent_id), @@ -2440,7 +2450,9 @@ impl RequestContext { pub fn use_temp_role(&mut self, _app: &AppConfig, prompt: &str) -> Result<()> { let mut role = Role::new(TEMP_ROLE_NAME, prompt); role.set_model(self.current_model().clone()); - self.use_role_obj(role) + self.use_role_obj(role)?; + self.refresh_mcp_tool_filters(); + Ok(()) } pub fn edit_config(&self) -> Result<()> { @@ -3964,9 +3976,70 @@ impl RequestContext { mcp_runtime, tool_tracker, }; + self.refresh_mcp_tool_filters(); Ok(()) } + /// In-place full recompute of `tool_scope.mcp_runtime.tool_filters` from + /// the current declarative state: mcp.json `allowedTools`, app config, + /// active role, agent, session, one layer per loaded skill, and the + /// active graph-node layer (always applied last). Never incremental, so + /// detach and unload paths need no layer-removal logic. + pub fn refresh_mcp_tool_filters(&mut self) { + self.tool_scope.mcp_runtime.tool_filters = self.compute_mcp_tool_filters(); + } + + fn compute_mcp_tool_filters(&self) -> HashMap { + let Some(mcp_config) = self.app.mcp_config.as_ref() else { + return HashMap::new(); + }; + let app = &self.app.config; + let session_map = self.session.as_ref().and_then(|s| s.mcp_tools()); + let agent_map = self.agent.as_ref().and_then(|a| a.mcp_tools()); + let agent = self + .agent + .as_ref() + .zip(agent_map.as_ref()) + .map(|(a, map)| (a.name(), map)); + let role_map = self.role.as_ref().and_then(|r| r.mcp_tools()); + let role = self + .role + .as_ref() + .zip(role_map.as_ref()) + .map(|(r, map)| (r.name(), map)); + let skills: Vec = self + .skill_registry + .loaded_skills() + .filter_map(|skill| { + let mcp_tools = skill.mcp_tools()?.clone(); + Some(SkillMcpLayer { + name: skill.name().to_string(), + enabled_servers: expand_enabled_mcp_server_ids( + app, + mcp_config, + skill.enabled_mcp_servers().unwrap_or_default(), + ), + mcp_tools, + }) + }) + .collect(); + let node = self + .active_node_mcp_tools + .as_ref() + .map(|(id, map)| (id.as_str(), map)); + + McpToolPolicy::effective( + mcp_config, + session_map.as_ref(), + agent, + role, + app.mcp_tools.as_ref(), + &skills, + node, + &app.mapping_mcp_servers, + ) + } + pub async fn refresh_tool_scope(&mut self, abort_signal: AbortSignal) -> Result<()> { let app = (*self.app.config).clone(); let base_mcps = if app.mcp_server_support { @@ -4141,6 +4214,7 @@ impl RequestContext { } } self.session = session; + self.refresh_mcp_tool_filters(); self.init_agent_session_variables(new_session)?; Ok(()) } @@ -4261,6 +4335,7 @@ impl RequestContext { .is_some() .then(|| RagKey::Agent(agent.name().to_string())); self.agent = Some(agent); + self.refresh_mcp_tool_filters(); if let Some(old) = self.supervisor.as_ref() { old.read().cancel_recursive(); } @@ -8012,4 +8087,191 @@ mod tests { ); } } + + fn mcp_app_state(servers: &[&str]) -> Arc { + let mcp_servers = servers + .iter() + .map(|name| { + ( + name.to_string(), + McpServer { + transport_type: McpTransportType::Stdio, + command: Some("echo".to_string()), + args: None, + env: None, + cwd: None, + url: None, + headers: None, + oauth: None, + allowed_tools: None, + }, + ) + }) + .collect(); + Arc::new(AppState { + config: Arc::new(AppConfig::default()), + vault: Arc::new(Vault::default()), + mcp_factory: Arc::new(McpFactory::default()), + rag_cache: Arc::new(RagCache::default()), + mcp_config: Some(McpServersConfig { mcp_servers }), + mcp_log_path: None, + mcp_registry: None, + functions: Functions::default(), + }) + } + + fn gh_get_only_map() -> IndexMap> { + IndexMap::from([("gh".to_string(), vec!["get_*".to_string()])]) + } + + #[test] + #[serial] + fn rebuild_tool_scope_populates_filters_from_a_filtered_role() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = RequestContext::new(mcp_app_state(&["gh"]), WorkingMode::Cmd); + let mut role = Role::new("dev", "prompt"); + role.set_mcp_tools(Some(gh_get_only_map())); + ctx.role = Some(role); + + run_async(ctx.refresh_tool_scope(utils::create_abort_signal())).unwrap(); + + let filter = ctx + .tool_scope + .mcp_runtime + .tool_filters + .get("gh") + .expect("a rebuild must never leave a filtered role unfiltered"); + assert!(filter.allows("get_issue")); + assert!(!filter.allows("delete_repo")); + } + + #[test] + #[serial] + fn mid_node_skill_load_keeps_node_filter_layer() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = RequestContext::new(mcp_app_state(&["gh"]), WorkingMode::Cmd); + ctx.update_app_config(|app| app.mcp_server_support = false); + ctx.active_node_mcp_tools = Some(("n1".to_string(), gh_get_only_map())); + ctx.refresh_mcp_tool_filters(); + assert!(!ctx.tool_scope.mcp_runtime.tool_filters["gh"].allows("list_prs")); + + ctx.skill_registry + .insert(Skill::new( + "sec", + "---\nenabled_mcp_servers: gh\nmcp_tools:\n gh: [get_*, list_*]\n---\nBody", + )) + .unwrap(); + run_async(ctx.refresh_tool_scope(utils::create_abort_signal())).unwrap(); + + let filter = &ctx.tool_scope.mcp_runtime.tool_filters["gh"]; + assert!(filter.allows("get_issue")); + assert!( + !filter.allows("list_prs"), + "the node layer must survive a mid-node skill load" + ); + + ctx.active_node_mcp_tools = None; + ctx.refresh_mcp_tool_filters(); + let filter = &ctx.tool_scope.mcp_runtime.tool_filters["gh"]; + assert!( + filter.allows("list_prs"), + "the node layer must not outlive the node" + ); + assert!(!filter.allows("delete_repo")); + } + + #[test] + #[serial] + fn set_skills_enabled_does_not_drop_filters() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = RequestContext::new(mcp_app_state(&["gh"]), WorkingMode::Cmd); + let mut role = Role::new("dev", "prompt"); + role.set_mcp_tools(Some(gh_get_only_map())); + ctx.role = Some(role); + ctx.refresh_mcp_tool_filters(); + assert!(ctx.tool_scope.mcp_runtime.tool_filters.contains_key("gh")); + + run_async(ctx.update("skills_enabled false", utils::create_abort_signal())).unwrap(); + + assert!(ctx.tool_scope.mcp_runtime.tool_filters.contains_key("gh")); + } + + #[test] + #[serial] + fn use_session_applies_persisted_mcp_tools_immediately() { + // Mirrors use_session's ordering hazard: the tool-scope rebuild runs + // before `self.session` is assigned, so a re-attached session's + // persisted map is only enforced by the post-assignment refresh. + // (Session::load_from_ctx needs live model resolution, so the + // persisted session is round-tripped through serde directly.) + let _guard = TestConfigDirGuard::new(); + let mut ctx = RequestContext::new(mcp_app_state(&["gh"]), WorkingMode::Cmd); + let mut persisted = Session::default(); + persisted.set_mcp_tools(Some(gh_get_only_map())); + let yaml = serde_yaml::to_string(&persisted).unwrap(); + let reloaded: Session = serde_yaml::from_str(&yaml).unwrap(); + + run_async(ctx.refresh_tool_scope(utils::create_abort_signal())).unwrap(); + assert!( + !ctx.tool_scope.mcp_runtime.tool_filters.contains_key("gh"), + "the pre-assignment rebuild cannot see the session layer" + ); + ctx.session = Some(reloaded); + ctx.refresh_mcp_tool_filters(); + + let filter = ctx + .tool_scope + .mcp_runtime + .tool_filters + .get("gh") + .expect("a re-attached session's persisted map must apply immediately"); + assert!(filter.allows("get_issue")); + assert!(!filter.allows("delete_repo")); + } + + #[test] + #[serial] + fn use_agent_applies_agent_filters_without_a_session() { + let _guard = TestConfigDirGuard::new(); + let config_path = paths::agent_config_file("filterer"); + ensure_parent_exists(&config_path).unwrap(); + write( + &config_path, + "name: filterer\ninstructions: hi\nmcp_tools:\n gh:\n - get_*\n", + ) + .unwrap(); + let mut ctx = RequestContext::new(mcp_app_state(&["gh"]), WorkingMode::Cmd); + let app = ctx.app.config.clone(); + + run_async(ctx.use_agent(&app, "filterer", None, utils::create_abort_signal())).unwrap(); + + let filter = ctx + .tool_scope + .mcp_runtime + .tool_filters + .get("gh") + .expect("the agent layer must apply on the no-session use_agent path"); + assert!(filter.allows("get_issue")); + assert!(!filter.allows("delete_repo")); + } + + #[test] + #[serial] + fn use_temp_role_recomputes_filters() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = RequestContext::new(mcp_app_state(&["gh"]), WorkingMode::Cmd); + let mut role = Role::new("dev", "prompt"); + role.set_mcp_tools(Some(gh_get_only_map())); + ctx.role = Some(role); + ctx.refresh_mcp_tool_filters(); + assert!(ctx.tool_scope.mcp_runtime.tool_filters.contains_key("gh")); + + let app = ctx.app.config.clone(); + ctx.use_temp_role(&app, "temp prompt").unwrap(); + + assert!( + ctx.tool_scope.mcp_runtime.tool_filters.is_empty(), + "a temp role must clear the replaced role's filter layer" + ); + } } diff --git a/src/config/skill_registry.rs b/src/config/skill_registry.rs index 63706c3..5d5d0dc 100644 --- a/src/config/skill_registry.rs +++ b/src/config/skill_registry.rs @@ -34,6 +34,10 @@ impl SkillRegistry { self.loaded.keys().cloned().collect() } + pub fn loaded_skills(&self) -> impl Iterator { + self.loaded.values() + } + pub fn loaded_mcp_servers(&self) -> BTreeSet { let mut out = BTreeSet::new(); for skill in self.loaded.values() { diff --git a/src/config/tool_scope.rs b/src/config/tool_scope.rs index 9939bd7..a2488e7 100644 --- a/src/config/tool_scope.rs +++ b/src/config/tool_scope.rs @@ -1,3 +1,4 @@ +use super::mcp_tool_policy::ToolFilter; use crate::function::{Functions, ToolCallTracker}; use crate::mcp::{CatalogItem, CatalogItemKind, ConnectedServer, McpRegistry, McpServerFeatures}; @@ -44,6 +45,8 @@ pub enum McpPromptCompletion { #[derive(Default, Clone)] pub struct McpRuntime { pub servers: HashMap>, + /// Per-server effective tool allowlists; a server absent here is unfiltered. + pub tool_filters: HashMap, } impl McpRuntime { @@ -100,12 +103,24 @@ impl McpRuntime { if features.tools { match server_handle.list_all_tools().await { - Ok(tools) => merge_catalog_items( - &mut items, - tools - .into_iter() - .map(|tool| tool_catalog_item(server, tool)), - ), + Ok(mut tools) => { + if let Some(filter) = self.tool_filters.get(server) { + let advertised: Vec = + tools.iter().map(|tool| tool.name.to_string()).collect(); + for (source, pattern) in filter.dead_context_patterns(&advertised) { + warn!( + "MCP tool pattern '{pattern}' from {source} matches no allowed tools on server '{server}'" + ); + } + tools.retain(|tool| filter.allows(&tool.name)); + } + merge_catalog_items( + &mut items, + tools + .into_iter() + .map(|tool| tool_catalog_item(server, tool)), + ) + } Err(e) => warn!("Failed to list tools on MCP server {server}: {e}"), } } @@ -195,6 +210,11 @@ impl McpRuntime { match kind { "tool" => { + if let Some(filter) = self.tool_filters.get(server) + && !filter.allows(tool) + { + return Err(anyhow!("{tool} not found in {server} MCP server catalog")); + } let tool_schema = server_handle .list_all_tools() .await? @@ -297,6 +317,12 @@ impl McpRuntime { .cloned() .with_context(|| format!("Invoked MCP server does not exist: {server}"))?; + if let Some(filter) = self.tool_filters.get(server) + && !filter.allows(tool) + { + return Err(anyhow!("{tool} not found in {server} MCP server catalog")); + } + let mut request = CallToolRequestParams::new(tool.to_owned()); request.arguments = arguments.as_object().cloned(); @@ -869,6 +895,7 @@ mod tests { FIXTURE_ANNOTATED_URI, FixtureServer, add_fixture_server, fixture_runtime, }; use super::*; + use crate::config::mcp_tool_policy::LayerSource; use crate::function::ToolCall; use log::{Level, LevelFilter, Log, Metadata, Record}; use std::sync::atomic::Ordering; @@ -1643,4 +1670,157 @@ mod tests { assert_eq!(prompt, "summarize"); assert_eq!(typed_keys, vec!["path".to_string()]); } + + fn single_layer_filter(source: LayerSource, patterns: &[&str]) -> ToolFilter { + let mut filter = ToolFilter::default(); + filter.push_layer( + source, + &patterns.iter().map(|s| s.to_string()).collect::>(), + ); + filter + } + + fn deny_all_filter() -> ToolFilter { + single_layer_filter(LayerSource::Global, &[]) + } + + #[tokio::test] + async fn catalog_items_drops_filtered_tools_but_keeps_other_kinds() { + let fixture = FixtureServer { + resources_capability: true, + prompts_capability: true, + ..Default::default() + }; + let (mut runtime, _server) = fixture_runtime(fixture).await; + runtime + .tool_filters + .insert("fixture".to_string(), deny_all_filter()); + + let items = runtime.catalog_items("fixture").await.unwrap(); + + assert!(!items.contains_key("tool:dup")); + assert!(items.contains_key("resource:dup")); + assert!(items.contains_key("resource_template:file:///{path}/{name}")); + assert!(items.contains_key("prompt:summarize")); + } + + #[tokio::test] + async fn catalog_items_keeps_tools_the_filter_allows() { + let (mut runtime, _server) = fixture_runtime(FixtureServer::default()).await; + runtime.tool_filters.insert( + "fixture".to_string(), + single_layer_filter(LayerSource::Global, &["d*"]), + ); + + let items = runtime.catalog_items("fixture").await.unwrap(); + + assert!(items.contains_key("tool:dup")); + } + + #[tokio::test] + async fn search_never_surfaces_filtered_tools() { + let fixture = FixtureServer { + resources_capability: true, + ..Default::default() + }; + let (mut runtime, _server) = fixture_runtime(fixture).await; + runtime + .tool_filters + .insert("fixture".to_string(), deny_all_filter()); + + let results = runtime.search("fixture", "dup", 10).await.unwrap(); + + assert!( + results + .iter() + .all(|item| item.kind != CatalogItemKind::Tool) + ); + assert!( + results + .iter() + .any(|item| item.kind == CatalogItemKind::Resource) + ); + } + + #[tokio::test] + async fn describe_blocked_tool_is_indistinguishable_from_missing() { + let (unfiltered, _server) = fixture_runtime(FixtureServer::default()).await; + let missing = unfiltered + .describe("fixture", "tool", "ghost") + .await + .unwrap_err() + .to_string(); + + let (mut filtered, _other_server) = fixture_runtime(FixtureServer::default()).await; + filtered + .tool_filters + .insert("fixture".to_string(), deny_all_filter()); + let blocked = filtered + .describe("fixture", "tool", "dup") + .await + .unwrap_err() + .to_string(); + + assert_eq!(blocked, "dup not found in fixture MCP server catalog"); + assert_eq!(blocked, missing.replace("ghost", "dup")); + } + + #[tokio::test] + async fn invoke_blocked_tool_errors_like_describe_and_never_reaches_server() { + let fixture = FixtureServer::default(); + let call_tool_calls = Arc::clone(&fixture.call_tool_calls); + let (mut runtime, _server) = fixture_runtime(fixture).await; + runtime + .tool_filters + .insert("fixture".to_string(), deny_all_filter()); + + let invoke_err = runtime + .invoke("fixture", "dup", json!({})) + .await + .unwrap_err() + .to_string(); + let describe_err = runtime + .describe("fixture", "tool", "dup") + .await + .unwrap_err() + .to_string(); + + assert_eq!(invoke_err, "dup not found in fixture MCP server catalog"); + assert_eq!(invoke_err, describe_err); + assert_eq!(call_tool_calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn invoke_allowed_tool_still_reaches_the_server() { + let fixture = FixtureServer::default(); + let call_tool_calls = Arc::clone(&fixture.call_tool_calls); + let (mut runtime, _server) = fixture_runtime(fixture).await; + runtime.tool_filters.insert( + "fixture".to_string(), + single_layer_filter(LayerSource::Global, &["d*"]), + ); + + let _ = runtime.invoke("fixture", "dup", json!({})).await; + + assert_eq!(call_tool_calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn catalog_items_warns_on_dead_context_patterns() { + install_warn_collector(); + let (mut runtime, _server) = fixture_runtime(FixtureServer::default()).await; + let mut filter = single_layer_filter(LayerSource::Global, &["*"]); + filter.push_layer(LayerSource::Session, &["zzz_*".to_string()]); + runtime.tool_filters.insert("fixture".to_string(), filter); + + runtime.catalog_items("fixture").await.unwrap(); + + let messages = warn_messages().lock().unwrap(); + assert!( + messages.iter().any(|msg| msg.contains("'zzz_*'") + && msg.contains("session (.set)") + && msg.contains("'fixture'")), + "missing dead-pattern warning in: {messages:?}" + ); + } } diff --git a/src/function/agents.rs b/src/function/agents.rs index e7803ba..da795dd 100644 --- a/src/function/agents.rs +++ b/src/function/agents.rs @@ -686,6 +686,7 @@ pub async fn run_agent_for_graph( sync_agent_functions_to_ctx(&mut child_ctx)?; } else { populate_agent_mcp_runtime(&mut child_ctx, &agent_mcp_servers).await?; + child_ctx.refresh_mcp_tool_filters(); sync_agent_functions_to_ctx(&mut child_ctx)?; child_ctx.init_agent_shared_variables()?; } @@ -869,6 +870,7 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result { sync_agent_functions_to_ctx(&mut child_ctx)?; } else { populate_agent_mcp_runtime(&mut child_ctx, &agent_mcp_servers).await?; + child_ctx.refresh_mcp_tool_filters(); sync_agent_functions_to_ctx(&mut child_ctx)?; child_ctx.init_agent_shared_variables()?; } @@ -1604,6 +1606,7 @@ mod tests { use crate::config::test_fixtures::{FixtureServer, fixture_runtime}; use crate::config::{AgentConfig, AppState, WorkingMode}; use crate::function::jobs::RingBuf; + use crate::mcp::{McpServer, McpServersConfig, McpTransportType}; use crate::supervisor::escalation::{EscalationQueue, EscalationRequest}; use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus}; use parking_lot::Mutex; @@ -1791,6 +1794,61 @@ mod tests { assert!(!functions.contains("mcp_invoke_fixture")); } + fn app_state_with_fixture_mcp_config() -> Arc { + let mut state = AppState::test_default(); + state.mcp_config = Some(McpServersConfig { + mcp_servers: [( + "fixture".to_string(), + McpServer { + transport_type: McpTransportType::Stdio, + command: Some("echo".to_string()), + args: None, + env: None, + cwd: None, + url: None, + headers: None, + oauth: None, + allowed_tools: None, + }, + )] + .into_iter() + .collect(), + }); + Arc::new(state) + } + + #[tokio::test] + async fn spawned_child_runtime_enforces_child_agent_filters() { + use std::sync::atomic::Ordering; + + let config = AgentConfig { + mcp_tools: Some(IndexMap::from([( + "fixture".to_string(), + vec!["get_*".to_string()], + )])), + ..Default::default() + }; + let mut ctx = RequestContext::new(app_state_with_fixture_mcp_config(), WorkingMode::Cmd); + ctx.agent = Some(Agent::test_new(config)); + let fixture = FixtureServer::default(); + let call_tool_calls = Arc::clone(&fixture.call_tool_calls); + let (runtime, _server) = fixture_runtime(fixture).await; + ctx.tool_scope.mcp_runtime = runtime; + + populate_agent_mcp_runtime(&mut ctx, &[]).await.unwrap(); + ctx.refresh_mcp_tool_filters(); + + let err = ctx + .tool_scope + .mcp_runtime + .invoke("fixture", "dup", json!({})) + .await + .unwrap_err() + .to_string(); + assert_eq!(err, "dup not found in fixture MCP server catalog"); + assert_eq!(call_tool_calls.load(Ordering::SeqCst), 0); + } + #[test] fn handle_list_running_empty_supervisor() { let mut ctx = ctx_with_supervisor(4, 3); diff --git a/src/function/jobs.rs b/src/function/jobs.rs index ce2c425..9963710 100644 --- a/src/function/jobs.rs +++ b/src/function/jobs.rs @@ -457,6 +457,11 @@ async fn handle_start(ctx: &mut RequestContext, args: &Value) -> Result { .unwrap_or_else(|| json!({})); let mut mcp_runtime = McpRuntime::new(); mcp_runtime.insert(server.clone(), Arc::clone(server_handle)); + if let Some(filter) = ctx.tool_scope.mcp_runtime.tool_filters.get(&server) { + mcp_runtime + .tool_filters + .insert(server.clone(), filter.clone()); + } let job_ctx = JobCtx { mcp_runtime, current_depth: ctx.current_depth, @@ -1263,7 +1268,9 @@ fn tail_chars(text: &str, max_chars: usize) -> Option { #[cfg(test)] mod tests { use super::*; + use crate::config::test_fixtures::{FixtureServer, fixture_runtime}; use crate::config::{AppConfig, AppState, WorkingMode}; + use crate::config::{LayerSource, ToolFilter}; use crate::function::agents::{ GuardrailAction, check_pending_tasks_guardrail, handle_agent_tool, }; @@ -1697,6 +1704,43 @@ mod tests { }); } + #[test] + fn job_runtime_carries_the_servers_tool_filter() { + run_async(async { + let mut ctx = plain_ctx(); + ctx.declared_function_names + .insert("mcp_invoke_fixture".into()); + let fixture = FixtureServer::default(); + let call_tool_calls = Arc::clone(&fixture.call_tool_calls); + let (mut runtime, _server) = fixture_runtime(fixture).await; + let mut filter = ToolFilter::default(); + filter.push_layer(LayerSource::Global, &[]); + runtime.tool_filters.insert("fixture".to_string(), filter); + ctx.tool_scope.mcp_runtime = runtime; + + let started = handle_start( + &mut ctx, + &json!({"tool": "mcp_invoke_fixture", "arguments": {"tool": "dup"}}), + ) + .await + .unwrap(); + let job_id = started["job_id"].as_str().unwrap().to_string(); + + let collected = handle_collect(&ctx, &json!({"id": job_id})).await.unwrap(); + + assert_eq!(collected["status"], "failed"); + assert!( + collected["error"] + .as_str() + .unwrap() + .contains("dup not found in fixture MCP server catalog"), + "unexpected error: {}", + collected["error"] + ); + assert_eq!(call_tool_calls.load(std::sync::atomic::Ordering::SeqCst), 0); + }); + } + #[test] fn handle_start_rejects_unconnected_mcp_server() { let mut ctx = plain_ctx(); diff --git a/src/graph/executor.rs b/src/graph/executor.rs index 4c2366b..0d63daf 100644 --- a/src/graph/executor.rs +++ b/src/graph/executor.rs @@ -424,7 +424,7 @@ async fn step( Ok(StepResult::Continue(vec![next])) } NodeType::Llm(llm_node) => { - let outcome = LlmNodeExecutor::execute(llm_node, state, ctx).await?; + let outcome = LlmNodeExecutor::execute(current, llm_node, state, ctx).await?; let targets = match outcome { LlmExecutionOutcome::Continue => static_next_targets(node, current, "llm")?, LlmExecutionOutcome::FellBack(target) => vec![target], diff --git a/src/graph/llm.rs b/src/graph/llm.rs index bf87870..0b9faae 100644 --- a/src/graph/llm.rs +++ b/src/graph/llm.rs @@ -30,11 +30,12 @@ pub struct LlmNodeExecutor; impl LlmNodeExecutor { pub(super) async fn execute( + node_id: &str, node: &LlmNode, state_manager: &mut StateManager, parent_ctx: &mut RequestContext, ) -> Result { - let result = run(node, state_manager, parent_ctx).await; + let result = run(node_id, node, state_manager, parent_ctx).await; let (output, failure_reason) = match result { Ok(raw) => match &node.output_schema { Some(schema) => match structured::extract(&raw, schema, parent_ctx).await { @@ -79,6 +80,7 @@ fn outcome_from( } async fn run( + node_id: &str, node: &LlmNode, state_manager: &mut StateManager, parent_ctx: &mut RequestContext, @@ -177,6 +179,13 @@ async fn run( // Jobs are node-local: everything job__start registers while this node // runs is recorded here and reaped on every exit path below. let saved_job_scope = parent_ctx.node_job_scope.replace(Vec::new()); + // The node's tool filter layer lives in tracked context state so any + // mid-node filter recompute (e.g. a skill load) re-applies it last. + let saved_node_mcp_tools = std::mem::replace( + &mut parent_ctx.active_node_mcp_tools, + node.mcp_tools.clone().map(|map| (node_id.to_string(), map)), + ); + parent_ctx.refresh_mcp_tool_filters(); let result = match node.timeout { Some(secs) => match timeout( Duration::from_secs(secs), @@ -193,6 +202,8 @@ async fn run( let node_jobs = std::mem::replace(&mut parent_ctx.node_job_scope, saved_job_scope).unwrap_or_default(); reap_jobs(parent_ctx.supervisor.as_ref(), &node_jobs).await; + parent_ctx.active_node_mcp_tools = saved_node_mcp_tools; + parent_ctx.refresh_mcp_tool_filters(); restore_agent_skill_policy(parent_ctx, saved_agent_skill_state); result } diff --git a/src/graph/map.rs b/src/graph/map.rs index 403011a..1222887 100644 --- a/src/graph/map.rs +++ b/src/graph/map.rs @@ -85,9 +85,11 @@ impl MapNodeExecutor { let mut ctx = sub_ctx; let exec_result: Result<()> = match &branch_clone.node_type { - NodeType::Llm(n) => LlmNodeExecutor::execute(n, &mut state, &mut ctx) - .await - .map(|_| ()), + NodeType::Llm(n) => { + LlmNodeExecutor::execute(&branch_clone.id, n, &mut state, &mut ctx) + .await + .map(|_| ()) + } NodeType::Agent(n) => AgentNodeExecutor::execute(n, &mut state, &mut ctx) .await .map(|_| ()), diff --git a/src/graph/validator.rs b/src/graph/validator.rs index 262ca5e..c484d5d 100644 --- a/src/graph/validator.rs +++ b/src/graph/validator.rs @@ -215,11 +215,7 @@ impl GraphValidator { }; let expand_alias = |name: &str| { - ctx.app_config - .mapping_mcp_servers - .get(name) - .into_iter() - .flat_map(|mapped| mapped.split(',').map(|s| s.trim().to_string())) + crate::config::expand_mcp_server_alias(&ctx.app_config.mapping_mcp_servers, name) }; let mut enabled_servers: HashSet = ctx.mcp_servers.clone(); for server in &ctx.mcp_servers { @@ -255,7 +251,9 @@ impl GraphValidator { { for key in mcp_tools.keys() { let enabled = enabled_servers.contains(key) - || expand_alias(key).any(|id| enabled_servers.contains(&id)); + || expand_alias(key) + .iter() + .any(|id| enabled_servers.contains(id)); if !enabled { result.error(ValidationError::with_node( node_id,