feat(mcp): enforce per-server tool allowlists across runtime, jobs, agents, and graph nodes

This commit is contained in:
Dark-Alex-17 committed 2026-08-27 14:43:18 -06:00
1 parent a92ebb5b94
commit 92d9464b38
11 files changed
+695 -42

No files matched your search

+1 -1
View File
@@ -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],
+12 -1
View File
@@ -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<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 {
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
}
+5 -3
View File
@@ -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(|_| ()),
+4 -6
View File
@@ -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<String> = 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,