feat(swarm): agent_swarm — one prompt over many items, paced as a fleet

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.
This commit is contained in:
2026-07-27 01:22:52 -04:00
parent ed8049cf77
commit 9edb8729ef
22 changed files with 2232 additions and 16 deletions
@@ -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;
@@ -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<String>,
}
/// 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<Vec<MemberSpec>, 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<MemberSpec> = 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("<agent_swarm_result>\n");
out.push_str(&format!(
"<summary>completed: {completed}, failed: {failed}, aborted: {aborted}, \
still running: {backgrounded}</summary>\n"
));
if backgrounded > 0 {
out.push_str(
"<still_running>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.</still_running>\n",
);
}
if results.iter().any(SwarmMemberResult::is_resumable) {
out.push_str(
"<resume_hint>Call agent_swarm again with resume_agent_ids mapping the agent_id \
values below to a follow-up prompt to continue unfinished work.</resume_hint>\n",
);
}
for result in results {
out.push_str("<member");
if let Some(id) = &result.agent_id {
out.push_str(&format!(" agent_id=\"{}\"", escape_attr(id)));
}
out.push_str(&format!(
" item=\"{}\" state=\"{}\" outcome=\"{}\"",
escape_attr(&result.item),
if result.started() {
"started"
} else {
"not_started"
},
result.outcome.as_str()
));
if result.resumed {
out.push_str(" mode=\"resume\"");
}
out.push_str(">\n");
out.push_str(clamp_summary(result.summary.trim()).trim());
out.push_str("\n</member>\n");
}
out.push_str("</agent_swarm_result>");
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('&', "&amp;")
.replace('<', "&lt;")
.replace('>', "&gt;")
.replace('"', "&quot;")
}
#[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<String> = (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("&quot;"), "{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());
}
}
@@ -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<String>,
pub model: Option<String>,
pub cwd: Option<String>,
pub max_concurrency: Option<usize>,
}
/// 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<Instant>,
}
struct Finished {
agent_id: Option<String>,
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<dyn SubagentBackend>,
live: Arc<std::sync::Mutex<std::collections::BTreeSet<String>>>,
}
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<String> = 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<dyn SubagentBackend>,
specs: Vec<MemberSpec>,
config: SwarmRunConfig,
) -> Vec<SwarmMemberResult> {
let started_at = Instant::now();
let total = specs.len();
let labels: Vec<String> = specs.iter().map(|s| s.item.clone()).collect();
let resumed: Vec<bool> = specs.iter().map(|s| s.resume_from.is_some()).collect();
let mut queue: Vec<Pending> = 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<usize, Finished> = 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<Option<String>> = 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<SubagentResult, kigi_tool_runtime::ToolError>,
);
async fn collect_one(
in_flight: &mut FuturesUnordered<impl Future<Output = InFlight>>,
guard: &InFlightGuard,
pacer: &mut LaunchPacer,
queue: &mut Vec<Pending>,
done: &mut HashMap<usize, Finished>,
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<Pending>,
done: &mut HashMap<usize, Finished>,
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<Pending>) -> Option<Pending> {
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<Duration> {
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<String>,
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<Vec<String>>,
/// prompt -> remaining rate-limit rejections before it succeeds.
refuse: Mutex<HashMap<String, u32>>,
/// prompts that always fail outright.
fail: Mutex<Vec<String>>,
/// prompts the coordinator detaches instead of finishing.
background: Mutex<Vec<String>>,
/// prompt -> how long it occupies its slot, so members can overlap.
duration_ms: Mutex<HashMap<String, u64>>,
cancelled: Mutex<Vec<String>>,
peak_in_flight: AtomicUsize,
in_flight: AtomicUsize,
}
impl FakeBackend {
fn prompts(&self) -> Vec<String> {
self.seen.lock().unwrap().clone()
}
}
#[async_trait::async_trait]
impl SubagentBackend for FakeBackend {
async fn spawn(
&self,
request: SubagentRequest,
) -> Result<SubagentResult, kigi_tool_runtime::ToolError> {
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<u64>) -> Option<SubagentSnapshot> {
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<MemberSpec> {
items
.iter()
.map(|item| MemberSpec {
item: (*item).to_string(),
prompt: format!("work on {item}"),
resume_from: None,
})
.collect()
}
fn config(max_concurrency: Option<usize>) -> 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<_>>(),
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);
}
}
@@ -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<Option<usize>, String> {
let Some(raw) = raw.map(str::trim).filter(|v| !v.is_empty()) else {
return Ok(None);
};
match raw.parse::<usize>() {
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<usize>,
/// `None` until the first rate limit: before that the ramp alone paces us.
capacity: Option<usize>,
last_launch: Option<Duration>,
last_shrink: Option<Duration>,
last_rate_limit: Option<Duration>,
}
impl LaunchPacer {
pub fn new(queued: usize, hard_cap: Option<usize>) -> 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<usize> {
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());
}
}
@@ -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<ToolRequirement> {
// 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<ToolOutput, kigi_tool_runtime::ToolError> {
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::<SubagentDepthCounter>().map(|d| d.0).unwrap_or(0);
let model_validator = res.get::<TaskModelValidator>().cloned();
let backend = res
.get::<SubagentBackendResource>()
.ok_or_else(|| {
kigi_tool_runtime::ToolError::custom(
"missing_resource",
"SubagentBackendResource (subagent support not initialized)",
)
})?
.clone();
let parent_session_id = res
.get::<SessionIdResource>()
.map(|s| s.0.clone())
.unwrap_or_default();
let parent_prompt_id = res
.get::<CurrentPromptIdResource>()
.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<ToolOutput, String>, Arc<Seen>) {
let (tx, mut rx) = mpsc::unbounded_channel::<SubagentEvent>();
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}");
}
}
@@ -1,5 +1,6 @@
//! Tool implementations built on the `NewTool` trait; the sibling
//! `implementations/<tool>/` 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};
@@ -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,
@@ -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::{
@@ -103,6 +103,7 @@ pub fn canonical_input(input: &ToolInput) -> Option<serde_json::Value> {
| ToolInput::WaitTasks(_)
| ToolInput::KillTask(_)
| ToolInput::Task(_)
| ToolInput::AgentSwarm(_)
| ToolInput::WebSearch(_)
| ToolInput::WebFetch(_)
| ToolInput::ApplyPatch(_)
@@ -650,6 +650,7 @@ impl ToolRegistryBuilder {
b.register::<kigi::GetTerminalCommandOutputTool>();
b.register::<kigi::WaitTasksTool>();
b.register::<kigi::TaskTool>();
b.register::<kigi::AgentSwarmTool>();
b.register::<kigi::WebSearchTool>();
b.register_with_params::<kigi::WebFetchTool, kigi::web_fetch::WebFetchParams>();
b.register::<kigi::LspTool>();
@@ -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
@@ -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,
@@ -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),