feat(mcp): enforce per-server tool allowlists across runtime, jobs, agents, and graph nodes
This commit is contained in:
11 files changed
+695
-42
No files matched your search
@@ -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
@@ -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
@@ -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(|_| ()),
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in new issue
Block a user