From 5b9852d80e859ff2d4c51f0c54caee11af6577d3 Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Tue, 15 Sep 2026 11:13:12 -0600 Subject: [PATCH] fix: Require injection of user interaction functions for LLM nodes to be explicit --- src/config/request_context.rs | 232 +++++++++++++++++++++++++++++++++- src/graph/llm.rs | 83 ++++++++++-- src/graph/validator.rs | 108 ++++++++++++++++ 3 files changed, 410 insertions(+), 13 deletions(-) diff --git a/src/config/request_context.rs b/src/config/request_context.rs index 99714ca..4a56409 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -327,6 +327,7 @@ pub struct RequestContext { /// 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 in_graph_llm_node: bool, pub supervisor: Option>>, pub parent_supervisor: Option>>, @@ -378,6 +379,7 @@ impl RequestContext { declared_function_names: Default::default(), node_job_scope: None, active_node_mcp_tools: None, + in_graph_llm_node: false, supervisor: None, parent_supervisor: None, self_agent_id: None, @@ -448,6 +450,7 @@ impl RequestContext { declared_function_names: Default::default(), node_job_scope: None, active_node_mcp_tools: None, + in_graph_llm_node: false, supervisor: None, parent_supervisor: None, self_agent_id: None, @@ -507,6 +510,7 @@ impl RequestContext { declared_function_names: self.declared_function_names.clone(), node_job_scope: None, active_node_mcp_tools: self.active_node_mcp_tools.clone(), + in_graph_llm_node: self.in_graph_llm_node, supervisor: self.supervisor.clone(), parent_supervisor: self.parent_supervisor.clone(), self_agent_id: self.self_agent_id.clone(), @@ -562,6 +566,7 @@ impl RequestContext { declared_function_names: Default::default(), node_job_scope: None, active_node_mcp_tools: None, + in_graph_llm_node: false, supervisor: None, parent_supervisor: parent.supervisor.clone(), self_agent_id: Some(self_agent_id), @@ -2485,10 +2490,11 @@ impl RequestContext { .declarations() .iter() .filter(|v| { - (v.name.starts_with(USER_FUNCTION_PREFIX) + ((!self.in_graph_llm_node && v.name.starts_with(USER_FUNCTION_PREFIX)) || (!matches!(role.skills_enabled(), Some(false)) && v.name.starts_with(SKILL_FUNCTION_PREFIX)) - || (self.auto_continue_config().enabled + || (!self.in_graph_llm_node + && self.auto_continue_config().enabled && v.name.starts_with(TODO_FUNCTION_PREFIX)) || v.name.starts_with(RAG_FUNCTION_PREFIX) || v.name.starts_with(JOB_FUNCTION_PREFIX)) @@ -2513,8 +2519,9 @@ impl RequestContext { tool_names.contains(&v.name) || (!matches!(agent.skills_enabled(), Some(false)) && v.name.starts_with(SKILL_FUNCTION_PREFIX)) - || v.name.starts_with(USER_FUNCTION_PREFIX) - || v.name.starts_with(TODO_FUNCTION_PREFIX) + || (!self.in_graph_llm_node + && (v.name.starts_with(USER_FUNCTION_PREFIX) + || v.name.starts_with(TODO_FUNCTION_PREFIX))) || v.name.starts_with(AGENT_FUNCTION_PREFIX) || v.name.starts_with(MEMORY_FUNCTION_PREFIX) || v.name.starts_with(RAG_FUNCTION_PREFIX) @@ -6436,6 +6443,19 @@ mod tests { )); } + #[test] + fn fork_for_branch_copies_in_graph_llm_node() { + let mut ctx = create_test_ctx(); + ctx.in_graph_llm_node = true; + + let branch = ctx.fork_for_branch(); + + assert!( + branch.in_graph_llm_node, + "parallel map branches inside a node must keep the node's tool gating" + ); + } + fn app_state_with_mcp_config(mcp_server_support: bool, server_names: &[&str]) -> Arc { app_state_with_mcp_command(mcp_server_support, server_names, "echo") } @@ -7033,6 +7053,210 @@ mod tests { ); } + #[test] + #[serial] + fn select_functions_graph_llm_node_gates_user_and_todo_tools_under_agent_filter() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = create_test_ctx(); + let app = ctx.app.config.clone(); + let agent_name = format!( + "test_node_gate_agent_{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + ); + let agent_dir = paths::agent_data_dir(&agent_name); + create_dir_all(&agent_dir).unwrap(); + write( + agent_dir.join("config.yaml"), + format!( + "name: {agent_name}\ninstructions: hi\nauto_continue: true\ncan_spawn_agents: true\n" + ), + ) + .unwrap(); + + let abort = utils::create_abort_signal(); + run_async(ctx.use_agent(&app, &agent_name, None, abort)).unwrap(); + + let mut role = Role::new("r", "p"); + role.set_enabled_tools(Some(vec!["foo".to_string()])); + + ctx.in_graph_llm_node = true; + let fns = ctx.select_functions(&role).unwrap(); + let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect(); + assert!( + !names.iter().any(|n| n.starts_with("user__")), + "user__ tools must not leak into a graph llm node, got: {names:?}" + ); + assert!( + !names.iter().any(|n| n.starts_with("todo__")), + "todo__ tools must not leak into a graph llm node, got: {names:?}" + ); + assert!( + names.contains(&"agent__spawn"), + "agent__ tools keep implicit availability inside a node, got: {names:?}" + ); + + ctx.in_graph_llm_node = false; + let fns = ctx.select_functions(&role).unwrap(); + let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect(); + assert!( + names.contains(&"user__select"), + "outside a node, subagent escalation via user__ must keep working, got: {names:?}" + ); + assert!( + names.contains(&"todo__init"), + "outside a node, todo__ tools must survive the agent filter, got: {names:?}" + ); + } + + #[test] + #[serial] + fn select_functions_graph_llm_node_user_tool_opt_in_is_exact() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = create_test_ctx(); + let app = ctx.app.config.clone(); + let agent_name = format!( + "test_node_optin_agent_{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + ); + let agent_dir = paths::agent_data_dir(&agent_name); + create_dir_all(&agent_dir).unwrap(); + write( + agent_dir.join("config.yaml"), + format!("name: {agent_name}\ninstructions: hi\n"), + ) + .unwrap(); + + let abort = utils::create_abort_signal(); + run_async(ctx.use_agent(&app, &agent_name, None, abort)).unwrap(); + + let mut role = Role::new("r", "p"); + role.set_enabled_tools(Some(vec!["user__select".to_string()])); + + ctx.in_graph_llm_node = true; + let fns = ctx.select_functions(&role).unwrap(); + let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect(); + assert!( + names.contains(&"user__select"), + "an explicitly listed user tool must be offered, got: {names:?}" + ); + assert!( + !names.contains(&"user__confirm"), + "opt-in must not drag in sibling user tools, got: {names:?}" + ); + } + + #[test] + #[serial] + fn select_functions_graph_llm_node_todo_tool_opt_in_is_exact() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = create_test_ctx(); + let app = ctx.app.config.clone(); + let agent_name = format!( + "test_node_todo_optin_agent_{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + ); + let agent_dir = paths::agent_data_dir(&agent_name); + create_dir_all(&agent_dir).unwrap(); + write( + agent_dir.join("config.yaml"), + format!("name: {agent_name}\ninstructions: hi\nauto_continue: true\n"), + ) + .unwrap(); + + let abort = utils::create_abort_signal(); + run_async(ctx.use_agent(&app, &agent_name, None, abort)).unwrap(); + + let mut role = Role::new("r", "p"); + role.set_enabled_tools(Some(vec!["todo__init".to_string()])); + + ctx.in_graph_llm_node = true; + let fns = ctx.select_functions(&role).unwrap(); + let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect(); + assert!( + names.contains(&"todo__init"), + "an explicitly listed todo tool must be offered, got: {names:?}" + ); + assert!( + !names.contains(&"todo__pause"), + "opt-in must not drag in sibling todo tools, got: {names:?}" + ); + assert!( + !names.contains(&"user__select"), + "a todo opt-in must not expose user tools, got: {names:?}" + ); + } + + #[test] + fn select_functions_all_enabled_tools_exposes_user_tools_despite_node_flag() { + // `all` is a role/session `enabled_tools` semantic — graph llm nodes + // cannot declare it (validate_tools_subset rejects it at run()). The + // flag being set documents that the `all` expansion path bypasses the + // in-node gate by design. + let mut ctx = create_test_ctx(); + ctx.tool_scope.functions.append_user_interaction_functions(); + ctx.in_graph_llm_node = true; + + let mut role = Role::new("r", "p"); + role.set_enabled_tools(Some(vec!["all".to_string()])); + + let fns = ctx.select_functions(&role).unwrap(); + let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect(); + assert!( + names.contains(&"user__select"), + "`all` enabled_tools must expose user tools even inside a node, got: {names:?}" + ); + } + + #[test] + fn select_functions_graph_llm_node_suppresses_user_tools_without_agent() { + let mut ctx = create_test_ctx(); + ctx.tool_scope.functions.append_user_interaction_functions(); + + let mut role = Role::new("r", "p"); + role.set_enabled_tools(Some(vec!["foo".to_string()])); + + ctx.in_graph_llm_node = true; + assert!( + ctx.select_functions(&role).is_none(), + "user__ tools must not leak into a graph llm node without an agent" + ); + + ctx.in_graph_llm_node = false; + let fns = ctx.select_functions(&role).unwrap(); + let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect(); + assert!(names.contains(&"user__select")); + } + + #[test] + fn select_functions_graph_llm_node_suppresses_todo_tools_without_agent() { + let mut ctx = create_test_ctx(); + ctx.update_app_config(|app| app.auto_continue = true); + ctx.tool_scope.functions.append_todo_functions(); + + let mut role = Role::new("r", "p"); + role.set_enabled_tools(Some(vec!["foo".to_string()])); + + ctx.in_graph_llm_node = true; + assert!( + ctx.select_functions(&role).is_none(), + "todo__ tools must not leak into a graph llm node without an agent" + ); + + ctx.in_graph_llm_node = false; + let fns = ctx.select_functions(&role).unwrap(); + let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect(); + assert!(names.contains(&"todo__init")); + } + #[test] #[serial] fn select_functions_preserves_job_tools_under_agent_filter() { diff --git a/src/graph/llm.rs b/src/graph/llm.rs index 6e6b7b3..00ded58 100644 --- a/src/graph/llm.rs +++ b/src/graph/llm.rs @@ -38,7 +38,15 @@ impl LlmNodeExecutor { parent_ctx: &mut RequestContext, abort: &AbortSignal, ) -> Result { - let result = run(node_id, node, state_manager, parent_ctx, abort).await; + let result = run( + node_id, + node, + state_manager, + parent_ctx, + abort, + &mut default_completion_runner(), + ) + .await; finish(node, state_manager, parent_ctx, result, abort).await } } @@ -94,12 +102,15 @@ fn outcome_from( } } +/// The runner is a parameter for the same reason `finish` is split out: +/// tests drive the node scope without a model. async fn run( node_id: &str, node: &LlmNode, state_manager: &mut StateManager, parent_ctx: &mut RequestContext, abort: &AbortSignal, + runner: &mut CompletionRunner<'_>, ) -> Result { let mut instructions: Option = match &node.instructions { Some(s) => Some( @@ -202,14 +213,11 @@ async fn run( node.mcp_tools.clone().map(|map| (node_id.to_string(), map)), ); parent_ctx.refresh_mcp_tool_filters(); - let result = bounded_run( - node, - &prompt, - parent_ctx, - abort, - &mut default_completion_runner(), - ) - .await; + // Inside the node, user__/todo__ builtins require explicit opt-in via the + // node's `tools:` list; mem::replace so nested execution restores cleanly. + let saved_in_node = std::mem::replace(&mut parent_ctx.in_graph_llm_node, true); + let result = bounded_run(node, &prompt, parent_ctx, abort, runner).await; + parent_ctx.in_graph_llm_node = saved_in_node; parent_ctx.role = saved_role; let node_jobs = std::mem::replace(&mut parent_ctx.node_job_scope, saved_job_scope).unwrap_or_default(); @@ -1205,4 +1213,61 @@ mod tests { "{captured}" ); } + + #[tokio::test] + async fn node_scope_flag_is_set_during_the_run_and_restored_on_ok() { + let node = node_with(None); + let mut state = manager_with(&[]); + let mut ctx = plain_ctx(); + ctx.agent = Some(Agent::test_new(AgentConfig::default())); + let abort = create_abort_signal(); + + let out = run( + "n", + &node, + &mut state, + &mut ctx, + &abort, + &mut completion_runner(|_input, ctx, _abort| { + assert!( + ctx.in_graph_llm_node, + "flag must be set while the node runs" + ); + boxed_completion(async { Ok(("done".to_string(), vec![])) }) + }), + ) + .await + .unwrap(); + + assert_eq!(out, "done"); + assert!(!ctx.in_graph_llm_node, "flag must not outlive the node"); + } + + #[tokio::test] + async fn node_scope_flag_restores_prior_value_when_the_run_fails() { + let node = node_with(None); + let mut state = manager_with(&[]); + let mut ctx = plain_ctx(); + ctx.agent = Some(Agent::test_new(AgentConfig::default())); + ctx.in_graph_llm_node = true; + let abort = create_abort_signal(); + + run( + "n", + &node, + &mut state, + &mut ctx, + &abort, + &mut completion_runner(|_input, _ctx, _abort| { + boxed_completion(async { Err(anyhow!("boom")) }) + }), + ) + .await + .expect_err("the runner failure must surface"); + + assert!( + ctx.in_graph_llm_node, + "a nested run must restore the outer scope's flag" + ); + } } diff --git a/src/graph/validator.rs b/src/graph/validator.rs index 69aa6d4..0e941d6 100644 --- a/src/graph/validator.rs +++ b/src/graph/validator.rs @@ -3,6 +3,7 @@ use super::types::{ConcurrencyCap, Graph, NextTargets, Node, NodeType}; use crate::client::{Model, ModelType}; use crate::config; use crate::config::{Agent, AppConfig, paths}; +use crate::function::user_interaction::USER_FUNCTION_PREFIX; use crate::rag::{GraphRagConfig, RagData}; use anyhow::{Result, bail}; use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; @@ -179,6 +180,7 @@ impl GraphValidator { self.validate_max_attempts(graph, &mut result); self.validate_fallback_capture(graph, &mut result); self.validate_output_schema_properties(graph, &mut result); + self.validate_structured_output_user_tools(graph, &mut result); self.validate_max_concurrency(graph, &mut result); self.validate_max_concurrency_template(graph, &mut result); self.validate_orchestration_limits(graph, &mut result); @@ -458,6 +460,29 @@ impl GraphValidator { } } + fn validate_structured_output_user_tools(&self, graph: &Graph, result: &mut ValidationResult) { + for (node_id, node) in &graph.nodes { + let NodeType::Llm(llm) = &node.node_type else { + continue; + }; + if llm.output_schema.is_none() { + continue; + } + let exposes_user_tools = llm.tools.as_ref().is_some_and(|tools| { + tools + .iter() + .any(|t| t.trim().starts_with(USER_FUNCTION_PREFIX)) + }); + if exposes_user_tools { + result.warning(ValidationError::with_node( + node_id, + "structured-output llm node exposes interactive user tools (`user__*`); \ + in headless runs a user__ call blocks the node as an escalation", + )); + } + } + } + fn validate_node_references(&self, graph: &Graph, result: &mut ValidationResult) { for (node_id, node) in &graph.nodes { for (target, label) in declared_targets(node) { @@ -4685,6 +4710,89 @@ mod tests { ); } + #[test] + fn structured_output_node_with_user_tools_warns() { + let mut node = llm_with_output_schema("l", &["verdict"], Some("end")); + if let NodeType::Llm(ref mut n) = node.node_type { + n.tools = Some(vec!["user__select".into()]); + } + let graph = graph_with(vec![("l", node), ("end", end_node("end"))], "l"); + + let result = validator().validate(&graph); + + assert!( + result + .warnings + .iter() + .any(|e| e.node_id.as_deref() == Some("l") + && e.message.contains("interactive user tools")), + "expected user-tool warning: {:?}", + result.warnings + ); + } + + #[test] + fn structured_output_node_with_all_tools_does_not_warn() { + // `all` is not a valid llm-node tools entry (validate_tools_subset + // rejects it at run()), so this advisory stays silent about it. + let mut node = llm_with_output_schema("l", &["verdict"], Some("end")); + if let NodeType::Llm(ref mut n) = node.node_type { + n.tools = Some(vec!["all".into()]); + } + let graph = graph_with(vec![("l", node), ("end", end_node("end"))], "l"); + + let result = validator().validate(&graph); + + assert!( + !result + .warnings + .iter() + .any(|e| e.message.contains("interactive user tools")), + "unexpected user-tool warning: {:?}", + result.warnings + ); + } + + #[test] + fn structured_output_node_without_user_tools_does_not_warn() { + let mut node = llm_with_output_schema("l", &["verdict"], Some("end")); + if let NodeType::Llm(ref mut n) = node.node_type { + n.tools = Some(vec![]); + } + let graph = graph_with(vec![("l", node), ("end", end_node("end"))], "l"); + + let result = validator().validate(&graph); + + assert!( + !result + .warnings + .iter() + .any(|e| e.message.contains("interactive user tools")), + "unexpected user-tool warning: {:?}", + result.warnings + ); + } + + #[test] + fn non_structured_node_with_user_tools_does_not_warn() { + let mut node = llm_node("l", None, Some("end")); + if let NodeType::Llm(ref mut n) = node.node_type { + n.tools = Some(vec!["user__select".into()]); + } + let graph = graph_with(vec![("l", node), ("end", end_node("end"))], "l"); + + let result = validator().validate(&graph); + + assert!( + !result + .warnings + .iter() + .any(|e| e.message.contains("interactive user tools")), + "unexpected user-tool warning: {:?}", + result.warnings + ); + } + #[test] fn fallback_with_output_schema_and_empty_state_updates_warns() { let mut node = llm_with_output_schema("l", &["summary"], Some("end"));