Compare commits
3
Commits
v0.1.9
...
9edb8729ef
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9edb8729ef | ||
|
|
ed8049cf77 | ||
|
|
867b3e110b |
@@ -711,9 +711,19 @@ impl AgentBuilder {
|
|||||||
kigi_tools::types::tool::ToolNamespace::Kigi,
|
kigi_tools::types::tool::ToolNamespace::Kigi,
|
||||||
"task"
|
"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;
|
let mut task_stripped = false;
|
||||||
if !self.subagents_enabled {
|
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;
|
task_stripped = true;
|
||||||
} else {
|
} else {
|
||||||
let subagents = crate::discovery::all_subagents_with_plugins(
|
let subagents = crate::discovery::all_subagents_with_plugins(
|
||||||
@@ -722,7 +732,9 @@ impl AgentBuilder {
|
|||||||
self.plugin_registry.as_deref(),
|
self.plugin_registry.as_deref(),
|
||||||
);
|
);
|
||||||
if subagents.is_empty() {
|
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;
|
task_stripped = true;
|
||||||
} else if self.prompt_audience == crate::prompt::context::PromptAudience::Subagent {
|
} else if self.prompt_audience == crate::prompt::context::PromptAudience::Subagent {
|
||||||
if let Some(task_tc) = tool_config
|
if let Some(task_tc) = tool_config
|
||||||
@@ -810,7 +822,13 @@ impl AgentBuilder {
|
|||||||
.tools
|
.tools
|
||||||
.iter()
|
.iter()
|
||||||
.any(|t| AGENT_TASK_CLASSIFIER_RE.is_match(t));
|
.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 registered_tool_ids = tool_bridge_builder.known_tool_ids();
|
||||||
let present_kinds: std::collections::HashSet<ToolKind> =
|
let present_kinds: std::collections::HashSet<ToolKind> =
|
||||||
tool_config.tools.iter().filter_map(|tc| tc.kind).collect();
|
tool_config.tools.iter().filter_map(|tc| tc.kind).collect();
|
||||||
@@ -925,7 +943,13 @@ impl AgentBuilder {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
if definition.allowed_subagent_types.as_deref() == Some(&[]) {
|
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
|
tool_config
|
||||||
.tools
|
.tools
|
||||||
.retain(|tc| !task_deps.contains(&short_tool_name(&tc.id)));
|
.retain(|tc| !task_deps.contains(&short_tool_name(&tc.id)));
|
||||||
@@ -1601,6 +1625,11 @@ mod tests {
|
|||||||
has_task, *subagents,
|
has_task, *subagents,
|
||||||
"[{label}] spawn_subagent presence should match subagents_enabled={subagents}; got tools: {names:?}"
|
"[{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!(
|
assert!(
|
||||||
names.contains(&"enter_plan_mode"),
|
names.contains(&"enter_plan_mode"),
|
||||||
"[{label}] enter_plan_mode must always be present (TUI plan-mode keybind needs it); got tools: {names:?}"
|
"[{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_name("spawn_subagent")
|
||||||
.with_param_rename("run_in_background", "background")
|
.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:
|
/// Task output tool renamed for clarity:
|
||||||
/// `get_task_output` → `get_command_or_subagent_output`.
|
/// `get_task_output` → `get_command_or_subagent_output`.
|
||||||
fn task_output_tool_config() -> ToolConfig {
|
fn task_output_tool_config() -> ToolConfig {
|
||||||
@@ -270,6 +275,7 @@ fn default_kigi_toolset() -> ToolServerConfig {
|
|||||||
task_output_tool_config(),
|
task_output_tool_config(),
|
||||||
wait_tasks_tool_config(),
|
wait_tasks_tool_config(),
|
||||||
task_tool_config(),
|
task_tool_config(),
|
||||||
|
agent_swarm_tool_config(),
|
||||||
(&kigi::SchedulerCreateTool).into(),
|
(&kigi::SchedulerCreateTool).into(),
|
||||||
(&kigi::SchedulerDeleteTool).into(),
|
(&kigi::SchedulerDeleteTool).into(),
|
||||||
(&kigi::SchedulerListTool).into(),
|
(&kigi::SchedulerListTool).into(),
|
||||||
@@ -318,6 +324,7 @@ pub fn kigi_hashline_toolset(
|
|||||||
task_output_tool_config(),
|
task_output_tool_config(),
|
||||||
wait_tasks_tool_config(),
|
wait_tasks_tool_config(),
|
||||||
task_tool_config(),
|
task_tool_config(),
|
||||||
|
agent_swarm_tool_config(),
|
||||||
(&kigi::WebSearchTool).into(),
|
(&kigi::WebSearchTool).into(),
|
||||||
(&kigi::SchedulerCreateTool).into(),
|
(&kigi::SchedulerCreateTool).into(),
|
||||||
(&kigi::SchedulerDeleteTool).into(),
|
(&kigi::SchedulerDeleteTool).into(),
|
||||||
@@ -400,6 +407,7 @@ fn kigi_plan_toolset() -> ToolServerConfig {
|
|||||||
(&kigi::TodoWriteTool).into(),
|
(&kigi::TodoWriteTool).into(),
|
||||||
task_output_tool_config(),
|
task_output_tool_config(),
|
||||||
task_tool_config(),
|
task_tool_config(),
|
||||||
|
agent_swarm_tool_config(),
|
||||||
(&kigi::SchedulerCreateTool).into(),
|
(&kigi::SchedulerCreateTool).into(),
|
||||||
(&kigi::SchedulerDeleteTool).into(),
|
(&kigi::SchedulerDeleteTool).into(),
|
||||||
(&kigi::SchedulerListTool).into(),
|
(&kigi::SchedulerListTool).into(),
|
||||||
@@ -428,6 +436,7 @@ fn orchestrator_toolset() -> ToolServerConfig {
|
|||||||
(&kigi::ListDirTool).into(),
|
(&kigi::ListDirTool).into(),
|
||||||
(&kigi::GrepTool).into(),
|
(&kigi::GrepTool).into(),
|
||||||
task_tool_config(),
|
task_tool_config(),
|
||||||
|
agent_swarm_tool_config(),
|
||||||
task_output_tool_config(),
|
task_output_tool_config(),
|
||||||
wait_tasks_tool_config(),
|
wait_tasks_tool_config(),
|
||||||
kill_task_tool_config(),
|
kill_task_tool_config(),
|
||||||
@@ -497,6 +506,7 @@ fn kigi_ask_user_toolset() -> ToolServerConfig {
|
|||||||
task_output_tool_config(),
|
task_output_tool_config(),
|
||||||
wait_tasks_tool_config(),
|
wait_tasks_tool_config(),
|
||||||
task_tool_config(),
|
task_tool_config(),
|
||||||
|
agent_swarm_tool_config(),
|
||||||
(&kigi::SchedulerCreateTool).into(),
|
(&kigi::SchedulerCreateTool).into(),
|
||||||
(&kigi::SchedulerDeleteTool).into(),
|
(&kigi::SchedulerDeleteTool).into(),
|
||||||
(&kigi::SchedulerListTool).into(),
|
(&kigi::SchedulerListTool).into(),
|
||||||
|
|||||||
@@ -18,6 +18,76 @@ use super::embedding::EmbeddingProvider as _;
|
|||||||
use super::storage::MemoryStorage;
|
use super::storage::MemoryStorage;
|
||||||
use super::watcher::MemoryFileWatcher;
|
use super::watcher::MemoryFileWatcher;
|
||||||
|
|
||||||
|
/// The session's embedding credentials, bound to the one endpoint they may
|
||||||
|
/// reach. Only [`Self::for_endpoint`] retains a live handle; the default fails
|
||||||
|
/// closed.
|
||||||
|
///
|
||||||
|
/// `embed_base_url` is the CURRENT MODEL's endpoint, so a session on a BYOK or
|
||||||
|
/// subscription-OAuth model aims memory embeddings at that provider's host —
|
||||||
|
/// and [`kigi_auth::AuthRetryMiddleware`] stamps `Authorization` on every
|
||||||
|
/// request it wraps, with no idea where the request is going. The caller that
|
||||||
|
/// owns the credential rule decides `trusted` once, here; a platform's own
|
||||||
|
/// `embed_api_key` is unaffected and keeps serving its own endpoint.
|
||||||
|
#[derive(Clone, Default)]
|
||||||
|
pub struct EndpointScopedCredentials {
|
||||||
|
endpoint: Option<reqwest::Url>,
|
||||||
|
auth_credentials: Option<Arc<dyn kigi_auth::AuthCredentialProvider>>,
|
||||||
|
}
|
||||||
|
|
||||||
|
// Redacts the credential handles; only their presence is printable.
|
||||||
|
impl std::fmt::Debug for EndpointScopedCredentials {
|
||||||
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||||
|
f.debug_struct("EndpointScopedCredentials")
|
||||||
|
.field("endpoint", &self.endpoint)
|
||||||
|
.field("has_auth_credentials", &self.auth_credentials.is_some())
|
||||||
|
.finish()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl EndpointScopedCredentials {
|
||||||
|
/// No session credential may ride — the state every non-session caller wants.
|
||||||
|
pub fn none() -> Self {
|
||||||
|
Self::default()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Retains the handles only for a `trusted`, parsable `endpoint`.
|
||||||
|
pub fn for_endpoint(
|
||||||
|
endpoint: &str,
|
||||||
|
trusted: bool,
|
||||||
|
auth_credentials: Option<Arc<dyn kigi_auth::AuthCredentialProvider>>,
|
||||||
|
) -> Self {
|
||||||
|
if trusted && let Ok(url) = reqwest::Url::parse(endpoint) {
|
||||||
|
return Self {
|
||||||
|
endpoint: Some(url),
|
||||||
|
auth_credentials,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
if auth_credentials.is_some() {
|
||||||
|
tracing::info!(
|
||||||
|
target: kigi_log::memory_log::TARGET,
|
||||||
|
endpoint,
|
||||||
|
"memory embeddings: session credentials withheld from this endpoint; \
|
||||||
|
its own key, if any, still applies"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
Self::none()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn is_empty(&self) -> bool {
|
||||||
|
self.auth_credentials.is_none()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Enforced at request-build time, in release too: [`MemoryBackendParams`]
|
||||||
|
/// is `Clone` and callers rewrite fields on the copy, so construction-time
|
||||||
|
/// scoping alone would not survive a rewritten `embed_base_url`.
|
||||||
|
fn approved_for(&self, base_url: &str) -> bool {
|
||||||
|
match &self.endpoint {
|
||||||
|
None => self.is_empty(),
|
||||||
|
Some(endpoint) => reqwest::Url::parse(base_url).is_ok_and(|url| &url == endpoint),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// All configuration needed to build a fully-wired [`MemoryBackendImpl`] for a live session.
|
/// All configuration needed to build a fully-wired [`MemoryBackendImpl`] for a live session.
|
||||||
///
|
///
|
||||||
/// Grouping these in one struct ensures every call site — ToolBridge, first-turn
|
/// Grouping these in one struct ensures every call site — ToolBridge, first-turn
|
||||||
@@ -30,7 +100,8 @@ pub struct MemoryBackendParams {
|
|||||||
pub session_id: String,
|
pub session_id: String,
|
||||||
/// Embedding provider config — `None` forces FTS-only fallback everywhere.
|
/// Embedding provider config — `None` forces FTS-only fallback everywhere.
|
||||||
pub embed_config: Option<kigi_config_types::MemoryEmbeddingConfig>,
|
pub embed_config: Option<kigi_config_types::MemoryEmbeddingConfig>,
|
||||||
/// Base URL for embedding API calls (CLI proxy).
|
/// Base URL for embedding API calls (CLI proxy). Must match the endpoint
|
||||||
|
/// `embedding_credentials` was scoped to; a mismatch fails closed.
|
||||||
pub embed_base_url: String,
|
pub embed_base_url: String,
|
||||||
/// API key for embedding API calls.
|
/// API key for embedding API calls.
|
||||||
pub embed_api_key: Option<String>,
|
pub embed_api_key: Option<String>,
|
||||||
@@ -47,10 +118,8 @@ pub struct MemoryBackendParams {
|
|||||||
/// - `"injection"` — first-turn memory context injection
|
/// - `"injection"` — first-turn memory context injection
|
||||||
/// - `"compaction_recovery"` — post-compaction context re-injection
|
/// - `"compaction_recovery"` — post-compaction context re-injection
|
||||||
pub search_source: &'static str,
|
pub search_source: &'static str,
|
||||||
/// Dynamic API key provider — when set, `make_embedding_provider()` resolves
|
/// The session credentials, and the single endpoint they may reach.
|
||||||
/// the key per-call instead of using the static `embed_api_key`.
|
pub embedding_credentials: EndpointScopedCredentials,
|
||||||
pub api_key_provider: Option<kigi_tools::types::SharedApiKeyProvider>,
|
|
||||||
pub auth_credentials: Option<Arc<dyn kigi_auth::AuthCredentialProvider>>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl MemoryBackendParams {
|
impl MemoryBackendParams {
|
||||||
@@ -59,8 +128,7 @@ impl MemoryBackendParams {
|
|||||||
pub async fn make_embedding_provider(&self) -> Option<super::embedding::ApiEmbeddingProvider> {
|
pub async fn make_embedding_provider(&self) -> Option<super::embedding::ApiEmbeddingProvider> {
|
||||||
build_embedding_provider(
|
build_embedding_provider(
|
||||||
self.embed_config.as_ref(),
|
self.embed_config.as_ref(),
|
||||||
self.auth_credentials.as_ref(),
|
&self.embedding_credentials,
|
||||||
self.api_key_provider.as_ref(),
|
|
||||||
self.embed_api_key.as_deref(),
|
self.embed_api_key.as_deref(),
|
||||||
&self.embed_base_url,
|
&self.embed_base_url,
|
||||||
)
|
)
|
||||||
@@ -70,8 +138,7 @@ impl MemoryBackendParams {
|
|||||||
|
|
||||||
async fn build_embedding_provider(
|
async fn build_embedding_provider(
|
||||||
config: Option<&kigi_config_types::MemoryEmbeddingConfig>,
|
config: Option<&kigi_config_types::MemoryEmbeddingConfig>,
|
||||||
auth_credentials: Option<&Arc<dyn kigi_auth::AuthCredentialProvider>>,
|
credentials: &EndpointScopedCredentials,
|
||||||
api_key_provider: Option<&kigi_tools::types::SharedApiKeyProvider>,
|
|
||||||
static_api_key: Option<&str>,
|
static_api_key: Option<&str>,
|
||||||
base_url: &str,
|
base_url: &str,
|
||||||
) -> Option<super::embedding::ApiEmbeddingProvider> {
|
) -> Option<super::embedding::ApiEmbeddingProvider> {
|
||||||
@@ -80,9 +147,19 @@ async fn build_embedding_provider(
|
|||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let approved = credentials.approved_for(base_url);
|
||||||
|
if !approved {
|
||||||
|
tracing::error!(
|
||||||
|
target: kigi_log::memory_log::TARGET,
|
||||||
|
base_url,
|
||||||
|
approved_endpoint = ?credentials.endpoint,
|
||||||
|
"memory embeddings: scoped credentials do not match the request URL; dropping them"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
// Prefer the refresh-capable credential provider — the middleware gives
|
// Prefer the refresh-capable credential provider — the middleware gives
|
||||||
// 401 retry for free without any per-call key resolution.
|
// 401 retry for free without any per-call key resolution.
|
||||||
if let Some(creds) = auth_credentials {
|
if approved && let Some(creds) = credentials.auth_credentials.as_ref() {
|
||||||
let client = super::embedding::build_middleware_client(creds.clone());
|
let client = super::embedding::build_middleware_client(creds.clone());
|
||||||
return super::embedding::ApiEmbeddingProvider::from_config(
|
return super::embedding::ApiEmbeddingProvider::from_config(
|
||||||
config,
|
config,
|
||||||
@@ -91,14 +168,13 @@ async fn build_embedding_provider(
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Fallback: resolve API key per-call, wrap in a static middleware client
|
// The platform's own configured key, wrapped in a static middleware client
|
||||||
// (no 401 refresh, but auth header is still stamped by middleware).
|
// (no 401 refresh, but the auth header is still stamped by middleware).
|
||||||
let api_key = match api_key_provider {
|
super::embedding::ApiEmbeddingProvider::from_session(
|
||||||
Some(p) => p.current_api_key_async().await,
|
config,
|
||||||
None => None,
|
base_url.to_owned(),
|
||||||
}
|
static_api_key?.to_owned(),
|
||||||
.or_else(|| static_api_key.map(|s| s.to_owned()))?;
|
)
|
||||||
super::embedding::ApiEmbeddingProvider::from_session(config, base_url.to_owned(), api_key)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// `MemoryBackend` implementation backed by hybrid search (FTS5 + vector KNN).
|
/// `MemoryBackend` implementation backed by hybrid search (FTS5 + vector KNN).
|
||||||
@@ -129,10 +205,8 @@ pub struct MemoryBackendImpl {
|
|||||||
/// Only the ToolBridge backend's counter is shared back to the session actor;
|
/// Only the ToolBridge backend's counter is shared back to the session actor;
|
||||||
/// injection and compaction-recovery backends use their own local counters.
|
/// injection and compaction-recovery backends use their own local counters.
|
||||||
pub search_counter: std::sync::Arc<std::sync::atomic::AtomicU64>,
|
pub search_counter: std::sync::Arc<std::sync::atomic::AtomicU64>,
|
||||||
/// Dynamic API key provider for embedding requests.
|
/// The session credentials, and the single endpoint they may reach.
|
||||||
api_key_provider: Option<kigi_tools::types::SharedApiKeyProvider>,
|
embedding_credentials: EndpointScopedCredentials,
|
||||||
/// Refresh-capable credential provider for embedding HTTP middleware.
|
|
||||||
auth_credentials: Option<Arc<dyn kigi_auth::AuthCredentialProvider>>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl MemoryBackendImpl {
|
impl MemoryBackendImpl {
|
||||||
@@ -150,8 +224,7 @@ impl MemoryBackendImpl {
|
|||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
session_id: String::new(),
|
session_id: String::new(),
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
search_counter: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)),
|
search_counter: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -200,8 +273,7 @@ impl MemoryBackendImpl {
|
|||||||
async fn make_embedding_provider(&self) -> Option<super::embedding::ApiEmbeddingProvider> {
|
async fn make_embedding_provider(&self) -> Option<super::embedding::ApiEmbeddingProvider> {
|
||||||
build_embedding_provider(
|
build_embedding_provider(
|
||||||
self.embed_config.as_ref(),
|
self.embed_config.as_ref(),
|
||||||
self.auth_credentials.as_ref(),
|
&self.embedding_credentials,
|
||||||
self.api_key_provider.as_ref(),
|
|
||||||
self.embed_api_key.as_deref(),
|
self.embed_api_key.as_deref(),
|
||||||
&self.embed_base_url,
|
&self.embed_base_url,
|
||||||
)
|
)
|
||||||
@@ -232,8 +304,7 @@ impl MemoryBackendImpl {
|
|||||||
if let Some(w) = ¶ms.watcher {
|
if let Some(w) = ¶ms.watcher {
|
||||||
backend = backend.with_watcher(w.clone(), params.stale_claim_secs);
|
backend = backend.with_watcher(w.clone(), params.stale_claim_secs);
|
||||||
}
|
}
|
||||||
backend.api_key_provider = params.api_key_provider.clone();
|
backend.embedding_credentials = params.embedding_credentials.clone();
|
||||||
backend.auth_credentials = params.auth_credentials.clone();
|
|
||||||
backend
|
backend
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -479,8 +550,7 @@ mod factory_tests {
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1038,72 +1108,128 @@ mod factory_tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Regression: provider build must use `current_api_key_async`,
|
struct ProbeCredentials;
|
||||||
/// never sync. Prevents memory_search 401s on rotated tokens.
|
impl kigi_auth::HttpAuth for ProbeCredentials {
|
||||||
#[tokio::test]
|
fn apply(
|
||||||
async fn make_embedding_provider_uses_async_api_key_resolution() {
|
&self,
|
||||||
use kigi_tools::types::ApiKeyProvider;
|
builder: reqwest::RequestBuilder,
|
||||||
use std::sync::atomic::{AtomicU32, Ordering};
|
_base_url: &str,
|
||||||
|
) -> reqwest::RequestBuilder {
|
||||||
struct AsyncProbe {
|
builder.bearer_auth("session-bearer")
|
||||||
sync_calls: Arc<AtomicU32>,
|
|
||||||
async_calls: Arc<AtomicU32>,
|
|
||||||
}
|
}
|
||||||
impl ApiKeyProvider for AsyncProbe {
|
}
|
||||||
fn current_api_key(&self) -> Option<String> {
|
#[async_trait::async_trait]
|
||||||
self.sync_calls.fetch_add(1, Ordering::SeqCst);
|
impl kigi_auth::AuthCredentialProvider for ProbeCredentials {
|
||||||
Some("sync-stale".into())
|
fn snapshot(&self) -> kigi_auth::CredentialSnapshot {
|
||||||
}
|
kigi_auth::CredentialSnapshot {
|
||||||
fn current_api_key_async(
|
token: Some("session-bearer".into()),
|
||||||
&self,
|
..Default::default()
|
||||||
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Option<String>> + Send + '_>>
|
|
||||||
{
|
|
||||||
let counter = self.async_calls.clone();
|
|
||||||
Box::pin(async move {
|
|
||||||
counter.fetch_add(1, Ordering::SeqCst);
|
|
||||||
Some("async-fresh".into())
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
async fn refresh_after_unauthorized(&self) -> bool {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
let sync_calls = Arc::new(AtomicU32::new(0));
|
const SESSION_ENDPOINT: &str = "https://api.kimi.com/coding/v1";
|
||||||
let async_calls = Arc::new(AtomicU32::new(0));
|
const FOREIGN_ENDPOINT: &str = "https://api.anthropic.com/v1";
|
||||||
let probe: kigi_tools::types::SharedApiKeyProvider = Arc::new(AsyncProbe {
|
|
||||||
sync_calls: sync_calls.clone(),
|
|
||||||
async_calls: async_calls.clone(),
|
|
||||||
});
|
|
||||||
|
|
||||||
let params = MemoryBackendParams {
|
fn params_at(base_url: &str, credentials: EndpointScopedCredentials) -> MemoryBackendParams {
|
||||||
|
MemoryBackendParams {
|
||||||
session_id: "s1".into(),
|
session_id: "s1".into(),
|
||||||
embed_config: Some(MemoryEmbeddingConfig {
|
embed_config: Some(MemoryEmbeddingConfig {
|
||||||
model: Some("test-embed-model".into()),
|
model: Some("test-embed-model".into()),
|
||||||
..Default::default()
|
..Default::default()
|
||||||
}),
|
}),
|
||||||
embed_base_url: "http://example/v1".into(),
|
embed_base_url: base_url.into(),
|
||||||
embed_api_key: Some("static-fallback".into()),
|
embed_api_key: None,
|
||||||
search_config: MemorySearchConfig::default(),
|
search_config: MemorySearchConfig::default(),
|
||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: Some(probe),
|
embedding_credentials: credentials,
|
||||||
// No auth_credentials — forces the api_key_provider fallback path.
|
}
|
||||||
auth_credentials: None,
|
}
|
||||||
};
|
|
||||||
|
|
||||||
let provider = params.make_embedding_provider().await;
|
/// The session bearer must never ride to a third-party embedding host.
|
||||||
|
///
|
||||||
|
/// `embed_base_url` is the CURRENT MODEL's endpoint, so a session on a BYOK
|
||||||
|
/// or subscription-OAuth model aims memory embeddings at that provider —
|
||||||
|
/// and the auth middleware stamps `Authorization` unconditionally.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn session_credentials_are_withheld_from_a_foreign_endpoint() {
|
||||||
|
let scoped = EndpointScopedCredentials::for_endpoint(
|
||||||
|
FOREIGN_ENDPOINT,
|
||||||
|
false,
|
||||||
|
Some(Arc::new(ProbeCredentials)),
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
scoped.is_empty(),
|
||||||
|
"an untrusted endpoint must drop both handles"
|
||||||
|
);
|
||||||
|
|
||||||
|
let provider = params_at(FOREIGN_ENDPOINT, scoped)
|
||||||
|
.make_embedding_provider()
|
||||||
|
.await;
|
||||||
|
assert!(
|
||||||
|
provider.is_none(),
|
||||||
|
"no credential may ride to a foreign endpoint, and there is no static key to fall back to"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// The trusted-endpoint path still builds a credentialed provider.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn session_credentials_ride_their_own_endpoint() {
|
||||||
|
let scoped = EndpointScopedCredentials::for_endpoint(
|
||||||
|
SESSION_ENDPOINT,
|
||||||
|
true,
|
||||||
|
Some(Arc::new(ProbeCredentials)),
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!scoped.is_empty(),
|
||||||
|
"a trusted endpoint must retain the handles"
|
||||||
|
);
|
||||||
|
|
||||||
|
let provider = params_at(SESSION_ENDPOINT, scoped)
|
||||||
|
.make_embedding_provider()
|
||||||
|
.await;
|
||||||
assert!(
|
assert!(
|
||||||
provider.is_some(),
|
provider.is_some(),
|
||||||
"provider must be built when model is set"
|
"the session endpoint keeps its refresh-capable credential"
|
||||||
);
|
);
|
||||||
assert_eq!(
|
}
|
||||||
async_calls.load(Ordering::SeqCst),
|
|
||||||
1,
|
/// The runtime re-check, not construction alone, is what guards the wire:
|
||||||
"must call current_api_key_async exactly once per provider build"
|
/// `MemoryBackendParams` is `Clone` and callers rewrite fields on the copy.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn a_cloned_param_set_cannot_redirect_scoped_credentials() {
|
||||||
|
let scoped = EndpointScopedCredentials::for_endpoint(
|
||||||
|
SESSION_ENDPOINT,
|
||||||
|
true,
|
||||||
|
Some(Arc::new(ProbeCredentials)),
|
||||||
);
|
);
|
||||||
assert_eq!(
|
let redirected = MemoryBackendParams {
|
||||||
sync_calls.load(Ordering::SeqCst),
|
embed_base_url: FOREIGN_ENDPOINT.into(),
|
||||||
0,
|
..params_at(SESSION_ENDPOINT, scoped)
|
||||||
"sync current_api_key must NOT be called — the async path is the contract"
|
};
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
redirected.make_embedding_provider().await.is_none(),
|
||||||
|
"credentials scoped to one endpoint must not follow a rewritten base_url"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// A platform's own API key is not a session credential: it is resolved for
|
||||||
|
/// that platform and must keep serving that platform's endpoint.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn a_platform_api_key_still_serves_its_own_endpoint() {
|
||||||
|
let params = MemoryBackendParams {
|
||||||
|
embed_api_key: Some("platform-key".into()),
|
||||||
|
..params_at(FOREIGN_ENDPOINT, EndpointScopedCredentials::none())
|
||||||
|
};
|
||||||
|
assert!(
|
||||||
|
params.make_embedding_provider().await.is_some(),
|
||||||
|
"withholding the session credential must not disable BYOK embeddings"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -31,7 +31,7 @@ pub mod storage;
|
|||||||
pub mod text_utils;
|
pub mod text_utils;
|
||||||
pub mod watcher;
|
pub mod watcher;
|
||||||
|
|
||||||
pub use backend::{MemoryBackendImpl, MemoryBackendParams};
|
pub use backend::{EndpointScopedCredentials, MemoryBackendImpl, MemoryBackendParams};
|
||||||
pub use index::{MemoryIndex, init_sqlite_vec};
|
pub use index::{MemoryIndex, init_sqlite_vec};
|
||||||
pub use storage::{MemoryScope, MemoryStorage};
|
pub use storage::{MemoryScope, MemoryStorage};
|
||||||
|
|
||||||
|
|||||||
@@ -396,11 +396,14 @@ pub(crate) async fn handle_subagent_request(
|
|||||||
let child_depth = ctx.parent_depth + 1;
|
let child_depth = ctx.parent_depth + 1;
|
||||||
if child_depth >= MAX_SUBAGENT_DEPTH {
|
if child_depth >= MAX_SUBAGENT_DEPTH {
|
||||||
let before = definition.tool_config.tools.len();
|
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 {
|
if definition.tool_config.tools.len() < before {
|
||||||
tracing::info!(
|
tracing::info!(
|
||||||
subagent_id = % request.id, child_depth, max_depth =
|
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);
|
prune_orphaned_background_task_tools(&mut definition.tool_config);
|
||||||
@@ -1322,6 +1325,7 @@ pub(crate) async fn handle_subagent_request(
|
|||||||
SubagentResult {
|
SubagentResult {
|
||||||
success: false,
|
success: false,
|
||||||
cancelled: true,
|
cancelled: true,
|
||||||
|
rate_limited: false,
|
||||||
error: Some(reason),
|
error: Some(reason),
|
||||||
output: if final_text.is_empty() {
|
output: if final_text.is_empty() {
|
||||||
std::sync::Arc::from(
|
std::sync::Arc::from(
|
||||||
@@ -1359,6 +1363,7 @@ pub(crate) async fn handle_subagent_request(
|
|||||||
SubagentResult {
|
SubagentResult {
|
||||||
success: false,
|
success: false,
|
||||||
cancelled: true,
|
cancelled: true,
|
||||||
|
rate_limited: false,
|
||||||
error: Some(format!("max turns reached (limit: {limit})")),
|
error: Some(format!("max turns reached (limit: {limit})")),
|
||||||
output: if final_text.is_empty() {
|
output: if final_text.is_empty() {
|
||||||
std::sync::Arc::from(
|
std::sync::Arc::from(
|
||||||
@@ -1413,6 +1418,9 @@ pub(crate) async fn handle_subagent_request(
|
|||||||
SubagentResult {
|
SubagentResult {
|
||||||
success: false,
|
success: false,
|
||||||
cancelled: was_cancelled,
|
cancelled: was_cancelled,
|
||||||
|
rate_limited: !was_cancelled
|
||||||
|
&& i32::from(e.code)
|
||||||
|
== crate::sampling::error::RATE_LIMITED_ERROR_CODE,
|
||||||
error: Some(
|
error: Some(
|
||||||
if was_cancelled {
|
if was_cancelled {
|
||||||
"Subagent was cancelled".to_string()
|
"Subagent was cancelled".to_string()
|
||||||
|
|||||||
@@ -83,6 +83,35 @@ impl AuthCredentialProvider for ShellAuthCredentialProvider {
|
|||||||
self.auth_manager.try_recover_unauthorized().await
|
self.auth_manager.try_recover_unauthorized().await
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
/// The memory-embedding credentials for `embed_base_url`, decided by the ONE
|
||||||
|
/// authority rather than re-derived here (C1).
|
||||||
|
///
|
||||||
|
/// `embed_base_url` is the session model's own endpoint, so on a BYOK or
|
||||||
|
/// subscription-OAuth model it points at that provider's host — and
|
||||||
|
/// [`kigi_auth::AuthRetryMiddleware`] stamps `Authorization` on every request
|
||||||
|
/// it wraps. Asking [`CredentialAuthority::manager_for`] answers both halves at
|
||||||
|
/// once: whether a session credential may ride there at all, and WHICH manager
|
||||||
|
/// governs it (a subscription-OAuth platform's pooled manager at its own host,
|
||||||
|
/// the primary at the session's coding endpoint). No manager means no session
|
||||||
|
/// credential at all; the platform's own `embed_api_key` is untouched and keeps
|
||||||
|
/// serving its own endpoint.
|
||||||
|
///
|
||||||
|
/// The session's `SharedApiKeyProvider` is deliberately NOT forwarded: it is
|
||||||
|
/// hard-wired to the PRIMARY manager, so at a pooled platform's host — where a
|
||||||
|
/// credential may ride, but only that platform's own — it would resolve the
|
||||||
|
/// wrong bearer.
|
||||||
|
pub(crate) fn embedding_session_credentials(
|
||||||
|
embed_base_url: &str,
|
||||||
|
platform: Option<kigi_models::PlatformId>,
|
||||||
|
authority: &crate::auth::credential_authority::CredentialAuthority,
|
||||||
|
) -> kigi_memory::EndpointScopedCredentials {
|
||||||
|
let auth_credentials = authority.manager_for(platform, embed_base_url).map(|am| {
|
||||||
|
Arc::new(ShellAuthCredentialProvider::new(am, None, None))
|
||||||
|
as Arc<dyn AuthCredentialProvider>
|
||||||
|
});
|
||||||
|
let may_ride = auth_credentials.is_some();
|
||||||
|
kigi_memory::EndpointScopedCredentials::for_endpoint(embed_base_url, may_ride, auth_credentials)
|
||||||
|
}
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -253,6 +282,38 @@ mod tests {
|
|||||||
"snapshot must reflect refreshed token for subsequent apply() calls"
|
"snapshot must reflect refreshed token for subsequent apply() calls"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
/// C1: memory embeddings follow the credential authority, not the session.
|
||||||
|
///
|
||||||
|
/// A session whose model is a BYOK platform aims `embed_base_url` at that
|
||||||
|
/// provider's host; the authority answers "no manager governs a credential
|
||||||
|
/// there", so nothing rides. The session's own coding endpoint still does.
|
||||||
|
#[test]
|
||||||
|
fn embedding_credentials_follow_the_credential_authority() {
|
||||||
|
let _guard = EarlyInvalidationGuard::pin_to_default();
|
||||||
|
let dir = tempfile::tempdir().unwrap();
|
||||||
|
let mgr = make_manager(
|
||||||
|
&dir,
|
||||||
|
Some(make_auth("session-bearer", ChronoDuration::hours(1))),
|
||||||
|
);
|
||||||
|
let endpoints = crate::agent::config::EndpointsConfig::default();
|
||||||
|
let coding_endpoint = endpoints.proxy_url();
|
||||||
|
let authority =
|
||||||
|
crate::auth::credential_authority::CredentialAuthority::new(endpoints, Some(mgr));
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
!embedding_session_credentials(&coding_endpoint, None, &authority).is_empty(),
|
||||||
|
"the session's own coding endpoint keeps its credential"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
embedding_session_credentials(
|
||||||
|
"https://api.anthropic.com/v1",
|
||||||
|
kigi_models::PlatformId::parse("anthropic"),
|
||||||
|
&authority,
|
||||||
|
)
|
||||||
|
.is_empty(),
|
||||||
|
"an API-key platform's host must receive no session credential"
|
||||||
|
);
|
||||||
|
}
|
||||||
/// Deployment-key path has no recovery (operator owns the bearer).
|
/// Deployment-key path has no recovery (operator owns the bearer).
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn refresh_after_unauthorized_is_noop_for_deployment_key() {
|
async fn refresh_after_unauthorized_is_noop_for_deployment_key() {
|
||||||
|
|||||||
@@ -311,6 +311,18 @@ pub(crate) async fn spawn_session_actor(
|
|||||||
};
|
};
|
||||||
let embed_base_url = sampling_config.base_url.clone();
|
let embed_base_url = sampling_config.base_url.clone();
|
||||||
let embed_api_key = sampling_config.api_key.clone();
|
let embed_api_key = sampling_config.api_key.clone();
|
||||||
|
// The platform behind the endpoint memory embeddings will call, resolved
|
||||||
|
// through the SAME session-key disambiguation the actor is seeded with, so
|
||||||
|
// a slug that collides across platforms cannot resolve to the twin (H-b).
|
||||||
|
let embed_platform = {
|
||||||
|
let models = models_manager.models();
|
||||||
|
crate::agent::models::platform_for_slug(
|
||||||
|
&models,
|
||||||
|
crate::agent::models::selected_catalog_key_for_spawn(&models, &session_model_id)
|
||||||
|
.as_deref(),
|
||||||
|
&sampling_config.model,
|
||||||
|
)
|
||||||
|
};
|
||||||
let session_pruning_config: crate::config::PruningConfig = memory_config.as_ref().map_or_else(
|
let session_pruning_config: crate::config::PruningConfig = memory_config.as_ref().map_or_else(
|
||||||
|| crate::config::PruningConfig {
|
|| crate::config::PruningConfig {
|
||||||
enabled: false,
|
enabled: false,
|
||||||
@@ -618,16 +630,11 @@ pub(crate) async fn spawn_session_actor(
|
|||||||
watcher,
|
watcher,
|
||||||
stale_claim_secs: watcher_config.stale_claim_secs,
|
stale_claim_secs: watcher_config.stale_claim_secs,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: api_key_provider.clone(),
|
embedding_credentials: crate::auth::credential_provider::embedding_session_credentials(
|
||||||
auth_credentials: auth_manager.as_ref().map(|am| {
|
&embed_base_url,
|
||||||
std::sync::Arc::new(
|
embed_platform,
|
||||||
crate::auth::credential_provider::ShellAuthCredentialProvider::new(
|
&models_manager.credential_authority(),
|
||||||
am.clone(),
|
),
|
||||||
None,
|
|
||||||
None,
|
|
||||||
),
|
|
||||||
) as std::sync::Arc<dyn kigi_auth::AuthCredentialProvider>
|
|
||||||
}),
|
|
||||||
};
|
};
|
||||||
let backend = crate::session::memory::MemoryBackendImpl::from_session_params(
|
let backend = crate::session::memory::MemoryBackendImpl::from_session_params(
|
||||||
storage.clone(),
|
storage.clone(),
|
||||||
@@ -1287,8 +1294,11 @@ pub(crate) async fn spawn_session_actor(
|
|||||||
.map(|mc| mc.embedding.clone())
|
.map(|mc| mc.embedding.clone())
|
||||||
.unwrap_or_default();
|
.unwrap_or_default();
|
||||||
let embed_dims = embed_config.dimensions;
|
let embed_dims = embed_config.dimensions;
|
||||||
let sampling_base_url = embed_base_url.clone();
|
// The session's own params, so the background reindex embeds through
|
||||||
let sampling_api_key = embed_api_key.clone();
|
// the same endpoint-scoped credential as every foreground path — a
|
||||||
|
// second locally-built provider would re-derive credentials outside
|
||||||
|
// the chokepoint and would carry a static key with no 401 refresh.
|
||||||
|
let reindex_params = session.memory.backend_params.clone();
|
||||||
let session_id_for_reindex = session_info.id.to_string();
|
let session_id_for_reindex = session_info.id.to_string();
|
||||||
let chunks_added_counter = session.memory.chunks_added.clone();
|
let chunks_added_counter = session.memory.chunks_added.clone();
|
||||||
tokio::task::spawn_local(async move {
|
tokio::task::spawn_local(async move {
|
||||||
@@ -1311,13 +1321,8 @@ pub(crate) async fn spawn_session_actor(
|
|||||||
target : kigi_log::memory_log::TARGET, files = files.len(),
|
target : kigi_log::memory_log::TARGET, files = files.len(),
|
||||||
"MEMORY_REINDEX: background reindex complete"
|
"MEMORY_REINDEX: background reindex complete"
|
||||||
);
|
);
|
||||||
if let Some(api_key) = sampling_api_key
|
if let Some(ref params) = reindex_params
|
||||||
&& let Some(provider) =
|
&& let Some(provider) = params.make_embedding_provider().await
|
||||||
crate::session::memory::embedding::ApiEmbeddingProvider::from_session(
|
|
||||||
&embed_config,
|
|
||||||
sampling_base_url,
|
|
||||||
api_key,
|
|
||||||
)
|
|
||||||
{
|
{
|
||||||
crate::session::memory::embed_missing_chunks(&index, &provider).await;
|
crate::session::memory::embed_missing_chunks(&index, &provider).await;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1546,6 +1546,18 @@ impl SessionActor {
|
|||||||
vec![],
|
vec![],
|
||||||
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(_) => (
|
ToolInput::EnterPlanMode(_) => (
|
||||||
"Plan: Enter".to_string(),
|
"Plan: Enter".to_string(),
|
||||||
acp::ToolKind::Other,
|
acp::ToolKind::Other,
|
||||||
|
|||||||
+1
-2
@@ -555,8 +555,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
};
|
};
|
||||||
let (event_tx, _event_rx) = tokio::sync::mpsc::unbounded_channel::<SessionEvent>();
|
let (event_tx, _event_rx) = tokio::sync::mpsc::unbounded_channel::<SessionEvent>();
|
||||||
let actor = Arc::new(SessionActor {
|
let actor = Arc::new(SessionActor {
|
||||||
|
|||||||
+2
-4
@@ -393,8 +393,7 @@ fn initial_injection_backend_params_use_override_min_score() {
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
};
|
};
|
||||||
let initial_injection = crate::config::MemoryInitialInjectionConfig {
|
let initial_injection = crate::config::MemoryInitialInjectionConfig {
|
||||||
enabled: true,
|
enabled: true,
|
||||||
@@ -422,8 +421,7 @@ fn initial_injection_backend_params_preserve_default_zero_min_score() {
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
};
|
};
|
||||||
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
|
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
|
||||||
¶ms,
|
¶ms,
|
||||||
|
|||||||
@@ -18,8 +18,7 @@ fn initial_injection_backend_params_use_override_min_score() {
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
};
|
};
|
||||||
let initial_injection = crate::config::MemoryInitialInjectionConfig {
|
let initial_injection = crate::config::MemoryInitialInjectionConfig {
|
||||||
enabled: true,
|
enabled: true,
|
||||||
@@ -47,8 +46,7 @@ fn initial_injection_backend_params_preserve_default_zero_min_score() {
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
};
|
};
|
||||||
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
|
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
|
||||||
¶ms,
|
¶ms,
|
||||||
@@ -524,8 +522,7 @@ async fn create_injection_ready_actor(
|
|||||||
watcher: None,
|
watcher: None,
|
||||||
stale_claim_secs: 60,
|
stale_claim_secs: 60,
|
||||||
search_source: "tool",
|
search_source: "tool",
|
||||||
api_key_provider: None,
|
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
|
||||||
auth_credentials: None,
|
|
||||||
});
|
});
|
||||||
actor
|
actor
|
||||||
.chat_state_handle
|
.chat_state_handle
|
||||||
|
|||||||
@@ -36,7 +36,7 @@
|
|||||||
],
|
],
|
||||||
"definitions": {
|
"definitions": {
|
||||||
"ToolKind": {
|
"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"
|
"type": "string"
|
||||||
},
|
},
|
||||||
"ToolNamespace": {
|
"ToolNamespace": {
|
||||||
|
|||||||
@@ -148,7 +148,6 @@ mod linux {
|
|||||||
}
|
}
|
||||||
let err = std::io::Error::last_os_error();
|
let err = std::io::Error::last_os_error();
|
||||||
match err.kind() {
|
match err.kind() {
|
||||||
// The fd is non-blocking: the queue is empty.
|
|
||||||
std::io::ErrorKind::WouldBlock => break,
|
std::io::ErrorKind::WouldBlock => break,
|
||||||
std::io::ErrorKind::Interrupted => continue,
|
std::io::ErrorKind::Interrupted => continue,
|
||||||
_ => return Err(err),
|
_ => return Err(err),
|
||||||
|
|||||||
@@ -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
|
//! Tool implementations built on the `NewTool` trait; the sibling
|
||||||
//! `implementations/<tool>/` modules hold the `Tool`-trait counterparts.
|
//! `implementations/<tool>/` modules hold the `Tool`-trait counterparts.
|
||||||
|
pub mod agent_swarm;
|
||||||
pub mod ask_user_question;
|
pub mod ask_user_question;
|
||||||
pub mod bash;
|
pub mod bash;
|
||||||
#[path = "deploy_app_stub.rs"]
|
#[path = "deploy_app_stub.rs"]
|
||||||
@@ -21,6 +22,7 @@ pub mod todo;
|
|||||||
pub mod update_goal;
|
pub mod update_goal;
|
||||||
pub mod web_fetch;
|
pub mod web_fetch;
|
||||||
pub mod web_search;
|
pub mod web_search;
|
||||||
|
pub use agent_swarm::AgentSwarmTool;
|
||||||
pub use ask_user_question::AskUserQuestionTool;
|
pub use ask_user_question::AskUserQuestionTool;
|
||||||
pub use bash::BashTool;
|
pub use bash::BashTool;
|
||||||
pub use deploy_app::{AppBuilderDeployerConfig, DEPLOY_APP_TOOL_NAME};
|
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) {
|
pub fn prune_orphaned_background_task_tools(config: &mut crate::registry::types::ToolServerConfig) {
|
||||||
use crate::types::tool::ToolKind;
|
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
|
let has_task_tool = config
|
||||||
.tools
|
.tools
|
||||||
.iter()
|
.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);
|
let has_background_capable_bash = config.tools.iter().any(is_background_capable_bash_tool);
|
||||||
if has_task_tool || has_background_capable_bash {
|
if has_task_tool || has_background_capable_bash {
|
||||||
return;
|
return;
|
||||||
@@ -209,6 +211,7 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode {
|
|||||||
ToolKind::BackgroundTaskAction,
|
ToolKind::BackgroundTaskAction,
|
||||||
ToolKind::KillTaskAction,
|
ToolKind::KillTaskAction,
|
||||||
ToolKind::Task,
|
ToolKind::Task,
|
||||||
|
ToolKind::AgentSwarm,
|
||||||
ToolKind::EnterPlan,
|
ToolKind::EnterPlan,
|
||||||
ToolKind::ExitPlan,
|
ToolKind::ExitPlan,
|
||||||
ToolKind::AskUser,
|
ToolKind::AskUser,
|
||||||
@@ -232,6 +235,7 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode {
|
|||||||
ToolKind::BackgroundTaskAction,
|
ToolKind::BackgroundTaskAction,
|
||||||
ToolKind::KillTaskAction,
|
ToolKind::KillTaskAction,
|
||||||
ToolKind::Task,
|
ToolKind::Task,
|
||||||
|
ToolKind::AgentSwarm,
|
||||||
ToolKind::EnterPlan,
|
ToolKind::EnterPlan,
|
||||||
ToolKind::ExitPlan,
|
ToolKind::ExitPlan,
|
||||||
ToolKind::AskUser,
|
ToolKind::AskUser,
|
||||||
@@ -252,6 +256,7 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode {
|
|||||||
ToolKind::BackgroundTaskAction,
|
ToolKind::BackgroundTaskAction,
|
||||||
ToolKind::KillTaskAction,
|
ToolKind::KillTaskAction,
|
||||||
ToolKind::Task,
|
ToolKind::Task,
|
||||||
|
ToolKind::AgentSwarm,
|
||||||
ToolKind::EnterPlan,
|
ToolKind::EnterPlan,
|
||||||
ToolKind::ExitPlan,
|
ToolKind::ExitPlan,
|
||||||
ToolKind::AskUser,
|
ToolKind::AskUser,
|
||||||
@@ -276,6 +281,7 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode {
|
|||||||
ToolKind::BackgroundTaskAction,
|
ToolKind::BackgroundTaskAction,
|
||||||
ToolKind::KillTaskAction,
|
ToolKind::KillTaskAction,
|
||||||
ToolKind::Task,
|
ToolKind::Task,
|
||||||
|
ToolKind::AgentSwarm,
|
||||||
ToolKind::EnterPlan,
|
ToolKind::EnterPlan,
|
||||||
ToolKind::ExitPlan,
|
ToolKind::ExitPlan,
|
||||||
ToolKind::AskUser,
|
ToolKind::AskUser,
|
||||||
@@ -315,6 +321,12 @@ pub struct SubagentResult {
|
|||||||
/// `get_command_or_subagent_output`), so the tool returns a `task_id` notice
|
/// `get_command_or_subagent_output`), so the tool returns a `task_id` notice
|
||||||
/// instead of a completion. Never set for natively backgrounded subagents.
|
/// instead of a completion. Never set for natively backgrounded subagents.
|
||||||
pub backgrounded: bool,
|
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 {
|
impl Default for SubagentResult {
|
||||||
@@ -324,6 +336,7 @@ impl Default for SubagentResult {
|
|||||||
output: Arc::from(""),
|
output: Arc::from(""),
|
||||||
error: None,
|
error: None,
|
||||||
cancelled: false,
|
cancelled: false,
|
||||||
|
rate_limited: false,
|
||||||
subagent_id: String::new(),
|
subagent_id: String::new(),
|
||||||
child_session_id: String::new(),
|
child_session_id: String::new(),
|
||||||
tool_calls: 0,
|
tool_calls: 0,
|
||||||
|
|||||||
@@ -15,9 +15,9 @@ pub mod use_tool;
|
|||||||
pub mod web_search;
|
pub mod web_search;
|
||||||
pub use kigi::bash::{BashError, BashToolInput};
|
pub use kigi::bash::{BashError, BashToolInput};
|
||||||
pub use kigi::{
|
pub use kigi::{
|
||||||
AskUserQuestionTool, BashTool, EnterPlanModeTool, ExitPlanModeTool, GrepTool, KillTaskTool,
|
AgentSwarmTool, AskUserQuestionTool, BashTool, EnterPlanModeTool, ExitPlanModeTool, GrepTool,
|
||||||
ListDirTool, ReadFileTool, SearchReplaceTool, TaskOutputTool, TaskTool, TodoWriteTool,
|
KillTaskTool, ListDirTool, ReadFileTool, SearchReplaceTool, TaskOutputTool, TaskTool,
|
||||||
WaitTasksTool, WebFetchTool, WebSearchTool,
|
TodoWriteTool, WaitTasksTool, WebFetchTool, WebSearchTool,
|
||||||
};
|
};
|
||||||
pub use memory::{MemoryGetImpl, MemorySearchImpl};
|
pub use memory::{MemoryGetImpl, MemorySearchImpl};
|
||||||
pub use opencode::{
|
pub use opencode::{
|
||||||
|
|||||||
@@ -103,6 +103,7 @@ pub fn canonical_input(input: &ToolInput) -> Option<serde_json::Value> {
|
|||||||
| ToolInput::WaitTasks(_)
|
| ToolInput::WaitTasks(_)
|
||||||
| ToolInput::KillTask(_)
|
| ToolInput::KillTask(_)
|
||||||
| ToolInput::Task(_)
|
| ToolInput::Task(_)
|
||||||
|
| ToolInput::AgentSwarm(_)
|
||||||
| ToolInput::WebSearch(_)
|
| ToolInput::WebSearch(_)
|
||||||
| ToolInput::WebFetch(_)
|
| ToolInput::WebFetch(_)
|
||||||
| ToolInput::ApplyPatch(_)
|
| ToolInput::ApplyPatch(_)
|
||||||
|
|||||||
@@ -650,6 +650,7 @@ impl ToolRegistryBuilder {
|
|||||||
b.register::<kigi::GetTerminalCommandOutputTool>();
|
b.register::<kigi::GetTerminalCommandOutputTool>();
|
||||||
b.register::<kigi::WaitTasksTool>();
|
b.register::<kigi::WaitTasksTool>();
|
||||||
b.register::<kigi::TaskTool>();
|
b.register::<kigi::TaskTool>();
|
||||||
|
b.register::<kigi::AgentSwarmTool>();
|
||||||
b.register::<kigi::WebSearchTool>();
|
b.register::<kigi::WebSearchTool>();
|
||||||
b.register_with_params::<kigi::WebFetchTool, kigi::web_fetch::WebFetchParams>();
|
b.register_with_params::<kigi::WebFetchTool, kigi::web_fetch::WebFetchParams>();
|
||||||
b.register::<kigi::LspTool>();
|
b.register::<kigi::LspTool>();
|
||||||
|
|||||||
@@ -56,6 +56,7 @@ impl ToolKind {
|
|||||||
ToolKind::MemorySearch => "Memory Search",
|
ToolKind::MemorySearch => "Memory Search",
|
||||||
ToolKind::MemoryGet => "Memory Read",
|
ToolKind::MemoryGet => "Memory Read",
|
||||||
ToolKind::Task => "Subagent",
|
ToolKind::Task => "Subagent",
|
||||||
|
ToolKind::AgentSwarm => "Subagent Swarm",
|
||||||
ToolKind::EnterPlan => "Enter Plan Mode",
|
ToolKind::EnterPlan => "Enter Plan Mode",
|
||||||
ToolKind::ExitPlan => "Exit Plan Mode",
|
ToolKind::ExitPlan => "Exit Plan Mode",
|
||||||
ToolKind::AskUser => "Ask User",
|
ToolKind::AskUser => "Ask User",
|
||||||
@@ -96,6 +97,7 @@ impl ToolKind {
|
|||||||
| ToolKind::KillTaskAction
|
| ToolKind::KillTaskAction
|
||||||
| ToolKind::Skill
|
| ToolKind::Skill
|
||||||
| ToolKind::Task
|
| ToolKind::Task
|
||||||
|
| ToolKind::AgentSwarm
|
||||||
| ToolKind::DeployApp
|
| ToolKind::DeployApp
|
||||||
| ToolKind::SearchTool
|
| ToolKind::SearchTool
|
||||||
| ToolKind::UseTool
|
| ToolKind::UseTool
|
||||||
|
|||||||
@@ -88,6 +88,11 @@ pub enum ToolKind {
|
|||||||
MemorySearch,
|
MemorySearch,
|
||||||
MemoryGet,
|
MemoryGet,
|
||||||
Task,
|
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,
|
EnterPlan,
|
||||||
ExitPlan,
|
ExitPlan,
|
||||||
AskUser,
|
AskUser,
|
||||||
|
|||||||
@@ -32,6 +32,7 @@ use crate::implementations::opencode::write::WriteInput;
|
|||||||
use crate::implementations::search_tool::SearchToolInput;
|
use crate::implementations::search_tool::SearchToolInput;
|
||||||
use crate::implementations::skills::skill::SkillInput;
|
use crate::implementations::skills::skill::SkillInput;
|
||||||
use crate::implementations::use_tool::UseToolInput;
|
use crate::implementations::use_tool::UseToolInput;
|
||||||
|
use kigi_tool_types::AgentSwarmToolInput;
|
||||||
use kigi_tool_types::KillTaskToolInput;
|
use kigi_tool_types::KillTaskToolInput;
|
||||||
use kigi_tool_types::TaskOutputToolInput;
|
use kigi_tool_types::TaskOutputToolInput;
|
||||||
use kigi_tool_types::TaskToolInput;
|
use kigi_tool_types::TaskToolInput;
|
||||||
@@ -67,6 +68,7 @@ pub enum ToolInput {
|
|||||||
WaitTasks(WaitTasksToolInput),
|
WaitTasks(WaitTasksToolInput),
|
||||||
KillTask(KillTaskToolInput),
|
KillTask(KillTaskToolInput),
|
||||||
Task(TaskToolInput),
|
Task(TaskToolInput),
|
||||||
|
AgentSwarm(AgentSwarmToolInput),
|
||||||
WebSearch(WebSearchInput),
|
WebSearch(WebSearchInput),
|
||||||
WebFetch(WebFetchInput),
|
WebFetch(WebFetchInput),
|
||||||
Write(WriteInput),
|
Write(WriteInput),
|
||||||
|
|||||||
@@ -2157,7 +2157,10 @@ fn extract_variant(tc: &acp::ToolCall) -> Option<&str> {
|
|||||||
}
|
}
|
||||||
/// Twin without the optional-toolset spelling.
|
/// Twin without the optional-toolset spelling.
|
||||||
fn is_task_variant(variant: Option<&str>) -> bool {
|
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.
|
/// Twin without the optional-toolset spelling.
|
||||||
fn is_write_variant(variant: Option<&str>) -> bool {
|
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
|
/// SubagentSpawned notification) provides better visibility. Covers the
|
||||||
/// `task` / `Task` / `spawn_subagent` ids and Task-family variant tags.
|
/// `task` / `Task` / `spawn_subagent` ids and Task-family variant tags.
|
||||||
fn is_task_tool(tc: &acp::ToolCall) -> bool {
|
fn is_task_tool(tc: &acp::ToolCall) -> bool {
|
||||||
matches!(tc.title.as_str(), "task" | "Task" | "spawn_subagent")
|
matches!(
|
||||||
|| is_task_variant(extract_variant(tc))
|
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).
|
/// Check if a tool call is a scheduler tool (scheduler_create/delete/list).
|
||||||
///
|
///
|
||||||
|
|||||||
@@ -2333,11 +2333,7 @@ async fn drain_and_process(
|
|||||||
}
|
}
|
||||||
|
|
||||||
// On terminals without bracketed paste, try to capture more events
|
// On terminals without bracketed paste, try to capture more events
|
||||||
// that may still be in transit from the input reader thread. A batch
|
// that may still be in transit from the input reader thread.
|
||||||
// that already holds a key burst skips the short detection window:
|
|
||||||
// its 2 ms budget is below one Windows scheduler quantum, so a
|
|
||||||
// mid-paste gap would end collection before it starts and the paste
|
|
||||||
// tail would arrive as a separate fragment.
|
|
||||||
if should_extend_for_paste(&raw_events)
|
if should_extend_for_paste(&raw_events)
|
||||||
&& (is_paste_burst(&raw_events) || detect_paste(&mut raw_events, input_rx).await)
|
&& (is_paste_burst(&raw_events) || detect_paste(&mut raw_events, input_rx).await)
|
||||||
{
|
{
|
||||||
@@ -2542,18 +2538,13 @@ async fn drain_and_process(
|
|||||||
const PASTE_DETECT_TIMEOUT: Duration = Duration::from_millis(2);
|
const PASTE_DETECT_TIMEOUT: Duration = Duration::from_millis(2);
|
||||||
|
|
||||||
/// Timeout for subsequent rounds once paste has been detected. Must exceed
|
/// Timeout for subsequent rounds once paste has been detected. Must exceed
|
||||||
/// one Windows scheduler quantum (~15.6 ms): ConPTY delivers a paste as a
|
/// one Windows scheduler quantum (~15.6 ms), or a routine mid-burst gap
|
||||||
/// per-character key burst, and a mid-burst gap of one quantum is routine.
|
/// splits the paste and repaste-to-expand duplicates instead of expanding.
|
||||||
/// A smaller window ends collection mid-paste; the tail then arrives as a
|
|
||||||
/// second synthetic paste, whose content can never match the paste chip, so
|
|
||||||
/// repaste-to-expand inserts a duplicate instead of expanding.
|
|
||||||
const PASTE_CONTINUE_TIMEOUT: Duration = Duration::from_millis(25);
|
const PASTE_CONTINUE_TIMEOUT: Duration = Duration::from_millis(25);
|
||||||
|
|
||||||
/// Safety cap on events accumulated in one extension pass. On the
|
/// Safety cap on events accumulated in one extension pass. Pastes arrive
|
||||||
/// no-bracketed-paste path a paste is one key event per character, so the
|
/// one key event per character here, so the cap must exceed any real paste;
|
||||||
/// cap must exceed any real paste (~200 KB of text); hitting it splits the
|
/// hitting it splits the paste like a timeout miss would.
|
||||||
/// paste with the same duplicate-on-repaste effect as a timeout split. The
|
|
||||||
/// idle timeout above, not this cap, is the normal terminator.
|
|
||||||
const PASTE_EXTEND_MAX_EVENTS: usize = 200_000;
|
const PASTE_EXTEND_MAX_EVENTS: usize = 200_000;
|
||||||
|
|
||||||
/// Returns `true` when the batch contains pasteable key events but no
|
/// Returns `true` when the batch contains pasteable key events but no
|
||||||
@@ -2564,12 +2555,11 @@ fn should_extend_for_paste(events: &[Event]) -> bool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Minimum pasteable key events already in one batch to classify it as a
|
/// Minimum pasteable key events already in one batch to classify it as a
|
||||||
/// paste burst in progress. Human typing and key auto-repeat deliver one
|
/// paste burst; typing and auto-repeat deliver only a couple per batch.
|
||||||
/// or two events per batch; only a paste chunk lands more at once.
|
|
||||||
const PASTE_BURST_BATCH_EVENTS: usize = 8;
|
const PASTE_BURST_BATCH_EVENTS: usize = 8;
|
||||||
|
|
||||||
/// True when the batch alone proves a paste is in progress, without
|
/// True when the batch alone proves a paste is in progress, so collection
|
||||||
/// waiting on [`PASTE_DETECT_TIMEOUT`].
|
/// can start without the [`PASTE_DETECT_TIMEOUT`] round-trip.
|
||||||
fn is_paste_burst(events: &[Event]) -> bool {
|
fn is_paste_burst(events: &[Event]) -> bool {
|
||||||
events.iter().filter(|e| is_pasteable_key_event(e)).count() >= PASTE_BURST_BATCH_EVENTS
|
events.iter().filter(|e| is_pasteable_key_event(e)).count() >= PASTE_BURST_BATCH_EVENTS
|
||||||
}
|
}
|
||||||
@@ -3188,11 +3178,9 @@ mod tests {
|
|||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn typing_and_release_storms_are_not_a_paste_burst() {
|
fn typing_and_release_storms_are_not_a_paste_burst() {
|
||||||
// A short typed run stays below the burst threshold.
|
|
||||||
let typed = vec![press(KeyCode::Char('h')), press(KeyCode::Char('i'))];
|
let typed = vec![press(KeyCode::Char('h')), press(KeyCode::Char('i'))];
|
||||||
assert!(!is_paste_burst(&typed));
|
assert!(!is_paste_burst(&typed));
|
||||||
|
|
||||||
// Release events carry no content and must not count toward it.
|
|
||||||
let releases: Vec<Event> = "releases only!"
|
let releases: Vec<Event> = "releases only!"
|
||||||
.chars()
|
.chars()
|
||||||
.map(|c| release(KeyCode::Char(c)))
|
.map(|c| release(KeyCode::Char(c)))
|
||||||
@@ -3200,10 +3188,8 @@ mod tests {
|
|||||||
assert!(!is_paste_burst(&releases));
|
assert!(!is_paste_burst(&releases));
|
||||||
}
|
}
|
||||||
|
|
||||||
/// A key-burst paste whose tail trickles in with sub-window gaps must
|
/// A burst tail trickling in with one-quantum gaps must reassemble into
|
||||||
/// reassemble into ONE synthetic paste. A split here is what made
|
/// ONE paste — a split is what duplicated chips on Windows.
|
||||||
/// repaste-to-expand insert a duplicate chip on Windows: the fragment
|
|
||||||
/// could never byte-equal the chip content.
|
|
||||||
#[tokio::test(start_paused = true)]
|
#[tokio::test(start_paused = true)]
|
||||||
async fn burst_tail_with_scheduler_gaps_reassembles_into_one_paste() {
|
async fn burst_tail_with_scheduler_gaps_reassembles_into_one_paste() {
|
||||||
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
|
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
|
||||||
@@ -3213,8 +3199,6 @@ mod tests {
|
|||||||
let mut batch: Vec<Event> = head.chars().map(char_or_enter).collect();
|
let mut batch: Vec<Event> = head.chars().map(char_or_enter).collect();
|
||||||
assert!(is_paste_burst(&batch), "head chunk must classify as burst");
|
assert!(is_paste_burst(&batch), "head chunk must classify as burst");
|
||||||
|
|
||||||
// Trickle the tail with 15 ms gaps — one Windows scheduler quantum,
|
|
||||||
// larger than the old 10 ms window, smaller than the current one.
|
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
for c in tail.chars() {
|
for c in tail.chars() {
|
||||||
tokio::time::sleep(Duration::from_millis(15)).await;
|
tokio::time::sleep(Duration::from_millis(15)).await;
|
||||||
|
|||||||
@@ -91,6 +91,7 @@ pub(crate) const ALL_TOOL_KINDS: &[ToolKind] = &[
|
|||||||
ToolKind::MemorySearch,
|
ToolKind::MemorySearch,
|
||||||
ToolKind::MemoryGet,
|
ToolKind::MemoryGet,
|
||||||
ToolKind::Task,
|
ToolKind::Task,
|
||||||
|
ToolKind::AgentSwarm,
|
||||||
ToolKind::EnterPlan,
|
ToolKind::EnterPlan,
|
||||||
ToolKind::ExitPlan,
|
ToolKind::ExitPlan,
|
||||||
ToolKind::AskUser,
|
ToolKind::AskUser,
|
||||||
@@ -144,8 +145,9 @@ pub(crate) fn kind_allowed(mode: CapabilityMode, kind: ToolKind) -> bool {
|
|||||||
// Bash / shell.
|
// Bash / shell.
|
||||||
Execute => matches!(mode, M::Execute),
|
Execute => matches!(mode, M::Execute),
|
||||||
|
|
||||||
// Process control (background tasks, monitors).
|
// Process control (background tasks, monitors). A swarm is a fan-out of
|
||||||
BackgroundTaskAction | WaitTasksAction | KillTaskAction | Task | Monitor => {
|
// subagent spawns, so it gates exactly as `Task` does.
|
||||||
|
BackgroundTaskAction | WaitTasksAction | KillTaskAction | Task | AgentSwarm | Monitor => {
|
||||||
matches!(mode, M::Execute)
|
matches!(mode, M::Execute)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,116 @@
|
|||||||
|
//! Input/output types for the `agent_swarm` tool — one prompt template
|
||||||
|
//! expanded over a list of items into a fleet of subagents.
|
||||||
|
|
||||||
|
use schemars::JsonSchema;
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
/// The literal a `prompt_template` must contain; each expansion substitutes
|
||||||
|
/// one `items` entry for it.
|
||||||
|
pub const PROMPT_TEMPLATE_PLACEHOLDER: &str = "{{item}}";
|
||||||
|
|
||||||
|
/// Upper bound on members in one call, counting resumes.
|
||||||
|
pub const MAX_AGENT_SWARM_MEMBERS: usize = 128;
|
||||||
|
|
||||||
|
/// Input for the `agent_swarm` tool.
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
|
||||||
|
pub struct AgentSwarmToolInput {
|
||||||
|
#[schemars(description = "Short description of what the whole swarm is doing (3-7 words).")]
|
||||||
|
pub description: String,
|
||||||
|
|
||||||
|
/// Subagent type every item-spawned member runs as.
|
||||||
|
#[schemars(
|
||||||
|
description = "Name of the subagent type every member runs as. Built-in types: \"general-purpose\", \"explore\", \"plan\"."
|
||||||
|
)]
|
||||||
|
#[serde(default = "default_subagent_type")]
|
||||||
|
pub subagent_type: String,
|
||||||
|
|
||||||
|
/// Prompt shared by every member; must contain `{{item}}`.
|
||||||
|
#[schemars(
|
||||||
|
description = "Prompt shared by every member. Must contain the literal {{item}}, which is replaced by each entry of `items`. Required whenever `items` is given."
|
||||||
|
)]
|
||||||
|
#[serde(default)]
|
||||||
|
pub prompt_template: Option<String>,
|
||||||
|
|
||||||
|
/// The work units. Each expands `prompt_template` into one member.
|
||||||
|
#[schemars(
|
||||||
|
description = "One entry per member: each is substituted into `prompt_template`. Give every member a distinct scope so members never edit the same file. At least 2 entries unless `resume_agent_ids` is used."
|
||||||
|
)]
|
||||||
|
#[serde(default)]
|
||||||
|
pub items: Vec<String>,
|
||||||
|
|
||||||
|
/// Continue named subagents from a previous swarm: id → follow-up prompt.
|
||||||
|
#[schemars(
|
||||||
|
description = "Continue previously spawned subagents: a map of agent_id (from an earlier agent_swarm result) to the follow-up prompt for that member."
|
||||||
|
)]
|
||||||
|
#[serde(default)]
|
||||||
|
pub resume_agent_ids: std::collections::BTreeMap<String, String>,
|
||||||
|
|
||||||
|
/// Model slug every member runs on; omitted inherits the caller's.
|
||||||
|
#[schemars(
|
||||||
|
description = "Model every member runs on. Omit to inherit the caller's current model."
|
||||||
|
)]
|
||||||
|
#[serde(default)]
|
||||||
|
pub model: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn default_subagent_type() -> String {
|
||||||
|
"general-purpose".to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// How a member's run ended.
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||||
|
#[serde(rename_all = "lowercase")]
|
||||||
|
pub enum SwarmMemberOutcome {
|
||||||
|
Completed,
|
||||||
|
Failed,
|
||||||
|
Aborted,
|
||||||
|
/// Still running: it outlived the foreground await budget and the subagent
|
||||||
|
/// coordinator detached it. Distinct from `Failed` because the member is
|
||||||
|
/// alive and still writing — relaunching its item would put a second agent
|
||||||
|
/// on the same files, and resuming it is refused while it runs.
|
||||||
|
Backgrounded,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SwarmMemberOutcome {
|
||||||
|
pub fn as_str(self) -> &'static str {
|
||||||
|
match self {
|
||||||
|
Self::Completed => "completed",
|
||||||
|
Self::Failed => "failed",
|
||||||
|
Self::Aborted => "aborted",
|
||||||
|
Self::Backgrounded => "backgrounded",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Whether re-running this member's work is safe to suggest.
|
||||||
|
fn is_resumable(self) -> bool {
|
||||||
|
matches!(self, Self::Failed | Self::Aborted)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// One member's contribution to the aggregate result.
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct SwarmMemberResult {
|
||||||
|
/// The `items` entry (or the resumed agent id) this member was given —
|
||||||
|
/// what the caller needs to retry exactly the members that did not finish.
|
||||||
|
pub item: String,
|
||||||
|
/// Present once the member started; absent means it never launched.
|
||||||
|
pub agent_id: Option<String>,
|
||||||
|
pub resumed: bool,
|
||||||
|
pub outcome: SwarmMemberOutcome,
|
||||||
|
pub summary: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl SwarmMemberResult {
|
||||||
|
/// Whether the member ever reached the backend. A member that never
|
||||||
|
/// started has nothing to resume.
|
||||||
|
pub fn started(&self) -> bool {
|
||||||
|
self.agent_id.is_some()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Whether the caller may be told to continue this member. A member that
|
||||||
|
/// is still running must not be offered: the coordinator refuses to resume
|
||||||
|
/// a live subagent, and relaunching its item duplicates its writes.
|
||||||
|
pub fn is_resumable(&self) -> bool {
|
||||||
|
self.started() && self.outcome.is_resumable()
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -1,10 +1,15 @@
|
|||||||
//! Canonical, extensible tool types.
|
//! Canonical, extensible tool types.
|
||||||
|
mod agent_swarm;
|
||||||
mod ext;
|
mod ext;
|
||||||
mod schema_utils;
|
mod schema_utils;
|
||||||
pub mod serde_lenient;
|
pub mod serde_lenient;
|
||||||
mod task;
|
mod task;
|
||||||
mod types;
|
mod types;
|
||||||
|
|
||||||
|
pub use agent_swarm::{
|
||||||
|
AgentSwarmToolInput, MAX_AGENT_SWARM_MEMBERS, PROMPT_TEMPLATE_PLACEHOLDER, SwarmMemberOutcome,
|
||||||
|
SwarmMemberResult,
|
||||||
|
};
|
||||||
pub use ext::Extensions;
|
pub use ext::Extensions;
|
||||||
pub use schema_utils::parse_arguments_from_schema_lossy;
|
pub use schema_utils::parse_arguments_from_schema_lossy;
|
||||||
pub use serde_lenient::{
|
pub use serde_lenient::{
|
||||||
|
|||||||
Reference in New Issue
Block a user