From 9edb8729ef1f97f01a6b5741dae24aadd361f55a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Mon, 27 Jul 2026 01:22:52 -0400 Subject: [PATCH] =?UTF-8?q?feat(swarm):=20agent=5Fswarm=20=E2=80=94=20one?= =?UTF-8?q?=20prompt=20over=20many=20items,=20paced=20as=20a=20fleet?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports kimi-code's AgentSwarm: a `prompt_template` containing `{{item}}` expanded over an `items` list into up to 128 subagents, run to completion and returned as one aggregate. Kigi already exceeds upstream on planning, verification, isolation and merge via /graph; what it lacked was cheap immediate fan-out. Entirely client-side — no new backend surface. The engine is three pure pieces plus a runner: `plan` validates and expands (every fault reported before a single member starts — a half-launched swarm is expensive to unwind), `schedule` is the launch ramp as testable arithmetic, `run` drives it against the existing `SubagentBackend`. Reuses the single-spawn coordinator rather than inventing a batch API. Load-bearing decisions, each the result of a defect found in review: - `backgrounded` is its own outcome. A member that outlives the 600s foreground budget is detached by the coordinator and KEEPS RUNNING; reporting it as failed invites the model to relaunch its item, putting a second agent on the same files. It is never offered for resume. - `InFlightGuard` cancels live members on Drop. Send-now cancels the turn WITHOUT cancelling subagents and aborts the task; the dropped receivers read as "parent gone" and each child re-attaches itself. There is no cooperative path to use instead — `Cancellation` is constructed nowhere in the tree — so Drop is the only seam that fires. - Retries and wall clock are both bounded. The swarm blocks the caller's turn, so every wait needs a ceiling it cannot argue past; stragglers at the deadline are reported as still-running, with their ids. - `ToolKind::AgentSwarm` is its own variant: `TemplateRenderer`'s `by_kind` map holds one tool name per kind, so sharing `Task` would silently redirect `${{ tools.by_kind.task }}` in other tools' prompts. - An explicitly requested model that cannot be validated is refused, as the task tool already does — one loud error beats `items.len()` quiet ones. Depth stays capped at 1: upstream's unlimited nesting is a hazard, not a feature. - `SubagentResult.rate_limited` is classified where the typed ACP error code is still in hand; a scheduler re-deriving it from a formatted string would stop adapting the day the wording changed. - Aggregate output is clamped per member (head+tail, loss stated): native tool output is truncated nowhere downstream. 42 agent_swarm tests. The fake backend awaits, so the concurrency and ordering assertions can actually fail; the cap test also proves the fixture can exceed the cap. --- crates/codegen/kigi-agent/src/builder.rs | 37 +- crates/codegen/kigi-agent/src/config.rs | 10 + .../src/agent/subagent/handle_request.rs | 12 +- .../session/acp_session_impl/tool_calls.rs | 12 + .../kigi-tools/schema/tool_meta.schema.json | 2 +- .../implementations/kigi/agent_swarm/mod.rs | 7 + .../implementations/kigi/agent_swarm/plan.rs | 377 +++++++++ .../implementations/kigi/agent_swarm/run.rs | 773 ++++++++++++++++++ .../kigi/agent_swarm/schedule.rs | 348 ++++++++ .../implementations/kigi/agent_swarm/tool.rs | 498 +++++++++++ .../src/implementations/kigi/mod.rs | 2 + .../src/implementations/kigi/task/types.rs | 15 +- .../kigi-tools/src/implementations/mod.rs | 6 +- .../codegen/kigi-tools/src/normalization.rs | 1 + .../codegen/kigi-tools/src/registry/types.rs | 1 + .../codegen/kigi-tools/src/tool_taxonomy.rs | 2 + crates/codegen/kigi-tools/src/types/tool.rs | 5 + .../codegen/kigi-tools/src/types/tool_io.rs | 2 + crates/codegen/kigi-tui/src/acp/tracker.rs | 11 +- .../codegen/kigi-workspace/src/capability.rs | 6 +- .../common/kigi-tool-types/src/agent_swarm.rs | 116 +++ crates/common/kigi-tool-types/src/lib.rs | 5 + 22 files changed, 2232 insertions(+), 16 deletions(-) create mode 100644 crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/mod.rs create mode 100644 crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/plan.rs create mode 100644 crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/run.rs create mode 100644 crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/schedule.rs create mode 100644 crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/tool.rs create mode 100644 crates/common/kigi-tool-types/src/agent_swarm.rs diff --git a/crates/codegen/kigi-agent/src/builder.rs b/crates/codegen/kigi-agent/src/builder.rs index b7793c0..4e136bc 100644 --- a/crates/codegen/kigi-agent/src/builder.rs +++ b/crates/codegen/kigi-agent/src/builder.rs @@ -711,9 +711,19 @@ impl AgentBuilder { kigi_tools::types::tool::ToolNamespace::Kigi, "task" ); + // The swarm is a fan-out of subagent spawns, so it lives and dies with + // the task tool: any condition that leaves no subagent to spawn leaves + // the swarm with nothing to fan out to. + let swarm_tool_id = format!( + "{}:{}", + kigi_tools::types::tool::ToolNamespace::Kigi, + "agent_swarm" + ); let mut task_stripped = false; if !self.subagents_enabled { - tool_config.tools.retain(|tc| tc.id != task_tool_id); + tool_config + .tools + .retain(|tc| tc.id != task_tool_id && tc.id != swarm_tool_id); task_stripped = true; } else { let subagents = crate::discovery::all_subagents_with_plugins( @@ -722,7 +732,9 @@ impl AgentBuilder { self.plugin_registry.as_deref(), ); if subagents.is_empty() { - tool_config.tools.retain(|tc| tc.id != task_tool_id); + tool_config + .tools + .retain(|tc| tc.id != task_tool_id && tc.id != swarm_tool_id); task_stripped = true; } else if self.prompt_audience == crate::prompt::context::PromptAudience::Subagent { if let Some(task_tc) = tool_config @@ -810,7 +822,13 @@ impl AgentBuilder { .tools .iter() .any(|t| AGENT_TASK_CLASSIFIER_RE.is_match(t)); - let task_deps = ["task", "get_task_output", "kill_task", "wait_tasks"]; + let task_deps = [ + "task", + "agent_swarm", + "get_task_output", + "kill_task", + "wait_tasks", + ]; let registered_tool_ids = tool_bridge_builder.known_tool_ids(); let present_kinds: std::collections::HashSet = tool_config.tools.iter().filter_map(|tc| tc.kind).collect(); @@ -925,7 +943,13 @@ impl AgentBuilder { } } if definition.allowed_subagent_types.as_deref() == Some(&[]) { - let task_deps = ["task", "get_task_output", "kill_task", "wait_tasks"]; + let task_deps = [ + "task", + "agent_swarm", + "get_task_output", + "kill_task", + "wait_tasks", + ]; tool_config .tools .retain(|tc| !task_deps.contains(&short_tool_name(&tc.id))); @@ -1601,6 +1625,11 @@ mod tests { has_task, *subagents, "[{label}] spawn_subagent presence should match subagents_enabled={subagents}; got tools: {names:?}" ); + let has_swarm = names.contains(&"agent_swarm"); + assert_eq!( + has_swarm, *subagents, + "[{label}] agent_swarm fans out to subagents, so it must follow the same gate as spawn_subagent; got tools: {names:?}" + ); assert!( names.contains(&"enter_plan_mode"), "[{label}] enter_plan_mode must always be present (TUI plan-mode keybind needs it); got tools: {names:?}" diff --git a/crates/codegen/kigi-agent/src/config.rs b/crates/codegen/kigi-agent/src/config.rs index fff2cf2..d0fda20 100644 --- a/crates/codegen/kigi-agent/src/config.rs +++ b/crates/codegen/kigi-agent/src/config.rs @@ -154,6 +154,11 @@ fn task_tool_config() -> ToolConfig { .with_name("spawn_subagent") .with_param_rename("run_in_background", "background") } +/// Swarm tool. Keeps its registry name: unlike `task` it has no CLI-specific +/// alias, and the name is what the model is told to reuse for `resume_agent_ids`. +fn agent_swarm_tool_config() -> ToolConfig { + ToolConfig::from(&kigi::AgentSwarmTool) +} /// Task output tool renamed for clarity: /// `get_task_output` → `get_command_or_subagent_output`. fn task_output_tool_config() -> ToolConfig { @@ -270,6 +275,7 @@ fn default_kigi_toolset() -> ToolServerConfig { task_output_tool_config(), wait_tasks_tool_config(), task_tool_config(), + agent_swarm_tool_config(), (&kigi::SchedulerCreateTool).into(), (&kigi::SchedulerDeleteTool).into(), (&kigi::SchedulerListTool).into(), @@ -318,6 +324,7 @@ pub fn kigi_hashline_toolset( task_output_tool_config(), wait_tasks_tool_config(), task_tool_config(), + agent_swarm_tool_config(), (&kigi::WebSearchTool).into(), (&kigi::SchedulerCreateTool).into(), (&kigi::SchedulerDeleteTool).into(), @@ -400,6 +407,7 @@ fn kigi_plan_toolset() -> ToolServerConfig { (&kigi::TodoWriteTool).into(), task_output_tool_config(), task_tool_config(), + agent_swarm_tool_config(), (&kigi::SchedulerCreateTool).into(), (&kigi::SchedulerDeleteTool).into(), (&kigi::SchedulerListTool).into(), @@ -428,6 +436,7 @@ fn orchestrator_toolset() -> ToolServerConfig { (&kigi::ListDirTool).into(), (&kigi::GrepTool).into(), task_tool_config(), + agent_swarm_tool_config(), task_output_tool_config(), wait_tasks_tool_config(), kill_task_tool_config(), @@ -497,6 +506,7 @@ fn kigi_ask_user_toolset() -> ToolServerConfig { task_output_tool_config(), wait_tasks_tool_config(), task_tool_config(), + agent_swarm_tool_config(), (&kigi::SchedulerCreateTool).into(), (&kigi::SchedulerDeleteTool).into(), (&kigi::SchedulerListTool).into(), diff --git a/crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs b/crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs index 246997f..e82b4db 100644 --- a/crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs +++ b/crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs @@ -396,11 +396,14 @@ pub(crate) async fn handle_subagent_request( let child_depth = ctx.parent_depth + 1; if child_depth >= MAX_SUBAGENT_DEPTH { let before = definition.tool_config.tools.len(); - definition.tool_config.tools.retain(|tc| tc.kind != Some(ToolKind::Task)); + definition + .tool_config + .tools + .retain(|tc| !matches!(tc.kind, Some(ToolKind::Task | ToolKind::AgentSwarm))); if definition.tool_config.tools.len() < before { tracing::info!( subagent_id = % request.id, child_depth, max_depth = - MAX_SUBAGENT_DEPTH, "Stripped task tool from child at max depth" + MAX_SUBAGENT_DEPTH, "Stripped subagent-spawning tools from child at max depth" ); } prune_orphaned_background_task_tools(&mut definition.tool_config); @@ -1322,6 +1325,7 @@ pub(crate) async fn handle_subagent_request( SubagentResult { success: false, cancelled: true, + rate_limited: false, error: Some(reason), output: if final_text.is_empty() { std::sync::Arc::from( @@ -1359,6 +1363,7 @@ pub(crate) async fn handle_subagent_request( SubagentResult { success: false, cancelled: true, + rate_limited: false, error: Some(format!("max turns reached (limit: {limit})")), output: if final_text.is_empty() { std::sync::Arc::from( @@ -1413,6 +1418,9 @@ pub(crate) async fn handle_subagent_request( SubagentResult { success: false, cancelled: was_cancelled, + rate_limited: !was_cancelled + && i32::from(e.code) + == crate::sampling::error::RATE_LIMITED_ERROR_CODE, error: Some( if was_cancelled { "Subagent was cancelled".to_string() diff --git a/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs b/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs index 4a34834..b291d8e 100644 --- a/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs +++ b/crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs @@ -1546,6 +1546,18 @@ impl SessionActor { vec![], vec![], ), + // Without an explicit arm the catch-all below titles this "Tool + // call" — for an entry that can hold the turn for the whole swarm. + ToolInput::AgentSwarm(swarm) => ( + format!( + "{} ({} members)", + swarm.description, + swarm.items.len() + swarm.resume_agent_ids.len() + ), + acp::ToolKind::Other, + vec![], + vec![], + ), ToolInput::EnterPlanMode(_) => ( "Plan: Enter".to_string(), acp::ToolKind::Other, diff --git a/crates/codegen/kigi-tools/schema/tool_meta.schema.json b/crates/codegen/kigi-tools/schema/tool_meta.schema.json index f6c9e16..44b538a 100644 --- a/crates/codegen/kigi-tools/schema/tool_meta.schema.json +++ b/crates/codegen/kigi-tools/schema/tool_meta.schema.json @@ -36,7 +36,7 @@ ], "definitions": { "ToolKind": { - "description": "Categorizes what a tool does at a high level. Open set — consumers must tolerate unknown values (Rust deserializes them to `other` via `#[serde(other)]`). Known values: `read`, `edit`, `delete`, `list_dir`, `write`, `move`, `search`, `lsp`, `execute`, `plan`, `web_search`, `web_fetch`, `background_task_action`, `wait_tasks_action`, `kill_task_action`, `list`, `skill`, `memory_search`, `memory_get`, `task`, `enter_plan`, `exit_plan`, `ask_user`, `deploy_app`, `search_tool`, `use_tool`, `monitor`, `goal_update`, `other`.", + "description": "Categorizes what a tool does at a high level. Open set — consumers must tolerate unknown values (Rust deserializes them to `other` via `#[serde(other)]`). Known values: `read`, `edit`, `delete`, `list_dir`, `write`, `move`, `search`, `lsp`, `execute`, `plan`, `web_search`, `web_fetch`, `background_task_action`, `wait_tasks_action`, `kill_task_action`, `list`, `skill`, `memory_search`, `memory_get`, `task`, `agent_swarm`, `enter_plan`, `exit_plan`, `ask_user`, `deploy_app`, `search_tool`, `use_tool`, `monitor`, `goal_update`, `other`.", "type": "string" }, "ToolNamespace": { diff --git a/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/mod.rs b/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/mod.rs new file mode 100644 index 0000000..e08d013 --- /dev/null +++ b/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/mod.rs @@ -0,0 +1,7 @@ +//! `agent_swarm` tool — one prompt template over many items, run as a fleet. +pub mod plan; +pub mod run; +pub mod schedule; +pub mod tool; + +pub use tool::AgentSwarmTool; diff --git a/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/plan.rs b/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/plan.rs new file mode 100644 index 0000000..26888f5 --- /dev/null +++ b/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/plan.rs @@ -0,0 +1,377 @@ +//! Turning `agent_swarm` input into member specs, and member results back +//! into one tool result. Pure: no spawning, no I/O. + +use kigi_tool_types::{ + AgentSwarmToolInput, MAX_AGENT_SWARM_MEMBERS, PROMPT_TEMPLATE_PLACEHOLDER, SwarmMemberOutcome, + SwarmMemberResult, +}; + +/// One member's launch instructions. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MemberSpec { + /// The `items` entry, or the agent id for a resume — the label the caller + /// sees in the result and reuses to retry this exact member. + pub item: String, + pub prompt: String, + /// Set for a resume; `None` spawns a fresh member. + pub resume_from: Option, +} + +/// Everything wrong with the input is reported before anything is spawned: +/// a half-launched swarm is far more expensive to recover from than a refused +/// tool call. +pub fn plan_members(input: &AgentSwarmToolInput) -> Result, String> { + let resumes = input.resume_agent_ids.len(); + + if input.items.is_empty() && resumes == 0 { + return Err( + "agent_swarm needs `items` (with `prompt_template`) or `resume_agent_ids`; \ + use the task tool for a single subagent" + .to_string(), + ); + } + if resumes == 0 && input.items.len() < 2 { + return Err( + "agent_swarm runs 2 or more members; use the task tool for a single subagent" + .to_string(), + ); + } + if input.items.len() + resumes > MAX_AGENT_SWARM_MEMBERS { + return Err(format!( + "agent_swarm runs at most {MAX_AGENT_SWARM_MEMBERS} members, got {}", + input.items.len() + resumes + )); + } + + // Models emit `""`/`"null"`/`"none"` where they mean "no id"; accepting one + // spawns a resume that can only die inside the coordinator. + if let Some(bad) = input + .resume_agent_ids + .keys() + .find(|id| !crate::implementations::kigi::task::types::is_valid_resume_id(id)) + { + return Err(format!( + "`resume_agent_ids` key {bad:?} is not a subagent id; use the agent_id \ + values from a previous agent_swarm result" + )); + } + + // Resumes first: they already hold context, so they reach the provider + // before the fresh members compete for the same rate-limit budget. + let mut specs: Vec = input + .resume_agent_ids + .iter() + .map(|(agent_id, prompt)| MemberSpec { + item: agent_id.clone(), + prompt: prompt.clone(), + resume_from: Some(agent_id.clone()), + }) + .collect(); + + if !input.items.is_empty() { + let template = input + .prompt_template + .as_deref() + .map(str::trim) + .filter(|t| !t.is_empty()) + .ok_or("agent_swarm requires `prompt_template` whenever `items` is given")?; + if !template.contains(PROMPT_TEMPLATE_PLACEHOLDER) { + return Err(format!( + "`prompt_template` must contain the literal {PROMPT_TEMPLATE_PLACEHOLDER}, \ + which is replaced by each entry of `items`" + )); + } + for item in &input.items { + if item.trim().is_empty() { + return Err("`items` entries must be non-empty".to_string()); + } + specs.push(MemberSpec { + item: item.clone(), + prompt: template.replace(PROMPT_TEMPLATE_PLACEHOLDER, item), + resume_from: None, + }); + } + } + + // Two members handed the same prompt do the same work twice and, if it + // writes, race each other over the same files. + let mut seen = std::collections::HashSet::with_capacity(specs.len()); + for spec in &specs { + if !seen.insert(spec.prompt.as_str()) { + return Err(format!( + "two members would run an identical prompt (from item {:?}); \ + give every member a distinct scope", + spec.item + )); + } + } + + Ok(specs) +} + +/// Per-member ceiling on rendered summary text. +/// +/// 128 members' full outputs concatenated can exceed the context window at +/// exactly the moment the results matter. Native tool output is not truncated +/// anywhere downstream — only the MCP dispatcher does that — so the budget has +/// to live here. Each member keeps its head and its tail: the head says what it +/// did, the tail usually holds the verdict. +const MEMBER_SUMMARY_BUDGET: usize = 4_000; + +/// Trims to [`MEMBER_SUMMARY_BUDGET`] on a char boundary, keeping both ends and +/// saying plainly how much was dropped. +fn clamp_summary(summary: &str) -> String { + if summary.len() <= MEMBER_SUMMARY_BUDGET { + return summary.to_string(); + } + let keep = MEMBER_SUMMARY_BUDGET / 2; + let head_end = (0..=keep) + .rev() + .find(|i| summary.is_char_boundary(*i)) + .unwrap_or(0); + let tail_start = (summary.len() - keep..summary.len()) + .find(|i| summary.is_char_boundary(*i)) + .unwrap_or(summary.len()); + format!( + "{}\n… {} bytes omitted; read this member's full output with the task-output tool …\n{}", + &summary[..head_end], + summary.len() - head_end - (summary.len() - tail_start), + &summary[tail_start..] + ) +} + +/// Renders the fleet's results as one tool result. +/// +/// A failed member is reported inside the aggregate rather than failing the +/// call: the caller needs the members that DID succeed, and needs to know +/// precisely which ones to retry. +pub fn render_results(results: &[SwarmMemberResult]) -> String { + let count = |wanted: SwarmMemberOutcome| results.iter().filter(|r| r.outcome == wanted).count(); + let completed = count(SwarmMemberOutcome::Completed); + let failed = count(SwarmMemberOutcome::Failed); + let aborted = count(SwarmMemberOutcome::Aborted); + let backgrounded = count(SwarmMemberOutcome::Backgrounded); + + let mut out = String::from("\n"); + out.push_str(&format!( + "completed: {completed}, failed: {failed}, aborted: {aborted}, \ + still running: {backgrounded}\n" + )); + if backgrounded > 0 { + out.push_str( + "Some members outlived the foreground budget and are still \ + working. Do NOT re-launch their items — a second agent on the same files \ + corrupts both. Poll them with the task-output tool instead.\n", + ); + } + + if results.iter().any(SwarmMemberResult::is_resumable) { + out.push_str( + "Call agent_swarm again with resume_agent_ids mapping the agent_id \ + values below to a follow-up prompt to continue unfinished work.\n", + ); + } + + for result in results { + out.push_str("\n"); + out.push_str(clamp_summary(result.summary.trim()).trim()); + out.push_str("\n\n"); + } + + out.push_str(""); + out +} + +/// Attribute-safe: an item is a model-supplied string and routinely contains +/// quotes or angle brackets (file paths, globs, shell fragments). +fn escape_attr(value: &str) -> String { + value + .replace('&', "&") + .replace('<', "<") + .replace('>', ">") + .replace('"', """) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn input(items: &[&str], template: Option<&str>) -> AgentSwarmToolInput { + AgentSwarmToolInput { + description: "test swarm".into(), + subagent_type: "general-purpose".into(), + prompt_template: template.map(str::to_string), + items: items.iter().map(|s| s.to_string()).collect(), + resume_agent_ids: Default::default(), + model: None, + } + } + + #[test] + fn each_item_becomes_one_member_with_the_placeholder_substituted() { + let specs = plan_members(&input(&["a.rs", "b.rs"], Some("Review {{item}} for bugs"))) + .expect("valid input"); + assert_eq!( + specs, + vec![ + MemberSpec { + item: "a.rs".into(), + prompt: "Review a.rs for bugs".into(), + resume_from: None, + }, + MemberSpec { + item: "b.rs".into(), + prompt: "Review b.rs for bugs".into(), + resume_from: None, + }, + ] + ); + } + + #[test] + fn a_single_item_is_refused_in_favour_of_the_task_tool() { + let err = plan_members(&input(&["only.rs"], Some("Do {{item}}"))).unwrap_err(); + assert!(err.contains("2 or more"), "{err}"); + } + + #[test] + fn a_template_without_the_placeholder_is_refused() { + let err = plan_members(&input(&["a", "b"], Some("Do the work"))).unwrap_err(); + assert!(err.contains(PROMPT_TEMPLATE_PLACEHOLDER), "{err}"); + } + + #[test] + fn items_without_a_template_are_refused() { + let err = plan_members(&input(&["a", "b"], None)).unwrap_err(); + assert!(err.contains("prompt_template"), "{err}"); + } + + #[test] + fn duplicate_expansions_are_refused_before_anything_spawns() { + let err = plan_members(&input(&["same", "same"], Some("Do {{item}}"))).unwrap_err(); + assert!(err.contains("identical prompt"), "{err}"); + } + + #[test] + fn more_members_than_the_ceiling_are_refused() { + let items: Vec = (0..=MAX_AGENT_SWARM_MEMBERS) + .map(|i| format!("item-{i}")) + .collect(); + let mut spec = input(&[], Some("Do {{item}}")); + spec.items = items; + let err = plan_members(&spec).unwrap_err(); + assert!(err.contains(&MAX_AGENT_SWARM_MEMBERS.to_string()), "{err}"); + } + + #[test] + fn resumes_are_planned_before_fresh_members() { + let mut spec = input(&["fresh.rs"], Some("Do {{item}}")); + spec.resume_agent_ids + .insert("agent-1".into(), "keep going".into()); + let specs = plan_members(&spec).expect("a resume lifts the two-item floor"); + assert_eq!(specs[0].resume_from.as_deref(), Some("agent-1")); + assert_eq!(specs[0].prompt, "keep going"); + assert_eq!(specs[1].item, "fresh.rs"); + } + + #[test] + fn an_empty_call_is_refused() { + let err = plan_members(&input(&[], None)).unwrap_err(); + assert!(err.contains("resume_agent_ids"), "{err}"); + } + + fn member(item: &str, outcome: SwarmMemberOutcome, id: Option<&str>) -> SwarmMemberResult { + SwarmMemberResult { + item: item.into(), + agent_id: id.map(str::to_string), + resumed: false, + outcome, + summary: format!("summary for {item}"), + } + } + + #[test] + fn the_aggregate_counts_outcomes_and_labels_every_member() { + let rendered = render_results(&[ + member("a.rs", SwarmMemberOutcome::Completed, Some("id-a")), + member("b.rs", SwarmMemberOutcome::Failed, Some("id-b")), + ]); + assert!(rendered.contains("completed: 1, failed: 1, aborted: 0, still running: 0")); + assert!( + rendered.contains(r#"agent_id="id-a" item="a.rs" state="started" outcome="completed""#) + ); + assert!(rendered.contains(r#"outcome="failed""#)); + assert!(rendered.contains("summary for b.rs")); + } + + #[test] + fn the_resume_hint_appears_only_when_something_resumable_is_unfinished() { + let all_done = render_results(&[member("a", SwarmMemberOutcome::Completed, Some("id-a"))]); + assert!(!all_done.contains("resume_hint")); + + let never_started = render_results(&[member("a", SwarmMemberOutcome::Aborted, None)]); + assert!( + !never_started.contains("resume_hint"), + "a member with no agent_id has nothing to resume" + ); + + let retryable = render_results(&[member("a", SwarmMemberOutcome::Failed, Some("id-a"))]); + assert!(retryable.contains("resume_hint")); + } + + #[test] + fn an_item_containing_markup_cannot_break_out_of_its_attribute() { + let rendered = render_results(&[member( + r#"a" onload="x"#, + SwarmMemberOutcome::Completed, + None, + )]); + assert!(!rendered.contains(r#"item="a" onload="#), "{rendered}"); + assert!(rendered.contains("""), "{rendered}"); + } +} + +#[cfg(test)] +mod budget_tests { + use super::*; + + #[test] + fn a_huge_member_output_is_clamped_but_keeps_both_ends() { + let summary = format!("HEAD{}TAIL", "x".repeat(MEMBER_SUMMARY_BUDGET * 2)); + let clamped = clamp_summary(&summary); + assert!(clamped.len() < summary.len(), "must shrink"); + assert!(clamped.starts_with("HEAD"), "the head says what it did"); + assert!(clamped.ends_with("TAIL"), "the tail holds the verdict"); + assert!(clamped.contains("bytes omitted"), "the loss must be stated"); + } + + #[test] + fn a_summary_inside_the_budget_is_untouched() { + assert_eq!(clamp_summary("short"), "short"); + } + + #[test] + fn clamping_never_splits_a_character() { + // Multi-byte throughout, so a naive byte slice would panic. + let summary = "café ".repeat(MEMBER_SUMMARY_BUDGET); + let clamped = clamp_summary(&summary); + assert!(clamped.len() < summary.len()); + } +} diff --git a/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/run.rs b/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/run.rs new file mode 100644 index 0000000..1c1ba73 --- /dev/null +++ b/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/run.rs @@ -0,0 +1,773 @@ +//! Drives planned members through the backend at the pace +//! [`super::schedule::LaunchPacer`] allows, and collects their results. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +// One clock for the whole runner: `tokio::time::Instant` is the clock the +// sleeps below advance. Mixing it with `std::time::Instant` makes the retry +// windows and the timer disagree — under a paused clock they never converge +// and the loop spins forever. +use tokio::time::Instant; + +use futures::stream::{FuturesUnordered, StreamExt}; +use kigi_tool_types::{SwarmMemberOutcome, SwarmMemberResult}; + +use super::plan::MemberSpec; +use super::schedule::{ + LAUNCH_INTERVAL, LaunchDecision, LaunchPacer, MAX_RATE_LIMIT_RETRIES, MAX_SWARM_RUNTIME, + retry_backoff, +}; +use crate::implementations::kigi::task::backend::SubagentBackend; +use crate::implementations::kigi::task::types::{ + ModelOverrideProvenance, SubagentRequest, SubagentResult, SubagentRuntimeOverrides, +}; + +/// Everything the runner needs that is not per-member. +pub struct SwarmRunConfig { + pub subagent_type: String, + pub description: String, + pub parent_session_id: String, + pub parent_prompt_id: Option, + pub model: Option, + pub cwd: Option, + pub max_concurrency: Option, +} + +/// A member waiting for its turn, plus how often the provider has refused it. +struct Pending { + index: usize, + spec: MemberSpec, + attempts: u32, + /// Earliest instant this member may be retried after a rate limit. + not_before: Option, +} + +struct Finished { + agent_id: Option, + result: SubagentResult, +} + +/// Cancels members that are still running if the swarm's future goes away. +/// +/// A send-now interrupt cancels the turn WITHOUT cancelling subagents, then +/// aborts the turn task — which drops this future and closes every member's +/// result channel. The coordinator reads a closed channel as "parent gone" and +/// re-attaches the child as a background task, so without this guard an +/// interrupt silently leaves up to `MAX_AGENT_SWARM_MEMBERS` agents editing the +/// caller's tree with no way to list them. `Drop` cannot await, so the cancels +/// are handed to the runtime; if there is no runtime left to hand them to, +/// nothing can be done and the ids are logged instead of lost silently. +struct InFlightGuard { + backend: Arc, + live: Arc>>, +} + +impl InFlightGuard { + fn track(&self, id: &str) { + self.live + .lock() + .expect("not poisoned") + .insert(id.to_string()); + } + + fn release(&self, id: &str) { + self.live.lock().expect("not poisoned").remove(id); + } +} + +impl Drop for InFlightGuard { + fn drop(&mut self) { + let ids: Vec = self + .live + .lock() + .map(|live| live.iter().cloned().collect()) + .unwrap_or_default(); + if ids.is_empty() { + return; + } + let backend = self.backend.clone(); + match tokio::runtime::Handle::try_current() { + Ok(handle) => { + handle.spawn(async move { + for id in ids { + backend.cancel(&id).await; + } + }); + } + Err(_) => tracing::warn!( + orphans = ?ids, + "agent_swarm dropped with no runtime to cancel its members on" + ), + } + } +} + +/// Runs every member to a terminal state and returns their results in the +/// order they were planned. +pub async fn run_swarm( + backend: Arc, + specs: Vec, + config: SwarmRunConfig, +) -> Vec { + let started_at = Instant::now(); + let total = specs.len(); + let labels: Vec = specs.iter().map(|s| s.item.clone()).collect(); + let resumed: Vec = specs.iter().map(|s| s.resume_from.is_some()).collect(); + + let mut queue: Vec = specs + .into_iter() + .enumerate() + .map(|(index, spec)| Pending { + index, + spec, + attempts: 0, + not_before: None, + }) + .collect(); + queue.reverse(); // pop() takes the earliest-planned member first + + let mut pacer = LaunchPacer::new(total, config.max_concurrency); + let mut in_flight = FuturesUnordered::new(); + let mut done: HashMap = HashMap::with_capacity(total); + // The id each member was last launched under, so one abandoned at the wall + // clock is still reportable as a live agent rather than an anonymous gap. + let mut launched: Vec> = vec![None; total]; + let guard = InFlightGuard { + backend: backend.clone(), + live: Arc::new(std::sync::Mutex::new(std::collections::BTreeSet::new())), + }; + + while done.len() < total { + let now = started_at.elapsed(); + if now >= MAX_SWARM_RUNTIME { + tracing::warn!( + elapsed_s = now.as_secs(), + unfinished = total - done.len(), + "agent_swarm hit its wall clock; reporting unfinished members" + ); + break; + } + let decision = if queue.is_empty() { + LaunchDecision::Drained + } else { + pacer.poll(now) + }; + + match decision { + LaunchDecision::Launch => { + // A member cooling off after a rate limit is not eligible yet; + // rotate it behind one that is rather than idling the fleet. + let Some(pending) = take_ready(&mut queue) else { + // Every queued member is cooling off. Whichever comes first + // — a slot freeing or the soonest retry opening — is the + // event worth waking for; awaiting only the in-flight side + // would park a member with a 3s backoff behind one with ten + // minutes left to run. + let wait = soonest_retry(&queue); + if in_flight.is_empty() { + match wait { + Some(wait) => tokio::time::sleep(wait).await, + // Unreachable: a member that is not ready has a + // `not_before`. Break rather than spin if it ever is. + None => break, + } + continue; + } + tokio::select! { + () = tokio::time::sleep(wait.unwrap_or(LAUNCH_INTERVAL)) => {} + Some(item) = in_flight.next() => { + settle(item, &guard, &mut pacer, &mut queue, &mut done, started_at); + } + } + continue; + }; + let request = build_request(&pending, &config); + let agent_id = request.id.clone(); + let index = pending.index; + let attempts = pending.attempts; + let spec = pending.spec; + let backend = backend.clone(); + guard.track(&agent_id); + launched[index] = Some(agent_id.clone()); + pacer.on_launched(now); + in_flight.push(async move { + let result = backend.spawn(request).await; + (index, agent_id, attempts, spec, result) + }); + } + LaunchDecision::Wait(delay) => { + if in_flight.is_empty() { + tokio::time::sleep(delay).await; + } else { + // A finishing member frees a slot sooner than the timer. + tokio::select! { + () = tokio::time::sleep(delay) => {} + Some(item) = in_flight.next() => { + settle(item, &guard, &mut pacer, &mut queue, &mut done, started_at); + } + } + } + } + LaunchDecision::Drained => { + if in_flight.is_empty() { + break; + } + collect_one( + &mut in_flight, + &guard, + &mut pacer, + &mut queue, + &mut done, + started_at, + ) + .await; + } + } + } + + (0..total) + .map(|index| match done.remove(&index) { + Some(finished) => to_member_result( + labels[index].clone(), + resumed[index], + finished.agent_id, + finished.result, + ), + // Reached when the swarm hit its wall clock: the member is alive + // and unaccounted for, which is exactly `Backgrounded` — never + // offer it for resume, and never invite a relaunch of its item. + None => SwarmMemberResult { + item: labels[index].clone(), + agent_id: launched[index].clone(), + resumed: resumed[index], + outcome: SwarmMemberOutcome::Backgrounded, + summary: "Still running when the swarm reached its time limit.".to_string(), + }, + }) + .collect() +} + +type InFlight = ( + usize, + String, + u32, + MemberSpec, + Result, +); + +async fn collect_one( + in_flight: &mut FuturesUnordered>, + guard: &InFlightGuard, + pacer: &mut LaunchPacer, + queue: &mut Vec, + done: &mut HashMap, + started_at: Instant, +) { + if let Some(item) = in_flight.next().await { + settle(item, guard, pacer, queue, done, started_at); + } +} + +/// Files one finished member: either terminal, or re-queued because the +/// provider — not the work — refused it. +fn settle( + (index, agent_id, attempts, spec, outcome): InFlight, + guard: &InFlightGuard, + pacer: &mut LaunchPacer, + queue: &mut Vec, + done: &mut HashMap, + started_at: Instant, +) { + guard.release(&agent_id); + let now = started_at.elapsed(); + let mut transport_failed = false; + let result = match outcome { + Ok(result) => result, + Err(err) => { + transport_failed = true; + SubagentResult { + success: false, + error: Some(err.to_string()), + ..Default::default() + } + } + }; + + if result.rate_limited && attempts < MAX_RATE_LIMIT_RETRIES { + pacer.on_rate_limited(now); + queue.push(Pending { + index, + spec, + attempts: attempts + 1, + not_before: Some(Instant::now() + retry_backoff(attempts)), + }); + return; + } + + pacer.on_finished(); + // A transport failure means no child was ever created, so there is no id + // to resume; prefer the coordinator's own id when it minted one. + let agent_id = match (transport_failed, result.subagent_id.as_str()) { + (true, _) => None, + (false, "") => Some(agent_id), + (false, minted) => Some(minted.to_string()), + }; + done.insert(index, Finished { agent_id, result }); +} + +/// The earliest-planned member whose retry window has opened. +fn take_ready(queue: &mut Vec) -> Option { + let now = Instant::now(); + let position = (0..queue.len()) + .rev() + .find(|&i| queue[i].not_before.is_none_or(|at| at <= now))?; + Some(queue.remove(position)) +} + +fn soonest_retry(queue: &[Pending]) -> Option { + let now = Instant::now(); + queue + .iter() + .filter_map(|p| p.not_before) + .map(|at| at.saturating_duration_since(now)) + .min() +} + +fn build_request(pending: &Pending, config: &SwarmRunConfig) -> SubagentRequest { + let (result_tx, _) = tokio::sync::oneshot::channel(); + let resume_from = pending.spec.resume_from.clone(); + SubagentRequest { + id: uuid::Uuid::now_v7().to_string(), + prompt: pending.spec.prompt.clone(), + description: config.description.clone(), + subagent_type: config.subagent_type.clone(), + parent_session_id: config.parent_session_id.clone(), + parent_prompt_id: config.parent_prompt_id.clone(), + cwd: config.cwd.clone(), + runtime_overrides: SubagentRuntimeOverrides { + // A resume inherits the source member's model, so an override here + // would be dropped by the coordinator anyway. + model: resume_from + .is_none() + .then(|| config.model.clone()) + .flatten(), + model_override_provenance: ModelOverrideProvenance::Tool, + reasoning_effort: None, + persona: None, + capability_mode: None, + // Members share the caller's tree: 128 worktrees is not viable, and + // the distinct-prompt rule is what keeps them off each other's files. + isolation: None, + harness_agent_type: None, + }, + resume_from, + // The swarm owns its members' lifetimes: it awaits every one of them + // before returning, so none may outlive the tool call. + run_in_background: false, + surface_completion: true, + fork_context: false, + result_tx, + } +} + +fn to_member_result( + item: String, + resumed: bool, + agent_id: Option, + result: SubagentResult, +) -> SwarmMemberResult { + let outcome = if result.backgrounded { + // Checked FIRST: the coordinator reports a detached member with + // `success: false` and no output, which is indistinguishable from a + // failure by every other field. + SwarmMemberOutcome::Backgrounded + } else if result.success { + SwarmMemberOutcome::Completed + } else if result.cancelled { + SwarmMemberOutcome::Aborted + } else { + SwarmMemberOutcome::Failed + }; + + let output = result.output.trim(); + let summary = match (&result.error, outcome) { + (_, SwarmMemberOutcome::Backgrounded) => format!( + "Still running in the background; its result is not part of this call.{}", + if output.is_empty() { + String::new() + } else { + format!("\nProgress so far:\n{output}") + } + ), + // Why it ended is the load-bearing half for anything that did not + // complete — "max turns reached" must not be swallowed by whatever + // text the member happened to emit last. + (Some(error), SwarmMemberOutcome::Failed | SwarmMemberOutcome::Aborted) => { + if output.is_empty() { + error.clone() + } else { + format!("{error}\n{output}") + } + } + _ if output.is_empty() => "Member produced no output.".to_string(), + _ => output.to_string(), + }; + + SwarmMemberResult { + item, + agent_id, + resumed, + outcome, + summary, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::implementations::kigi::task::backend::SubagentBackend; + use crate::implementations::kigi::task::types::{ + SubagentCancelOutcome, SubagentDescribeOutcome, SubagentSnapshot, + SubagentValidateTypeOutcome, + }; + use std::sync::Mutex; + use std::sync::atomic::{AtomicUsize, Ordering}; + + /// Records every prompt it is asked to run, and answers from a script + /// keyed by prompt so a member can be refused once and then succeed. + #[derive(Default)] + struct FakeBackend { + seen: Mutex>, + /// prompt -> remaining rate-limit rejections before it succeeds. + refuse: Mutex>, + /// prompts that always fail outright. + fail: Mutex>, + /// prompts the coordinator detaches instead of finishing. + background: Mutex>, + /// prompt -> how long it occupies its slot, so members can overlap. + duration_ms: Mutex>, + cancelled: Mutex>, + peak_in_flight: AtomicUsize, + in_flight: AtomicUsize, + } + + impl FakeBackend { + fn prompts(&self) -> Vec { + self.seen.lock().unwrap().clone() + } + } + + #[async_trait::async_trait] + impl SubagentBackend for FakeBackend { + async fn spawn( + &self, + request: SubagentRequest, + ) -> Result { + let live = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; + self.peak_in_flight.fetch_max(live, Ordering::SeqCst); + self.seen.lock().unwrap().push(request.prompt.clone()); + + let refuse_now = { + let mut refuse = self.refuse.lock().unwrap(); + match refuse.get_mut(&request.prompt) { + Some(remaining) if *remaining > 0 => { + *remaining -= 1; + true + } + _ => false, + } + }; + let fails = self.fail.lock().unwrap().contains(&request.prompt); + let backgrounds = self.background.lock().unwrap().contains(&request.prompt); + let hold = self + .duration_ms + .lock() + .unwrap() + .get(&request.prompt) + .copied() + .unwrap_or(0); + // Yield for a real (virtual) interval so `FuturesUnordered` can + // actually interleave: without an await point every member runs to + // completion inside its own poll and nothing overlaps. + tokio::time::sleep(Duration::from_millis(hold.max(1))).await; + self.in_flight.fetch_sub(1, Ordering::SeqCst); + + if backgrounds { + return Ok(SubagentResult { + success: false, + backgrounded: true, + subagent_id: request.id.clone(), + child_session_id: "child".into(), + ..Default::default() + }); + } + + if refuse_now { + return Ok(SubagentResult { + success: false, + rate_limited: true, + child_session_id: "child".into(), + error: Some("Session error: Rate limited".into()), + ..Default::default() + }); + } + if fails { + return Ok(SubagentResult { + success: false, + error: Some("it broke".into()), + ..Default::default() + }); + } + Ok(SubagentResult { + success: true, + output: Arc::from(format!("done: {}", request.prompt)), + subagent_id: request.id.clone(), + child_session_id: "child".into(), + ..Default::default() + }) + } + + async fn query(&self, _: &str, _: bool, _: Option) -> Option { + None + } + async fn cancel(&self, id: &str) -> SubagentCancelOutcome { + self.cancelled.lock().unwrap().push(id.to_string()); + SubagentCancelOutcome::NotFound + } + async fn validate_type(&self, _: &str, _: &str) -> SubagentValidateTypeOutcome { + SubagentValidateTypeOutcome::Ok + } + async fn describe_subagent_type( + &self, + _: &str, + _: Option<&str>, + _: &str, + ) -> SubagentDescribeOutcome { + SubagentDescribeOutcome::Unavailable + } + } + + fn specs(items: &[&str]) -> Vec { + items + .iter() + .map(|item| MemberSpec { + item: (*item).to_string(), + prompt: format!("work on {item}"), + resume_from: None, + }) + .collect() + } + + fn config(max_concurrency: Option) -> SwarmRunConfig { + SwarmRunConfig { + subagent_type: "general-purpose".into(), + description: "test swarm".into(), + parent_session_id: "parent".into(), + parent_prompt_id: None, + model: None, + cwd: None, + max_concurrency, + } + } + + #[tokio::test(start_paused = true)] + async fn every_member_runs_and_results_come_back_in_plan_order() { + let backend = Arc::new(FakeBackend::default()); + // Finish in the exact reverse of plan order, so a runner that reported + // completion order instead would fail this. + for (item, ms) in [("a", 300), ("b", 200), ("c", 100)] { + backend + .duration_ms + .lock() + .unwrap() + .insert(format!("work on {item}"), ms); + } + let results = run_swarm(backend.clone(), specs(&["a", "b", "c"]), config(None)).await; + + assert_eq!( + results.iter().map(|r| r.item.as_str()).collect::>(), + vec!["a", "b", "c"], + "results must be ordered as planned, not by completion" + ); + assert!( + results + .iter() + .all(|r| r.outcome == SwarmMemberOutcome::Completed) + ); + assert_eq!(backend.prompts().len(), 3); + } + + #[tokio::test(start_paused = true)] + async fn a_failing_member_does_not_sink_the_others() { + let backend = Arc::new(FakeBackend::default()); + backend.fail.lock().unwrap().push("work on b".into()); + + let results = run_swarm(backend, specs(&["a", "b", "c"]), config(None)).await; + assert_eq!(results[0].outcome, SwarmMemberOutcome::Completed); + assert_eq!(results[1].outcome, SwarmMemberOutcome::Failed); + assert_eq!(results[2].outcome, SwarmMemberOutcome::Completed); + assert!(results[1].summary.contains("it broke")); + } + + #[tokio::test(start_paused = true)] + async fn a_rate_limited_member_is_retried_not_discarded() { + let backend = Arc::new(FakeBackend::default()); + backend.refuse.lock().unwrap().insert("work on b".into(), 1); + + let results = run_swarm(backend.clone(), specs(&["a", "b"]), config(None)).await; + assert!( + results + .iter() + .all(|r| r.outcome == SwarmMemberOutcome::Completed), + "the refused member must succeed on its retry: {results:?}" + ); + assert_eq!( + backend + .prompts() + .iter() + .filter(|p| *p == "work on b") + .count(), + 2, + "the refused member must be attempted exactly twice" + ); + } + + #[tokio::test(start_paused = true)] + async fn the_operator_cap_bounds_concurrency() { + let hold_all = || { + let backend = Arc::new(FakeBackend::default()); + for item in ["a", "b", "c", "d"] { + backend + .duration_ms + .lock() + .unwrap() + .insert(format!("work on {item}"), 100); + } + backend + }; + + let capped = hold_all(); + run_swarm( + capped.clone(), + specs(&["a", "b", "c", "d"]), + config(Some(2)), + ) + .await; + assert!( + capped.peak_in_flight.load(Ordering::SeqCst) <= 2, + "cap of 2 exceeded: peak was {}", + capped.peak_in_flight.load(Ordering::SeqCst) + ); + + let uncapped = hold_all(); + run_swarm(uncapped.clone(), specs(&["a", "b", "c", "d"]), config(None)).await; + assert!( + uncapped.peak_in_flight.load(Ordering::SeqCst) > 2, + "the fixture must be able to exceed the cap, or the assertion above proves nothing" + ); + } + + /// A member the coordinator detached is still running: reporting it as a + /// failure invites the model to relaunch its item, putting a second agent + /// on the same files. + #[tokio::test(start_paused = true)] + async fn a_backgrounded_member_is_not_reported_as_failed_or_offered_for_resume() { + let backend = Arc::new(FakeBackend::default()); + backend.background.lock().unwrap().push("work on b".into()); + + let results = run_swarm(backend, specs(&["a", "b"]), config(None)).await; + assert_eq!(results[1].outcome, SwarmMemberOutcome::Backgrounded); + assert!( + !results[1].is_resumable(), + "a live member must never be offered for resume" + ); + assert!( + results[1].summary.contains("Still running"), + "{}", + results[1].summary + ); + } + + /// Dropping the runner is what a send-now interrupt does; the members must + /// be cancelled rather than silently detached onto the user's tree. + #[tokio::test(start_paused = true)] + async fn dropping_the_swarm_cancels_its_live_members() { + let backend = Arc::new(FakeBackend::default()); + for item in ["a", "b"] { + backend + .duration_ms + .lock() + .unwrap() + .insert(format!("work on {item}"), 10_000); + } + + tokio::select! { + _ = run_swarm(backend.clone(), specs(&["a", "b"]), config(None)) => { + panic!("members hold their slots for 10s; the swarm cannot finish first") + } + () = tokio::time::sleep(Duration::from_millis(50)) => {} + } + // The guard hands its cancels to the runtime; let them run. + tokio::task::yield_now().await; + tokio::time::sleep(Duration::from_millis(10)).await; + + assert!( + !backend.cancelled.lock().unwrap().is_empty(), + "a dropped swarm must cancel the members it started" + ); + } + + #[tokio::test(start_paused = true)] + async fn a_permanently_refused_member_fails_instead_of_hanging_the_turn() { + let backend = Arc::new(FakeBackend::default()); + // Refuses far more often than the retry bound allows. + backend + .refuse + .lock() + .unwrap() + .insert("work on a".into(), 100); + + let results = run_swarm(backend.clone(), specs(&["a", "b"]), config(None)).await; + assert_eq!( + results[0].outcome, + SwarmMemberOutcome::Failed, + "the swarm blocks the caller's turn, so retries must be bounded" + ); + assert_eq!( + results[1].outcome, + SwarmMemberOutcome::Completed, + "one exhausted member must not sink its siblings" + ); + assert_eq!( + backend + .prompts() + .iter() + .filter(|p| *p == "work on a") + .count() as u32, + MAX_RATE_LIMIT_RETRIES + 1, + "the first attempt plus exactly the retry budget" + ); + } + + #[tokio::test(start_paused = true)] + async fn a_resume_member_carries_its_source_id() { + let backend = Arc::new(FakeBackend::default()); + let specs = vec![ + MemberSpec { + item: "agent-7".into(), + prompt: "keep going".into(), + resume_from: Some("agent-7".into()), + }, + MemberSpec { + item: "fresh".into(), + prompt: "start here".into(), + resume_from: None, + }, + ]; + let results = run_swarm(backend, specs, config(None)).await; + assert!(results[0].resumed); + assert!(!results[1].resumed); + } +} diff --git a/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/schedule.rs b/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/schedule.rs new file mode 100644 index 0000000..a2c7738 --- /dev/null +++ b/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/schedule.rs @@ -0,0 +1,348 @@ +//! Launch pacing for a swarm, as pure state: no I/O, no clock, no tasks. +//! +//! A fan-out of N members hits ONE provider at once, so the launch order is +//! the difference between a swarm that runs and a swarm that 429s itself to +//! death. The runner asks [`LaunchPacer`] when it may start the next member +//! and reports rate limits back; everything here is decided arithmetically so +//! the policy is unit-testable without spawning an agent. + +use std::time::Duration; + +/// Members allowed to start with no wait at all. +pub const INITIAL_LAUNCH_BURST: usize = 5; + +/// Spacing between launches once the burst is spent. +pub const LAUNCH_INTERVAL: Duration = Duration::from_millis(700); + +/// First wait before re-attempting a rate-limited member. +pub const RETRY_MIN_BACKOFF: Duration = Duration::from_secs(3); + +/// Cap on a single member's retry wait; a provider window outlasting this is +/// better spent letting other members through than sleeping longer. +pub const RETRY_MAX_BACKOFF: Duration = Duration::from_secs(120); + +/// Quiet period after which the fleet regains one lost capacity slot. +pub const CAPACITY_RECOVERY: Duration = Duration::from_secs(180); + +/// Shortest gap between two capacity reductions, so one provider window that +/// rejects several members in a burst costs one slot, not all of them. +pub const CAPACITY_SHRINK_COOLDOWN: Duration = Duration::from_secs(2); + +/// Env override for the concurrency ceiling; unset means the ramp is the only +/// brake. A value that does not parse as a positive integer is a hard error: +/// silently ignoring it would run an unbounded fan-out the operator forbade. +pub const MAX_CONCURRENCY_ENV: &str = "KIGI_AGENT_SWARM_MAX_CONCURRENCY"; + +/// Resolves [`MAX_CONCURRENCY_ENV`], failing loudly on a malformed value. +pub fn max_concurrency_from_env(raw: Option<&str>) -> Result, String> { + let Some(raw) = raw.map(str::trim).filter(|v| !v.is_empty()) else { + return Ok(None); + }; + match raw.parse::() { + Ok(0) | Err(_) => Err(format!( + "{MAX_CONCURRENCY_ENV} must be a positive integer, got {raw:?}" + )), + Ok(value) => Ok(Some(value)), + } +} + +/// What the runner should do right now. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LaunchDecision { + /// Start the next queued member immediately. + Launch, + /// Nothing may start yet; wait this long and ask again. + Wait(Duration), + /// Every member has been started. + Drained, +} + +/// The launch ramp plus the fleet's rate-limit response. +/// +/// Time is injected as a monotonic `now` so tests drive it directly. +#[derive(Debug)] +pub struct LaunchPacer { + queued: usize, + in_flight: usize, + started: usize, + interval: Duration, + hard_cap: Option, + /// `None` until the first rate limit: before that the ramp alone paces us. + capacity: Option, + last_launch: Option, + last_shrink: Option, + last_rate_limit: Option, +} + +impl LaunchPacer { + pub fn new(queued: usize, hard_cap: Option) -> Self { + Self { + queued, + in_flight: 0, + started: 0, + interval: LAUNCH_INTERVAL, + hard_cap, + capacity: None, + last_launch: None, + last_shrink: None, + last_rate_limit: None, + } + } + + /// The ceiling in force now: the operator's cap and the rate-limit-derived + /// capacity both apply, whichever is lower. + fn ceiling(&self) -> Option { + match (self.hard_cap, self.capacity) { + (Some(a), Some(b)) => Some(a.min(b)), + (only, None) | (None, only) => only, + } + } + + pub fn poll(&mut self, now: Duration) -> LaunchDecision { + self.recover_capacity(now); + + if self.queued == 0 { + return LaunchDecision::Drained; + } + if let Some(ceiling) = self.ceiling() + && self.in_flight >= ceiling + { + // Held by capacity, not by the clock: only a member finishing (or + // the recovery timer) can release this, so poll on the recovery + // grain rather than spinning. + return LaunchDecision::Wait(self.recovery_wait(now)); + } + if self.started < INITIAL_LAUNCH_BURST { + return LaunchDecision::Launch; + } + match self.last_launch { + Some(last) if now.saturating_sub(last) < self.interval => { + LaunchDecision::Wait(self.interval - now.saturating_sub(last)) + } + _ => LaunchDecision::Launch, + } + } + + /// Record that the runner acted on a [`LaunchDecision::Launch`]. + pub fn on_launched(&mut self, now: Duration) { + self.queued = self.queued.saturating_sub(1); + self.in_flight += 1; + self.started += 1; + self.last_launch = Some(now); + } + + /// Record that a member reached a terminal state. + pub fn on_finished(&mut self) { + self.in_flight = self.in_flight.saturating_sub(1); + } + + /// Record that a member was rejected for rate limiting and re-queued. + /// + /// The fleet loses one slot (never below one) at most once per + /// [`CAPACITY_SHRINK_COOLDOWN`]. + /// + /// Upstream also doubles the launch interval when the rejected member never + /// reached the provider. Kigi cannot observe that: every rejection arrives + /// through a child session that DID start, so the branch would be dead code + /// — and it ratchets one way, with no path back from a two-minute interval. + pub fn on_rate_limited(&mut self, now: Duration) { + self.queued += 1; + self.in_flight = self.in_flight.saturating_sub(1); + self.last_rate_limit = Some(now); + + let cooling = self + .last_shrink + .is_some_and(|last| now.saturating_sub(last) < CAPACITY_SHRINK_COOLDOWN); + if !cooling { + let current = self.capacity.unwrap_or(self.in_flight.max(1)); + self.capacity = Some(current.saturating_sub(1).max(1)); + self.last_shrink = Some(now); + } + } + + /// One slot back per quiet [`CAPACITY_RECOVERY`] window, until the cap is + /// no longer the binding constraint. + fn recover_capacity(&mut self, now: Duration) { + let (Some(capacity), Some(last)) = (self.capacity, self.last_rate_limit) else { + return; + }; + if now.saturating_sub(last) < CAPACITY_RECOVERY { + return; + } + self.last_rate_limit = Some(now); + self.capacity = Some(capacity + 1); + } + + fn recovery_wait(&self, now: Duration) -> Duration { + let elapsed = self + .last_rate_limit + .map(|last| now.saturating_sub(last)) + .unwrap_or_default(); + CAPACITY_RECOVERY.saturating_sub(elapsed).max(self.interval) + } +} + +/// Ceiling on the whole swarm's wall clock. +/// +/// Per-member bounds do not bound the fleet: a member may hold its slot for the +/// full foreground budget and then be re-queued, so a large swarm can otherwise +/// hold the caller's turn for hours. At the deadline the runner stops launching +/// and reports whatever has not finished, with ids, rather than waiting on. +pub const MAX_SWARM_RUNTIME: Duration = Duration::from_secs(30 * 60); + +/// Rejections a single member may absorb before the runner gives up on it. +/// +/// The swarm blocks the caller's turn, so every retry path needs a bound it +/// cannot argue its way past: a provider that refuses one member indefinitely +/// must surface as a failed member the caller can retry deliberately, never as +/// a turn that hangs. +pub const MAX_RATE_LIMIT_RETRIES: u32 = 5; + +/// Per-member exponential backoff, capped. `attempt` counts prior rejections. +pub fn retry_backoff(attempt: u32) -> Duration { + RETRY_MIN_BACKOFF + .saturating_mul(2u32.saturating_pow(attempt.min(6))) + .min(RETRY_MAX_BACKOFF) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ms(v: u64) -> Duration { + Duration::from_millis(v) + } + + #[test] + fn the_first_five_members_launch_without_waiting() { + let mut pacer = LaunchPacer::new(10, None); + for _ in 0..INITIAL_LAUNCH_BURST { + assert_eq!(pacer.poll(ms(0)), LaunchDecision::Launch); + pacer.on_launched(ms(0)); + } + assert_eq!( + pacer.poll(ms(0)), + LaunchDecision::Wait(LAUNCH_INTERVAL), + "the sixth member must wait out the ramp interval" + ); + } + + #[test] + fn the_ramp_admits_one_member_per_interval() { + let mut pacer = LaunchPacer::new(10, None); + for _ in 0..INITIAL_LAUNCH_BURST { + pacer.on_launched(ms(0)); + } + assert_eq!(pacer.poll(ms(699)), LaunchDecision::Wait(ms(1))); + assert_eq!(pacer.poll(ms(700)), LaunchDecision::Launch); + } + + #[test] + fn an_operator_cap_binds_before_the_ramp() { + let mut pacer = LaunchPacer::new(10, Some(2)); + pacer.on_launched(ms(0)); + pacer.on_launched(ms(0)); + assert!( + matches!(pacer.poll(ms(0)), LaunchDecision::Wait(_)), + "a cap of 2 must not admit a third member even inside the burst" + ); + pacer.on_finished(); + assert_eq!(pacer.poll(ms(0)), LaunchDecision::Launch); + } + + #[test] + fn a_rate_limit_requeues_the_member_and_costs_one_slot() { + let mut pacer = LaunchPacer::new(4, None); + for _ in 0..4 { + pacer.on_launched(ms(0)); + } + pacer.on_rate_limited(ms(1_000)); + assert_eq!(pacer.capacity, Some(2), "3 in flight, minus the lost slot"); + assert!( + matches!(pacer.poll(ms(1_000)), LaunchDecision::Wait(_)), + "3 in flight against a capacity of 2 must not admit the requeued member" + ); + } + + #[test] + fn a_burst_of_rejections_costs_one_slot_not_all_of_them() { + let mut pacer = LaunchPacer::new(6, None); + for _ in 0..6 { + pacer.on_launched(ms(0)); + } + pacer.on_rate_limited(ms(1_000)); + let after_first = pacer.capacity; + pacer.on_rate_limited(ms(1_500)); + assert_eq!( + pacer.capacity, after_first, + "a second rejection inside the cooldown must not shrink again" + ); + pacer.on_rate_limited(ms(4_000)); + assert_eq!( + pacer.capacity, + after_first.map(|c| c - 1), + "past the cooldown the fleet gives up another slot" + ); + } + + #[test] + fn capacity_never_reaches_zero() { + let mut pacer = LaunchPacer::new(3, None); + pacer.on_launched(ms(0)); + for i in 0..10 { + pacer.on_rate_limited(Duration::from_secs(10 * (i + 1))); + } + assert_eq!( + pacer.capacity, + Some(1), + "a fleet with no slots could never make progress" + ); + } + + #[test] + fn quiet_time_returns_a_lost_slot() { + let mut pacer = LaunchPacer::new(4, None); + for _ in 0..3 { + pacer.on_launched(ms(0)); + } + pacer.on_rate_limited(ms(1_000)); + let shrunk = pacer.capacity.expect("shrunk"); + pacer.poll(ms(1_000) + CAPACITY_RECOVERY); + assert_eq!(pacer.capacity, Some(shrunk + 1)); + } + + #[test] + fn the_retry_bound_is_reachable_within_the_backoff_cap() { + // The bound must terminate in bounded time, not merely be finite. + let worst: Duration = (0..MAX_RATE_LIMIT_RETRIES).map(retry_backoff).sum(); + assert!( + worst <= RETRY_MAX_BACKOFF * MAX_RATE_LIMIT_RETRIES, + "worst-case retry time {worst:?} must stay inside the per-attempt cap" + ); + } + + #[test] + fn draining_is_reported_once_every_member_has_started() { + let mut pacer = LaunchPacer::new(1, None); + pacer.on_launched(ms(0)); + assert_eq!(pacer.poll(ms(0)), LaunchDecision::Drained); + } + + #[test] + fn backoff_grows_then_stops_at_the_cap() { + assert_eq!(retry_backoff(0), RETRY_MIN_BACKOFF); + assert_eq!(retry_backoff(1), RETRY_MIN_BACKOFF * 2); + assert_eq!(retry_backoff(2), RETRY_MIN_BACKOFF * 4); + assert_eq!(retry_backoff(30), RETRY_MAX_BACKOFF); + } + + #[test] + fn a_malformed_concurrency_cap_is_refused_not_ignored() { + assert_eq!(max_concurrency_from_env(None), Ok(None)); + assert_eq!(max_concurrency_from_env(Some(" ")), Ok(None)); + assert_eq!(max_concurrency_from_env(Some("4")), Ok(Some(4))); + assert!(max_concurrency_from_env(Some("0")).is_err()); + assert!(max_concurrency_from_env(Some("many")).is_err()); + assert!(max_concurrency_from_env(Some("-1")).is_err()); + } +} diff --git a/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/tool.rs b/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/tool.rs new file mode 100644 index 0000000..f1f1f26 --- /dev/null +++ b/crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/tool.rs @@ -0,0 +1,498 @@ +//! The `agent_swarm` tool: validate, then run every member to completion. + +use kigi_tool_types::AgentSwarmToolInput; + +use super::plan::{plan_members, render_results}; +use super::run::{SwarmRunConfig, run_swarm}; +use super::schedule::{MAX_CONCURRENCY_ENV, max_concurrency_from_env}; +use crate::implementations::kigi::task::MAX_SUBAGENT_DEPTH; +use crate::implementations::kigi::task::backend::SubagentBackendResource; +use crate::implementations::kigi::task::types::{ + CurrentPromptIdResource, SessionIdResource, SubagentDepthCounter, SubagentValidateTypeOutcome, + TaskModelValidator, +}; +use crate::types::output::ToolOutput; +use crate::types::requirements::{Expr, ToolRequirement}; +use crate::types::tool::{ToolKind, ToolNamespace}; + +const DESCRIPTION: &str = "\ +Run one prompt over many independent work items at once, as a fleet of subagents. + +Give a `prompt_template` containing the literal {{item}} and an `items` list; each entry \ +becomes one subagent whose prompt is the template with {{item}} substituted. The call \ +returns only when every member has finished, with each member's result labelled by its item. + +Use this when the work splits into 2 or more INDEPENDENT scopes — separate files, separate \ +directories, separate questions. Every member must have a distinct scope: members share one \ +working tree with no isolation, so two members told to edit the same file will corrupt each \ +other's work. Read-only scopes may overlap freely. + +For a single item, use the subagent (task) tool instead. To continue members from an earlier \ +swarm, pass `resume_agent_ids` mapping the agent_id values from that swarm's result to a \ +follow-up prompt."; + +#[derive(Debug, Default)] +pub struct AgentSwarmTool; + +impl crate::types::tool_metadata::ToolMetadata for AgentSwarmTool { + fn kind(&self) -> ToolKind { + ToolKind::AgentSwarm + } + + fn tool_namespace(&self) -> ToolNamespace { + ToolNamespace::Kigi + } + + fn description_template(&self) -> &str { + DESCRIPTION + } + + fn requires_expr(&self) -> Expr { + // Members are subagents, so the same background-task management tools + // the `task` tool depends on must be present. + Expr::And(vec![ + Expr::Value(ToolRequirement::tool_kind(ToolKind::BackgroundTaskAction)), + Expr::Value(ToolRequirement::tool_kind(ToolKind::KillTaskAction)), + ]) + } + + fn is_read_only(&self) -> bool { + false + } +} + +impl kigi_tool_runtime::Tool for AgentSwarmTool { + type Args = AgentSwarmToolInput; + type Output = ToolOutput; + + fn id(&self) -> kigi_tool_protocol::ToolId { + kigi_tool_protocol::ToolId::new("agent_swarm").expect("valid tool id") + } + + fn description( + &self, + _ctx: &::kigi_tool_runtime::ListToolsContext, + ) -> kigi_tool_types::ToolDescription { + kigi_tool_types::ToolDescription::new("agent_swarm", DESCRIPTION) + } + + fn capabilities(&self) -> kigi_tool_protocol::ToolCapabilities { + kigi_tool_protocol::ToolCapabilities { + is_read_only: false, + tool_scope: Some(kigi_tool_protocol::ToolScope::Write), + ..Default::default() + } + } + + #[tracing::instrument( + name = "tool.agent_swarm", + skip_all, + fields( + subagent_type = %input.subagent_type, + members = input.items.len() + input.resume_agent_ids.len(), + ) + )] + async fn run( + &self, + ctx: kigi_tool_runtime::ToolCallContext, + input: AgentSwarmToolInput, + ) -> Result { + use crate::types::tool_metadata::shared_resources; + let resources = shared_resources(&ctx)?; + + let (depth, backend, model_validator, parent_session_id, parent_prompt_id) = { + let res = resources.lock().await; + let depth = res.get::().map(|d| d.0).unwrap_or(0); + let model_validator = res.get::().cloned(); + let backend = res + .get::() + .ok_or_else(|| { + kigi_tool_runtime::ToolError::custom( + "missing_resource", + "SubagentBackendResource (subagent support not initialized)", + ) + })? + .clone(); + let parent_session_id = res + .get::() + .map(|s| s.0.clone()) + .unwrap_or_default(); + let parent_prompt_id = res + .get::() + .map(|p| p.0.clone()) + .filter(|prompt_id| !prompt_id.is_empty()); + ( + depth, + backend, + model_validator, + parent_session_id, + parent_prompt_id, + ) + }; + + if depth >= MAX_SUBAGENT_DEPTH { + return Err(kigi_tool_runtime::ToolError::invalid_arguments(format!( + "Subagent depth limit exceeded (current depth: {depth}, max: {MAX_SUBAGENT_DEPTH}). \ + A subagent cannot start a swarm." + ))); + } + + let max_concurrency = + max_concurrency_from_env(std::env::var(MAX_CONCURRENCY_ENV).ok().as_deref()) + .map_err(kigi_tool_runtime::ToolError::invalid_arguments)?; + + // Every input fault is reported before a single member starts: a + // half-launched swarm costs real tokens to unwind, and one bad model + // slug would otherwise fan out into as many failures as there are items. + let model = kigi_tool_types::sanitize_optional_arg(input.model.clone()); + if let Some(requested) = model.as_deref() { + // Same contract as the `task` tool: an explicitly requested model + // that cannot be checked is refused, not waved through. Skipping + // silently would trade one loud error for `items.len()` quiet ones. + let validator = model_validator.ok_or_else(|| { + kigi_tool_runtime::ToolError::custom( + "validation_unavailable", + "Cannot validate agent_swarm.model: model catalog validator is unavailable.", + ) + })?; + if let Some(error) = validator.error_for(requested) { + return Err(kigi_tool_runtime::ToolError::invalid_arguments(error)); + } + } + let specs = + plan_members(&input).map_err(kigi_tool_runtime::ToolError::invalid_arguments)?; + + match backend + .0 + .validate_type(&input.subagent_type, &parent_session_id) + .await + { + SubagentValidateTypeOutcome::Ok => {} + SubagentValidateTypeOutcome::Unknown { available } => { + let suffix = if available.is_empty() { + String::new() + } else { + format!(". Available types: {}", available.join(", ")) + }; + return Err(kigi_tool_runtime::ToolError::invalid_arguments(format!( + "Unknown subagent type: {}{suffix}", + input.subagent_type + ))); + } + SubagentValidateTypeOutcome::Disabled => { + return Err(kigi_tool_runtime::ToolError::invalid_arguments(format!( + "Subagent '{}' is disabled via [subagents.toggle] in config.toml", + input.subagent_type + ))); + } + SubagentValidateTypeOutcome::NotAllowed { allowed } => { + return Err(kigi_tool_runtime::ToolError::invalid_arguments(format!( + "agent can only spawn: {}; '{}' not allowed", + allowed.join(", "), + input.subagent_type + ))); + } + SubagentValidateTypeOutcome::ValidationUnavailable => { + // `custom` (not `invalid_arguments`) so the model doesn't + // retry with a different name on transport faults. + return Err(kigi_tool_runtime::ToolError::custom( + "validation_unavailable", + format!( + "Cannot validate subagent type '{}': the subagent coordinator is \ + unreachable. Retry shortly or notify ops.", + input.subagent_type + ), + )); + } + } + + let config = SwarmRunConfig { + subagent_type: input.subagent_type.clone(), + description: input.description.clone(), + parent_session_id, + parent_prompt_id, + model, + // Members share the caller's tree; see `run::build_request`. + cwd: None, + max_concurrency, + }; + + let results = run_swarm(backend.0.clone(), specs, config).await; + Ok(ToolOutput::Text(render_results(&results).into())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::implementations::kigi::task::backend::ChannelBackend; + use crate::implementations::kigi::task::types::{SubagentEvent, SubagentResult}; + use crate::types::resources::Resources; + use crate::types::tool_metadata::test_ctx; + use kigi_env::EnvVarGuard; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use tokio::sync::mpsc; + + /// What the coordinator was actually asked to do — a refusal that still + /// reached the backend is the failure these tests exist to catch. + #[derive(Default)] + struct Seen { + validations: AtomicUsize, + spawns: AtomicUsize, + } + + impl Seen { + fn validations(&self) -> usize { + self.validations.load(Ordering::SeqCst) + } + fn spawns(&self) -> usize { + self.spawns.load(Ordering::SeqCst) + } + } + + /// `run` resolves the operator cap from the process environment, which every + /// test in this binary shares. Each test therefore pins the variable for its + /// duration through [`EnvVarGuard`], whose lock also serializes them against + /// each other. The tests are synchronous so that lock never spans an await. + fn cap_unset() -> EnvVarGuard { + EnvVarGuard::remove(MAX_CONCURRENCY_ENV) + } + + fn swarm_input(items: &[&str]) -> AgentSwarmToolInput { + AgentSwarmToolInput { + description: "test swarm".into(), + subagent_type: "general-purpose".into(), + prompt_template: Some("Review {{item}} for bugs".into()), + items: items.iter().map(|s| (*s).to_string()).collect(), + resume_agent_ids: Default::default(), + model: None, + } + } + + /// Drives one whole tool call against a live channel backend that answers + /// `ValidateType` with `validate` and completes every spawn. + fn call( + depth: u32, + validate: SubagentValidateTypeOutcome, + input: AgentSwarmToolInput, + ) -> (Result, Arc) { + let (tx, mut rx) = mpsc::unbounded_channel::(); + let mut resources = Resources::new(); + resources.insert(SubagentBackendResource(Arc::new(ChannelBackend::new(tx)))); + resources.insert(SubagentDepthCounter(depth)); + resources.insert(SessionIdResource("parent-session".to_string())); + resources.insert(CurrentPromptIdResource("prompt-1".to_string())); + + let seen = Arc::new(Seen::default()); + let recorder = seen.clone(); + let result = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime") + .block_on(async move { + let pump = tokio::spawn(async move { + while let Some(event) = rx.recv().await { + match event { + SubagentEvent::ValidateType(req) => { + recorder.validations.fetch_add(1, Ordering::SeqCst); + let _ = req.respond_to.send(validate.clone()); + } + SubagentEvent::Spawn(req) => { + recorder.spawns.fetch_add(1, Ordering::SeqCst); + let _ = req.result_tx.send(SubagentResult { + success: true, + output: Arc::from(format!("done: {}", req.prompt)), + subagent_id: req.id.clone(), + child_session_id: req.id.clone(), + ..Default::default() + }); + } + _ => {} + } + } + }); + let out = kigi_tool_runtime::Tool::run( + &AgentSwarmTool, + test_ctx(resources.into_shared()), + input, + ) + .await + .map_err(|e| e.to_string()); + pump.abort(); + out + }); + (result, seen) + } + + #[test] + fn every_member_runs_and_the_call_returns_one_aggregate() { + let _cap = cap_unset(); + let (result, seen) = call( + 0, + SubagentValidateTypeOutcome::Ok, + swarm_input(&["a.rs", "b.rs"]), + ); + + match result.expect("a valid swarm runs") { + ToolOutput::Text(text) => { + assert!( + text.text + .contains("completed: 2, failed: 0, aborted: 0, still running: 0"), + "{}", + text.text + ); + assert!(text.text.contains(r#"item="a.rs""#), "{}", text.text); + assert!(text.text.contains(r#"item="b.rs""#), "{}", text.text); + } + other => panic!("expected Text output, got {other:?}"), + } + assert_eq!(seen.spawns(), 2, "one member per item"); + } + + /// A swarm child must never start its own swarm: kigi caps subagent nesting + /// at one level, and a fan-out tool is exactly how that cap would be lost — + /// 128 members each starting 128 more. + #[test] + fn a_child_at_the_depth_ceiling_cannot_start_a_swarm() { + let _cap = cap_unset(); + let (result, seen) = call( + MAX_SUBAGENT_DEPTH, + SubagentValidateTypeOutcome::Ok, + swarm_input(&["a.rs", "b.rs"]), + ); + + let err = result.expect_err("a child must not fan out"); + assert!(err.contains("depth limit exceeded"), "error: {err}"); + assert_eq!(seen.spawns(), 0, "no member may reach the coordinator"); + assert_eq!( + seen.validations(), + 0, + "the coordinator must not even be asked to validate" + ); + } + + /// An operator who set a concurrency ceiling must never silently get an + /// unbounded fan-out because the value failed to parse. + #[test] + fn a_malformed_operator_cap_fails_the_call_rather_than_being_ignored() { + let _cap = EnvVarGuard::set(MAX_CONCURRENCY_ENV, "lots"); + let (result, seen) = call( + 0, + SubagentValidateTypeOutcome::Ok, + swarm_input(&["a.rs", "b.rs"]), + ); + + let err = result.expect_err("a cap that does not parse must reject the call"); + assert!(err.contains(MAX_CONCURRENCY_ENV), "error: {err}"); + assert!(err.contains("positive integer"), "error: {err}"); + assert_eq!( + seen.spawns(), + 0, + "an unparseable cap must stop the swarm before anything spawns" + ); + } + + /// A well-formed cap is honoured rather than rejected — the fail-fast path + /// above must not swallow valid operator configuration. + #[test] + fn a_well_formed_operator_cap_still_runs_the_swarm() { + let _cap = EnvVarGuard::set(MAX_CONCURRENCY_ENV, "1"); + let (result, seen) = call( + 0, + SubagentValidateTypeOutcome::Ok, + swarm_input(&["a.rs", "b.rs"]), + ); + + assert!(result.is_ok(), "a valid cap must not fail the call"); + assert_eq!(seen.spawns(), 2, "every member still runs, just serially"); + } + + /// `plan_members` rejections reach the model as invalid_arguments, and + /// nothing spawns. + /// An explicitly requested model that cannot be checked is refused, so a + /// bad slug fails once here instead of once per member. + #[test] + fn an_unvalidatable_model_fails_the_call_before_any_member_spawns() { + let _guard = cap_unset(); + let input = AgentSwarmToolInput { + model: Some("nonexistent/model".into()), + ..swarm_input(&["a.rs", "b.rs"]) + }; + let (result, seen) = call(0, SubagentValidateTypeOutcome::Ok, input); + let err = result.expect_err("no validator resource is registered in this fixture"); + assert!(err.contains("validate"), "{err}"); + assert_eq!( + seen.spawns(), + 0, + "nothing may spawn on an unvalidated model" + ); + } + + /// The discriminating partner: with no model requested the same call runs. + #[test] + fn omitting_the_model_leaves_the_swarm_runnable() { + let _guard = cap_unset(); + let (result, seen) = call( + 0, + SubagentValidateTypeOutcome::Ok, + swarm_input(&["a.rs", "b.rs"]), + ); + assert!(result.is_ok(), "{result:?}"); + assert_eq!(seen.spawns(), 2); + } + + #[test] + fn a_single_item_is_refused_in_favour_of_the_task_tool() { + let _cap = cap_unset(); + let (result, seen) = call( + 0, + SubagentValidateTypeOutcome::Ok, + swarm_input(&["only.rs"]), + ); + + let err = result.expect_err("one item is a task, not a swarm"); + assert!(err.contains("task tool"), "error: {err}"); + assert_eq!(seen.spawns(), 0, "nothing may spawn"); + } + + #[test] + fn an_unknown_subagent_type_is_refused_before_any_member_spawns() { + let _cap = cap_unset(); + let mut input = swarm_input(&["a.rs", "b.rs"]); + input.subagent_type = "invented-agent".into(); + let (result, seen) = call( + 0, + SubagentValidateTypeOutcome::Unknown { + available: vec!["general-purpose".to_string(), "explore".to_string()], + }, + input, + ); + + let err = result.expect_err("an unknown type must reject"); + assert!( + err.contains("Unknown subagent type: invented-agent"), + "error: {err}" + ); + assert!(err.contains("explore"), "error: {err}"); + assert_eq!(seen.validations(), 1, "the type is validated exactly once"); + assert_eq!(seen.spawns(), 0, "no member may spawn"); + } + + #[test] + fn a_missing_backend_is_reported_rather_than_silently_skipped() { + let _cap = cap_unset(); + let result = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("test runtime") + .block_on(kigi_tool_runtime::Tool::run( + &AgentSwarmTool, + test_ctx(Resources::new().into_shared()), + swarm_input(&["a.rs", "b.rs"]), + )); + + let err = result.expect_err("no backend must error").to_string(); + assert!(err.contains("SubagentBackendResource"), "error: {err}"); + } +} diff --git a/crates/codegen/kigi-tools/src/implementations/kigi/mod.rs b/crates/codegen/kigi-tools/src/implementations/kigi/mod.rs index 4755fe6..78b5723 100644 --- a/crates/codegen/kigi-tools/src/implementations/kigi/mod.rs +++ b/crates/codegen/kigi-tools/src/implementations/kigi/mod.rs @@ -1,5 +1,6 @@ //! Tool implementations built on the `NewTool` trait; the sibling //! `implementations//` modules hold the `Tool`-trait counterparts. +pub mod agent_swarm; pub mod ask_user_question; pub mod bash; #[path = "deploy_app_stub.rs"] @@ -21,6 +22,7 @@ pub mod todo; pub mod update_goal; pub mod web_fetch; pub mod web_search; +pub use agent_swarm::AgentSwarmTool; pub use ask_user_question::AskUserQuestionTool; pub use bash::BashTool; pub use deploy_app::{AppBuilderDeployerConfig, DEPLOY_APP_TOOL_NAME}; diff --git a/crates/codegen/kigi-tools/src/implementations/kigi/task/types.rs b/crates/codegen/kigi-tools/src/implementations/kigi/task/types.rs index d7af3cb..0c67ac0 100644 --- a/crates/codegen/kigi-tools/src/implementations/kigi/task/types.rs +++ b/crates/codegen/kigi-tools/src/implementations/kigi/task/types.rs @@ -149,10 +149,12 @@ pub trait SubagentCapabilityModeExt { pub fn prune_orphaned_background_task_tools(config: &mut crate::registry::types::ToolServerConfig) { use crate::types::tool::ToolKind; + // A swarm member can be auto-backgrounded by the coordinator just like a + // `task` child, so either spawner keeps the lifecycle tools alive. let has_task_tool = config .tools .iter() - .any(|tc| tc.kind == Some(ToolKind::Task)); + .any(|tc| matches!(tc.kind, Some(ToolKind::Task | ToolKind::AgentSwarm))); let has_background_capable_bash = config.tools.iter().any(is_background_capable_bash_tool); if has_task_tool || has_background_capable_bash { return; @@ -209,6 +211,7 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode { ToolKind::BackgroundTaskAction, ToolKind::KillTaskAction, ToolKind::Task, + ToolKind::AgentSwarm, ToolKind::EnterPlan, ToolKind::ExitPlan, ToolKind::AskUser, @@ -232,6 +235,7 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode { ToolKind::BackgroundTaskAction, ToolKind::KillTaskAction, ToolKind::Task, + ToolKind::AgentSwarm, ToolKind::EnterPlan, ToolKind::ExitPlan, ToolKind::AskUser, @@ -252,6 +256,7 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode { ToolKind::BackgroundTaskAction, ToolKind::KillTaskAction, ToolKind::Task, + ToolKind::AgentSwarm, ToolKind::EnterPlan, ToolKind::ExitPlan, ToolKind::AskUser, @@ -276,6 +281,7 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode { ToolKind::BackgroundTaskAction, ToolKind::KillTaskAction, ToolKind::Task, + ToolKind::AgentSwarm, ToolKind::EnterPlan, ToolKind::ExitPlan, ToolKind::AskUser, @@ -315,6 +321,12 @@ pub struct SubagentResult { /// `get_command_or_subagent_output`), so the tool returns a `task_id` notice /// instead of a completion. Never set for natively backgrounded subagents. pub backgrounded: bool, + /// The child's turn ended because the PROVIDER refused it for rate + /// limiting, not because the work failed. Classified where the typed ACP + /// error code is still in hand: a fleet scheduler that had to re-derive + /// this from the formatted `error` string would silently stop adapting + /// the day that wording changes. + pub rate_limited: bool, } impl Default for SubagentResult { @@ -324,6 +336,7 @@ impl Default for SubagentResult { output: Arc::from(""), error: None, cancelled: false, + rate_limited: false, subagent_id: String::new(), child_session_id: String::new(), tool_calls: 0, diff --git a/crates/codegen/kigi-tools/src/implementations/mod.rs b/crates/codegen/kigi-tools/src/implementations/mod.rs index 16a2682..79118c8 100644 --- a/crates/codegen/kigi-tools/src/implementations/mod.rs +++ b/crates/codegen/kigi-tools/src/implementations/mod.rs @@ -15,9 +15,9 @@ pub mod use_tool; pub mod web_search; pub use kigi::bash::{BashError, BashToolInput}; pub use kigi::{ - AskUserQuestionTool, BashTool, EnterPlanModeTool, ExitPlanModeTool, GrepTool, KillTaskTool, - ListDirTool, ReadFileTool, SearchReplaceTool, TaskOutputTool, TaskTool, TodoWriteTool, - WaitTasksTool, WebFetchTool, WebSearchTool, + AgentSwarmTool, AskUserQuestionTool, BashTool, EnterPlanModeTool, ExitPlanModeTool, GrepTool, + KillTaskTool, ListDirTool, ReadFileTool, SearchReplaceTool, TaskOutputTool, TaskTool, + TodoWriteTool, WaitTasksTool, WebFetchTool, WebSearchTool, }; pub use memory::{MemoryGetImpl, MemorySearchImpl}; pub use opencode::{ diff --git a/crates/codegen/kigi-tools/src/normalization.rs b/crates/codegen/kigi-tools/src/normalization.rs index 93ab656..69d9837 100644 --- a/crates/codegen/kigi-tools/src/normalization.rs +++ b/crates/codegen/kigi-tools/src/normalization.rs @@ -103,6 +103,7 @@ pub fn canonical_input(input: &ToolInput) -> Option { | ToolInput::WaitTasks(_) | ToolInput::KillTask(_) | ToolInput::Task(_) + | ToolInput::AgentSwarm(_) | ToolInput::WebSearch(_) | ToolInput::WebFetch(_) | ToolInput::ApplyPatch(_) diff --git a/crates/codegen/kigi-tools/src/registry/types.rs b/crates/codegen/kigi-tools/src/registry/types.rs index e57aa50..4e0b40d 100644 --- a/crates/codegen/kigi-tools/src/registry/types.rs +++ b/crates/codegen/kigi-tools/src/registry/types.rs @@ -650,6 +650,7 @@ impl ToolRegistryBuilder { b.register::(); b.register::(); b.register::(); + b.register::(); b.register::(); b.register_with_params::(); b.register::(); diff --git a/crates/codegen/kigi-tools/src/tool_taxonomy.rs b/crates/codegen/kigi-tools/src/tool_taxonomy.rs index 3ec760b..9d80f88 100644 --- a/crates/codegen/kigi-tools/src/tool_taxonomy.rs +++ b/crates/codegen/kigi-tools/src/tool_taxonomy.rs @@ -56,6 +56,7 @@ impl ToolKind { ToolKind::MemorySearch => "Memory Search", ToolKind::MemoryGet => "Memory Read", ToolKind::Task => "Subagent", + ToolKind::AgentSwarm => "Subagent Swarm", ToolKind::EnterPlan => "Enter Plan Mode", ToolKind::ExitPlan => "Exit Plan Mode", ToolKind::AskUser => "Ask User", @@ -96,6 +97,7 @@ impl ToolKind { | ToolKind::KillTaskAction | ToolKind::Skill | ToolKind::Task + | ToolKind::AgentSwarm | ToolKind::DeployApp | ToolKind::SearchTool | ToolKind::UseTool diff --git a/crates/codegen/kigi-tools/src/types/tool.rs b/crates/codegen/kigi-tools/src/types/tool.rs index f3ac7bc..4ac168e 100644 --- a/crates/codegen/kigi-tools/src/types/tool.rs +++ b/crates/codegen/kigi-tools/src/types/tool.rs @@ -88,6 +88,11 @@ pub enum ToolKind { MemorySearch, MemoryGet, Task, + /// Fan-out sibling of [`ToolKind::Task`]: its own kind because + /// `TemplateRenderer`'s `by_kind` map holds ONE tool name per kind, so + /// sharing `Task` would let the swarm win that slot and silently redirect + /// every `${{ tools.by_kind.task }}` reference in other tools' prompts. + AgentSwarm, EnterPlan, ExitPlan, AskUser, diff --git a/crates/codegen/kigi-tools/src/types/tool_io.rs b/crates/codegen/kigi-tools/src/types/tool_io.rs index 49427cd..d62a33a 100644 --- a/crates/codegen/kigi-tools/src/types/tool_io.rs +++ b/crates/codegen/kigi-tools/src/types/tool_io.rs @@ -32,6 +32,7 @@ use crate::implementations::opencode::write::WriteInput; use crate::implementations::search_tool::SearchToolInput; use crate::implementations::skills::skill::SkillInput; use crate::implementations::use_tool::UseToolInput; +use kigi_tool_types::AgentSwarmToolInput; use kigi_tool_types::KillTaskToolInput; use kigi_tool_types::TaskOutputToolInput; use kigi_tool_types::TaskToolInput; @@ -67,6 +68,7 @@ pub enum ToolInput { WaitTasks(WaitTasksToolInput), KillTask(KillTaskToolInput), Task(TaskToolInput), + AgentSwarm(AgentSwarmToolInput), WebSearch(WebSearchInput), WebFetch(WebFetchInput), Write(WriteInput), diff --git a/crates/codegen/kigi-tui/src/acp/tracker.rs b/crates/codegen/kigi-tui/src/acp/tracker.rs index 5d85bc8..adc4c3d 100644 --- a/crates/codegen/kigi-tui/src/acp/tracker.rs +++ b/crates/codegen/kigi-tui/src/acp/tracker.rs @@ -2157,7 +2157,10 @@ fn extract_variant(tc: &acp::ToolCall) -> Option<&str> { } /// Twin without the optional-toolset spelling. fn is_task_variant(variant: Option<&str>) -> bool { - matches!(variant, Some("Task")) + // The swarm belongs here for the same reason `Task` does: its members each + // raise their own SubagentBlock, and it must register the blocking wait or + // a call that holds the turn until every member finishes shows no spinner. + matches!(variant, Some("Task" | "AgentSwarm")) } /// Twin without the optional-toolset spelling. fn is_write_variant(variant: Option<&str>) -> bool { @@ -2190,8 +2193,10 @@ fn is_goal_tool(tc: &acp::ToolCall) -> bool { /// SubagentSpawned notification) provides better visibility. Covers the /// `task` / `Task` / `spawn_subagent` ids and Task-family variant tags. fn is_task_tool(tc: &acp::ToolCall) -> bool { - matches!(tc.title.as_str(), "task" | "Task" | "spawn_subagent") - || is_task_variant(extract_variant(tc)) + matches!( + tc.title.as_str(), + "task" | "Task" | "spawn_subagent" | "agent_swarm" + ) || is_task_variant(extract_variant(tc)) } /// Check if a tool call is a scheduler tool (scheduler_create/delete/list). /// diff --git a/crates/codegen/kigi-workspace/src/capability.rs b/crates/codegen/kigi-workspace/src/capability.rs index 288f574..f12cdb5 100644 --- a/crates/codegen/kigi-workspace/src/capability.rs +++ b/crates/codegen/kigi-workspace/src/capability.rs @@ -91,6 +91,7 @@ pub(crate) const ALL_TOOL_KINDS: &[ToolKind] = &[ ToolKind::MemorySearch, ToolKind::MemoryGet, ToolKind::Task, + ToolKind::AgentSwarm, ToolKind::EnterPlan, ToolKind::ExitPlan, ToolKind::AskUser, @@ -144,8 +145,9 @@ pub(crate) fn kind_allowed(mode: CapabilityMode, kind: ToolKind) -> bool { // Bash / shell. Execute => matches!(mode, M::Execute), - // Process control (background tasks, monitors). - BackgroundTaskAction | WaitTasksAction | KillTaskAction | Task | Monitor => { + // Process control (background tasks, monitors). A swarm is a fan-out of + // subagent spawns, so it gates exactly as `Task` does. + BackgroundTaskAction | WaitTasksAction | KillTaskAction | Task | AgentSwarm | Monitor => { matches!(mode, M::Execute) } diff --git a/crates/common/kigi-tool-types/src/agent_swarm.rs b/crates/common/kigi-tool-types/src/agent_swarm.rs new file mode 100644 index 0000000..36eb8c0 --- /dev/null +++ b/crates/common/kigi-tool-types/src/agent_swarm.rs @@ -0,0 +1,116 @@ +//! Input/output types for the `agent_swarm` tool — one prompt template +//! expanded over a list of items into a fleet of subagents. + +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// The literal a `prompt_template` must contain; each expansion substitutes +/// one `items` entry for it. +pub const PROMPT_TEMPLATE_PLACEHOLDER: &str = "{{item}}"; + +/// Upper bound on members in one call, counting resumes. +pub const MAX_AGENT_SWARM_MEMBERS: usize = 128; + +/// Input for the `agent_swarm` tool. +#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] +pub struct AgentSwarmToolInput { + #[schemars(description = "Short description of what the whole swarm is doing (3-7 words).")] + pub description: String, + + /// Subagent type every item-spawned member runs as. + #[schemars( + description = "Name of the subagent type every member runs as. Built-in types: \"general-purpose\", \"explore\", \"plan\"." + )] + #[serde(default = "default_subagent_type")] + pub subagent_type: String, + + /// Prompt shared by every member; must contain `{{item}}`. + #[schemars( + description = "Prompt shared by every member. Must contain the literal {{item}}, which is replaced by each entry of `items`. Required whenever `items` is given." + )] + #[serde(default)] + pub prompt_template: Option, + + /// The work units. Each expands `prompt_template` into one member. + #[schemars( + description = "One entry per member: each is substituted into `prompt_template`. Give every member a distinct scope so members never edit the same file. At least 2 entries unless `resume_agent_ids` is used." + )] + #[serde(default)] + pub items: Vec, + + /// Continue named subagents from a previous swarm: id → follow-up prompt. + #[schemars( + description = "Continue previously spawned subagents: a map of agent_id (from an earlier agent_swarm result) to the follow-up prompt for that member." + )] + #[serde(default)] + pub resume_agent_ids: std::collections::BTreeMap, + + /// Model slug every member runs on; omitted inherits the caller's. + #[schemars( + description = "Model every member runs on. Omit to inherit the caller's current model." + )] + #[serde(default)] + pub model: Option, +} + +fn default_subagent_type() -> String { + "general-purpose".to_string() +} + +/// How a member's run ended. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum SwarmMemberOutcome { + Completed, + Failed, + Aborted, + /// Still running: it outlived the foreground await budget and the subagent + /// coordinator detached it. Distinct from `Failed` because the member is + /// alive and still writing — relaunching its item would put a second agent + /// on the same files, and resuming it is refused while it runs. + Backgrounded, +} + +impl SwarmMemberOutcome { + pub fn as_str(self) -> &'static str { + match self { + Self::Completed => "completed", + Self::Failed => "failed", + Self::Aborted => "aborted", + Self::Backgrounded => "backgrounded", + } + } + + /// Whether re-running this member's work is safe to suggest. + fn is_resumable(self) -> bool { + matches!(self, Self::Failed | Self::Aborted) + } +} + +/// One member's contribution to the aggregate result. +#[derive(Debug, Clone)] +pub struct SwarmMemberResult { + /// The `items` entry (or the resumed agent id) this member was given — + /// what the caller needs to retry exactly the members that did not finish. + pub item: String, + /// Present once the member started; absent means it never launched. + pub agent_id: Option, + pub resumed: bool, + pub outcome: SwarmMemberOutcome, + pub summary: String, +} + +impl SwarmMemberResult { + /// Whether the member ever reached the backend. A member that never + /// started has nothing to resume. + pub fn started(&self) -> bool { + self.agent_id.is_some() + } + + /// Whether the caller may be told to continue this member. A member that + /// is still running must not be offered: the coordinator refuses to resume + /// a live subagent, and relaunching its item duplicates its writes. + pub fn is_resumable(&self) -> bool { + self.started() && self.outcome.is_resumable() + } +} diff --git a/crates/common/kigi-tool-types/src/lib.rs b/crates/common/kigi-tool-types/src/lib.rs index 4da01ee..bc18b08 100644 --- a/crates/common/kigi-tool-types/src/lib.rs +++ b/crates/common/kigi-tool-types/src/lib.rs @@ -1,10 +1,15 @@ //! Canonical, extensible tool types. +mod agent_swarm; mod ext; mod schema_utils; pub mod serde_lenient; mod task; mod types; +pub use agent_swarm::{ + AgentSwarmToolInput, MAX_AGENT_SWARM_MEMBERS, PROMPT_TEMPLATE_PLACEHOLDER, SwarmMemberOutcome, + SwarmMemberResult, +}; pub use ext::Extensions; pub use schema_utils::parse_arguments_from_schema_lossy; pub use serde_lenient::{