From 5711432ac484bbf68fb486a50e9c725a99823bcc Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 27 Aug 2026 13:37:46 -0600 Subject: [PATCH 01/11] feat(mcp): add per-server tool allowlist policy module and allowedTools config field --- src/config/install_remote.rs | 43 ++ src/config/mcp_factory.rs | 2 + src/config/mcp_tool_policy.rs | 759 ++++++++++++++++++++++++++++++++++ src/config/mod.rs | 1 + src/config/request_context.rs | 1 + src/mcp/manage.rs | 2 + src/mcp/mod.rs | 24 +- 7 files changed, 831 insertions(+), 1 deletion(-) create mode 100644 src/config/mcp_tool_policy.rs diff --git a/src/config/install_remote.rs b/src/config/install_remote.rs index 5013171..84c9d9e 100644 --- a/src/config/install_remote.rs +++ b/src/config/install_remote.rs @@ -4790,6 +4790,49 @@ mod tests { let _ = fs::remove_dir_all(&dir); } + #[test] + fn uninstall_mcp_preserves_allowed_tools_on_surviving_entries() { + let dir = fresh_temp_dir("uninst-mcp-allowed-tools-"); + let mut store = BundleStore::load_from(dir.join("installed-bundles.yaml")).unwrap(); + store + .upsert_bundle("omc", test_metadata("https://github.com/x/omc")) + .unwrap(); + let mcp = dir.join("mcp.json"); + write_mcp( + &mcp, + r#"{"mcpServers": { + "srv": {"type": "stdio", "command": "echo"}, + "user-srv": {"type": "stdio", "command": "mine", "allowedTools": ["get_*", "list_issues"]} + }}"#, + ); + let parsed: McpServersConfig = + serde_json::from_str(&fs::read_to_string(&mcp).unwrap()).unwrap(); + let hash = hash_bytes( + serde_json::to_string(parsed.mcp_servers.get("srv").unwrap()) + .unwrap() + .as_bytes(), + ); + store + .record_mcp_servers( + "omc", + vec![mcp_server_record("srv", McpAction::Added, Some(hash))], + ) + .unwrap(); + let servers = store.get("omc").unwrap().mcp_servers.clone(); + + let summary = uninstall_mcp_entries(&mut store, "omc", &servers, &mcp, true).unwrap(); + + assert_eq!(summary.removed, vec!["srv"]); + let raw = fs::read_to_string(&mcp).unwrap(); + assert!(raw.contains("allowedTools")); + let written: McpServersConfig = serde_json::from_str(&raw).unwrap(); + assert_eq!( + written.mcp_servers.get("user-srv").unwrap().allowed_tools, + Some(vec!["get_*".to_string(), "list_issues".to_string()]) + ); + let _ = fs::remove_dir_all(&dir); + } + #[test] fn uninstall_mcp_reports_referenced_secrets_without_removing_them() { let dir = fresh_temp_dir("uninst-mcp-secrets-"); diff --git a/src/config/mcp_factory.rs b/src/config/mcp_factory.rs index 961847d..23dba66 100644 --- a/src/config/mcp_factory.rs +++ b/src/config/mcp_factory.rs @@ -139,6 +139,7 @@ mod tests { url: None, headers: None, oauth: None, + allowed_tools: None, } } @@ -156,6 +157,7 @@ mod tests { url: Some(url.to_string()), headers, oauth: None, + allowed_tools: None, } } diff --git a/src/config/mcp_tool_policy.rs b/src/config/mcp_tool_policy.rs new file mode 100644 index 0000000..48581b1 --- /dev/null +++ b/src/config/mcp_tool_policy.rs @@ -0,0 +1,759 @@ +use crate::mcp::McpServersConfig; + +use fancy_regex::Regex; +use indexmap::IndexMap; +use log::warn; +use std::collections::HashMap; +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, + AppConfig, + Role(String), + Agent(String), + Session, + Skill(String), + Node(String), +} + +impl fmt::Display for LayerSource { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + LayerSource::Global => write!(f, "global (mcp.json)"), + LayerSource::AppConfig => write!(f, "config (config.yaml)"), + LayerSource::Role(name) => write!(f, "role ({name})"), + LayerSource::Agent(name) => write!(f, "agent ({name})"), + LayerSource::Session => write!(f, "session (.set)"), + LayerSource::Skill(name) => write!(f, "skill ({name})"), + LayerSource::Node(id) => write!(f, "node ({id})"), + } + } +} + +#[allow(dead_code)] +#[derive(Debug)] +pub struct CompiledPatterns { + source: LayerSource, + raw: Vec, + regexes: Vec, +} + +#[allow(dead_code)] +#[derive(Debug, 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 { + source, + raw: patterns.to_vec(), + regexes: patterns.iter().map(|p| compile_glob(p)).collect(), + }); + } + + /// A tool is allowed iff it matches at least one pattern in every layer. + pub fn allows(&self, tool: &str) -> bool { + self.layers.iter().all(|layer| { + layer + .regexes + .iter() + .any(|regex| regex.is_match(tool).unwrap_or(false)) + }) + } + + /// The first matching raw pattern per layer, in layer order, or the + /// source of the first layer with no match. + pub fn allows_explain(&self, tool: &str) -> Result, &LayerSource> { + let mut matched = Vec::with_capacity(self.layers.len()); + for layer in &self.layers { + // fancy_regex can fail at match time (backtracking limits); + // treat that as a non-match rather than allowing the tool. + match layer + .regexes + .iter() + .position(|regex| regex.is_match(tool).unwrap_or(false)) + { + Some(index) => matched.push((&layer.source, layer.raw[index].as_str())), + None => return Err(&layer.source), + } + } + Ok(matched) + } +} + +/// 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!( + "^{}$", + fancy_regex::escape(pattern) + .replace("\\*", ".*") + .replace("\\?", ".") + ); + Regex::new(&translated).unwrap_or_else(|error| { + warn!("Invalid MCP tool pattern '{pattern}': {error}. It will match nothing."); + never_matching_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( + mcp_config: &McpServersConfig, + session: Option<&IndexMap>>, + agent: Option<(&str, &IndexMap>)>, + role: Option<(&str, &IndexMap>)>, + global: Option<&IndexMap>>, + skills: &[SkillMcpLayer], + node: Option<(&str, &IndexMap>)>, + aliases: &IndexMap, + ) -> HashMap { + let mut filters: HashMap = HashMap::new(); + + for (server, spec) in &mcp_config.mcp_servers { + if let Some(patterns) = &spec.allowed_tools { + filters + .entry(server.clone()) + .or_default() + .push_layer(LayerSource::Global, patterns); + } + } + + if let Some(map) = global { + push_level( + &mut filters, + mcp_config, + aliases, + &LayerSource::AppConfig, + map, + None, + ); + } + if let Some((name, map)) = role { + push_level( + &mut filters, + mcp_config, + aliases, + &LayerSource::Role(name.to_string()), + map, + None, + ); + } + if let Some((name, map)) = agent { + push_level( + &mut filters, + mcp_config, + aliases, + &LayerSource::Agent(name.to_string()), + map, + None, + ); + } + if let Some(map) = session { + push_level( + &mut filters, + mcp_config, + aliases, + &LayerSource::Session, + map, + None, + ); + } + for skill in skills { + push_level( + &mut filters, + mcp_config, + aliases, + &LayerSource::Skill(skill.name.clone()), + &skill.mcp_tools, + Some(&skill.enabled_servers), + ); + } + if let Some((id, map)) = node { + push_level( + &mut filters, + mcp_config, + aliases, + &LayerSource::Node(id.to_string()), + map, + None, + ); + } + + filters + } +} + +#[allow(dead_code)] +fn push_level( + filters: &mut HashMap, + mcp_config: &McpServersConfig, + aliases: &IndexMap, + source: &LayerSource, + map: &IndexMap>, + enabled_servers: Option<&[String]>, +) { + for (server, patterns) in expand_server_keys(mcp_config, aliases, map) { + if let Some(enabled) = enabled_servers + && !enabled.iter().any(|id| id == &server) + { + continue; + } + filters + .entry(server) + .or_default() + .push_layer(source.clone(), &patterns); + } +} + +/// Expands map keys to server ids: a key that is a server id maps to itself; +/// 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, + map: &IndexMap>, +) -> IndexMap> { + let mut expanded: IndexMap> = IndexMap::new(); + for (key, patterns) in map { + let key = key.trim(); + if mcp_config.mcp_servers.contains_key(key) { + expanded + .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) { + expanded + .entry(mapped_id.to_string()) + .or_default() + .extend(patterns.iter().cloned()); + } + } + } + } + expanded +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mcp::{McpServer, McpServersConfig, McpTransportType}; + + fn spec(allowed_tools: Option<&[&str]>) -> McpServer { + McpServer { + transport_type: McpTransportType::Stdio, + command: Some("echo".to_string()), + args: None, + env: None, + cwd: None, + url: None, + headers: None, + oauth: None, + allowed_tools: allowed_tools.map(list), + } + } + + fn config(servers: &[(&str, Option<&[&str]>)]) -> McpServersConfig { + McpServersConfig { + mcp_servers: servers + .iter() + .map(|(name, tools)| (name.to_string(), spec(*tools))) + .collect(), + } + } + + fn list(items: &[&str]) -> Vec { + items.iter().map(|s| s.to_string()).collect() + } + + fn tool_map(entries: &[(&str, &[&str])]) -> IndexMap> { + entries + .iter() + .map(|(server, patterns)| (server.to_string(), list(patterns))) + .collect() + } + + fn single_layer(patterns: &[&str]) -> ToolFilter { + layered(&[(LayerSource::Global, patterns)]) + } + + fn layered(layers: &[(LayerSource, &[&str])]) -> ToolFilter { + let mut filter = ToolFilter::default(); + for (source, patterns) in layers { + filter.push_layer(source.clone(), &list(patterns)); + } + filter + } + + fn no_aliases() -> IndexMap { + IndexMap::new() + } + + fn aliases(entries: &[(&str, &str)]) -> IndexMap { + entries + .iter() + .map(|(key, value)| (key.to_string(), value.to_string())) + .collect() + } + + fn resolve( + config: &McpServersConfig, + session: Option<&IndexMap>>, + role: Option<(&str, &IndexMap>)>, + ) -> HashMap { + McpToolPolicy::effective(config, session, None, role, None, &[], None, &no_aliases()) + } + + #[test] + fn literal_pattern_matches_only_the_exact_name() { + let filter = single_layer(&["get_issue"]); + + assert!(filter.allows("get_issue")); + assert!(!filter.allows("get_issues")); + assert!(!filter.allows("get_issu")); + assert!(!filter.allows("xget_issue")); + } + + #[test] + fn star_matches_any_run_of_characters() { + let filter = single_layer(&["get_*"]); + assert!(filter.allows("get_issue")); + assert!(filter.allows("get_")); + assert!(!filter.allows("set_issue")); + + let filter = single_layer(&["*_issue"]); + assert!(filter.allows("create_issue")); + assert!(!filter.allows("create_pr")); + + let filter = single_layer(&["get*sue"]); + assert!(filter.allows("get_issue")); + assert!(filter.allows("getsue")); + + let filter = single_layer(&["*"]); + assert!(filter.allows("")); + assert!(filter.allows("anything_at_all")); + } + + #[test] + fn question_mark_matches_exactly_one_character() { + let filter = single_layer(&["get_?"]); + + assert!(filter.allows("get_a")); + assert!(!filter.allows("get_")); + assert!(!filter.allows("get_ab")); + } + + #[test] + fn regex_metacharacters_are_matched_literally() { + let filter = single_layer(&["get.issue"]); + assert!(filter.allows("get.issue")); + assert!(!filter.allows("getXissue")); + + for pattern in ["a(b", "a[b", "a+b", "a|b", "a$b"] { + let filter = single_layer(&[pattern]); + assert!(filter.allows(pattern), "'{pattern}' should match itself"); + assert!(!filter.allows("ab"), "'{pattern}' should not match 'ab'"); + } + } + + #[test] + fn backslash_is_literal_and_star_still_wildcards() { + let filter = single_layer(&["a\\b"]); + assert!(filter.allows("a\\b")); + assert!(!filter.allows("ab")); + + let filter = single_layer(&["a\\*b"]); + assert!(filter.allows("a\\b")); + assert!(filter.allows("a\\xyzb")); + assert!(!filter.allows("ab")); + } + + #[test] + fn the_never_matching_placeholder_matches_nothing() { + let regex = never_matching_regex(); + + assert!(!regex.is_match("").unwrap()); + assert!(!regex.is_match("anything").unwrap()); + } + + #[test] + fn within_a_layer_any_pattern_may_match() { + let filter = single_layer(&["get_*", "set_*"]); + + assert!(filter.allows("get_x")); + assert!(filter.allows("set_x")); + assert!(!filter.allows("delete_x")); + } + + #[test] + fn across_layers_every_layer_must_match() { + let filter = layered(&[ + (LayerSource::Global, &["get_*"]), + (LayerSource::Session, &["*_issue"]), + ]); + + assert!(filter.allows("get_issue")); + assert!(!filter.allows("get_pr")); + assert!(!filter.allows("create_issue")); + } + + #[test] + fn an_empty_layer_blocks_everything() { + let filter = layered(&[(LayerSource::Global, &["*"]), (LayerSource::Session, &[])]); + + assert!(!filter.allows("anything")); + assert_eq!( + filter.allows_explain("anything"), + Err(&LayerSource::Session) + ); + } + + #[test] + fn allows_explain_reports_the_first_matching_pattern_per_layer() { + let filter = layered(&[ + (LayerSource::Global, &["x_*", "get_*"]), + (LayerSource::Session, &["*"]), + ]); + + assert_eq!( + filter.allows_explain("get_issue").unwrap(), + vec![ + (&LayerSource::Global, "get_*"), + (&LayerSource::Session, "*") + ] + ); + } + + #[test] + fn allows_explain_reports_the_first_layer_without_a_match() { + let filter = layered(&[ + (LayerSource::Global, &["get_*"]), + (LayerSource::Session, &["*"]), + ]); + assert_eq!( + filter.allows_explain("delete_repo"), + Err(&LayerSource::Global) + ); + + let filter = layered(&[ + (LayerSource::Global, &["*"]), + (LayerSource::Session, &["get_*"]), + ]); + assert_eq!( + filter.allows_explain("delete_repo"), + Err(&LayerSource::Session) + ); + } + + #[test] + fn global_allowed_tools_from_mcp_json_is_the_first_layer() { + let config = config(&[("gh", Some(&["get_*"]))]); + let session_map = tool_map(&[("gh", &["*"])]); + + let filters = resolve(&config, Some(&session_map), None); + + assert_eq!( + filters["gh"].allows_explain("get_issue").unwrap(), + vec![ + (&LayerSource::Global, "get_*"), + (&LayerSource::Session, "*") + ] + ); + } + + #[test] + fn servers_without_patterns_at_any_level_are_absent() { + let config = config(&[("gh", None)]); + + let filters = resolve(&config, None, None); + + assert!(filters.is_empty()); + } + + #[test] + fn server_absent_from_a_level_map_gets_no_layer_from_it() { + let config = config(&[("gh", Some(&["get_*"])), ("gl", None)]); + let role_map = tool_map(&[("gl", &["x_*"])]); + + let filters = resolve(&config, None, Some(("dev", &role_map))); + + assert!(filters["gh"].allows("get_issue")); + assert!(!filters["gh"].allows("delete_repo")); + assert_eq!(filters["gh"].allows_explain("get_issue").unwrap().len(), 1); + assert!(filters["gl"].allows("x_1")); + assert!(!filters["gl"].allows("y_1")); + } + + #[test] + fn empty_pattern_list_at_a_level_blocks_all_tools_for_that_server() { + let config = config(&[("gh", Some(&["get_*"]))]); + let session_map = tool_map(&[("gh", &[])]); + + let filters = resolve(&config, Some(&session_map), None); + + assert!(!filters["gh"].allows("get_issue")); + assert_eq!( + filters["gh"].allows_explain("get_issue"), + Err(&LayerSource::Session) + ); + } + + #[test] + fn session_cannot_widen_a_role_restriction() { + let config = config(&[("gh", None)]); + let role_map = tool_map(&[("gh", &["get_*"])]); + let session_map = tool_map(&[("gh", &["*"])]); + + let filters = resolve(&config, Some(&session_map), Some(("dev", &role_map))); + + assert!(filters["gh"].allows("get_issue")); + assert!(!filters["gh"].allows("delete_repo")); + } + + #[test] + fn app_config_map_contributes_its_own_layer() { + let config = config(&[("gh", None)]); + let app_map = tool_map(&[("gh", &["get_*"])]); + + let filters = McpToolPolicy::effective( + &config, + None, + None, + None, + Some(&app_map), + &[], + None, + &no_aliases(), + ); + + assert_eq!( + filters["gh"].allows_explain("get_issue").unwrap(), + vec![(&LayerSource::AppConfig, "get_*")] + ); + assert!(!filters["gh"].allows("delete_repo")); + } + + #[test] + fn skill_layer_applies_only_to_its_enabled_servers() { + let config = config(&[("gh", None), ("gl", None)]); + let skill = SkillMcpLayer { + name: "reviewer".to_string(), + enabled_servers: vec!["gh".to_string()], + mcp_tools: tool_map(&[("gh", &["get_*"]), ("gl", &["*"])]), + }; + + let filters = McpToolPolicy::effective( + &config, + None, + None, + None, + None, + &[skill], + None, + &no_aliases(), + ); + + assert!(filters.contains_key("gh")); + assert!(!filters.contains_key("gl")); + } + + #[test] + fn two_skills_naming_the_same_server_stack_independent_layers() { + let config = config(&[("gh", None)]); + let skills = vec![ + SkillMcpLayer { + name: "a".to_string(), + enabled_servers: vec!["gh".to_string()], + mcp_tools: tool_map(&[("gh", &["get_*"])]), + }, + SkillMcpLayer { + name: "b".to_string(), + enabled_servers: vec!["gh".to_string()], + mcp_tools: tool_map(&[("gh", &["*_issue"])]), + }, + ]; + + let filters = McpToolPolicy::effective( + &config, + None, + None, + None, + None, + &skills, + None, + &no_aliases(), + ); + + assert!(filters["gh"].allows("get_issue")); + assert!(!filters["gh"].allows("get_pr")); + assert!(!filters["gh"].allows("create_issue")); + assert_eq!(filters["gh"].allows_explain("get_issue").unwrap().len(), 2); + } + + #[test] + fn layers_stack_in_documented_order_with_node_last() { + let config = config(&[("gh", Some(&["*"]))]); + let app_map = tool_map(&[("gh", &["*"])]); + let role_map = tool_map(&[("gh", &["*"])]); + let agent_map = tool_map(&[("gh", &["*"])]); + let session_map = tool_map(&[("gh", &["*"])]); + let skills = vec![SkillMcpLayer { + name: "reviewer".to_string(), + enabled_servers: vec!["gh".to_string()], + mcp_tools: tool_map(&[("gh", &["*"])]), + }]; + let node_map = tool_map(&[("gh", &["*"])]); + + let filters = McpToolPolicy::effective( + &config, + Some(&session_map), + Some(("worker", &agent_map)), + Some(("dev", &role_map)), + Some(&app_map), + &skills, + Some(("n1", &node_map)), + &no_aliases(), + ); + + let sources: Vec = filters["gh"] + .allows_explain("anything") + .unwrap() + .iter() + .map(|(source, _)| source.to_string()) + .collect(); + assert_eq!( + sources, + vec![ + "global (mcp.json)", + "config (config.yaml)", + "role (dev)", + "agent (worker)", + "session (.set)", + "skill (reviewer)", + "node (n1)", + ] + ); + } + + #[test] + fn alias_key_expands_to_all_mapped_servers() { + let config = config(&[("github", None), ("gitlab", None)]); + let role_map = tool_map(&[("gh", &["get_*"])]); + + let filters = McpToolPolicy::effective( + &config, + None, + None, + Some(("dev", &role_map)), + None, + &[], + None, + &aliases(&[("gh", "github,gitlab")]), + ); + + assert!(filters["github"].allows("get_issue")); + assert!(!filters["github"].allows("delete_repo")); + assert!(filters["gitlab"].allows("get_issue")); + assert!(!filters["gitlab"].allows("delete_repo")); + } + + #[test] + fn alias_ids_missing_from_the_config_are_skipped() { + let config = config(&[("github", None)]); + let role_map = tool_map(&[("gh", &["get_*"])]); + + let filters = McpToolPolicy::effective( + &config, + None, + None, + Some(("dev", &role_map)), + None, + &[], + None, + &aliases(&[("gh", "github,missing")]), + ); + + assert_eq!(filters.len(), 1); + assert!(filters.contains_key("github")); + } + + #[test] + fn unknown_map_keys_are_dropped() { + let config = config(&[("github", None)]); + let role_map = tool_map(&[("nope", &["get_*"])]); + + let filters = resolve(&config, None, Some(("dev", &role_map))); + + assert!(filters.is_empty()); + } + + #[test] + fn alias_and_direct_key_for_the_same_server_merge_into_one_layer() { + let config = config(&[("github", None)]); + let role_map = tool_map(&[("gh", &["get_*"]), ("github", &["set_*"])]); + + let filters = McpToolPolicy::effective( + &config, + None, + None, + Some(("dev", &role_map)), + None, + &[], + None, + &aliases(&[("gh", "github")]), + ); + + assert!(filters["github"].allows("get_issue")); + assert!(filters["github"].allows("set_topic")); + assert!(!filters["github"].allows("delete_repo")); + assert_eq!( + filters["github"].allows_explain("get_issue").unwrap().len(), + 1 + ); + } + + #[test] + fn layer_source_display() { + assert_eq!(LayerSource::Global.to_string(), "global (mcp.json)"); + assert_eq!(LayerSource::AppConfig.to_string(), "config (config.yaml)"); + assert_eq!(LayerSource::Role("dev".into()).to_string(), "role (dev)"); + assert_eq!( + LayerSource::Agent("worker".into()).to_string(), + "agent (worker)" + ); + assert_eq!(LayerSource::Session.to_string(), "session (.set)"); + assert_eq!( + LayerSource::Skill("review".into()).to_string(), + "skill (review)" + ); + assert_eq!(LayerSource::Node("n1".into()).to_string(), "node (n1)"); + } +} diff --git a/src/config/mod.rs b/src/config/mod.rs index e89caee..4832c9e 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -8,6 +8,7 @@ pub(crate) mod instructions; mod macro_policy; mod macros; mod mcp_factory; +mod mcp_tool_policy; pub(crate) mod memory; pub(crate) mod paths; pub(crate) mod prompts; diff --git a/src/config/request_context.rs b/src/config/request_context.rs index b8a7368..94092c8 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -5599,6 +5599,7 @@ mod tests { url: None, headers: None, oauth: None, + allowed_tools: None, }, ); } diff --git a/src/mcp/manage.rs b/src/mcp/manage.rs index c9d4a53..aae7598 100644 --- a/src/mcp/manage.rs +++ b/src/mcp/manage.rs @@ -242,6 +242,7 @@ fn build_stdio(cli: &Cli, has_url: bool) -> Result { url: None, headers: None, oauth: None, + allowed_tools: None, }) } @@ -300,6 +301,7 @@ fn build_remote(cli: &Cli, transport: McpTransportType, has_command: bool) -> Re url: Some(url), headers: (!headers.is_empty()).then_some(headers), oauth, + allowed_tools: None, }) } diff --git a/src/mcp/mod.rs b/src/mcp/mod.rs index 99d8272..72f3930 100644 --- a/src/mcp/mod.rs +++ b/src/mcp/mod.rs @@ -6,7 +6,7 @@ mod sse_transport; use crate::config::AppConfig; use crate::config::paths; -use crate::utils::{AbortSignal, abortable_run_with_spinner}; +use crate::utils::{AbortSignal, abortable_run_with_spinner, dimmed_text}; use crate::vault::Vault; use crate::vault::interpolate_secrets; use anyhow::Error; @@ -166,6 +166,8 @@ pub(crate) struct McpServer { pub headers: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub oauth: Option, + #[serde(rename = "allowedTools", skip_serializing_if = "Option::is_none")] + pub allowed_tools: Option>, } impl McpServer { @@ -177,6 +179,15 @@ impl McpServer { } pub fn validate(&self, name: &str) -> Result<()> { + if let Some(tools) = &self.allowed_tools + && tools.is_empty() + { + let message = format!( + "MCP server '{name}' has an empty \"allowedTools\" list, so none of its tools will be callable" + ); + warn!("{message}"); + eprintln!("{}", dimmed_text(&message)); + } if self.is_remote() { let type_label = match self.transport_type { McpTransportType::Http => "http", @@ -802,6 +813,7 @@ mod tests { url: None, headers: None, oauth: None, + allowed_tools: None, } } @@ -815,6 +827,7 @@ mod tests { url: Some(url.to_string()), headers: None, oauth: None, + allowed_tools: None, } } @@ -828,6 +841,7 @@ mod tests { url: Some(url.to_string()), headers: None, oauth: None, + allowed_tools: None, } } @@ -860,6 +874,7 @@ mod tests { url: None, headers: None, oauth: None, + allowed_tools: None, }; let err = spec.validate("test").unwrap_err(); @@ -878,6 +893,7 @@ mod tests { url: Some("http://localhost".into()), headers: None, oauth: None, + allowed_tools: None, }; let err = spec.validate("test").unwrap_err(); @@ -898,6 +914,7 @@ mod tests { url: None, headers: Some(headers), oauth: None, + allowed_tools: None, }; let err = spec.validate("test").unwrap_err(); @@ -923,6 +940,7 @@ mod tests { url: None, headers: None, oauth: None, + allowed_tools: None, }; let err = spec.validate("test").unwrap_err(); @@ -941,6 +959,7 @@ mod tests { url: Some("http://localhost".into()), headers: None, oauth: None, + allowed_tools: None, }; let err = spec.validate("test").unwrap_err(); @@ -959,6 +978,7 @@ mod tests { url: Some("http://localhost".into()), headers: None, oauth: None, + allowed_tools: None, }; let err = spec.validate("test").unwrap_err(); @@ -977,6 +997,7 @@ mod tests { url: Some("http://localhost".into()), headers: None, oauth: None, + allowed_tools: None, }; let err = spec.validate("test").unwrap_err(); @@ -1002,6 +1023,7 @@ mod tests { url: None, headers: None, oauth: None, + allowed_tools: None, }; let err = spec.validate("test").unwrap_err(); From a92ebb5b94ec905f3c64abebb5a1dd418bddfcb0 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 27 Aug 2026 14:00:04 -0600 Subject: [PATCH 02/11] feat(mcp): add mcp_tools allowlist config surfaces across roles, sessions, agents, graphs, and skills --- src/config/agent.rs | 38 +++++++++++ src/config/app_config.rs | 19 ++++++ src/config/mod.rs | 13 ++++ src/config/role.rs | 97 ++++++++++++++++++++++++++ src/config/session.rs | 62 +++++++++++++++++ src/config/skill.rs | 61 +++++++++++++++++ src/graph/llm.rs | 1 + src/graph/types.rs | 47 +++++++++++++ src/graph/validator.rs | 142 +++++++++++++++++++++++++++++++++++++++ 9 files changed, 480 insertions(+) diff --git a/src/config/agent.rs b/src/config/agent.rs index 15638ea..87d6fe6 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -682,6 +682,10 @@ impl RoleLike for Agent { Some(self.config.mcp_servers.clone()) } + fn mcp_tools(&self) -> Option>> { + self.config.mcp_tools.clone() + } + fn set_model(&mut self, model: Model) { self.config.model_id = Some(model.id()); self.model = model; @@ -723,6 +727,10 @@ impl RoleLike for Agent { } } } + + fn set_mcp_tools(&mut self, value: Option>>) { + self.config.mcp_tools = value; + } } #[derive(Debug, Clone, Default, Deserialize, Serialize)] @@ -774,6 +782,8 @@ pub struct AgentConfig { pub version: String, #[serde(default)] pub mcp_servers: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub mcp_tools: Option>>, #[serde(default)] pub global_tools: Vec, #[serde(skip_serializing_if = "Option::is_none")] @@ -846,6 +856,7 @@ impl AgentConfig { description: graph.description.clone(), global_tools: graph.global_tools.clone(), mcp_servers: graph.mcp_servers.clone(), + mcp_tools: graph.mcp_tools.clone(), skills_enabled: graph.skills_enabled, enabled_skills: graph.enabled_skills.clone(), inject_skill_instructions: graph.inject_skill_instructions.unwrap_or(true), @@ -1285,6 +1296,33 @@ variables: assert_eq!(config.enabled_macros, None); } + #[test] + fn agent_config_parses_mcp_tools() { + let yaml = + "name: minimal\ninstructions: hi\nmcp_tools:\n github:\n - get_*\n - list_*\n"; + let config: AgentConfig = serde_yaml::from_str(yaml).unwrap(); + + let mcp_tools = config.mcp_tools.unwrap(); + assert_eq!( + mcp_tools.get("github"), + Some(&vec!["get_*".to_string(), "list_*".to_string()]) + ); + } + + #[test] + fn agent_mcp_tools_role_like_round_trip() { + let config: AgentConfig = + serde_yaml::from_str("name: minimal\ninstructions: hi\n").unwrap(); + let mut agent = Agent::test_new(config); + assert_eq!(agent.mcp_tools(), None); + + let mut mcp_tools = IndexMap::new(); + mcp_tools.insert("github".to_string(), vec!["get_*".to_string()]); + agent.set_mcp_tools(Some(mcp_tools.clone())); + + assert_eq!(agent.mcp_tools(), Some(mcp_tools)); + } + #[test] fn agent_config_enabled_macros_empty_list_is_some_empty() { let yaml = "name: minimal\ninstructions: hi\nenabled_macros: []\n"; diff --git a/src/config/app_config.rs b/src/config/app_config.rs index 4f4e4b1..6235c1d 100644 --- a/src/config/app_config.rs +++ b/src/config/app_config.rs @@ -50,6 +50,7 @@ pub struct AppConfig { pub mapping_mcp_servers: IndexMap, #[serde(default, deserialize_with = "super::deserialize_csv_or_vec")] pub enabled_mcp_servers: Option>, + pub mcp_tools: Option>>, pub auto_continue: bool, pub max_auto_continues: usize, @@ -136,6 +137,7 @@ impl Default for AppConfig { mcp_server_support: true, mapping_mcp_servers: Default::default(), enabled_mcp_servers: None, + mcp_tools: None, auto_continue: false, max_auto_continues: 10, @@ -223,6 +225,7 @@ impl AppConfig { mcp_server_support: config.mcp_server_support, mapping_mcp_servers: config.mapping_mcp_servers, enabled_mcp_servers: config.enabled_mcp_servers, + mcp_tools: config.mcp_tools, auto_continue: config.auto_continue, max_auto_continues: config.max_auto_continues, @@ -786,6 +789,22 @@ mod tests { ); } + #[test] + fn from_config_copies_mcp_tools() { + let mut mcp_tools = IndexMap::new(); + mcp_tools.insert("github".to_string(), vec!["get_*".to_string()]); + let cfg = Config { + model_id: "test-model".to_string(), + clients: vec![ClientConfig::default()], + mcp_tools: Some(mcp_tools.clone()), + ..Config::default() + }; + + let app = AppConfig::from_config(cfg).unwrap(); + + assert_eq!(app.mcp_tools, Some(mcp_tools)); + } + #[test] #[serial_test::serial] fn from_config_copies_enabled_macros() { diff --git a/src/config/mod.rs b/src/config/mod.rs index 4832c9e..d04db86 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -251,6 +251,7 @@ pub struct Config { pub mapping_mcp_servers: IndexMap, #[serde(default, deserialize_with = "deserialize_csv_or_vec")] pub enabled_mcp_servers: Option>, + pub mcp_tools: Option>>, pub auto_continue: bool, pub max_auto_continues: usize, @@ -334,6 +335,7 @@ impl Default for Config { mcp_server_support: true, mapping_mcp_servers: Default::default(), enabled_mcp_servers: None, + mcp_tools: None, auto_continue: false, max_auto_continues: 10, @@ -1129,6 +1131,17 @@ clients: assert!(validate_no_template_in_secrets_provider(yaml).is_ok()); } + #[test] + fn config_yaml_parses_mcp_tools() { + let cfg: Config = serde_yaml::from_str("mcp_tools:\n github:\n - get_*\n").unwrap(); + + assert_eq!( + cfg.mcp_tools.as_ref().unwrap().get("github"), + Some(&vec!["get_*".to_string()]) + ); + assert_eq!(Config::default().mcp_tools, None); + } + #[test] fn config_defaults_match_expected() { let cfg = Config::default(); diff --git a/src/config/role.rs b/src/config/role.rs index db82655..39acb4f 100644 --- a/src/config/role.rs +++ b/src/config/role.rs @@ -30,6 +30,7 @@ pub trait RoleLike { fn top_p(&self) -> Option; fn enabled_tools(&self) -> Option>; fn enabled_mcp_servers(&self) -> Option>; + fn mcp_tools(&self) -> Option>>; fn set_model(&mut self, model: Model); fn set_temperature(&mut self, value: Option); fn reasoning_effort(&self) -> Option; @@ -37,6 +38,7 @@ pub trait RoleLike { fn set_reasoning_effort(&mut self, value: Option); fn set_enabled_tools(&mut self, value: Option>); fn set_enabled_mcp_servers(&mut self, value: Option>); + fn set_mcp_tools(&mut self, value: Option>>); } #[derive(Debug, Clone, Default, Deserialize, Serialize)] @@ -67,6 +69,8 @@ pub struct Role { deserialize_with = "super::deserialize_csv_or_vec" )] enabled_mcp_servers: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + mcp_tools: Option>>, #[serde(skip_serializing_if = "Option::is_none")] skills_enabled: Option, #[serde( @@ -133,6 +137,7 @@ impl Role { "enabled_mcp_servers" => { role.enabled_mcp_servers = parse_string_or_array(value) } + "mcp_tools" => role.mcp_tools = parse_mcp_tools_map(value), "skills_enabled" => role.skills_enabled = value.as_bool(), "enabled_skills" => role.enabled_skills = parse_string_or_array(value), "enabled_macros" => role.enabled_macros = parse_string_or_array(value), @@ -196,6 +201,10 @@ impl Role { serde_json::to_string(enabled_mcp_servers).unwrap_or_else(|_| "[]".to_string()); metadata.push(format!("enabled_mcp_servers: {inline}")); } + if let Some(mcp_tools) = &self.mcp_tools { + let inline = serde_json::to_string(mcp_tools).unwrap_or_else(|_| "{}".to_string()); + metadata.push(format!("mcp_tools: {inline}")); + } if let Some(skills_enabled) = self.skills_enabled { metadata.push(format!("skills_enabled: {skills_enabled}")); } @@ -279,6 +288,10 @@ impl Role { enabled_tools, enabled_mcp_servers, ); + let mcp_tools = role_like.mcp_tools(); + if mcp_tools.is_some() { + self.set_mcp_tools(mcp_tools); + } } pub fn batch_set( @@ -453,6 +466,10 @@ impl RoleLike for Role { self.enabled_mcp_servers.clone() } + fn mcp_tools(&self) -> Option>> { + self.mcp_tools.clone() + } + fn set_model(&mut self, model: Model) { if !self.model().id().is_empty() { self.model_id = Some(model.id().to_string()); @@ -479,6 +496,10 @@ impl RoleLike for Role { fn set_enabled_mcp_servers(&mut self, value: Option>) { self.enabled_mcp_servers = value; } + + fn set_mcp_tools(&mut self, value: Option>>) { + self.mcp_tools = value; + } } fn parse_string_or_array(value: &Value) -> Option> { @@ -503,6 +524,19 @@ fn parse_string_or_array(value: &Value) -> Option> { None } +fn parse_mcp_tools_map(value: &Value) -> Option>> { + let map = value.as_object()?; + let mut mcp_tools = IndexMap::new(); + for (server, tools) in map { + if tools.is_null() { + mcp_tools.insert(server.clone(), Vec::new()); + } else if let Some(tools) = parse_string_or_array(tools) { + mcp_tools.insert(server.clone(), tools); + } + } + Some(mcp_tools) +} + fn parse_structure_prompt(prompt: &str) -> (&str, Vec<(&str, &str)>) { let mut text = prompt; let mut search_input = true; @@ -652,6 +686,69 @@ mod tests { assert_eq!(role.enabled_macros, None); } + #[test] + fn role_new_parses_mcp_tools_list_and_csv_values() { + let content = "---\nmcp_tools:\n github: [get_*, list_*, search_code]\n slack: conversations_history,conversations_replies\n---\nPrompt"; + + let role = Role::new("test", content); + + let mcp_tools = role.mcp_tools().unwrap(); + assert_eq!( + mcp_tools.get("github"), + Some(&vec![ + "get_*".to_string(), + "list_*".to_string(), + "search_code".to_string() + ]) + ); + assert_eq!( + mcp_tools.get("slack"), + Some(&vec![ + "conversations_history".to_string(), + "conversations_replies".to_string() + ]) + ); + } + + #[test] + fn role_new_mcp_tools_empty_list_server_is_some_empty() { + let role = Role::new("test", "---\nmcp_tools:\n github: []\n---\nPrompt"); + + assert_eq!(role.mcp_tools().unwrap().get("github"), Some(&vec![])); + } + + #[test] + fn role_new_mcp_tools_per_server_null_is_some_empty() { + let role = Role::new("test", "---\nmcp_tools:\n github:\n---\nPrompt"); + + assert_eq!(role.mcp_tools().unwrap().get("github"), Some(&vec![])); + } + + #[test] + fn role_new_mcp_tools_absent_is_none() { + let role = Role::new("test", "---\ntemperature: 0.5\n---\nPrompt"); + + assert_eq!(role.mcp_tools(), None); + } + + #[test] + fn role_new_mcp_tools_null_is_none() { + let role = Role::new("test", "---\nmcp_tools: null\n---\nPrompt"); + + assert_eq!(role.mcp_tools(), None); + } + + #[test] + fn role_export_mcp_tools_round_trips() { + let content = "---\nmcp_tools:\n github: [get_issue]\n slack: a,b\n---\nPrompt"; + let role = Role::new("test", content); + + let reparsed = Role::new("test", &role.export()); + + assert_eq!(reparsed.mcp_tools(), role.mcp_tools()); + assert!(role.mcp_tools().is_some()); + } + #[test] fn role_export_includes_enabled_macros() { let role = Role::new("test", "---\nenabled_macros: [a]\n---\nPrompt"); diff --git a/src/config/session.rs b/src/config/session.rs index 9022d09..b6c8d22 100644 --- a/src/config/session.rs +++ b/src/config/session.rs @@ -40,6 +40,8 @@ pub struct Session { deserialize_with = "super::deserialize_csv_or_vec" )] enabled_mcp_servers: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + mcp_tools: Option>>, #[serde(skip_serializing_if = "Option::is_none")] skills_enabled: Option, #[serde( @@ -249,6 +251,9 @@ impl Session { if let Some(enabled_mcp_servers) = self.enabled_mcp_servers() { data["enabled_mcp_servers"] = json!(enabled_mcp_servers); } + if let Some(mcp_tools) = self.mcp_tools() { + data["mcp_tools"] = json!(mcp_tools); + } if let Some(skills_enabled) = self.skills_enabled() { data["skills_enabled"] = skills_enabled.into(); } @@ -329,6 +334,13 @@ impl Session { items.push(("enabled_mcp_servers", enabled_mcp_servers.join(","))); } + if let Some(mcp_tools) = self.mcp_tools() { + items.push(( + "mcp_tools", + serde_json::to_string(&mcp_tools).unwrap_or_default(), + )); + } + if let Some(skills_enabled) = self.skills_enabled() { items.push(("skills_enabled", skills_enabled.to_string())); } @@ -870,6 +882,10 @@ impl RoleLike for Session { self.enabled_mcp_servers.clone() } + fn mcp_tools(&self) -> Option>> { + self.mcp_tools.clone() + } + fn set_model(&mut self, model: Model) { if self.model().id() != model.id() { self.model_id = model.id(); @@ -913,6 +929,13 @@ impl RoleLike for Session { self.dirty = true; } } + + fn set_mcp_tools(&mut self, value: Option>>) { + if self.mcp_tools != value { + self.mcp_tools = value; + self.dirty = true; + } + } } #[derive(Debug, Clone, Default)] @@ -1044,6 +1067,45 @@ mod tests { assert_eq!(session.enabled_macros, None); } + #[test] + fn session_mcp_tools_survives_yaml_round_trip() { + let mut session = Session::default(); + let mut mcp_tools = IndexMap::new(); + mcp_tools.insert("github".to_string(), vec!["get_*".to_string()]); + mcp_tools.insert("slack".to_string(), vec![]); + session.set_mcp_tools(Some(mcp_tools.clone())); + + let yaml = serde_yaml::to_string(&session).unwrap(); + let reloaded: Session = serde_yaml::from_str(&yaml).unwrap(); + + assert_eq!(reloaded.mcp_tools(), Some(mcp_tools)); + } + + #[test] + fn session_set_role_does_not_copy_mcp_tools() { + let role = Role::new( + "test", + "---\nmcp_tools:\n github: [get_issue]\n---\nPrompt", + ); + assert!(role.mcp_tools().is_some()); + let mut session = Session::default(); + + session.set_role(role); + + assert_eq!(session.mcp_tools(), None); + } + + #[test] + fn session_set_mcp_tools_marks_dirty() { + let mut session = Session::default(); + assert!(!session.dirty()); + + session.set_mcp_tools(Some(IndexMap::new())); + + assert!(session.dirty()); + assert_eq!(session.mcp_tools(), Some(IndexMap::new())); + } + #[test] fn session_enabled_macros_empty_list_is_some_empty() { let session: Session = diff --git a/src/config/skill.rs b/src/config/skill.rs index 971c810..e8ec760 100644 --- a/src/config/skill.rs +++ b/src/config/skill.rs @@ -37,6 +37,8 @@ pub struct Skill { #[serde(skip_serializing_if = "Option::is_none")] enabled_mcp_servers: Option>, #[serde(skip_serializing_if = "Option::is_none")] + mcp_tools: Option>>, + #[serde(skip_serializing_if = "Option::is_none")] auto_unload: Option, } @@ -74,6 +76,9 @@ impl Skill { "enabled_mcp_servers" => { skill.enabled_mcp_servers = parse_skill_string_or_array(value); } + "mcp_tools" => { + skill.mcp_tools = parse_skill_mcp_tools_map(value); + } "auto_unload" => { skill.auto_unload = value.as_bool(); } @@ -147,6 +152,11 @@ impl Skill { self.enabled_mcp_servers.as_deref() } + #[allow(dead_code)] + pub fn mcp_tools(&self) -> Option<&IndexMap>> { + self.mcp_tools.as_ref() + } + pub fn auto_unload(&self) -> bool { self.auto_unload.unwrap_or(false) } @@ -185,6 +195,19 @@ fn parse_skill_string_or_array(value: &Value) -> Option> { None } +fn parse_skill_mcp_tools_map(value: &Value) -> Option>> { + let map = value.as_object()?; + let mut mcp_tools = IndexMap::new(); + for (server, tools) in map { + if tools.is_null() { + mcp_tools.insert(server.clone(), Vec::new()); + } else if let Some(tools) = parse_skill_string_or_array(tools) { + mcp_tools.insert(server.clone(), tools); + } + } + Some(mcp_tools) +} + #[cfg(test)] mod tests { use super::*; @@ -198,6 +221,44 @@ mod tests { assert_eq!(skill.description(), ""); } + #[test] + fn skill_new_parses_mcp_tools_list_and_csv_values() { + let content = "---\nmcp_tools:\n github: [get_*, list_*]\n slack: a,b\n---\nBody"; + + let skill = Skill::new("test", content); + + let mcp_tools = skill.mcp_tools().unwrap(); + assert_eq!( + mcp_tools.get("github"), + Some(&vec!["get_*".to_string(), "list_*".to_string()]) + ); + assert_eq!( + mcp_tools.get("slack"), + Some(&vec!["a".to_string(), "b".to_string()]) + ); + } + + #[test] + fn skill_new_mcp_tools_absent_is_none() { + let skill = Skill::new("test", "---\ndescription: d\n---\nBody"); + + assert_eq!(skill.mcp_tools(), None); + } + + #[test] + fn skill_new_mcp_tools_empty_list_server_is_some_empty() { + let skill = Skill::new("test", "---\nmcp_tools:\n github: []\n---\nBody"); + + assert_eq!(skill.mcp_tools().unwrap().get("github"), Some(&vec![])); + } + + #[test] + fn skill_new_mcp_tools_per_server_null_is_some_empty() { + let skill = Skill::new("test", "---\nmcp_tools:\n github:\n---\nBody"); + + assert_eq!(skill.mcp_tools().unwrap().get("github"), Some(&vec![])); + } + #[test] fn skill_new_parses_full_metadata() { let content = "---\n\ diff --git a/src/graph/llm.rs b/src/graph/llm.rs index 623d110..bf87870 100644 --- a/src/graph/llm.rs +++ b/src/graph/llm.rs @@ -506,6 +506,7 @@ mod tests { instructions: Some("sys".into()), prompt: "user".into(), tools: None, + mcp_tools: None, model: None, temperature: None, top_p: None, diff --git a/src/graph/types.rs b/src/graph/types.rs index e45322b..f598f53 100644 --- a/src/graph/types.rs +++ b/src/graph/types.rs @@ -37,6 +37,9 @@ pub struct Graph { #[serde(default)] pub mcp_servers: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub mcp_tools: Option>>, + #[serde(default, skip_serializing_if = "Option::is_none")] pub skills_enabled: Option, @@ -285,6 +288,9 @@ pub struct LlmNode { #[serde(default, skip_serializing_if = "Option::is_none")] pub tools: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub mcp_tools: Option>>, + #[serde(default, skip_serializing_if = "Option::is_none")] pub model: Option, @@ -530,6 +536,47 @@ nodes: } } + #[test] + fn deserializes_mcp_tools_at_graph_and_node_level() { + let yaml = r#" +name: g +mcp_servers: [github] +mcp_tools: + github: + - get_* + - list_* +start: l +nodes: + l: + id: l + type: llm + prompt: hi + mcp_tools: + github: + - search_code + next: e + e: + id: e + type: end + output: done +"#; + let graph: Graph = serde_yaml::from_str(yaml).unwrap(); + + assert_eq!( + graph.mcp_tools.as_ref().unwrap().get("github"), + Some(&vec!["get_*".to_string(), "list_*".to_string()]) + ); + match &graph.get_node("l").unwrap().node_type { + NodeType::Llm(llm) => { + assert_eq!( + llm.mcp_tools.as_ref().unwrap().get("github"), + Some(&vec!["search_code".to_string()]) + ); + } + _ => panic!("expected Llm variant"), + } + } + #[test] fn deserializes_every_node_type() { let yaml = r#" diff --git a/src/graph/validator.rs b/src/graph/validator.rs index 4b7cc64..262ca5e 100644 --- a/src/graph/validator.rs +++ b/src/graph/validator.rs @@ -214,6 +214,19 @@ impl GraphValidator { return; }; + 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())) + }; + let mut enabled_servers: HashSet = ctx.mcp_servers.clone(); + for server in &ctx.mcp_servers { + enabled_servers.extend(expand_alias(server)); + } + let all_servers_enabled = ctx.mcp_servers.iter().any(|s| s.trim() == "all"); + for (node_id, node) in &graph.nodes { let NodeType::Llm(llm) = &node.node_type else { continue; @@ -237,6 +250,23 @@ impl GraphValidator { } } + if let Some(mcp_tools) = &llm.mcp_tools + && !all_servers_enabled + { + for key in mcp_tools.keys() { + let enabled = enabled_servers.contains(key) + || expand_alias(key).any(|id| enabled_servers.contains(&id)); + if !enabled { + result.error(ValidationError::with_node( + node_id, + format!( + "llm node 'mcp_tools' references MCP server '{key}' not enabled by this graph" + ), + )); + } + } + } + if let Some(model_id) = &llm.model && Model::retrieve_model(ctx.app_config.as_ref(), model_id, ModelType::Chat) .is_err() @@ -1001,6 +1031,7 @@ mod tests { max_concurrent_jobs: None, global_tools: Vec::new(), mcp_servers: Vec::new(), + mcp_tools: None, skills_enabled: None, enabled_skills: None, inject_skill_instructions: None, @@ -1099,6 +1130,7 @@ mod tests { instructions: None, prompt: "p".into(), tools: None, + mcp_tools: None, model: None, temperature: None, top_p: None, @@ -1257,6 +1289,19 @@ mod tests { node } + fn llm_node_with_mcp_tools(id: &str, servers: &[&str]) -> Node { + let mut node = llm_node(id, None, Some("end")); + if let NodeType::Llm(ref mut n) = node.node_type { + let mut mcp_tools = IndexMap::new(); + for server in servers { + mcp_tools.insert(server.to_string(), vec!["get_*".to_string()]); + } + n.mcp_tools = Some(mcp_tools); + } + + node + } + #[test] fn llm_node_unknown_tool_is_an_error() { let graph = graph_with( @@ -1340,6 +1385,103 @@ mod tests { assert!(result.is_valid()); } + #[test] + fn llm_node_mcp_tools_enabled_server_passes() { + let graph = graph_with( + vec![ + ("l", llm_node_with_mcp_tools("l", &["github"])), + ("end", end_node("end")), + ], + "l", + ); + + let result = validator() + .with_agent_context(agent_ctx(&[], &["github"])) + .validate(&graph); + + assert!(result.is_valid()); + } + + #[test] + fn llm_node_mcp_tools_unknown_server_is_an_error() { + let graph = graph_with( + vec![ + ("l", llm_node_with_mcp_tools("l", &["slack"])), + ("end", end_node("end")), + ], + "l", + ); + + let result = validator() + .with_agent_context(agent_ctx(&[], &["github"])) + .validate(&graph); + + assert!(!result.is_valid()); + assert!( + result + .errors + .iter() + .any(|e| e.message.contains("'slack' not enabled")) + ); + } + + #[test] + fn llm_node_mcp_tools_alias_key_passes() { + let graph = graph_with( + vec![ + ("l", llm_node_with_mcp_tools("l", &["gh"])), + ("end", end_node("end")), + ], + "l", + ); + let mut ctx = agent_ctx(&[], &["github-mcp"]); + let mut app = AppConfig::default(); + app.mapping_mcp_servers + .insert("gh".to_string(), "github-mcp".to_string()); + ctx.app_config = Arc::new(app); + + let result = validator().with_agent_context(ctx).validate(&graph); + + assert!(result.is_valid(), "errors: {:?}", result.errors); + } + + #[test] + fn llm_node_mcp_tools_key_matching_alias_expansion_passes() { + let graph = graph_with( + vec![ + ("l", llm_node_with_mcp_tools("l", &["github-mcp"])), + ("end", end_node("end")), + ], + "l", + ); + let mut ctx = agent_ctx(&[], &["gh"]); + let mut app = AppConfig::default(); + app.mapping_mcp_servers + .insert("gh".to_string(), "github-mcp".to_string()); + ctx.app_config = Arc::new(app); + + let result = validator().with_agent_context(ctx).validate(&graph); + + assert!(result.is_valid(), "errors: {:?}", result.errors); + } + + #[test] + fn llm_node_mcp_tools_with_all_sentinel_passes() { + let graph = graph_with( + vec![ + ("l", llm_node_with_mcp_tools("l", &["github"])), + ("end", end_node("end")), + ], + "l", + ); + + let result = validator() + .with_agent_context(agent_ctx(&[], &["all"])) + .validate(&graph); + + assert!(result.is_valid(), "errors: {:?}", result.errors); + } + #[test] fn llm_node_unknown_model_is_an_error() { let graph = graph_with( From 92d9464b387d595bc4edc00b7e98b9ad3d08994d Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 27 Aug 2026 14:43:18 -0600 Subject: [PATCH 03/11] 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, From c0224e20cbe3dd5bb45ff1f3d383aff554ffec34 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 27 Aug 2026 15:12:22 -0600 Subject: [PATCH 04/11] test(mcp): drive session re-attach and RAG attach filter guards through their real entry points - use_session_applies_persisted_mcp_tools_immediately now loads a session file from disk through the real use_session, proving the post-assignment filter refresh applies a re-attached session's persisted allowlist. - New use_rag_does_not_drop_role_filters loads a yaml-driver RAG through the real use_rag and proves the tool-scope rebuild recomputes the role's filter layer instead of dropping it. - Seed the process-wide client/model registries in a pre-main ctor (new ctor dev-dependency) so model resolution is deterministic across test orderings; the seed exposes only an embedding model so tests that assert 'no chat model available' keep their premise. --- Cargo.lock | 23 ++++++++ Cargo.toml | 3 +- src/config/request_context.rs | 102 ++++++++++++++++++++++++++++------ 3 files changed, 111 insertions(+), 17 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d856e4a..0208f01 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1685,6 +1685,7 @@ dependencies = [ "colored", "comfy-table", "crossterm 0.29.0", + "ctor", "dirs", "duckdb", "duct", @@ -1907,6 +1908,16 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "ctor" +version = "1.0.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "914a755b7c2d4af2bdcff7ce1739e2db9a1b81a9b07123d8015786ae03c0980d" +dependencies = [ + "link-section", + "linktime-proc-macro", +] + [[package]] name = "ctutils" version = "0.4.2" @@ -3833,6 +3844,18 @@ dependencies = [ "libc", ] +[[package]] +name = "link-section" +version = "0.19.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39c29a617ce3df32c08497bdc1ab6e2376e0b17948ac166a2fbe5977c5954cd9" + +[[package]] +name = "linktime-proc-macro" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e57c38c1e860fd37c604281cdfb1dd2216977fd76a50f85ba2f388ef3219616" + [[package]] name = "linux-raw-sys" version = "0.4.15" diff --git a/Cargo.toml b/Cargo.toml index fbfa196..784f512 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -143,6 +143,7 @@ arboard = { version = "3.3.0", default-features = false } libc = "0.2" [dev-dependencies] +ctor = "1.0.13" pretty_assertions = "1.4.0" rmcp = { version = "3.1.2", features = ["server"] } serial_test = "3" @@ -154,4 +155,4 @@ path = "src/main.rs" [profile.release] lto = true strip = true -opt-level = "z" \ No newline at end of file +opt-level = "z" diff --git a/src/config/request_context.rs b/src/config/request_context.rs index 8018bab..bbf711d 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -4901,6 +4901,39 @@ mod tests { use std::time::{Instant, SystemTime, UNIX_EPOCH}; use std::{env, mem}; + // `list_client_names` / `list_all_models` cache the first AppConfig they + // see in process-wide OnceLocks. Several tests reach them through configs + // with an empty client list, which would permanently pin every later + // model lookup in this process to "unknown model" and make any test that + // needs a resolvable model dependent on test ordering. Seed the caches + // before any test runs with one client, "test-seeded", exposing a single + // embedding model. Deliberately NO chat model: some tests assert that no + // chat model is available, and chat lookups don't need one — a + // "test-seeded:" chat id resolves through the create-from-name + // fallback because the client name is registered. + // + // `unsafe` is ctor's required acknowledgment that this runs before main; + // the body only allocates and initializes OnceLocks, both of which are + // sound pre-main. + #[ctor::ctor(unsafe)] + fn seed_model_registries() { + use crate::client::{ClientConfig, ModelData, list_all_models, list_client_names}; + + let mut client = ClientConfig::default(); + if let ClientConfig::OpenAIConfig(config) = &mut client { + config.name = Some("test-seeded".to_string()); + let mut embedder = ModelData::new("test-embedder"); + embedder.model_type = "embedding".to_string(); + config.models = vec![embedder]; + } + let config = AppConfig { + clients: vec![client], + ..AppConfig::default() + }; + let _ = list_client_names(&config); + let _ = list_all_models(&config); + } + struct TestConfigDirGuard { key: String, previous: Option, @@ -8199,25 +8232,23 @@ mod tests { #[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.) + // use_session rebuilds the tool scope BEFORE `self.session` is + // assigned, so a re-attached session's persisted allowlist is + // enforced only by the filter refresh that runs after the + // assignment. Drive the real use_session against a session file on + // disk to prove that refresh happens. 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(); + let session_path = ctx.session_file("persisted"); + ensure_parent_exists(&session_path).unwrap(); + write( + &session_path, + "model: test-seeded:test-chat\nmessages: []\nmcp_tools:\n gh:\n - get_*\n", + ) + .unwrap(); + let app = ctx.app.config.clone(); - 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(); + run_async(ctx.use_session(&app, Some("persisted"), utils::create_abort_signal())).unwrap(); let filter = ctx .tool_scope @@ -8229,6 +8260,45 @@ mod tests { assert!(!filter.allows("delete_repo")); } + #[test] + #[serial] + fn use_rag_does_not_drop_role_filters() { + // Attaching a RAG rebuilds the tool scope from scratch, and the + // rebuilt McpRuntime starts with no filters — so the rebuild must + // recompute the declarative filter layers or the active role's + // allowlist silently disappears. Uses the yaml driver so the load + // stays on the local filesystem; the externally-backed attach path + // needs a live vector store and cannot run here, but it funnels + // through the same tool-scope refresh. + 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 rag_path = ctx.rag_file("kb"); + ensure_parent_exists(&rag_path).unwrap(); + write( + &rag_path, + "driver: yaml\nembedding_model: test-seeded:test-embedder\nchunk_size: 512\nchunk_overlap: 64\ntop_k: 5\n", + ) + .unwrap(); + + run_async(ctx.use_rag(Some("kb"), utils::create_abort_signal())).unwrap(); + + assert!(ctx.rag.is_some(), "the RAG must actually load"); + let filter = ctx + .tool_scope + .mcp_runtime + .tool_filters + .get("gh") + .expect("attaching a RAG must not drop the role's filter layer"); + assert!(filter.allows("get_issue")); + assert!(!filter.allows("delete_repo")); + } + #[test] #[serial] fn use_agent_applies_agent_filters_without_a_session() { From e7811a6b872837517db7b0ed86ff6b47e1020190 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 27 Aug 2026 15:45:39 -0600 Subject: [PATCH 05/11] feat(mcp): add .info mcp-server, .set mcp_tools, and filtered-server listing to the REPL --- src/config/mcp_tool_policy.rs | 23 +- src/config/request_context.rs | 665 +++++++++++++++++++++++++++++++--- src/config/tool_scope.rs | 13 +- src/repl/mod.rs | 21 +- 4 files changed, 671 insertions(+), 51 deletions(-) diff --git a/src/config/mcp_tool_policy.rs b/src/config/mcp_tool_policy.rs index db48c5e..ae2fa0d 100644 --- a/src/config/mcp_tool_policy.rs +++ b/src/config/mcp_tool_policy.rs @@ -33,6 +33,21 @@ impl fmt::Display for LayerSource { } } +impl LayerSource { + /// The bare level keyword, for compact diagnostics. + pub fn short_label(&self) -> &'static str { + match self { + LayerSource::Global => "global", + LayerSource::AppConfig => "config", + LayerSource::Role(_) => "role", + LayerSource::Agent(_) => "agent", + LayerSource::Session => "session", + LayerSource::Skill(_) => "skill", + LayerSource::Node(_) => "node", + } + } +} + #[derive(Debug, Clone)] pub struct CompiledPatterns { source: LayerSource, @@ -54,6 +69,13 @@ impl ToolFilter { }); } + /// Each layer's source and raw patterns, in application order. + pub fn layers(&self) -> impl Iterator { + self.layers + .iter() + .map(|layer| (&layer.source, layer.raw.as_slice())) + } + /// A tool is allowed iff it matches at least one pattern in every layer. pub fn allows(&self, tool: &str) -> bool { self.layers.iter().all(|layer| { @@ -66,7 +88,6 @@ 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 { diff --git a/src/config/request_context.rs b/src/config/request_context.rs index bbf711d..1ad5d46 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -29,7 +29,8 @@ use crate::function::{ }; use crate::mcp::{ CatalogItem, MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpAuthReason, McpAuthRequired, - McpServersConfig, is_auth_required_error, is_mcp_meta_function, mcp_meta_function_names, + McpServerFeatures, McpServersConfig, McpTransportType, is_auth_required_error, + is_mcp_meta_function, mcp_meta_function_names, }; use crate::rag::Rag; use crate::supervisor::Supervisor; @@ -207,7 +208,7 @@ fn complete_skills_with_descriptions(names: Vec) -> Vec<(String, Option< .collect() } -const SET_COMPLETION_KEYS: [&str; 26] = [ +const SET_COMPLETION_KEYS: [&str; 27] = [ "auto_continue", "continuation_prompt", "temperature", @@ -220,6 +221,7 @@ const SET_COMPLETION_KEYS: [&str; 26] = [ "inject_skill_instructions", "skill_instructions", "max_auto_continues", + "mcp_tools", "memory", "save_session", "compression_threshold", @@ -773,6 +775,133 @@ impl RequestContext { } } + pub async fn mcp_server_info(&self, name: &str) -> Result { + let Some(spec) = self + .app + .mcp_config + .as_ref() + .and_then(|config| config.mcp_servers.get(name)) + else { + bail!( + "MCP server '{name}' is not configured. Run `.list mcp-servers` to see what's available" + ); + }; + let Some(handle) = self.tool_scope.mcp_runtime.servers.get(name).cloned() else { + bail!("MCP server '{name}' is not running. Enable it with `.mcp enable {name}`."); + }; + + let transport = match spec.transport_type { + McpTransportType::Stdio => "stdio", + McpTransportType::Http => "http", + McpTransportType::Sse => "sse", + }; + let info = handle.peer_info(); + let features = + McpServerFeatures::from_capabilities(name, info.as_ref().map(|i| &i.capabilities)); + let capabilities: Vec<&str> = [ + ("tools", features.tools), + ("resources", features.resources), + ("prompts", features.prompts), + ] + .iter() + .filter(|(_, supported)| *supported) + .map(|(label, _)| *label) + .collect(); + + const INFO_LABEL_WIDTH: usize = 15; + let mut out = String::new(); + out.push_str(&format!( + "{: = filter + .map(|f| { + f.layers() + .map(|(source, patterns)| (format!("{source}:"), patterns.join(" | "))) + .collect() + }) + .unwrap_or_default(); + if layers.is_empty() { + out.push_str(&format!( + "{: = tools.iter().map(|tool| tool.name.to_string()).collect(); + names.sort_unstable(); + let allowed = names + .iter() + .filter(|tool| filter.is_none_or(|f| f.allows(tool))) + .count(); + out.push_str(&format!( + "\ntools ({allowed} allowed / {} total)\n", + names.len() + )); + let name_width = names + .iter() + .map(|tool| tool.chars().count()) + .max() + .unwrap_or_default(); + for tool in &names { + let explained = filter.map(|f| f.allows_explain(tool)); + match explained { + None => out.push_str(&format!(" ✓ {tool}\n")), + Some(Ok(matches)) => { + let chain: Vec = matches + .iter() + .map(|(source, pattern)| format!("{pattern} ({})", source.short_label())) + .collect(); + if chain.is_empty() { + out.push_str(&format!(" ✓ {tool}\n")); + } else { + out.push_str(&format!(" ✓ {tool: out.push_str(&format!( + " ✗ {tool: