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:
@@ -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<ToolKind> =
|
||||
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:?}"
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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": {
|
||||
|
||||
@@ -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('&', "&")
|
||||
.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<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("""), "{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),
|
||||
|
||||
@@ -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).
|
||||
///
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user