feat(mcp): enforce per-server tool allowlists across runtime, jobs, agents, and graph nodes
This commit is contained in:
+108
-17
@@ -8,7 +8,6 @@ use std::fmt;
|
|||||||
|
|
||||||
/// The configuration level that contributed a layer of tool patterns for an
|
/// The configuration level that contributed a layer of tool patterns for an
|
||||||
/// MCP server, as rendered in diagnostics.
|
/// MCP server, as rendered in diagnostics.
|
||||||
#[allow(dead_code)]
|
|
||||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||||
pub enum LayerSource {
|
pub enum LayerSource {
|
||||||
Global,
|
Global,
|
||||||
@@ -34,21 +33,18 @@ impl fmt::Display for LayerSource {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[allow(dead_code)]
|
#[derive(Debug, Clone)]
|
||||||
#[derive(Debug)]
|
|
||||||
pub struct CompiledPatterns {
|
pub struct CompiledPatterns {
|
||||||
source: LayerSource,
|
source: LayerSource,
|
||||||
raw: Vec<String>,
|
raw: Vec<String>,
|
||||||
regexes: Vec<Regex>,
|
regexes: Vec<Regex>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[allow(dead_code)]
|
#[derive(Debug, Clone, Default)]
|
||||||
#[derive(Debug, Default)]
|
|
||||||
pub struct ToolFilter {
|
pub struct ToolFilter {
|
||||||
layers: Vec<CompiledPatterns>,
|
layers: Vec<CompiledPatterns>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[allow(dead_code)]
|
|
||||||
impl ToolFilter {
|
impl ToolFilter {
|
||||||
pub fn push_layer(&mut self, source: LayerSource, patterns: &[String]) {
|
pub fn push_layer(&mut self, source: LayerSource, patterns: &[String]) {
|
||||||
self.layers.push(CompiledPatterns {
|
self.layers.push(CompiledPatterns {
|
||||||
@@ -70,6 +66,7 @@ impl ToolFilter {
|
|||||||
|
|
||||||
/// The first matching raw pattern per layer, in layer order, or the
|
/// The first matching raw pattern per layer, in layer order, or the
|
||||||
/// source of the first layer with no match.
|
/// source of the first layer with no match.
|
||||||
|
#[allow(dead_code)]
|
||||||
pub fn allows_explain(&self, tool: &str) -> Result<Vec<(&LayerSource, &str)>, &LayerSource> {
|
pub fn allows_explain(&self, tool: &str) -> Result<Vec<(&LayerSource, &str)>, &LayerSource> {
|
||||||
let mut matched = Vec::with_capacity(self.layers.len());
|
let mut matched = Vec::with_capacity(self.layers.len());
|
||||||
for layer in &self.layers {
|
for layer in &self.layers {
|
||||||
@@ -86,11 +83,48 @@ impl ToolFilter {
|
|||||||
}
|
}
|
||||||
Ok(matched)
|
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)
|
/// Translates a glob pattern (`*` = any run of characters, `?` = exactly one)
|
||||||
/// into an anchored regex. Patterns that fail to compile match nothing.
|
/// into an anchored regex. Patterns that fail to compile match nothing.
|
||||||
#[allow(dead_code)]
|
|
||||||
fn compile_glob(pattern: &str) -> Regex {
|
fn compile_glob(pattern: &str) -> Regex {
|
||||||
let translated = format!(
|
let translated = format!(
|
||||||
"^{}$",
|
"^{}$",
|
||||||
@@ -104,22 +138,18 @@ fn compile_glob(pattern: &str) -> Regex {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
#[allow(dead_code)]
|
|
||||||
fn never_matching_regex() -> Regex {
|
fn never_matching_regex() -> Regex {
|
||||||
Regex::new("(?!)").expect("'(?!)' is a valid never-matching regex")
|
Regex::new("(?!)").expect("'(?!)' is a valid never-matching regex")
|
||||||
}
|
}
|
||||||
|
|
||||||
#[allow(dead_code)]
|
|
||||||
pub struct SkillMcpLayer {
|
pub struct SkillMcpLayer {
|
||||||
pub name: String,
|
pub name: String,
|
||||||
pub enabled_servers: Vec<String>,
|
pub enabled_servers: Vec<String>,
|
||||||
pub mcp_tools: IndexMap<String, Vec<String>>,
|
pub mcp_tools: IndexMap<String, Vec<String>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[allow(dead_code)]
|
|
||||||
pub struct McpToolPolicy;
|
pub struct McpToolPolicy;
|
||||||
|
|
||||||
#[allow(dead_code)]
|
|
||||||
impl McpToolPolicy {
|
impl McpToolPolicy {
|
||||||
#[allow(clippy::too_many_arguments)]
|
#[allow(clippy::too_many_arguments)]
|
||||||
pub fn effective(
|
pub fn effective(
|
||||||
@@ -208,7 +238,6 @@ impl McpToolPolicy {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[allow(dead_code)]
|
|
||||||
fn push_level(
|
fn push_level(
|
||||||
filters: &mut HashMap<String, ToolFilter>,
|
filters: &mut HashMap<String, ToolFilter>,
|
||||||
mcp_config: &McpServersConfig,
|
mcp_config: &McpServersConfig,
|
||||||
@@ -234,7 +263,6 @@ fn push_level(
|
|||||||
/// a key that is an alias expands to every configured id in its
|
/// a key that is an alias expands to every configured id in its
|
||||||
/// comma-separated value; anything else is dropped. Keys expanding to the
|
/// comma-separated value; anything else is dropped. Keys expanding to the
|
||||||
/// same server merge their pattern lists.
|
/// same server merge their pattern lists.
|
||||||
#[allow(dead_code)]
|
|
||||||
fn expand_server_keys(
|
fn expand_server_keys(
|
||||||
mcp_config: &McpServersConfig,
|
mcp_config: &McpServersConfig,
|
||||||
aliases: &IndexMap<String, String>,
|
aliases: &IndexMap<String, String>,
|
||||||
@@ -248,11 +276,11 @@ fn expand_server_keys(
|
|||||||
.entry(key.to_string())
|
.entry(key.to_string())
|
||||||
.or_default()
|
.or_default()
|
||||||
.extend(patterns.iter().cloned());
|
.extend(patterns.iter().cloned());
|
||||||
} else if let Some(mapped) = aliases.get(key) {
|
} else {
|
||||||
for mapped_id in mapped.split(',').map(str::trim) {
|
for mapped_id in expand_mcp_server_alias(aliases, key) {
|
||||||
if mcp_config.mcp_servers.contains_key(mapped_id) {
|
if mcp_config.mcp_servers.contains_key(&mapped_id) {
|
||||||
expanded
|
expanded
|
||||||
.entry(mapped_id.to_string())
|
.entry(mapped_id)
|
||||||
.or_default()
|
.or_default()
|
||||||
.extend(patterns.iter().cloned());
|
.extend(patterns.iter().cloned());
|
||||||
}
|
}
|
||||||
@@ -262,6 +290,25 @@ fn expand_server_keys(
|
|||||||
expanded
|
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<String, String>,
|
||||||
|
key: &str,
|
||||||
|
) -> Vec<String> {
|
||||||
|
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)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -756,4 +803,48 @@ mod tests {
|
|||||||
);
|
);
|
||||||
assert_eq!(LayerSource::Node("n1".into()).to_string(), "node (n1)");
|
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());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -42,6 +42,9 @@ pub use self::install_remote::{
|
|||||||
pub use self::macro_policy::{
|
pub use self::macro_policy::{
|
||||||
MacroAllowlistLevel, MacroPolicy, MacroSource, MacroState, RESERVED_MACRO_NAMES, ResolvedMacro,
|
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)]
|
#[allow(unused_imports)]
|
||||||
pub use self::request_context::{
|
pub use self::request_context::{
|
||||||
RenderMode, RequestContext, effective_max_concurrent_jobs, jobs_enabled,
|
RenderMode, RequestContext, effective_max_concurrent_jobs, jobs_enabled,
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
use super::bundles::BundleStore;
|
use super::bundles::BundleStore;
|
||||||
|
use super::mcp_tool_policy::{McpToolPolicy, SkillMcpLayer, ToolFilter, expand_mcp_server_alias};
|
||||||
use super::rag_cache::{RagCache, RagKey};
|
use super::rag_cache::{RagCache, RagKey};
|
||||||
use super::session::{INTERRUPTED_RESPONSE_TEXT, Session};
|
use super::session::{INTERRUPTED_RESPONSE_TEXT, Session};
|
||||||
use super::skill::{SKILL_SCAFFOLD, Skill};
|
use super::skill::{SKILL_SCAFFOLD, Skill};
|
||||||
@@ -51,7 +52,6 @@ use crate::graph;
|
|||||||
use anyhow::{Context, Error, Result, bail};
|
use anyhow::{Context, Error, Result, bail};
|
||||||
use colored::Colorize;
|
use colored::Colorize;
|
||||||
use gman::providers::SupportedProvider;
|
use gman::providers::SupportedProvider;
|
||||||
#[cfg(test)]
|
|
||||||
use indexmap::IndexMap;
|
use indexmap::IndexMap;
|
||||||
use indoc::formatdoc;
|
use indoc::formatdoc;
|
||||||
use inquire::{Confirm, MultiSelect, Text, list_option::ListOption, validator::Validation};
|
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()) {
|
for item in enabled_mcp_servers.iter().map(|s| s.trim()) {
|
||||||
if mcp_config.mcp_servers.contains_key(item) {
|
if mcp_config.mcp_servers.contains_key(item) {
|
||||||
ids.push(item.to_string());
|
ids.push(item.to_string());
|
||||||
} else if let Some(mapped) = app.mapping_mcp_servers.get(item) {
|
} else {
|
||||||
for mapped_id in mapped.split(',').map(|s| s.trim()) {
|
for mapped_id in expand_mcp_server_alias(&app.mapping_mcp_servers, item) {
|
||||||
if mcp_config.mcp_servers.contains_key(mapped_id) {
|
if mcp_config.mcp_servers.contains_key(&mapped_id) {
|
||||||
ids.push(mapped_id.to_string());
|
ids.push(mapped_id);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -333,6 +333,10 @@ pub struct RequestContext {
|
|||||||
/// context owns every job in its supervisor.
|
/// context owns every job in its supervisor.
|
||||||
pub node_job_scope: Option<Vec<String>>,
|
pub node_job_scope: Option<Vec<String>>,
|
||||||
|
|
||||||
|
/// 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<String, Vec<String>>)>,
|
||||||
|
|
||||||
pub supervisor: Option<Arc<RwLock<Supervisor>>>,
|
pub supervisor: Option<Arc<RwLock<Supervisor>>>,
|
||||||
pub parent_supervisor: Option<Arc<RwLock<Supervisor>>>,
|
pub parent_supervisor: Option<Arc<RwLock<Supervisor>>>,
|
||||||
pub self_agent_id: Option<String>,
|
pub self_agent_id: Option<String>,
|
||||||
@@ -369,6 +373,7 @@ impl RequestContext {
|
|||||||
tool_scope: ToolScope::default(),
|
tool_scope: ToolScope::default(),
|
||||||
declared_function_names: Default::default(),
|
declared_function_names: Default::default(),
|
||||||
node_job_scope: None,
|
node_job_scope: None,
|
||||||
|
active_node_mcp_tools: None,
|
||||||
supervisor: None,
|
supervisor: None,
|
||||||
parent_supervisor: None,
|
parent_supervisor: None,
|
||||||
self_agent_id: None,
|
self_agent_id: None,
|
||||||
@@ -410,7 +415,7 @@ impl RequestContext {
|
|||||||
mcp_runtime.sync_from_registry(registry);
|
mcp_runtime.sync_from_registry(registry);
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(Self {
|
let mut ctx = Self {
|
||||||
app,
|
app,
|
||||||
macro_flag: false,
|
macro_flag: false,
|
||||||
macro_non_isolated: false,
|
macro_non_isolated: false,
|
||||||
@@ -431,6 +436,7 @@ impl RequestContext {
|
|||||||
},
|
},
|
||||||
declared_function_names: Default::default(),
|
declared_function_names: Default::default(),
|
||||||
node_job_scope: None,
|
node_job_scope: None,
|
||||||
|
active_node_mcp_tools: None,
|
||||||
supervisor: None,
|
supervisor: None,
|
||||||
parent_supervisor: None,
|
parent_supervisor: None,
|
||||||
self_agent_id: None,
|
self_agent_id: None,
|
||||||
@@ -445,7 +451,9 @@ impl RequestContext {
|
|||||||
last_continuation_response: None,
|
last_continuation_response: None,
|
||||||
pending_prefill: None,
|
pending_prefill: None,
|
||||||
render_mode: RenderMode::default(),
|
render_mode: RenderMode::default(),
|
||||||
})
|
};
|
||||||
|
ctx.refresh_mcp_tool_filters();
|
||||||
|
Ok(ctx)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Forks the context for one parallel branch of a graph super-step.
|
/// Forks the context for one parallel branch of a graph super-step.
|
||||||
@@ -480,6 +488,7 @@ impl RequestContext {
|
|||||||
tool_scope: self.tool_scope.clone(),
|
tool_scope: self.tool_scope.clone(),
|
||||||
declared_function_names: self.declared_function_names.clone(),
|
declared_function_names: self.declared_function_names.clone(),
|
||||||
node_job_scope: None,
|
node_job_scope: None,
|
||||||
|
active_node_mcp_tools: self.active_node_mcp_tools.clone(),
|
||||||
supervisor: self.supervisor.clone(),
|
supervisor: self.supervisor.clone(),
|
||||||
parent_supervisor: self.parent_supervisor.clone(),
|
parent_supervisor: self.parent_supervisor.clone(),
|
||||||
self_agent_id: self.self_agent_id.clone(),
|
self_agent_id: self.self_agent_id.clone(),
|
||||||
@@ -527,6 +536,7 @@ impl RequestContext {
|
|||||||
},
|
},
|
||||||
declared_function_names: Default::default(),
|
declared_function_names: Default::default(),
|
||||||
node_job_scope: None,
|
node_job_scope: None,
|
||||||
|
active_node_mcp_tools: None,
|
||||||
supervisor: None,
|
supervisor: None,
|
||||||
parent_supervisor: parent.supervisor.clone(),
|
parent_supervisor: parent.supervisor.clone(),
|
||||||
self_agent_id: Some(self_agent_id),
|
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<()> {
|
pub fn use_temp_role(&mut self, _app: &AppConfig, prompt: &str) -> Result<()> {
|
||||||
let mut role = Role::new(TEMP_ROLE_NAME, prompt);
|
let mut role = Role::new(TEMP_ROLE_NAME, prompt);
|
||||||
role.set_model(self.current_model().clone());
|
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<()> {
|
pub fn edit_config(&self) -> Result<()> {
|
||||||
@@ -3964,9 +3976,70 @@ impl RequestContext {
|
|||||||
mcp_runtime,
|
mcp_runtime,
|
||||||
tool_tracker,
|
tool_tracker,
|
||||||
};
|
};
|
||||||
|
self.refresh_mcp_tool_filters();
|
||||||
Ok(())
|
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<String, ToolFilter> {
|
||||||
|
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<SkillMcpLayer> = 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<()> {
|
pub async fn refresh_tool_scope(&mut self, abort_signal: AbortSignal) -> Result<()> {
|
||||||
let app = (*self.app.config).clone();
|
let app = (*self.app.config).clone();
|
||||||
let base_mcps = if app.mcp_server_support {
|
let base_mcps = if app.mcp_server_support {
|
||||||
@@ -4141,6 +4214,7 @@ impl RequestContext {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
self.session = session;
|
self.session = session;
|
||||||
|
self.refresh_mcp_tool_filters();
|
||||||
self.init_agent_session_variables(new_session)?;
|
self.init_agent_session_variables(new_session)?;
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
@@ -4261,6 +4335,7 @@ impl RequestContext {
|
|||||||
.is_some()
|
.is_some()
|
||||||
.then(|| RagKey::Agent(agent.name().to_string()));
|
.then(|| RagKey::Agent(agent.name().to_string()));
|
||||||
self.agent = Some(agent);
|
self.agent = Some(agent);
|
||||||
|
self.refresh_mcp_tool_filters();
|
||||||
if let Some(old) = self.supervisor.as_ref() {
|
if let Some(old) = self.supervisor.as_ref() {
|
||||||
old.read().cancel_recursive();
|
old.read().cancel_recursive();
|
||||||
}
|
}
|
||||||
@@ -8012,4 +8087,191 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn mcp_app_state(servers: &[&str]) -> Arc<AppState> {
|
||||||
|
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<String, Vec<String>> {
|
||||||
|
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"
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -34,6 +34,10 @@ impl SkillRegistry {
|
|||||||
self.loaded.keys().cloned().collect()
|
self.loaded.keys().cloned().collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn loaded_skills(&self) -> impl Iterator<Item = &Skill> {
|
||||||
|
self.loaded.values()
|
||||||
|
}
|
||||||
|
|
||||||
pub fn loaded_mcp_servers(&self) -> BTreeSet<String> {
|
pub fn loaded_mcp_servers(&self) -> BTreeSet<String> {
|
||||||
let mut out = BTreeSet::new();
|
let mut out = BTreeSet::new();
|
||||||
for skill in self.loaded.values() {
|
for skill in self.loaded.values() {
|
||||||
|
|||||||
+182
-2
@@ -1,3 +1,4 @@
|
|||||||
|
use super::mcp_tool_policy::ToolFilter;
|
||||||
use crate::function::{Functions, ToolCallTracker};
|
use crate::function::{Functions, ToolCallTracker};
|
||||||
use crate::mcp::{CatalogItem, CatalogItemKind, ConnectedServer, McpRegistry, McpServerFeatures};
|
use crate::mcp::{CatalogItem, CatalogItemKind, ConnectedServer, McpRegistry, McpServerFeatures};
|
||||||
|
|
||||||
@@ -44,6 +45,8 @@ pub enum McpPromptCompletion {
|
|||||||
#[derive(Default, Clone)]
|
#[derive(Default, Clone)]
|
||||||
pub struct McpRuntime {
|
pub struct McpRuntime {
|
||||||
pub servers: HashMap<String, Arc<ConnectedServer>>,
|
pub servers: HashMap<String, Arc<ConnectedServer>>,
|
||||||
|
/// Per-server effective tool allowlists; a server absent here is unfiltered.
|
||||||
|
pub tool_filters: HashMap<String, ToolFilter>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl McpRuntime {
|
impl McpRuntime {
|
||||||
@@ -100,12 +103,24 @@ impl McpRuntime {
|
|||||||
|
|
||||||
if features.tools {
|
if features.tools {
|
||||||
match server_handle.list_all_tools().await {
|
match server_handle.list_all_tools().await {
|
||||||
Ok(tools) => merge_catalog_items(
|
Ok(mut tools) => {
|
||||||
|
if let Some(filter) = self.tool_filters.get(server) {
|
||||||
|
let advertised: Vec<String> =
|
||||||
|
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,
|
&mut items,
|
||||||
tools
|
tools
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|tool| tool_catalog_item(server, tool)),
|
.map(|tool| tool_catalog_item(server, tool)),
|
||||||
),
|
)
|
||||||
|
}
|
||||||
Err(e) => warn!("Failed to list tools on MCP server {server}: {e}"),
|
Err(e) => warn!("Failed to list tools on MCP server {server}: {e}"),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -195,6 +210,11 @@ impl McpRuntime {
|
|||||||
|
|
||||||
match kind {
|
match kind {
|
||||||
"tool" => {
|
"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
|
let tool_schema = server_handle
|
||||||
.list_all_tools()
|
.list_all_tools()
|
||||||
.await?
|
.await?
|
||||||
@@ -297,6 +317,12 @@ impl McpRuntime {
|
|||||||
.cloned()
|
.cloned()
|
||||||
.with_context(|| format!("Invoked MCP server does not exist: {server}"))?;
|
.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());
|
let mut request = CallToolRequestParams::new(tool.to_owned());
|
||||||
request.arguments = arguments.as_object().cloned();
|
request.arguments = arguments.as_object().cloned();
|
||||||
|
|
||||||
@@ -869,6 +895,7 @@ mod tests {
|
|||||||
FIXTURE_ANNOTATED_URI, FixtureServer, add_fixture_server, fixture_runtime,
|
FIXTURE_ANNOTATED_URI, FixtureServer, add_fixture_server, fixture_runtime,
|
||||||
};
|
};
|
||||||
use super::*;
|
use super::*;
|
||||||
|
use crate::config::mcp_tool_policy::LayerSource;
|
||||||
use crate::function::ToolCall;
|
use crate::function::ToolCall;
|
||||||
use log::{Level, LevelFilter, Log, Metadata, Record};
|
use log::{Level, LevelFilter, Log, Metadata, Record};
|
||||||
use std::sync::atomic::Ordering;
|
use std::sync::atomic::Ordering;
|
||||||
@@ -1643,4 +1670,157 @@ mod tests {
|
|||||||
assert_eq!(prompt, "summarize");
|
assert_eq!(prompt, "summarize");
|
||||||
assert_eq!(typed_keys, vec!["path".to_string()]);
|
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::<Vec<_>>(),
|
||||||
|
);
|
||||||
|
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:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -686,6 +686,7 @@ pub async fn run_agent_for_graph(
|
|||||||
sync_agent_functions_to_ctx(&mut child_ctx)?;
|
sync_agent_functions_to_ctx(&mut child_ctx)?;
|
||||||
} else {
|
} else {
|
||||||
populate_agent_mcp_runtime(&mut child_ctx, &agent_mcp_servers).await?;
|
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)?;
|
sync_agent_functions_to_ctx(&mut child_ctx)?;
|
||||||
child_ctx.init_agent_shared_variables()?;
|
child_ctx.init_agent_shared_variables()?;
|
||||||
}
|
}
|
||||||
@@ -869,6 +870,7 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
|
|||||||
sync_agent_functions_to_ctx(&mut child_ctx)?;
|
sync_agent_functions_to_ctx(&mut child_ctx)?;
|
||||||
} else {
|
} else {
|
||||||
populate_agent_mcp_runtime(&mut child_ctx, &agent_mcp_servers).await?;
|
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)?;
|
sync_agent_functions_to_ctx(&mut child_ctx)?;
|
||||||
child_ctx.init_agent_shared_variables()?;
|
child_ctx.init_agent_shared_variables()?;
|
||||||
}
|
}
|
||||||
@@ -1604,6 +1606,7 @@ mod tests {
|
|||||||
use crate::config::test_fixtures::{FixtureServer, fixture_runtime};
|
use crate::config::test_fixtures::{FixtureServer, fixture_runtime};
|
||||||
use crate::config::{AgentConfig, AppState, WorkingMode};
|
use crate::config::{AgentConfig, AppState, WorkingMode};
|
||||||
use crate::function::jobs::RingBuf;
|
use crate::function::jobs::RingBuf;
|
||||||
|
use crate::mcp::{McpServer, McpServersConfig, McpTransportType};
|
||||||
use crate::supervisor::escalation::{EscalationQueue, EscalationRequest};
|
use crate::supervisor::escalation::{EscalationQueue, EscalationRequest};
|
||||||
use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus};
|
use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus};
|
||||||
use parking_lot::Mutex;
|
use parking_lot::Mutex;
|
||||||
@@ -1791,6 +1794,61 @@ mod tests {
|
|||||||
assert!(!functions.contains("mcp_invoke_fixture"));
|
assert!(!functions.contains("mcp_invoke_fixture"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn app_state_with_fixture_mcp_config() -> Arc<AppState> {
|
||||||
|
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]
|
#[test]
|
||||||
fn handle_list_running_empty_supervisor() {
|
fn handle_list_running_empty_supervisor() {
|
||||||
let mut ctx = ctx_with_supervisor(4, 3);
|
let mut ctx = ctx_with_supervisor(4, 3);
|
||||||
|
|||||||
@@ -457,6 +457,11 @@ async fn handle_start(ctx: &mut RequestContext, args: &Value) -> Result<Value> {
|
|||||||
.unwrap_or_else(|| json!({}));
|
.unwrap_or_else(|| json!({}));
|
||||||
let mut mcp_runtime = McpRuntime::new();
|
let mut mcp_runtime = McpRuntime::new();
|
||||||
mcp_runtime.insert(server.clone(), Arc::clone(server_handle));
|
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 {
|
let job_ctx = JobCtx {
|
||||||
mcp_runtime,
|
mcp_runtime,
|
||||||
current_depth: ctx.current_depth,
|
current_depth: ctx.current_depth,
|
||||||
@@ -1263,7 +1268,9 @@ fn tail_chars(text: &str, max_chars: usize) -> Option<String> {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
use crate::config::test_fixtures::{FixtureServer, fixture_runtime};
|
||||||
use crate::config::{AppConfig, AppState, WorkingMode};
|
use crate::config::{AppConfig, AppState, WorkingMode};
|
||||||
|
use crate::config::{LayerSource, ToolFilter};
|
||||||
use crate::function::agents::{
|
use crate::function::agents::{
|
||||||
GuardrailAction, check_pending_tasks_guardrail, handle_agent_tool,
|
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]
|
#[test]
|
||||||
fn handle_start_rejects_unconnected_mcp_server() {
|
fn handle_start_rejects_unconnected_mcp_server() {
|
||||||
let mut ctx = plain_ctx();
|
let mut ctx = plain_ctx();
|
||||||
|
|||||||
@@ -424,7 +424,7 @@ async fn step(
|
|||||||
Ok(StepResult::Continue(vec![next]))
|
Ok(StepResult::Continue(vec![next]))
|
||||||
}
|
}
|
||||||
NodeType::Llm(llm_node) => {
|
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 {
|
let targets = match outcome {
|
||||||
LlmExecutionOutcome::Continue => static_next_targets(node, current, "llm")?,
|
LlmExecutionOutcome::Continue => static_next_targets(node, current, "llm")?,
|
||||||
LlmExecutionOutcome::FellBack(target) => vec![target],
|
LlmExecutionOutcome::FellBack(target) => vec![target],
|
||||||
|
|||||||
+12
-1
@@ -30,11 +30,12 @@ pub struct LlmNodeExecutor;
|
|||||||
|
|
||||||
impl LlmNodeExecutor {
|
impl LlmNodeExecutor {
|
||||||
pub(super) async fn execute(
|
pub(super) async fn execute(
|
||||||
|
node_id: &str,
|
||||||
node: &LlmNode,
|
node: &LlmNode,
|
||||||
state_manager: &mut StateManager,
|
state_manager: &mut StateManager,
|
||||||
parent_ctx: &mut RequestContext,
|
parent_ctx: &mut RequestContext,
|
||||||
) -> Result<LlmExecutionOutcome> {
|
) -> Result<LlmExecutionOutcome> {
|
||||||
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 {
|
let (output, failure_reason) = match result {
|
||||||
Ok(raw) => match &node.output_schema {
|
Ok(raw) => match &node.output_schema {
|
||||||
Some(schema) => match structured::extract(&raw, schema, parent_ctx).await {
|
Some(schema) => match structured::extract(&raw, schema, parent_ctx).await {
|
||||||
@@ -79,6 +80,7 @@ fn outcome_from(
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn run(
|
async fn run(
|
||||||
|
node_id: &str,
|
||||||
node: &LlmNode,
|
node: &LlmNode,
|
||||||
state_manager: &mut StateManager,
|
state_manager: &mut StateManager,
|
||||||
parent_ctx: &mut RequestContext,
|
parent_ctx: &mut RequestContext,
|
||||||
@@ -177,6 +179,13 @@ async fn run(
|
|||||||
// Jobs are node-local: everything job__start registers while this node
|
// Jobs are node-local: everything job__start registers while this node
|
||||||
// runs is recorded here and reaped on every exit path below.
|
// runs is recorded here and reaped on every exit path below.
|
||||||
let saved_job_scope = parent_ctx.node_job_scope.replace(Vec::new());
|
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 {
|
let result = match node.timeout {
|
||||||
Some(secs) => match timeout(
|
Some(secs) => match timeout(
|
||||||
Duration::from_secs(secs),
|
Duration::from_secs(secs),
|
||||||
@@ -193,6 +202,8 @@ async fn run(
|
|||||||
let node_jobs =
|
let node_jobs =
|
||||||
std::mem::replace(&mut parent_ctx.node_job_scope, saved_job_scope).unwrap_or_default();
|
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;
|
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);
|
restore_agent_skill_policy(parent_ctx, saved_agent_skill_state);
|
||||||
result
|
result
|
||||||
}
|
}
|
||||||
|
|||||||
+4
-2
@@ -85,9 +85,11 @@ impl MapNodeExecutor {
|
|||||||
let mut ctx = sub_ctx;
|
let mut ctx = sub_ctx;
|
||||||
|
|
||||||
let exec_result: Result<()> = match &branch_clone.node_type {
|
let exec_result: Result<()> = match &branch_clone.node_type {
|
||||||
NodeType::Llm(n) => LlmNodeExecutor::execute(n, &mut state, &mut ctx)
|
NodeType::Llm(n) => {
|
||||||
|
LlmNodeExecutor::execute(&branch_clone.id, n, &mut state, &mut ctx)
|
||||||
.await
|
.await
|
||||||
.map(|_| ()),
|
.map(|_| ())
|
||||||
|
}
|
||||||
NodeType::Agent(n) => AgentNodeExecutor::execute(n, &mut state, &mut ctx)
|
NodeType::Agent(n) => AgentNodeExecutor::execute(n, &mut state, &mut ctx)
|
||||||
.await
|
.await
|
||||||
.map(|_| ()),
|
.map(|_| ()),
|
||||||
|
|||||||
@@ -215,11 +215,7 @@ impl GraphValidator {
|
|||||||
};
|
};
|
||||||
|
|
||||||
let expand_alias = |name: &str| {
|
let expand_alias = |name: &str| {
|
||||||
ctx.app_config
|
crate::config::expand_mcp_server_alias(&ctx.app_config.mapping_mcp_servers, name)
|
||||||
.mapping_mcp_servers
|
|
||||||
.get(name)
|
|
||||||
.into_iter()
|
|
||||||
.flat_map(|mapped| mapped.split(',').map(|s| s.trim().to_string()))
|
|
||||||
};
|
};
|
||||||
let mut enabled_servers: HashSet<String> = ctx.mcp_servers.clone();
|
let mut enabled_servers: HashSet<String> = ctx.mcp_servers.clone();
|
||||||
for server in &ctx.mcp_servers {
|
for server in &ctx.mcp_servers {
|
||||||
@@ -255,7 +251,9 @@ impl GraphValidator {
|
|||||||
{
|
{
|
||||||
for key in mcp_tools.keys() {
|
for key in mcp_tools.keys() {
|
||||||
let enabled = enabled_servers.contains(key)
|
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 {
|
if !enabled {
|
||||||
result.error(ValidationError::with_node(
|
result.error(ValidationError::with_node(
|
||||||
node_id,
|
node_id,
|
||||||
|
|||||||
Reference in New Issue
Block a user