diff --git a/Cargo.lock b/Cargo.lock index 05028a1..d856e4a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1703,6 +1703,7 @@ dependencies = [ "inquire", "is-terminal", "json-patch", + "libc", "log", "log4rs", "nu-ansi-term", diff --git a/Cargo.toml b/Cargo.toml index d8473a6..fbfa196 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -139,6 +139,9 @@ arboard = { version = "3.3.0", default-features = false, features = [ [target.'cfg(not(any(target_os = "linux", target_os = "android", target_os = "emscripten")))'.dependencies] arboard = { version = "3.3.0", default-features = false } +[target.'cfg(unix)'.dependencies] +libc = "0.2" + [dev-dependencies] pretty_assertions = "1.4.0" rmcp = { version = "3.1.2", features = ["server"] } diff --git a/README.md b/README.md index 0d32e56..2b0b79a 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,7 @@ Coming from [AIChat](https://github.com/sigoden/aichat)? Follow the [migration g * [Skills](https://github.com/Dark-Alex-17/coyote/wiki/Skills): Modular knowledge or capability packs the LLM can load and unload mid-conversation. Multiple skills compose; instructions stack, tools and MCPs union. * [Agents](https://github.com/Dark-Alex-17/coyote/wiki/Agents): Leverage AI agents to perform complex tasks and workflows, including sub-agent spawning, teammate messaging, and user interaction tools. * [Graph Agents](https://github.com/Dark-Alex-17/coyote/wiki/Graph-Agents): Define an agent as a declarative, YAML-driven workflow. A directed graph of typed nodes (LLM calls, scripts, approvals, user input, RAG retrieval, sub-agent spawns). +* [Background Jobs](https://github.com/Dark-Alex-17/coyote/wiki/Background-Jobs): Run long tool calls (builds, test suites, slow MCP calls) in the background with the `job__*` tools while the model keeps working, and completion arrives as a push notification. * [Todo System](https://github.com/Dark-Alex-17/coyote/wiki/TODO-System): Built-in task tracking for improved LLM reliability with smaller models. * [Environment Variables](https://github.com/Dark-Alex-17/coyote/wiki/Environment-Variables): Override and customize your Coyote configuration at runtime with environment variables. * [Client Configurations](https://github.com/Dark-Alex-17/coyote/wiki/Clients): Configuration instructions for various LLM providers. diff --git a/assets/agents/architect/config.yaml b/assets/agents/architect/config.yaml index b4b365d..3cc1fab 100644 --- a/assets/agents/architect/config.yaml +++ b/assets/agents/architect/config.yaml @@ -261,7 +261,7 @@ instructions: | 3. **Wait for Sisyphus.** Do not poll `agent__collect` on a running agent — do non-overlapping work (e.g. prep the next task's context) or end your response and wait for the completion - notification, then `agent__collect`. + notification (a `system_notifications` entry on your next tool result), then `agent__collect`. 4. **Verify against the plan (divergence check).** When Sisyphus returns, do NOT trust its self-report — get an INDEPENDENT conformance verdict: diff --git a/assets/agents/sisyphus/config.yaml b/assets/agents/sisyphus/config.yaml index dbbd2b7..05dc267 100644 --- a/assets/agents/sisyphus/config.yaml +++ b/assets/agents/sisyphus/config.yaml @@ -233,7 +233,7 @@ instructions: | 1. Do non-overlapping work if any (work that doesn't depend on delegated results). 2. If none → **end your response.** Do not call `agent__collect` immediately. - 3. The system notifies you on completion. + 3. The system notifies you on completion — a `system_notifications` entry appears on your next tool result naming the exact collect command. 4. On notification, call `agent__collect` to retrieve results. ### Anti-duplication rule (BLOCKING) diff --git a/assets/functions/tools/execute_command.sh b/assets/functions/tools/execute_command.sh index 4105a85..c360aef 100755 --- a/assets/functions/tools/execute_command.sh +++ b/assets/functions/tools/execute_command.sh @@ -20,5 +20,12 @@ main() { trap "rm -f '$script'" EXIT # shellcheck disable=SC2154 printf '%s\n' "$argc_command" > "$script" - bash -e -o pipefail "$script" >> "$LLM_OUTPUT" + # No -e: the command gets standard interactive-shell semantics — the last + # statement decides the exit code, so trailing guards like `; exit 0` work + # and an intermediate non-zero status (grep with no matches, a failing + # test run being inspected) cannot abort the script mid-way. pipefail is + # kept so a failing pipeline stage still surfaces in the exit code. 2>&1: + # the harness only returns $LLM_OUTPUT on success, so without it stderr + # (git push, cargo progress, curl -v) vanishes from successful calls. + bash -o pipefail "$script" >> "$LLM_OUTPUT" 2>&1 } diff --git a/config.agent.example.yaml b/config.agent.example.yaml index c4562df..18c731c 100644 --- a/config.agent.example.yaml +++ b/config.agent.example.yaml @@ -41,6 +41,8 @@ can_spawn_agents: false # Enable the agent to spawn child agents # Graph agents (graph.yaml) ignore this; they declare spawn targets in agent nodes. max_concurrent_agents: 4 # Maximum number of agents that can run simultaneously max_agent_depth: 3 # Maximum nesting depth for sub-agents (prevents runaway spawning) +max_concurrent_jobs: 5 # Max background jobs (`job__*` tools) running at once for this agent + # (overrides the global setting; 0 disables background jobs for this agent) inject_spawn_instructions: true # Inject the default agent spawning instructions into the agent's system prompt summarization_model: null # Model to use for summarizing sub-agent output (e.g. 'openai:gpt-4o-mini'); defaults to current model summarization_threshold: 4000 # Character threshold above which sub-agent output is summarized before returning to parent diff --git a/config.example.yaml b/config.example.yaml index 247333a..40ab1c5 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -204,6 +204,7 @@ summary_context_prompt: > # The text prompt used for including the summar 'This is a summary of the chat history as a recap: ' compression_keep_last: 0 # Number of most-recent messages to keep visible after compression (0 = compress all messages) max_tool_result_chars: null # Cap on tool result characters forwarded to the model per call (null = no cap) +max_concurrent_jobs: 5 # Max background jobs (`job__*` tools) running at once per context (default: 5; 0 disables background jobs entirely) # ---- Memory ---- # See the [Memory documentation](https://github.com/Dark-Alex-17/coyote/wiki/Memory) for more information. diff --git a/graph.example.yaml b/graph.example.yaml index f1ea9ae..19a5d0b 100644 --- a/graph.example.yaml +++ b/graph.example.yaml @@ -36,6 +36,14 @@ top_p: null # Default sampling top-p for `llm` nodes reasoning_effort: null # Default reasoning effort for `llm` nodes that don't override it. # 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: 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 - fetch_url_via_curl.sh diff --git a/src/acp/server.rs b/src/acp/server.rs index 689a472..506c15e 100644 --- a/src/acp/server.rs +++ b/src/acp/server.rs @@ -1,7 +1,7 @@ use super::types::{METHOD_NOT_FOUND, PARSE_ERROR, Request, Response}; use crate::client::call_chat_completions_streaming; use crate::config::{Input, RenderMode, RequestContext}; -use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail}; +use crate::function::agents::{GuardrailAction, check_pending_tasks_guardrail}; use crate::utils; use crate::utils::AbortSignal; use anyhow::Result; @@ -211,7 +211,7 @@ async fn run_prompt_turn( input = input.merge_tool_results(output, tool_results); continue; } - match check_pending_agents_guardrail(ctx) { + match check_pending_tasks_guardrail(ctx) { GuardrailAction::Inject(prompt) => { input = Input::from_str(ctx, &prompt, None)?; } diff --git a/src/config/agent.rs b/src/config/agent.rs index 40771fa..15638ea 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -3,15 +3,19 @@ use super::*; use crate::{ client::Model, config::memory, - function::{Functions, run_llm_function}, + function::{ + Functions, + jobs::{DEFAULT_MAX_CONCURRENT_JOBS, JOB_FUNCTION_PREFIX}, + run_llm_function, + }, graph, rag, }; use super::rag_cache::RagKey; use crate::config::paths; use crate::config::prompts::{ - DEFAULT_SPAWN_INSTRUCTIONS, DEFAULT_TEAMMATE_INSTRUCTIONS, DEFAULT_TODO_INSTRUCTIONS, - DEFAULT_USER_INTERACTION_INSTRUCTIONS, + DEFAULT_JOB_INSTRUCTIONS, DEFAULT_SPAWN_INSTRUCTIONS, DEFAULT_TEAMMATE_INSTRUCTIONS, + DEFAULT_TODO_INSTRUCTIONS, DEFAULT_USER_INTERACTION_INSTRUCTIONS, }; use crate::graph::types::RagNode; use crate::graph::{Graph, GraphParser, NodeType}; @@ -225,6 +229,16 @@ impl Agent { functions.append_supervisor_functions(); } + if app.function_calling_support + && agent_config + .max_concurrent_jobs + .or(app.max_concurrent_jobs) + .unwrap_or(DEFAULT_MAX_CONCURRENT_JOBS) + > 0 + { + functions.append_job_functions(); + } + functions.append_teammate_functions(); functions.append_user_interaction_functions(); @@ -440,6 +454,18 @@ impl Agent { output.push_str(DEFAULT_SPAWN_INSTRUCTIONS); } + if self + .functions + .declarations() + .iter() + .any(|f| f.name.starts_with(JOB_FUNCTION_PREFIX)) + { + if !output.ends_with('\n') { + output.push('\n'); + } + output.push_str(DEFAULT_JOB_INSTRUCTIONS); + } + output.push_str(DEFAULT_TEAMMATE_INSTRUCTIONS); output.push_str(DEFAULT_USER_INTERACTION_INSTRUCTIONS); @@ -561,6 +587,10 @@ impl Agent { self.config.max_tool_result_chars } + pub fn max_concurrent_jobs(&self) -> Option { + self.config.max_concurrent_jobs + } + pub fn compression_keep_last(&self) -> Option { self.config.compression_keep_last } @@ -735,6 +765,8 @@ pub struct AgentConfig { #[serde(default, skip_serializing_if = "Option::is_none")] pub max_tool_result_chars: Option, #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_concurrent_jobs: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub compression_keep_last: Option, #[serde(default)] pub description: String, @@ -822,6 +854,7 @@ impl AgentConfig { variables: graph.variables.clone(), can_spawn_agents: graph.has_agent_node(), max_concurrent_agents: default_max_concurrent_agents(), + max_concurrent_jobs: graph.max_concurrent_jobs, max_agent_depth: default_max_agent_depth(), escalation_timeout: default_escalation_timeout(), ..AgentConfig::default() @@ -1296,6 +1329,7 @@ variables: model: claude:claude-sonnet-4-6 temperature: 0.3 top_p: 0.8 + max_concurrent_jobs: 2 global_tools: - fetch_pdf.sh mcp_servers: @@ -1318,6 +1352,7 @@ variables: assert_eq!(config.model_id.as_deref(), Some("claude:claude-sonnet-4-6")); assert_eq!(config.temperature, Some(0.3)); assert_eq!(config.top_p, Some(0.8)); + assert_eq!(config.max_concurrent_jobs, Some(2)); assert_eq!(config.global_tools, vec!["fetch_pdf.sh"]); assert_eq!(config.mcp_servers, vec!["pubmed-search"]); assert_eq!(config.conversation_starters, vec!["Start here"]); @@ -1481,4 +1516,51 @@ nodes: {} assert_eq!(config.top_k, Some(7)); assert_eq!(config.embedding_model.as_deref(), Some("some:model")); } + + #[test] + fn interpolated_instructions_without_job_declarations_is_byte_identical_across_job_settings() { + let agent = |max_concurrent_jobs| { + Agent::test_new(AgentConfig { + instructions: "hi".to_string(), + max_concurrent_jobs, + ..AgentConfig::default() + }) + }; + + let baseline = agent(None).interpolated_instructions(); + assert!( + !baseline.contains(DEFAULT_JOB_INSTRUCTIONS), + "no job guidance may be injected without job__ declarations" + ); + assert_eq!(baseline, agent(Some(0)).interpolated_instructions()); + assert_eq!(baseline, agent(Some(7)).interpolated_instructions()); + + let mut with_unrelated = agent(None); + with_unrelated.functions.append_todo_functions(); + assert_eq!( + baseline, + with_unrelated.interpolated_instructions(), + "job guidance injection must key strictly on the job__ prefix" + ); + } + + #[test] + fn interpolated_instructions_with_job_declarations_appends_job_guidance() { + let config = AgentConfig { + instructions: "hi".to_string(), + ..AgentConfig::default() + }; + let baseline = Agent::test_new(config.clone()).interpolated_instructions(); + + let mut agent = Agent::test_new(config); + agent.functions.append_job_functions(); + let output = agent.interpolated_instructions(); + + assert!(output.contains(DEFAULT_JOB_INSTRUCTIONS)); + let expected = format!( + "hi\n{DEFAULT_JOB_INSTRUCTIONS}{}", + baseline.strip_prefix("hi").unwrap() + ); + assert_eq!(output, expected); + } } diff --git a/src/config/app_config.rs b/src/config/app_config.rs index 5dd6b2d..4f4e4b1 100644 --- a/src/config/app_config.rs +++ b/src/config/app_config.rs @@ -68,6 +68,7 @@ pub struct AppConfig { pub summarization_prompt: Option, pub summary_context_prompt: Option, pub max_tool_result_chars: Option, + pub max_concurrent_jobs: Option, pub memory: Option, pub memory_cap_with_tools: Option, @@ -153,6 +154,7 @@ impl Default for AppConfig { summarization_prompt: None, summary_context_prompt: None, max_tool_result_chars: None, + max_concurrent_jobs: None, memory: None, memory_cap_with_tools: None, @@ -239,6 +241,7 @@ impl AppConfig { summarization_prompt: config.summarization_prompt, summary_context_prompt: config.summary_context_prompt, max_tool_result_chars: config.max_tool_result_chars, + max_concurrent_jobs: config.max_concurrent_jobs, memory: config.memory, memory_cap_with_tools: config.memory_cap_with_tools, @@ -574,6 +577,9 @@ impl AppConfig { { self.compression_threshold = v; } + if let Some(v) = super::read_env_value::(&get_env_name("max_concurrent_jobs")) { + self.max_concurrent_jobs = v; + } if let Some(v) = super::read_env_value::(&get_env_name("summarization_prompt")) { self.summarization_prompt = v; } @@ -838,16 +844,59 @@ mod tests { unsafe { match prev { - Some(v) => std::env::set_var(&env_name, v), - None => std::env::remove_var(&env_name), + Some(v) => env::set_var(&env_name, v), + None => env::remove_var(&env_name), + } + } + } + + #[test] + fn from_config_copies_max_concurrent_jobs() { + let cfg = Config { + model_id: "test-model".to_string(), + max_concurrent_jobs: Some(3), + clients: vec![ClientConfig::default()], + ..Config::default() + }; + + let app = AppConfig::from_config(cfg).unwrap(); + + assert_eq!(app.max_concurrent_jobs, Some(3)); + } + + #[test] + #[serial_test::serial] + fn load_envs_overrides_max_concurrent_jobs() { + let env_name = get_env_name("max_concurrent_jobs"); + let prev = env::var_os(&env_name); + + let mut app = AppConfig::default(); + + unsafe { env::set_var(&env_name, "7") }; + app.load_envs(); + assert_eq!(app.max_concurrent_jobs, Some(7)); + + unsafe { env::set_var(&env_name, "0") }; + app.load_envs(); + assert_eq!(app.max_concurrent_jobs, Some(0)); + + unsafe { env::remove_var(&env_name) }; + app.max_concurrent_jobs = Some(2); + app.load_envs(); + assert_eq!(app.max_concurrent_jobs, Some(2)); + + unsafe { + match prev { + Some(v) => env::set_var(&env_name, v), + None => env::remove_var(&env_name), } } } #[test] fn editor_returns_configured_value() { - let configured = cached_editor() - .unwrap_or_else(|| std::env::current_exe().unwrap().display().to_string()); + let configured = + cached_editor().unwrap_or_else(|| env::current_exe().unwrap().display().to_string()); let app = AppConfig { editor: Some(configured.clone()), ..AppConfig::default() @@ -864,9 +913,9 @@ mod tests { return; } - let expected = std::env::current_exe().unwrap().display().to_string(); + let expected = env::current_exe().unwrap().display().to_string(); unsafe { - std::env::set_var("VISUAL", &expected); + env::set_var("VISUAL", &expected); } let app = AppConfig::default(); @@ -934,7 +983,7 @@ mod tests { let app = AppConfig::from_config(cfg).unwrap(); let ua = app.user_agent.as_deref().unwrap(); - assert!(ua != "auto", "user_agent should have been resolved"); + assert_ne!(ua, "auto", "user_agent should have been resolved"); assert!(ua.contains('/'), "user_agent should be '/'"); } diff --git a/src/config/app_state.rs b/src/config/app_state.rs index 7a6cb0e..5cbd452 100644 --- a/src/config/app_state.rs +++ b/src/config/app_state.rs @@ -1,6 +1,7 @@ use super::mcp_factory::{McpFactory, McpServerKey}; use super::rag_cache::RagCache; use crate::config::AppConfig; +use crate::config::jobs_enabled; use crate::function::Functions; use crate::mcp::{McpRegistry, McpServersConfig}; use crate::utils::AbortSignal; @@ -73,6 +74,10 @@ impl AppState { functions.append_mcp_meta_functions(mcp_registry.server_features()); } + if jobs_enabled(None, &config) { + functions.append_job_functions(); + } + let mcp_registry = if mcp_registry.is_empty() { None } else { diff --git a/src/config/input.rs b/src/config/input.rs index 71f82d5..52b389c 100644 --- a/src/config/input.rs +++ b/src/config/input.rs @@ -9,7 +9,12 @@ use crate::utils::{AbortSignal, base64_encode, is_loader_protocol, sha256}; use anyhow::{Context, Result, bail}; use indexmap::IndexSet; -use std::{collections::HashMap, fs::File, io::Read, sync::Arc}; +use std::{ + collections::{HashMap, HashSet}, + fs::File, + io::Read, + sync::Arc, +}; use unicode_width::{UnicodeWidthChar, UnicodeWidthStr}; const IMAGE_EXTS: [&str; 5] = ["png", "jpeg", "jpg", "webp", "gif"]; @@ -158,6 +163,13 @@ impl Input { self.data_urls.clone() } + pub fn declared_function_names(&self) -> HashSet { + self.functions + .as_ref() + .map(|functions| functions.iter().map(|f| f.name.clone()).collect()) + .unwrap_or_default() + } + pub fn tool_calls(&self) -> &Option { &self.tool_calls } @@ -593,6 +605,8 @@ mod tests { use super::*; use crate::config::request_context::RequestContext; use crate::config::{AppState, WorkingMode}; + use crate::function::ToolCall; + use serde_json::json; use std::fs; use std::sync::Arc; use std::time::SystemTime; @@ -973,4 +987,70 @@ mod tests { )); assert!(result.is_err()); } + + fn tool_result(id: &str, output: &str) -> ToolResult { + ToolResult::new( + ToolCall::new("t".into(), json!({}), Some(id.to_string())), + json!(output), + ) + } + + #[test] + fn merge_tool_results_first_merge_creates_container() { + let ctx = create_test_ctx(); + let input = Input::from_str(&ctx, "test", None).unwrap(); + + let input = + input.merge_tool_results("assistant text".into(), vec![tool_result("id-1", "ok")]); + + let tool_calls = input.tool_calls().as_ref().unwrap(); + assert_eq!(tool_calls.text, "assistant text"); + assert!(!tool_calls.sequence); + assert_eq!(tool_calls.tool_results.len(), 1); + assert!(tool_calls.tool_results[0].text.is_none()); + } + + #[test] + fn merge_tool_results_second_merge_marks_sequence_and_tags_text() { + let ctx = create_test_ctx(); + let input = Input::from_str(&ctx, "test", None) + .unwrap() + .merge_tool_results("assistant text".into(), vec![tool_result("id-1", "ok")]); + + let input = + input.merge_tool_results("second text".into(), vec![tool_result("id-2", "ok2")]); + + let tool_calls = input.tool_calls().as_ref().unwrap(); + assert!(tool_calls.sequence); + assert_eq!(tool_calls.tool_results.len(), 2); + assert_eq!(tool_calls.text, "assistant text"); + assert!(tool_calls.tool_results[0].text.is_none()); + assert_eq!( + tool_calls.tool_results[1].text, + Some("second text".to_string()) + ); + } + + #[test] + fn build_messages_wraps_tool_results_in_single_assistant_message() { + let ctx = create_test_ctx(); + let input = Input::from_str(&ctx, "test", None) + .unwrap() + .merge_tool_results("assistant text".into(), vec![tool_result("id-1", "ok")]) + .merge_tool_results("second text".into(), vec![tool_result("id-2", "ok2")]); + + let messages = input.build_messages().unwrap(); + + let tool_call_messages: Vec<_> = messages + .iter() + .filter(|m| matches!(m.content, MessageContent::ToolCalls(_))) + .collect(); + assert_eq!(tool_call_messages.len(), 1); + let message = tool_call_messages[0]; + assert!(matches!(message.role, MessageRole::Assistant)); + let MessageContent::ToolCalls(tool_calls) = &message.content else { + unreachable!(); + }; + assert_eq!(tool_calls.tool_results.len(), 2); + } } diff --git a/src/config/mod.rs b/src/config/mod.rs index 08992fd..e89caee 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -42,7 +42,10 @@ pub use self::macro_policy::{ MacroAllowlistLevel, MacroPolicy, MacroSource, MacroState, RESERVED_MACRO_NAMES, ResolvedMacro, }; #[allow(unused_imports)] -pub use self::request_context::{RenderMode, RequestContext, should_inject_skill_instructions}; +pub use self::request_context::{ + RenderMode, RequestContext, effective_max_concurrent_jobs, jobs_enabled, + should_inject_skill_instructions, +}; pub use self::role::{ CODE_ROLE, CREATE_TITLE_ROLE, EXPLAIN_SHELL_ROLE, Role, RoleLike, SHELL_ROLE, }; @@ -56,7 +59,8 @@ pub use self::skill_registry::SkillRegistry; #[cfg(test)] pub(crate) use self::tool_scope::test_fixtures; pub use self::tool_scope::{ - McpPromptCompletion, flatten_prompt_messages, resolve_prompt_args, sanitize_display_text, + McpPromptCompletion, McpRuntime, flatten_prompt_messages, resolve_prompt_args, + sanitize_display_text, }; pub use self::update::run_self_update; use crate::client::{ @@ -264,6 +268,7 @@ pub struct Config { pub summarization_prompt: Option, pub summary_context_prompt: Option, pub max_tool_result_chars: Option, + pub max_concurrent_jobs: Option, pub memory: Option, pub memory_cap_with_tools: Option, @@ -346,6 +351,7 @@ impl Default for Config { summarization_prompt: None, summary_context_prompt: None, max_tool_result_chars: None, + max_concurrent_jobs: None, memory: None, memory_cap_with_tools: None, diff --git a/src/config/prompts.rs b/src/config/prompts.rs index 8af226b..db1e3ce 100644 --- a/src/config/prompts.rs +++ b/src/config/prompts.rs @@ -82,7 +82,7 @@ pub(in crate::config) const DEFAULT_SPAWN_INSTRUCTIONS: &str = indoc! {" | Tool | Purpose | |------|----------| | `agent__spawn` | Spawn a subagent in the background. Returns an `id` immediately. | - | `agent__check` | Non-blocking check: is the agent done yet? Returns PENDING or result. | + | `agent__check` | Non-blocking status probe: running or finished. Never returns/consumes the result — use `agent__collect`. | | `agent__collect` | Blocking wait: wait for an agent to finish, return its output. | | `agent__list_available` | List all agent types you can spawn (name + description). Use this to discover specialists before calling `agent__spawn`. | | `agent__list_running` | List all subagents YOU have spawned, with their status. | @@ -112,9 +112,10 @@ pub(in crate::config) const DEFAULT_SPAWN_INSTRUCTIONS: &str = indoc! {" ### CRITICAL: Never end your turn with pending agents - Spawned agents do NOT report back on their own. They run in the background until you - actively reclaim them with `agent__collect` (to get their output) or `agent__cancel` - (to discard them). If you spawn agents and then emit a final message without reclaiming + Spawned agents do NOT deliver their results on their own. When one finishes, a + `system_notifications` entry appears on your next tool result naming the exact collect + command — but the output is only retrieved when you actively reclaim it with `agent__collect` + (or discard it with `agent__cancel`). If you spawn agents and then emit a final message without reclaiming them, the system will detect the unreclaimed agents and reject the turn-end, injecting a reminder forcing you to handle them. After several such reminders, the system will auto-cancel them and warn you that work was lost. @@ -190,6 +191,23 @@ pub(in crate::config) const DEFAULT_SPAWN_INSTRUCTIONS: &str = indoc! {" 4. **Respond promptly**; the child agent is blocked and waiting (5-minute timeout). "}; +pub(in crate::config) const DEFAULT_JOB_INSTRUCTIONS: &str = indoc! {" + ## Background Jobs + + For long-running tool calls (builds, test suites, slow commands), call `job__start` and keep + working instead of blocking — completion arrives as a `system_notifications` entry on your + next tool result. Check progress with `job__check` (sparingly), block on the result with + `job__collect` (only when you have nothing else to do), cancel with `job__cancel`, and list + 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, 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. +" +}; + pub(in crate::config) const DEFAULT_TEAMMATE_INSTRUCTIONS: &str = indoc! {" ## Teammate Messaging diff --git a/src/config/request_context.rs b/src/config/request_context.rs index 2be3d37..8fd9ab4 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -17,9 +17,13 @@ use super::{ use super::{MessageContentToolCalls, prompts}; use crate::client::{Model, ModelType, list_models}; use crate::function::{ - FunctionDeclaration, Functions, ToolCallTracker, ToolResult, memory::MEMORY_FUNCTION_PREFIX, - rag_query::RAG_FUNCTION_PREFIX, skill::SKILL_FUNCTION_PREFIX, - supervisor::SUPERVISOR_FUNCTION_PREFIX, todo::TODO_FUNCTION_PREFIX, + FunctionDeclaration, Functions, ToolCallTracker, ToolResult, + agents::AGENT_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, + todo::TODO_FUNCTION_PREFIX, user_interaction::USER_FUNCTION_PREFIX, }; use crate::mcp::{ @@ -30,6 +34,7 @@ use crate::rag::Rag; use crate::supervisor::Supervisor; use crate::supervisor::escalation::EscalationQueue; use crate::supervisor::mailbox::Inbox; +use crate::supervisor::notification::NotificationQueue; use crate::utils::{ AbortSignal, abortable_run_with_spinner, edit_file, fuzzy_filter, get_env_name, list_file_names, now, render_prompt, temp_file, @@ -138,6 +143,17 @@ pub fn should_inject_skill_instructions(app: &AppConfig, policy: &SkillPolicy) - app.function_calling_support && policy.skills_enabled && !policy.compatible_enabled.is_empty() } +pub fn effective_max_concurrent_jobs(agent: Option<&Agent>, app: &AppConfig) -> usize { + agent + .and_then(|a| a.max_concurrent_jobs()) + .or(app.max_concurrent_jobs) + .unwrap_or(DEFAULT_MAX_CONCURRENT_JOBS) +} + +pub fn jobs_enabled(agent: Option<&Agent>, app: &AppConfig) -> bool { + app.function_calling_support && effective_max_concurrent_jobs(agent, app) > 0 +} + fn print_asset_names(kind: &str, names: &[String]) -> Result<()> { if names.is_empty() { println!("No {kind} found."); @@ -308,14 +324,24 @@ pub struct RequestContext { pub tool_scope: ToolScope, + 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, pub inbox: Option>, pub escalation_queue: Option>, + pub notification_queue: Arc, pub current_depth: usize, pub auto_continue_count: usize, - pub pending_agents_guardrail_count: u32, + pub pending_tasks_guardrail_count: u32, pub todo_list: TodoList, pub skill_registry: SkillRegistry, pub last_continuation_response: Option, @@ -341,14 +367,17 @@ impl RequestContext { agent: None, 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, inbox: None, escalation_queue: None, + notification_queue: Arc::new(NotificationQueue::new()), current_depth: 0, auto_continue_count: 0, - pending_agents_guardrail_count: 0, + pending_tasks_guardrail_count: 0, todo_list: TodoList::default(), skill_registry: SkillRegistry::default(), last_continuation_response: None, @@ -400,14 +429,17 @@ impl RequestContext { mcp_runtime, tool_tracker: ToolCallTracker::default(), }, + declared_function_names: Default::default(), + node_job_scope: None, supervisor: None, parent_supervisor: None, self_agent_id: None, inbox: None, escalation_queue: None, + notification_queue: Arc::new(NotificationQueue::new()), current_depth: 0, auto_continue_count: 0, - pending_agents_guardrail_count: 0, + pending_tasks_guardrail_count: 0, todo_list: TodoList::default(), skill_registry: SkillRegistry::default(), last_continuation_response: None, @@ -446,14 +478,17 @@ impl RequestContext { agent: self.agent.clone(), 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(), inbox: self.inbox.clone(), escalation_queue: self.escalation_queue.clone(), + notification_queue: self.notification_queue.clone(), current_depth: self.current_depth, auto_continue_count: 0, - pending_agents_guardrail_count: 0, + pending_tasks_guardrail_count: 0, todo_list: self.todo_list.clone(), skill_registry: self.skill_registry.clone(), last_continuation_response: None, @@ -490,14 +525,17 @@ impl RequestContext { mcp_runtime: McpRuntime::default(), 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), inbox: Some(inbox), escalation_queue: parent.escalation_queue.clone(), + notification_queue: Arc::new(NotificationQueue::new()), current_depth, auto_continue_count: 0, - pending_agents_guardrail_count: 0, + pending_tasks_guardrail_count: 0, todo_list: TodoList::default(), skill_registry: SkillRegistry::default(), last_continuation_response: None, @@ -894,6 +932,15 @@ impl RequestContext { } pub fn before_chat_completion(&mut self, input: &Input) -> Result<()> { + // `job__start` validates against exactly what was declared to the + // model for THIS request; refresh it every time. + // + // This is necessary to prevent the model from invoking functions it + // otherwise wouldn't have access to by going through the free `tool` + // argument of `job__start`. If a function is disabled, the model + // shouldn't be able to invoke it at all in any way. This prevents + // that backdoor. + self.declared_function_names = input.declared_function_names(); self.last_message = Some(LastMessage::new(input.clone(), String::new())); Ok(()) } @@ -1302,6 +1349,7 @@ impl RequestContext { && !v.name.starts_with("memory__") && !v.name.starts_with("skill__") && !v.name.starts_with("rag__") + && !v.name.starts_with("job__") }) .map(|v| v.name.clone()) .collect() @@ -2119,7 +2167,8 @@ impl RequestContext { && v.name.starts_with(SKILL_FUNCTION_PREFIX)) || (self.auto_continue_config().enabled && v.name.starts_with(TODO_FUNCTION_PREFIX)) - || v.name.starts_with(RAG_FUNCTION_PREFIX)) + || v.name.starts_with(RAG_FUNCTION_PREFIX) + || v.name.starts_with(JOB_FUNCTION_PREFIX)) && !existing.contains(&v.name) }) .cloned() @@ -2143,9 +2192,10 @@ impl RequestContext { && v.name.starts_with(SKILL_FUNCTION_PREFIX)) || v.name.starts_with(USER_FUNCTION_PREFIX) || v.name.starts_with(TODO_FUNCTION_PREFIX) - || v.name.starts_with(SUPERVISOR_FUNCTION_PREFIX) + || v.name.starts_with(AGENT_FUNCTION_PREFIX) || v.name.starts_with(MEMORY_FUNCTION_PREFIX) || v.name.starts_with(RAG_FUNCTION_PREFIX) + || v.name.starts_with(JOB_FUNCTION_PREFIX) }); } @@ -2278,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 @@ -2286,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()) { @@ -3858,6 +3941,9 @@ impl RequestContext { { functions.append_rag_query_functions(); } + if self.agent.is_none() && jobs_enabled(None, app) { + functions.append_job_functions(); + } let tool_tracker = self.tool_scope.tool_tracker.clone(); self.tool_scope = ToolScope { @@ -4136,11 +4222,21 @@ impl RequestContext { ); } - let should_init_supervisor = agent.can_spawn_agents(); - let max_concurrent = agent.max_concurrent_agents(); + let jobs_enabled = jobs_enabled(Some(&agent), app); + let should_init_supervisor = agent.can_spawn_agents() || jobs_enabled; + let max_concurrent_agents = if agent.can_spawn_agents() { + agent.max_concurrent_agents() + } else { + 0 + }; let max_depth = agent.max_agent_depth(); - let supervisor = should_init_supervisor - .then(|| Arc::new(RwLock::new(Supervisor::new(max_concurrent, max_depth)))); + let max_jobs = effective_max_concurrent_jobs(Some(&agent), app); + let supervisor = should_init_supervisor.then(|| { + Arc::new(RwLock::new( + Supervisor::new(max_concurrent_agents, max_depth) + .with_max_concurrent_jobs(max_jobs), + )) + }); self.rag = agent.rag(); // Keep `rag_key` in lockstep with `rag`. Agent RAGs are cached under @@ -4152,9 +4248,13 @@ impl RequestContext { .is_some() .then(|| RagKey::Agent(agent.name().to_string())); self.agent = Some(agent); + if let Some(old) = self.supervisor.as_ref() { + old.read().cancel_recursive(); + } self.supervisor = supervisor; self.inbox = None; self.escalation_queue = None; + self.notification_queue = Arc::new(NotificationQueue::new()); self.self_agent_id = None; self.parent_supervisor = None; self.current_depth = 0; @@ -4178,6 +4278,9 @@ impl RequestContext { if self.working_mode.is_repl() { functions.append_user_interaction_functions(); } + if jobs_enabled(None, app) { + functions.append_job_functions(); + } let tool_tracker = self.tool_scope.tool_tracker.clone(); self.tool_scope = ToolScope { functions, @@ -4194,9 +4297,10 @@ impl RequestContext { self.self_agent_id = None; self.inbox = None; self.escalation_queue = None; + self.notification_queue = Arc::new(NotificationQueue::new()); self.current_depth = 0; self.auto_continue_count = 0; - self.pending_agents_guardrail_count = 0; + self.pending_tasks_guardrail_count = 0; self.todo_list = TodoList::default(); self.rag.take(); // Cleared alongside `rag` so the pair never disagrees: an agent RAG is @@ -4692,18 +4796,22 @@ mod tests { use super::*; use crate::config::AppState; use crate::config::agent::AgentConfig; + use crate::function::jobs::RingBuf; use crate::function::{ToolCall, skill}; use crate::mcp::{McpServer, McpServerFeatures, McpServersConfig, McpTransportType}; + use crate::supervisor::{ + AgentExitStatus, AgentHandle, AgentResult, JobHandle, JobResult, JobState, JobStatus, + }; use crate::utils; use crate::utils::get_env_name; use crate::vault::Vault; use rmcp::model::PromptArgument; use serde_json::json; use serial_test::serial; - use std::env; use std::fs::{create_dir_all, remove_dir_all, write}; use std::path::PathBuf; - use std::time::{SystemTime, UNIX_EPOCH}; + use std::time::{Instant, SystemTime, UNIX_EPOCH}; + use std::{env, mem}; struct TestConfigDirGuard { key: String, @@ -4755,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(), @@ -5217,6 +5334,171 @@ mod tests { assert_eq!(ctx.rag_key, None); } + #[test] + fn effective_max_concurrent_jobs_resolution_precedence() { + let mut app = AppConfig::default(); + assert_eq!(effective_max_concurrent_jobs(None, &app), 5); + + app.max_concurrent_jobs = Some(9); + assert_eq!(effective_max_concurrent_jobs(None, &app), 9); + + let agent = Agent::test_new(AgentConfig { + max_concurrent_jobs: Some(2), + ..AgentConfig::default() + }); + assert_eq!(effective_max_concurrent_jobs(Some(&agent), &app), 2); + } + + #[test] + fn jobs_enabled_requires_function_calling_and_nonzero_capacity() { + let mut app = AppConfig::default(); + assert!(jobs_enabled(None, &app)); + + app.max_concurrent_jobs = Some(0); + assert!(!jobs_enabled(None, &app)); + + app.max_concurrent_jobs = None; + app.function_calling_support = false; + assert!(!jobs_enabled(None, &app)); + + app.function_calling_support = true; + let agent = Agent::test_new(AgentConfig { + max_concurrent_jobs: Some(0), + ..AgentConfig::default() + }); + assert!(!jobs_enabled(Some(&agent), &app)); + } + + #[test] + #[serial] + fn use_agent_cancels_previous_supervisor() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = create_test_ctx(); + let app = ctx.app.config.clone(); + let agent_name = format!( + "test_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 old_sig = utils::create_abort_signal(); + + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + let join_handle = tokio::spawn(async { + Ok(AgentResult { + id: "a1".into(), + agent_name: "explore".into(), + output: String::new(), + exit_status: AgentExitStatus::Completed, + }) + }); + let handle = AgentHandle { + id: "a1".to_string(), + agent_name: "explore".to_string(), + depth: 1, + inbox: Arc::new(Inbox::new()), + abort_signal: old_sig.clone(), + join_handle, + child_supervisor: None, + }; + let old_sup = Arc::new(RwLock::new(Supervisor::new(4, 3))); + old_sup.write().register(handle).unwrap(); + ctx.supervisor = Some(old_sup); + + ctx.use_agent(&app, &agent_name, None, utils::create_abort_signal()) + .await + .unwrap(); + }); + + assert!(old_sig.aborted()); + assert!(ctx.supervisor.is_some()); + } + + #[test] + #[serial] + fn use_agent_inits_job_capable_supervisor_without_spawning() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = create_test_ctx(); + let app = ctx.app.config.clone(); + let agent_name = format!( + "test_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(); + + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + ctx.use_agent(&app, &agent_name, None, utils::create_abort_signal()) + .await + .unwrap(); + }); + + let supervisor = ctx.supervisor.as_ref().expect("supervisor for jobs"); + let supervisor = supervisor.read(); + assert_eq!(supervisor.max_concurrent(), 0); + assert_eq!(supervisor.max_concurrent_jobs(), 5); + } + + #[test] + #[serial] + fn use_agent_skips_supervisor_when_jobs_disabled() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = create_test_ctx(); + let mut app = ctx.app.config.as_ref().clone(); + app.max_concurrent_jobs = Some(0); + let agent_name = format!( + "test_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(); + + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + ctx.use_agent(&app, &agent_name, None, utils::create_abort_signal()) + .await + .unwrap(); + }); + + assert!(ctx.supervisor.is_none()); + } + #[test] fn current_depth_default_is_zero() { let ctx = create_test_ctx(); @@ -5248,6 +5530,32 @@ mod tests { assert!(ctx.root_escalation_queue().is_none()); } + #[test] + fn new_for_child_gets_fresh_notification_queue() { + let parent = create_test_ctx(); + let child = RequestContext::new_for_child( + Arc::clone(&parent.app), + &parent, + 1, + Arc::new(Inbox::new()), + "agent_test_1".to_string(), + ); + assert!( + !Arc::ptr_eq(&parent.notification_queue, &child.notification_queue), + "each child owns its notifications; a shared queue would race drains" + ); + } + + #[test] + fn fork_for_branch_shares_notification_queue() { + let ctx = create_test_ctx(); + let branch = ctx.fork_for_branch(); + assert!(Arc::ptr_eq( + &ctx.notification_queue, + &branch.notification_queue + )); + } + 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") } @@ -5520,6 +5828,146 @@ mod tests { assert!(ctx.select_functions(&role).is_none()); } + #[test] + fn select_functions_hides_job_functions_without_backgroundable_tools() { + let mut ctx = create_test_ctx(); + ctx.tool_scope.functions.append_job_functions(); + + assert!( + ctx.select_functions(&Role::default()).is_none(), + "job__ tools must not be declared when nothing backgroundable is declared" + ); + } + + #[test] + 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!["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 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" + ); + } + + #[test] + fn concrete_tool_names_excludes_job_functions() { + let mut ctx = create_test_ctx(); + ctx.tool_scope.functions.append_job_functions(); + + assert!(ctx.concrete_tool_names().is_empty()); + } + + #[test] + fn before_chat_completion_refreshes_declared_function_names() { + 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()])); + let input = Input::from_str(&ctx, "hello", Some(role)).unwrap(); + ctx.before_chat_completion(&input).unwrap(); + + 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 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!( + ctx.declared_function_names.is_empty(), + "stash must be refreshed on every request" + ); + } + + #[test] + #[serial] + fn rebuild_tool_scope_gates_job_functions_on_jobs_enabled() { + let _guard = TestConfigDirGuard::new(); + let app_state = app_state_with_mcp_config(false, &[]); + let mut ctx = RequestContext::new(app_state, WorkingMode::Cmd); + let app = ctx.app.config.clone(); + let abort = utils::create_abort_signal(); + + run_async(ctx.rebuild_tool_scope(&app, None, abort.clone())).unwrap(); + assert!(ctx.tool_scope.functions.contains("job__start")); + + let jobs_off = AppConfig { + max_concurrent_jobs: Some(0), + ..(*app).clone() + }; + run_async(ctx.rebuild_tool_scope(&jobs_off, None, abort.clone())).unwrap(); + assert!( + !ctx.tool_scope + .functions + .declarations() + .iter() + .any(|f| f.name.starts_with("job__")) + ); + + let fc_off = AppConfig { + function_calling_support: false, + ..(*app).clone() + }; + run_async(ctx.rebuild_tool_scope(&fc_off, None, abort)).unwrap(); + assert!( + !ctx.tool_scope + .functions + .declarations() + .iter() + .any(|f| f.name.starts_with("job__")) + ); + } + + #[test] + #[serial] + fn exit_agent_rebuild_retains_job_functions_when_enabled() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = create_test_ctx(); + let app = ctx.app.config.clone(); + + ctx.exit_agent(&app).unwrap(); + assert!(ctx.tool_scope.functions.contains("job__start")); + + let jobs_off = AppConfig { + max_concurrent_jobs: Some(0), + ..(*app).clone() + }; + ctx.exit_agent(&jobs_off).unwrap(); + assert!(!ctx.tool_scope.functions.contains("job__start")); + } + #[test] fn select_functions_all_enabled_tools_returns_all_non_mcp() { let mut ctx = create_test_ctx(); @@ -5686,6 +6134,45 @@ mod tests { ); } + #[test] + #[serial] + fn select_functions_preserves_job_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_job_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(); + 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()])); + + let fns = ctx.select_functions(&role).unwrap(); + 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 that declares a backgroundable tool, got: {names:?}" + ); + assert!(names.contains(&"job__collect")); + } + #[test] fn fork_for_branch_clones_skill_registry() { let mut ctx = create_test_ctx(); @@ -7150,4 +7637,308 @@ mod tests { "global config" ); } + + #[test] + 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__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_hides_job_tools_under_empty_agent_filter() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = create_test_ctx(); + let app = ctx.app.config.clone(); + let agent_name = format!( + "test_job_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![])); + + let names: Vec = ctx + .select_functions(&role) + .unwrap_or_default() + .iter() + .map(|f| f.name.clone()) + .collect(); + assert!( + !names.iter().any(|n| n.starts_with("job__")), + "job__ tools must be hidden under an empty agent filter, got: {names:?}" + ); + } + + #[test] + #[serial] + fn select_functions_when_jobs_disabled_is_byte_identical_to_no_jobs_baseline() { + let _guard = TestConfigDirGuard::new(); + let app_state = app_state_with_mcp_config(false, &[]); + let mut ctx = RequestContext::new(app_state, WorkingMode::Repl); + let app = ctx.app.config.clone(); + let abort = utils::create_abort_signal(); + + let mut role = Role::new("r", "p"); + role.set_enabled_tools(Some(vec!["all".to_string()])); + + let jobs_off = AppConfig { + max_concurrent_jobs: Some(0), + ..(*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__"), + "no job__ declarations may leak when jobs are disabled, got: {without_jobs}" + ); + + 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 + .into_iter() + .filter(|f| !f.name.starts_with("job__")) + .collect(); + + assert_eq!( + without_jobs, + serde_json::to_string(&Some(stripped)).unwrap(), + "jobs-disabled tool list must be byte-identical to the jobs-enabled list minus job__ declarations" + ); + } + + #[test] + fn select_functions_returns_none_when_no_tools_enabled_and_jobs_disabled() { + let app_state = { + let config = AppConfig { + max_concurrent_jobs: Some(0), + ..AppConfig::default() + }; + Arc::new(AppState { + config: Arc::new(config), + vault: Arc::new(Vault::default()), + mcp_factory: Arc::new(McpFactory::default()), + rag_cache: Arc::new(RagCache::default()), + mcp_config: None, + mcp_log_path: None, + mcp_registry: None, + functions: Functions::default(), + }) + }; + let ctx = RequestContext::new(app_state, WorkingMode::Cmd); + assert!(ctx.select_functions(&Role::default()).is_none()); + } + + #[test] + 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(); + + for name in [ + "job__start", + "job__check", + "job__collect", + "job__cancel", + "job__list", + ] { + assert!( + info.contains(name), + "expected {name} in output, got: {info}" + ); + } + } + + fn make_running_job(abort_signal: utils::AbortSignal) -> JobHandle { + // Leak the runtime so the spawned task is never polled and the job + // stays running for the duration of the test. + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let join_handle = rt.spawn(async { + Ok(JobResult { + output: serde_json::Value::Null, + exit_code: Some(0), + output_bytes_captured: 0, + }) + }); + mem::forget(rt); + JobHandle { + id: "j1".to_string(), + tool: "execute_command".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal, + state: Arc::new(parking_lot::Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })), + output_buf: Arc::new(parking_lot::Mutex::new(RingBuf::default())), + no_change_checks: 0, + last_check_state: None, + } + } + + #[test] + #[serial] + fn use_agent_cancels_running_jobs_of_previous_supervisor() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = create_test_ctx(); + let app = ctx.app.config.clone(); + let agent_name = format!( + "test_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 job_sig = utils::create_abort_signal(); + let old_sup = Arc::new(RwLock::new( + Supervisor::new(4, 3).with_max_concurrent_jobs(1), + )); + old_sup + .write() + .register(make_running_job(job_sig.clone())) + .unwrap(); + ctx.supervisor = Some(old_sup); + + run_async(ctx.use_agent(&app, &agent_name, None, utils::create_abort_signal())).unwrap(); + + assert!( + job_sig.aborted(), + "running jobs of the previous supervisor must be cancelled" + ); + assert!(ctx.supervisor.is_some()); + } + + #[test] + #[serial] + fn exit_agent_cancels_running_jobs() { + let _guard = TestConfigDirGuard::new(); + let mut ctx = create_test_ctx(); + let app = ctx.app.config.clone(); + + let job_sig = utils::create_abort_signal(); + let sup = Arc::new(RwLock::new( + Supervisor::new(4, 3).with_max_concurrent_jobs(1), + )); + sup.write() + .register(make_running_job(job_sig.clone())) + .unwrap(); + ctx.agent = Some(Agent::test_new(AgentConfig::default())); + ctx.supervisor = Some(sup); + + ctx.exit_agent(&app).unwrap(); + + assert!(job_sig.aborted(), "exit_agent must cancel running jobs"); + assert!(ctx.supervisor.is_none()); + } + + #[test] + fn toggle_tool_rejects_job_tools_as_unknown() { + let mut ctx = create_test_ctx(); + ctx.tool_scope.functions.append_job_functions(); + + for action in ["enable", "disable"] { + let err = ctx.toggle_tool(action, "job__start").unwrap_err(); + assert!( + err.to_string().contains("Unknown tool 'job__start'"), + "expected job__start to be rejected on {action}, got: {err}" + ); + } + } } diff --git a/src/function/supervisor.rs b/src/function/agents.rs similarity index 65% rename from src/function/supervisor.rs rename to src/function/agents.rs index 32c5d48..e7803ba 100644 --- a/src/function/supervisor.rs +++ b/src/function/agents.rs @@ -1,17 +1,19 @@ use super::{FunctionDeclaration, JsonSchema}; use crate::client::{Model, ModelType, call_chat_completions}; use crate::config::{ - Agent, AppState, Input, RequestContext, Role, RoleLike, list_agents_with_descriptions, + Agent, AppState, Input, RequestContext, Role, RoleLike, effective_max_concurrent_jobs, + jobs_enabled, list_agents_with_descriptions, }; use crate::supervisor::mailbox::{Envelope, EnvelopePayload, Inbox}; -use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult, Supervisor}; +use crate::supervisor::notification::agent_notification; +use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult, Supervisor, TaskKind}; use crate::utils::{AbortSignal, create_abort_signal, wait_abort_signal}; use crate::graph; use anyhow::{Context, Result, anyhow, bail}; use chrono::Utc; use indexmap::IndexMap; -use log::debug; +use log::{debug, warn}; use parking_lot::RwLock; use serde_json::{Value, json}; use std::pin::Pin; @@ -21,9 +23,9 @@ use tokio::time; use tokio::time::Instant; use uuid::Uuid; -pub const SUPERVISOR_FUNCTION_PREFIX: &str = "agent__"; +pub const AGENT_FUNCTION_PREFIX: &str = "agent__"; -pub const PENDING_AGENTS_GUARDRAIL_MAX: u32 = 3; +pub const PENDING_TASKS_GUARDRAIL_MAX: u32 = 3; fn agent_permitted(whitelist: Option<&[String]>, target: &str) -> bool { match whitelist { @@ -32,60 +34,151 @@ fn agent_permitted(whitelist: Option<&[String]>, target: &str) -> bool { } } +fn is_job_task(supervisor: Option<&Arc>>, id: &str) -> bool { + id.starts_with("job_") || supervisor.is_some_and(|sup| sup.read().has_job(id)) +} + +fn job_id_teaching_error(id: &str) -> Value { + json!({ + "status": "error", + "message": format!( + "'{id}' is a background job, not an agent — use job__check / job__collect / job__cancel" + ), + }) +} + pub enum GuardrailAction { NoAction, Inject(String), ForceTerminate(Vec), } -pub fn pending_agent_ids(ctx: &RequestContext) -> Vec { +pub struct PendingTask { + pub id: String, + pub kind: TaskKind, + pub finished: bool, +} + +pub fn pending_tasks(ctx: &RequestContext) -> Vec { let Some(sup) = ctx.supervisor.as_ref() else { return Vec::new(); }; - let sup = sup.read(); - sup.list_agents() + let mut tasks: Vec = sup + .read() + .list_tasks() .into_iter() - .filter_map(|(id, _)| match sup.is_finished(id) { - Some(false) => Some(id.to_string()), - _ => None, + .map(|(id, kind, finished)| PendingTask { + id: id.to_string(), + kind, + finished, }) - .collect() + .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 } -pub fn build_pending_agents_guardrail_prompt(ids: &[String]) -> String { - let count = ids.len(); - let id_list = ids - .iter() - .map(|id| format!("- {id}")) - .collect::>() - .join("\n"); +pub fn build_pending_tasks_guardrail_prompt(tasks: &[PendingTask]) -> String { + let running: Vec<&PendingTask> = tasks.iter().filter(|t| !t.finished).collect(); + let finished: Vec<&PendingTask> = tasks.iter().filter(|t| t.finished).collect(); + + let mut sections = Vec::new(); + if !running.is_empty() { + let id_list = running + .iter() + .map(|t| { + let (kind, collect, cancel) = match t.kind { + TaskKind::Agent => ("agent", "agent__collect", "agent__cancel"), + TaskKind::Job => ("job", "job__collect", "job__cancel"), + }; + format!( + "- {id} ({kind}): call `{collect}` (blocks until done, returns output) or \ + `{cancel}` (discards)", + id = t.id + ) + }) + .collect::>() + .join("\n"); + sections.push(format!( + "Still running ({count}):\n{id_list}\n\nThese will be abandoned if your turn ends \ + now. You MUST reclaim each one before ending your turn. Do NOT emit a text-only \ + response expecting them to 'report back' — they will not.", + count = running.len() + )); + } + + if !finished.is_empty() { + let cmd_list = finished + .iter() + .map(|t| { + let collect = match t.kind { + TaskKind::Agent => "agent__collect", + TaskKind::Job => "job__collect", + }; + format!("- `{collect} --id {id}`", id = t.id) + }) + .collect::>() + .join("\n"); + sections.push(format!( + "Completed but UNCOLLECTED — collect NOW ({count}):\n{cmd_list}\n\nCollect returns \ + instantly on a finished task. Their results are LOST if your turn ends without \ + collecting.", + count = finished.len() + )); + } + format!( - "[SYSTEM GUARDRAIL] You attempted to end your turn while {count} spawned background agent(s) \ - are still running:\n{id_list}\n\nThese agents will be abandoned if your turn ends now. You MUST \ - reclaim each one before ending your turn. For each agent: call `agent__collect` (blocks until \ - done, returns output) or `agent__cancel` (discards). Do NOT emit a text-only response \ - expecting them to 'report back' — they will not." + "[SYSTEM GUARDRAIL] You attempted to end your turn with {count} unreclaimed background \ + task(s).\n\n{body}", + count = tasks.len(), + body = sections.join("\n\n") ) } -pub fn check_pending_agents_guardrail(ctx: &mut RequestContext) -> GuardrailAction { - let pending = pending_agent_ids(ctx); +pub fn check_pending_tasks_guardrail(ctx: &mut RequestContext) -> GuardrailAction { + let pending = pending_tasks(ctx); if pending.is_empty() { - ctx.pending_agents_guardrail_count = 0; + ctx.pending_tasks_guardrail_count = 0; return GuardrailAction::NoAction; } - if ctx.pending_agents_guardrail_count >= PENDING_AGENTS_GUARDRAIL_MAX { + if ctx.pending_tasks_guardrail_count >= PENDING_TASKS_GUARDRAIL_MAX { if let Some(sup) = ctx.supervisor.as_ref().cloned() { sup.read().cancel_recursive(); + let finished: Vec<&PendingTask> = pending.iter().filter(|t| t.finished).collect(); + if !finished.is_empty() { + let ids: Vec<&str> = finished.iter().map(|t| t.id.as_str()).collect(); + warn!( + "Turn-end guardrail: discarding uncollected result(s) for finished task(s) \ + after max reminders: {ids:?}" + ); + let mut sup = sup.write(); + for task in &finished { + match task.kind { + TaskKind::Agent => { + let _ = sup.take(&task.id); + } + TaskKind::Job => { + let _ = sup.take_job(&task.id); + } + } + } + } } - ctx.pending_agents_guardrail_count = 0; + ctx.pending_tasks_guardrail_count = 0; - return GuardrailAction::ForceTerminate(pending); + return GuardrailAction::ForceTerminate(pending.into_iter().map(|t| t.id).collect()); } - ctx.pending_agents_guardrail_count += 1; - let mut prompt = build_pending_agents_guardrail_prompt(&pending); + ctx.pending_tasks_guardrail_count += 1; + let mut prompt = build_pending_tasks_guardrail_prompt(&pending); if let Some(queue) = ctx.root_escalation_queue() && queue.has_pending() { @@ -100,8 +193,9 @@ pub fn check_pending_agents_guardrail(ctx: &mut RequestContext) -> GuardrailActi pub fn escalation_function_declarations() -> Vec { vec![FunctionDeclaration { - name: format!("{SUPERVISOR_FUNCTION_PREFIX}reply_escalation"), - description: "Reply to a pending escalation from a child agent. The child is blocked waiting for this reply. Use this after seeing pending_escalations notifications.".to_string(), + name: format!("{AGENT_FUNCTION_PREFIX}reply_escalation"), + description: "Reply to a pending escalation from a child agent. The child is blocked waiting for this reply. \ + Use this after seeing pending_escalations notifications.".to_string(), parameters: JsonSchema { type_value: Some("object".to_string()), properties: Some(IndexMap::from([ @@ -117,7 +211,8 @@ pub fn escalation_function_declarations() -> Vec { "reply".to_string(), JsonSchema { type_value: Some("string".to_string()), - description: Some("Your answer to the child agent's question. For ask/confirm questions, use the exact option text. For input questions, provide the text response.".into()), + description: Some("Your answer to the child agent's question. For ask/confirm questions, use \ + the exact option text. For input questions, provide the text response.".into()), ..Default::default() }, ), @@ -129,10 +224,10 @@ pub fn escalation_function_declarations() -> Vec { }] } -pub fn supervisor_function_declarations() -> Vec { +pub fn agent_function_declarations() -> Vec { vec![ FunctionDeclaration { - name: format!("{SUPERVISOR_FUNCTION_PREFIX}spawn"), + name: format!("{AGENT_FUNCTION_PREFIX}spawn"), description: "Spawn a subagent to run in the background. Returns an `id` immediately so you can continue \ working in parallel. CRITICAL: every spawned agent MUST be reclaimed before you end your \ turn — call `agent__collect` to retrieve its output, or `agent__cancel` if you no longer \ @@ -172,8 +267,9 @@ pub fn supervisor_function_declarations() -> Vec { agent: false, }, FunctionDeclaration { - name: format!("{SUPERVISOR_FUNCTION_PREFIX}check"), - description: "Check if a spawned agent has finished. Non-blocking; returns PENDING if still running, or the result if complete.".to_string(), + name: format!("{AGENT_FUNCTION_PREFIX}check"), + description: "Non-blocking status probe: reports whether a spawned agent is still running or finished. \ + NEVER returns or consumes the result — when finished, call agent__collect to retrieve it.".to_string(), parameters: JsonSchema { type_value: Some("object".to_string()), properties: Some(IndexMap::from([( @@ -190,7 +286,7 @@ pub fn supervisor_function_declarations() -> Vec { agent: false, }, FunctionDeclaration { - name: format!("{SUPERVISOR_FUNCTION_PREFIX}collect"), + name: format!("{AGENT_FUNCTION_PREFIX}collect"), description: "Block until the named spawned agent finishes and return its result. This is your primary \ wait primitive — it pauses your execution until the agent completes (or you are interrupted). \ Call this for every agent you spawned before ending your turn. Do NOT end your turn assuming \ @@ -212,7 +308,7 @@ pub fn supervisor_function_declarations() -> Vec { agent: false, }, FunctionDeclaration { - name: format!("{SUPERVISOR_FUNCTION_PREFIX}list_running"), + name: format!("{AGENT_FUNCTION_PREFIX}list_running"), description: "List all subagents YOU have spawned that are still tracked by the supervisor, with their \ status. Use this to see which of your background agents are still active. To discover which \ agent types you can spawn in the first place, use `agent__list_available` instead.".to_string(), @@ -224,7 +320,7 @@ pub fn supervisor_function_declarations() -> Vec { agent: false, }, FunctionDeclaration { - name: format!("{SUPERVISOR_FUNCTION_PREFIX}list_available"), + name: format!("{AGENT_FUNCTION_PREFIX}list_available"), description: "List all agent types installed and available to spawn (name + description). Use this to \ discover what specialists exist before calling `agent__spawn` — especially when you're unsure \ which agent to delegate to. This is the discovery counterpart to `agent__list_running` \ @@ -237,7 +333,7 @@ pub fn supervisor_function_declarations() -> Vec { agent: false, }, FunctionDeclaration { - name: format!("{SUPERVISOR_FUNCTION_PREFIX}cancel"), + name: format!("{AGENT_FUNCTION_PREFIX}cancel"), description: "Cancel a running subagent by its ID. Use this when an agent's output is no longer needed \ (e.g. you changed direction, or you're about to end your turn and don't want to wait). \ Cancellation cascades: all of the cancelled agent's own descendants are also cancelled. This \ @@ -258,7 +354,7 @@ pub fn supervisor_function_declarations() -> Vec { agent: false, }, FunctionDeclaration { - name: format!("{SUPERVISOR_FUNCTION_PREFIX}task_create"), + name: format!("{AGENT_FUNCTION_PREFIX}task_create"), description: "Create a task in the task queue. Returns the task ID.".to_string(), parameters: JsonSchema { type_value: Some("object".to_string()), @@ -314,7 +410,7 @@ pub fn supervisor_function_declarations() -> Vec { agent: false, }, FunctionDeclaration { - name: format!("{SUPERVISOR_FUNCTION_PREFIX}task_list"), + name: format!("{AGENT_FUNCTION_PREFIX}task_list"), description: "List all tasks in the task queue with their status and dependencies.".to_string(), parameters: JsonSchema { type_value: Some("object".to_string()), @@ -324,7 +420,7 @@ pub fn supervisor_function_declarations() -> Vec { agent: false, }, FunctionDeclaration { - name: format!("{SUPERVISOR_FUNCTION_PREFIX}task_complete"), + name: format!("{AGENT_FUNCTION_PREFIX}task_complete"), description: "Mark a task as completed. Returns any newly unblocked task IDs.".to_string(), parameters: JsonSchema { type_value: Some("object".to_string()), @@ -342,7 +438,7 @@ pub fn supervisor_function_declarations() -> Vec { agent: false, }, FunctionDeclaration { - name: format!("{SUPERVISOR_FUNCTION_PREFIX}task_fail"), + name: format!("{AGENT_FUNCTION_PREFIX}task_fail"), description: "Mark a task as failed. Dependents will remain blocked.".to_string(), parameters: JsonSchema { type_value: Some("object".to_string()), @@ -365,7 +461,7 @@ pub fn supervisor_function_declarations() -> Vec { pub fn teammate_function_declarations() -> Vec { vec![ FunctionDeclaration { - name: format!("{SUPERVISOR_FUNCTION_PREFIX}send_message"), + name: format!("{AGENT_FUNCTION_PREFIX}send_message"), description: "Send a text message to a sibling or child agent's inbox. Use to share cross-cutting findings or coordinate with teammates.".to_string(), parameters: JsonSchema { type_value: Some("object".to_string()), @@ -393,7 +489,7 @@ pub fn teammate_function_declarations() -> Vec { agent: false, }, FunctionDeclaration { - name: format!("{SUPERVISOR_FUNCTION_PREFIX}check_inbox"), + name: format!("{AGENT_FUNCTION_PREFIX}check_inbox"), description: "Check for and drain all pending messages in your inbox from sibling agents or your parent.".to_string(), parameters: JsonSchema { type_value: Some("object".to_string()), @@ -405,13 +501,13 @@ pub fn teammate_function_declarations() -> Vec { ] } -pub async fn handle_supervisor_tool( +pub async fn handle_agent_tool( ctx: &mut RequestContext, cmd_name: &str, args: &Value, ) -> Result { let action = cmd_name - .strip_prefix(SUPERVISOR_FUNCTION_PREFIX) + .strip_prefix(AGENT_FUNCTION_PREFIX) .unwrap_or(cmd_name); match action { @@ -428,7 +524,7 @@ pub async fn handle_supervisor_tool( "task_complete" => handle_task_complete(ctx, args).await, "task_fail" => handle_task_fail(ctx, args), "reply_escalation" => handle_reply_escalation(ctx, args), - _ => bail!("Unknown supervisor action: {action}"), + _ => bail!("Unknown agent action: {action}"), } } @@ -475,10 +571,10 @@ pub fn run_child_agent( } if tool_results.is_empty() { - match check_pending_agents_guardrail(&mut child_ctx) { + match check_pending_tasks_guardrail(&mut child_ctx) { GuardrailAction::NoAction => break, GuardrailAction::ForceTerminate(ids) => { - log::warn!( + warn!( "Pending-agent guardrail force-cancelled {} agent(s) after max reminders: {:?}", ids.len(), ids @@ -557,9 +653,15 @@ pub async fn run_agent_for_graph( let agent_mcp_servers = agent.mcp_server_names().to_vec(); let session = agent.agent_session().map(|v| v.to_string()); - let should_init_supervisor = agent.can_spawn_agents(); - let agent_max_concurrent = agent.max_concurrent_agents(); + let child_jobs_enabled = jobs_enabled(Some(&agent), app_config.as_ref()); + let should_init_supervisor = agent.can_spawn_agents() || child_jobs_enabled; + let agent_max_concurrent_subagents = if agent.can_spawn_agents() { + agent.max_concurrent_agents() + } else { + 0 + }; let agent_max_depth = agent.max_agent_depth(); + let agent_max_jobs = effective_max_concurrent_jobs(Some(&agent), app_config.as_ref()); let mut child_ctx = RequestContext::new_for_child( Arc::clone(&child_app_state), @@ -571,10 +673,10 @@ pub async fn run_agent_for_graph( child_ctx.rag = agent.rag(); child_ctx.agent = Some(agent); if should_init_supervisor { - child_ctx.supervisor = Some(Arc::new(RwLock::new(Supervisor::new( - agent_max_concurrent, - agent_max_depth, - )))); + child_ctx.supervisor = Some(Arc::new(RwLock::new( + Supervisor::new(agent_max_concurrent_subagents, agent_max_depth) + .with_max_concurrent_jobs(agent_max_jobs), + ))); } if let Some(session) = session { @@ -736,9 +838,15 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result { let agent_mcp_servers = agent.mcp_server_names().to_vec(); let session = agent.agent_session().map(|v| v.to_string()); - let should_init_supervisor = agent.can_spawn_agents(); - let max_concurrent = agent.max_concurrent_agents(); + let child_jobs_enabled = jobs_enabled(Some(&agent), app_config.as_ref()); + let should_init_supervisor = agent.can_spawn_agents() || child_jobs_enabled; + let max_concurrent_agents = if agent.can_spawn_agents() { + agent.max_concurrent_agents() + } else { + 0 + }; let max_depth = agent.max_agent_depth(); + let max_jobs = effective_max_concurrent_jobs(Some(&agent), app_config.as_ref()); let mut child_ctx = RequestContext::new_for_child( Arc::clone(&child_app_state), ctx, @@ -749,10 +857,9 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result { child_ctx.rag = agent.rag(); child_ctx.agent = Some(agent); if should_init_supervisor { - child_ctx.supervisor = Some(Arc::new(RwLock::new(Supervisor::new( - max_concurrent, - max_depth, - )))); + child_ctx.supervisor = Some(Arc::new(RwLock::new( + Supervisor::new(max_concurrent_agents, max_depth).with_max_concurrent_jobs(max_jobs), + ))); } if let Some(session) = session { @@ -773,25 +880,34 @@ async fn handle_spawn(ctx: &mut RequestContext, args: &Value) -> Result { let spawn_agent_id = agent_id.clone(); let spawn_agent_name = agent_name.clone(); let spawn_abort = child_abort.clone(); + let spawn_notifications = Arc::clone(&ctx.notification_queue); let child_supervisor = child_ctx.supervisor.clone(); let join_handle = tokio::spawn(async move { let result = run_child_agent(child_ctx, input, spawn_abort).await; - match result { - Ok(output) => Ok(AgentResult { + let agent_result = match result { + Ok(output) => AgentResult { id: spawn_agent_id, agent_name: spawn_agent_name, output, exit_status: AgentExitStatus::Completed, - }), - Err(e) => Ok(AgentResult { + }, + Err(e) => AgentResult { id: spawn_agent_id, agent_name: spawn_agent_name, output: String::new(), exit_status: AgentExitStatus::Failed(e.to_string()), - }), - } + }, + }; + let success = agent_result.exit_status == AgentExitStatus::Completed; + spawn_notifications.push(agent_notification( + &agent_result.id, + &agent_result.agent_name, + success, + )); + + Ok(agent_result) }); let handle = AgentHandle { @@ -841,7 +957,15 @@ async fn handle_check(ctx: &mut RequestContext, args: &Value) -> Result { }; match is_finished { - Some(true) => handle_collect(ctx, args).await, + Some(true) => Ok(json!({ + "status": "finished", + "id": id, + "message": format!( + "Agent '{id}' has finished; its result is ready and has NOT been consumed. \ + Call `agent__collect --id {id}` to retrieve it (returns instantly on a \ + finished agent). The handle stays registered until collected." + ), + })), Some(false) => { let mut result = json!({ "status": "pending", @@ -861,10 +985,16 @@ async fn handle_check(ctx: &mut RequestContext, args: &Value) -> Result { Ok(result) } - None => Ok(json!({ - "status": "error", - "message": format!("No agent found with id '{id}'") - })), + None => { + if is_job_task(ctx.supervisor.as_ref(), id) { + return Ok(job_id_teaching_error(id)); + } + + Ok(json!({ + "status": "error", + "message": format!("No agent found with id '{id}'") + })) + } } } @@ -883,6 +1013,10 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result let target_abort = { let sup = supervisor.read(); if sup.is_finished(id).is_none() { + if id.starts_with("job_") || sup.has_job(id) { + return Ok(job_id_teaching_error(id)); + } + return Ok(json!({ "status": "error", "message": format!("Agent '{id}' not found. Use agent__check to verify it exists and is finished.") @@ -950,7 +1084,7 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result .map_err(|e| anyhow!("Agent failed: {e}"))?; let output = summarize_output(ctx, &result.agent_name, &result.output).await?; - ctx.pending_agents_guardrail_count = 0; + ctx.pending_tasks_guardrail_count = 0; Ok(json!({ "status": "completed", @@ -1051,7 +1185,7 @@ async fn handle_cancel(ctx: &mut RequestContext, args: &Value) -> Result let cleanup = tokio::time::timeout(Duration::from_secs(5), handle.join_handle).await; - ctx.pending_agents_guardrail_count = 0; + ctx.pending_tasks_guardrail_count = 0; let message = match cleanup { Ok(_) => format!("Cancelled agent '{agent_name}' and waited for cleanup."), @@ -1065,10 +1199,16 @@ async fn handle_cancel(ctx: &mut RequestContext, args: &Value) -> Result "message": message, })) } - None => Ok(json!({ - "status": "error", - "message": format!("No agent found with id '{id}'"), - })), + None => { + if is_job_task(ctx.supervisor.as_ref(), id) { + return Ok(job_id_teaching_error(id)); + } + + Ok(json!({ + "status": "error", + "message": format!("No agent found with id '{id}'"), + })) + } } } @@ -1115,10 +1255,18 @@ fn handle_send_message(ctx: &mut RequestContext, args: &Value) -> Result "message": format!("Message delivered to agent '{id}'"), })) } - None => Ok(json!({ - "status": "error", - "message": format!("No agent found with id '{id}'. Agent may not exist or may have already completed."), - })), + None => { + if is_job_task(ctx.supervisor.as_ref(), id) + || is_job_task(ctx.parent_supervisor.as_ref(), id) + { + return Ok(job_id_teaching_error(id)); + } + + Ok(json!({ + "status": "error", + "message": format!("No agent found with id '{id}'. Agent may not exist or may have already completed."), + })) + } } } @@ -1455,9 +1603,13 @@ mod tests { use super::*; use crate::config::test_fixtures::{FixtureServer, fixture_runtime}; use crate::config::{AgentConfig, AppState, WorkingMode}; + use crate::function::jobs::RingBuf; use crate::supervisor::escalation::{EscalationQueue, EscalationRequest}; + use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus}; + use parking_lot::Mutex; use serde_json::json; use serial_test::serial; + use std::mem; fn default_app_state() -> Arc { Arc::new(AppState::test_default()) @@ -1472,19 +1624,83 @@ mod tests { ctx } + fn ctx_with_job_capable_supervisor() -> RequestContext { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + ctx.supervisor = Some(Arc::new(RwLock::new( + Supervisor::new(4, 3).with_max_concurrent_jobs(4), + ))); + ctx + } + + fn make_fake_job(id: &str) -> JobHandle { + let rt = tokio::runtime::Runtime::new().unwrap(); + let join_handle = rt.spawn(async { + Ok(JobResult { + output: json!(null), + exit_code: Some(0), + output_bytes_captured: 0, + }) + }); + mem::forget(rt); + JobHandle { + id: id.to_string(), + tool: "execute_command".to_string(), + started_at: std::time::Instant::now(), + join_handle, + abort_signal: create_abort_signal(), + state: Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })), + output_buf: Arc::new(Mutex::new(RingBuf::default())), + no_change_checks: 0, + last_check_state: None, + } + } + + fn register_fake_job(ctx: &mut RequestContext, id: &str) { + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(make_fake_job(id)) + .unwrap(); + } + + fn assert_job_teaching_error(result: &Value, id: &str) { + assert_eq!(result["status"], "error"); + let message = result["message"].as_str().unwrap(); + assert_eq!( + message, + format!( + "'{id}' is a background job, not an agent — use job__check / job__collect / job__cancel" + ) + ); + } + fn register_fake_agent(ctx: &mut RequestContext, id: &str, name: &str) { + register_fake_agent_with_output(ctx, id, name, "fake output"); + } + + fn register_fake_agent_with_output( + ctx: &mut RequestContext, + id: &str, + name: &str, + output: &str, + ) { let rt = tokio::runtime::Runtime::new().unwrap(); let id_owned = id.to_string(); let name_owned = name.to_string(); + let output_owned = output.to_string(); let join_handle = rt.spawn(async move { Ok(AgentResult { id: id_owned, agent_name: name_owned, - output: "fake output".into(), + output: output_owned, exit_status: AgentExitStatus::Completed, }) }); - std::mem::forget(rt); + mem::forget(rt); let handle = AgentHandle { id: id.to_string(), @@ -1511,6 +1727,48 @@ mod tests { .block_on(f) } + fn register_running_agent(ctx: &mut RequestContext, id: &str, name: &str) -> AbortSignal { + let abort = create_abort_signal(); + let id_owned = id.to_string(); + let name_owned = name.to_string(); + let join_handle = tokio::spawn(async move { + time::sleep(Duration::from_secs(60)).await; + Ok(AgentResult { + id: id_owned, + agent_name: name_owned, + output: String::new(), + exit_status: AgentExitStatus::Completed, + }) + }); + let handle = AgentHandle { + id: id.to_string(), + agent_name: name.to_string(), + depth: 1, + inbox: Arc::new(Inbox::new()), + abort_signal: abort.clone(), + join_handle, + child_supervisor: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + abort + } + + fn wait_until_finished(ctx: &RequestContext, id: &str) { + let deadline = std::time::Instant::now() + Duration::from_secs(5); + while ctx.supervisor.as_ref().unwrap().read().is_finished(id) != Some(true) { + assert!( + std::time::Instant::now() < deadline, + "agent '{id}' never finished" + ); + std::thread::sleep(Duration::from_millis(10)); + } + } + #[tokio::test] async fn sync_agent_functions_gates_meta_functions_on_live_capabilities() { let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); @@ -1886,20 +2144,20 @@ mod tests { #[test] fn dispatch_unknown_action_errors() { let mut ctx = ctx_with_supervisor(4, 3); - let result = run_async(handle_supervisor_tool(&mut ctx, "agent__bogus", &json!({}))); + let result = run_async(handle_agent_tool(&mut ctx, "agent__bogus", &json!({}))); assert!(result.is_err()); assert!( result .unwrap_err() .to_string() - .contains("Unknown supervisor action") + .contains("Unknown agent action") ); } #[test] fn dispatch_routes_list_running() { let mut ctx = ctx_with_supervisor(4, 3); - let result = run_async(handle_supervisor_tool( + let result = run_async(handle_agent_tool( &mut ctx, "agent__list_running", &json!({}), @@ -1911,7 +2169,7 @@ mod tests { #[test] fn dispatch_routes_list_available() { let mut ctx = ctx_with_supervisor(4, 3); - let result = run_async(handle_supervisor_tool( + let result = run_async(handle_agent_tool( &mut ctx, "agent__list_available", &json!({}), @@ -1924,12 +2182,8 @@ mod tests { #[test] fn dispatch_routes_task_list() { let mut ctx = ctx_with_supervisor(4, 3); - let result = run_async(handle_supervisor_tool( - &mut ctx, - "agent__task_list", - &json!({}), - )) - .unwrap(); + let result = + run_async(handle_agent_tool(&mut ctx, "agent__task_list", &json!({}))).unwrap(); assert!(result["tasks"].is_array()); } @@ -2075,7 +2329,7 @@ mod tests { reply_tx: tx, }); - match check_pending_agents_guardrail(&mut ctx) { + match check_pending_tasks_guardrail(&mut ctx) { GuardrailAction::Inject(prompt) => { assert!(prompt.contains("agent__reply_escalation")); assert!(prompt.contains("esc_9")); @@ -2119,7 +2373,7 @@ mod tests { .register(handle) .unwrap(); - match check_pending_agents_guardrail(&mut ctx) { + match check_pending_tasks_guardrail(&mut ctx) { GuardrailAction::Inject(prompt) => { assert!(!prompt.contains("agent__reply_escalation")); } @@ -2127,4 +2381,649 @@ mod tests { } }); } + + #[test] + fn handle_collect_finished_agent_returns_output_and_consumes_handle() { + let mut ctx = ctx_with_supervisor(4, 3); + register_fake_agent(&mut ctx, "a1", "explore"); + ctx.pending_tasks_guardrail_count = 2; + + let result = run_async(handle_collect(&mut ctx, &json!({"id": "a1"}))).unwrap(); + + assert_eq!(result["status"], "completed"); + assert_eq!(result["id"], "a1"); + assert_eq!(result["agent"], "explore"); + assert_eq!(result["exit_status"], "Completed"); + assert_eq!(result["output"], "fake output"); + assert_eq!(ctx.pending_tasks_guardrail_count, 0); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().is_finished("a1"), + None + ); + } + + #[test] + fn handle_collect_pending_escalations_early_out_keeps_handle() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + rt.block_on(async { + let mut ctx = ctx_with_supervisor(4, 3); + let _abort = register_running_agent(&mut ctx, "slow", "test"); + let queue = ctx.ensure_root_escalation_queue(); + let (tx, _rx) = tokio::sync::oneshot::channel(); + queue.submit(EscalationRequest { + id: "esc_1".into(), + from_agent_id: "a1".into(), + from_agent_name: "explore".into(), + question: "What do?".into(), + options: None, + reply_tx: tx, + }); + + let result = handle_collect(&mut ctx, &json!({"id": "slow"})) + .await + .unwrap(); + + assert_eq!(result["status"], "pending"); + assert!(result["pending_escalations"].is_array()); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().is_finished("slow"), + Some(false) + ); + }); + } + + #[test] + fn handle_collect_unknown_agent_errors() { + let mut ctx = ctx_with_supervisor(4, 3); + + let result = run_async(handle_collect(&mut ctx, &json!({"id": "missing"}))).unwrap(); + + assert_eq!(result["status"], "error"); + assert!(result["message"].as_str().unwrap().contains("not found")); + } + + #[test] + fn handle_collect_without_agent_passes_long_output_through_verbatim() { + let mut ctx = ctx_with_supervisor(4, 3); + let long_output = "x".repeat(10_000); + register_fake_agent_with_output(&mut ctx, "a1", "explore", &long_output); + + let result = run_async(handle_collect(&mut ctx, &json!({"id": "a1"}))).unwrap(); + + assert_eq!(result["output"], long_output); + } + + #[test] + fn handle_collect_output_below_agent_threshold_passes_through() { + let mut ctx = ctx_with_supervisor(4, 3); + ctx.agent = Some(Agent::test_new(AgentConfig { + summarization_threshold: 1_000_000, + ..Default::default() + })); + register_fake_agent(&mut ctx, "a1", "explore"); + + let result = run_async(handle_collect(&mut ctx, &json!({"id": "a1"}))).unwrap(); + + assert_eq!(result["status"], "completed"); + assert_eq!(result["output"], "fake output"); + } + + #[test] + fn handle_collect_over_threshold_with_unknown_summarization_model_errors() { + let mut ctx = ctx_with_supervisor(4, 3); + ctx.agent = Some(Agent::test_new(AgentConfig { + summarization_threshold: 1, + summarization_model: Some("nonexistent_client:model".into()), + ..Default::default() + })); + register_fake_agent(&mut ctx, "a1", "explore"); + + let err = run_async(handle_collect(&mut ctx, &json!({"id": "a1"}))).unwrap_err(); + + assert!(err.to_string().contains("nonexistent_client")); + } + + #[test] + fn guardrail_no_supervisor_is_no_action_and_resets_counter() { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + ctx.pending_tasks_guardrail_count = 2; + + assert!(matches!( + check_pending_tasks_guardrail(&mut ctx), + GuardrailAction::NoAction + )); + assert_eq!(ctx.pending_tasks_guardrail_count, 0); + } + + /// A finished-but-uncollected agent counts as pending: the turn-end + /// guardrail tells the model to collect it instead of letting the result + /// be silently dropped, and the handle stays registered. + #[test] + fn guardrail_surfaces_finished_but_uncollected_agents() { + let mut ctx = ctx_with_supervisor(4, 3); + register_fake_agent(&mut ctx, "a1", "explore"); + wait_until_finished(&ctx, "a1"); + ctx.pending_tasks_guardrail_count = 2; + + match check_pending_tasks_guardrail(&mut ctx) { + GuardrailAction::Inject(prompt) => { + assert!(prompt.contains("a1")); + assert!(prompt.contains("agent__collect --id a1")); + assert!(prompt.contains("Completed but UNCOLLECTED")); + } + _ => panic!("expected Inject action"), + } + assert_eq!(ctx.pending_tasks_guardrail_count, 3); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().is_finished("a1"), + Some(true) + ); + } + + #[test] + fn guardrail_force_terminate_discards_finished_uncollected_handles() { + let mut ctx = ctx_with_supervisor(4, 3); + register_fake_agent(&mut ctx, "a1", "explore"); + wait_until_finished(&ctx, "a1"); + ctx.pending_tasks_guardrail_count = PENDING_TASKS_GUARDRAIL_MAX; + + match check_pending_tasks_guardrail(&mut ctx) { + GuardrailAction::ForceTerminate(ids) => { + assert_eq!(ids, vec!["a1".to_string()]); + } + _ => panic!("expected ForceTerminate action"), + } + assert_eq!(ctx.pending_tasks_guardrail_count, 0); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().is_finished("a1"), + None + ); + } + + #[test] + fn guardrail_prompt_renders_running_and_finished_sections() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + rt.block_on(async { + let mut ctx = ctx_with_supervisor(4, 3); + let _abort = register_running_agent(&mut ctx, "slow", "test"); + register_fake_agent(&mut ctx, "a1", "explore"); + wait_until_finished(&ctx, "a1"); + + match check_pending_tasks_guardrail(&mut ctx) { + GuardrailAction::Inject(prompt) => { + assert!(prompt.contains("Still running")); + assert!(prompt.contains("slow (agent)")); + assert!(prompt.contains("Completed but UNCOLLECTED")); + assert!(prompt.contains("agent__collect --id a1")); + } + _ => panic!("expected Inject action"), + } + }); + } + + #[test] + fn guardrail_prompt_is_kind_aware_for_jobs() { + let tasks = vec![ + PendingTask { + id: "job_1".into(), + kind: TaskKind::Job, + finished: false, + }, + PendingTask { + id: "job_2".into(), + kind: TaskKind::Job, + finished: true, + }, + ]; + + let prompt = build_pending_tasks_guardrail_prompt(&tasks); + + assert!(prompt.contains("job_1 (job)")); + assert!(prompt.contains("job__cancel")); + assert!(prompt.contains("job__collect --id job_2")); + } + + #[test] + fn pending_tasks_includes_registered_jobs() { + let mut ctx = ctx_with_job_capable_supervisor(); + register_fake_job(&mut ctx, "job_1"); + + let tasks = pending_tasks(&ctx); + + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0].id, "job_1"); + 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() + .enable_all() + .build() + .unwrap(); + + rt.block_on(async { + let mut ctx = ctx_with_supervisor(4, 3); + let abort = register_running_agent(&mut ctx, "slow", "test"); + ctx.pending_tasks_guardrail_count = PENDING_TASKS_GUARDRAIL_MAX; + + match check_pending_tasks_guardrail(&mut ctx) { + GuardrailAction::ForceTerminate(ids) => { + assert_eq!(ids, vec!["slow".to_string()]); + } + _ => panic!("expected ForceTerminate action"), + } + assert_eq!(ctx.pending_tasks_guardrail_count, 0); + assert!(abort.aborted()); + }); + } + + #[test] + fn guardrail_injects_prompt_and_increments_counter_below_max() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + rt.block_on(async { + let mut ctx = ctx_with_supervisor(4, 3); + let _abort = register_running_agent(&mut ctx, "slow", "test"); + ctx.pending_tasks_guardrail_count = 1; + + match check_pending_tasks_guardrail(&mut ctx) { + GuardrailAction::Inject(prompt) => { + assert!(prompt.contains("slow")); + assert!(prompt.contains("agent__collect")); + } + _ => panic!("expected Inject action"), + } + assert_eq!(ctx.pending_tasks_guardrail_count, 2); + }); + } + + #[test] + fn handle_cancel_resets_guardrail_counter() { + let mut ctx = ctx_with_supervisor(4, 3); + register_fake_agent(&mut ctx, "a1", "explore"); + ctx.pending_tasks_guardrail_count = 2; + + let result = run_async(handle_cancel(&mut ctx, &json!({"id": "a1"}))).unwrap(); + + assert_eq!(result["status"], "ok"); + assert_eq!(ctx.pending_tasks_guardrail_count, 0); + } + + #[test] + fn handle_spawn_missing_agent_arg_errors() { + let mut ctx = ctx_with_supervisor(4, 3); + let err = run_async(handle_spawn(&mut ctx, &json!({}))).unwrap_err(); + assert!(err.to_string().contains("'agent' is required")); + } + + #[test] + fn handle_spawn_missing_prompt_arg_errors() { + let mut ctx = ctx_with_supervisor(4, 3); + let err = run_async(handle_spawn(&mut ctx, &json!({"agent": "explore"}))).unwrap_err(); + assert!(err.to_string().contains("'prompt' is required")); + } + + #[test] + fn handle_spawn_rejects_agent_outside_whitelist() { + let mut ctx = ctx_with_supervisor(4, 3); + ctx.agent = Some(Agent::test_new(AgentConfig { + spawnable_agents: Some(vec!["allowed".into()]), + ..Default::default() + })); + + let result = run_async(handle_spawn( + &mut ctx, + &json!({"agent": "notallowed", "prompt": "p"}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("spawnable_agents") + ); + } + + #[test] + fn handle_spawn_at_capacity_errors() { + let mut ctx = ctx_with_supervisor(1, 3); + register_fake_agent(&mut ctx, "a1", "explore"); + + let result = run_async(handle_spawn( + &mut ctx, + &json!({"agent": "x", "prompt": "p"}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert_eq!( + result["message"], + "At capacity: 1/1 agents running. Wait for one to finish or cancel one." + ); + } + + #[test] + fn handle_spawn_exceeding_depth_errors() { + let mut ctx = ctx_with_supervisor(4, 0); + + let result = run_async(handle_spawn( + &mut ctx, + &json!({"agent": "x", "prompt": "p"}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("Max agent depth exceeded") + ); + } + + #[test] + fn handle_spawn_no_supervisor_errors() { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + let err = run_async(handle_spawn( + &mut ctx, + &json!({"agent": "x", "prompt": "p"}), + )) + .unwrap_err(); + assert!(err.to_string().contains("No supervisor active")); + } + + /// Checking a finished agent is a pure status probe: it reports the + /// agent as finished, points at agent__collect, and leaves the handle + /// registered so a subsequent collect still returns the result. + #[test] + fn handle_check_finished_agent_reports_status_and_keeps_handle() { + let mut ctx = ctx_with_supervisor(4, 3); + register_fake_agent(&mut ctx, "a1", "explore"); + wait_until_finished(&ctx, "a1"); + + let result = run_async(handle_check(&mut ctx, &json!({"id": "a1"}))).unwrap(); + + assert_eq!(result["status"], "finished"); + assert_eq!(result["id"], "a1"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("agent__collect") + ); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().is_finished("a1"), + Some(true) + ); + + let collected = run_async(handle_collect(&mut ctx, &json!({"id": "a1"}))).unwrap(); + + assert_eq!(collected["status"], "completed"); + assert_eq!(collected["output"], "fake output"); + } + + #[test] + fn handle_cancel_running_agent_aborts_and_waits_for_cleanup() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + rt.block_on(async { + let mut ctx = ctx_with_supervisor(4, 3); + let sig = create_abort_signal(); + let sig2 = sig.clone(); + let join_handle = tokio::spawn(async move { + loop { + if sig2.aborted() { + return Ok(AgentResult { + id: "a1".into(), + agent_name: "explore".into(), + output: String::new(), + exit_status: AgentExitStatus::Completed, + }); + } + time::sleep(Duration::from_millis(10)).await; + } + }); + let handle = AgentHandle { + id: "a1".into(), + agent_name: "explore".into(), + depth: 1, + inbox: Arc::new(Inbox::new()), + abort_signal: sig.clone(), + join_handle, + child_supervisor: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + ctx.pending_tasks_guardrail_count = 2; + + let result = handle_cancel(&mut ctx, &json!({"id": "a1"})).await.unwrap(); + + assert_eq!(result["status"], "ok"); + let message = result["message"].as_str().unwrap(); + assert!(message.contains("Cancelled agent 'explore'")); + assert!(message.contains("waited for cleanup")); + assert!(sig.aborted()); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().is_finished("a1"), + None + ); + assert_eq!(ctx.pending_tasks_guardrail_count, 0); + }); + } + + #[test] + fn handle_check_registered_job_id_teaches_job_tools() { + run_async(async { + let mut ctx = ctx_with_job_capable_supervisor(); + register_fake_job(&mut ctx, "bg_1"); + + let result = handle_check(&mut ctx, &json!({"id": "bg_1"})) + .await + .unwrap(); + + assert_job_teaching_error(&result, "bg_1"); + }); + } + + #[test] + fn handle_check_job_prefixed_id_teaches_job_tools() { + run_async(async { + let mut ctx = ctx_with_supervisor(4, 3); + + let result = handle_check(&mut ctx, &json!({"id": "job_deadbeef"})) + .await + .unwrap(); + + assert_job_teaching_error(&result, "job_deadbeef"); + }); + } + + #[test] + fn handle_collect_registered_job_id_teaches_job_tools() { + run_async(async { + let mut ctx = ctx_with_job_capable_supervisor(); + register_fake_job(&mut ctx, "bg_1"); + + let result = handle_collect(&mut ctx, &json!({"id": "bg_1"})) + .await + .unwrap(); + + assert_job_teaching_error(&result, "bg_1"); + }); + } + + #[test] + fn handle_collect_job_prefixed_id_teaches_job_tools() { + run_async(async { + let mut ctx = ctx_with_supervisor(4, 3); + + let result = handle_collect(&mut ctx, &json!({"id": "job_deadbeef"})) + .await + .unwrap(); + + assert_job_teaching_error(&result, "job_deadbeef"); + }); + } + + #[test] + fn handle_cancel_registered_job_id_teaches_job_tools_and_keeps_job() { + run_async(async { + let mut ctx = ctx_with_job_capable_supervisor(); + register_fake_job(&mut ctx, "bg_1"); + + let result = handle_cancel(&mut ctx, &json!({"id": "bg_1"})) + .await + .unwrap(); + + assert_job_teaching_error(&result, "bg_1"); + assert!(ctx.supervisor.as_ref().unwrap().read().has_job("bg_1")); + }); + } + + #[test] + fn handle_send_message_registered_job_id_teaches_job_tools() { + run_async(async { + let mut ctx = ctx_with_job_capable_supervisor(); + register_fake_job(&mut ctx, "bg_1"); + + let result = + handle_send_message(&mut ctx, &json!({"id": "bg_1", "message": "hi"})).unwrap(); + + assert_job_teaching_error(&result, "bg_1"); + }); + } + + #[test] + fn handle_send_message_job_in_parent_supervisor_teaches_job_tools() { + run_async(async { + let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd); + let mut parent_sup = Supervisor::new(4, 3).with_max_concurrent_jobs(4); + parent_sup.register(make_fake_job("bg_p")).unwrap(); + ctx.parent_supervisor = Some(Arc::new(RwLock::new(parent_sup))); + + let result = + handle_send_message(&mut ctx, &json!({"id": "bg_p", "message": "hi"})).unwrap(); + + assert_job_teaching_error(&result, "bg_p"); + }); + } + + #[test] + fn guardrail_burns_bounded_injects_then_force_terminates_running_job() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + + rt.block_on(async { + let mut ctx = ctx_with_job_capable_supervisor(); + let abort = create_abort_signal(); + let join_handle = tokio::spawn(async { + time::sleep(Duration::from_secs(60)).await; + Ok(JobResult { + output: json!(null), + exit_code: Some(0), + output_bytes_captured: 0, + }) + }); + let handle = JobHandle { + id: "job_1".to_string(), + tool: "execute_command".to_string(), + started_at: std::time::Instant::now(), + join_handle, + abort_signal: abort.clone(), + state: Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })), + output_buf: Arc::new(Mutex::new(RingBuf::default())), + no_change_checks: 0, + last_check_state: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + + for expected_count in 1..=PENDING_TASKS_GUARDRAIL_MAX { + match check_pending_tasks_guardrail(&mut ctx) { + GuardrailAction::Inject(prompt) => assert!(prompt.contains("job_1")), + _ => panic!("expected Inject below max"), + } + assert_eq!(ctx.pending_tasks_guardrail_count, expected_count); + } + + match check_pending_tasks_guardrail(&mut ctx) { + GuardrailAction::ForceTerminate(ids) => { + assert_eq!(ids, vec!["job_1".to_string()]); + } + _ => panic!("expected ForceTerminate at max"), + } + assert_eq!(ctx.pending_tasks_guardrail_count, 0); + assert!(abort.aborted()); + }); + } + + #[test] + fn guardrail_force_terminate_discards_finished_uncollected_job() { + let mut ctx = ctx_with_job_capable_supervisor(); + register_fake_job(&mut ctx, "job_1"); + let deadline = std::time::Instant::now() + Duration::from_secs(5); + while !pending_tasks(&ctx).iter().any(|t| t.finished) { + assert!( + std::time::Instant::now() < deadline, + "job 'job_1' never finished" + ); + std::thread::sleep(Duration::from_millis(10)); + } + ctx.pending_tasks_guardrail_count = PENDING_TASKS_GUARDRAIL_MAX; + + match check_pending_tasks_guardrail(&mut ctx) { + GuardrailAction::ForceTerminate(ids) => { + assert_eq!(ids, vec!["job_1".to_string()]); + } + _ => panic!("expected ForceTerminate action"), + } + assert_eq!(ctx.pending_tasks_guardrail_count, 0); + assert!(!ctx.supervisor.as_ref().unwrap().read().has_job("job_1")); + } } diff --git a/src/function/jobs.rs b/src/function/jobs.rs new file mode 100644 index 0000000..262b3ef --- /dev/null +++ b/src/function/jobs.rs @@ -0,0 +1,2889 @@ +use super::agents::AGENT_FUNCTION_PREFIX; +use super::memory::MEMORY_FUNCTION_PREFIX; +use super::rag_query::RAG_FUNCTION_PREFIX; +use super::skill::SKILL_FUNCTION_PREFIX; +use super::todo::TODO_FUNCTION_PREFIX; +use super::user_interaction::USER_FUNCTION_PREFIX; +use super::{FunctionDeclaration, JsonSchema, PATH_SEP, mcp_error_display, render_tool_result}; +use crate::config::{ + McpRuntime, RequestContext, effective_max_concurrent_jobs, jobs_enabled, paths, +}; +use crate::graph; +use crate::mcp::{ + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX, + MCP_PROMPT_META_FUNCTION_NAME_PREFIX, MCP_READ_META_FUNCTION_NAME_PREFIX, + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, +}; +use crate::supervisor::notification::job_notification; +use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus, Supervisor}; +use crate::utils::{create_abort_signal, muted_warning_text, temp_file, wait_abort_signal}; + +use anyhow::{Context, Result, anyhow, bail}; +use indexmap::IndexMap; +use parking_lot::{Mutex, RwLock}; +use serde_json::{Value, json}; +use std::collections::HashMap; +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use std::{env, fs}; +use tokio::io::AsyncReadExt; +use tokio::time; +use uuid::Uuid; + +pub const JOB_FUNCTION_PREFIX: &str = "job__"; + +pub const DEFAULT_MAX_CONCURRENT_JOBS: usize = 5; + +const JOB_RESULT_TAIL_CAP_CHARS: usize = 50_000; + +const JOB_KILL_GRACE: Duration = Duration::from_secs(5); + +const JOB_PUMP_DRAIN_GRACE: Duration = Duration::from_secs(2); + +pub fn is_agent_task(supervisor: Option<&Arc>>, id: &str) -> bool { + id.starts_with("agent_") + || id.starts_with("graph_agent_") + || supervisor.is_some_and(|sup| sup.read().has_agent(id)) +} + +pub struct RingBuf { + buf: Vec, + capacity: usize, + write_pos: usize, + total_written: u64, +} + +impl RingBuf { + pub fn new(capacity: usize) -> Self { + Self { + buf: Vec::new(), + capacity, + write_pos: 0, + total_written: 0, + } + } + + pub fn push(&mut self, bytes: &[u8]) { + self.total_written += bytes.len() as u64; + if self.capacity == 0 { + return; + } + + let src = if bytes.len() > self.capacity { + &bytes[bytes.len() - self.capacity..] + } else { + bytes + }; + + for &byte in src { + if self.buf.len() < self.capacity { + self.buf.push(byte); + } else { + self.buf[self.write_pos] = byte; + } + + self.write_pos = (self.write_pos + 1) % self.capacity; + } + } + + pub fn total_written(&self) -> u64 { + self.total_written + } + + pub fn tail(&self) -> Vec { + if self.buf.len() < self.capacity { + return self.buf.clone(); + } + + let mut out = Vec::with_capacity(self.capacity); + out.extend_from_slice(&self.buf[self.write_pos..]); + out.extend_from_slice(&self.buf[..self.write_pos]); + out + } +} + +impl Default for RingBuf { + fn default() -> Self { + Self::new(64 * 1024) + } +} + +/// Everything a detached process job needs, frozen at `job__start`: config, +/// env, and PATH changes made afterwards do not affect a running job. +pub struct JobEnvSnapshot { + cmd_name: String, + display_name: String, + cmd_args: Vec, + envs: HashMap, + output_file: PathBuf, + timeout_secs: u64, +} + +/// The complete context an MCP job task owns. The runtime holds ONLY the +/// validated server's handle, so the detached task cannot reach any other +/// server even by bug. +pub struct JobCtx { + mcp_runtime: McpRuntime, + current_depth: usize, +} + +pub fn job_function_declarations() -> Vec { + vec![ + FunctionDeclaration { + name: format!("{JOB_FUNCTION_PREFIX}start"), + description: "Run a tool call as a background job and return immediately with a job id, so you can \ + keep working while it runs. `arguments` is the same object the tool takes when called \ + directly. Backgroundable tools: external command tools (e.g. execute_command) and \ + `mcp_invoke_*` calls; built-in `agent__`/`job__`/`user__`/`todo__`/`memory__`/`skill__` \ + 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. 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([ + ( + "tool".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + description: Some("Name of the tool to run in the background, exactly as it appears in your tool catalog".into()), + ..Default::default() + }, + ), + ( + "arguments".to_string(), + JsonSchema { + type_value: Some("object".to_string()), + description: Some("The arguments object the tool takes when called directly".into()), + ..Default::default() + }, + ), + ])), + required: Some(vec!["tool".to_string()]), + ..Default::default() + }, + agent: false, + }, + FunctionDeclaration { + name: format!("{JOB_FUNCTION_PREFIX}check"), + description: "Non-blocking status probe for a background job. Returns status, elapsed time, and a tail \ + of the output captured so far; it NEVER consumes the result — use `job__collect` for \ + that. Call sparingly: if repeated checks show no change, do other work instead — you \ + will be notified when the job completes. `output_bytes_captured` reports the total \ + output size so far — use it to decide how to collect (`tail_lines`, `full_result`, or \ + having the command write to a file).".to_string(), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(IndexMap::from([( + "id".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + description: Some("The job ID returned by job__start".into()), + ..Default::default() + }, + )])), + required: Some(vec!["id".to_string()]), + ..Default::default() + }, + agent: false, + }, + FunctionDeclaration { + name: format!("{JOB_FUNCTION_PREFIX}collect"), + description: "Block until the named background job finishes, then return its result and remove the \ + job. The result keeps the LAST 50,000 chars by default (failures land at the tail of \ + build logs); pass `tail_lines` to keep only the last N lines instead, or \ + `full_result: true` to skip the cap entirely (the session-wide tool-output limit still \ + applies). Collecting is consume-once — decide first via `job__check`'s \ + `output_bytes_captured`. For very large outputs, prefer having the command write to a \ + file and paging it with `fs_read`.".to_string(), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(IndexMap::from([ + ( + "id".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + description: Some("The job ID returned by job__start".into()), + ..Default::default() + }, + ), + ( + "tail_lines".to_string(), + JsonSchema { + type_value: Some("number".to_string()), + description: Some("Keep only the last N lines of the result".into()), + ..Default::default() + }, + ), + ( + "full_result".to_string(), + JsonSchema { + type_value: Some("boolean".to_string()), + description: Some("Return the complete result, skipping the default 50,000-char tail cap (default: false)".into()), + ..Default::default() + }, + ), + ])), + required: Some(vec!["id".to_string()]), + ..Default::default() + }, + agent: false, + }, + FunctionDeclaration { + name: format!("{JOB_FUNCTION_PREFIX}cancel"), + description: "Cancel a background job: kills its process (SIGTERM, then SIGKILL after a 5s grace) and \ + discards the handle. Returns any partial output captured so far. The id cannot be used \ + afterwards.".to_string(), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(IndexMap::from([( + "id".to_string(), + JsonSchema { + type_value: Some("string".to_string()), + description: Some("The job ID returned by job__start".into()), + ..Default::default() + }, + )])), + required: Some(vec!["id".to_string()]), + ..Default::default() + }, + agent: false, + }, + FunctionDeclaration { + name: format!("{JOB_FUNCTION_PREFIX}list"), + description: "List all background jobs you have started that are still registered, with status, \ + elapsed time, and bytes of output captured.".to_string(), + parameters: JsonSchema { + type_value: Some("object".to_string()), + properties: Some(IndexMap::new()), + ..Default::default() + }, + agent: false, + }, + ] +} + +pub async fn handle_job_tool( + ctx: &mut RequestContext, + cmd_name: &str, + args: &Value, +) -> Result { + let action = cmd_name + .strip_prefix(JOB_FUNCTION_PREFIX) + .unwrap_or(cmd_name); + + match action { + "start" => handle_start(ctx, args).await, + "check" => handle_check(ctx, args), + "collect" => handle_collect(ctx, args).await, + "cancel" => handle_cancel(ctx, args).await, + "list" => handle_list(ctx), + _ => bail!("Unknown job action: {action}"), + } +} + +fn job_status_str(status: JobStatus) -> &'static str { + match status { + JobStatus::Running => "running", + JobStatus::Completed => "completed", + JobStatus::Failed => "failed", + } +} + +fn job_miss_error(supervisor: Option<&Arc>>, id: &str) -> Value { + if is_agent_task(supervisor, id) { + json!({ + "status": "error", + "message": format!( + "'{id}' is a spawned agent, not a background job — use agent__check / agent__collect / agent__cancel" + ), + }) + } else { + json!({ + "status": "error", + "message": format!( + "No job '{id}' is registered — it may have already been collected or cancelled. job__list shows active jobs." + ), + }) + } +} + +fn whitelist_rejection(tool: &str) -> Option { + let non_invoke_mcp_prefixes = [ + MCP_SEARCH_META_FUNCTION_NAME_PREFIX, + MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, + MCP_READ_META_FUNCTION_NAME_PREFIX, + MCP_PROMPT_META_FUNCTION_NAME_PREFIX, + ]; + let reason = if tool.starts_with(AGENT_FUNCTION_PREFIX) || tool.starts_with(JOB_FUNCTION_PREFIX) + { + Some(format!( + "'{tool}' is already asynchronous — call it directly. Agents may start jobs, but jobs never start agents or other jobs." + )) + } else if tool.starts_with(USER_FUNCTION_PREFIX) { + Some(format!( + "'{tool}' is interactive and must run in-turn — a background job cannot touch the terminal. Call it directly." + )) + } else if tool.starts_with(TODO_FUNCTION_PREFIX) + || tool.starts_with(MEMORY_FUNCTION_PREFIX) + || tool.starts_with(SKILL_FUNCTION_PREFIX) + || tool.starts_with(RAG_FUNCTION_PREFIX) + { + Some(format!( + "'{tool}' mutates agent/session state and must run in-turn. Call it directly." + )) + } else if tool.starts_with("fs_") || tool == "ast_grep" { + Some(format!( + "'{tool}' is fast — invoke it directly instead of backgrounding it." + )) + } else if non_invoke_mcp_prefixes + .iter() + .any(|prefix| tool.starts_with(prefix)) + { + Some(format!( + "'{tool}' is a sub-second call; invoke it directly." + )) + } else { + None + }; + + reason.map(|why| { + json!({ + "status": "error", + "message": format!( + "{why} Backgroundable tools: external command tools (e.g. execute_command) and mcp_invoke_* calls." + ), + }) + }) +} + +/// 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!({ + "status": "error", + "message": "Background jobs are disabled in this context (max_concurrent_jobs is 0).", + })); + } + + let tool = args + .get("tool") + .and_then(Value::as_str) + .map(str::trim) + .filter(|s| !s.is_empty()) + .ok_or_else(|| anyhow!("'tool' is required"))? + .to_string(); + let arguments = args.get("arguments").cloned().unwrap_or_else(|| json!({})); + + if let Some(rejection) = whitelist_rejection(&tool) { + return Ok(rejection); + } + + if !ctx.declared_function_names.contains(&tool) { + return Ok(json!({ + "status": "error", + "message": format!( + "'{tool}' is not enabled in this context — job__start can only background tools declared to you in this \ + request. Use the exact name of a tool from your current catalog." + ), + })); + } + + let supervisor = match ctx.supervisor.as_ref() { + Some(sup) => Arc::clone(sup), + None => { + let max_jobs = effective_max_concurrent_jobs(ctx.agent.as_ref(), &ctx.app.config); + let sup = Arc::new(RwLock::new( + Supervisor::new(0, 0).with_max_concurrent_jobs(max_jobs), + )); + ctx.supervisor = Some(Arc::clone(&sup)); + sup + } + }; + + { + let sup = supervisor.read(); + if sup.active_job_count() >= sup.max_concurrent_jobs() { + return Ok(json!({ + "status": "error", + "message": format!( + "At capacity: {}/{} jobs running. Collect or cancel one first.", + sup.active_job_count(), + sup.max_concurrent_jobs() + ), + })); + } + } + + let short_uuid = &Uuid::new_v4().to_string()[..8]; + let job_id = format!("job_{short_uuid}"); + let state = Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })); + let output_buf = Arc::new(Mutex::new(RingBuf::default())); + + let join_handle = if tool.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) { + let server = tool + .strip_prefix(&format!("{MCP_INVOKE_META_FUNCTION_NAME_PREFIX}_")) + .ok_or_else(|| anyhow!("Malformed MCP invoke function name: {tool}"))? + .to_string(); + let Some(server_handle) = ctx.tool_scope.mcp_runtime.get(&server) else { + return Ok(json!({ + "status": "error", + "message": format!("MCP server '{server}' is not connected in this context."), + })); + }; + let inner_tool = arguments + .get("tool") + .and_then(Value::as_str) + .ok_or_else(|| anyhow!("Missing 'tool' in arguments"))? + .to_string(); + let inner_args = arguments + .get("arguments") + .cloned() + .unwrap_or_else(|| json!({})); + let mut mcp_runtime = McpRuntime::new(); + mcp_runtime.insert(server.clone(), Arc::clone(server_handle)); + let job_ctx = JobCtx { + mcp_runtime, + current_depth: ctx.current_depth, + }; + let task_state = Arc::clone(&state); + let task_notifications = Arc::clone(&ctx.notification_queue); + let notify_id = job_id.clone(); + let notify_tool = tool.clone(); + tokio::spawn(async move { + let result = run_mcp_job(job_ctx, server, inner_tool, inner_args).await; + let success = result.is_ok(); + task_state.lock().status = if success { + JobStatus::Completed + } else { + JobStatus::Failed + }; + task_notifications.push(job_notification(¬ify_id, ¬ify_tool, success)); + result + }) + } else { + let snapshot = build_env_snapshot(ctx, &tool, &arguments)?; + let task_state = Arc::clone(&state); + let task_buf = Arc::clone(&output_buf); + let task_notifications = Arc::clone(&ctx.notification_queue); + let notify_id = job_id.clone(); + let notify_tool = tool.clone(); + tokio::spawn(async move { + let result = run_process_job(snapshot, Arc::clone(&task_state), task_buf).await; + let success = matches!(&result, Ok(job_result) if job_result.exit_code == Some(0)); + let mut job_state = task_state.lock(); + job_state.pgid = None; + job_state.status = if success { + JobStatus::Completed + } else { + JobStatus::Failed + }; + + drop(job_state); + + task_notifications.push(job_notification(¬ify_id, ¬ify_tool, success)); + result + }) + }; + + let handle = JobHandle { + id: job_id.clone(), + tool: tool.clone(), + started_at: Instant::now(), + join_handle, + abort_signal: create_abort_signal(), + state, + output_buf, + no_change_checks: 0, + last_check_state: None, + }; + + // On a capacity race the handle is dropped here, which kills the process + // group and aborts the task. + if let Err(e) = supervisor.write().register(handle) { + return Ok(json!({ + "status": "error", + "message": format!("{e}"), + })); + } + + if let Some(scope) = ctx.node_job_scope.as_mut() { + scope.push(job_id.clone()); + } + + Ok(json!({ + "status": "ok", + "job_id": job_id, + "tool": tool, + "message": "Running in background. Check with job__check, block with job__collect, cancel with job__cancel. \ + You will receive a system_notifications entry on completion. Jobs do not survive coyote exiting.", + })) +} + +fn handle_check(ctx: &RequestContext, args: &Value) -> Result { + let id = args + .get("id") + .and_then(Value::as_str) + .ok_or_else(|| anyhow!("'id' is required"))?; + + let Some(supervisor) = ctx.supervisor.as_ref() else { + return Ok(job_miss_error(None, id)); + }; + let mut sup = supervisor.write(); + let Some(job) = sup.job_mut(id) else { + drop(sup); + return Ok(job_miss_error(ctx.supervisor.as_ref(), id)); + }; + + let status = job.state.lock().status; + let (tail, total_written) = { + let buf = job.output_buf.lock(); + (buf.tail(), buf.total_written()) + }; + let check_state = (status, total_written); + + if job.last_check_state == Some(check_state) { + job.no_change_checks += 1; + } else { + job.no_change_checks = 0; + job.last_check_state = Some(check_state); + } + + let tail_truncated = (tail.len() as u64) < total_written; + let mut result = json!({ + "status": job_status_str(status), + "id": id, + "tool": job.tool, + "elapsed_secs": job.started_at.elapsed().as_secs(), + "output_tail": String::from_utf8_lossy(&tail).to_string(), + "output_bytes_captured": total_written, + "tail_truncated": tail_truncated, + }); + + if matches!(status, JobStatus::Running) { + result["message"] = json!( + "Job is still running. Call job__collect to block for the result, or do other work — you will be notified on completion." + ); + + if job.no_change_checks >= 3 { + result["hint"] = json!( + "No change across repeated checks — still running; call job__collect to block, or do other work — a system notification will fire on completion." + ); + } + } else { + result["message"] = json!(format!( + "Job finished — retrieve the result with job__collect --id {id}" + )); + } + + Ok(result) +} + +async fn handle_collect(ctx: &RequestContext, args: &Value) -> Result { + let id = args + .get("id") + .and_then(Value::as_str) + .ok_or_else(|| anyhow!("'id' is required"))?; + let tail_lines = args + .get("tail_lines") + .and_then(Value::as_u64) + .map(|n| n as usize); + let full_result = args + .get("full_result") + .and_then(Value::as_bool) + .unwrap_or(false); + + let Some(supervisor) = ctx.supervisor.as_ref().cloned() else { + return Ok(job_miss_error(None, id)); + }; + + let target_abort = { + let sup = supervisor.read(); + let Some(job) = sup.job(id) else { + drop(sup); + return Ok(job_miss_error(ctx.supervisor.as_ref(), id)); + }; + job.abort_signal.clone() + }; + + loop { + let is_finished = { + let sup = supervisor.read(); + sup.job(id).is_none_or(|job| job.join_handle.is_finished()) + }; + + if is_finished { + break; + } + + if let Some(queue) = ctx.root_escalation_queue() + && queue.has_pending() + { + let summary = queue.pending_summary(); + return Ok(json!({ + "status": "pending", + "id": id, + "message": format!("Job '{id}' is still running, but child agents have pending escalations that need your reply. Reply via agent__reply_escalation, then call job__collect again."), + "pending_escalations": summary, + })); + } + + if target_abort.aborted() { + let deadline = time::Instant::now() + Duration::from_secs(2); + while time::Instant::now() < deadline { + let is_finished = { + let sup = supervisor.read(); + sup.job(id).is_none_or(|job| job.join_handle.is_finished()) + }; + + if is_finished { + break; + } + + time::sleep(Duration::from_millis(50)).await; + } + + break; + } + + tokio::select! { + _ = time::sleep(Duration::from_millis(200)) => {} + _ = wait_abort_signal(&target_abort) => {} + } + } + + let handle = { + let mut sup = supervisor.write(); + sup.take_job(id) + }; + + let Some(mut handle) = handle else { + return Ok(json!({ + "status": "error", + "message": format!("Job '{id}' completed but could not be collected. It may have been collected by another call."), + })); + }; + + // Ctrl-C/exit teardown SIGTERMs the group without escalating, so a + // TERM-ignoring process would hang this join forever. Bound it and + // escalate to a group SIGKILL, gated on the pgid still being set. + let joined = match time::timeout(JOB_KILL_GRACE, &mut handle.join_handle).await { + Ok(joined) => joined, + Err(_) => { + #[cfg(unix)] + if let Some(pgid) = handle.state.lock().pgid { + unsafe { libc::killpg(pgid, libc::SIGKILL) }; + } + + match time::timeout(JOB_KILL_GRACE, &mut handle.join_handle).await { + Ok(joined) => joined, + Err(_) => { + handle.join_handle.abort(); + (&mut handle.join_handle).await + } + } + } + }; + let tool = handle.tool.clone(); + let elapsed_secs = handle.started_at.elapsed().as_secs(); + let status = handle.state.lock().status; + let (tail, total_written) = { + let buf = handle.output_buf.lock(); + (buf.tail(), buf.total_written()) + }; + let output_tail = String::from_utf8_lossy(&tail).to_string(); + + let job_result = match joined { + Err(join_err) => { + return Ok(json!({ + "status": "failed", + "id": id, + "tool": tool, + "error": format!("Job task panicked: {join_err}"), + "output_tail": output_tail, + "output_bytes_captured": total_written, + })); + } + Ok(Err(e)) => { + return Ok(json!({ + "status": "failed", + "id": id, + "tool": tool, + "error": format!("{e}"), + "output_tail": output_tail, + "output_bytes_captured": total_written, + })); + } + Ok(Ok(job_result)) => job_result, + }; + + let (result_value, result_truncated) = cap_result(job_result.output, tail_lines, full_result); + let mut response = json!({ + "status": job_status_str(status), + "id": id, + "tool": tool, + "elapsed_secs": elapsed_secs, + "result": result_value, + "output_tail": output_tail, + "output_bytes_captured": job_result.output_bytes_captured, + }); + + if let Some(exit_code) = job_result.exit_code { + response["exit_code"] = json!(exit_code); + } + if result_truncated { + response["result_truncated"] = json!(true); + } + + Ok(response) +} + +async fn handle_cancel(ctx: &RequestContext, args: &Value) -> Result { + let id = args + .get("id") + .and_then(Value::as_str) + .ok_or_else(|| anyhow!("'id' is required"))?; + + let Some(supervisor) = ctx.supervisor.as_ref() else { + return Ok(job_miss_error(None, id)); + }; + + let handle = { + let mut sup = supervisor.write(); + sup.take_job(id) + }; + + let Some(mut handle) = handle else { + return Ok(job_miss_error(ctx.supervisor.as_ref(), id)); + }; + + handle.abort_signal.set_ctrlc(); + kill_job_with_grace(&mut handle).await; + + let tool = handle.tool.clone(); + let (tail, total_written) = { + let buf = handle.output_buf.lock(); + (buf.tail(), buf.total_written()) + }; + + Ok(json!({ + "status": "cancelled", + "id": id, + "tool": tool, + "output_tail": String::from_utf8_lossy(&tail).to_string(), + "output_bytes_captured": total_written, + })) +} + +/// SIGTERM the process group, give it a grace period, then SIGKILL. Every +/// group kill is gated on `state.pgid` still being set: the job task clears +/// it right after `wait()` reaps the child, and killing after the reap could +/// signal an innocent recycled pid. MCP jobs (no pgid) fall through to a +/// plain task abort. Only the SIGKILL is re-gated; the SIGTERM fires after +/// the pgid read drops the lock, so a reap in that window could still hit a +/// recycled pid — accepted residual risk. +async fn kill_job_with_grace(handle: &mut JobHandle) { + #[cfg(unix)] + { + let pgid = handle.state.lock().pgid; + if let Some(pgid) = pgid { + unsafe { libc::killpg(pgid, libc::SIGTERM) }; + if time::timeout(JOB_KILL_GRACE, &mut handle.join_handle) + .await + .is_ok() + { + return; + } + if handle.state.lock().pgid.is_some() { + unsafe { libc::killpg(pgid, libc::SIGKILL) }; + } + } + } + handle.join_handle.abort(); + 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!({ + "active_jobs": 0, + "max_concurrent_jobs": effective_max_concurrent_jobs(ctx.agent.as_ref(), &ctx.app.config), + "jobs": [], + })); + }; + let sup = supervisor.read(); + + let jobs: Vec = sup + .jobs() + .map(|job| { + let status = job.state.lock().status; + json!({ + "id": job.id, + "tool": job.tool, + "status": job_status_str(status), + "elapsed_secs": job.started_at.elapsed().as_secs(), + "output_bytes_captured": job.output_buf.lock().total_written(), + }) + }) + .collect(); + + Ok(json!({ + "active_jobs": sup.active_job_count(), + "max_concurrent_jobs": sup.max_concurrent_jobs(), + "jobs": jobs, + })) +} + +/// Mirrors the foreground `extract_call_config` + `run_llm_function` env +/// assembly, resolved eagerly so the detached task owns everything it needs. +fn build_env_snapshot( + ctx: &RequestContext, + tool: &str, + arguments: &Value, +) -> Result { + let agent = ctx.agent.as_ref(); + let (cmd_name, mut cmd_args, mut envs) = match agent { + Some(agent) => match agent.functions().find(tool) { + Some(declaration) if declaration.agent => ( + format!("{}-{tool}", agent.name()), + vec![tool.to_string()], + agent.variable_envs(), + ), + Some(_) => (tool.to_string(), vec![], agent.variable_envs()), + None => (tool.to_string(), vec![], HashMap::new()), + }, + None => (tool.to_string(), vec![], HashMap::new()), + }; + + let mut bin_dirs: Vec = vec![]; + if let Some(agent) = agent { + let dir = paths::agent_bin_dir(agent.name()); + if dir.exists() { + bin_dirs.push(dir); + } + if graph::agent_has_graph(agent.name()) { + envs.insert("AUTO_CONFIRM".into(), "true".into()); + } + } else { + bin_dirs.push(paths::functions_bin_dir()); + } + let current_path = env::var("PATH").context("No PATH environment variable")?; + let prepend_path = bin_dirs + .iter() + .map(|v| format!("{}{PATH_SEP}", v.display())) + .collect::>() + .join(""); + envs.insert("PATH".into(), format!("{prepend_path}{current_path}")); + + let output_file = temp_file("-job-", ""); + envs.insert("LLM_OUTPUT".into(), output_file.display().to_string()); + envs.insert("CLICOLOR_FORCE".into(), "1".into()); + envs.insert("FORCE_COLOR".into(), "1".into()); + + cmd_args.push(arguments.to_string()); + + #[cfg(windows)] + let cmd_name = super::polyfill_cmd_name(&cmd_name, &bin_dirs); + + #[cfg(windows)] + let cmd_args = { + let mut args = cmd_args; + if let Some(json_data) = args.pop() { + let tool_data_file = temp_file("-tool-data-", ".json"); + fs::write(&tool_data_file, &json_data)?; + envs.insert( + "LLM_TOOL_DATA_FILE".into(), + tool_data_file.display().to_string(), + ); + } + args + }; + + let timeout_secs = env::var("COYOTE_TOOL_TIMEOUT") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(1800); + + Ok(JobEnvSnapshot { + cmd_name, + display_name: tool.to_string(), + cmd_args, + envs, + output_file, + timeout_secs, + }) +} + +async fn pump_into_ring(mut reader: impl AsyncReadExt + Unpin, output_buf: Arc>) { + let mut chunk = [0u8; 1024]; + loop { + match reader.read(&mut chunk).await { + Ok(0) | Err(_) => break, + Ok(n) => output_buf.lock().push(&chunk[..n]), + } + } +} + +/// Deletes the env snapshot's temp files however the job task exits, +/// including a cancel/abort dropping the future mid-await. +struct TempFileGuard(Vec); + +impl Drop for TempFileGuard { + fn drop(&mut self) { + for path in &self.0 { + let _ = fs::remove_file(path); + } + } +} + +/// A grandchild that inherits the pipes keeps them open past the child's +/// exit; don't let that hold the job task past its own timeout. +async fn drain_pump(mut pump: tokio::task::JoinHandle<()>) { + if time::timeout(JOB_PUMP_DRAIN_GRACE, &mut pump) + .await + .is_err() + { + pump.abort(); + } +} + +async fn run_process_job( + snapshot: JobEnvSnapshot, + state: Arc>, + output_buf: Arc>, +) -> Result { + let mut temp_files = vec![snapshot.output_file.clone()]; + if let Some(tool_data_file) = snapshot.envs.get("LLM_TOOL_DATA_FILE") { + temp_files.push(PathBuf::from(tool_data_file)); + } + let _temp_guard = TempFileGuard(temp_files); + + let mut command = tokio::process::Command::new(&snapshot.cmd_name); + command + .args(&snapshot.cmd_args) + .envs(&snapshot.envs) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + #[cfg(unix)] + command.process_group(0); + + let mut child = command + .spawn() + .map_err(|err| anyhow!("Unable to run {}, {err}", snapshot.display_name))?; + + #[cfg(unix)] + if let Some(pid) = child.id() { + state.lock().pgid = Some(pid as i32); + } + + let stdout = child + .stdout + .take() + .ok_or_else(|| anyhow!("Failed to capture stdout"))?; + let stderr = child + .stderr + .take() + .ok_or_else(|| anyhow!("Failed to capture stderr"))?; + let stdout_pump = tokio::spawn(pump_into_ring(stdout, Arc::clone(&output_buf))); + let stderr_pump = tokio::spawn(pump_into_ring(stderr, Arc::clone(&output_buf))); + + let wait_result = if snapshot.timeout_secs > 0 { + match time::timeout(Duration::from_secs(snapshot.timeout_secs), child.wait()).await { + Ok(wait_result) => wait_result, + Err(_) => { + kill_expired_job(&mut child, &state).await; + state.lock().pgid = None; + drain_pump(stdout_pump).await; + drain_pump(stderr_pump).await; + let output_bytes_captured = output_buf.lock().total_written(); + let message = format!( + "Tool call '{}' timed out after {}s and was killed (set COYOTE_TOOL_TIMEOUT to adjust; 0 = unlimited)", + snapshot.display_name, snapshot.timeout_secs + ); + + return Ok(JobResult { + output: json!({"tool_call_error": message}), + exit_code: None, + output_bytes_captured, + }); + } + } + } else { + child.wait().await + }; + let status = match wait_result { + Ok(status) => status, + Err(err) => { + stdout_pump.abort(); + stderr_pump.abort(); + bail!("Unable to run {}, {err}", snapshot.display_name); + } + }; + // pid-reuse guard: the child is reaped, so a later group kill against + // this pgid could hit an innocent recycled pid. + state.lock().pgid = None; + drain_pump(stdout_pump).await; + drain_pump(stderr_pump).await; + let output_bytes_captured = output_buf.lock().total_written(); + + let exit_code = status.code(); + if exit_code != Some(0) { + let message = match exit_code { + Some(code) => format!( + "Tool call '{}' exited with code {code}", + snapshot.display_name + ), + None => format!( + "Tool call '{}' was terminated by a signal", + snapshot.display_name + ), + }; + let mut error_json = json!({"tool_call_error": message}); + if let Ok(contents) = fs::read_to_string(&snapshot.output_file) + && !contents.trim().is_empty() + { + error_json["output"] = json!(contents); + } + + return Ok(JobResult { + output: error_json, + exit_code, + output_bytes_captured, + }); + } + + let mut output = Value::Null; + if snapshot.output_file.exists() { + let contents = fs::read_to_string(&snapshot.output_file) + .context("Failed to retrieve tool call output")?; + if !contents.is_empty() { + output = serde_json::from_str(&contents) + .ok() + .unwrap_or_else(|| json!({"output": contents})); + } + } + + Ok(JobResult { + output, + exit_code, + output_bytes_captured, + }) +} + +/// Same kill discipline as `kill_job_with_grace`, driven through the owned +/// `Child`. The SIGTERM here fires after the pgid read drops the lock and is +/// not re-gated — the same accepted pid-reuse window. +async fn kill_expired_job(child: &mut tokio::process::Child, state: &Arc>) { + #[cfg(unix)] + { + let pgid = state.lock().pgid; + if let Some(pgid) = pgid { + unsafe { libc::killpg(pgid, libc::SIGTERM) }; + if time::timeout(JOB_KILL_GRACE, child.wait()).await.is_err() { + unsafe { libc::killpg(pgid, libc::SIGKILL) }; + let _ = child.wait().await; + } + return; + } + } + #[cfg(not(unix))] + let _ = state; + let _ = child.start_kill(); + let _ = child.wait().await; +} + +async fn run_mcp_job( + job_ctx: JobCtx, + server: String, + tool: String, + arguments: Value, +) -> Result { + let raw = match job_ctx.mcp_runtime.invoke(&server, &tool, arguments).await { + Ok(raw) => raw, + Err(e) => { + if job_ctx.current_depth == 0 { + let error_msg = format!("MCP job invocation failed: {e}"); + eprintln!("{}", muted_warning_text(&mcp_error_display(&error_msg))); + } + return Err(e); + } + }; + let output = render_tool_result(serde_json::to_value(raw)?, &server)?; + + Ok(JobResult { + output, + exit_code: None, + output_bytes_captured: 0, + }) +} + +/// Tail-biased result capping owned by the collect handler: keeps the LAST +/// `tail_lines`/50,000 chars (build failures land at the tail), always cutting +/// on a char boundary. `full_result` lifts the 50,000-char ceiling (the +/// session-wide tool-output limit still applies downstream); `tail_lines` is +/// honored either way. +fn cap_result(output: Value, tail_lines: Option, full_result: bool) -> (Value, bool) { + if output.is_null() { + return (json!("DONE"), false); + } + let mut text = match &output { + Value::String(s) => s.clone(), + other => other.to_string(), + }; + let mut truncated = false; + if let Some(n) = tail_lines { + let lines: Vec<&str> = text.lines().collect(); + if lines.len() > n { + text = lines[lines.len() - n..].join("\n"); + truncated = true; + } + } + if !full_result && let Some(capped) = tail_chars(&text, JOB_RESULT_TAIL_CAP_CHARS) { + text = capped; + truncated = true; + } + + if truncated { + (json!(text), true) + } else { + (output, false) + } +} + +fn tail_chars(text: &str, max_chars: usize) -> Option { + let total = text.chars().count(); + if total <= max_chars { + return None; + } + + let cut = text + .char_indices() + .nth(total - max_chars) + .map(|(i, _)| i) + .unwrap_or(0); + + Some(format!( + "[truncated: kept last {max_chars} of {total} chars — the rest was not retained; next time \ + collect with full_result: true, or have the command write its output to a file]\n{}", + &text[cut..] + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::{AppConfig, AppState, WorkingMode}; + use crate::function::agents::{ + GuardrailAction, check_pending_tasks_guardrail, handle_agent_tool, + }; + use crate::supervisor::mailbox::Inbox; + use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult}; + use std::future::Future; + use std::mem; + + fn default_app_state() -> Arc { + Arc::new(AppState::test_default()) + } + + fn app_state_with_config(update: impl FnOnce(&mut AppConfig)) -> Arc { + let mut state = AppState::test_default(); + let mut config = (*state.config).clone(); + update(&mut config); + state.config = Arc::new(config); + Arc::new(state) + } + + fn plain_ctx() -> RequestContext { + RequestContext::new(default_app_state(), WorkingMode::Cmd) + } + + fn ctx_with_job_supervisor(max_jobs: usize) -> RequestContext { + let mut ctx = plain_ctx(); + ctx.supervisor = Some(Arc::new(RwLock::new( + Supervisor::new(0, 3).with_max_concurrent_jobs(max_jobs), + ))); + ctx + } + + fn make_running_job(id: &str) -> JobHandle { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let join_handle = rt.spawn(async { + Ok(JobResult { + output: Value::Null, + exit_code: Some(0), + output_bytes_captured: 0, + }) + }); + mem::forget(rt); + JobHandle { + id: id.to_string(), + tool: "execute_command".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal: create_abort_signal(), + state: Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })), + output_buf: Arc::new(Mutex::new(RingBuf::default())), + no_change_checks: 0, + last_check_state: None, + } + } + + fn run_async(f: F) -> F::Output { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(f) + } + + #[cfg(unix)] + fn test_snapshot(cmd: &str, args: &[&str], timeout_secs: u64) -> JobEnvSnapshot { + let output_file = temp_file("-job-test-", ""); + let mut envs = HashMap::new(); + envs.insert("PATH".to_string(), env::var("PATH").unwrap()); + envs.insert("LLM_OUTPUT".to_string(), output_file.display().to_string()); + JobEnvSnapshot { + cmd_name: cmd.to_string(), + display_name: cmd.to_string(), + cmd_args: args.iter().map(|s| s.to_string()).collect(), + envs, + output_file, + timeout_secs, + } + } + + #[test] + fn ring_buf_returns_contents_below_capacity() { + let mut buf = RingBuf::new(8); + + buf.push(b"abc"); + buf.push(b"de"); + + assert_eq!(buf.tail(), b"abcde"); + assert_eq!(buf.total_written(), 5); + } + + #[test] + fn ring_buf_exact_fit_keeps_everything() { + let mut buf = RingBuf::new(5); + + buf.push(b"abcde"); + + assert_eq!(buf.tail(), b"abcde"); + assert_eq!(buf.total_written(), 5); + } + + #[test] + fn ring_buf_wrap_around_keeps_newest_bytes() { + let mut buf = RingBuf::new(5); + + buf.push(b"abcde"); + buf.push(b"fg"); + + assert_eq!(buf.tail(), b"cdefg"); + assert_eq!(buf.total_written(), 7); + } + + #[test] + fn ring_buf_oversize_push_keeps_last_capacity_bytes() { + let mut buf = RingBuf::new(4); + + buf.push(b"abcdefghij"); + + assert_eq!(buf.tail(), b"ghij"); + assert_eq!(buf.total_written(), 10); + } + + #[test] + fn ring_buf_default_capacity_is_64_kib() { + let mut buf = RingBuf::default(); + let payload = vec![b'x'; 64 * 1024 + 1]; + + buf.push(&payload); + + assert_eq!(buf.tail().len(), 64 * 1024); + assert_eq!(buf.total_written(), 64 * 1024 + 1); + } + + #[test] + fn is_agent_task_matches_agent_prefixes() { + assert!(is_agent_task(None, "agent_explore_a1b2c3d4")); + assert!(is_agent_task(None, "graph_agent_explore_a1b2c3d4")); + assert!(!is_agent_task(None, "job_deadbeef")); + } + + #[test] + fn is_agent_task_matches_registered_agents() { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let join_handle = rt.spawn(async { + Ok(AgentResult { + id: "a1".into(), + agent_name: "explore".into(), + output: String::new(), + exit_status: AgentExitStatus::Completed, + }) + }); + mem::forget(rt); + let handle = AgentHandle { + id: "a1".to_string(), + agent_name: "explore".to_string(), + depth: 1, + inbox: Arc::new(Inbox::new()), + abort_signal: create_abort_signal(), + join_handle, + child_supervisor: None, + }; + let mut sup = Supervisor::new(4, 3); + sup.register(handle).unwrap(); + let sup = Arc::new(RwLock::new(sup)); + + assert!(is_agent_task(Some(&sup), "a1")); + assert!(!is_agent_task(Some(&sup), "missing")); + } + + #[test] + fn job_function_declarations_cover_all_five_actions() { + let names: Vec = job_function_declarations() + .into_iter() + .map(|d| d.name) + .collect(); + + assert_eq!( + names, + vec![ + "job__start", + "job__check", + "job__collect", + "job__cancel", + "job__list" + ] + ); + } + + #[test] + fn whitelist_rejects_state_mutating_tools() { + for tool in ["memory__write", "todo__add", "skill__load", "rag__query"] { + let rejection = whitelist_rejection(tool).unwrap(); + + let message = rejection["message"].as_str().unwrap(); + assert!( + message.contains("mutates agent/session state"), + "unexpected message for {tool}: {message}" + ); + } + } + + #[test] + fn whitelist_rejects_async_and_interactive_tools() { + for tool in ["agent__spawn", "job__check"] { + let message = whitelist_rejection(tool).unwrap()["message"] + .as_str() + .unwrap() + .to_string(); + + assert!( + message.contains("already asynchronous"), + "unexpected message for {tool}: {message}" + ); + } + let message = whitelist_rejection("user__select").unwrap()["message"] + .as_str() + .unwrap() + .to_string(); + assert!(message.contains("interactive")); + } + + #[test] + fn whitelist_rejects_fast_mcp_meta_tools() { + for tool in [ + "mcp_search_github", + "mcp_describe_github", + "mcp_read_github", + "mcp_prompt_github", + ] { + let message = whitelist_rejection(tool).unwrap()["message"] + .as_str() + .unwrap() + .to_string(); + + assert!( + message.contains("sub-second"), + "unexpected message for {tool}: {message}" + ); + } + } + + #[test] + fn whitelist_rejects_fast_file_builtins() { + for tool in ["fs_read", "fs_cat", "fs_grep", "ast_grep"] { + let message = whitelist_rejection(tool).unwrap()["message"] + .as_str() + .unwrap() + .to_string(); + + assert!( + message.contains("is fast"), + "unexpected message for {tool}: {message}" + ); + } + } + + #[test] + fn whitelist_allows_external_and_mcp_invoke_tools() { + assert!(whitelist_rejection("execute_command").is_none()); + assert!(whitelist_rejection("mcp_invoke_github").is_none()); + } + + #[test] + fn handle_start_rejects_fast_builtins_without_spawn() { + let mut ctx = plain_ctx(); + ctx.declared_function_names.insert("fs_read".into()); + ctx.declared_function_names.insert("ast_grep".into()); + + for tool in ["fs_read", "ast_grep"] { + let result = run_async(handle_start( + &mut ctx, + &json!({"tool": tool, "arguments": {}}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!(result["message"].as_str().unwrap().contains("is fast")); + } + assert!(ctx.supervisor.is_none(), "no job may be spawned"); + } + + #[test] + fn handle_start_rejects_non_whitelisted_tool_without_spawn() { + let mut ctx = plain_ctx(); + ctx.declared_function_names.insert("memory__write".into()); + + let result = run_async(handle_start( + &mut ctx, + &json!({"tool": "memory__write", "arguments": {}}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("mutates agent/session state") + ); + assert!(ctx.supervisor.is_none(), "no job may be spawned"); + } + + #[test] + fn handle_start_rejects_undeclared_tool_without_spawn() { + let mut ctx = plain_ctx(); + + let result = run_async(handle_start( + &mut ctx, + &json!({"tool": "execute_command", "arguments": {}}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("not enabled in this context") + ); + assert!(ctx.supervisor.is_none(), "no job may be spawned"); + } + + #[test] + fn handle_start_rejects_when_jobs_disabled() { + let app_state = app_state_with_config(|config| config.max_concurrent_jobs = Some(0)); + let mut ctx = RequestContext::new(app_state, WorkingMode::Cmd); + ctx.declared_function_names.insert("execute_command".into()); + + let result = run_async(handle_start( + &mut ctx, + &json!({"tool": "execute_command", "arguments": {}}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!(result["message"].as_str().unwrap().contains("disabled")); + assert!(ctx.supervisor.is_none()); + } + + #[test] + fn handle_start_rejects_at_capacity_without_spawn() { + let mut ctx = ctx_with_job_supervisor(1); + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(make_running_job("j1")) + .unwrap(); + ctx.declared_function_names.insert("execute_command".into()); + + let result = run_async(handle_start( + &mut ctx, + &json!({"tool": "execute_command", "arguments": {}}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("At capacity: 1/1") + ); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().active_job_count(), + 1 + ); + } + + #[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(); + ctx.declared_function_names + .insert("mcp_invoke_github".into()); + + let result = run_async(handle_start( + &mut ctx, + &json!({"tool": "mcp_invoke_github", "arguments": {"tool": "search"}}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("not connected") + ); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().active_job_count(), + 0 + ); + } + + #[cfg(unix)] + #[test] + fn handle_start_lazy_inits_supervisor_and_collect_returns_result() { + run_async(async { + let mut ctx = plain_ctx(); + ctx.declared_function_names.insert("echo".into()); + + let started = handle_start(&mut ctx, &json!({"tool": "echo", "arguments": {}})) + .await + .unwrap(); + + assert_eq!(started["status"], "ok"); + let job_id = started["job_id"].as_str().unwrap().to_string(); + assert!(job_id.starts_with("job_")); + assert!( + ctx.supervisor.is_some(), + "plain sessions lazily init a supervisor" + ); + + let collected = handle_collect(&ctx, &json!({"id": job_id})).await.unwrap(); + + assert_eq!(collected["status"], "completed"); + assert_eq!(collected["result"], "DONE"); + assert_eq!(collected["exit_code"], 0); + assert!(collected["output_tail"].as_str().unwrap().contains("{}")); + assert!(!ctx.supervisor.as_ref().unwrap().read().has_job(&job_id)); + }); + } + + #[test] + fn handle_check_unknown_id_teaches_job_list() { + let ctx = ctx_with_job_supervisor(4); + + let result = handle_check(&ctx, &json!({"id": "job_x"})).unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("No job 'job_x' is registered") + ); + } + + #[cfg(unix)] + #[test] + fn job_completion_pushes_notification_for_own_context() { + run_async(async { + let mut ctx = plain_ctx(); + 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(); + + let supervisor = ctx.supervisor.clone().unwrap(); + let deadline = time::Instant::now() + Duration::from_secs(5); + while time::Instant::now() < deadline { + let finished = supervisor + .read() + .job(&job_id) + .is_none_or(|job| job.join_handle.is_finished()); + if finished { + break; + } + time::sleep(Duration::from_millis(10)).await; + } + + let events = ctx.notification_queue.drain(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].event, "job_completed"); + assert_eq!(events[0].id, job_id); + assert_eq!(events[0].tool_or_agent, "echo"); + assert_eq!(events[0].status, "success"); + assert_eq!( + events[0].next_action, + format!("job__collect --id {job_id} for output") + ); + }); + } + + #[cfg(unix)] + #[test] + fn job_failure_pushes_failed_notification() { + run_async(async { + let mut ctx = plain_ctx(); + ctx.declared_function_names.insert("false".into()); + + let started = handle_start(&mut ctx, &json!({"tool": "false", "arguments": {}})) + .await + .unwrap(); + let job_id = started["job_id"].as_str().unwrap().to_string(); + + let supervisor = ctx.supervisor.clone().unwrap(); + let deadline = time::Instant::now() + Duration::from_secs(5); + while time::Instant::now() < deadline { + let finished = supervisor + .read() + .job(&job_id) + .is_none_or(|job| job.join_handle.is_finished()); + if finished { + break; + } + time::sleep(Duration::from_millis(10)).await; + } + + let events = ctx.notification_queue.drain(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].event, "job_failed"); + assert_eq!(events[0].status, "failed"); + }); + } + + #[test] + fn cancelled_job_notification_is_suppressed_at_drain() { + run_async(async { + let ctx = ctx_with_job_supervisor(4); + let join_handle = tokio::spawn(async { + time::sleep(Duration::from_secs(30)).await; + Ok(JobResult { + output: Value::Null, + exit_code: Some(0), + output_bytes_captured: 0, + }) + }); + let handle = JobHandle { + id: "j1".to_string(), + tool: "execute_command".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal: create_abort_signal(), + state: Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })), + output_buf: Arc::new(Mutex::new(RingBuf::default())), + no_change_checks: 0, + last_check_state: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + ctx.notification_queue + .push(job_notification("j1", "execute_command", false)); + + let result = handle_cancel(&ctx, &json!({"id": "j1"})).await.unwrap(); + + assert_eq!(result["status"], "cancelled"); + assert!( + super::super::drain_live_notifications(&ctx).is_empty(), + "events for a cancelled job must never reach the model" + ); + }); + } + + #[test] + fn already_collected_job_notification_is_suppressed_at_drain() { + let ctx = ctx_with_job_supervisor(4); + ctx.notification_queue + .push(job_notification("j1", "execute_command", true)); + + assert!( + super::super::drain_live_notifications(&ctx).is_empty(), + "events for an already-collected job must be dropped" + ); + } + + #[test] + fn job_handlers_teach_cross_kind_for_agent_ids() { + let ctx = ctx_with_job_supervisor(4); + for result in [ + handle_check(&ctx, &json!({"id": "agent_explore_1"})).unwrap(), + run_async(handle_collect(&ctx, &json!({"id": "agent_explore_1"}))).unwrap(), + run_async(handle_cancel(&ctx, &json!({"id": "agent_explore_1"}))).unwrap(), + ] { + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("is a spawned agent, not a background job") + ); + } + } + + #[test] + fn job_handlers_miss_without_supervisor() { + let ctx = plain_ctx(); + + let result = handle_check(&ctx, &json!({"id": "job_x"})).unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("No job 'job_x'") + ); + } + + #[test] + fn handle_check_reports_running_job_tail() { + let ctx = ctx_with_job_supervisor(4); + let job = make_running_job("j1"); + job.output_buf.lock().push(b"hello"); + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(job) + .unwrap(); + + let result = handle_check(&ctx, &json!({"id": "j1"})).unwrap(); + + assert_eq!(result["status"], "running"); + assert_eq!(result["tool"], "execute_command"); + assert_eq!(result["output_tail"], "hello"); + assert_eq!(result["output_bytes_captured"], 5); + assert_eq!(result["tail_truncated"], false); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("still running") + ); + } + + #[test] + fn handle_check_finished_job_points_at_collect_without_consuming() { + let ctx = ctx_with_job_supervisor(4); + let job = make_running_job("j1"); + job.state.lock().status = JobStatus::Completed; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(job) + .unwrap(); + + let result = handle_check(&ctx, &json!({"id": "j1"})).unwrap(); + + assert_eq!(result["status"], "completed"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("job__collect --id j1") + ); + assert!(ctx.supervisor.as_ref().unwrap().read().has_job("j1")); + } + + #[test] + fn handle_check_hints_after_repeated_unchanged_checks() { + let ctx = ctx_with_job_supervisor(4); + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(make_running_job("j1")) + .unwrap(); + + for _ in 0..3 { + let result = handle_check(&ctx, &json!({"id": "j1"})).unwrap(); + assert!(result.get("hint").is_none()); + } + let result = handle_check(&ctx, &json!({"id": "j1"})).unwrap(); + assert_eq!( + result["hint"], + "No change across repeated checks — still running; call job__collect to block, or do other work — a system notification will fire on completion." + ); + } + + #[test] + fn handle_check_no_change_counter_resets_on_output_change() { + let ctx = ctx_with_job_supervisor(4); + let job = make_running_job("j1"); + let output_buf = Arc::clone(&job.output_buf); + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(job) + .unwrap(); + + for _ in 0..3 { + handle_check(&ctx, &json!({"id": "j1"})).unwrap(); + } + assert!( + handle_check(&ctx, &json!({"id": "j1"})) + .unwrap() + .get("hint") + .is_some() + ); + + output_buf.lock().push(b"more"); + for _ in 0..3 { + let result = handle_check(&ctx, &json!({"id": "j1"})).unwrap(); + assert!(result.get("hint").is_none()); + } + assert!( + handle_check(&ctx, &json!({"id": "j1"})) + .unwrap() + .get("hint") + .is_some() + ); + } + + #[test] + fn handle_check_finished_job_never_hints() { + let ctx = ctx_with_job_supervisor(4); + let job = make_running_job("j1"); + job.state.lock().status = JobStatus::Completed; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(job) + .unwrap(); + + for _ in 0..5 { + let result = handle_check(&ctx, &json!({"id": "j1"})).unwrap(); + assert!(result.get("hint").is_none()); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("job__collect --id j1") + ); + } + } + + #[test] + fn handle_collect_applies_tail_lines() { + run_async(async { + let ctx = ctx_with_job_supervisor(4); + let join_handle = tokio::spawn(async { + Ok(JobResult { + output: json!("l1\nl2\nl3"), + exit_code: Some(0), + output_bytes_captured: 0, + }) + }); + let handle = JobHandle { + id: "j1".to_string(), + tool: "execute_command".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal: create_abort_signal(), + state: Arc::new(Mutex::new(JobState { + status: JobStatus::Completed, + pgid: None, + })), + output_buf: Arc::new(Mutex::new(RingBuf::default())), + no_change_checks: 0, + last_check_state: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + + let result = handle_collect(&ctx, &json!({"id": "j1", "tail_lines": 2})) + .await + .unwrap(); + + assert_eq!(result["status"], "completed"); + assert_eq!(result["result"], "l2\nl3"); + assert_eq!(result["result_truncated"], true); + }); + } + + #[test] + fn handle_collect_full_result_returns_uncapped() { + run_async(async { + let ctx = ctx_with_job_supervisor(4); + let big = "x".repeat(JOB_RESULT_TAIL_CAP_CHARS + 10); + let payload = big.clone(); + let join_handle = tokio::spawn(async move { + Ok(JobResult { + output: json!(payload), + exit_code: Some(0), + output_bytes_captured: 0, + }) + }); + let handle = JobHandle { + id: "j1".to_string(), + tool: "execute_command".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal: create_abort_signal(), + state: Arc::new(Mutex::new(JobState { + status: JobStatus::Completed, + pgid: None, + })), + output_buf: Arc::new(Mutex::new(RingBuf::default())), + no_change_checks: 0, + last_check_state: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + + let result = handle_collect(&ctx, &json!({"id": "j1", "full_result": true})) + .await + .unwrap(); + + assert_eq!(result["status"], "completed"); + assert_eq!(result["result"], json!(big)); + assert!(result.get("result_truncated").is_none()); + }); + } + + #[test] + fn handle_collect_maps_panic_to_failed_with_ring_content() { + run_async(async { + let ctx = ctx_with_job_supervisor(4); + let join_handle: tokio::task::JoinHandle> = + tokio::spawn(async { panic!("boom") }); + let output_buf = Arc::new(Mutex::new(RingBuf::default())); + output_buf.lock().push(b"partial"); + let handle = JobHandle { + id: "j1".to_string(), + tool: "execute_command".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal: create_abort_signal(), + state: Arc::new(Mutex::new(JobState { + status: JobStatus::Failed, + pgid: None, + })), + output_buf, + no_change_checks: 0, + last_check_state: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + + let result = handle_collect(&ctx, &json!({"id": "j1"})).await.unwrap(); + + assert_eq!(result["status"], "failed"); + assert!(result["error"].as_str().unwrap().contains("panicked")); + assert_eq!(result["output_tail"], "partial"); + }); + } + + /// The completion notification is pushed from inside the job task, after + /// the run — a panic unwinds past the push, and neither the supervisor nor + /// collect synthesizes a notification for a panicked job. + #[test] + fn panicked_job_task_skips_the_completion_notification() { + run_async(async { + let ctx = ctx_with_job_supervisor(4); + let join_handle: tokio::task::JoinHandle> = + tokio::spawn(async { panic!("boom") }); + let handle = JobHandle { + id: "j1".to_string(), + tool: "execute_command".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal: create_abort_signal(), + state: Arc::new(Mutex::new(JobState { + // A panic never reaches the status update — the cell + // stays Running, which is the real post-panic state. + status: JobStatus::Running, + pgid: None, + })), + output_buf: Arc::new(Mutex::new(RingBuf::default())), + no_change_checks: 0, + last_check_state: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + + let result = handle_collect(&ctx, &json!({"id": "j1"})).await.unwrap(); + + assert_eq!(result["status"], "failed"); + assert!(result["error"].as_str().unwrap().contains("panicked")); + assert!( + ctx.notification_queue.drain().is_empty(), + "a panicked job must never produce a completion notification" + ); + }); + } + + #[cfg(unix)] + #[test] + fn handle_cancel_kills_running_process_job() { + run_async(async { + let ctx = ctx_with_job_supervisor(4); + let state = Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })); + let output_buf = Arc::new(Mutex::new(RingBuf::default())); + let snapshot = test_snapshot("sleep", &["30"], 0); + let task_state = Arc::clone(&state); + let task_buf = Arc::clone(&output_buf); + let join_handle = + tokio::spawn(async move { run_process_job(snapshot, task_state, task_buf).await }); + + let deadline = time::Instant::now() + Duration::from_secs(2); + while state.lock().pgid.is_none() && time::Instant::now() < deadline { + time::sleep(Duration::from_millis(10)).await; + } + let pgid = state.lock().pgid.expect("runner must record the pgid"); + + let handle = JobHandle { + id: "j1".to_string(), + tool: "sleep".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal: create_abort_signal(), + state: Arc::clone(&state), + output_buf, + no_change_checks: 0, + last_check_state: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + + let result = handle_cancel(&ctx, &json!({"id": "j1"})).await.unwrap(); + + assert_eq!(result["status"], "cancelled"); + assert_eq!(result["tool"], "sleep"); + assert!(!ctx.supervisor.as_ref().unwrap().read().has_job("j1")); + assert_eq!( + unsafe { libc::killpg(pgid, 0) }, + -1, + "process group must be dead after cancel" + ); + assert_eq!( + std::io::Error::last_os_error().raw_os_error(), + Some(libc::ESRCH) + ); + }); + } + + #[test] + fn handle_list_reports_jobs_and_capacity() { + let ctx = ctx_with_job_supervisor(4); + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(make_running_job("j1")) + .unwrap(); + + let result = handle_list(&ctx).unwrap(); + + assert_eq!(result["active_jobs"], 1); + assert_eq!(result["max_concurrent_jobs"], 4); + assert_eq!(result["jobs"][0]["id"], "j1"); + assert_eq!(result["jobs"][0]["status"], "running"); + } + + #[test] + fn handle_list_without_supervisor_reports_empty() { + let ctx = plain_ctx(); + + let result = handle_list(&ctx).unwrap(); + + assert_eq!(result["active_jobs"], 0); + assert_eq!(result["max_concurrent_jobs"], 5); + assert_eq!(result["jobs"].as_array().unwrap().len(), 0); + } + + #[test] + fn cap_result_normalizes_null_to_done() { + let (value, truncated) = cap_result(Value::Null, None, false); + + assert_eq!(value, json!("DONE")); + assert!(!truncated); + } + + #[test] + fn cap_result_preserves_small_values() { + let (value, truncated) = cap_result(json!({"a": 1}), None, false); + + assert_eq!(value, json!({"a": 1})); + assert!(!truncated); + } + + #[test] + fn cap_result_keeps_last_chars_on_char_boundary() { + let text = "é".repeat(JOB_RESULT_TAIL_CAP_CHARS + 10); + let (value, truncated) = cap_result(json!(text), None, false); + assert!(truncated); + let capped = value.as_str().unwrap(); + assert!(capped.starts_with(&format!( + "[truncated: kept last {} of {} chars — ", + JOB_RESULT_TAIL_CAP_CHARS, + JOB_RESULT_TAIL_CAP_CHARS + 10 + ))); + } + + #[test] + fn cap_result_full_result_skips_char_cap() { + let text = "a".repeat(JOB_RESULT_TAIL_CAP_CHARS + 10); + let (value, truncated) = cap_result(json!(text), None, true); + + assert!(!truncated); + assert_eq!(value, json!(text)); + } + + #[test] + fn cap_result_full_result_still_applies_tail_lines() { + let (value, truncated) = cap_result(json!("l1\nl2\nl3"), Some(2), true); + + assert!(truncated); + assert_eq!(value, json!("l2\nl3")); + } + + #[test] + fn tail_chars_floors_to_char_boundary() { + let capped = tail_chars("aébc", 2).unwrap(); + + assert!(capped.ends_with("bc")); + assert!(capped.starts_with("[truncated: kept last 2 of 4 chars — ")); + assert!(capped.contains("full_result: true")); + assert!(tail_chars("abc", 3).is_none()); + } + + #[cfg(unix)] + #[test] + fn run_process_job_times_out_and_clears_pgid() { + run_async(async { + let state = Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })); + let output_buf = Arc::new(Mutex::new(RingBuf::default())); + let snapshot = test_snapshot("sleep", &["30"], 1); + + let result = run_process_job(snapshot, Arc::clone(&state), output_buf) + .await + .unwrap(); + + assert_eq!(result.exit_code, None); + assert!( + result.output["tool_call_error"] + .as_str() + .unwrap() + .contains("timed out after 1s") + ); + assert!( + state.lock().pgid.is_none(), + "pid-reuse guard must clear pgid" + ); + }); + } + + #[cfg(unix)] + #[test] + fn run_process_job_reads_llm_output_and_captures_ring() { + run_async(async { + let state = Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })); + let output_buf = Arc::new(Mutex::new(RingBuf::default())); + let snapshot = test_snapshot( + "sh", + &["-c", "printf hi > \"$LLM_OUTPUT\"; echo captured"], + 0, + ); + + let result = run_process_job(snapshot, state, Arc::clone(&output_buf)) + .await + .unwrap(); + + assert_eq!(result.exit_code, Some(0)); + assert_eq!(result.output, json!({"output": "hi"})); + let tail = String::from_utf8_lossy(&output_buf.lock().tail()).to_string(); + assert!(tail.contains("captured")); + assert_eq!(result.output_bytes_captured, 9); + }); + } + + #[cfg(unix)] + #[test] + fn run_process_job_reports_nonzero_exit_with_partial_output() { + run_async(async { + let state = Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })); + let output_buf = Arc::new(Mutex::new(RingBuf::default())); + let snapshot = + test_snapshot("sh", &["-c", "printf partial > \"$LLM_OUTPUT\"; exit 3"], 0); + + let result = run_process_job(snapshot, state, output_buf).await.unwrap(); + + assert_eq!(result.exit_code, Some(3)); + assert!( + result.output["tool_call_error"] + .as_str() + .unwrap() + .contains("exited with code 3") + ); + assert_eq!(result.output["output"], "partial"); + }); + } + + #[cfg(unix)] + #[test] + fn run_process_job_reports_signal_death_as_error() { + run_async(async { + let state = Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })); + let output_buf = Arc::new(Mutex::new(RingBuf::default())); + let snapshot = test_snapshot("sh", &["-c", "kill -KILL $$"], 0); + + let result = run_process_job(snapshot, state, output_buf).await.unwrap(); + + assert_eq!(result.exit_code, None); + assert!( + result.output["tool_call_error"] + .as_str() + .unwrap() + .contains("terminated by a signal") + ); + }); + } + + #[cfg(unix)] + #[test] + fn run_process_job_removes_output_temp_file() { + run_async(async { + let state = Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })); + let output_buf = Arc::new(Mutex::new(RingBuf::default())); + let snapshot = test_snapshot("sh", &["-c", "printf hi > \"$LLM_OUTPUT\""], 0); + let output_file = snapshot.output_file.clone(); + + let result = run_process_job(snapshot, state, output_buf).await.unwrap(); + + assert_eq!(result.exit_code, Some(0)); + assert_eq!(result.output, json!({"output": "hi"})); + assert!(!output_file.exists(), "temp file must be removed"); + }); + } + + #[cfg(unix)] + #[test] + fn handle_collect_returns_after_term_ignoring_job_is_aborted() { + run_async(async { + let ctx = ctx_with_job_supervisor(4); + let state = Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })); + let output_buf = Arc::new(Mutex::new(RingBuf::default())); + let snapshot = + test_snapshot("sh", &["-c", "trap '' TERM; while :; do sleep 1; done"], 0); + let task_state = Arc::clone(&state); + let task_buf = Arc::clone(&output_buf); + let join_handle = tokio::spawn(async move { + let result = run_process_job(snapshot, Arc::clone(&task_state), task_buf).await; + let mut job_state = task_state.lock(); + job_state.pgid = None; + job_state.status = match &result { + Ok(job_result) if job_result.exit_code == Some(0) => JobStatus::Completed, + _ => JobStatus::Failed, + }; + drop(job_state); + result + }); + + let deadline = time::Instant::now() + Duration::from_secs(2); + while state.lock().pgid.is_none() && time::Instant::now() < deadline { + time::sleep(Duration::from_millis(10)).await; + } + let pgid = state.lock().pgid.expect("runner must record the pgid"); + + let abort_signal = create_abort_signal(); + let handle = JobHandle { + id: "j1".to_string(), + tool: "sh".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal: abort_signal.clone(), + state: Arc::clone(&state), + output_buf, + no_change_checks: 0, + last_check_state: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + + // Mimic Ctrl-C teardown: SIGTERM the group and flag the abort + // signal, leaving the handle registered. + abort_signal.set_ctrlc(); + unsafe { libc::killpg(pgid, libc::SIGTERM) }; + + let result = time::timeout( + Duration::from_secs(20), + handle_collect(&ctx, &json!({"id": "j1"})), + ) + .await + .expect("collect must not hang on a TERM-ignoring job") + .unwrap(); + + assert_eq!(result["status"], "failed"); + // killpg(pgid, 0) is unreliable here: the orphaned sleep + // grandchild lingers as an unreaped zombie and keeps the group + // id alive. Signal death proves the escalated SIGKILL landed. + assert!( + result["result"]["tool_call_error"] + .as_str() + .unwrap() + .contains("terminated by a signal") + ); + }); + } + + #[test] + fn jobs_only_supervisor_rejects_agent_spawn_at_capacity_zero() { + let mut ctx = ctx_with_job_supervisor(5); + + let result = run_async(handle_agent_tool( + &mut ctx, + "agent__spawn", + &json!({"agent": "explore", "prompt": "x"}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("At capacity: 0/0") + ); + } + + #[test] + fn jobs_only_supervisor_guardrail_surfaces_running_job() { + let mut ctx = ctx_with_job_supervisor(5); + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(make_running_job("j1")) + .unwrap(); + + match check_pending_tasks_guardrail(&mut ctx) { + GuardrailAction::Inject(prompt) => { + assert!(prompt.contains("j1")); + assert!(prompt.contains("job__collect")); + } + _ => panic!("expected Inject for a running job"), + } + assert_eq!(ctx.pending_tasks_guardrail_count, 1); + + let empty_ctx = &mut ctx_with_job_supervisor(5); + assert!(matches!( + check_pending_tasks_guardrail(empty_ctx), + GuardrailAction::NoAction + )); + } + + #[test] + fn jobs_only_supervisor_agent_surfaces_stay_functional() { + let mut ctx = ctx_with_job_supervisor(5); + + let listed = run_async(handle_agent_tool( + &mut ctx, + "agent__list_running", + &json!({}), + )) + .unwrap(); + assert_eq!(listed["active_count"], 0); + assert_eq!(listed["max_concurrent"], 0); + + let created = run_async(handle_agent_tool( + &mut ctx, + "agent__task_create", + &json!({"subject": "research"}), + )) + .unwrap(); + assert_eq!(created["status"], "ok"); + + let tasks = run_async(handle_agent_tool(&mut ctx, "agent__task_list", &json!({}))).unwrap(); + assert_eq!(tasks["tasks"].as_array().unwrap().len(), 1); + } + + #[cfg(unix)] + #[test] + fn handle_cancel_kills_grandchild_process() { + run_async(async { + let ctx = ctx_with_job_supervisor(4); + let state = Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })); + let output_buf = Arc::new(Mutex::new(RingBuf::default())); + let snapshot = test_snapshot("sh", &["-c", "sleep 30 & echo CHILD:$!; wait"], 0); + let task_state = Arc::clone(&state); + let task_buf = Arc::clone(&output_buf); + let join_handle = + tokio::spawn(async move { run_process_job(snapshot, task_state, task_buf).await }); + + let deadline = time::Instant::now() + Duration::from_secs(5); + let grandchild_pid = loop { + let tail = String::from_utf8_lossy(&output_buf.lock().tail()).to_string(); + if let Some(rest) = tail.split("CHILD:").nth(1) + && let Some(line_end) = rest.find('\n') + { + break rest[..line_end].trim().parse::().unwrap(); + } + assert!( + time::Instant::now() < deadline, + "grandchild pid never appeared in the ring buffer" + ); + time::sleep(Duration::from_millis(10)).await; + }; + + let handle = JobHandle { + id: "j1".to_string(), + tool: "sh".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal: create_abort_signal(), + state: Arc::clone(&state), + output_buf, + no_change_checks: 0, + last_check_state: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + + let result = handle_cancel(&ctx, &json!({"id": "j1"})).await.unwrap(); + + assert_eq!(result["status"], "cancelled"); + // kill(pid, 0) alone can't observe the death: the orphaned + // grandchild lingers as an unreaped zombie under init/launchd, + // so a Z state also proves the group kill landed. + fn grandchild_is_dead(pid: i32) -> bool { + let esrch = unsafe { libc::kill(pid, 0) } == -1 + && std::io::Error::last_os_error().raw_os_error() == Some(libc::ESRCH); + if esrch { + return true; + } + let stat = std::process::Command::new("ps") + .args(["-o", "stat=", "-p", &pid.to_string()]) + .output() + .map(|out| String::from_utf8_lossy(&out.stdout).trim().to_string()) + .unwrap_or_default(); + stat.is_empty() || stat.starts_with('Z') + } + let deadline = time::Instant::now() + Duration::from_secs(5); + while !grandchild_is_dead(grandchild_pid) { + assert!( + time::Instant::now() < deadline, + "grandchild must die with the process group" + ); + time::sleep(Duration::from_millis(10)).await; + } + }); + } + + #[cfg(unix)] + #[test] + fn run_process_job_clears_pgid_after_normal_completion() { + run_async(async { + let state = Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })); + let output_buf = Arc::new(Mutex::new(RingBuf::default())); + let snapshot = test_snapshot("echo", &["done"], 0); + + let result = run_process_job(snapshot, Arc::clone(&state), output_buf) + .await + .unwrap(); + + assert_eq!(result.exit_code, Some(0)); + assert!( + state.lock().pgid.is_none(), + "pid-reuse guard must clear pgid" + ); + }); + } + + #[test] + fn handle_start_rejects_shell_and_path_shaped_names_without_spawn() { + let mut ctx = plain_ctx(); + + for tool in ["bash", "./script.sh", "/usr/bin/env", "ls"] { + let result = run_async(handle_start( + &mut ctx, + &json!({"tool": tool, "arguments": {}}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("not enabled in this context") + ); + } + assert!(ctx.supervisor.is_none(), "no job may be spawned"); + } + + #[cfg(unix)] + #[test] + fn handle_start_rejects_context_filtered_tool_and_accepts_in_filter() { + run_async(async { + let mut ctx = plain_ctx(); + ctx.declared_function_names.insert("echo".into()); + + let rejected = handle_start(&mut ctx, &json!({"tool": "git_command", "arguments": {}})) + .await + .unwrap(); + + assert_eq!(rejected["status"], "error"); + assert!( + rejected["message"] + .as_str() + .unwrap() + .contains("not enabled in this context") + ); + assert!(ctx.supervisor.is_none(), "no job may be spawned"); + + let started = handle_start(&mut ctx, &json!({"tool": "echo", "arguments": {}})) + .await + .unwrap(); + + assert_eq!(started["status"], "ok"); + 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"], "completed"); + }); + } + + #[test] + fn handle_start_rejects_undeclared_mcp_invoke_without_spawn() { + let mut ctx = plain_ctx(); + + let result = run_async(handle_start( + &mut ctx, + &json!({"tool": "mcp_invoke_someserver", "arguments": {"tool": "search"}}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("not enabled in this context") + ); + assert!(ctx.supervisor.is_none(), "no job may be spawned"); + } + + #[test] + fn handle_start_rejects_declared_but_non_whitelisted_tools_without_spawn() { + let mut ctx = plain_ctx(); + for tool in ["memory__write", "fs_read", "agent__spawn", "user__select"] { + ctx.declared_function_names.insert(tool.into()); + } + + for (tool, category) in [ + ("memory__write", "mutates agent/session state"), + ("fs_read", "is fast"), + ("agent__spawn", "already asynchronous"), + ("user__select", "interactive"), + ] { + let result = run_async(handle_start( + &mut ctx, + &json!({"tool": tool, "arguments": {}}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!(result["message"].as_str().unwrap().contains(category)); + } + assert!(ctx.supervisor.is_none(), "no job may be spawned"); + } + + #[test] + fn handle_start_rejects_mapping_tools_alias() { + let app_state = app_state_with_config(|config| { + config + .mapping_tools + .insert("shell".into(), "execute_command".into()); + }); + let mut ctx = RequestContext::new(app_state, WorkingMode::Cmd); + ctx.declared_function_names.insert("execute_command".into()); + + let result = run_async(handle_start( + &mut ctx, + &json!({"tool": "shell", "arguments": {}}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("not enabled in this context") + ); + assert!(ctx.supervisor.is_none(), "no job may be spawned"); + } + + #[test] + fn child_context_cannot_reach_parent_job_ids() { + run_async(async { + let parent = ctx_with_job_supervisor(4); + parent + .supervisor + .as_ref() + .unwrap() + .write() + .register(make_running_job("job_p1")) + .unwrap(); + let child = RequestContext::new_for_child( + default_app_state(), + &parent, + 1, + Arc::new(Inbox::new()), + "c1".into(), + ); + assert!(child.supervisor.is_none()); + + let checked = handle_check(&child, &json!({"id": "job_p1"})).unwrap(); + let collected = handle_collect(&child, &json!({"id": "job_p1"})) + .await + .unwrap(); + let cancelled = handle_cancel(&child, &json!({"id": "job_p1"})) + .await + .unwrap(); + + for result in [checked, collected, cancelled] { + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("No job 'job_p1' is registered") + ); + } + assert!(parent.supervisor.as_ref().unwrap().read().has_job("job_p1")); + }); + } + + #[test] + fn handle_start_ignores_mid_batch_tool_scope_additions() { + let mut ctx = plain_ctx(); + ctx.declared_function_names.insert("job__start".into()); + ctx.tool_scope + .functions + .declarations + .push(FunctionDeclaration { + name: "late_external_tool".into(), + description: String::new(), + parameters: JsonSchema::default(), + agent: false, + }); + + let result = run_async(handle_start( + &mut ctx, + &json!({"tool": "late_external_tool", "arguments": {}}), + )) + .unwrap(); + + assert_eq!(result["status"], "error"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("not enabled in this context") + ); + assert!(ctx.supervisor.is_none(), "no job may be spawned"); + } +} diff --git a/src/function/mod.rs b/src/function/mod.rs index d32a24c..38e72a9 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -1,7 +1,8 @@ +pub(crate) mod agents; +pub(crate) mod jobs; pub(crate) mod memory; pub(crate) mod rag_query; pub(crate) mod skill; -pub(crate) mod supervisor; pub(crate) mod todo; pub(crate) mod user_interaction; @@ -23,10 +24,12 @@ use crate::mcp::{ McpServersConfig, is_mcp_meta_function, render, }; use crate::parsers::{bash, python, typescript}; +use agents::AGENT_FUNCTION_PREFIX; use anyhow::{Context, Result, anyhow, bail}; use futures_util::future; use indexmap::IndexMap; use indoc::formatdoc; +use jobs::JOB_FUNCTION_PREFIX; use memory::MEMORY_FUNCTION_PREFIX; use rag_query::RAG_FUNCTION_PREFIX; use rust_embed::Embed; @@ -46,7 +49,6 @@ use std::{ time::{Duration, Instant}, }; use strum_macros::AsRefStr; -use supervisor::SUPERVISOR_FUNCTION_PREFIX; use todo::TODO_FUNCTION_PREFIX; use user_interaction::USER_FUNCTION_PREFIX; @@ -339,12 +341,17 @@ pub async fn eval_tool_calls( } } - if ctx.current_depth == 0 - && let Some(queue) = ctx.root_escalation_queue() - && queue.has_pending() - && let Some(last) = output.last_mut() - { - inject_escalation_notification(last, queue.pending_summary()); + if let Some(last) = output.last_mut() { + let escalations = if ctx.current_depth == 0 { + ctx.root_escalation_queue() + .filter(|queue| queue.has_pending()) + .map(|queue| queue.pending_summary()) + .unwrap_or_default() + } else { + vec![] + }; + let notifications = drain_live_notifications(ctx); + merge_system_channel(last, escalations, notifications); } Ok(output) @@ -363,22 +370,79 @@ fn normalize_tool_result(result: Value) -> Value { } } -fn inject_escalation_notification(last: &mut ToolResult, summary: Vec) { - let instruction = "Child agents are BLOCKED waiting for your reply. \ - Call agent__reply_escalation for each pending escalation to unblock them."; - match &mut last.output { - Value::Object(map) => { - map.insert("pending_escalations".into(), json!(summary)); - map.insert("escalation_instruction".into(), json!(instruction)); - } - other => { - *other = json!({ - "output": other.take(), - "pending_escalations": summary, - "escalation_instruction": instruction, - }); - } +/// Drains this context's own notification queue and drops events whose +/// handle is no longer registered with the supervisor (already collected or +/// cancelled), so the model is never pointed at a dead id. +fn drain_live_notifications(ctx: &RequestContext) -> Vec { + let events = ctx.notification_queue.drain(); + if events.is_empty() { + return vec![]; } + + let Some(supervisor) = ctx.supervisor.as_ref() else { + return vec![]; + }; + let sup = supervisor.read(); + events + .into_iter() + .filter(|event| sup.has_job(&event.id) || sup.has_agent(&event.id)) + .map(|event| event.to_value()) + .collect() +} + +/// Single-pass merge of both system channels onto the last tool result of a +/// batch: pending escalations (children are blocked; listed first) and +/// background-task completion notifications. A single pass is mandatory — +/// two independent mergers would each apply the non-object wrap and nest the +/// output twice. With both channels empty this is a no-op, and with only +/// escalations it produces exactly the pre-notification output shape. +fn merge_system_channel(last: &mut ToolResult, escalations: Vec, notifications: Vec) { + if escalations.is_empty() && notifications.is_empty() { + return; + } + + let escalation_instruction = "Child agents are BLOCKED waiting for your reply. \ + Call agent__reply_escalation for each pending escalation to unblock them."; + let notification_instruction = + "Background tasks have finished; collect each result with its next_action command."; + + let map = match &mut last.output { + Value::Object(map) => map, + other => { + let mut map = serde_json::Map::new(); + map.insert("output".into(), other.take()); + *other = Value::Object(map); + match other { + Value::Object(map) => map, + _ => unreachable!(), + } + } + }; + + if !escalations.is_empty() { + map.insert("pending_escalations".into(), json!(escalations)); + map.insert( + "escalation_instruction".into(), + json!(escalation_instruction), + ); + } + + if !notifications.is_empty() { + map.insert("system_notifications".into(), json!(notifications)); + map.insert( + "notification_instruction".into(), + json!(notification_instruction), + ); + } +} + +/// Escalation-only entry point retained so the characterization tests that +/// pinned the pre-merger output shape keep proving, unmodified, that +/// `merge_system_channel` with no notifications is byte-identical to the +/// injection behavior they were written against. +#[cfg(test)] +fn inject_escalation_notification(last: &mut ToolResult, summary: Vec) { + merge_system_channel(last, summary, vec![]); } #[derive(Debug, Clone, Deserialize, Serialize)] @@ -404,7 +468,12 @@ impl ToolResult { pub fn truncate_if_needed(mut self, max_chars: usize) -> Self { let s = self.output.to_string(); if s.len() > max_chars { - let prefix = s.get(..max_chars).unwrap_or(s.as_str()); + let mut cut = max_chars; + while !s.is_char_boundary(cut) { + cut -= 1; + } + + let prefix = &s[..cut]; self.output = json!(format!( "[truncated: tool output exceeded {max_chars} chars]\n{prefix}" )); @@ -615,14 +684,23 @@ impl Functions { pub fn append_supervisor_functions(&mut self) { self.declarations - .extend(supervisor::supervisor_function_declarations()); + .extend(agents::agent_function_declarations()); self.declarations - .extend(supervisor::escalation_function_declarations()); + .extend(agents::escalation_function_declarations()); + } + + pub fn append_job_functions(&mut self) { + self.declarations.extend(jobs::job_function_declarations()); + } + + #[cfg(test)] + 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()); + .extend(agents::teammate_function_declarations()); } pub fn append_user_interaction_functions(&mut self) { @@ -1516,11 +1594,11 @@ impl ToolCall { json!({"tool_call_error": error_msg}) }) } - _ if cmd_name.starts_with(SUPERVISOR_FUNCTION_PREFIX) => { - supervisor::handle_supervisor_tool(ctx, &cmd_name, &json_data) + _ if cmd_name.starts_with(AGENT_FUNCTION_PREFIX) => { + agents::handle_agent_tool(ctx, &cmd_name, &json_data) .await .unwrap_or_else(|e| { - let error_msg = format!("Supervisor tool failed: {e}"); + let error_msg = format!("Agent tool failed: {e}"); eprintln!("{}", muted_warning_text(&format!("⚠️ {error_msg} ⚠️"))); json!({"tool_call_error": error_msg}) }) @@ -1543,6 +1621,15 @@ impl ToolCall { json!({"tool_call_error": error_msg}) }) } + _ if cmd_name.starts_with(JOB_FUNCTION_PREFIX) => { + jobs::handle_job_tool(ctx, &cmd_name, &json_data) + .await + .unwrap_or_else(|e| { + let error_msg = format!("Job tool failed: {e}"); + eprintln!("{}", muted_warning_text(&format!("⚠️ {error_msg} ⚠️"))); + json!({"tool_call_error": error_msg}) + }) + } _ => match run_llm_function(cmd_name, cmd_args, envs, agent_name) { Ok(Some(contents)) => serde_json::from_str(&contents) .ok() @@ -2332,6 +2419,19 @@ fn polyfill_cmd_name>(cmd_name: &str, bin_dir: &[T]) -> String { cmd_name } +// Polling tools are expected to repeat; recording them would also let them +// break up detection of a real loop in the calls they interleave with. +const LOOP_TRACKER_EXEMPT_TOOLS: [&str; 4] = [ + "job__check", + "job__list", + "agent__check", + "agent__list_running", +]; + +fn is_loop_tracker_exempt(name: &str) -> bool { + LOOP_TRACKER_EXEMPT_TOOLS.contains(&name) +} + #[derive(Debug, Clone)] pub struct ToolCallTracker { last_calls: VecDeque, @@ -2353,6 +2453,10 @@ impl ToolCallTracker { } pub fn check_loop(&self, new_call: &ToolCall) -> Option { + if is_loop_tracker_exempt(&new_call.name) { + return None; + } + if self.last_calls.len() < self.max_repeats { return None; } @@ -2419,6 +2523,10 @@ impl ToolCallTracker { } pub fn record_call(&mut self, call: ToolCall) { + if is_loop_tracker_exempt(&call.name) { + return; + } + if self.last_calls.len() >= self.chain_len * self.max_repeats { self.last_calls.pop_front(); } @@ -2474,15 +2582,22 @@ mod tests { FIXTURE_ANNOTATED_TEXT, FIXTURE_ANNOTATED_URI, FIXTURE_BLOB_BYTES, FIXTURE_BLOB_URI, FIXTURE_LOG_TEXT, FIXTURE_LOG_URI, FixtureServer, fixture_runtime, }; - use crate::config::{AppState, WorkingMode}; + use crate::config::{Agent, AgentConfig, AppConfig, AppState, WorkingMode}; use crate::supervisor::escalation::{EscalationQueue, EscalationRequest}; + use crate::supervisor::mailbox::Inbox; + use crate::supervisor::notification::{agent_notification, job_notification}; + use crate::supervisor::{ + AgentExitStatus, AgentHandle, AgentResult, JobHandle, JobResult, JobState, JobStatus, + Supervisor, + }; use base64::Engine; use base64::engine::general_purpose::STANDARD; + use jobs::RingBuf; use rmcp::model::{CallToolResult, ContentBlock}; use serde_json::json; use serial_test::serial; - use std::process; use std::sync::Arc; + use std::{mem, process}; fn call(name: &str, id: Option<&str>) -> ToolCall { ToolCall::new(name.to_string(), json!({}), id.map(|s| s.to_string())) @@ -2560,6 +2675,256 @@ mod tests { assert!(result.output["escalation_instruction"].is_string()); } + fn ctx_with_registered_job(id: &str) -> RequestContext { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let join_handle = rt.spawn(async { + Ok(JobResult { + output: Value::Null, + exit_code: Some(0), + output_bytes_captured: 0, + }) + }); + mem::forget(rt); + let handle = JobHandle { + id: id.to_string(), + tool: "execute_command".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal: create_abort_signal(), + state: Arc::new(parking_lot::Mutex::new(JobState { + status: JobStatus::Completed, + pgid: None, + })), + output_buf: Arc::new(parking_lot::Mutex::new(RingBuf::default())), + no_change_checks: 0, + last_check_state: None, + }; + let mut sup = Supervisor::new(0, 3).with_max_concurrent_jobs(4); + sup.register(handle).unwrap(); + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.supervisor = Some(Arc::new(parking_lot::RwLock::new(sup))); + ctx + } + + #[test] + fn merge_system_channel_noop_when_both_channels_empty() { + let mut object_result = ToolResult::new(call("t", Some("id-1")), json!({"status": "ok"})); + merge_system_channel(&mut object_result, vec![], vec![]); + assert_eq!(object_result.output, json!({"status": "ok"})); + + let mut plain_result = ToolResult::new(call("t", Some("id-1")), json!("DONE")); + merge_system_channel(&mut plain_result, vec![], vec![]); + assert_eq!(plain_result.output, json!("DONE")); + } + + #[test] + fn merge_system_channel_escalations_only_matches_legacy_wrap_bytes() { + let summary = vec![json!({"escalation_id": "esc_1"})]; + let expected = json!({ + "output": "DONE", + "pending_escalations": summary, + "escalation_instruction": "Child agents are BLOCKED waiting for your reply. \ + Call agent__reply_escalation for each pending escalation to unblock them.", + }); + + let mut result = ToolResult::new(call("t", Some("id-1")), json!("DONE")); + merge_system_channel(&mut result, summary, vec![]); + + assert_eq!( + serde_json::to_string(&result.output).unwrap(), + serde_json::to_string(&expected).unwrap() + ); + } + + #[test] + fn merge_system_channel_adds_notifications_without_escalation_keys() { + let mut result = ToolResult::new(call("t", Some("id-1")), json!({"status": "ok"})); + + merge_system_channel(&mut result, vec![], vec![json!({"id": "job_1"})]); + + assert_eq!(result.output["status"], "ok"); + assert_eq!( + result.output["system_notifications"], + json!([{"id": "job_1"}]) + ); + assert!( + result.output["notification_instruction"] + .as_str() + .unwrap() + .contains("next_action") + ); + assert!(result.output.get("pending_escalations").is_none()); + assert!(result.output.get("escalation_instruction").is_none()); + } + + #[test] + fn merge_system_channel_wraps_non_object_once_with_both_channels() { + let mut result = ToolResult::new(call("t", Some("id-1")), json!("DONE")); + + merge_system_channel( + &mut result, + vec![json!({"escalation_id": "esc_1"})], + vec![json!({"id": "job_1"})], + ); + + assert_eq!(result.output["output"], json!("DONE")); + assert_eq!( + result.output["pending_escalations"][0]["escalation_id"], + "esc_1" + ); + assert_eq!(result.output["system_notifications"][0]["id"], "job_1"); + let keys: Vec<&str> = result + .output + .as_object() + .unwrap() + .keys() + .map(|k| k.as_str()) + .collect(); + assert_eq!( + keys, + vec![ + "output", + "pending_escalations", + "escalation_instruction", + "system_notifications", + "notification_instruction" + ] + ); + } + + #[test] + fn drain_live_notifications_drops_unregistered_ids() { + let ctx = ctx_with_registered_job("job_live"); + ctx.notification_queue + .push(job_notification("job_live", "execute_command", true)); + ctx.notification_queue + .push(job_notification("job_gone", "execute_command", true)); + + let live = drain_live_notifications(&ctx); + + assert_eq!(live.len(), 1); + assert_eq!(live[0]["id"], "job_live"); + assert!( + ctx.notification_queue.drain().is_empty(), + "drain must consume the queue" + ); + } + + #[test] + fn drain_live_notifications_without_supervisor_drops_everything() { + let ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.notification_queue + .push(job_notification("job_x", "execute_command", true)); + assert!(drain_live_notifications(&ctx).is_empty()); + } + + fn ctx_with_registered_agent(id: &str) -> RequestContext { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let agent_id = id.to_string(); + let join_handle = rt.spawn(async move { + Ok(AgentResult { + id: agent_id, + agent_name: "explore".into(), + output: String::new(), + exit_status: AgentExitStatus::Completed, + }) + }); + mem::forget(rt); + let handle = AgentHandle { + id: id.to_string(), + agent_name: "explore".to_string(), + depth: 1, + inbox: Arc::new(Inbox::new()), + abort_signal: create_abort_signal(), + join_handle, + child_supervisor: None, + }; + let mut sup = Supervisor::new(4, 3); + sup.register(handle).unwrap(); + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.supervisor = Some(Arc::new(parking_lot::RwLock::new(sup))); + ctx + } + + #[test] + fn drain_live_notifications_keeps_registered_agent_events() { + let ctx = ctx_with_registered_agent("agent_explore_1"); + ctx.notification_queue + .push(agent_notification("agent_explore_1", "explore", true)); + + let live = drain_live_notifications(&ctx); + + assert_eq!(live.len(), 1); + assert_eq!(live[0]["event"], "agent_completed"); + assert_eq!( + live[0]["next_action"], + "agent__collect --id agent_explore_1 for output" + ); + } + + #[test] + fn drain_live_notifications_drops_collected_agent_events() { + let ctx = ctx_with_registered_agent("agent_explore_1"); + ctx.supervisor + .as_ref() + .unwrap() + .write() + .take("agent_explore_1") + .unwrap(); + ctx.notification_queue + .push(agent_notification("agent_explore_1", "explore", true)); + + assert!(drain_live_notifications(&ctx).is_empty()); + } + + #[test] + fn eval_tool_calls_merges_notifications_at_depth_without_escalations() { + let mut ctx = ctx_with_registered_job("job_n1"); + ctx.current_depth = 1; + let queue = ctx.ensure_root_escalation_queue(); + submit_escalation(&queue, "esc_1"); + ctx.notification_queue + .push(job_notification("job_n1", "execute_command", true)); + + let calls = vec![call("unknown_tool", Some("id-1"))]; + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + let out = &results[0].output; + assert_eq!(out["system_notifications"][0]["id"], "job_n1"); + assert_eq!(out["system_notifications"][0]["event"], "job_completed"); + assert!(out["notification_instruction"].is_string()); + assert!( + out.get("pending_escalations").is_none(), + "escalations are root-only" + ); + } + + #[test] + fn eval_tool_calls_merges_both_channels_onto_last_result() { + let mut ctx = ctx_with_registered_job("job_n1"); + let queue = ctx.ensure_root_escalation_queue(); + submit_escalation(&queue, "esc_1"); + ctx.notification_queue + .push(job_notification("job_n1", "execute_command", false)); + + let calls = vec![call("unknown_tool", Some("id-1"))]; + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + let out = &results[0].output; + assert_eq!(out["pending_escalations"][0]["escalation_id"], "esc_1"); + assert_eq!(out["system_notifications"][0]["event"], "job_failed"); + assert!( + out.get("output").is_none(), + "object outputs are extended in place, never wrapped" + ); + } + #[test] fn eval_tool_calls_soft_fails_unknown_tool() { let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); @@ -2597,6 +2962,51 @@ mod tests { ); } + #[cfg(unix)] + #[test] + fn job_finished_earlier_drains_notification_on_later_batch() { + run_async(async { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.declared_function_names.insert("echo".into()); + + let started = jobs::handle_job_tool( + &mut ctx, + "job__start", + &json!({"tool": "echo", "arguments": {}}), + ) + .await + .unwrap(); + assert_eq!(started["status"], "ok"); + let job_id = started["job_id"].as_str().unwrap().to_string(); + + let supervisor = ctx.supervisor.clone().unwrap(); + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + while tokio::time::Instant::now() < deadline { + let finished = supervisor + .read() + .job(&job_id) + .is_none_or(|job| job.join_handle.is_finished()); + if finished { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + + let calls = vec![call("unknown_tool", Some("id-2"))]; + let results = eval_tool_calls(&mut ctx, calls).await.unwrap(); + + let out = &results.last().unwrap().output; + assert_eq!(out["system_notifications"][0]["id"], job_id); + assert_eq!(out["system_notifications"][0]["event"], "job_completed"); + assert!( + out["notification_instruction"] + .as_str() + .unwrap() + .contains("next_action") + ); + }); + } + #[test] fn normalize_tool_result_preserves_non_null_values() { assert_eq!( @@ -2782,10 +3192,69 @@ mod tests { assert!(msg.contains("repeat_tool")); } + #[test] + fn loop_tracker_exempt_list_is_exactly_the_polling_tools() { + let actual: HashSet<&str> = LOOP_TRACKER_EXEMPT_TOOLS.iter().copied().collect(); + let expected: HashSet<&str> = [ + "job__check", + "job__list", + "agent__check", + "agent__list_running", + ] + .into_iter() + .collect(); + assert_eq!(LOOP_TRACKER_EXEMPT_TOOLS.len(), 4); + assert_eq!(actual, expected); + } + + #[test] + fn tracker_exempt_tools_never_trip() { + for name in LOOP_TRACKER_EXEMPT_TOOLS { + let mut tracker = ToolCallTracker::default(); + let exempt = call_with_args(name, json!({"id": "j1"})); + tracker.record_call(exempt.clone()); + tracker.record_call(exempt.clone()); + assert!(tracker.check_loop(&exempt).is_none()); + + let other = call_with_args("execute_command", json!({"command": "ls"})); + tracker.record_call(other.clone()); + assert!( + tracker.check_loop(&other).is_none(), + "exempt calls must not count toward the repeat threshold" + ); + tracker.record_call(other.clone()); + assert!(tracker.check_loop(&other).is_some()); + } + } + + #[test] + fn tracker_exempt_interleave_does_not_mask_real_loop() { + let mut tracker = ToolCallTracker::default(); + let x = call_with_args("execute_command", json!({"command": "ls"})); + + tracker.record_call(call_with_args("job__check", json!({"id": "j1"}))); + tracker.record_call(x.clone()); + tracker.record_call(call_with_args("job__check", json!({"id": "j1"}))); + tracker.record_call(x.clone()); + + assert!(tracker.check_loop(&x).is_some()); + } + + #[test] + fn tracker_non_exempt_behavior_unchanged() { + let mut tracker = ToolCallTracker::default(); + let c = call_with_args("fs_cat", json!({"path": "a.txt"})); + + tracker.record_call(c.clone()); + tracker.record_call(c.clone()); + + assert!(tracker.check_loop(&c).is_some()); + } + #[test] fn prefix_constants_are_correct() { assert_eq!(TODO_FUNCTION_PREFIX, "todo__"); - assert_eq!(SUPERVISOR_FUNCTION_PREFIX, "agent__"); + assert_eq!(AGENT_FUNCTION_PREFIX, "agent__"); assert_eq!(USER_FUNCTION_PREFIX, "user__"); assert_eq!(MCP_INVOKE_META_FUNCTION_NAME_PREFIX, "mcp_invoke"); assert_eq!(MCP_SEARCH_META_FUNCTION_NAME_PREFIX, "mcp_search"); @@ -2853,6 +3322,43 @@ mod tests { assert!(f.contains("agent__reply_escalation")); } + #[test] + fn functions_append_job_adds_declarations() { + let mut f = Functions::default(); + + f.append_job_functions(); + + assert!(f.contains("job__start")); + assert!(f.contains("job__check")); + assert!(f.contains("job__collect")); + assert!(f.contains("job__cancel")); + assert!(f.contains("job__list")); + } + + #[test] + fn eval_routes_declared_job_calls_to_job_handlers() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.functions.append_job_functions(); + let calls = vec![call("job__list", Some("id-1"))]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].output["active_jobs"], 0); + assert_eq!(results[0].output["jobs"], json!([])); + } + + #[test] + fn eval_soft_fails_job_calls_when_jobs_not_declared() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + let calls = vec![call("job__start", Some("id-1"))]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + let err = results[0].output["tool_call_error"].as_str().unwrap(); + assert!(err.contains("Unexpected call")); + } + #[test] fn functions_append_teammate_adds_declarations() { let mut f = Functions::default(); @@ -4004,4 +4510,349 @@ mod tests { assert!(dir.is_dir()); fs::remove_dir_all(&dir).unwrap(); } + + #[test] + fn eval_tool_calls_partitions_mcp_and_sequential_then_resorts() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + let calls = vec![ + call("unknown_first", Some("id-1")), + ToolCall::new( + "mcp_search_foo".into(), + json!({"query": "q"}), + Some("id-2".into()), + ), + call("unknown_last", Some("id-3")), + ]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + assert_eq!(results.len(), 3); + assert_eq!(results[0].call.name, "unknown_first"); + assert_eq!(results[1].call.name, "mcp_search_foo"); + assert_eq!(results[2].call.name, "unknown_last"); + + for sequential in [&results[0], &results[2]] { + let err = sequential.output["tool_call_error"].as_str().unwrap(); + assert!( + err.contains("use only tools listed in your catalog"), + "{err}" + ); + } + let mcp_err = results[1].output["tool_call_error"].as_str().unwrap(); + assert!(mcp_err.starts_with("MCP search failed"), "{mcp_err}"); + assert!(!mcp_err.contains("use only tools listed in your catalog")); + } + + #[test] + fn eval_tool_calls_isolates_failures_within_a_batch() { + let app = AppState { + config: Arc::new(AppConfig { + auto_continue: true, + ..Default::default() + }), + ..AppState::test_default() + }; + let mut ctx = RequestContext::new(Arc::new(app), WorkingMode::Cmd); + ctx.tool_scope.functions.append_todo_functions(); + let calls = vec![ + ToolCall::new( + "todo__init".into(), + json!({"goal": "ship it"}), + Some("id-1".into()), + ), + call("unknown_tool", Some("id-2")), + ]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].call.name, "todo__init"); + assert_eq!(results[0].output["status"], "ok"); + assert!(results[0].output.get("tool_call_error").is_none()); + let err = results[1].output["tool_call_error"].as_str().unwrap(); + assert!(err.contains("Unexpected call"), "{err}"); + } + + #[test] + fn eval_tool_calls_reports_loop_alert_without_executing() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + let looped = call_with_args("looped_tool", json!({"a": 1})); + ctx.tool_scope.tool_tracker.record_call(looped.clone()); + ctx.tool_scope.tool_tracker.record_call(looped.clone()); + let calls = vec![ + looped, + ToolCall::new("other_tool".into(), json!({}), Some("id-2".into())), + ]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + assert_eq!(results.len(), 2); + let alert = results[0].output.as_str().unwrap(); + assert!(alert.starts_with("{\"tool_call_loop_alert\":"), "{alert}"); + let err = results[1].output["tool_call_error"].as_str().unwrap(); + assert!(err.contains("Unexpected call"), "{err}"); + } + + #[test] + fn eval_tool_calls_truncates_with_global_max_chars() { + let app = AppState { + config: Arc::new(AppConfig { + max_tool_result_chars: Some(50), + ..Default::default() + }), + ..AppState::test_default() + }; + let mut ctx = RequestContext::new(Arc::new(app), WorkingMode::Cmd); + let calls = vec![ToolCall::new( + "unknown_tool".into(), + json!({"padding": "x".repeat(200)}), + Some("id-1".into()), + )]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + let out = results[0].output.as_str().unwrap(); + assert!( + out.starts_with("[truncated: tool output exceeded 50 chars]\n"), + "{out}" + ); + } + + #[test] + fn eval_tool_calls_agent_max_chars_overrides_global() { + let app = AppState { + config: Arc::new(AppConfig { + max_tool_result_chars: Some(5000), + ..Default::default() + }), + ..AppState::test_default() + }; + let mut ctx = RequestContext::new(Arc::new(app), WorkingMode::Cmd); + ctx.agent = Some(Agent::test_new(AgentConfig { + max_tool_result_chars: Some(30), + ..Default::default() + })); + let calls = vec![ToolCall::new( + "unknown_tool".into(), + json!({"padding": "x".repeat(200)}), + Some("id-1".into()), + )]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + let out = results[0].output.as_str().unwrap(); + assert!( + out.starts_with("[truncated: tool output exceeded 30 chars]\n"), + "{out}" + ); + } + + #[test] + fn eval_tool_calls_zero_max_chars_disables_truncation() { + let app = AppState { + config: Arc::new(AppConfig { + max_tool_result_chars: Some(0), + ..Default::default() + }), + ..AppState::test_default() + }; + let mut ctx = RequestContext::new(Arc::new(app), WorkingMode::Cmd); + let calls = vec![ToolCall::new( + "unknown_tool".into(), + json!({"padding": "x".repeat(200)}), + Some("id-1".into()), + )]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + assert!(results[0].output["tool_call_error"].is_string()); + assert!(!results[0].output.to_string().contains("[truncated")); + } + + #[test] + fn eval_tool_calls_no_max_chars_configured_never_truncates() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + let calls = vec![ToolCall::new( + "unknown_tool".into(), + json!({"padding": "x".repeat(200)}), + Some("id-1".into()), + )]; + + let results = run_async(eval_tool_calls(&mut ctx, calls)).unwrap(); + + assert!(results[0].output["tool_call_error"].is_string()); + assert!(!results[0].output.to_string().contains("[truncated")); + } + + /// When the char cap lands inside a multi-byte UTF-8 character of the + /// serialized output, the cut is floored to the previous char boundary + /// so the output actually shrinks. + #[test] + fn truncate_if_needed_floors_cut_to_char_boundary() { + let serialized = json!("aé").to_string(); + let result = ToolResult::new(call("t", Some("id-1")), json!("aé")); + + let truncated = result.truncate_if_needed(3); + + let out = truncated.output.as_str().unwrap(); + assert_eq!(out, "[truncated: tool output exceeded 3 chars]\n\"a"); + assert!(out.len() < "[truncated: tool output exceeded 3 chars]\n".len() + serialized.len()); + } + + #[test] + fn truncate_if_needed_cap_on_char_boundary_truncates_normally() { + let result = ToolResult::new(call("t", Some("id-1")), json!("aé")); + + let truncated = result.truncate_if_needed(2); + + assert_eq!( + truncated.output.as_str().unwrap(), + "[truncated: tool output exceeded 2 chars]\n\"a" + ); + } + + #[test] + fn truncate_if_needed_cap_zero_yields_marker_and_empty_prefix() { + let result = ToolResult::new(call("t", Some("id-1")), json!("aé")); + + let truncated = result.truncate_if_needed(0); + + assert_eq!( + truncated.output.as_str().unwrap(), + "[truncated: tool output exceeded 0 chars]\n" + ); + } + + #[test] + fn truncate_if_needed_cap_at_or_above_length_leaves_output_unchanged() { + for cap in [5, 100] { + let result = ToolResult::new(call("t", Some("id-1")), json!("aé")); + + let truncated = result.truncate_if_needed(cap); + + assert_eq!(truncated.output, json!("aé")); + } + } + + #[test] + fn eval_routes_agent_prefix_to_supervisor_handler() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.functions.append_supervisor_functions(); + + let out = + run_async(call_with_args("agent__check", json!({"id": "x"})).eval(&mut ctx)).unwrap(); + + let err = out["tool_call_error"].as_str().unwrap(); + assert!(err.starts_with("Agent tool failed"), "{err}"); + assert!(err.contains("No supervisor active"), "{err}"); + } + + #[test] + fn eval_routes_todo_prefix_to_todo_handler() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.functions.append_todo_functions(); + + let out = run_async(call_with_args("todo__list", json!({})).eval(&mut ctx)).unwrap(); + + let err = out["tool_call_error"].as_str().unwrap(); + assert!(err.starts_with("Todo tool failed"), "{err}"); + assert!(err.contains("Auto-continue is not enabled"), "{err}"); + } + + #[test] + fn eval_routes_memory_prefix_to_memory_handler() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.functions.append_memory_functions(); + + let out = run_async(call_with_args("memory__read", json!({})).eval(&mut ctx)).unwrap(); + + let err = out["tool_call_error"].as_str().unwrap(); + assert!(err.starts_with("Memory tool failed"), "{err}"); + assert!( + err.contains("name is required") || err.contains("Memory tools are disabled"), + "expected a memory-handler-owned error regardless of host memory files: {err}" + ); + } + + #[test] + fn eval_routes_skill_prefix_to_skill_handler() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.functions.append_skill_functions(); + + let out = run_async(call_with_args("skill__load", json!({})).eval(&mut ctx)).unwrap(); + + assert_eq!(out["error"], "name is required"); + } + + #[test] + fn eval_routes_user_prefix_to_user_handler() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.functions.append_user_interaction_functions(); + + let out = run_async(call_with_args("user__confirm", json!({})).eval(&mut ctx)).unwrap(); + + let err = out["tool_call_error"].as_str().unwrap(); + assert!(err.starts_with("User interaction failed"), "{err}"); + assert!(err.contains("'question' is required"), "{err}"); + } + + #[test] + fn eval_routes_rag_prefix_to_rag_handler() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.tool_scope.functions.append_rag_query_functions(); + + let out = + run_async(call_with_args("rag__query", json!({"query": "x"})).eval(&mut ctx)).unwrap(); + + let err = out["tool_call_error"].as_str().unwrap(); + assert!(err.starts_with("RAG query failed"), "{err}"); + assert!(err.contains("No RAG is attached"), "{err}"); + } + + #[test] + fn eval_unknown_name_errors_with_unexpected_call() { + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + + let err = run_async(call_with_args("nope", json!({})).eval(&mut ctx)).unwrap_err(); + + assert!(err.to_string().contains("Unexpected call"), "{err}"); + } + + #[test] + fn eval_mcp_empty_runtime_returns_distinct_error_per_prefix() { + let ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + let cases = [ + ( + "mcp_invoke_ghost", + json!({"tool": "t"}), + "MCP tool invocation failed", + ), + ( + "mcp_search_ghost", + json!({"query": "q"}), + "MCP search failed", + ), + ( + "mcp_describe_ghost", + json!({"tool": "t"}), + "MCP describe failed", + ), + ( + "mcp_read_ghost", + json!({"uri": "file:///x"}), + "MCP read failed", + ), + ( + "mcp_prompt_ghost", + json!({"prompt": "p"}), + "MCP prompt failed", + ), + ]; + + for (name, args, expected) in cases { + let out = run_async(call_with_args(name, args).eval_mcp(&ctx)).unwrap(); + let err = out["tool_call_error"].as_str().unwrap(); + assert!(err.starts_with(expected), "{name}: {err}"); + } + } } diff --git a/src/graph/agent.rs b/src/graph/agent.rs index de95cfd..9bde501 100644 --- a/src/graph/agent.rs +++ b/src/graph/agent.rs @@ -2,7 +2,7 @@ use super::state::StateManager; use super::structured; use super::types::AgentNode; use crate::config::RequestContext; -use crate::function::supervisor::run_agent_for_graph; +use crate::function::agents::run_agent_for_graph; use anyhow::{Context, Result}; use serde_json::Value; use std::time::Duration; diff --git a/src/graph/executor.rs b/src/graph/executor.rs index a53441d..4c2366b 100644 --- a/src/graph/executor.rs +++ b/src/graph/executor.rs @@ -563,8 +563,14 @@ mod tests { mod integration_tests { use super::*; use crate::config::{AppState, WorkingMode}; + #[cfg(unix)] + use crate::function::jobs::RingBuf; + #[cfg(unix)] + use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus, Supervisor, notification}; use crate::utils::{create_abort_signal, temp_file}; use std::fs; + #[cfg(unix)] + use std::mem; fn cmd_available(name: &str) -> bool { which::which(name).is_ok() @@ -856,4 +862,83 @@ nodes: ); assert!(err.contains("sleeper"), "error should name frontier: {err}"); } + + #[cfg(unix)] + #[tokio::test] + async fn background_job_survives_graph_node_execution() { + if !cmd_available("bash") { + eprintln!("skipping: bash not available"); + return; + } + let ws = TestWorkspace::new(); + ws.write_script("noop.sh", "#!/bin/bash\necho '{}'\n"); + + let yaml = r#" +name: background_job_survival_test +start: noop +nodes: + noop: + type: script + script: noop.sh + state_updates: {} + next: done + done: + type: end + output: "done" +"#; + let graph: Graph = serde_yaml::from_str(yaml).unwrap(); + + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let join_handle = rt.spawn(async { + Ok(JobResult { + output: Value::Null, + exit_code: Some(0), + output_bytes_captured: 0, + }) + }); + mem::forget(rt); + let handle = JobHandle { + id: "job_bg".to_string(), + tool: "execute_command".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal: create_abort_signal(), + state: Arc::new(parking_lot::Mutex::new(JobState { + status: JobStatus::Completed, + pgid: None, + })), + output_buf: Arc::new(parking_lot::Mutex::new(RingBuf::default())), + no_change_checks: 0, + last_check_state: None, + }; + let mut sup = Supervisor::new(0, 3).with_max_concurrent_jobs(4); + sup.register(handle).unwrap(); + + let mut ctx = make_ctx(); + ctx.supervisor = Some(Arc::new(parking_lot::RwLock::new(sup))); + ctx.notification_queue.push(notification::job_notification( + "job_bg", + "execute_command", + true, + )); + + let abort = create_abort_signal(); + let result = GraphExecutor::new(graph, &ws.dir) + .execute(&mut ctx, abort) + .await + .unwrap_or_else(|e| panic!("executor failed: {e:#}")); + assert_eq!(result, "done"); + + assert!( + ctx.supervisor.as_ref().unwrap().read().has_job("job_bg"), + "graph execution must not touch registered job handles" + ); + let events = ctx.notification_queue.drain(); + assert_eq!(events.len(), 1, "queued notification must survive the run"); + assert_eq!(events[0].id, "job_bg"); + assert_eq!(events[0].event, "job_completed"); + } } diff --git a/src/graph/llm.rs b/src/graph/llm.rs index b62f4a2..623d110 100644 --- a/src/graph/llm.rs +++ b/src/graph/llm.rs @@ -6,8 +6,9 @@ use crate::config::prompts::DEFAULT_SKILL_INSTRUCTIONS; use crate::config::{ Input, RequestContext, Role, RoleLike, SkillPolicy, should_inject_skill_instructions, }; +use crate::function::agents::{GuardrailAction, check_pending_tasks_guardrail}; +use crate::function::jobs::reap_jobs; use crate::function::skill::skill_function_declarations; -use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail}; use crate::utils::create_abort_signal; use anyhow::{Context, Error, Result, anyhow, bail}; use log::warn; @@ -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 } @@ -268,7 +275,7 @@ async fn run_chat_loop(node: &LlmNode, prompt: &str, ctx: &mut RequestContext) - } if tool_results.is_empty() { - match check_pending_agents_guardrail(ctx) { + match check_pending_tasks_guardrail(ctx) { GuardrailAction::NoAction => return Ok(accumulated), GuardrailAction::ForceTerminate(ids) => { warn!( diff --git a/src/graph/types.rs b/src/graph/types.rs index 18bdf5d..e45322b 100644 --- a/src/graph/types.rs +++ b/src/graph/types.rs @@ -28,6 +28,9 @@ pub struct Graph { #[serde(default, skip_serializing_if = "Option::is_none")] pub reasoning_effort: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_concurrent_jobs: Option, + #[serde(default)] pub global_tools: Vec, @@ -895,6 +898,7 @@ nodes: assert!(graph.model.is_none()); assert!(graph.temperature.is_none()); assert!(graph.top_p.is_none()); + assert!(graph.max_concurrent_jobs.is_none()); assert!(graph.global_tools.is_empty()); assert!(graph.mcp_servers.is_empty()); assert!(graph.conversation_starters.is_empty()); diff --git a/src/graph/validator.rs b/src/graph/validator.rs index fe438f7..4b7cc64 100644 --- a/src/graph/validator.rs +++ b/src/graph/validator.rs @@ -998,6 +998,7 @@ mod tests { temperature: None, top_p: None, reasoning_effort: None, + max_concurrent_jobs: None, global_tools: Vec::new(), mcp_servers: Vec::new(), skills_enabled: None, diff --git a/src/main.rs b/src/main.rs index 3e5d9b0..9309149 100644 --- a/src/main.rs +++ b/src/main.rs @@ -29,7 +29,7 @@ use crate::config::{ install_builtins, list_agents, load_env_file, macro_execute, sync_models, }; use crate::config::{memory, paths}; -use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail}; +use crate::function::agents::{GuardrailAction, check_pending_tasks_guardrail}; use crate::mcp::McpServersConfig; use crate::render::{prompt_theme, render_error}; use crate::repl::Repl; @@ -595,7 +595,7 @@ async fn start_directive( ) .await?; } else { - match check_pending_agents_guardrail(ctx) { + match check_pending_tasks_guardrail(ctx) { GuardrailAction::Inject(prompt) => { let guardrail_input = Input::from_str(ctx, &prompt, None)?; return start_directive(ctx, guardrail_input, code_mode, abort_signal).await; diff --git a/src/repl/mod.rs b/src/repl/mod.rs index f3b0d0f..c53202c 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -16,7 +16,7 @@ use crate::config::{ StateFlags, flatten_prompt_messages, macro_execute, resolve_prompt_args, sanitize_display_text, }; use crate::config::{AssetCategory, paths}; -use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail}; +use crate::function::agents::{GuardrailAction, check_pending_tasks_guardrail}; use crate::render::render_error; use crate::utils::{ AbortSignal, SHELL, abortable_run_with_spinner, create_abort_signal, dimmed_text, @@ -602,7 +602,7 @@ pub async fn run_repl_command( abort_signal: AbortSignal, mut line: &str, ) -> Result { - ctx.pending_agents_guardrail_count = 0; + ctx.pending_tasks_guardrail_count = 0; if let Ok(Some(captures)) = MULTILINE_RE.captures(line) && let Some(text_match) = captures.get(1) { @@ -1475,7 +1475,7 @@ async fn ask( ) .await } else { - match check_pending_agents_guardrail(ctx) { + match check_pending_tasks_guardrail(ctx) { GuardrailAction::Inject(prompt) => { let guardrail_input = Input::from_str(ctx, &prompt, None)?; return ask(ctx, abort_signal, guardrail_input, false).await; diff --git a/src/supervisor/mod.rs b/src/supervisor/mod.rs index dc55d21..d3e0487 100644 --- a/src/supervisor/mod.rs +++ b/src/supervisor/mod.rs @@ -1,17 +1,21 @@ pub mod escalation; pub mod mailbox; +pub mod notification; pub mod taskqueue; +use crate::function::jobs::RingBuf; use crate::utils::AbortSignal; use fmt::{Debug, Formatter}; use mailbox::Inbox; -use parking_lot::RwLock; +use parking_lot::{Mutex, RwLock}; use taskqueue::TaskQueue; use anyhow::{Result, bail}; +use serde_json::Value; use std::collections::HashMap; use std::fmt; use std::sync::Arc; +use std::time::Instant; use tokio::task::JoinHandle; #[derive(Debug, Clone, PartialEq, Eq)] @@ -37,11 +41,85 @@ pub struct AgentHandle { pub child_supervisor: Option>>, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum JobStatus { + Running, + Completed, + Failed, +} + +pub struct JobState { + pub status: JobStatus, + pub pgid: Option, +} + +pub struct JobResult { + pub output: Value, + pub exit_code: Option, + pub output_bytes_captured: u64, +} + +pub struct JobHandle { + pub id: String, + pub tool: String, + pub started_at: Instant, + pub join_handle: JoinHandle>, + pub abort_signal: AbortSignal, + pub state: Arc>, + pub output_buf: Arc>, + pub no_change_checks: u32, + pub last_check_state: Option<(JobStatus, u64)>, +} + +impl JobHandle { + // pgid == child pid under process_group(0); after wait() reaps the child + // the pid can be recycled, so never kill unless pgid is still set. + fn kill_process_group(&self) { + #[cfg(unix)] + if let Some(pgid) = self.state.lock().pgid { + unsafe { + libc::killpg(pgid, libc::SIGTERM); + } + } + } +} + +impl Drop for JobHandle { + fn drop(&mut self) { + self.kill_process_group(); + self.join_handle.abort(); + } +} + +pub enum TaskHandle { + Agent(AgentHandle), + Job(JobHandle), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TaskKind { + Agent, + Job, +} + +impl From for TaskHandle { + fn from(handle: AgentHandle) -> Self { + Self::Agent(handle) + } +} + +impl From for TaskHandle { + fn from(handle: JobHandle) -> Self { + Self::Job(handle) + } +} + pub struct Supervisor { - handles: HashMap, + handles: HashMap, task_queue: TaskQueue, max_concurrent: usize, max_depth: usize, + max_concurrent_jobs: usize, } impl Supervisor { @@ -51,17 +129,64 @@ impl Supervisor { task_queue: TaskQueue::new(), max_concurrent, max_depth, + max_concurrent_jobs: 0, } } + pub fn with_max_concurrent_jobs(mut self, max_concurrent_jobs: usize) -> Self { + self.max_concurrent_jobs = max_concurrent_jobs; + self + } + + fn agent(&self, id: &str) -> Option<&AgentHandle> { + match self.handles.get(id) { + Some(TaskHandle::Agent(handle)) => Some(handle), + _ => None, + } + } + + fn agents(&self) -> impl Iterator { + self.handles.values().filter_map(|handle| match handle { + TaskHandle::Agent(handle) => Some(handle), + TaskHandle::Job(_) => None, + }) + } + + pub fn job(&self, id: &str) -> Option<&JobHandle> { + match self.handles.get(id) { + Some(TaskHandle::Job(handle)) => Some(handle), + _ => None, + } + } + + pub fn job_mut(&mut self, id: &str) -> Option<&mut JobHandle> { + match self.handles.get_mut(id) { + Some(TaskHandle::Job(handle)) => Some(handle), + _ => None, + } + } + + pub fn jobs(&self) -> impl Iterator { + self.handles.values().filter_map(|handle| match handle { + TaskHandle::Job(handle) => Some(handle), + TaskHandle::Agent(_) => None, + }) + } + pub fn active_count(&self) -> usize { - self.handles.len() + self.agents().count() } pub fn effective_active_count(&self) -> usize { + self.agents() + .filter(|h| !h.join_handle.is_finished()) + .count() + } + + pub fn active_job_count(&self) -> usize { self.handles .values() - .filter(|h| !h.join_handle.is_finished()) + .filter(|h| matches!(h, TaskHandle::Job(job) if !job.join_handle.is_finished())) .count() } @@ -73,6 +198,10 @@ impl Supervisor { self.max_depth } + pub fn max_concurrent_jobs(&self) -> usize { + self.max_concurrent_jobs + } + pub fn task_queue(&self) -> &TaskQueue { &self.task_queue } @@ -81,59 +210,128 @@ impl Supervisor { &mut self.task_queue } - pub fn register(&mut self, handle: AgentHandle) -> Result<()> { - if self.effective_active_count() >= self.max_concurrent { - bail!( - "Cannot spawn agent: at capacity ({}/{})", - self.effective_active_count(), - self.max_concurrent - ); + pub fn register(&mut self, handle: impl Into) -> Result<()> { + match handle.into() { + TaskHandle::Agent(handle) => { + if self.effective_active_count() >= self.max_concurrent { + bail!( + "Cannot spawn agent: at capacity ({}/{})", + self.effective_active_count(), + self.max_concurrent + ); + } + if handle.depth > self.max_depth { + bail!( + "Cannot spawn agent: max depth exceeded ({}/{})", + handle.depth, + self.max_depth + ); + } + self.handles + .insert(handle.id.clone(), TaskHandle::Agent(handle)); + } + TaskHandle::Job(handle) => { + if self.active_job_count() >= self.max_concurrent_jobs { + bail!( + "Cannot start job: at capacity ({}/{})", + self.active_job_count(), + self.max_concurrent_jobs + ); + } + self.handles + .insert(handle.id.clone(), TaskHandle::Job(handle)); + } } - if handle.depth > self.max_depth { - bail!( - "Cannot spawn agent: max depth exceeded ({}/{})", - handle.depth, - self.max_depth - ); - } - self.handles.insert(handle.id.clone(), handle); Ok(()) } pub fn is_finished(&self, id: &str) -> Option { - self.handles.get(id).map(|h| h.join_handle.is_finished()) + self.agent(id).map(|h| h.join_handle.is_finished()) } pub fn take(&mut self, id: &str) -> Option { - self.handles.remove(id) + self.agent(id)?; + match self.handles.remove(id) { + Some(TaskHandle::Agent(handle)) => Some(handle), + _ => None, + } + } + + pub fn take_job(&mut self, id: &str) -> Option { + if !self.has_job(id) { + return None; + } + match self.handles.remove(id) { + Some(TaskHandle::Job(handle)) => Some(handle), + _ => None, + } + } + + pub fn has_job(&self, id: &str) -> bool { + matches!(self.handles.get(id), Some(TaskHandle::Job(_))) + } + + pub fn has_agent(&self, id: &str) -> bool { + self.agent(id).is_some() } pub fn inbox(&self, id: &str) -> Option<&Arc> { - self.handles.get(id).map(|h| &h.inbox) + self.agent(id).map(|h| &h.inbox) } pub fn abort_signal_for(&self, id: &str) -> Option { - self.handles.get(id).map(|h| h.abort_signal.clone()) + self.agent(id).map(|h| h.abort_signal.clone()) } pub fn list_agents(&self) -> Vec<(&str, &str)> { + self.agents() + .map(|h| (h.id.as_str(), h.agent_name.as_str())) + .collect() + } + + pub fn list_tasks(&self) -> Vec<(&str, TaskKind, bool)> { self.handles .values() - .map(|h| (h.id.as_str(), h.agent_name.as_str())) + .map(|handle| match handle { + TaskHandle::Agent(agent) => ( + agent.id.as_str(), + TaskKind::Agent, + agent.join_handle.is_finished(), + ), + TaskHandle::Job(job) => ( + job.id.as_str(), + TaskKind::Job, + job.join_handle.is_finished(), + ), + }) .collect() } pub fn cancel_all(&self) { for handle in self.handles.values() { - handle.abort_signal.set_ctrlc(); + match handle { + TaskHandle::Agent(agent) => agent.abort_signal.set_ctrlc(), + TaskHandle::Job(job) => { + job.abort_signal.set_ctrlc(); + job.kill_process_group(); + } + } } } pub fn cancel_recursive(&self) { for handle in self.handles.values() { - handle.abort_signal.set_ctrlc(); - if let Some(child_sup) = handle.child_supervisor.as_ref() { - child_sup.read().cancel_recursive(); + match handle { + TaskHandle::Agent(agent) => { + agent.abort_signal.set_ctrlc(); + if let Some(child_sup) = agent.child_supervisor.as_ref() { + child_sup.read().cancel_recursive(); + } + } + TaskHandle::Job(job) => { + job.abort_signal.set_ctrlc(); + job.kill_process_group(); + } } } } @@ -142,7 +340,7 @@ impl Supervisor { impl Debug for Supervisor { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { f.debug_struct("Supervisor") - .field("active_agents", &self.handles.len()) + .field("active_agents", &self.active_count()) .field("max_concurrent", &self.max_concurrent) .field("max_depth", &self.max_depth) .finish() @@ -154,6 +352,7 @@ mod tests { use super::*; use crate::utils::create_abort_signal; use anyhow::Error; + use std::mem; use tokio::runtime::Builder; fn make_handle(id: &str, agent_name: &str, depth: usize) -> AgentHandle { @@ -177,6 +376,34 @@ mod tests { } } + fn make_job(id: &str, abort_signal: AbortSignal) -> JobHandle { + // Keep the runtime alive so the spawned task is never polled and the + // job counts as running for capacity checks. + let rt = Builder::new_current_thread().enable_all().build().unwrap(); + let join_handle = rt.spawn(async { + Ok(JobResult { + output: Value::Null, + exit_code: Some(0), + output_bytes_captured: 0, + }) + }); + mem::forget(rt); + JobHandle { + id: id.to_string(), + tool: "execute_command".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal, + state: Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })), + output_buf: Arc::new(Mutex::new(RingBuf::default())), + no_change_checks: 0, + last_check_state: None, + } + } + #[test] fn supervisor_new_empty() { let sup = Supervisor::new(4, 3); @@ -294,4 +521,111 @@ mod tests { AgentExitStatus::Failed("x".into()) ); } + + #[test] + fn cancel_recursive_aborts_nested_supervisors() { + let child_sig = create_abort_signal(); + let mut child_handle = make_handle("c1", "worker", 2); + child_handle.abort_signal = child_sig.clone(); + let mut child_sup = Supervisor::new(4, 3); + child_sup.register(child_handle).unwrap(); + + let parent_sig = create_abort_signal(); + let mut parent_handle = make_handle("a1", "explore", 1); + parent_handle.abort_signal = parent_sig.clone(); + parent_handle.child_supervisor = Some(Arc::new(RwLock::new(child_sup))); + let mut sup = Supervisor::new(4, 3); + sup.register(parent_handle).unwrap(); + + sup.cancel_recursive(); + + assert!(parent_sig.aborted()); + assert!(child_sig.aborted()); + } + + #[test] + fn job_registration_rejects_when_job_capacity_zero() { + let mut sup = Supervisor::new(4, 3); + + let result = sup.register(make_job("j1", create_abort_signal())); + + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("at capacity")); + } + + #[test] + fn job_registration_rejects_at_job_capacity() { + let mut sup = Supervisor::new(4, 3).with_max_concurrent_jobs(1); + sup.register(make_job("j1", create_abort_signal())).unwrap(); + + let result = sup.register(make_job("j2", create_abort_signal())); + + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("at capacity")); + } + + #[test] + fn job_capacity_is_independent_of_agent_capacity() { + let mut sup = Supervisor::new(1, 3).with_max_concurrent_jobs(1); + + sup.register(make_job("j1", create_abort_signal())).unwrap(); + sup.register(make_handle("a1", "explore", 1)).unwrap(); + + assert_eq!(sup.active_job_count(), 1); + assert_eq!(sup.active_count(), 1); + assert_eq!(sup.max_concurrent_jobs(), 1); + } + + #[test] + fn agent_accessors_ignore_jobs() { + let mut sup = Supervisor::new(4, 3).with_max_concurrent_jobs(2); + + sup.register(make_job("j1", create_abort_signal())).unwrap(); + + assert_eq!(sup.active_count(), 0); + assert_eq!(sup.effective_active_count(), 0); + assert!(sup.list_agents().is_empty()); + assert_eq!(sup.is_finished("j1"), None); + assert!(sup.inbox("j1").is_none()); + assert!(sup.abort_signal_for("j1").is_none()); + assert!(sup.take("j1").is_none()); + assert!(sup.has_job("j1")); + assert!(!sup.has_agent("j1")); + assert_eq!(sup.active_job_count(), 1); + } + + #[test] + fn take_job_removes_job_but_not_agents() { + let mut sup = Supervisor::new(4, 3).with_max_concurrent_jobs(2); + + sup.register(make_job("j1", create_abort_signal())).unwrap(); + sup.register(make_handle("a1", "explore", 1)).unwrap(); + + assert!(sup.take_job("a1").is_none()); + assert!(sup.has_agent("a1")); + assert!(sup.take_job("j1").is_some()); + assert_eq!(sup.active_job_count(), 0); + } + + #[test] + fn cancel_recursive_aborts_jobs() { + let sig = create_abort_signal(); + let mut sup = Supervisor::new(4, 3).with_max_concurrent_jobs(1); + sup.register(make_job("j1", sig.clone())).unwrap(); + + sup.cancel_recursive(); + + assert!(sig.aborted()); + } + + #[test] + fn cancel_all_aborts_jobs() { + let sig = create_abort_signal(); + let mut sup = Supervisor::new(4, 3).with_max_concurrent_jobs(1); + sup.register(make_job("j1", sig.clone())).unwrap(); + + sup.cancel_all(); + + assert!(sig.aborted()); + } } diff --git a/src/supervisor/notification.rs b/src/supervisor/notification.rs new file mode 100644 index 0000000..4f9261a --- /dev/null +++ b/src/supervisor/notification.rs @@ -0,0 +1,168 @@ +use fmt::{Debug, Formatter}; +use serde_json::{Value, json}; +use std::fmt; + +/// One background-task completion event, delivered to the context that +/// started the task by merging a `system_notifications` entry onto the last +/// tool result of a batch. +#[derive(Clone)] +pub struct SystemNotification { + pub event: &'static str, + pub id: String, + pub tool_or_agent: String, + pub status: &'static str, + pub next_action: String, +} + +impl SystemNotification { + pub fn to_value(&self) -> Value { + json!({ + "event": self.event, + "id": self.id, + "tool_or_agent": self.tool_or_agent, + "status": self.status, + "next_action": self.next_action, + }) + } +} + +pub fn job_notification(id: &str, tool: &str, success: bool) -> SystemNotification { + SystemNotification { + event: if success { + "job_completed" + } else { + "job_failed" + }, + id: id.to_string(), + tool_or_agent: tool.to_string(), + status: if success { "success" } else { "failed" }, + next_action: format!("job__collect --id {id} for output"), + } +} + +pub fn agent_notification(id: &str, agent_name: &str, success: bool) -> SystemNotification { + SystemNotification { + event: if success { + "agent_completed" + } else { + "agent_failed" + }, + id: id.to_string(), + tool_or_agent: agent_name.to_string(), + status: if success { "success" } else { "failed" }, + next_action: format!("agent__collect --id {id} for output"), + } +} + +/// Completion events for background work started by ONE context. Unlike the +/// escalation queue (shared, root-owned), every context owns a fresh queue: +/// a queue shared between parent and child would race their drains and +/// deliver one context's events into the other's transcript. +pub struct NotificationQueue { + pending: parking_lot::Mutex>, +} + +impl NotificationQueue { + pub fn new() -> Self { + Self { + pending: parking_lot::Mutex::new(Vec::new()), + } + } + + pub fn push(&self, notification: SystemNotification) { + self.pending.lock().push(notification); + } + + pub fn drain(&self) -> Vec { + std::mem::take(&mut *self.pending.lock()) + } +} + +impl Default for NotificationQueue { + fn default() -> Self { + Self::new() + } +} + +impl Debug for NotificationQueue { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let count = self.pending.lock().len(); + f.debug_struct("NotificationQueue") + .field("pending_count", &count) + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn job_notification_success_shape() { + let event = job_notification("job_a1b2", "execute_command", true); + assert_eq!( + event.to_value(), + json!({ + "event": "job_completed", + "id": "job_a1b2", + "tool_or_agent": "execute_command", + "status": "success", + "next_action": "job__collect --id job_a1b2 for output", + }) + ); + } + + #[test] + fn job_notification_failure_shape() { + let event = job_notification("job_a1b2", "execute_command", false); + assert_eq!(event.event, "job_failed"); + assert_eq!(event.status, "failed"); + assert_eq!(event.next_action, "job__collect --id job_a1b2 for output"); + } + + #[test] + fn agent_notification_success_shape() { + let event = agent_notification("agent_explore_a1b2", "explore", true); + assert_eq!( + event.to_value(), + json!({ + "event": "agent_completed", + "id": "agent_explore_a1b2", + "tool_or_agent": "explore", + "status": "success", + "next_action": "agent__collect --id agent_explore_a1b2 for output", + }) + ); + } + + #[test] + fn agent_notification_failure_shape() { + let event = agent_notification("agent_explore_a1b2", "explore", false); + assert_eq!(event.event, "agent_failed"); + assert_eq!(event.status, "failed"); + assert_eq!( + event.next_action, + "agent__collect --id agent_explore_a1b2 for output" + ); + } + + #[test] + fn drain_empties_queue_and_preserves_order() { + let queue = NotificationQueue::new(); + queue.push(job_notification("job_1", "execute_command", true)); + queue.push(job_notification("job_2", "execute_command", false)); + + let drained = queue.drain(); + + assert_eq!(drained.len(), 2); + assert_eq!(drained[0].id, "job_1"); + assert_eq!(drained[1].id, "job_2"); + assert!(queue.drain().is_empty()); + } + + #[test] + fn drain_on_empty_queue_is_a_noop() { + let queue = NotificationQueue::default(); + assert!(queue.drain().is_empty()); + } +}