diff --git a/graph.example.yaml b/graph.example.yaml index f027538..19a5d0b 100644 --- a/graph.example.yaml +++ b/graph.example.yaml @@ -37,10 +37,12 @@ reasoning_effort: null # Default reasoning effort for `llm` nodes th # Only valid when the model declares reasoning_levels. max_concurrent_jobs: 5 # Max background jobs (`job__*` tools) running at once across the - # whole graph run. Jobs live in the run-wide supervisor and outlive - # the `llm` node that started them, so the budget is graph-wide — - # there is no per-node override. Overrides the global setting; - # 0 disables background jobs for this graph agent. + # whole graph run: every `llm` node (including parallel branches) + # draws from this one pool, so the budget is graph-wide — there is + # no per-node override. Jobs themselves are node-local: the node + # that starts a job must collect or cancel it before it ends, and + # anything left running at node exit is cancelled. Overrides the + # global setting; 0 disables background jobs for this graph agent. global_tools: # Tool universe an `llm` node's `tools:` whitelist draws from - web_search_coyote.sh diff --git a/src/config/prompts.rs b/src/config/prompts.rs index beb4485..db1e3ce 100644 --- a/src/config/prompts.rs +++ b/src/config/prompts.rs @@ -201,10 +201,10 @@ pub(in crate::config) const DEFAULT_JOB_INSTRUCTIONS: &str = indoc! {" jobs with `job__list`. Collected results over 50,000 chars are tail-capped; collecting is consume-once, so when you need the complete output pass `full_result: true` (or have the command write to a file). Collect or cancel every job you started before ending your turn. In - graph LLM nodes, collect or cancel your jobs before ending your final node turn — an - uncollected job at node turn-end burns node iterations via the guardrail and can fail the - node. Jobs run against a snapshot of the current config/environment and do not survive - coyote exiting. + graph LLM nodes, jobs are node-local: collect or cancel every job you start before the node + ends — an uncollected job burns node iterations via the guardrail, and anything still + running when the node exits is cancelled with its result discarded. Jobs run against a + snapshot of the current config/environment and do not survive coyote exiting. " }; diff --git a/src/config/request_context.rs b/src/config/request_context.rs index fd00b11..fd28018 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -18,7 +18,7 @@ use super::{MessageContentToolCalls, prompts}; use crate::client::{Model, ModelType, list_models}; use crate::function::{ FunctionDeclaration, Functions, ToolCallTracker, ToolResult, - jobs::{DEFAULT_MAX_CONCURRENT_JOBS, JOB_FUNCTION_PREFIX}, + jobs::{DEFAULT_MAX_CONCURRENT_JOBS, JOB_FUNCTION_PREFIX, is_backgroundable_tool}, memory::MEMORY_FUNCTION_PREFIX, rag_query::RAG_FUNCTION_PREFIX, skill::SKILL_FUNCTION_PREFIX, @@ -326,6 +326,13 @@ pub struct RequestContext { pub declared_function_names: HashSet, + /// Ids of jobs started by the currently executing graph LLM node. + /// `Some` only while a node runs: `job__start` records into it, the + /// turn-end guardrail scopes its nag to it, and the node executor reaps + /// whatever is left in it on exit. `None` outside graph nodes — there the + /// context owns every job in its supervisor. + pub node_job_scope: Option>, + pub supervisor: Option>>, pub parent_supervisor: Option>>, pub self_agent_id: Option, @@ -361,6 +368,7 @@ impl RequestContext { last_message: None, tool_scope: ToolScope::default(), declared_function_names: Default::default(), + node_job_scope: None, supervisor: None, parent_supervisor: None, self_agent_id: None, @@ -422,6 +430,7 @@ impl RequestContext { tool_tracker: ToolCallTracker::default(), }, declared_function_names: Default::default(), + node_job_scope: None, supervisor: None, parent_supervisor: None, self_agent_id: None, @@ -470,6 +479,7 @@ impl RequestContext { last_message: self.last_message.clone(), tool_scope: self.tool_scope.clone(), declared_function_names: self.declared_function_names.clone(), + node_job_scope: None, supervisor: self.supervisor.clone(), parent_supervisor: self.parent_supervisor.clone(), self_agent_id: self.self_agent_id.clone(), @@ -516,6 +526,7 @@ impl RequestContext { tool_tracker: tool_call_tracker, }, declared_function_names: Default::default(), + node_job_scope: None, supervisor: None, parent_supervisor: parent.supervisor.clone(), self_agent_id: Some(self_agent_id), @@ -2317,6 +2328,7 @@ impl RequestContext { let mut functions = vec![]; functions.extend(self.select_enabled_functions(role)); functions.extend(self.select_enabled_mcp_servers(role)); + self.apply_job_tool_visibility(&mut functions); if functions.is_empty() { None @@ -2325,6 +2337,38 @@ impl RequestContext { } } + /// Node-local job-ownership visibility rule: the `job__*` family is only + /// declared where it can do something. `job__start` requires at least one + /// backgroundable tool among this request's declarations; the lifecycle + /// verbs (`check`/`collect`/`cancel`/`list`) additionally survive while + /// the context still owns registered jobs, so a job started before a + /// filter change stays reachable. + fn apply_job_tool_visibility(&self, functions: &mut Vec) { + let has_backgroundable = functions.iter().any(|f| is_backgroundable_tool(&f.name)); + if has_backgroundable { + return; + } + let owns_jobs = self.owns_active_jobs(); + let start_name = format!("{JOB_FUNCTION_PREFIX}start"); + functions.retain(|f| { + !f.name.starts_with(JOB_FUNCTION_PREFIX) || (owns_jobs && f.name != start_name) + }); + } + + /// Whether this context has registered jobs it is responsible for: + /// inside a graph LLM node, only the jobs that node started; everywhere + /// else, any job in the context's supervisor. + pub fn owns_active_jobs(&self) -> bool { + let Some(supervisor) = self.supervisor.as_ref() else { + return false; + }; + let sup = supervisor.read(); + match self.node_job_scope.as_ref() { + Some(ids) => ids.iter().any(|id| sup.job(id).is_some()), + None => sup.jobs().next().is_some(), + } + } + pub fn retrieve_role(&self, app: &AppConfig, name: &str) -> Result { let names = paths::list_roles(false); let mut role = if names.contains(&name.to_string()) { @@ -4819,6 +4863,15 @@ mod tests { RequestContext::new(default_app_state(), WorkingMode::Cmd) } + fn test_decl(name: &str) -> FunctionDeclaration { + FunctionDeclaration { + name: name.to_string(), + description: String::new(), + parameters: Default::default(), + agent: false, + } + } + fn tools_only_features(name: &str) -> McpServerFeatures { McpServerFeatures { name: name.to_string(), @@ -5776,37 +5829,49 @@ mod tests { } #[test] - fn select_functions_returns_job_functions_even_with_no_enabled_tools() { + fn select_functions_hides_job_functions_without_backgroundable_tools() { let mut ctx = create_test_ctx(); ctx.tool_scope.functions.append_job_functions(); - let fns = ctx.select_functions(&Role::default()).unwrap(); - let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect(); - - assert_eq!( - names, - vec![ - "job__start", - "job__check", - "job__collect", - "job__cancel", - "job__list" - ] + assert!( + ctx.select_functions(&Role::default()).is_none(), + "job__ tools must not be declared when nothing backgroundable is declared" ); } #[test] - fn select_functions_preserves_job_tools_under_role_filter() { + fn select_functions_keeps_job_tools_when_filter_includes_backgroundable_tool() { let mut ctx = create_test_ctx(); ctx.tool_scope.functions.append_job_functions(); + ctx.tool_scope + .functions + .append_declaration(test_decl("my_build_tool")); let mut role = Role::new("r", "p"); - role.set_enabled_tools(Some(vec!["foo".to_string()])); + role.set_enabled_tools(Some(vec!["my_build_tool".to_string()])); let fns = ctx.select_functions(&role).unwrap(); assert!( fns.iter().any(|f| f.name == "job__start"), - "job__ tools must survive a role tool filter" + "job__ tools must survive a role tool filter that declares a backgroundable tool" + ); + } + + #[test] + fn select_functions_hides_job_tools_when_filter_has_only_non_backgroundable_tools() { + let mut ctx = create_test_ctx(); + ctx.tool_scope.functions.append_job_functions(); + ctx.tool_scope + .functions + .append_declaration(test_decl("fs_cat")); + + let mut role = Role::new("r", "p"); + role.set_enabled_tools(Some(vec!["fs_cat".to_string()])); + + let fns = ctx.select_functions(&role).unwrap(); + assert!( + !fns.iter().any(|f| f.name.starts_with("job__")), + "job__ tools must be hidden when no declared tool is backgroundable" ); } @@ -5822,14 +5887,22 @@ mod tests { fn before_chat_completion_refreshes_declared_function_names() { let mut ctx = create_test_ctx(); ctx.tool_scope.functions.append_job_functions(); - let input = Input::from_str(&ctx, "hello", None).unwrap(); + ctx.tool_scope + .functions + .append_declaration(test_decl("echo")); + let mut role = Role::new("r", "p"); + role.set_enabled_tools(Some(vec!["echo".to_string()])); + let input = Input::from_str(&ctx, "hello", Some(role)).unwrap(); ctx.before_chat_completion(&input).unwrap(); - assert_eq!(ctx.declared_function_names.len(), 5); + assert_eq!(ctx.declared_function_names.len(), 6); assert!(ctx.declared_function_names.contains("job__start")); + assert!(ctx.declared_function_names.contains("echo")); ctx.tool_scope = ToolScope::default(); - let input = Input::from_str(&ctx, "hello again", None).unwrap(); + let mut role = Role::new("r", "p"); + role.set_enabled_tools(Some(vec!["echo".to_string()])); + let input = Input::from_str(&ctx, "hello again", Some(role)).unwrap(); ctx.before_chat_completion(&input).unwrap(); assert!( @@ -6084,6 +6157,9 @@ mod tests { let abort = utils::create_abort_signal(); run_async(ctx.use_agent(&app, &agent_name, None, abort)).unwrap(); + ctx.tool_scope + .functions + .append_declaration(test_decl("foo")); let mut role = Role::new("r", "p"); role.set_enabled_tools(Some(vec!["foo".to_string()])); @@ -6092,7 +6168,7 @@ mod tests { let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect(); assert!( names.contains(&"job__start"), - "job__ tools must survive an agent tool filter, got: {names:?}" + "job__ tools must survive an agent tool filter that declares a backgroundable tool, got: {names:?}" ); assert!(names.contains(&"job__collect")); } @@ -7563,31 +7639,72 @@ mod tests { } #[test] - fn select_functions_preserves_job_tools_under_empty_role_filter() { + fn select_functions_hides_job_tools_under_empty_role_filter() { let mut ctx = create_test_ctx(); ctx.tool_scope.functions.append_job_functions(); let mut role = Role::new("r", "p"); role.set_enabled_tools(Some(vec![])); + assert!( + ctx.select_functions(&role).is_none(), + "an empty tool filter declares nothing backgroundable, so job__ tools must be hidden" + ); + } + + #[test] + fn select_functions_keeps_lifecycle_job_tools_when_context_owns_jobs() { + let mut ctx = create_test_ctx(); + ctx.tool_scope.functions.append_job_functions(); + let sup = Arc::new(RwLock::new( + Supervisor::new(4, 3).with_max_concurrent_jobs(1), + )); + sup.write() + .register(make_running_job(utils::create_abort_signal())) + .unwrap(); + ctx.supervisor = Some(sup); + + let mut role = Role::new("r", "p"); + role.set_enabled_tools(Some(vec![])); + let fns = ctx.select_functions(&role).unwrap(); let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect(); assert_eq!( names, - vec![ - "job__start", - "job__check", - "job__collect", - "job__cancel", - "job__list" - ], - "job__ tools must survive an empty role tool filter" + vec!["job__check", "job__collect", "job__cancel", "job__list"], + "lifecycle verbs must stay reachable while the context owns a job; job__start must not" ); } + #[test] + fn owns_active_jobs_respects_node_scope() { + let mut ctx = create_test_ctx(); + let sup = Arc::new(RwLock::new( + Supervisor::new(4, 3).with_max_concurrent_jobs(1), + )); + sup.write() + .register(make_running_job(utils::create_abort_signal())) + .unwrap(); + ctx.supervisor = Some(sup); + + assert!( + ctx.owns_active_jobs(), + "outside a node, the context owns every registry job" + ); + + ctx.node_job_scope = Some(vec![]); + assert!( + !ctx.owns_active_jobs(), + "a node owns only jobs it started, not other registry entries" + ); + + ctx.node_job_scope = Some(vec!["j1".to_string()]); + assert!(ctx.owns_active_jobs()); + } + #[test] #[serial] - fn select_functions_preserves_job_tools_under_empty_agent_filter() { + fn select_functions_hides_job_tools_under_empty_agent_filter() { let _guard = TestConfigDirGuard::new(); let mut ctx = create_test_ctx(); let app = ctx.app.config.clone(); @@ -7612,13 +7729,16 @@ mod tests { let mut role = Role::new("r", "p"); role.set_enabled_tools(Some(vec![])); - let fns = ctx.select_functions(&role).unwrap(); - let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect(); + let names: Vec = ctx + .select_functions(&role) + .unwrap_or_default() + .iter() + .map(|f| f.name.clone()) + .collect(); assert!( - names.contains(&"job__start"), - "job__ tools must survive an empty agent tool filter, got: {names:?}" + !names.iter().any(|n| n.starts_with("job__")), + "job__ tools must be hidden under an empty agent filter, got: {names:?}" ); - assert!(names.contains(&"job__collect")); } #[test] @@ -7638,6 +7758,9 @@ mod tests { ..(*app).clone() }; run_async(ctx.rebuild_tool_scope(&jobs_off, None, abort.clone())).unwrap(); + ctx.tool_scope + .functions + .append_declaration(test_decl("echo")); let without_jobs = serde_json::to_string(&ctx.select_functions(&role)).unwrap(); assert!( !without_jobs.contains("job__"), @@ -7645,6 +7768,9 @@ mod tests { ); run_async(ctx.rebuild_tool_scope(&app, None, abort)).unwrap(); + ctx.tool_scope + .functions + .append_declaration(test_decl("echo")); let with_jobs = ctx.select_functions(&role).unwrap(); assert!(with_jobs.iter().any(|f| f.name.starts_with("job__"))); let stripped: Vec = with_jobs @@ -7685,6 +7811,12 @@ mod tests { fn tools_info_lists_job_tools_when_enabled() { let mut ctx = create_test_ctx(); ctx.tool_scope.functions.append_job_functions(); + ctx.tool_scope + .functions + .append_declaration(test_decl("echo")); + let mut role = Role::new("r", "p"); + role.set_enabled_tools(Some(vec!["echo".to_string()])); + ctx.role = Some(role); let info = ctx.tools_info().unwrap(); diff --git a/src/function/jobs.rs b/src/function/jobs.rs index 0b5fba4..135a97e 100644 --- a/src/function/jobs.rs +++ b/src/function/jobs.rs @@ -140,7 +140,8 @@ pub fn job_function_declarations() -> Vec { tools cannot be backgrounded. The job runs against a snapshot of the current config and \ environment; later changes do not affect it. Process jobs honor COYOTE_TOOL_TIMEOUT; MCP \ jobs have NO timeout — cancel a hung one with `job__cancel`. Jobs do not survive coyote \ - exiting.".to_string(), + exiting. In graph LLM nodes, jobs are node-local: collect or cancel every job you start \ + before the node ends — leftovers are cancelled at node exit.".to_string(), parameters: JsonSchema { type_value: Some("object".to_string()), properties: Some(IndexMap::from([ @@ -359,6 +360,13 @@ fn whitelist_rejection(tool: &str) -> Option { }) } +/// Whether a declared tool could be run as a background job. This is the +/// declare-side twin of `whitelist_rejection`: a tool is backgroundable +/// exactly when `job__start` would not reject it by name. +pub fn is_backgroundable_tool(tool: &str) -> bool { + whitelist_rejection(tool).is_none() +} + async fn handle_start(ctx: &mut RequestContext, args: &Value) -> Result { if !jobs_enabled(ctx.agent.as_ref(), &ctx.app.config) { return Ok(json!({ @@ -511,6 +519,10 @@ async fn handle_start(ctx: &mut RequestContext, args: &Value) -> Result { })); } + if let Some(scope) = ctx.node_job_scope.as_mut() { + scope.push(job_id.clone()); + } + Ok(json!({ "status": "ok", "job_id": job_id, @@ -803,6 +815,36 @@ async fn kill_job_with_grace(handle: &mut JobHandle) { let _ = time::timeout(JOB_KILL_GRACE, &mut handle.join_handle).await; } +/// Cancel and deregister the named jobs, if still registered. Used by graph +/// LLM nodes to enforce node-local job ownership: any job the node started +/// but did not collect or cancel by the time it exits is killed here, on +/// every exit path. Returns the ids actually reaped. +pub async fn reap_jobs( + supervisor: Option<&Arc>>, + ids: &[String], +) -> Vec { + let Some(supervisor) = supervisor else { + return Vec::new(); + }; + let mut reaped = Vec::new(); + for id in ids { + let handle = { + let mut sup = supervisor.write(); + sup.take_job(id) + }; + if let Some(mut handle) = handle { + handle.abort_signal.set_ctrlc(); + kill_job_with_grace(&mut handle).await; + warn!( + "Reaped background job '{id}' ({}): left unreclaimed at graph node exit", + handle.tool + ); + reaped.push(id.clone()); + } + } + reaped +} + fn handle_list(ctx: &RequestContext) -> Result { let Some(supervisor) = ctx.supervisor.as_ref() else { return Ok(json!({ @@ -1552,6 +1594,62 @@ mod tests { ); } + #[test] + fn is_backgroundable_tool_matches_start_whitelist() { + assert!(is_backgroundable_tool("execute_command")); + assert!(is_backgroundable_tool("my_custom_tool.sh")); + assert!(is_backgroundable_tool("mcp_invoke_github")); + assert!(!is_backgroundable_tool("job__start")); + assert!(!is_backgroundable_tool("agent__spawn")); + assert!(!is_backgroundable_tool("user__confirm")); + assert!(!is_backgroundable_tool("todo__add")); + assert!(!is_backgroundable_tool("fs_read")); + assert!(!is_backgroundable_tool("ast_grep")); + assert!(!is_backgroundable_tool("mcp_search_github")); + } + + #[cfg(unix)] + #[test] + fn handle_start_records_job_in_node_scope() { + run_async(async { + let mut ctx = plain_ctx(); + ctx.node_job_scope = Some(Vec::new()); + ctx.declared_function_names.insert("echo".into()); + + let started = handle_start(&mut ctx, &json!({"tool": "echo", "arguments": {}})) + .await + .unwrap(); + + let job_id = started["job_id"].as_str().unwrap().to_string(); + assert_eq!(ctx.node_job_scope.clone().unwrap(), vec![job_id.clone()]); + + let collected = handle_collect(&ctx, &json!({"id": job_id})).await.unwrap(); + assert_eq!(collected["status"], "completed"); + }); + } + + #[test] + fn reap_jobs_kills_registered_jobs_and_reports_ids() { + run_async(async { + let ctx = ctx_with_job_supervisor(4); + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(make_running_job("j1")) + .unwrap(); + + let reaped = reap_jobs( + ctx.supervisor.as_ref(), + &["j1".to_string(), "missing".to_string()], + ) + .await; + + assert_eq!(reaped, vec!["j1".to_string()]); + assert!(!ctx.supervisor.as_ref().unwrap().read().has_job("j1")); + }); + } + #[test] fn handle_start_rejects_unconnected_mcp_server() { let mut ctx = plain_ctx(); diff --git a/src/function/mod.rs b/src/function/mod.rs index 3969aaf..87d6aad 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -693,6 +693,10 @@ impl Functions { self.declarations.extend(jobs::job_function_declarations()); } + pub fn append_declaration(&mut self, declaration: FunctionDeclaration) { + self.declarations.push(declaration); + } + pub fn append_teammate_functions(&mut self) { self.declarations .extend(supervisor::teammate_function_declarations()); diff --git a/src/function/supervisor.rs b/src/function/supervisor.rs index 5e05d2c..52a4081 100644 --- a/src/function/supervisor.rs +++ b/src/function/supervisor.rs @@ -74,6 +74,13 @@ pub fn pending_tasks(ctx: &RequestContext) -> Vec { }) .collect(); + // Inside a graph LLM node, jobs are node-owned: the guardrail must only + // nag about jobs this node started. Jobs belonging to a parallel branch + // live in the same shared registry but are that branch's to reclaim. + if let Some(scope) = ctx.node_job_scope.as_ref() { + tasks.retain(|t| t.kind != TaskKind::Job || scope.contains(&t.id)); + } + tasks.sort_by(|a, b| a.id.cmp(&b.id)); tasks } @@ -2596,6 +2603,22 @@ mod tests { assert_eq!(tasks[0].kind, TaskKind::Job); } + #[test] + fn pending_tasks_scopes_jobs_to_node_scope() { + let mut ctx = ctx_with_job_capable_supervisor(); + register_fake_job(&mut ctx, "job_mine"); + register_fake_job(&mut ctx, "job_other"); + + ctx.node_job_scope = Some(vec!["job_mine".to_string()]); + let tasks = pending_tasks(&ctx); + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0].id, "job_mine"); + assert_eq!(tasks[0].kind, TaskKind::Job); + + ctx.node_job_scope = None; + assert_eq!(pending_tasks(&ctx).len(), 2); + } + #[test] fn guardrail_force_terminates_at_max_and_cancels_agents() { let rt = tokio::runtime::Builder::new_current_thread() diff --git a/src/graph/llm.rs b/src/graph/llm.rs index 7e19a19..36a17dd 100644 --- a/src/graph/llm.rs +++ b/src/graph/llm.rs @@ -6,6 +6,7 @@ use crate::config::prompts::DEFAULT_SKILL_INSTRUCTIONS; use crate::config::{ Input, RequestContext, Role, RoleLike, SkillPolicy, should_inject_skill_instructions, }; +use crate::function::jobs::reap_jobs; use crate::function::skill::skill_function_declarations; use crate::function::supervisor::{GuardrailAction, check_pending_tasks_guardrail}; use crate::utils::create_abort_signal; @@ -173,6 +174,9 @@ async fn run( let saved_role = parent_ctx.role.clone(); parent_ctx.role = Some(composed_role); + // 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()); let result = match node.timeout { Some(secs) => match timeout( Duration::from_secs(secs), @@ -186,6 +190,9 @@ async fn run( None => run_with_retries(node, &prompt, parent_ctx).await, }; parent_ctx.role = saved_role; + 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; restore_agent_skill_policy(parent_ctx, saved_agent_skill_state); result }