ZacharyZhang-NY/Kigi-CLI
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Compare commits

base: ZacharyZhang-NY/Kigi-CLI:v0.1.9
Branches Tags
ZacharyZhang-NY/Kigi-CLI:main
ZacharyZhang-NY/Kigi-CLI:v0.1.11 ZacharyZhang-NY/Kigi-CLI:v0.1.10 ZacharyZhang-NY/Kigi-CLI:v0.1.9 ZacharyZhang-NY/Kigi-CLI:v0.1.8 ZacharyZhang-NY/Kigi-CLI:v0.1.7 ZacharyZhang-NY/Kigi-CLI:v0.1.6 ZacharyZhang-NY/Kigi-CLI:v0.1.5 ZacharyZhang-NY/Kigi-CLI:v0.1.4 ZacharyZhang-NY/Kigi-CLI:v0.1.3 ZacharyZhang-NY/Kigi-CLI:v0.1.2 ZacharyZhang-NY/Kigi-CLI:v0.1.1 ZacharyZhang-NY/Kigi-CLI:v0.1.0
...
compare: ZacharyZhang-NY/Kigi-CLI:main
Branches Tags
ZacharyZhang-NY/Kigi-CLI:main
ZacharyZhang-NY/Kigi-CLI:v0.1.11 ZacharyZhang-NY/Kigi-CLI:v0.1.10 ZacharyZhang-NY/Kigi-CLI:v0.1.9 ZacharyZhang-NY/Kigi-CLI:v0.1.8 ZacharyZhang-NY/Kigi-CLI:v0.1.7 ZacharyZhang-NY/Kigi-CLI:v0.1.6 ZacharyZhang-NY/Kigi-CLI:v0.1.5 ZacharyZhang-NY/Kigi-CLI:v0.1.4 ZacharyZhang-NY/Kigi-CLI:v0.1.3 ZacharyZhang-NY/Kigi-CLI:v0.1.2 ZacharyZhang-NY/Kigi-CLI:v0.1.1 ZacharyZhang-NY/Kigi-CLI:v0.1.0
13 Commits
v0.1.9 ... main
Author SHA1 Message Date
ZacharyZhang-NY 553ec1c620 chore(release): bump workspace version to 0.1.11
Warm build cache / warm (aarch64-apple-darwin) (push) Waiting to run
Details
Warm build cache / warm (x86_64-apple-darwin) (push) Waiting to run
Details
Warm build cache / warm (aarch64-unknown-linux-gnu) (push) Waiting to run
Details
Warm build cache / warm (x86_64-pc-windows-msvc) (push) Waiting to run
Details
Release / build (aarch64-apple-darwin) (push) Waiting to run
Details
Release / build (x86_64-apple-darwin) (push) Waiting to run
Details
Release / build (aarch64-unknown-linux-gnu) (push) Waiting to run
Details
Release / build (x86_64-pc-windows-msvc) (push) Waiting to run
Details
Release / publish GitHub Release (push) Blocked by required conditions
Details
Warm build cache / warm (x86_64-unknown-linux-gnu) (push) Failing after 1m36s
Details
Release / build (x86_64-unknown-linux-gnu) (push) Failing after 6s
Details 
All 62 crates inherit `workspace.package.version`, so the single edit
plus the refreshed lockfile carries the release. The `v0.1.11` tag is
what triggers the build; release.yml gates the tag against this value.

Ships the sessionless-model-switch fix (f656627): `/model` from a
non-project launch directory now starts the session its deferred switch
drains into instead of dangling silently.
 2026-07-27 16:04:51 -04:00
ZacharyZhang-NY f656627c10 fix(tui): start the session a sessionless model switch defers into 
Launching kigi in a non-project directory (~/Downloads, ~, /tmp — anywhere
`is_project_dir` rejects) leaves the agent view session-less behind the
project-picker question, which only a PLAIN prompt can open. Slash
commands still execute there, so `/model <name> <effort>` stashed its
switch in `deferred_model_switch` — a stash that assumes a create is in
flight — and with none pending it dangled forever with zero feedback:
the picker rendered, the user chose model and effort, and nothing
changed. Not a 0.1.10 regression: 0.1.9 pty-reproduces identically; the
report correlated with the upgrade only because every earlier launch
happened to be from a project directory.

Apply the QueueCommand precedent ("queued slash work bypasses the
picker, so create the deferred session or it never drains") to both
deferral sites — Action::SwitchModel in the router and
set_default_model's no-session branch — via
skip_picker_and_create_session, whose in-flight guard already makes it a
no-op while a create is pending, so the racing-create case is unchanged.
SessionCreated then applies the stash through the existing
apply_deferred_model_switch path; pty-verified end-to-end from
~/Downloads (session created, model changed, effort applied).
 2026-07-27 15:55:57 -04:00
ZacharyZhang-NY 11c3ca9803 chore(release): bump workspace version to 0.1.10
Release / build (aarch64-apple-darwin) (push) Waiting to run
Details
Release / build (x86_64-apple-darwin) (push) Waiting to run
Details
Release / build (aarch64-unknown-linux-gnu) (push) Waiting to run
Details
Release / build (x86_64-pc-windows-msvc) (push) Waiting to run
Details
Release / publish GitHub Release (push) Blocked by required conditions
Details
Release / build (x86_64-unknown-linux-gnu) (push) Failing after 7s
Details 
All 62 crates inherit `workspace.package.version`, so the single edit
plus the refreshed lockfile carries the release. The `v0.1.10` tag is
what triggers the build; release.yml gates the tag against this value.

Also records the shared-SSRF-policy invariant in AGENTS.md, so a future
port does not reintroduce a second copy. The zero-egress entry is
compressed by the same number of lines to keep the file at 500.
 2026-07-27 13:26:25 -04:00
ZacharyZhang-NY bac8470c80 fix(hooks): share one SSRF policy instead of a second, stale copy 
The hook runner carried its own `is_blocked_ip`, a line-for-line copy of
the pre-hardening web_fetch predicate: loopback allowed unconditionally,
and no TEST-NET, 198.18/15, 240/4, 0.0.0.0/8, multicast, IPv6 site-local
or embedded-v4 wrapper coverage. Hook URLs come from settings, and
project settings from an untrusted repo are loaded today, so the gap is
reachable.

`kigi-hooks` already depends on `kigi-tools`, so the fix is to delete the
copy and call the shared predicate — no new crate, and no third
implementation to drift.

`allow_local` is on for hooks: a loopback hook receiver is a legitimate
local setup, which the old copy also allowed. It is now allowed only
when the URL names the host literally, so a public name resolving to
loopback is refused.

Range coverage now lives with the policy; the runner's tests pin what it
adds on top. The URL-scrubbing regression test moves off TEST-NET-1,
which the policy now blocks, onto a closed loopback port — faster, and
no real egress from a test.
 2026-07-27 13:01:50 -04:00
ZacharyZhang-NY 74402f3078 fix(pager): stop a sibling test flattening the diff-band assertion 
`committed_edit_keeps_diff_line_backgrounds` passed alone and
single-threaded but failed under parallel execution, so the workspace
gate was red.

`terminal_native_lock_paints_only_native_colors` engages the
process-global terminal-native lock. While it is held `Theme::current()`
returns `terminal_default()`, whose `diff_*_bg` are all `Color::Reset`,
so no band is painted at all — and `current_kind()` reports a nominal
`KigiNight`, which is what made the failure read as a theme mismatch.

Two changes, both at the cause:
- hold the shared theme lock via `theme_cache::pin_theme()`, the helper
  written for this and until now unused, so the lock cannot be engaged
  mid-test
- assert the invariant (two distinct non-Reset bands) instead of exact
  RGB: the entry line cache is keyed on the GLOBAL theme kind, so an
  exact-color assertion races by construction. Span-level colors stay
  pinned by the `tool::edit` tests that own them

The failure message now reports the theme, height, and the bands it
actually found, which is how the cause was located.
 2026-07-27 12:38:14 -04:00
ZacharyZhang-NY 649c5f5641 fix(models): drop the ultra reasoning effort, which no backend accepts 
The hardcoded Codex catalog advertised `ultra` for gpt-5.6-sol and
gpt-5.6-terra, so selecting it sent `reasoning.effort = "ultra"` and the
Responses endpoint answered 400 with its real menu, which ends at `max`.

The tier was kigi's invention: the enum documented it as codex-only, the
catalog was its only producer, and the upstream model snapshot never
listed it. Removing the variant as well as the catalog entries means
`/effort ultra` is rejected at the TUI instead of parsing and failing on
the wire; a persisted `ultra` degrades to the model default through
`lenient_reasoning_effort_opt`, as designed for vocabulary changes.
 2026-07-27 12:12:14 -04:00
ZacharyZhang-NY ad3840f9ec fix(web_fetch): block non-public targets by default and gate every hop 
kigi allowed loopback unconditionally and missed several non-public
ranges, and the SSRF check ran only on the initial URL.

Policy (ssrf.rs):
- loopback is blocked unless `[toolset.web_fetch] allow_local` (or
  KIGI_WEB_FETCH_ALLOW_LOCAL) is on, AND the URL names it explicitly,
  so a public name resolving to loopback stays blocked (DNS rebinding)
- add 0.0.0.0/8, 100.64/10, 192.0.0.0/24, TEST-NET-1/2/3, 198.18/15,
  240/4, IPv6 site-local and documentation prefixes
- inherit the IPv4 verdict through mapped, compatible, NAT64 and 6to4
  wrappers; network-specific NAT64 prefixes remain uncovered (see doc)

Plumbing (client.rs), where the exploitable half lived:
- re-check every redirect hop, not just the first
- compare hosts exactly; a `www` sibling has its own A records, so it
  is a cross-host redirect rather than an auto-followed hop
- run the check before the fetch service, so a blocked URL is never
  posted to an endpoint that egresses elsewhere
- exempt explicit local hosts from the https upgrade and from the
  single-label filter, and re-upgrade each followed hop

Wiring: allow_local reaches WebFetchParams from both construction
paths; documented in the config guide and the README env table.
 2026-07-27 11:44:20 -04:00
ZacharyZhang-NY ac23ebc9a1 feat(turn): halt a turn stuck repeating one identical tool call 
Ports grok-build's action stationarity breaker. A polling loop looks like
progress from inside the turn — every step succeeds — so nothing stops it
and it burns the whole budget re-reading one file. Nudge at 8 identical
batches, halt at 16, halt a pure no-op run at 4.

Three defects were found and fixed before this landed, none of which the
gates caught:

- The check runs AFTER execute_tool_calls, not before. Before it, the
  assistant's tool_calls are recorded with no results, so pushing the
  nudge triggers repair_dangling_tool_calls: synthetic "cancelled by the
  user" results get spliced in, the real results land after a user
  message, and dedup misses them — the orphan tool_result shape that
  draws a provider 400. It is Ok-gated too, since one `?` inside can
  leave a call resultless.
- `TurnOutcome::StationarityHalted` is its own variant grouped with
  Completed, not a reuse of Cancelled. As a cancellation it would report
  StopReason::Cancelled, fire the abort lifecycle, kill the turn's
  subagents, grow the goal harness back-off streak until three halts
  paused a running goal, and let completion-requirement recovery re-run
  the very turn halted for looping.
- The no-op gate is size-based, not name-based. kigi registers the shell
  tool as `run_terminal_cmd` and renames it `run_terminal_command`, so
  matching on "bash" silently disabled the tighter ceiling while every
  test still passed.
 2026-07-27 10:37:41 -04:00
ZacharyZhang-NY a07d889ad9 fix(swarm): label each member with its own item 
Every member was launched with the swarm's description, so the pager drew
N identical subagent blocks: the user could not tell which member was
running, and a failure could not be attributed to an item.

This is also why K6 does not port upstream's progress grid. A member IS
an ordinary subagent, so the pager already renders one block per member
with live status, duration, tool-call and turn counts — richer than a
grid cell. The gap was the label, not the widget; upstream's grid (and
the estimator behind it, which infers progress from tool-call rate)
answers a different TUI's shortcomings. The parent tool call already
reads "<description> (N members)" and `is_task_variant` already knows
`AgentSwarm`, so the blocking-wait spinner registers.
 2026-07-27 03:01:42 -04:00
ZacharyZhang-NY 16328e55f9 feat(swarm): /swarm mode — a standing instruction to fan work out 
The `agent_swarm` tool works with the mode off; what the mode adds is the
doctrine — decompose finely, give every member a disjoint scope, do not
do the work yourself. `/swarm`, `/swarm on|off`, `/swarm <task>`, gated
on the tool actually being in the toolset.

Two triggers, not upstream's three. Upstream's `tool` trigger exists so
invoking the tool makes the doctrine appear; kigi's tool carries its own
description and works with the mode off, so that trigger would add a
state only the tool can reach. `Manual` persists until switched off,
`Task` expires at turn end — both are real user intents.

Load-bearing decisions, each from a defect found in review:

- Expiry is a DROP guard, not a post-loop call. A user interrupt aborts
  the turn future rather than cancelling it, so post-loop code never
  runs; three `?` paths skip it too. Either way a `/swarm <task>` mode
  would leak into the next, unrelated prompt.
- `enter` is a total no-op while armed, so the `/swarm <task>` shorthand
  cannot downgrade a standing `/swarm on` into a per-turn mode that then
  disarms itself. This matches upstream; the first cut diverged, and the
  test asserting the divergence was inverted with the fix.
- An explicit `/swarm off` retracts UNCONDITIONALLY. The doctrine rides
  the conversation and so survives a resume, a fork and a compaction that
  an in-memory flag does not; trusting the flag left a session reading
  "off" with the instruction still steering and no way to clear it. That
  also retires the `reminder_live` field, whose doc comment claimed to
  know something the code cannot.
- The mode is deliberately NOT persisted, unlike /goal and /graph: those
  strand real work when lost, this is a prompt hint whose recovery is
  retyping one command. Recorded in AGENTS.md so the omission reads as a
  decision, not a gap.
- The pre-session gate is `subagents_enabled`, not a hardcoded `true`:
  the builder strips `agent_swarm` wherever it strips `task`, and
  advertising it then offers a menu entry that resolves to literal text.
- The exit reminder is as emphatic as the doctrine it revokes; a single
  weak clause is the likelier of the two to be summarised away.

The doctrine rides the existing `push_system_reminder` channel rather
than a second injection path of its own.
 2026-07-27 02:39:55 -04:00
ZacharyZhang-NY 9edb8729ef feat(swarm): agent_swarm — one prompt over many items, paced as a fleet 
Ports kimi-code's AgentSwarm: a `prompt_template` containing `{{item}}`
expanded over an `items` list into up to 128 subagents, run to completion
and returned as one aggregate. Kigi already exceeds upstream on planning,
verification, isolation and merge via /graph; what it lacked was cheap
immediate fan-out. Entirely client-side — no new backend surface.

The engine is three pure pieces plus a runner: `plan` validates and
expands (every fault reported before a single member starts — a
half-launched swarm is expensive to unwind), `schedule` is the launch
ramp as testable arithmetic, `run` drives it against the existing
`SubagentBackend`. Reuses the single-spawn coordinator rather than
inventing a batch API.

Load-bearing decisions, each the result of a defect found in review:

- `backgrounded` is its own outcome. A member that outlives the 600s
  foreground budget is detached by the coordinator and KEEPS RUNNING;
  reporting it as failed invites the model to relaunch its item, putting
  a second agent on the same files. It is never offered for resume.
- `InFlightGuard` cancels live members on Drop. Send-now cancels the turn
  WITHOUT cancelling subagents and aborts the task; the dropped receivers
  read as "parent gone" and each child re-attaches itself. There is no
  cooperative path to use instead — `Cancellation` is constructed nowhere
  in the tree — so Drop is the only seam that fires.
- Retries and wall clock are both bounded. The swarm blocks the caller's
  turn, so every wait needs a ceiling it cannot argue past; stragglers at
  the deadline are reported as still-running, with their ids.
- `ToolKind::AgentSwarm` is its own variant: `TemplateRenderer`'s
  `by_kind` map holds one tool name per kind, so sharing `Task` would
  silently redirect `${{ tools.by_kind.task }}` in other tools' prompts.
- An explicitly requested model that cannot be validated is refused, as
  the task tool already does — one loud error beats `items.len()` quiet
  ones. Depth stays capped at 1: upstream's unlimited nesting is a
  hazard, not a feature.
- `SubagentResult.rate_limited` is classified where the typed ACP error
  code is still in hand; a scheduler re-deriving it from a formatted
  string would stop adapting the day the wording changed.
- Aggregate output is clamped per member (head+tail, loss stated):
  native tool output is truncated nowhere downstream.

42 agent_swarm tests. The fake backend awaits, so the concurrency and
ordering assertions can actually fail; the cap test also proves the
fixture can exceed the cap.
 2026-07-27 01:22:52 -04:00
ZacharyZhang-NY ed8049cf77 fix(memory): scope embedding credentials to the endpoint that may receive them 
`MemoryBackendParams` carried the primary session `AuthManager` and the
session api-key provider unconditionally, while `embed_base_url` is the
CURRENT MODEL's endpoint and `AuthRetryMiddleware` stamps `Authorization`
on every request it wraps. A user who enables `[memory.embedding] model`
while running a BYOK or subscription-OAuth model therefore sent the Kimi
session bearer to that third party.

`EndpointScopedCredentials` binds the credential to the one endpoint it
may reach: `for_endpoint` drops the handle unless the caller vouches for
the URL, and `approved_for` re-checks at provider-build time in release
too, because `MemoryBackendParams` is `Clone` and callers rewrite fields
on the copy.

The shell decides through `CredentialAuthority::manager_for` rather than
a second URL predicate — it answers both whether a credential may ride
and which manager governs it, so a subscription-OAuth platform gets its
own pooled manager. The session's `SharedApiKeyProvider` is not
forwarded at all: it is hard-wired to the primary manager, so at a
pooled platform's host it would resolve the wrong bearer. A platform's
own `embed_api_key` is untouched and keeps serving its own endpoint.

The background reindex built a second provider straight from
`ApiEmbeddingProvider::from_session`, outside the chokepoint and without
401 refresh; it now embeds through the session's own params.

Test strength verified by mutation: with the guard reverted, exactly
`session_credentials_are_withheld_from_a_foreign_endpoint` and
`a_cloned_param_set_cannot_redirect_scoped_credentials` fail.
 2026-07-26 23:41:45 -04:00
ZacharyZhang-NY 867b3e110b docs(comments): trim paste-fix commentary to the crucial constraints 2026-07-24 13:09:35 -04:00
70 changed files with 4207 additions and 549 deletions
Expand all files Collapse all files
AGENTS.md
+58 -8
View File
@@ -16,14 +16,14 @@ import) or any `KIMI_*` env var.
## Hard constraints
- **Zero egress**: outbound connections are limited to
`auth.kimi.com`, `api.kimi.com`, `api.moonshot.cn`, `api.moonshot.ai`,
GitHub Releases domains, user-configured MCP servers, the endpoints of
provider platforms the user has credentialed, and `models.dev` (model
metadata refresh — reached ONLY when an enabled platform's `/models` wire
lacks metadata, `wire_serves_metadata=false`; Kimi/Moonshot never trigger
it; `KIGI_MODELS_DEV_URL=0` disables). No telemetry, no analytics, ever.
`crates/codegen/kigi-env` is the single home of first-party endpoints.
- **Zero egress**: only `auth.kimi.com`, `api.kimi.com`, `api.moonshot.cn`,
`api.moonshot.ai`, GitHub Releases, user-configured MCP servers, credentialed
provider endpoints, and `models.dev` (metadata refresh, ONLY when an enabled
platform's `/models` lacks it — `wire_serves_metadata=false`, never
Kimi/Moonshot, `KIGI_MODELS_DEV_URL=0` disables). No telemetry, no analytics,
ever. `crates/codegen/kigi-env` is the single home of first-party endpoints.
- **One SSRF policy**: `ssrf::is_blocked_for_host` gates every model/settings
URL (`web_fetch` + hooks), per redirect hop. Never write a second copy.
- **Toolchain**: Rust 1.97.0 (rust-toolchain.toml), edition 2024.
- **Gates** (all must stay green):
`cargo check --workspace --all-targets`,
@@ -164,6 +164,56 @@ edges stay deterministic Rust. The harness appends a terminal
the replan cap; `{"ops": []}` is a respected free no-op; failures
degrade.
## Swarm (`agent_swarm` tool + `/swarm` mode)
Cheap, unstructured fan-out — the one thing `/graph` does not offer. One
`prompt_template` containing `{{item}}` is expanded over an `items` list
into up to `MAX_AGENT_SWARM_MEMBERS` (128) ordinary subagents, run to
completion, and returned as ONE aggregate. Ported from kimi-code; entirely
client-side, no backend surface.
- Engine (`kigi-tools/.../kigi/agent_swarm/`): `plan.rs` validates and
expands (every fault is reported BEFORE a member starts; expanded
prompts must be pairwise distinct), `schedule.rs` is the launch ramp as
pure arithmetic (5 immediate, then 1 per 700ms; capacity shrinks on a
provider rate limit and recovers after a quiet window;
`KIGI_AGENT_SWARM_MAX_CONCURRENCY` caps it and a malformed value is a
hard error), `run.rs` drives it against the EXISTING single-spawn
`SubagentBackend`, `tool.rs` is the tool.
- `SwarmMemberOutcome::Backgrounded` is load-bearing: a member that
outlives the 600s foreground budget is detached by the coordinator and
KEEPS RUNNING. It is never reported `Failed` and never offered for
resume — relaunching its item would put a second agent on the same
files. Members share the caller's tree with NO isolation; the
distinct-prompt rule is the only thing keeping them apart.
- `InFlightGuard` (run.rs) cancels live members when the runner's future
is dropped. Send-now cancels the turn WITHOUT cancelling subagents and
aborts the task; the dropped receivers read as "parent gone" and each
child re-attaches itself. There is no cooperative path to use instead —
`kigi_tool_runtime::Cancellation` is constructed nowhere.
- Both retries (`MAX_RATE_LIMIT_RETRIES`) and wall clock
(`MAX_SWARM_RUNTIME`) are bounded: the swarm blocks the caller's turn,
so every wait needs a ceiling. Stragglers at the deadline are reported
as still-running WITH their ids.
- `ToolKind::AgentSwarm` is its own variant because `TemplateRenderer`'s
`by_kind` map holds ONE tool name per kind — sharing `Task` would
silently redirect `${{ tools.by_kind.task }}` in other tools' prompts.
It gates exactly as `Task` does (`capability.rs`), and `builder.rs`
strips it wherever `task` is stripped. `MAX_SUBAGENT_DEPTH` stays 1:
upstream's unlimited nesting is a hazard, not a feature.
- `/swarm | /swarm on|off | /swarm <task>` (`BuiltinGate::Swarm` = the
tool is in the toolset) arms a doctrine reminder via the existing
`push_system_reminder`. Two triggers only: `Manual` persists until
switched off, `Task` auto-expires at turn end via `SwarmTurnGuard` — a
DROP guard, because a user interrupt aborts the turn future and never
reaches post-loop code. `enter` is a total no-op while armed so the
`/swarm <task>` shorthand cannot downgrade a standing `/swarm on`.
- The mode is deliberately NOT persisted (unlike `/goal` and `/graph`,
which strand real work if lost): it is a prompt hint whose recovery is
retyping one command. The injected doctrine IS durable, so an explicit
`/swarm off` retracts UNCONDITIONALLY — a resumed or compacted session
can read as "off" with the instruction still in context.
## Provider registry & API-key auth (post-0.1.3 expansion)
- The platform registry is compiled-in spec rows in `kigi-models`
Cargo.lock
Generated
+62 -62
View File
@@ -5442,7 +5442,7 @@ dependencies = [
[[package]]
name = "kigi-acp-lib"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"agent-client-protocol",
"async-trait",
@@ -5456,7 +5456,7 @@ dependencies = [
[[package]]
name = "kigi-agent"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"chrono",
"dirs 6.0.0",
@@ -5486,7 +5486,7 @@ dependencies = [
[[package]]
name = "kigi-agent-lifecycle"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"async-trait",
"tokio",
@@ -5495,7 +5495,7 @@ dependencies = [
[[package]]
name = "kigi-auth"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"async-trait",
"http 1.4.2",
@@ -5508,7 +5508,7 @@ dependencies = [
[[package]]
name = "kigi-bin"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anyhow",
"clap",
@@ -5543,7 +5543,7 @@ dependencies = [
[[package]]
name = "kigi-chat-state"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"indexmap",
"kigi-compaction",
@@ -5560,7 +5560,7 @@ dependencies = [
[[package]]
name = "kigi-codebase-graph"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"ahash",
"clap",
@@ -5596,7 +5596,7 @@ dependencies = [
[[package]]
name = "kigi-compaction"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anyhow",
"async-trait",
@@ -5609,7 +5609,7 @@ dependencies = [
[[package]]
name = "kigi-config"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"base64",
"blake3",
@@ -5632,7 +5632,7 @@ dependencies = [
[[package]]
name = "kigi-config-types"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"agent-client-protocol",
"indexmap",
@@ -5646,7 +5646,7 @@ dependencies = [
[[package]]
name = "kigi-crash-handler"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"backtrace",
"libc",
@@ -5657,7 +5657,7 @@ dependencies = [
[[package]]
name = "kigi-env"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"tracing",
"url",
@@ -5665,7 +5665,7 @@ dependencies = [
[[package]]
name = "kigi-fast-worktree"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anyhow",
"bytes",
@@ -5697,7 +5697,7 @@ dependencies = [
[[package]]
name = "kigi-file-utils"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anyhow",
"aws-config",
@@ -5721,7 +5721,7 @@ dependencies = [
[[package]]
name = "kigi-fsnotify"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"criterion",
"dunce",
@@ -5742,7 +5742,7 @@ dependencies = [
[[package]]
name = "kigi-gix-status"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"gix",
"kigi-test-utils",
@@ -5752,7 +5752,7 @@ dependencies = [
[[package]]
name = "kigi-hooks"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"fastrand",
"kigi-config",
@@ -5771,7 +5771,7 @@ dependencies = [
[[package]]
name = "kigi-hooks-plugins-types"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"serde",
"serde_json",
@@ -5779,7 +5779,7 @@ dependencies = [
[[package]]
name = "kigi-http"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"kigi-auth",
"kigi-log",
@@ -5794,7 +5794,7 @@ dependencies = [
[[package]]
name = "kigi-hunk-tracker"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"chrono",
"dunce",
@@ -5815,14 +5815,14 @@ dependencies = [
[[package]]
name = "kigi-interjection-core"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"serde",
]
[[package]]
name = "kigi-log"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anyhow",
"chrono",
@@ -5840,7 +5840,7 @@ dependencies = [
[[package]]
name = "kigi-markdown"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anstyle",
"anstyle-lossy",
@@ -5864,14 +5864,14 @@ dependencies = [
[[package]]
name = "kigi-markdown-core"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"pulldown-cmark",
]
[[package]]
name = "kigi-mcp"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"agent-client-protocol",
"async-trait",
@@ -5908,7 +5908,7 @@ dependencies = [
[[package]]
name = "kigi-memory"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anyhow",
"arc-swap",
@@ -5942,7 +5942,7 @@ dependencies = [
[[package]]
name = "kigi-mermaid"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"fontdb",
"image",
@@ -5960,7 +5960,7 @@ dependencies = [
[[package]]
name = "kigi-models"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"kigi-env",
"serde",
@@ -5970,7 +5970,7 @@ dependencies = [
[[package]]
name = "kigi-pager-minimal"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"chrono",
"crossterm",
@@ -5987,7 +5987,7 @@ dependencies = [
[[package]]
name = "kigi-pager-pty-harness"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"alacritty_terminal",
"anyhow",
@@ -6012,7 +6012,7 @@ dependencies = [
[[package]]
name = "kigi-pager-render"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"agent-client-protocol",
"anstyle",
@@ -6064,7 +6064,7 @@ dependencies = [
[[package]]
name = "kigi-paths"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"camino",
"serde",
@@ -6074,7 +6074,7 @@ dependencies = [
[[package]]
name = "kigi-prompt-queue"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"serde",
"serde_json",
@@ -6082,7 +6082,7 @@ dependencies = [
[[package]]
name = "kigi-proto-build"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anyhow",
"pbjson-build",
@@ -6093,7 +6093,7 @@ dependencies = [
[[package]]
name = "kigi-ratatui-inline"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"ansi-width",
"anstyle-parse 0.2.7",
@@ -6110,7 +6110,7 @@ dependencies = [
[[package]]
name = "kigi-ratatui-textarea"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"arboard",
"chrono",
@@ -6131,7 +6131,7 @@ dependencies = [
[[package]]
name = "kigi-sampler"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"async-openai",
"async-stream",
@@ -6154,7 +6154,7 @@ dependencies = [
[[package]]
name = "kigi-sampling-types"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"assert_matches",
"async-openai",
@@ -6171,7 +6171,7 @@ dependencies = [
[[package]]
name = "kigi-sandbox"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anyhow",
"chrono",
@@ -6192,7 +6192,7 @@ dependencies = [
[[package]]
name = "kigi-secrets"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"regex",
"serde_json",
@@ -6230,7 +6230,7 @@ dependencies = [
[[package]]
name = "kigi-shell"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"agent-client-protocol",
"anyhow",
@@ -6367,7 +6367,7 @@ dependencies = [
[[package]]
name = "kigi-shell-base"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anyhow",
"chrono",
@@ -6392,7 +6392,7 @@ dependencies = [
[[package]]
name = "kigi-sqlite-journal"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"libc",
"rusqlite",
@@ -6403,7 +6403,7 @@ dependencies = [
[[package]]
name = "kigi-subagent-resolution"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"kigi-sampling-types",
"kigi-tool-types",
@@ -6418,7 +6418,7 @@ dependencies = [
[[package]]
name = "kigi-system-power"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"windows-sys 0.59.0",
"zbus",
@@ -6426,7 +6426,7 @@ dependencies = [
[[package]]
name = "kigi-test-support"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"agent-client-protocol",
"anyhow",
@@ -6448,7 +6448,7 @@ dependencies = [
[[package]]
name = "kigi-test-utils"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"runfiles",
"tracing",
@@ -6457,11 +6457,11 @@ dependencies = [
[[package]]
name = "kigi-token-estimation"
version = "0.1.9"
version = "0.1.11"
[[package]]
name = "kigi-tool-protocol"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"kigi-tool-types",
"serde",
@@ -6472,7 +6472,7 @@ dependencies = [
[[package]]
name = "kigi-tool-runtime"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anyhow",
"async-trait",
@@ -6490,7 +6490,7 @@ dependencies = [
[[package]]
name = "kigi-tool-types"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"minijinja",
"schemars 1.2.1",
@@ -6500,7 +6500,7 @@ dependencies = [
[[package]]
name = "kigi-tools"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anyhow",
"arc-swap",
@@ -6577,7 +6577,7 @@ dependencies = [
[[package]]
name = "kigi-tools-api"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"kigi-proto-build",
"kigi-tool-protocol",
@@ -6590,11 +6590,11 @@ dependencies = [
[[package]]
name = "kigi-tracing-macros"
version = "0.1.9"
version = "0.1.11"
[[package]]
name = "kigi-tty-utils"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"libc",
"nix 0.30.1",
@@ -6604,7 +6604,7 @@ dependencies = [
[[package]]
name = "kigi-tui"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"agent-client-protocol",
"ansi-to-tui",
@@ -6691,7 +6691,7 @@ dependencies = [
[[package]]
name = "kigi-update"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anyhow",
"dunce",
@@ -6720,14 +6720,14 @@ dependencies = [
[[package]]
name = "kigi-version"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"semver",
]
[[package]]
name = "kigi-workspace"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"agent-client-protocol",
"anyhow",
@@ -6806,7 +6806,7 @@ dependencies = [
[[package]]
name = "kigi-workspace-types"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"base64",
"chrono",
@@ -8840,7 +8840,7 @@ dependencies = [
[[package]]
name = "ptyctl"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"alacritty_terminal",
"anyhow",
@@ -8858,7 +8858,7 @@ dependencies = [
[[package]]
name = "ptyctl-cli"
version = "0.1.9"
version = "0.1.11"
dependencies = [
"anyhow",
"axum",
Cargo.toml
+1 -1
View File
@@ -76,7 +76,7 @@ members = [
]
[workspace.package]
version = "0.1.9"
version = "0.1.11"
edition = "2024"
license = "Apache-2.0"
crates/codegen/kigi-agent/src/builder.rs
+33 -4
View File
@@ -711,9 +711,19 @@ impl AgentBuilder {
kigi_tools::types::tool::ToolNamespace::Kigi,
"task"
);
// The swarm is a fan-out of subagent spawns, so it lives and dies with
// the task tool: any condition that leaves no subagent to spawn leaves
// the swarm with nothing to fan out to.
let swarm_tool_id = format!(
"{}:{}",
kigi_tools::types::tool::ToolNamespace::Kigi,
"agent_swarm"
);
let mut task_stripped = false;
if !self.subagents_enabled {
tool_config.tools.retain(|tc| tc.id != task_tool_id);
tool_config
.tools
.retain(|tc| tc.id != task_tool_id && tc.id != swarm_tool_id);
task_stripped = true;
} else {
let subagents = crate::discovery::all_subagents_with_plugins(
@@ -722,7 +732,9 @@ impl AgentBuilder {
self.plugin_registry.as_deref(),
);
if subagents.is_empty() {
tool_config.tools.retain(|tc| tc.id != task_tool_id);
tool_config
.tools
.retain(|tc| tc.id != task_tool_id && tc.id != swarm_tool_id);
task_stripped = true;
} else if self.prompt_audience == crate::prompt::context::PromptAudience::Subagent {
if let Some(task_tc) = tool_config
@@ -810,7 +822,13 @@ impl AgentBuilder {
.tools
.iter()
.any(|t| AGENT_TASK_CLASSIFIER_RE.is_match(t));
let task_deps = ["task", "get_task_output", "kill_task", "wait_tasks"];
let task_deps = [
"task",
"agent_swarm",
"get_task_output",
"kill_task",
"wait_tasks",
];
let registered_tool_ids = tool_bridge_builder.known_tool_ids();
let present_kinds: std::collections::HashSet<ToolKind> =
tool_config.tools.iter().filter_map(|tc| tc.kind).collect();
@@ -925,7 +943,13 @@ impl AgentBuilder {
}
}
if definition.allowed_subagent_types.as_deref() == Some(&[]) {
let task_deps = ["task", "get_task_output", "kill_task", "wait_tasks"];
let task_deps = [
"task",
"agent_swarm",
"get_task_output",
"kill_task",
"wait_tasks",
];
tool_config
.tools
.retain(|tc| !task_deps.contains(&short_tool_name(&tc.id)));
@@ -1601,6 +1625,11 @@ mod tests {
has_task, *subagents,
"[{label}] spawn_subagent presence should match subagents_enabled={subagents}; got tools: {names:?}"
);
let has_swarm = names.contains(&"agent_swarm");
assert_eq!(
has_swarm, *subagents,
"[{label}] agent_swarm fans out to subagents, so it must follow the same gate as spawn_subagent; got tools: {names:?}"
);
assert!(
names.contains(&"enter_plan_mode"),
"[{label}] enter_plan_mode must always be present (TUI plan-mode keybind needs it); got tools: {names:?}"
crates/codegen/kigi-agent/src/config.rs
+10
View File
@@ -154,6 +154,11 @@ fn task_tool_config() -> ToolConfig {
.with_name("spawn_subagent")
.with_param_rename("run_in_background", "background")
}
/// Swarm tool. Keeps its registry name: unlike `task` it has no CLI-specific
/// alias, and the name is what the model is told to reuse for `resume_agent_ids`.
fn agent_swarm_tool_config() -> ToolConfig {
ToolConfig::from(&kigi::AgentSwarmTool)
}
/// Task output tool renamed for clarity:
/// `get_task_output` → `get_command_or_subagent_output`.
fn task_output_tool_config() -> ToolConfig {
@@ -270,6 +275,7 @@ fn default_kigi_toolset() -> ToolServerConfig {
task_output_tool_config(),
wait_tasks_tool_config(),
task_tool_config(),
agent_swarm_tool_config(),
(&kigi::SchedulerCreateTool).into(),
(&kigi::SchedulerDeleteTool).into(),
(&kigi::SchedulerListTool).into(),
@@ -318,6 +324,7 @@ pub fn kigi_hashline_toolset(
task_output_tool_config(),
wait_tasks_tool_config(),
task_tool_config(),
agent_swarm_tool_config(),
(&kigi::WebSearchTool).into(),
(&kigi::SchedulerCreateTool).into(),
(&kigi::SchedulerDeleteTool).into(),
@@ -400,6 +407,7 @@ fn kigi_plan_toolset() -> ToolServerConfig {
(&kigi::TodoWriteTool).into(),
task_output_tool_config(),
task_tool_config(),
agent_swarm_tool_config(),
(&kigi::SchedulerCreateTool).into(),
(&kigi::SchedulerDeleteTool).into(),
(&kigi::SchedulerListTool).into(),
@@ -428,6 +436,7 @@ fn orchestrator_toolset() -> ToolServerConfig {
(&kigi::ListDirTool).into(),
(&kigi::GrepTool).into(),
task_tool_config(),
agent_swarm_tool_config(),
task_output_tool_config(),
wait_tasks_tool_config(),
kill_task_tool_config(),
@@ -497,6 +506,7 @@ fn kigi_ask_user_toolset() -> ToolServerConfig {
task_output_tool_config(),
wait_tasks_tool_config(),
task_tool_config(),
agent_swarm_tool_config(),
(&kigi::SchedulerCreateTool).into(),
(&kigi::SchedulerDeleteTool).into(),
(&kigi::SchedulerListTool).into(),
crates/codegen/kigi-hooks/src/runner/http.rs
+56 -135
View File
@@ -26,68 +26,14 @@ struct HttpHookOutput {
reason: Option<String>,
}
/// CWE-918: Returns `true` if an IP address is in a private, link-local,
/// or cloud metadata range that should be blocked to prevent SSRF attacks.
/// CWE-918: whether this address is blocked for a hook.
///
/// Loopback (`127.x` / `::1`) is allowed for local development servers.
fn is_blocked_ip(ip: &IpAddr) -> bool {
match ip {
IpAddr::V4(v4) => {
let octets = v4.octets();
if octets[0] == 127 {
// loopback — allowed for local dev
return false;
}
if octets[0] == 10 {
// RFC 1918: 10.0.0.0/8
return true;
}
if octets[0] == 172 && (16..=31).contains(&octets[1]) {
// RFC 1918: 172.16.0.0/12
return true;
}
if octets[0] == 192 && octets[1] == 168 {
// RFC 1918: 192.168.0.0/16
return true;
}
if octets[0] == 169 && octets[1] == 254 {
// RFC 3927: 169.254.0.0/16 (link-local, cloud metadata)
return true;
}
if octets[0] == 100 && (64..=127).contains(&octets[1]) {
// RFC 6598: 100.64.0.0/10 (CGNAT)
return true;
}
if v4.is_unspecified() {
// 0.0.0.0
return true;
}
false
}
IpAddr::V6(v6) => {
if v6.is_loopback() {
// ::1 — allowed for local dev
return false;
}
if v6.is_unspecified() {
// ::
return true;
}
if let Some(v4) = v6.to_ipv4_mapped() {
return is_blocked_ip(&IpAddr::V4(v4));
}
let segments = v6.segments();
if segments[0] & 0xffc0 == 0xfe80 {
// fe80::/10 — link-local
return true;
}
if segments[0] & 0xfe00 == 0xfc00 {
// fc00::/7 — unique local (ULA)
return true;
}
false
}
}
/// Delegates to the `web_fetch` policy, so one implementation
/// governs every outbound URL. `allow_local` is on: a loopback hook
/// receiver is legitimate, but only when named literally, so
/// rebinding through a public name stays blocked.
fn is_blocked_ip(ip: &IpAddr, host: &str) -> bool {
kigi_tools::implementations::kigi::web_fetch::ssrf::is_blocked_for_host(ip, host, true)
}
/// CWE-918: Validate a hook URL to prevent SSRF.
@@ -112,7 +58,7 @@ async fn validate_hook_url(url: &str) -> Result<(), String> {
// If host is a literal IP, check it directly.
if let Ok(ip) = host.parse::<IpAddr>() {
if is_blocked_ip(&ip) {
if is_blocked_ip(&ip, host) {
return Err(format!("URL resolves to blocked private/internal IP: {ip}"));
}
return Ok(());
@@ -131,7 +77,7 @@ async fn validate_hook_url(url: &str) -> Result<(), String> {
}
for addr in &addrs {
if is_blocked_ip(&addr.ip()) {
if is_blocked_ip(&addr.ip(), host) {
return Err(format!(
"URL host {host} resolves to blocked private/internal IP: {}",
addr.ip()
@@ -571,76 +517,53 @@ mod tests {
}
}
// SSRF protection: is_blocked_ip tests
// SSRF protection: hook-side policy tests
//
// Range coverage lives with the shared predicate in `web_fetch::ssrf`;
// these pin what the runner adds on top.
/// Every range the shared policy knows is refused here.
#[test]
fn ssrf_blocks_rfc1918_10x() {
assert!(is_blocked_ip(&"10.0.0.1".parse().unwrap()));
assert!(is_blocked_ip(&"10.255.255.255".parse().unwrap()));
fn ssrf_delegates_to_the_shared_policy() {
for ip in [
"10.0.0.1",
"172.16.0.1",
"192.168.0.1",
"169.254.169.254",
"100.64.0.1",
"0.0.0.0",
"::",
"fe80::1",
"fc00::1",
"::ffff:10.0.0.1",
// Ranges the old hand-rolled copy missed entirely.
"198.18.0.1",
"192.0.2.1",
"203.0.113.1",
"240.0.0.1",
"64:ff9b::a9fe:a9fe",
] {
let ip: IpAddr = ip.parse().unwrap();
assert!(is_blocked_ip(&ip, &ip.to_string()), "{ip}");
}
for ip in ["1.1.1.1", "8.8.8.8", "172.32.0.1", "100.63.0.1"] {
let ip: IpAddr = ip.parse().unwrap();
assert!(!is_blocked_ip(&ip, &ip.to_string()), "{ip}");
}
}
/// A local hook receiver stays reachable when named literally.
#[test]
fn ssrf_blocks_rfc1918_172x() {
assert!(is_blocked_ip(&"172.16.0.1".parse().unwrap()));
assert!(is_blocked_ip(&"172.31.255.255".parse().unwrap()));
assert!(!is_blocked_ip(&"172.15.0.1".parse().unwrap()));
assert!(!is_blocked_ip(&"172.32.0.1".parse().unwrap()));
fn ssrf_allows_a_literal_loopback_hook_target() {
assert!(!is_blocked_ip(&"127.0.0.1".parse().unwrap(), "127.0.0.1"));
assert!(!is_blocked_ip(&"::1".parse().unwrap(), "localhost"));
}
/// A public name resolving to loopback is rebinding.
#[test]
fn ssrf_blocks_rfc1918_192168() {
assert!(is_blocked_ip(&"192.168.0.1".parse().unwrap()));
assert!(is_blocked_ip(&"192.168.255.255".parse().unwrap()));
}
#[test]
fn ssrf_blocks_link_local_metadata() {
assert!(is_blocked_ip(&"169.254.0.1".parse().unwrap()));
assert!(is_blocked_ip(&"169.254.169.254".parse().unwrap()));
}
#[test]
fn ssrf_blocks_cgnat() {
assert!(is_blocked_ip(&"100.64.0.1".parse().unwrap()));
assert!(is_blocked_ip(&"100.127.255.255".parse().unwrap()));
assert!(!is_blocked_ip(&"100.63.0.1".parse().unwrap()));
}
#[test]
fn ssrf_blocks_unspecified() {
assert!(is_blocked_ip(&"0.0.0.0".parse().unwrap()));
assert!(is_blocked_ip(&"::".parse().unwrap()));
}
#[test]
fn ssrf_allows_loopback() {
assert!(!is_blocked_ip(&"127.0.0.1".parse().unwrap()));
assert!(!is_blocked_ip(&"::1".parse().unwrap()));
}
#[test]
fn ssrf_allows_public_ips() {
assert!(!is_blocked_ip(&"1.1.1.1".parse().unwrap()));
assert!(!is_blocked_ip(&"8.8.8.8".parse().unwrap()));
}
#[test]
fn ssrf_blocks_ipv6_link_local() {
assert!(is_blocked_ip(&"fe80::1".parse().unwrap()));
}
#[test]
fn ssrf_blocks_ipv6_unique_local() {
assert!(is_blocked_ip(&"fc00::1".parse().unwrap()));
assert!(is_blocked_ip(&"fd00::1".parse().unwrap()));
}
#[test]
fn ssrf_blocks_ipv4_mapped_ipv6_private() {
assert!(is_blocked_ip(&"::ffff:10.0.0.1".parse::<IpAddr>().unwrap()));
assert!(is_blocked_ip(
&"::ffff:192.168.1.1".parse::<IpAddr>().unwrap()
));
fn ssrf_blocks_a_public_name_that_resolves_to_loopback() {
let ip: IpAddr = "127.0.0.1".parse().unwrap();
assert!(is_blocked_ip(&ip, "evil.example.com"));
}
// SSRF protection: validate_hook_url tests
@@ -839,14 +762,15 @@ mod tests {
/// the secret does NOT appear in the returned error message.
#[tokio::test]
async fn run_http_hook_scrubs_url_from_reqwest_error() {
// Use a TEST-NET-1 host (RFC 5737, "MUST NOT be used in
// public networks"). It is not RFC1918 so SSRF validation
// will let it through, but no real DNS or connection will
// succeed -- reqwest will surface a connection error whose
// default Display includes the URL.
// A closed loopback port: allowed as a literal local host,
// and refused at once. TEST-NET-1 is now blocked by policy.
let dead = {
let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
l.local_addr().unwrap()
};
let secret = "ghp_VERY_REAL_SECRET_TOKEN_42";
let mut extra_env = std::collections::HashMap::new();
extra_env.insert("RUNTIME_HOST".to_string(), "192.0.2.1".to_string());
extra_env.insert("RUNTIME_HOST".to_string(), dead.to_string());
extra_env.insert("MY_TOKEN".to_string(), secret.to_string());
let raw = "https://${RUNTIME_HOST}/check?token=${MY_TOKEN}";
@@ -928,10 +852,7 @@ mod tests {
// debugging). The wire-DTO consumer must prefer raw_url for
// display -- documented in the HttpInfo rustdoc.
let info = info.expect("HttpInfo should be present for connection failures too");
assert_eq!(
info.url,
"https://192.0.2.1/check?token=ghp_VERY_REAL_SECRET_TOKEN_42"
);
assert_eq!(info.url, format!("https://{dead}/check?token={secret}"));
assert_eq!(info.raw_url.as_deref(), Some(raw));
}
crates/codegen/kigi-memory/src/backend.rs
+203 -77
View File
@@ -18,6 +18,76 @@ use super::embedding::EmbeddingProvider as _;
use super::storage::MemoryStorage;
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.
///
/// Grouping these in one struct ensures every call site — ToolBridge, first-turn
@@ -30,7 +100,8 @@ pub struct MemoryBackendParams {
pub session_id: String,
/// Embedding provider config — `None` forces FTS-only fallback everywhere.
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,
/// API key for embedding API calls.
pub embed_api_key: Option<String>,
@@ -47,10 +118,8 @@ pub struct MemoryBackendParams {
/// - `"injection"` — first-turn memory context injection
/// - `"compaction_recovery"` — post-compaction context re-injection
pub search_source: &'static str,
/// Dynamic API key provider — when set, `make_embedding_provider()` resolves
/// the key per-call instead of using the static `embed_api_key`.
pub api_key_provider: Option<kigi_tools::types::SharedApiKeyProvider>,
pub auth_credentials: Option<Arc<dyn kigi_auth::AuthCredentialProvider>>,
/// The session credentials, and the single endpoint they may reach.
pub embedding_credentials: EndpointScopedCredentials,
}
impl MemoryBackendParams {
@@ -59,8 +128,7 @@ impl MemoryBackendParams {
pub async fn make_embedding_provider(&self) -> Option<super::embedding::ApiEmbeddingProvider> {
build_embedding_provider(
self.embed_config.as_ref(),
self.auth_credentials.as_ref(),
self.api_key_provider.as_ref(),
&self.embedding_credentials,
self.embed_api_key.as_deref(),
&self.embed_base_url,
)
@@ -70,8 +138,7 @@ impl MemoryBackendParams {
async fn build_embedding_provider(
config: Option<&kigi_config_types::MemoryEmbeddingConfig>,
auth_credentials: Option<&Arc<dyn kigi_auth::AuthCredentialProvider>>,
api_key_provider: Option<&kigi_tools::types::SharedApiKeyProvider>,
credentials: &EndpointScopedCredentials,
static_api_key: Option<&str>,
base_url: &str,
) -> Option<super::embedding::ApiEmbeddingProvider> {
@@ -80,9 +147,19 @@ async fn build_embedding_provider(
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
// 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());
return super::embedding::ApiEmbeddingProvider::from_config(
config,
@@ -91,14 +168,13 @@ async fn build_embedding_provider(
);
}
// Fallback: resolve API key per-call, wrap in a static middleware client
// (no 401 refresh, but auth header is still stamped by middleware).
let api_key = match api_key_provider {
Some(p) => p.current_api_key_async().await,
None => None,
}
.or_else(|| static_api_key.map(|s| s.to_owned()))?;
super::embedding::ApiEmbeddingProvider::from_session(config, base_url.to_owned(), api_key)
// The platform's own configured key, wrapped in a static middleware client
// (no 401 refresh, but the auth header is still stamped by middleware).
super::embedding::ApiEmbeddingProvider::from_session(
config,
base_url.to_owned(),
static_api_key?.to_owned(),
)
}
/// `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;
/// injection and compaction-recovery backends use their own local counters.
pub search_counter: std::sync::Arc<std::sync::atomic::AtomicU64>,
/// Dynamic API key provider for embedding requests.
api_key_provider: Option<kigi_tools::types::SharedApiKeyProvider>,
/// Refresh-capable credential provider for embedding HTTP middleware.
auth_credentials: Option<Arc<dyn kigi_auth::AuthCredentialProvider>>,
/// The session credentials, and the single endpoint they may reach.
embedding_credentials: EndpointScopedCredentials,
}
impl MemoryBackendImpl {
@@ -150,8 +224,7 @@ impl MemoryBackendImpl {
stale_claim_secs: 60,
session_id: String::new(),
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: EndpointScopedCredentials::none(),
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> {
build_embedding_provider(
self.embed_config.as_ref(),
self.auth_credentials.as_ref(),
self.api_key_provider.as_ref(),
&self.embedding_credentials,
self.embed_api_key.as_deref(),
&self.embed_base_url,
)
@@ -232,8 +304,7 @@ impl MemoryBackendImpl {
if let Some(w) = &params.watcher {
backend = backend.with_watcher(w.clone(), params.stale_claim_secs);
}
backend.api_key_provider = params.api_key_provider.clone();
backend.auth_credentials = params.auth_credentials.clone();
backend.embedding_credentials = params.embedding_credentials.clone();
backend
}
}
@@ -479,8 +550,7 @@ mod factory_tests {
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: EndpointScopedCredentials::none(),
}
}
@@ -1038,72 +1108,128 @@ mod factory_tests {
);
}
/// Regression: provider build must use `current_api_key_async`,
/// never sync. Prevents memory_search 401s on rotated tokens.
#[tokio::test]
async fn make_embedding_provider_uses_async_api_key_resolution() {
use kigi_tools::types::ApiKeyProvider;
use std::sync::atomic::{AtomicU32, Ordering};
struct AsyncProbe {
sync_calls: Arc<AtomicU32>,
async_calls: Arc<AtomicU32>,
struct ProbeCredentials;
impl kigi_auth::HttpAuth for ProbeCredentials {
fn apply(
&self,
builder: reqwest::RequestBuilder,
_base_url: &str,
) -> reqwest::RequestBuilder {
builder.bearer_auth("session-bearer")
}
impl ApiKeyProvider for AsyncProbe {
fn current_api_key(&self) -> Option<String> {
self.sync_calls.fetch_add(1, Ordering::SeqCst);
Some("sync-stale".into())
}
fn current_api_key_async(
&self,
) -> 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_trait::async_trait]
impl kigi_auth::AuthCredentialProvider for ProbeCredentials {
fn snapshot(&self) -> kigi_auth::CredentialSnapshot {
kigi_auth::CredentialSnapshot {
token: Some("session-bearer".into()),
..Default::default()
}
}
async fn refresh_after_unauthorized(&self) -> bool {
false
}
}
let sync_calls = Arc::new(AtomicU32::new(0));
let async_calls = Arc::new(AtomicU32::new(0));
let probe: kigi_tools::types::SharedApiKeyProvider = Arc::new(AsyncProbe {
sync_calls: sync_calls.clone(),
async_calls: async_calls.clone(),
});
const SESSION_ENDPOINT: &str = "https://api.kimi.com/coding/v1";
const FOREIGN_ENDPOINT: &str = "https://api.anthropic.com/v1";
let params = MemoryBackendParams {
fn params_at(base_url: &str, credentials: EndpointScopedCredentials) -> MemoryBackendParams {
MemoryBackendParams {
session_id: "s1".into(),
embed_config: Some(MemoryEmbeddingConfig {
model: Some("test-embed-model".into()),
..Default::default()
}),
embed_base_url: "http://example/v1".into(),
embed_api_key: Some("static-fallback".into()),
embed_base_url: base_url.into(),
embed_api_key: None,
search_config: MemorySearchConfig::default(),
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: Some(probe),
// No auth_credentials — forces the api_key_provider fallback path.
auth_credentials: None,
};
embedding_credentials: credentials,
}
}
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!(
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,
"must call current_api_key_async exactly once per provider build"
}
/// The runtime re-check, not construction alone, is what guards the wire:
/// `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!(
sync_calls.load(Ordering::SeqCst),
0,
"sync current_api_key must NOT be called — the async path is the contract"
let redirected = MemoryBackendParams {
embed_base_url: FOREIGN_ENDPOINT.into(),
..params_at(SESSION_ENDPOINT, scoped)
};
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"
);
}
}
crates/codegen/kigi-memory/src/lib.rs
+1 -1
View File
@@ -31,7 +31,7 @@ pub mod storage;
pub mod text_utils;
pub mod watcher;
pub use backend::{MemoryBackendImpl, MemoryBackendParams};
pub use backend::{EndpointScopedCredentials, MemoryBackendImpl, MemoryBackendParams};
pub use index::{MemoryIndex, init_sqlite_vec};
pub use storage::{MemoryScope, MemoryStorage};
crates/codegen/kigi-models/src/lib.rs
+6 -6
View File
@@ -1574,7 +1574,7 @@ impl PlatformId {
/// `models_cache.json`, the `visibility=="list"` AND `supported_in_api==true`
/// set) because OpenAI exposes no stable public models endpoint for the
/// ChatGPT Codex backend. Each entry carries context window + per-model
/// selectable reasoning efforts (incl. the codex-only `xhigh`/`max`/`ultra`
/// selectable reasoning efforts (incl. the codex-only `xhigh`/`max`
/// tiers), so the fetch path maps them through the SAME
/// `platform_wire_model_to_entry` output as a live listing — no new type.
pub fn hardcoded_catalog(self) -> Option<Vec<WireModel>> {
@@ -1617,13 +1617,13 @@ fn openai_codex_wire_models() -> Vec<WireModel> {
codex_wire_model(
"gpt-5.6-sol",
"GPT-5.6-Sol",
&["low", "medium", "high", "xhigh", "max", "ultra"],
&["low", "medium", "high", "xhigh", "max"],
"low",
),
codex_wire_model(
"gpt-5.6-terra",
"GPT-5.6-Terra",
&["low", "medium", "high", "xhigh", "max", "ultra"],
&["low", "medium", "high", "xhigh", "max"],
"medium",
),
codex_wire_model(
@@ -2655,7 +2655,7 @@ mod tests {
/// The HARDCODED openai-codex catalog is exactly the 4 supported+listed
/// models, keyed by slug, ctx 272000, each exposing its exact supported
/// efforts (incl. the codex-only `xhigh`/`max`/`ultra` tiers). The
/// efforts (incl. the codex-only `xhigh`/`max` tiers). The
/// list-but-broken / hidden models are absent. Every other platform serves
/// NO hardcoded catalog (its models come from the live wire).
#[test]
@@ -2702,11 +2702,11 @@ mod tests {
};
assert_eq!(
efforts("gpt-5.6-sol"),
["low", "medium", "high", "xhigh", "max", "ultra"]
["low", "medium", "high", "xhigh", "max"]
);
assert_eq!(
efforts("gpt-5.6-terra"),
["low", "medium", "high", "xhigh", "max", "ultra"]
["low", "medium", "high", "xhigh", "max"]
);
assert_eq!(
efforts("gpt-5.6-luna"),
crates/codegen/kigi-pager-minimal/src/commit.rs
+16 -11
View File
@@ -1251,6 +1251,10 @@ mod tests {
use ratatui::layout::Rect;
use similar::ChangeTag;
// A sibling test engages the terminal-native lock, which flattens
// every diff background process-wide. Hold the shared theme lock.
let _theme = kigi_tui::theme::cache::pin_theme();
let hunk = vec![
DiffLine {
text: "let x = 1;\n".into(),
@@ -1293,23 +1297,24 @@ mod tests {
// The committed edit uses a flat background (terminal transparency), but
// must still paint the per-line diff backgrounds — otherwise an added /
// removed line is indistinguishable from context.
let mut saw_insert = false;
let mut saw_delete = false;
//
// Theme-agnostic: the line cache is keyed on the global theme,
// so an exact-RGB match would race. Colors pinned in tool::edit.
let mut bands = std::collections::BTreeSet::new();
for y in 0..h {
for x in 0..width {
if let Some(cell) = buf.cell((x, y)) {
saw_insert |= cell.bg == theme.diff_insert_bg;
saw_delete |= cell.bg == theme.diff_delete_bg;
if let Some(cell) = buf.cell((x, y))
&& cell.bg != ratatui::style::Color::Reset
{
bands.insert(format!("{:?}", cell.bg));
}
}
}
assert!(
saw_insert,
"committed edit lost the insert (green) diff background"
);
assert!(
saw_delete,
"committed edit lost the delete (red) diff background"
bands.len() >= 2,
"committed edit lost its insert/delete diff bands: \
theme={:?} h={h} bands={bands:?}",
Theme::current_kind(),
);
}
crates/codegen/kigi-sampling-types/src/types.rs
+12 -36
View File
@@ -911,11 +911,6 @@ pub enum ReasoningEffort {
/// Messages both accept `xhigh` AND `max` as separate levels in 2026;
/// the Kimi wire spells its top tier `max` with no `xhigh`).
Max,
/// Codex-only top tier above `max` (the ChatGPT Codex backend exposes an
/// `ultra` reasoning effort on its flagship models). Reachable ONLY via a
/// model's server-declared effort menu (openai-codex); no built-in fallback
/// menu offers it, so other providers never emit it.
Ultra,
}
impl ReasoningEffort {
@@ -944,15 +939,12 @@ impl ReasoningEffort {
Self::High => "high",
Self::Xhigh => "xhigh",
Self::Max => "max",
Self::Ultra => "ultra",
}
}
/// Anthropic Messages API effort string; `None` for unsupported variants.
/// `xhigh` and `max` are distinct levels on the 2026 Messages API (both
/// appear in `GET /v1/models` `capabilities.effort`). `ultra` is codex-only
/// and never selected on an Anthropic model, but maps to its own string for
/// completeness (the Responses path writes effort via `as_str`, not this).
/// appear in `GET /v1/models` `capabilities.effort`).
pub fn to_messages_api(self) -> Option<&'static str> {
match self {
Self::None | Self::Minimal => None,
@@ -961,7 +953,6 @@ impl ReasoningEffort {
Self::High => Some("high"),
Self::Xhigh => Some("xhigh"),
Self::Max => Some("max"),
Self::Ultra => Some("ultra"),
}
}
}
@@ -984,9 +975,8 @@ impl std::str::FromStr for ReasoningEffort {
"high" => Ok(Self::High),
"xhigh" => Ok(Self::Xhigh),
"max" => Ok(Self::Max),
"ultra" => Ok(Self::Ultra),
_ => Err(format!(
"invalid reasoning effort: {s:?} (expected one of: none, minimal, low, medium, high, xhigh, max, ultra)"
"invalid reasoning effort: {s:?} (expected one of: none, minimal, low, medium, high, xhigh, max)"
)),
}
}
@@ -1878,30 +1868,16 @@ mod tests {
assert_eq!(ReasoningEffort::Max.to_messages_api(), Some("max"));
}
/// The codex-only `ultra` tier parses, serializes, and patches onto a
/// Responses body as `reasoning.effort = "ultra"` (the crux of surfacing a
/// codex model's full thinking menu). It is a DISTINCT level above `max`.
/// `ultra` is not a tier any backend accepts.
///
/// The Codex Responses endpoint rejects it with a 400 listing its menu,
/// which tops out at `max`. Parsing it would only send it again.
#[test]
fn reasoning_effort_ultra_is_a_distinct_codex_tier() {
assert_eq!(
"ultra".parse::<ReasoningEffort>().unwrap(),
ReasoningEffort::Ultra
);
assert_eq!(
"ULTRA".parse::<ReasoningEffort>().unwrap(),
ReasoningEffort::Ultra
);
assert_ne!(ReasoningEffort::Ultra, ReasoningEffort::Max);
assert_eq!(ReasoningEffort::Ultra.as_str(), "ultra");
let json = serde_json::to_string(&ReasoningEffort::Ultra).unwrap();
assert_eq!(json, "\"ultra\"");
assert_eq!(
serde_json::from_str::<ReasoningEffort>("\"ultra\"").unwrap(),
ReasoningEffort::Ultra
);
let mut body = serde_json::json!({ "model": "gpt-5.6-sol" });
patch_reasoning_effort(&mut body, Some(ReasoningEffort::Ultra));
assert_eq!(body["reasoning"]["effort"], "ultra");
fn reasoning_effort_ultra_is_not_a_tier() {
assert!("ultra".parse::<ReasoningEffort>().is_err());
assert!("ULTRA".parse::<ReasoningEffort>().is_err());
assert!(serde_json::from_str::<ReasoningEffort>("\"ultra\"").is_err());
assert_eq!(ReasoningEffort::Max.as_str(), "max");
}
/// The account id is decoded STATELESSLY from the bearer JWT payload's
@@ -2122,7 +2098,7 @@ mod tests {
);
let bad_type = as_map(serde_json::json!({"reasoningEffort": 3}));
assert_eq!(parse_reasoning_effort_meta(Some(&bad_type)), None);
// `ultra` is a real codex tier; a genuinely-unknown token still None.
// A genuinely-unknown token is dropped, not guessed at.
let unknown = as_map(serde_json::json!({"reasoningEffort": "MEGA"}));
assert_eq!(parse_reasoning_effort_meta(Some(&unknown)), None);
}
crates/codegen/kigi-shell/README.md
+2
View File
@@ -1309,6 +1309,7 @@ output_byte_limit = 65536 # max output size (64KB)
[toolset.web_fetch]
proxy_endpoint = "https://proxy.example.com" # egress proxy URL (all requests routed through it)
allowed_domains = ["docs.rs", "x.ai"] # override the built-in ~84-domain allowlist
allow_local = false # true = reach an explicit localhost / 127.0.0.0/8 / ::1 URL
[shortcuts]
send = ["Enter"]
@@ -2384,6 +2385,7 @@ The agent persists all session updates automatically. Clients can reconnect and
| `KIGI_AGENT` | Custom agent definition path or name (see [Agent Profiles](#agent-profiles)) |
| `KIGI_WEB_FETCH` | Enable (`1`) or disable (`0`) the `web_fetch` tool |
| `KIGI_WEB_FETCH_PROXY` | Egress proxy URL for `web_fetch` requests (overridden by `[toolset.web_fetch] proxy_endpoint`) |
| `KIGI_WEB_FETCH_ALLOW_LOCAL` | `1` lets `web_fetch` reach an explicit loopback URL; private and metadata ranges stay blocked |
| `KIGI_RESPECT_GITIGNORE` | Disable `.gitignore` filtering in tools when set to `0` |
| `KIGI_FEEDBACK_ENABLED` | Enable (`1`) or disable (`0`) feedback system independently from telemetry |
| `KIGI_DEPLOYMENT_KEY` | Management API key for enterprise deployments |
crates/codegen/kigi-shell/src/agent/models_fetch.rs
+4 -4
View File
@@ -1475,7 +1475,7 @@ mod tests {
/// short-circuits BEFORE any HTTP — there is NO mock `/models` server, yet
/// the fetch returns exactly the 4 compiled-in models keyed
/// `openai-codex/<slug>` on the Responses backend, ctx 272000, each exposing
/// its exact reasoning efforts (incl. the codex-only `xhigh`/`max`/`ultra`).
/// its exact reasoning efforts (incl. the codex-only `xhigh`/`max`).
/// A BOGUS base URL confirms no live `/models` request is attempted (it would
/// otherwise fail against an unroutable host).
#[tokio::test(flavor = "multi_thread")]
@@ -1556,10 +1556,10 @@ mod tests {
.iter()
.map(|o| o.id.as_str())
.collect::<Vec<_>>(),
vec!["low", "medium", "high", "xhigh", "max", "ultra"],
"sol exposes the full codex effort menu incl. ultra"
vec!["low", "medium", "high", "xhigh", "max"],
"sol exposes the full codex effort menu"
);
// gpt-5.5 tops out at xhigh (no max/ultra).
// gpt-5.5 tops out at xhigh (no max).
let five_five = result
.models
.iter()
crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs
+5
View File
@@ -357,6 +357,11 @@ impl MvpAgent {
// time, so advertise pre-session; the in-session path
// re-checks the live toolset.
graph: goal && self.cfg.borrow().resolve_graph().value,
// Tool-dependent, so fail closed like every other tool gate: the
// builder strips `agent_swarm` whenever subagents are unavailable,
// and advertising it then offers a menu entry that resolves to
// literal prompt text.
swarm: self.cfg.borrow().subagents_enabled,
..crate::session::slash_commands::CommandAvailability::default()
}
}
crates/codegen/kigi-shell/src/agent/session_config.rs
-1
View File
@@ -60,7 +60,6 @@ fn effort_label(effort: ReasoningEffort) -> String {
ReasoningEffort::High => "High",
ReasoningEffort::Xhigh => "X-High",
ReasoningEffort::Max => "Max",
ReasoningEffort::Ultra => "Ultra",
}
.to_string()
}
crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs
+10 -2
View File
@@ -396,11 +396,14 @@ pub(crate) async fn handle_subagent_request(
let child_depth = ctx.parent_depth + 1;
if child_depth >= MAX_SUBAGENT_DEPTH {
let before = definition.tool_config.tools.len();
definition.tool_config.tools.retain(|tc| tc.kind != Some(ToolKind::Task));
definition
.tool_config
.tools
.retain(|tc| !matches!(tc.kind, Some(ToolKind::Task | ToolKind::AgentSwarm)));
if definition.tool_config.tools.len() < before {
tracing::info!(
subagent_id = % request.id, child_depth, max_depth =
MAX_SUBAGENT_DEPTH, "Stripped task tool from child at max depth"
MAX_SUBAGENT_DEPTH, "Stripped subagent-spawning tools from child at max depth"
);
}
prune_orphaned_background_task_tools(&mut definition.tool_config);
@@ -1322,6 +1325,7 @@ pub(crate) async fn handle_subagent_request(
SubagentResult {
success: false,
cancelled: true,
rate_limited: false,
error: Some(reason),
output: if final_text.is_empty() {
std::sync::Arc::from(
@@ -1359,6 +1363,7 @@ pub(crate) async fn handle_subagent_request(
SubagentResult {
success: false,
cancelled: true,
rate_limited: false,
error: Some(format!("max turns reached (limit: {limit})")),
output: if final_text.is_empty() {
std::sync::Arc::from(
@@ -1413,6 +1418,9 @@ pub(crate) async fn handle_subagent_request(
SubagentResult {
success: false,
cancelled: was_cancelled,
rate_limited: !was_cancelled
&& i32::from(e.code)
== crate::sampling::error::RATE_LIMITED_ERROR_CODE,
error: Some(
if was_cancelled {
"Subagent was cancelled".to_string()
crates/codegen/kigi-shell/src/auth/credential_provider.rs
+61
View File
@@ -83,6 +83,35 @@ impl AuthCredentialProvider for ShellAuthCredentialProvider {
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)]
mod tests {
use super::*;
@@ -253,6 +282,38 @@ mod tests {
"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).
#[tokio::test]
async fn refresh_after_unauthorized_is_noop_for_deployment_key() {
crates/codegen/kigi-shell/src/session/acp_session.rs
+10
View File
@@ -121,6 +121,10 @@ mod model_switch;
mod prompt_queue;
#[path = "acp_session_impl/slash_exec.rs"]
mod slash_exec;
#[path = "acp_session_impl/stationarity_seam.rs"]
mod stationarity_seam;
#[path = "acp_session_impl/swarm.rs"]
mod swarm;
use super::PromptOrigin;
use super::acp_types;
use super::chat_persistence;
@@ -616,6 +620,9 @@ pub(crate) struct SessionActor {
/// layered over the goal engine. Modeled after `goal_tracker` above;
/// all graph state logic lives in `graph_tracker.rs`.
pub(crate) graph_tracker: Arc<parking_lot::Mutex<crate::session::graph_tracker::GraphTracker>>,
/// Swarm mode: a standing instruction to fan work out, independent of the
/// goal/graph engines — it steers tool choice, it does not drive turns.
pub(crate) swarm_mode: std::cell::Cell<crate::session::swarm_mode::SwarmMode>,
/// Max graph nodes running concurrently (1 = serial G0 behavior).
/// Cached at actor construction from `resolve_graph_concurrency`.
pub(crate) graph_concurrency: u32,
@@ -1033,6 +1040,9 @@ impl SessionActor {
// Graph rides the goal harness: nodes execute as goals, so
// `/graph` is only real when `/goal` is.
graph: self.graph_enabled && goal,
swarm: tool_names
.iter()
.any(|n| n == kigi_tools::implementations::kigi::AGENT_SWARM_TOOL_NAME),
}
}
/// Names of every tool registered with the session's tool bridge.
crates/codegen/kigi-shell/src/session/acp_session_impl/hook_dispatch.rs
+6 -1
View File
@@ -8,7 +8,12 @@ pub(super) fn turn_result_to_hook_outcome(
) -> kigi_tool_protocol::turn_hook::TurnHookOutcome {
use kigi_tool_protocol::turn_hook::TurnHookOutcome;
match result {
Ok(TurnOutcome::Completed { .. }) => TurnHookOutcome::Completed,
// A stationarity halt is a completed turn for hook purposes: nobody
// cancelled it, and a `Cancelled` outcome would tell every Stop hook
// the user interrupted the model.
Ok(TurnOutcome::Completed { .. }) | Ok(TurnOutcome::StationarityHalted { .. }) => {
TurnHookOutcome::Completed
}
Ok(TurnOutcome::Cancelled { .. }) | Ok(TurnOutcome::MaxTurnsReached { .. }) => {
TurnHookOutcome::Cancelled
}
crates/codegen/kigi-shell/src/session/acp_session_impl/slash_exec.rs
+30
View File
@@ -709,6 +709,36 @@ impl SessionActor {
BuiltinAction::GoalSet { .. } => {
unreachable!("GoalSet is intercepted in handle_prompt")
}
BuiltinAction::SwarmSet { enabled } => {
let msg = self.apply_swarm_mode(
enabled.then_some(crate::session::swarm_mode::SwarmTrigger::Manual),
);
self.send_slash_command_output(&msg).await;
ok_end_turn(0, None)
}
BuiltinAction::SwarmToggle => {
let turning_on = !self.swarm_mode.get().is_active();
let msg = self.apply_swarm_mode(
turning_on.then_some(crate::session::swarm_mode::SwarmTrigger::Manual),
);
self.send_slash_command_output(&msg).await;
ok_end_turn(0, None)
}
// `/swarm <task>` is handled before dispatch (it has to seed the
// turn with the task text); reaching here means the interception
// was bypassed, so report rather than silently dropping the task.
BuiltinAction::SwarmTask { prompt } => {
tracing::warn!(
prompt_len = prompt.len(),
"/swarm <task> reached the builtin executor; the turn seam did not intercept it"
);
self.send_slash_command_output(
"Could not start a swarm turn for that task. Run `/swarm on` and send the \
task as a normal message.",
)
.await;
ok_end_turn(0, None)
}
BuiltinAction::GoalStatus => {
let current_tokens = self.chat_state_handle.get_total_tokens().await as i64;
let goal_tokens = self.goal_tokens_used(current_tokens);
crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs
+25 -19
View File
@@ -311,6 +311,18 @@ pub(crate) async fn spawn_session_actor(
};
let embed_base_url = sampling_config.base_url.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(
|| crate::config::PruningConfig {
enabled: false,
@@ -618,16 +630,11 @@ pub(crate) async fn spawn_session_actor(
watcher,
stale_claim_secs: watcher_config.stale_claim_secs,
search_source: "tool",
api_key_provider: api_key_provider.clone(),
auth_credentials: auth_manager.as_ref().map(|am| {
std::sync::Arc::new(
crate::auth::credential_provider::ShellAuthCredentialProvider::new(
am.clone(),
None,
None,
),
) as std::sync::Arc<dyn kigi_auth::AuthCredentialProvider>
}),
embedding_credentials: crate::auth::credential_provider::embedding_session_credentials(
&embed_base_url,
embed_platform,
&models_manager.credential_authority(),
),
};
let backend = crate::session::memory::MemoryBackendImpl::from_session_params(
storage.clone(),
@@ -1108,6 +1115,7 @@ pub(crate) async fn spawn_session_actor(
goal_tracker,
graph_enabled,
graph_tracker,
swarm_mode: std::cell::Cell::new(Default::default()),
graph_concurrency: effective_config.resolve_graph_concurrency(),
graph_node_rounds: effective_config.resolve_graph_node_rounds(),
graph_replan_cap: effective_config.resolve_graph_replan_cap(),
@@ -1287,8 +1295,11 @@ pub(crate) async fn spawn_session_actor(
.map(|mc| mc.embedding.clone())
.unwrap_or_default();
let embed_dims = embed_config.dimensions;
let sampling_base_url = embed_base_url.clone();
let sampling_api_key = embed_api_key.clone();
// The session's own params, so the background reindex embeds through
// 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 chunks_added_counter = session.memory.chunks_added.clone();
tokio::task::spawn_local(async move {
@@ -1311,13 +1322,8 @@ pub(crate) async fn spawn_session_actor(
target : kigi_log::memory_log::TARGET, files = files.len(),
"MEMORY_REINDEX: background reindex complete"
);
if let Some(api_key) = sampling_api_key
&& let Some(provider) =
crate::session::memory::embedding::ApiEmbeddingProvider::from_session(
&embed_config,
sampling_base_url,
api_key,
)
if let Some(ref params) = reindex_params
&& let Some(provider) = params.make_embedding_provider().await
{
crate::session::memory::embed_missing_chunks(&index, &provider).await;
}
crates/codegen/kigi-shell/src/session/acp_session_impl/stationarity_seam.rs
+120
View File
@@ -0,0 +1,120 @@
//! Turn-loop seam for the stationarity detector.
use super::*;
use crate::session::stationarity::{
IdenticalToolCallRun, NUDGE_AFTER_IDENTICAL_TOOL_CALLS, command_is_true, hash_batch,
};
/// Sent once, naming the blocking wait a poller should use.
const STATIONARITY_NUDGE: &str = "\
You have called the same tool with the same arguments repeatedly and are in a \
polling loop. Stop repeating that call. If you are waiting on a background \
task, block on it with the wait-tasks tool instead of re-checking its output; \
if you are waiting on anything else, sleep once and check once. If you cannot \
make progress, stop and tell the user what you are waiting for. This turn will \
be halted automatically if the identical call keeps repeating.";
impl SessionActor {
/// Enforces the ceilings; `Some` halts the turn. Silent but logged.
pub(crate) async fn observe_tool_call_stationarity(
self: &Arc<Self>,
run: &mut IdenticalToolCallRun,
tool_calls: &[kigi_sampling_types::conversation::ToolCall],
loop_index: usize,
) -> Option<StationarityHalt> {
let batch_hash = hash_batch(
tool_calls
.iter()
.map(|tc| (tc.name.as_str(), tc.arguments.as_ref())),
);
let tool_name = tool_calls
.first()
.map(|tc| tc.name.clone())
.unwrap_or_default();
let is_true_noop = self.is_run_true_step(tool_calls).await;
let run_len = run.observe(batch_hash, &tool_name, is_true_noop);
if run_len == NUDGE_AFTER_IDENTICAL_TOOL_CALLS {
tracing::warn!(
tool_name = %run.tool_name(),
run_len,
loop_index,
"action stationarity: nudging a repeating tool call"
);
kigi_log::unified_log::warn(
"shell.turn.action_stationarity_nudge",
Some(self.session_info.id.0.as_ref()),
Some(serde_json::json!({
"tool_name": run.tool_name(),
"run_len": run_len,
"loop_index": loop_index,
})),
);
self.push_system_reminder(STATIONARITY_NUDGE);
}
if run_len < run.hard_stop_threshold() {
return None;
}
tracing::warn!(
tool_name = %run.tool_name(),
run_len,
loop_index,
true_noop = run.is_true_noop_run(),
"action stationarity: halting the turn"
);
kigi_log::unified_log::warn(
"shell.turn.action_stationarity_stop",
Some(self.session_info.id.0.as_ref()),
Some(serde_json::json!({
"tool_name": run.tool_name(),
"run_len": run_len,
"loop_index": loop_index,
"true_noop": run.is_true_noop_run(),
})),
);
Some(StationarityHalt {
tool_name: run.tool_name().to_string(),
run_len,
true_noop: run.is_true_noop_run(),
})
}
/// Whether this batch is a single shell call that does nothing.
///
/// Size-gated, not name-gated: the shell tool is renamed
/// `run_terminal_command`, so a name gate would silently disable this.
/// Parse failure fails open; multi-call batches use the ordinary ceiling.
async fn is_run_true_step(
&self,
tool_calls: &[kigi_sampling_types::conversation::ToolCall],
) -> bool {
/// Bound on `{"command":"true"}` plus sibling fields.
const MAX_NOOP_ARGS_BYTES: usize = 512;
let [tc] = tool_calls else {
return false;
};
if tc.arguments.as_ref().len() > MAX_NOOP_ARGS_BYTES {
return false;
}
let Ok(args) = serde_json::from_str::<serde_json::Value>(tc.arguments.as_ref()) else {
return false;
};
let Ok(input) = self.tool_bridge_handle().try_parse(&tc.name, args).await else {
return false;
};
matches!(
input,
kigi_tools::types::tool_io::ToolInput::Bash(ref b) if command_is_true(&b.command)
)
}
}
/// Why a turn was halted, carried to its outcome.
pub(crate) struct StationarityHalt {
pub(crate) tool_name: String,
pub(crate) run_len: u32,
pub(crate) true_noop: bool,
}
crates/codegen/kigi-shell/src/session/acp_session_impl/swarm.rs
+80
View File
@@ -0,0 +1,80 @@
//! Swarm-mode seam: arming/disarming the mode and the reminder it injects.
use super::*;
use crate::session::swarm_mode::{SWARM_ENTER_REMINDER, SWARM_EXIT_REMINDER, SwarmTrigger};
impl SessionActor {
/// Arms (`Some(trigger)`) or disarms (`None`) swarm mode, injecting or
/// retracting the doctrine exactly once, and returns the line to show.
pub(crate) fn apply_swarm_mode(self: &Arc<Self>, trigger: Option<SwarmTrigger>) -> String {
let mut mode = self.swarm_mode.get();
let message = match trigger {
Some(trigger) => {
if mode.enter(trigger) {
self.inject_swarm_reminder(SWARM_ENTER_REMINDER);
}
"Swarm mode on: work will be split across a fleet of subagents. \
`/swarm off` to stop."
}
None => {
// Unconditional, unlike the automatic expiry: the doctrine
// rides the conversation and therefore survives a resume, a
// fork and a compaction that the in-memory flag does not. If
// the user explicitly asks for it off, the retraction has to
// reach the model even when this session never saw it armed.
mode.exit();
self.inject_swarm_reminder(SWARM_EXIT_REMINDER);
"Swarm mode off."
}
};
self.swarm_mode.set(mode);
message.to_string()
}
/// A turn-scoped guard that disarms a per-turn swarm mode however the turn
/// ends.
///
/// The post-loop call site is not enough: a user interrupt ABORTS the turn
/// future (`cancel_running_task` → `JoinHandle::abort`), dropping it at its
/// current await point, and several `?` paths return before the loop's end.
/// Each of those leaks a `/swarm <task>` mode into the user's next,
/// unrelated prompt. The goal engine hit the same class and compensates
/// inside the cancel path; a guard is the version that cannot be forgotten
/// at a new exit.
pub(crate) fn swarm_turn_guard(self: &Arc<Self>) -> SwarmTurnGuard {
SwarmTurnGuard {
session: self.clone(),
}
}
/// Disarms at a turn boundary when the trigger was per-turn.
pub(crate) fn expire_swarm_mode_at_turn_end(self: &Arc<Self>) {
if !self.swarm_mode.get().expires_at_turn_end() {
return;
}
let mut mode = self.swarm_mode.get();
if mode.exit() {
self.inject_swarm_reminder(SWARM_EXIT_REMINDER);
}
self.swarm_mode.set(mode);
}
/// The doctrine rides the session's existing `<system-reminder>` channel,
/// so it is tagged the same way every other reminder is and needs no
/// second injection path of its own.
fn inject_swarm_reminder(self: &Arc<Self>, text: &str) {
self.push_system_reminder(text);
}
}
/// Runs [`SessionActor::expire_swarm_mode_at_turn_end`] on every turn exit,
/// including an aborted future.
pub(crate) struct SwarmTurnGuard {
session: Arc<SessionActor>,
}
impl Drop for SwarmTurnGuard {
fn drop(&mut self) {
self.session.expire_swarm_mode_at_turn_end();
}
}
crates/codegen/kigi-shell/src/session/acp_session_impl/tool_calls.rs
+12
View File
@@ -1546,6 +1546,18 @@ impl SessionActor {
vec![],
vec![],
),
// Without an explicit arm the catch-all below titles this "Tool
// call" — for an entry that can hold the turn for the whole swarm.
ToolInput::AgentSwarm(swarm) => (
format!(
"{} ({} members)",
swarm.description,
swarm.items.len() + swarm.resume_agent_ids.len()
),
acp::ToolKind::Other,
vec![],
vec![],
),
ToolInput::EnterPlanMode(_) => (
"Plan: Enter".to_string(),
acp::ToolKind::Other,
crates/codegen/kigi-shell/src/session/acp_session_impl/turn.rs
+75 -4
View File
@@ -217,6 +217,10 @@ impl SessionActor {
persist_ack: Option<oneshot::Sender<()>>,
) -> PromptTurnResult {
let handle_prompt_start = std::time::Instant::now();
// Armed before anything can arm swarm mode, so a `/swarm <task>` mode
// is disarmed on EVERY exit from this turn — including the abort a
// user interrupt performs, which never reaches post-loop code.
let _swarm_turn_guard = self.swarm_turn_guard();
let prompt_length: usize = prompt_blocks
.iter()
.map(|b| match b {
@@ -301,6 +305,13 @@ impl SessionActor {
span.record("command_source", "builtin");
}
match action {
// `/swarm <task>` arms the mode for this turn only and
// sends the task as the prompt, so the doctrine is in the
// conversation before the model reads the work.
BuiltinAction::SwarmTask { prompt } => {
self.apply_swarm_mode(Some(crate::session::swarm_mode::SwarmTrigger::Task));
vec![text_block(prompt)]
}
BuiltinAction::GoalSet {
objective,
token_budget,
@@ -779,7 +790,7 @@ impl SessionActor {
let turn_tool_count = self.events.tool_count_this_turn();
let bridge_outcome = turn_result_to_hook_outcome(&result);
match &result {
Ok(TurnOutcome::Completed { .. }) => {
Ok(TurnOutcome::Completed { .. }) | Ok(TurnOutcome::StationarityHalted { .. }) => {
self.emit_turn_ended(
crate::session::events::TurnOutcomeLabel::Completed,
None,
@@ -869,7 +880,9 @@ impl SessionActor {
let doom_tally = std::mem::take(&mut *self.doom_loop_turn_tally.lock());
doom_tally.fired();
let stop_reason_str = match &result {
Ok(TurnOutcome::Completed { .. }) => "end_turn",
Ok(TurnOutcome::Completed { .. }) | Ok(TurnOutcome::StationarityHalted { .. }) => {
"end_turn"
}
Ok(TurnOutcome::Cancelled { .. }) | Ok(TurnOutcome::MaxTurnsReached { .. }) => {
"cancelled"
}
@@ -885,7 +898,7 @@ impl SessionActor {
)
.await;
match &result {
Ok(TurnOutcome::Completed { .. }) => {
Ok(TurnOutcome::Completed { .. }) | Ok(TurnOutcome::StationarityHalted { .. }) => {
for contributor in self.extension_registry.turn_lifecycle_contributors() {
contributor
.on_turn_done(&kigi_agent_lifecycle::TurnDoneInput)
@@ -970,6 +983,26 @@ impl SessionActor {
PromptCompletionKind::MaxTurnsReached { limit },
None,
),
// `EndTurn`, not `Cancelled`: the client must not render
// an interrupted turn for something nobody interrupted.
// The halt's detail rides `completion_kind`, which is
// where a bug report can still read it.
TurnOutcome::StationarityHalted {
snapshot,
tool_name,
run_len,
true_noop,
..
} => (
acp::StopReason::EndTurn,
*snapshot,
PromptCompletionKind::StationarityHalted {
tool_name,
run_len,
true_noop,
},
None,
),
};
if let Some(snapshot) = snapshot.as_mut() {
self.apply_prompt_modes_to_snapshot(snapshot);
@@ -1271,7 +1304,11 @@ impl SessionActor {
let mut result = self
.process_conversation_turn(req_id, json_schema.clone())
.await;
if matches!(result, Ok(TurnOutcome::MaxTurnsReached { .. })) {
// Harness stopped the turn; retrying re-enters the same wall.
if matches!(
result,
Ok(TurnOutcome::MaxTurnsReached { .. }) | Ok(TurnOutcome::StationarityHalted { .. })
) {
return result;
}
if let Ok(TurnOutcome::Completed {
@@ -1588,6 +1625,10 @@ impl SessionActor {
self.record_turn_model().await;
let mut metrics_drop_guard = TurnMetrics::new();
let mut turn_tools_called: Vec<String> = Vec::new();
let mut identical_tool_calls =
crate::session::stationarity::IdenticalToolCallRun::default();
// Retained across execute: observed only after results land.
let mut last_batch: Vec<kigi_sampling_types::conversation::ToolCall> = Vec::new();
let mut tool_turn_count: usize = 1;
let mut loop_index: u32 = 0;
let mut todo_gate_fires: u32 = 0;
@@ -2004,6 +2045,8 @@ impl SessionActor {
}
turn_tools_called.push(tc.name.clone());
}
last_batch.clear();
last_batch.extend(tool_calls.iter().cloned());
let tool_call_responses: Vec<ToolCallResponse> = tool_calls
.into_iter()
.map(|tc| ToolCallResponse {
@@ -2046,6 +2089,34 @@ impl SessionActor {
}
_ => {}
}
// After execute: every call has a result, so nothing dangles
// and the nudge cannot trigger the "cancelled" repair.
// Ok-gated: one `?` inside can leave a call resultless.
if execute_tool_calls_result.is_ok()
&& let Some(halt) = self
.observe_tool_call_stationarity(
&mut identical_tool_calls,
&last_batch,
tool_turn_count,
)
.await
{
let snapshot = self
.finalize_turn_bookkeeping(
req_id,
conv_turn_start,
&turn_span_totals,
model_fingerprint.clone(),
)
.await;
return Ok(TurnOutcome::StationarityHalted {
snapshot: Box::new(snapshot),
tools_called: std::mem::take(&mut turn_tools_called),
tool_name: halt.tool_name,
run_len: halt.run_len,
true_noop: halt.true_noop,
});
}
let next_turn = tool_turn_count + 1;
if let Some(limit) = self.max_turns
&& next_turn > limit
crates/codegen/kigi-shell/src/session/acp_session_impl/types.rs
+13
View File
@@ -62,6 +62,19 @@ pub(crate) enum TurnOutcome {
},
/// The `--max-turns` limit was reached after a tool-execution cycle.
MaxTurnsReached { limit: usize },
/// One tool call repeated past its ceiling; the turn was halted.
///
/// Groups with [`Self::Completed`], NOT [`Self::Cancelled`]: nobody
/// cancelled anything. As a cancellation it would report
/// `StopReason::Cancelled`, fire the abort lifecycle, kill the turn's
/// subagents, grow the goal back-off streak, and let recovery re-run it.
StationarityHalted {
snapshot: Box<Option<TurnDeltaSnapshot>>,
tools_called: Vec<String>,
tool_name: String,
run_len: u32,
true_noop: bool,
},
}
#[derive(Debug)]
crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs
+5 -2
View File
@@ -106,6 +106,7 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
tokio_util::sync::CancellationToken::new(),
);
let actor = Arc::new(SessionActor {
swarm_mode: std::cell::Cell::new(Default::default()),
session_info,
auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None),
@@ -555,11 +556,11 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
};
let (event_tx, _event_rx) = tokio::sync::mpsc::unbounded_channel::<SessionEvent>();
let actor = Arc::new(SessionActor {
swarm_mode: std::cell::Cell::new(Default::default()),
session_info: session_info.clone(),
auth_method_id: test_auth_method_id("test-auth"),
model_auth_facts: std::cell::RefCell::new(None),
@@ -821,6 +822,7 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
)
.await;
let actor = SessionActor {
swarm_mode: std::cell::Cell::new(Default::default()),
session_info: SessionInfo {
id: acp::SessionId::new("test-cancel"),
cwd: cwd.as_str().to_string(),
@@ -1814,6 +1816,7 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
)
.await;
let actor = SessionActor {
swarm_mode: std::cell::Cell::new(Default::default()),
session_info: SessionInfo {
id: acp::SessionId::new("test-cancel-sampler"),
cwd: cwd.as_str().to_string(),
crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs
+1
View File
@@ -115,6 +115,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
});
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let actor = SessionActor {
swarm_mode: std::cell::Cell::new(Default::default()),
session_info: SessionInfo {
id: acp::SessionId::new("test-idle-resume"),
cwd: cwd.as_str().to_string(),
crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs
+5 -4
View File
@@ -64,6 +64,7 @@ async fn create_test_actor(
);
chat_state_handle.record_token_usage(total_tokens);
SessionActor {
swarm_mode: std::cell::Cell::new(Default::default()),
session_info: SessionInfo {
id: acp::SessionId::new("test-auto-compact"),
cwd: cwd.as_str().to_string(),
@@ -393,8 +394,7 @@ fn initial_injection_backend_params_use_override_min_score() {
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
};
let initial_injection = crate::config::MemoryInitialInjectionConfig {
enabled: true,
@@ -422,8 +422,7 @@ fn initial_injection_backend_params_preserve_default_zero_min_score() {
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
};
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
&params,
@@ -495,6 +494,7 @@ async fn create_test_actor_with_memory(
.as_ref()
.map_or_else(Default::default, |mc| mc.initial_injection.clone());
SessionActor {
swarm_mode: std::cell::Cell::new(Default::default()),
session_info: SessionInfo {
id: acp::SessionId::new("test-memory"),
cwd: cwd.as_str().to_string(),
@@ -1237,6 +1237,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
});
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let actor = SessionActor {
swarm_mode: std::cell::Cell::new(Default::default()),
session_info: SessionInfo {
id: acp::SessionId::new("test-idle-resume"),
cwd: cwd.as_str().to_string(),
crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs
+4 -6
View File
@@ -18,8 +18,7 @@ fn initial_injection_backend_params_use_override_min_score() {
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
};
let initial_injection = crate::config::MemoryInitialInjectionConfig {
enabled: true,
@@ -47,8 +46,7 @@ fn initial_injection_backend_params_preserve_default_zero_min_score() {
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
};
let (adjusted, effective_min_score) = build_initial_injection_backend_params(
&params,
@@ -121,6 +119,7 @@ async fn create_test_actor_with_memory(
.as_ref()
.map_or_else(Default::default, |mc| mc.initial_injection.clone());
SessionActor {
swarm_mode: std::cell::Cell::new(Default::default()),
session_info: SessionInfo {
id: acp::SessionId::new("test-memory"),
cwd: cwd.as_str().to_string(),
@@ -524,8 +523,7 @@ async fn create_injection_ready_actor(
watcher: None,
stale_claim_secs: 60,
search_source: "tool",
api_key_provider: None,
auth_credentials: None,
embedding_credentials: kigi_memory::EndpointScopedCredentials::none(),
});
actor
.chat_state_handle
crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs
+1
View File
@@ -72,6 +72,7 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture
});
let (event_tx, event_rx) = mpsc::unbounded_channel::<SessionEvent>();
let actor = SessionActor {
swarm_mode: std::cell::Cell::new(Default::default()),
session_info: SessionInfo {
id: acp::SessionId::new("test-session"),
cwd: cwd.as_str().to_string(),
crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs
+1
View File
@@ -182,6 +182,7 @@ pub(crate) async fn create_test_actor_ex(
chat_state_handle.record_token_usage(total_tokens);
let (goal_update_tx, goal_update_rx) = tokio::sync::mpsc::unbounded_channel();
let actor = SessionActor {
swarm_mode: std::cell::Cell::new(Default::default()),
session_info: SessionInfo {
id: acp::SessionId::new("test-actor"),
cwd: cwd.as_str().to_string(),
crates/codegen/kigi-shell/src/session/commands.rs
+7
View File
@@ -29,6 +29,13 @@ pub enum PromptCompletionKind {
MaxTurnsReached {
limit: usize,
},
/// One tool call repeated past its ceiling. Reported as `EndTurn`:
/// the turn ended, nobody interrupted it.
StationarityHalted {
tool_name: String,
run_len: u32,
true_noop: bool,
},
Rewound,
/// A queued prompt was removed (or cleared) from the server-authoritative
/// queue before it ever ran. Used to resolve the still-pending
crates/codegen/kigi-shell/src/session/compaction.rs
+1
View File
@@ -2154,6 +2154,7 @@ mod inline_auto_compact_flow_tests {
);
chat_state_handle.record_token_usage(total_tokens);
SessionActor {
swarm_mode: std::cell::Cell::new(Default::default()),
session_info: SessionInfo {
id: acp::SessionId::new("test-auto-compact"),
cwd: cwd.as_str().to_string(),
crates/codegen/kigi-shell/src/session/mod.rs
+2
View File
@@ -326,9 +326,11 @@ pub mod restore;
pub mod result;
pub mod signals;
pub(crate) mod slash_commands;
pub(crate) mod stationarity;
pub mod storage;
pub(crate) mod streaming_capture;
pub(crate) mod summary;
pub mod swarm_mode;
pub(crate) mod telemetry;
pub mod tool_index;
pub(crate) mod turn_completion;
crates/codegen/kigi-shell/src/session/slash_commands.rs
+90
View File
@@ -45,6 +45,10 @@ pub(crate) enum BuiltinGate {
/// available (graph nodes execute as goals, so `/graph` needs
/// everything `/goal` needs).
Graph,
/// The `agent_swarm` tool is in the session toolset. The mode is only a
/// standing instruction to use that tool, so without it the command would
/// advertise a doctrine the model has no way to follow.
Swarm,
}
/// All built-in slash commands. Order here = display order in autocomplete.
@@ -303,6 +307,26 @@ pub(super) const BUILTIN_COMMANDS: &[BuiltinCommand] = &[
}
},
},
BuiltinCommand {
name: "swarm",
description: "Delegate aggressively: split the work across a fleet of subagents",
argument_hint: Some("[on | off | <task>]"),
aliases: &[],
gate: BuiltinGate::Swarm,
resolve: |args| {
let trimmed = args.trim();
match trimmed.to_lowercase().as_str() {
"on" => BuiltinAction::SwarmSet { enabled: true },
"off" => BuiltinAction::SwarmSet { enabled: false },
"" => BuiltinAction::SwarmToggle,
// Anything else is the task itself: arm the mode for exactly
// this turn and send the text as the prompt.
_ => BuiltinAction::SwarmTask {
prompt: trimmed.to_string(),
},
}
},
},
];
/// Split a trailing `--budget <tokens>` flag off a `/goal` objective.
@@ -437,6 +461,8 @@ pub(crate) struct CommandAvailability {
/// `/graph` gate: the graph feature flag AND the goal harness (nodes
/// execute as goals) are both available.
pub graph: bool,
/// `/swarm` gate: the `agent_swarm` tool is in the active toolset.
pub swarm: bool,
}
impl CommandAvailability {
@@ -452,6 +478,7 @@ impl CommandAvailability {
BuiltinGate::Plugins => self.plugins,
BuiltinGate::Goal => self.goal,
BuiltinGate::Graph => self.graph,
BuiltinGate::Swarm => self.swarm,
}
}
@@ -468,6 +495,7 @@ impl CommandAvailability {
plugins: true,
goal: true,
graph: true,
swarm: true,
}
}
}
@@ -705,6 +733,17 @@ pub(super) enum BuiltinAction {
token_budget: Option<i64>,
},
GoalStatus,
/// `/swarm on|off` — arm or disarm the standing delegate-aggressively
/// instruction. Survives turns until switched off.
SwarmSet {
enabled: bool,
},
/// `/swarm` with no argument.
SwarmToggle,
/// `/swarm <task>` — arm for this turn only, then send `prompt`.
SwarmTask {
prompt: String,
},
GoalPause,
GoalResume,
GoalClear,
@@ -757,6 +796,9 @@ impl BuiltinAction {
| BuiltinAction::GraphPause
| BuiltinAction::GraphResume { .. }
| BuiltinAction::GraphClear => "graph",
BuiltinAction::SwarmSet { .. }
| BuiltinAction::SwarmToggle
| BuiltinAction::SwarmTask { .. } => "swarm",
}
}
@@ -795,6 +837,8 @@ impl BuiltinAction {
| BuiltinAction::GraphShow
| BuiltinAction::GraphPause
| BuiltinAction::GraphClear => false,
BuiltinAction::SwarmToggle => false,
BuiltinAction::SwarmSet { .. } | BuiltinAction::SwarmTask { .. } => true,
}
}
}
@@ -1584,6 +1628,7 @@ mod tests {
"feedback",
"goal",
"graph",
"swarm",
"loop",
"commit",
"deploy",
@@ -1694,6 +1739,51 @@ mod tests {
);
}
/// Without the `agent_swarm` tool the mode has nothing to steer toward, so
/// the command must fall through as ordinary prompt text rather than
/// arming a doctrine the model cannot act on.
#[test]
fn swarm_does_not_resolve_when_gate_off() {
let availability = CommandAvailability {
swarm: false,
..CommandAvailability::all_enabled()
};
assert!(
resolve(
vec![text_block("/swarm on")],
&[],
availability,
SkillSlashRewrite::default(),
)
.is_ok(),
"expected pass-through (Ok), got an outcome",
);
}
#[test]
fn swarm_resolves_each_form_to_its_own_action() {
assert!(matches!(
resolve_builtin("swarm", "on").expect("/swarm on must resolve"),
BuiltinAction::SwarmSet { enabled: true }
));
assert!(matches!(
resolve_builtin("swarm", "off").expect("/swarm off must resolve"),
BuiltinAction::SwarmSet { enabled: false }
));
assert!(matches!(
resolve_builtin("swarm", "").expect("bare /swarm must resolve"),
BuiltinAction::SwarmToggle
));
// Anything else is the task, NOT an unknown subcommand: mis-parsing it
// would silently drop the user's work instead of running it.
match resolve_builtin("swarm", "split the auth refactor")
.expect("/swarm <task> must resolve")
{
BuiltinAction::SwarmTask { prompt } => assert_eq!(prompt, "split the auth refactor"),
other => panic!("expected SwarmTask, got {}", other.command_name()),
}
}
#[test]
fn graph_resolves_subcommands_and_budget() {
let set = resolve_builtin("graph", "ship the feature --budget 5000")
crates/codegen/kigi-shell/src/session/stationarity.rs
+216
View File
@@ -0,0 +1,216 @@
//! Detects a model stuck repeating one identical tool call.
/// Consecutive identical batches after which the turn is halted.
pub(crate) const MAX_CONSECUTIVE_IDENTICAL_TOOL_CALLS: u32 = 16;
/// Nudge threshold; half the budget remains after it.
pub(crate) const NUDGE_AFTER_IDENTICAL_TOOL_CALLS: u32 = 8;
/// Below the nudge on purpose: no-op runs get none.
pub(crate) const MAX_CONSECUTIVE_TRUE_NOOPS: u32 = 4;
const _: () = assert!(NUDGE_AFTER_IDENTICAL_TOOL_CALLS < MAX_CONSECUTIVE_IDENTICAL_TOOL_CALLS);
const _: () = assert!(MAX_CONSECUTIVE_TRUE_NOOPS < NUDGE_AFTER_IDENTICAL_TOOL_CALLS);
/// A shell command that does nothing whatsoever.
pub(crate) fn command_is_true(cmd: &str) -> bool {
cmd.trim().eq_ignore_ascii_case("true")
}
/// Hashes name+args per call; separators prevent concatenation collisions.
pub(crate) fn hash_batch<'a>(calls: impl IntoIterator<Item = (&'a str, &'a str)>) -> u64 {
use std::hash::{Hash, Hasher};
let mut hasher = std::collections::hash_map::DefaultHasher::new();
for (name, args) in calls {
name.hash(&mut hasher);
'\u{1f}'.hash(&mut hasher);
args.hash(&mut hasher);
'\u{1e}'.hash(&mut hasher);
}
hasher.finish()
}
/// One identity, so re-spelling a no-op cannot reset it.
const TRUE_NOOP_HASH: u64 = u64::MAX;
/// How many times the current batch has repeated unchanged.
#[derive(Default)]
pub(crate) struct IdenticalToolCallRun {
/// Hash only: signatures hold raw tool arguments.
last_hash: Option<u64>,
tool_name: String,
run_len: u32,
is_true_noop_run: bool,
}
impl IdenticalToolCallRun {
/// Records one batch and returns the length of the run it belongs to.
pub(crate) fn observe(&mut self, batch_hash: u64, tool_name: &str, is_true_noop: bool) -> u32 {
let hash = if is_true_noop {
TRUE_NOOP_HASH
} else {
batch_hash
};
if self.last_hash == Some(hash) {
self.run_len += 1;
} else {
self.run_len = 1;
self.last_hash = Some(hash);
self.is_true_noop_run = is_true_noop;
}
self.tool_name = tool_name.to_string();
self.run_len
}
/// The run length at which this turn must be halted.
pub(crate) fn hard_stop_threshold(&self) -> u32 {
if self.is_true_noop_run {
MAX_CONSECUTIVE_TRUE_NOOPS
} else {
MAX_CONSECUTIVE_IDENTICAL_TOOL_CALLS
}
}
pub(crate) fn tool_name(&self) -> &str {
&self.tool_name
}
pub(crate) fn is_true_noop_run(&self) -> bool {
self.is_true_noop_run
}
}
#[cfg(test)]
mod tests {
use super::*;
fn call(name: &str, args: &str) -> u64 {
hash_batch([(name, args)])
}
#[test]
fn an_unchanged_batch_accumulates_a_run() {
let mut run = IdenticalToolCallRun::default();
let h = call("read_file", "a.rs");
assert_eq!(run.observe(h, "read_file", false), 1);
assert_eq!(run.observe(h, "read_file", false), 2);
assert_eq!(run.observe(h, "read_file", false), 3);
}
#[test]
fn any_change_in_the_arguments_restarts_the_run() {
let mut run = IdenticalToolCallRun::default();
run.observe(call("read_file", "a.rs"), "read_file", false);
run.observe(call("read_file", "a.rs"), "read_file", false);
assert_eq!(
run.observe(call("read_file", "b.rs"), "read_file", false),
1,
"a different argument is a different action"
);
}
#[test]
fn two_batches_cannot_collide_by_concatenation() {
// Unseparated these would hash identically.
assert_ne!(
hash_batch([("ab", "cd")]),
hash_batch([("a", "b"), ("c", "d")])
);
}
#[test]
fn no_ops_share_one_run_however_they_are_spelled() {
let mut run = IdenticalToolCallRun::default();
assert_eq!(run.observe(call("bash", "true"), "bash", true), 1);
assert_eq!(
run.observe(call("bash", " TRUE "), "bash", true),
2,
"re-spelling a no-op must not reset the tighter ceiling"
);
assert_eq!(run.hard_stop_threshold(), MAX_CONSECUTIVE_TRUE_NOOPS);
}
#[test]
fn a_real_call_after_a_noop_run_restores_the_ordinary_ceiling() {
let mut run = IdenticalToolCallRun::default();
run.observe(call("bash", "true"), "bash", true);
assert_eq!(run.hard_stop_threshold(), MAX_CONSECUTIVE_TRUE_NOOPS);
run.observe(call("read_file", "a.rs"), "read_file", false);
assert_eq!(
run.hard_stop_threshold(),
MAX_CONSECUTIVE_IDENTICAL_TOOL_CALLS
);
}
#[test]
fn an_ordinary_repeated_call_is_halted_at_its_ceiling_and_not_before() {
let mut run = IdenticalToolCallRun::default();
let h = call("read_file", "a.rs");
for expected in 1..MAX_CONSECUTIVE_IDENTICAL_TOOL_CALLS {
let len = run.observe(h, "read_file", false);
assert_eq!(len, expected);
assert!(len < run.hard_stop_threshold(), "must not halt early");
}
assert_eq!(
run.observe(h, "read_file", false),
run.hard_stop_threshold(),
"the turn halts on the 16th identical call"
);
}
#[test]
fn a_repeated_no_op_is_halted_far_sooner_and_without_a_nudge() {
let mut run = IdenticalToolCallRun::default();
let h = call("bash", "true");
let mut len = 0;
while len < run.hard_stop_threshold() {
len = run.observe(h, "bash", true);
}
assert_eq!(len, MAX_CONSECUTIVE_TRUE_NOOPS);
assert!(
len < NUDGE_AFTER_IDENTICAL_TOOL_CALLS,
"documented: a no-op run is halted before any nudge could fire"
);
}
#[test]
fn the_nudge_lands_with_budget_left_to_act_on_it() {
let mut run = IdenticalToolCallRun::default();
let h = call("read_file", "a.rs");
let mut len = 0;
while len < NUDGE_AFTER_IDENTICAL_TOOL_CALLS {
len = run.observe(h, "read_file", false);
}
assert!(
len < run.hard_stop_threshold(),
"a warning the model cannot act on is not a warning"
);
}
#[test]
fn work_interleaved_with_repeats_is_never_halted() {
let mut run = IdenticalToolCallRun::default();
for i in 0..MAX_CONSECUTIVE_IDENTICAL_TOOL_CALLS * 2 {
let h = if i % 2 == 0 {
call("read_file", "a.rs")
} else {
call("read_file", "b.rs")
};
let len = run.observe(h, "read_file", false);
assert!(
len < run.hard_stop_threshold(),
"alternating calls are progress, not a loop"
);
}
}
#[test]
fn only_a_bare_true_counts_as_a_no_op() {
assert!(command_is_true("true"));
assert!(command_is_true(" true "));
assert!(command_is_true("TRUE"));
assert!(!command_is_true("true && make"));
assert!(!command_is_true("truely"));
assert!(!command_is_true(""));
}
}
crates/codegen/kigi-shell/src/session/swarm_mode.rs
+178
View File
@@ -0,0 +1,178 @@
//! Swarm mode: a standing instruction to split work across a fleet.
//!
//! The `agent_swarm` tool works with the mode off; what the mode adds is the
//! doctrine — decompose finely, give every member a disjoint scope, do not do
//! the work yourself. Kept as pure state so the turn loop decides when to
//! inject and the session decides when to persist.
use serde::{Deserialize, Serialize};
/// Why the mode is on, which is what decides when it turns off.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SwarmTrigger {
/// `/swarm on` (or a bare `/swarm` toggle): stays until switched off.
Manual,
/// `/swarm <task>`: armed for exactly that turn.
Task,
}
/// The session's swarm-mode state. `None` means off.
///
/// Deliberately NOT persisted, unlike `/goal` and `/graph`: those drive
/// autonomous multi-turn work that is stranded if it is lost, whereas this is
/// a prompt hint whose worst-case recovery is typing `/swarm on` again. The
/// injected doctrine IS durable (it rides the conversation), so a resumed
/// session can read as "off" with the instruction still in context — which is
/// exactly why an explicit `/swarm off` always retracts (see
/// [`SessionActor::apply_swarm_mode`]) rather than trusting a remembered flag
/// that a restore, a compaction or a fork can each falsify.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct SwarmMode {
trigger: Option<SwarmTrigger>,
}
impl SwarmMode {
pub fn is_active(self) -> bool {
self.trigger.is_some()
}
pub fn trigger(self) -> Option<SwarmTrigger> {
self.trigger
}
/// Arms the mode, and does NOTHING if it is already armed.
///
/// Total no-op, not just "don't re-inject": overwriting the trigger would
/// let the `/swarm <task>` shorthand downgrade a standing `/swarm on` into
/// a per-turn mode, which then disarms itself at that turn's end — the
/// user's deliberate choice, silently undone.
///
/// Returns whether the caller should inject the enter reminder.
pub fn enter(&mut self, trigger: SwarmTrigger) -> bool {
if self.trigger.is_some() {
return false;
}
self.trigger = Some(trigger);
true
}
/// Disarms. Returns whether the mode had been armed — which is what an
/// AUTOMATIC expiry keys its retraction off. An explicit `/swarm off` must
/// retract regardless (see the type docs).
pub fn exit(&mut self) -> bool {
self.trigger.take().is_some()
}
/// Whether a turn ending now should disarm the mode.
///
/// Only the `Task` trigger auto-exits: `/swarm on` is a standing choice the
/// user made and a turn boundary is not a reason to undo it.
pub fn expires_at_turn_end(self) -> bool {
self.trigger == Some(SwarmTrigger::Task)
}
}
/// The doctrine injected when the mode is armed.
///
/// Deliberately short: it is re-read on every turn it is live, and the tool's
/// own description already carries the mechanics.
pub const SWARM_ENTER_REMINDER: &str = "\
Swarm mode is on. Explore only as far as you must to identify the work, then \
split it: use the agent_swarm tool with one item per independent scope rather \
than doing the work yourself. Decompose finely — do not try to conserve \
members. Every member must own a disjoint scope; members share one working \
tree, so two members told to touch the same file will corrupt each other. \
Read-only scopes may overlap. If the work genuinely does not split, say so and \
carry on alone.";
/// Injected when the mode is switched off mid-conversation, so the earlier
/// doctrine does not keep steering the model.
/// Deliberately as emphatic as the enter doctrine it revokes: a one-line
/// "mode is off" is the weaker of the two texts in context and the likelier to
/// be summarised away, leaving the fan-out directives still steering.
pub const SWARM_EXIT_REMINDER: &str = "\
Swarm mode is off. The swarm instructions above no longer apply — you are not \
required to split work across subagents, and you should not decompose a task \
just because they said to. Decide how to approach each new request from the \
request itself. Delegate only where it clearly helps.";
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_fresh_session_has_the_mode_off() {
let mode = SwarmMode::default();
assert!(!mode.is_active());
assert!(!mode.expires_at_turn_end());
}
#[test]
fn arming_asks_for_the_doctrine_once_not_once_per_command() {
let mut mode = SwarmMode::default();
assert!(mode.enter(SwarmTrigger::Manual), "first arm injects");
assert!(
!mode.enter(SwarmTrigger::Manual),
"a repeated /swarm on must not stack a second copy"
);
assert!(mode.is_active());
}
#[test]
fn only_the_task_trigger_expires_at_a_turn_boundary() {
let mut manual = SwarmMode::default();
manual.enter(SwarmTrigger::Manual);
assert!(
!manual.expires_at_turn_end(),
"/swarm on is a standing choice, not a per-turn one"
);
let mut task = SwarmMode::default();
task.enter(SwarmTrigger::Task);
assert!(task.expires_at_turn_end());
}
#[test]
fn exiting_retracts_exactly_once() {
let mut mode = SwarmMode::default();
mode.enter(SwarmTrigger::Manual);
assert!(mode.exit(), "the live doctrine must be retracted");
assert!(!mode.is_active());
assert!(
!mode.exit(),
"a second /swarm off has nothing left to retract"
);
}
/// `/swarm <task>` under a standing `/swarm on` must seed the turn and
/// nothing more — it must not convert the standing mode into a per-turn
/// one that disarms itself when that turn ends.
#[test]
fn the_task_shorthand_never_downgrades_a_standing_mode() {
let mut mode = SwarmMode::default();
mode.enter(SwarmTrigger::Manual);
assert!(
!mode.enter(SwarmTrigger::Task),
"the doctrine is already in the conversation"
);
assert_eq!(mode.trigger(), Some(SwarmTrigger::Manual));
assert!(
!mode.expires_at_turn_end(),
"the user's standing /swarm on must survive the turn"
);
}
/// An automatic expiry has nothing to retract once the mode is already
/// off; only an explicit `/swarm off` retracts unconditionally, and that
/// rule lives at the call site, not here.
#[test]
fn exit_reports_whether_it_actually_disarmed_something() {
let mut armed = SwarmMode::default();
armed.enter(SwarmTrigger::Task);
assert!(armed.exit());
let mut idle = SwarmMode::default();
assert!(!idle.exit());
}
}
crates/codegen/kigi-shell/src/tools/config.rs
+13
View File
@@ -108,6 +108,10 @@ pub struct WebFetchToolConfig {
/// default allowlist. An explicit empty list blocks all fetches.
/// Resolution: TOML > remote settings > built-in defaults.
pub allowed_domains: Option<Vec<String>>,
/// Allow fetches to explicit loopback hosts only (`localhost` /
/// `127.0.0.0/8` / `::1`). Private and metadata ranges stay blocked.
/// Resolution: TOML > `KIGI_WEB_FETCH_ALLOW_LOCAL` env > false.
pub allow_local: Option<bool>,
}
impl WebFetchToolConfig {
@@ -137,10 +141,15 @@ impl WebFetchToolConfig {
.cloned()
.or_else(|| remote_domains.map(|d| d.to_vec()));
let allow_local = self
.allow_local
.or_else(|| kigi_config::env_bool("KIGI_WEB_FETCH_ALLOW_LOCAL"));
kigi_tools::implementations::kigi::web_fetch::WebFetchParams {
proxy_endpoint,
allowed_domains,
context_window_tokens,
allow_local,
..Default::default()
}
}
@@ -484,6 +493,7 @@ mod tests {
let local = WebFetchToolConfig {
proxy_endpoint: Some("https://toml-proxy.example.com".to_owned()),
allowed_domains: Some(vec!["toml.example.com".to_owned()]),
allow_local: Some(true),
};
let params = local.resolve_params(
Some("https://remote-proxy.example.com"),
@@ -498,6 +508,7 @@ mod tests {
params.allowed_domains,
Some(vec!["toml.example.com".to_owned()])
);
assert!(params.allow_local(), "the opt-in must reach the tool");
}
#[test]
@@ -524,6 +535,7 @@ mod tests {
let params = local.resolve_params(None, None, None);
assert!(params.proxy_endpoint.is_none());
assert!(params.allowed_domains.is_none());
assert!(!params.allow_local(), "local access is off by default");
}
#[test]
@@ -531,6 +543,7 @@ mod tests {
let local = WebFetchToolConfig {
proxy_endpoint: None,
allowed_domains: Some(vec![]),
allow_local: None,
};
let params = local.resolve_params(None, Some(&["remote.example.com".to_owned()]), None);
assert_eq!(params.allowed_domains, Some(vec![]));
crates/codegen/kigi-tools/schema/tool_meta.schema.json
+1 -1
View File
@@ -36,7 +36,7 @@
],
"definitions": {
"ToolKind": {
"description": "Categorizes what a tool does at a high level. Open set — consumers must tolerate unknown values (Rust deserializes them to `other` via `#[serde(other)]`). Known values: `read`, `edit`, `delete`, `list_dir`, `write`, `move`, `search`, `lsp`, `execute`, `plan`, `web_search`, `web_fetch`, `background_task_action`, `wait_tasks_action`, `kill_task_action`, `list`, `skill`, `memory_search`, `memory_get`, `task`, `enter_plan`, `exit_plan`, `ask_user`, `deploy_app`, `search_tool`, `use_tool`, `monitor`, `goal_update`, `other`.",
"description": "Categorizes what a tool does at a high level. Open set — consumers must tolerate unknown values (Rust deserializes them to `other` via `#[serde(other)]`). Known values: `read`, `edit`, `delete`, `list_dir`, `write`, `move`, `search`, `lsp`, `execute`, `plan`, `web_search`, `web_fetch`, `background_task_action`, `wait_tasks_action`, `kill_task_action`, `list`, `skill`, `memory_search`, `memory_get`, `task`, `agent_swarm`, `enter_plan`, `exit_plan`, `ask_user`, `deploy_app`, `search_tool`, `use_tool`, `monitor`, `goal_update`, `other`.",
"type": "string"
},
"ToolNamespace": {
crates/codegen/kigi-tools/src/computer/local/cgroup.rs
-1
View File
@@ -148,7 +148,6 @@ mod linux {
}
let err = std::io::Error::last_os_error();
match err.kind() {
// The fd is non-blocking: the queue is empty.
std::io::ErrorKind::WouldBlock => break,
std::io::ErrorKind::Interrupted => continue,
_ => return Err(err),
crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/mod.rs
+7
View File
@@ -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::{AGENT_SWARM_TOOL_NAME, AgentSwarmTool};
crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/plan.rs
+377
View File
@@ -0,0 +1,377 @@
//! Turning `agent_swarm` input into member specs, and member results back
//! into one tool result. Pure: no spawning, no I/O.
use kigi_tool_types::{
AgentSwarmToolInput, MAX_AGENT_SWARM_MEMBERS, PROMPT_TEMPLATE_PLACEHOLDER, SwarmMemberOutcome,
SwarmMemberResult,
};
/// One member's launch instructions.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MemberSpec {
/// The `items` entry, or the agent id for a resume — the label the caller
/// sees in the result and reuses to retry this exact member.
pub item: String,
pub prompt: String,
/// Set for a resume; `None` spawns a fresh member.
pub resume_from: Option<String>,
}
/// Everything wrong with the input is reported before anything is spawned:
/// a half-launched swarm is far more expensive to recover from than a refused
/// tool call.
pub fn plan_members(input: &AgentSwarmToolInput) -> Result<Vec<MemberSpec>, String> {
let resumes = input.resume_agent_ids.len();
if input.items.is_empty() && resumes == 0 {
return Err(
"agent_swarm needs `items` (with `prompt_template`) or `resume_agent_ids`; \
use the task tool for a single subagent"
.to_string(),
);
}
if resumes == 0 && input.items.len() < 2 {
return Err(
"agent_swarm runs 2 or more members; use the task tool for a single subagent"
.to_string(),
);
}
if input.items.len() + resumes > MAX_AGENT_SWARM_MEMBERS {
return Err(format!(
"agent_swarm runs at most {MAX_AGENT_SWARM_MEMBERS} members, got {}",
input.items.len() + resumes
));
}
// Models emit `""`/`"null"`/`"none"` where they mean "no id"; accepting one
// spawns a resume that can only die inside the coordinator.
if let Some(bad) = input
.resume_agent_ids
.keys()
.find(|id| !crate::implementations::kigi::task::types::is_valid_resume_id(id))
{
return Err(format!(
"`resume_agent_ids` key {bad:?} is not a subagent id; use the agent_id \
values from a previous agent_swarm result"
));
}
// Resumes first: they already hold context, so they reach the provider
// before the fresh members compete for the same rate-limit budget.
let mut specs: Vec<MemberSpec> = input
.resume_agent_ids
.iter()
.map(|(agent_id, prompt)| MemberSpec {
item: agent_id.clone(),
prompt: prompt.clone(),
resume_from: Some(agent_id.clone()),
})
.collect();
if !input.items.is_empty() {
let template = input
.prompt_template
.as_deref()
.map(str::trim)
.filter(|t| !t.is_empty())
.ok_or("agent_swarm requires `prompt_template` whenever `items` is given")?;
if !template.contains(PROMPT_TEMPLATE_PLACEHOLDER) {
return Err(format!(
"`prompt_template` must contain the literal {PROMPT_TEMPLATE_PLACEHOLDER}, \
which is replaced by each entry of `items`"
));
}
for item in &input.items {
if item.trim().is_empty() {
return Err("`items` entries must be non-empty".to_string());
}
specs.push(MemberSpec {
item: item.clone(),
prompt: template.replace(PROMPT_TEMPLATE_PLACEHOLDER, item),
resume_from: None,
});
}
}
// Two members handed the same prompt do the same work twice and, if it
// writes, race each other over the same files.
let mut seen = std::collections::HashSet::with_capacity(specs.len());
for spec in &specs {
if !seen.insert(spec.prompt.as_str()) {
return Err(format!(
"two members would run an identical prompt (from item {:?}); \
give every member a distinct scope",
spec.item
));
}
}
Ok(specs)
}
/// Per-member ceiling on rendered summary text.
///
/// 128 members' full outputs concatenated can exceed the context window at
/// exactly the moment the results matter. Native tool output is not truncated
/// anywhere downstream — only the MCP dispatcher does that — so the budget has
/// to live here. Each member keeps its head and its tail: the head says what it
/// did, the tail usually holds the verdict.
const MEMBER_SUMMARY_BUDGET: usize = 4_000;
/// Trims to [`MEMBER_SUMMARY_BUDGET`] on a char boundary, keeping both ends and
/// saying plainly how much was dropped.
fn clamp_summary(summary: &str) -> String {
if summary.len() <= MEMBER_SUMMARY_BUDGET {
return summary.to_string();
}
let keep = MEMBER_SUMMARY_BUDGET / 2;
let head_end = (0..=keep)
.rev()
.find(|i| summary.is_char_boundary(*i))
.unwrap_or(0);
let tail_start = (summary.len() - keep..summary.len())
.find(|i| summary.is_char_boundary(*i))
.unwrap_or(summary.len());
format!(
"{}\n… {} bytes omitted; read this member's full output with the task-output tool …\n{}",
&summary[..head_end],
summary.len() - head_end - (summary.len() - tail_start),
&summary[tail_start..]
)
}
/// Renders the fleet's results as one tool result.
///
/// A failed member is reported inside the aggregate rather than failing the
/// call: the caller needs the members that DID succeed, and needs to know
/// precisely which ones to retry.
pub fn render_results(results: &[SwarmMemberResult]) -> String {
let count = |wanted: SwarmMemberOutcome| results.iter().filter(|r| r.outcome == wanted).count();
let completed = count(SwarmMemberOutcome::Completed);
let failed = count(SwarmMemberOutcome::Failed);
let aborted = count(SwarmMemberOutcome::Aborted);
let backgrounded = count(SwarmMemberOutcome::Backgrounded);
let mut out = String::from("<agent_swarm_result>\n");
out.push_str(&format!(
"<summary>completed: {completed}, failed: {failed}, aborted: {aborted}, \
still running: {backgrounded}</summary>\n"
));
if backgrounded > 0 {
out.push_str(
"<still_running>Some members outlived the foreground budget and are still \
working. Do NOT re-launch their items — a second agent on the same files \
corrupts both. Poll them with the task-output tool instead.</still_running>\n",
);
}
if results.iter().any(SwarmMemberResult::is_resumable) {
out.push_str(
"<resume_hint>Call agent_swarm again with resume_agent_ids mapping the agent_id \
values below to a follow-up prompt to continue unfinished work.</resume_hint>\n",
);
}
for result in results {
out.push_str("<member");
if let Some(id) = &result.agent_id {
out.push_str(&format!(" agent_id=\"{}\"", escape_attr(id)));
}
out.push_str(&format!(
" item=\"{}\" state=\"{}\" outcome=\"{}\"",
escape_attr(&result.item),
if result.started() {
"started"
} else {
"not_started"
},
result.outcome.as_str()
));
if result.resumed {
out.push_str(" mode=\"resume\"");
}
out.push_str(">\n");
out.push_str(clamp_summary(result.summary.trim()).trim());
out.push_str("\n</member>\n");
}
out.push_str("</agent_swarm_result>");
out
}
/// Attribute-safe: an item is a model-supplied string and routinely contains
/// quotes or angle brackets (file paths, globs, shell fragments).
fn escape_attr(value: &str) -> String {
value
.replace('&', "&amp;")
.replace('<', "&lt;")
.replace('>', "&gt;")
.replace('"', "&quot;")
}
#[cfg(test)]
mod tests {
use super::*;
fn input(items: &[&str], template: Option<&str>) -> AgentSwarmToolInput {
AgentSwarmToolInput {
description: "test swarm".into(),
subagent_type: "general-purpose".into(),
prompt_template: template.map(str::to_string),
items: items.iter().map(|s| s.to_string()).collect(),
resume_agent_ids: Default::default(),
model: None,
}
}
#[test]
fn each_item_becomes_one_member_with_the_placeholder_substituted() {
let specs = plan_members(&input(&["a.rs", "b.rs"], Some("Review {{item}} for bugs")))
.expect("valid input");
assert_eq!(
specs,
vec![
MemberSpec {
item: "a.rs".into(),
prompt: "Review a.rs for bugs".into(),
resume_from: None,
},
MemberSpec {
item: "b.rs".into(),
prompt: "Review b.rs for bugs".into(),
resume_from: None,
},
]
);
}
#[test]
fn a_single_item_is_refused_in_favour_of_the_task_tool() {
let err = plan_members(&input(&["only.rs"], Some("Do {{item}}"))).unwrap_err();
assert!(err.contains("2 or more"), "{err}");
}
#[test]
fn a_template_without_the_placeholder_is_refused() {
let err = plan_members(&input(&["a", "b"], Some("Do the work"))).unwrap_err();
assert!(err.contains(PROMPT_TEMPLATE_PLACEHOLDER), "{err}");
}
#[test]
fn items_without_a_template_are_refused() {
let err = plan_members(&input(&["a", "b"], None)).unwrap_err();
assert!(err.contains("prompt_template"), "{err}");
}
#[test]
fn duplicate_expansions_are_refused_before_anything_spawns() {
let err = plan_members(&input(&["same", "same"], Some("Do {{item}}"))).unwrap_err();
assert!(err.contains("identical prompt"), "{err}");
}
#[test]
fn more_members_than_the_ceiling_are_refused() {
let items: Vec<String> = (0..=MAX_AGENT_SWARM_MEMBERS)
.map(|i| format!("item-{i}"))
.collect();
let mut spec = input(&[], Some("Do {{item}}"));
spec.items = items;
let err = plan_members(&spec).unwrap_err();
assert!(err.contains(&MAX_AGENT_SWARM_MEMBERS.to_string()), "{err}");
}
#[test]
fn resumes_are_planned_before_fresh_members() {
let mut spec = input(&["fresh.rs"], Some("Do {{item}}"));
spec.resume_agent_ids
.insert("agent-1".into(), "keep going".into());
let specs = plan_members(&spec).expect("a resume lifts the two-item floor");
assert_eq!(specs[0].resume_from.as_deref(), Some("agent-1"));
assert_eq!(specs[0].prompt, "keep going");
assert_eq!(specs[1].item, "fresh.rs");
}
#[test]
fn an_empty_call_is_refused() {
let err = plan_members(&input(&[], None)).unwrap_err();
assert!(err.contains("resume_agent_ids"), "{err}");
}
fn member(item: &str, outcome: SwarmMemberOutcome, id: Option<&str>) -> SwarmMemberResult {
SwarmMemberResult {
item: item.into(),
agent_id: id.map(str::to_string),
resumed: false,
outcome,
summary: format!("summary for {item}"),
}
}
#[test]
fn the_aggregate_counts_outcomes_and_labels_every_member() {
let rendered = render_results(&[
member("a.rs", SwarmMemberOutcome::Completed, Some("id-a")),
member("b.rs", SwarmMemberOutcome::Failed, Some("id-b")),
]);
assert!(rendered.contains("completed: 1, failed: 1, aborted: 0, still running: 0"));
assert!(
rendered.contains(r#"agent_id="id-a" item="a.rs" state="started" outcome="completed""#)
);
assert!(rendered.contains(r#"outcome="failed""#));
assert!(rendered.contains("summary for b.rs"));
}
#[test]
fn the_resume_hint_appears_only_when_something_resumable_is_unfinished() {
let all_done = render_results(&[member("a", SwarmMemberOutcome::Completed, Some("id-a"))]);
assert!(!all_done.contains("resume_hint"));
let never_started = render_results(&[member("a", SwarmMemberOutcome::Aborted, None)]);
assert!(
!never_started.contains("resume_hint"),
"a member with no agent_id has nothing to resume"
);
let retryable = render_results(&[member("a", SwarmMemberOutcome::Failed, Some("id-a"))]);
assert!(retryable.contains("resume_hint"));
}
#[test]
fn an_item_containing_markup_cannot_break_out_of_its_attribute() {
let rendered = render_results(&[member(
r#"a" onload="x"#,
SwarmMemberOutcome::Completed,
None,
)]);
assert!(!rendered.contains(r#"item="a" onload="#), "{rendered}");
assert!(rendered.contains("&quot;"), "{rendered}");
}
}
#[cfg(test)]
mod budget_tests {
use super::*;
#[test]
fn a_huge_member_output_is_clamped_but_keeps_both_ends() {
let summary = format!("HEAD{}TAIL", "x".repeat(MEMBER_SUMMARY_BUDGET * 2));
let clamped = clamp_summary(&summary);
assert!(clamped.len() < summary.len(), "must shrink");
assert!(clamped.starts_with("HEAD"), "the head says what it did");
assert!(clamped.ends_with("TAIL"), "the tail holds the verdict");
assert!(clamped.contains("bytes omitted"), "the loss must be stated");
}
#[test]
fn a_summary_inside_the_budget_is_untouched() {
assert_eq!(clamp_summary("short"), "short");
}
#[test]
fn clamping_never_splits_a_character() {
// Multi-byte throughout, so a naive byte slice would panic.
let summary = "café ".repeat(MEMBER_SUMMARY_BUDGET);
let clamped = clamp_summary(&summary);
assert!(clamped.len() < summary.len());
}
}
crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/run.rs
+811
View File
@@ -0,0 +1,811 @@
//! 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(),
// The ITEM, not the swarm's description: every member would otherwise
// render as an identical subagent block and the user could not tell
// which one is running, or which one failed.
description: pending.spec.item.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
);
}
/// Each member must be identifiable while it runs: the TUI renders one
/// subagent block per member from this description, so a shared one leaves
/// the user staring at N identical rows.
#[tokio::test(start_paused = true)]
async fn each_member_is_labelled_with_its_own_item() {
let backend = Arc::new(FakeBackend::default());
let seen = Arc::new(Mutex::new(Vec::new()));
let recorder = seen.clone();
let specs = specs(&["a.rs", "b.rs"]);
let config = SwarmRunConfig {
description: "review files".into(),
..config(None)
};
// `build_request` is the only place the label is set, so assert on it
// directly rather than through the backend's prompt log.
for (index, spec) in specs.iter().enumerate() {
let pending = Pending {
index,
spec: spec.clone(),
attempts: 0,
not_before: None,
};
recorder
.lock()
.unwrap()
.push(build_request(&pending, &config).description);
}
assert_eq!(
*seen.lock().unwrap(),
vec!["a.rs".to_string(), "b.rs".to_string()],
"each member must carry its own item, not the swarm description"
);
drop(backend);
}
/// 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);
}
}
crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/schedule.rs
+348
View File
@@ -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());
}
}
crates/codegen/kigi-tools/src/implementations/kigi/agent_swarm/tool.rs
+502
View File
@@ -0,0 +1,502 @@
//! 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.";
/// Registry name, so gating code matches on one definition rather than a
/// literal that can drift from the tool id.
pub const AGENT_SWARM_TOOL_NAME: &str = "agent_swarm";
#[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_TOOL_NAME).expect("valid tool id")
}
fn description(
&self,
_ctx: &::kigi_tool_runtime::ListToolsContext,
) -> kigi_tool_types::ToolDescription {
kigi_tool_types::ToolDescription::new(AGENT_SWARM_TOOL_NAME, 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}");
}
}
crates/codegen/kigi-tools/src/implementations/kigi/mod.rs
+2
View File
@@ -1,5 +1,6 @@
//! Tool implementations built on the `NewTool` trait; the sibling
//! `implementations/<tool>/` modules hold the `Tool`-trait counterparts.
pub mod agent_swarm;
pub mod ask_user_question;
pub mod bash;
#[path = "deploy_app_stub.rs"]
@@ -21,6 +22,7 @@ pub mod todo;
pub mod update_goal;
pub mod web_fetch;
pub mod web_search;
pub use agent_swarm::{AGENT_SWARM_TOOL_NAME, AgentSwarmTool};
pub use ask_user_question::AskUserQuestionTool;
pub use bash::BashTool;
pub use deploy_app::{AppBuilderDeployerConfig, DEPLOY_APP_TOOL_NAME};
crates/codegen/kigi-tools/src/implementations/kigi/task/types.rs
+14 -1
View File
@@ -149,10 +149,12 @@ pub trait SubagentCapabilityModeExt {
pub fn prune_orphaned_background_task_tools(config: &mut crate::registry::types::ToolServerConfig) {
use crate::types::tool::ToolKind;
// A swarm member can be auto-backgrounded by the coordinator just like a
// `task` child, so either spawner keeps the lifecycle tools alive.
let has_task_tool = config
.tools
.iter()
.any(|tc| tc.kind == Some(ToolKind::Task));
.any(|tc| matches!(tc.kind, Some(ToolKind::Task | ToolKind::AgentSwarm)));
let has_background_capable_bash = config.tools.iter().any(is_background_capable_bash_tool);
if has_task_tool || has_background_capable_bash {
return;
@@ -209,6 +211,7 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode {
ToolKind::BackgroundTaskAction,
ToolKind::KillTaskAction,
ToolKind::Task,
ToolKind::AgentSwarm,
ToolKind::EnterPlan,
ToolKind::ExitPlan,
ToolKind::AskUser,
@@ -232,6 +235,7 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode {
ToolKind::BackgroundTaskAction,
ToolKind::KillTaskAction,
ToolKind::Task,
ToolKind::AgentSwarm,
ToolKind::EnterPlan,
ToolKind::ExitPlan,
ToolKind::AskUser,
@@ -252,6 +256,7 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode {
ToolKind::BackgroundTaskAction,
ToolKind::KillTaskAction,
ToolKind::Task,
ToolKind::AgentSwarm,
ToolKind::EnterPlan,
ToolKind::ExitPlan,
ToolKind::AskUser,
@@ -276,6 +281,7 @@ impl SubagentCapabilityModeExt for SubagentCapabilityMode {
ToolKind::BackgroundTaskAction,
ToolKind::KillTaskAction,
ToolKind::Task,
ToolKind::AgentSwarm,
ToolKind::EnterPlan,
ToolKind::ExitPlan,
ToolKind::AskUser,
@@ -315,6 +321,12 @@ pub struct SubagentResult {
/// `get_command_or_subagent_output`), so the tool returns a `task_id` notice
/// instead of a completion. Never set for natively backgrounded subagents.
pub backgrounded: bool,
/// The child's turn ended because the PROVIDER refused it for rate
/// limiting, not because the work failed. Classified where the typed ACP
/// error code is still in hand: a fleet scheduler that had to re-derive
/// this from the formatted `error` string would silently stop adapting
/// the day that wording changes.
pub rate_limited: bool,
}
impl Default for SubagentResult {
@@ -324,6 +336,7 @@ impl Default for SubagentResult {
output: Arc::from(""),
error: None,
cancelled: false,
rate_limited: false,
subagent_id: String::new(),
child_session_id: String::new(),
tool_calls: 0,
crates/codegen/kigi-tools/src/implementations/kigi/web_fetch/client.rs
+140 -16
View File
@@ -155,6 +155,9 @@ impl WebFetchClient {
}
}
// Before any egress: the service must not see this.
ssrf::check_ssrf(&url, self.params.allow_local()).await?;
// Kimi fetch service first (OAuth sessions); local pipeline is the
// fallback on any service failure (kimi-cli fetch.py `__call__`).
if let Some(service_url) = self.params.service_url.clone() {
@@ -182,10 +185,15 @@ impl WebFetchClient {
}
}
ssrf::check_ssrf(&url).await?;
let http = self.http.get_or_rebuild()?;
let result = match fetch_url(&http, &url, self.params.max_content_length()).await {
let result = match fetch_url(
&http,
&url,
self.params.max_content_length(),
self.params.allow_local(),
)
.await
{
Ok(result) => result,
Err(e @ WebFetchError::HttpRequest(_)) => {
self.http.invalidate();
@@ -380,6 +388,8 @@ fn validate_url(raw: &str) -> Result<Url, WebFetchError> {
if let Some(host) = parsed.host_str()
&& host.split('.').count() < 2
// `localhost` is single-label; SSRF still gates it on allow_local.
&& !ssrf::is_explicit_local_host(host)
{
return Err(WebFetchError::SingleLabelHost {
host: host.to_string(),
@@ -390,9 +400,16 @@ fn validate_url(raw: &str) -> Result<Url, WebFetchError> {
}
fn upgrade_to_https(url: &mut Url) {
if url.scheme() == "http" {
let _ = url.set_scheme("https");
if url.scheme() != "http" {
return;
}
// Local dev servers rarely serve TLS; SSRF still gates them.
if let Some(host) = url.host_str()
&& ssrf::is_explicit_local_host(host)
{
return;
}
let _ = url.set_scheme("https");
}
enum FetchResult {
@@ -409,15 +426,21 @@ enum FetchResult {
}
/// Fetch a URL with manual same-host redirect handling.
///
/// Every hop is re-checked, so a rebinding name cannot pass.
/// Partial: reqwest hides the peer IP of the live connection.
async fn fetch_url(
client: &reqwest::Client,
url: &Url,
max_content_length: usize,
allow_local: bool,
) -> Result<FetchResult, WebFetchError> {
let mut current_url = url.clone();
let mut hops = 0;
loop {
ssrf::check_ssrf(&current_url, allow_local).await?;
let resp = client
.get(current_url.as_str())
.header(USER_AGENT, USER_AGENT_STRING)
@@ -439,10 +462,13 @@ async fn fetch_url(
if let Some(location) = resp.headers().get("location") {
let location_str = location.to_str().unwrap_or("");
let next_url = current_url
let mut next_url = current_url
.join(location_str)
.map_err(|e| WebFetchError::InvalidRedirect(format!("{e}")))?;
if is_same_host(&current_url, &next_url) {
// An absolute `http://` Location would downgrade the hop.
upgrade_to_https(&mut next_url);
// check_ssrf runs at the top of the next iteration.
current_url = next_url;
continue;
}
@@ -479,13 +505,11 @@ async fn fetch_url(
}
}
/// Exact host equality, no `www.` stripping.
///
/// A `www` sibling has separate records, so it is cross-host.
fn is_same_host(a: &Url, b: &Url) -> bool {
fn strip_www(h: &str) -> &str {
h.strip_prefix("www.").unwrap_or(h)
}
let host_a = a.host_str().unwrap_or("");
let host_b = b.host_str().unwrap_or("");
strip_www(host_a) == strip_www(host_b)
a.host_str() == b.host_str()
}
fn require_media_session_folder(session_folder: Option<&Path>) -> Result<&Path, WebFetchError> {
@@ -963,6 +987,78 @@ mod tests {
assert!(matches!(err, WebFetchError::ServiceUnavailable(_)), "{err}");
}
/// A blocked target must never reach the remote fetch service.
///
/// It egresses elsewhere, so posting leaks an internal URL.
#[tokio::test]
async fn a_blocked_url_never_reaches_the_fetch_service() {
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/fetch"))
.respond_with(ResponseTemplate::new(200).set_body_string("{}"))
.mount(&server)
.await;
let provider = crate::types::api_key_provider::test_support::fixed_provider("t");
let params = WebFetchParams {
service_url: Some(format!("{}/fetch", server.uri())),
..WebFetchParams::default()
};
let client = WebFetchClient::new(&params, Some(provider)).unwrap();
let Err(err) = client
.fetch(
"http://localhost:8080/admin?token=secret",
"c",
None,
None,
None,
)
.await
else {
panic!("a loopback target must be blocked");
};
assert!(matches!(err, WebFetchError::SsrfBlocked { .. }), "{err}");
assert!(
server.received_requests().await.unwrap().is_empty(),
"the internal URL must never be posted anywhere"
);
}
/// `fetch_url` gates its own target, not trusting its caller.
///
/// Only hop one is covered: same-host hops share one verdict,
/// so rebinding between them needs a live resolver to observe.
#[tokio::test]
async fn fetch_url_blocks_a_loopback_target_without_allow_local() {
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/x"))
.respond_with(ResponseTemplate::new(200).set_body_string("<p>local</p>"))
.mount(&server)
.await;
let url = Url::parse(&format!("{}/x", server.uri())).unwrap();
let http = HttpClient::new(&WebFetchParams::default())
.unwrap()
.get_or_rebuild()
.unwrap();
let Err(err) = fetch_url(&http, &url, 1_000_000, false).await else {
panic!("a blocked host must not be fetched");
};
assert!(matches!(err, WebFetchError::SsrfBlocked { .. }), "{err}");
assert!(
server.received_requests().await.unwrap().is_empty(),
"a blocked host must not be contacted at all"
);
let ok = fetch_url(&http, &url, 1_000_000, true).await.unwrap();
assert!(matches!(ok, FetchResult::Content { .. }));
}
fn test_converter() -> htmd::HtmlToMarkdown {
htmd::HtmlToMarkdown::builder()
.skip_tags(vec![
@@ -1024,9 +1120,12 @@ mod tests {
#[test]
fn validate_url_rejects_single_label_hosts() {
assert!(validate_url("http://localhost:8080/foo").is_err());
assert!(validate_url("http://intranet/foo").is_err());
assert!(validate_url("http://metadata/computeMetadata").is_err());
assert!(
validate_url("http://localhost:8080/foo").is_ok(),
"localhost reaches the SSRF gate, which blocks it unless opted in"
);
}
#[test]
@@ -1068,6 +1167,19 @@ mod tests {
assert_eq!(url.scheme(), "https");
}
/// Upgrading a local host breaks the only target `allow_local` opens.
#[test]
fn upgrade_to_https_skips_explicit_local_hosts() {
for raw in ["http://127.0.0.1:8080/", "http://localhost:3000/"] {
let mut url = Url::parse(raw).unwrap();
upgrade_to_https(&mut url);
assert_eq!(url.scheme(), "http", "{raw}");
}
let mut public = Url::parse("http://example.com/").unwrap();
upgrade_to_https(&mut public);
assert_eq!(public.scheme(), "https");
}
#[test]
fn same_host_exact_match() {
let a = Url::parse("https://example.com/a").unwrap();
@@ -1075,12 +1187,24 @@ mod tests {
assert!(is_same_host(&a, &b));
}
/// An absolute `http://` Location must not downgrade a followed hop.
#[test]
fn same_host_www_stripping() {
fn same_host_redirect_location_reupgrades_http() {
let origin = Url::parse("https://example.com/start").unwrap();
let mut next = origin.join("http://example.com/next").unwrap();
assert_eq!(next.scheme(), "http");
assert!(is_same_host(&origin, &next));
upgrade_to_https(&mut next);
assert_eq!(next.as_str(), "https://example.com/next");
}
/// A `www` sibling is a separate name, with separate records.
#[test]
fn www_subdomain_is_cross_host() {
let a = Url::parse("https://example.com/a").unwrap();
let c = Url::parse("https://www.example.com/a").unwrap();
assert!(is_same_host(&a, &c));
assert!(is_same_host(&c, &a));
assert!(!is_same_host(&a, &c));
assert!(!is_same_host(&c, &a));
}
#[test]
crates/codegen/kigi-tools/src/implementations/kigi/web_fetch/config.rs
+8
View File
@@ -47,12 +47,20 @@ pub struct WebFetchParams {
/// on any failure (kimi-cli `tools/web/fetch.py FetchURL.__call__`).
#[serde(default)]
pub service_url: Option<String>,
/// Opt-in for loopback targets; off means no local access.
#[serde(default)]
pub allow_local: Option<bool>,
}
register_resource!("kigi", "WebFetch", WebFetchParams);
// Keep defaults here so call-sites don't have to manage unwrapping.
impl WebFetchParams {
/// From config or `KIGI_WEB_FETCH_ALLOW_LOCAL`, never tool input.
pub fn allow_local(&self) -> bool {
self.allow_local.unwrap_or(false)
}
pub fn cache_ttl_secs(&self) -> Duration {
Duration::from_secs(self.cache_ttl_secs.unwrap_or(15 * 60))
}
crates/codegen/kigi-tools/src/implementations/kigi/web_fetch/mod.rs
+1 -1
View File
@@ -14,7 +14,7 @@ pub mod domain;
pub mod error;
mod http;
pub(crate) mod overflow;
mod ssrf;
pub mod ssrf;
pub use client::WebFetchClient;
pub use config::WebFetchParams;
crates/codegen/kigi-tools/src/implementations/kigi/web_fetch/ssrf.rs
+284 -99
View File
@@ -1,85 +1,143 @@
//! SSRF (Server-Side Request Forgery) protection for `web_fetch`.
//! SSRF protection for `web_fetch`.
//!
//! Validates that resolved IP addresses are not in private, link-local, or
//! cloud metadata ranges before allowing outbound HTTP requests.
//! Non-public targets are blocked: loopback, RFC 1918, link-local,
//! CGNAT, TEST-NET, reserved, ULA. Loopback is opt-in via
//! `[toolset.web_fetch] allow_local` or `KIGI_WEB_FETCH_ALLOW_LOCAL`,
//! and even then only for a literal local host.
//!
//! Reference: [IANA IPv4 Special-Purpose Address Registry](https://www.iana.org/assignments/iana-ipv4-special-registry/)
use std::net::IpAddr;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use url::Url;
use super::error::WebFetchError;
/// Returns `true` if an IP address is in a private, link-local, or cloud
/// metadata range that should be blocked to prevent SSRF attacks.
/// Hosts allowed to reach loopback when local access is on.
///
/// **Allowed:** loopback (`127.x` / `::1`) for local development.
/// **Blocked:** RFC 1918, link-local, CGNAT/cloud metadata, unspecified.
pub(crate) fn is_blocked_ip(ip: &IpAddr) -> bool {
/// Names that merely RESOLVE to loopback are excluded: DNS rebinding.
pub fn is_explicit_local_host(host: &str) -> bool {
let host = host.trim().trim_end_matches('.').to_ascii_lowercase();
let host = host
.strip_prefix('[')
.and_then(|h| h.strip_suffix(']'))
.unwrap_or(&host);
// Drop an IPv6 zone id such as `fe80::1%lo0`.
let host = host.split('%').next().unwrap_or(host);
if host == "localhost" {
return true;
}
host.parse::<IpAddr>().is_ok_and(|ip| ip.is_loopback())
}
/// Whether an IP is not globally routable.
pub fn is_non_public_ip(ip: &IpAddr) -> bool {
match ip {
IpAddr::V4(v4) => {
let octets = v4.octets();
// Loopback (127.0.0.0/8) — allowed for local dev servers.
if octets[0] == 127 {
return false;
}
// RFC 1918: 10.0.0.0/8 — private network.
if octets[0] == 10 {
return true;
}
// RFC 1918: 172.16.0.0/12 — private network.
if octets[0] == 172 && (16..=31).contains(&octets[1]) {
return true;
}
// RFC 1918: 192.168.0.0/16 — private network.
if octets[0] == 192 && octets[1] == 168 {
return true;
}
// RFC 3927: 169.254.0.0/16 — link-local.
// Includes AWS/GCP/Azure metadata endpoint 169.254.169.254.
if octets[0] == 169 && octets[1] == 254 {
return true;
}
// RFC 6598: 100.64.0.0/10 — CGNAT / shared address space.
// Used by some cloud providers for internal metadata services.
if octets[0] == 100 && (64..=127).contains(&octets[1]) {
return true;
}
if v4.is_unspecified() {
return true;
}
false
}
IpAddr::V4(v4) => is_non_public_ipv4(*v4),
IpAddr::V6(v6) => is_non_public_ipv6(*v6),
}
}
fn is_non_public_ipv4(ip: Ipv4Addr) -> bool {
ip.is_loopback()
|| ip.is_private()
|| ip.is_link_local()
|| ip.is_unspecified()
|| ip.is_multicast()
|| ip.is_broadcast()
// "This network" (RFC 1122) 0.0.0.0/8
|| ipv4_in_cidr(ip, [0, 0, 0, 0], 8)
// CGNAT (RFC 6598) — some clouds serve metadata here
|| ipv4_in_cidr(ip, [100, 64, 0, 0], 10)
// IETF Protocol Assignments (RFC 6890)
|| ipv4_in_cidr(ip, [192, 0, 0, 0], 24)
// TEST-NET-1 (RFC 5737)
|| ipv4_in_cidr(ip, [192, 0, 2, 0], 24)
// Benchmarking (RFC 2544)
|| ipv4_in_cidr(ip, [198, 18, 0, 0], 15)
// TEST-NET-2 / TEST-NET-3
|| ipv4_in_cidr(ip, [198, 51, 100, 0], 24)
|| ipv4_in_cidr(ip, [203, 0, 113, 0], 24)
// Reserved (RFC 6890)
|| ipv4_in_cidr(ip, [240, 0, 0, 0], 4)
}
fn ipv4_in_cidr(ip: Ipv4Addr, base: [u8; 4], prefix: u8) -> bool {
debug_assert!(prefix <= 32, "IPv4 prefix out of range");
let ip = u32::from(ip);
let base = u32::from(Ipv4Addr::from(base));
let mask = if prefix == 0 {
0
} else {
u32::MAX << (32 - prefix)
};
(ip & mask) == (base & mask)
}
fn is_non_public_ipv6(ip: Ipv6Addr) -> bool {
// Identity wins: `::1` is not judged as `0.0.0.1`.
if ip.is_loopback() || ip.is_unspecified() || ip.is_multicast() {
return true;
}
if let Some(v4) = embedded_ipv4(ip) {
return is_non_public_ipv4(v4);
}
let seg = ip.segments();
ip.is_unique_local()
|| ip.is_unicast_link_local()
// Deprecated site-local (RFC 3879) fec0::/10
|| (seg[0] & 0xffc0) == 0xfec0
// Documentation (RFC 3849) 2001:db8::/32
|| (seg[0] == 0x2001 && seg[1] == 0x0db8)
}
/// IPv4 reachable through a known IPv6 wrapper, if any.
///
/// Covers mapped, compatible, well-known NAT64, and 6to4. Not complete:
/// network-specific NAT64 prefixes (RFC 6052) cannot be enumerated.
fn embedded_ipv4(ip: Ipv6Addr) -> Option<Ipv4Addr> {
let seg = ip.segments();
let embedded = |hi: u16, lo: u16| Ipv4Addr::from(u32::from(hi) << 16 | u32::from(lo));
if seg[0] == 0x0064 && seg[1] == 0xff9b && seg[2..6] == [0, 0, 0, 0] {
return Some(embedded(seg[6], seg[7]));
}
if seg[0] == 0x2002 {
return Some(embedded(seg[1], seg[2]));
}
// Covers `::ffff:a.b.c.d` and the deprecated `::a.b.c.d`.
ip.to_ipv4()
}
/// Loopback including IPv4-mapped forms like `::ffff:127.0.0.1`.
///
/// `IpAddr::is_loopback` is false for mapped addresses, so the opt-in path
/// cannot use it directly.
fn is_loopback_addr(ip: &IpAddr) -> bool {
match ip {
IpAddr::V4(v4) => v4.is_loopback(),
IpAddr::V6(v6) => {
// ::1 — loopback, allowed for local dev.
if v6.is_loopback() {
return false;
}
if v6.is_unspecified() {
return true;
}
// IPv4-mapped IPv6 (::ffff:x.x.x.x) — delegate to v4 checks.
if let Some(v4) = v6.to_ipv4_mapped() {
return is_blocked_ip(&IpAddr::V4(v4));
}
let segments = v6.segments();
// RFC 4291: fe80::/10 — link-local unicast.
if segments[0] & 0xffc0 == 0xfe80 {
return true;
}
// RFC 4193: fc00::/7 — unique local address (ULA).
if segments[0] & 0xfe00 == 0xfc00 {
return true;
}
false
v6.is_loopback() || v6.to_ipv4_mapped().is_some_and(|v4| v4.is_loopback())
}
}
}
/// Resolve hostname via DNS and verify none of the resolved addresses are
/// in blocked private/link-local ranges.
pub(crate) async fn check_ssrf(url: &Url) -> Result<(), WebFetchError> {
/// Dual gate: loopback opens only for an explicit local host.
///
/// Private and link-local never open through this flag.
/// Shared with the hook runner: one policy, every outbound URL.
pub fn is_blocked_for_host(ip: &IpAddr, host: &str, allow_local: bool) -> bool {
if !is_non_public_ip(ip) {
return false;
}
!(allow_local && is_loopback_addr(ip) && is_explicit_local_host(host))
}
/// Verifies no resolved address is blocked by the SSRF policy.
///
/// `allow_local` is config-only so the model cannot flip it.
pub(crate) async fn check_ssrf(url: &Url, allow_local: bool) -> Result<(), WebFetchError> {
let host = url
.host_str()
.ok_or_else(|| WebFetchError::SingleLabelHost {
@@ -87,7 +145,7 @@ pub(crate) async fn check_ssrf(url: &Url) -> Result<(), WebFetchError> {
})?;
if let Ok(ip) = host.parse::<IpAddr>() {
if is_blocked_ip(&ip) {
if is_blocked_for_host(&ip, host, allow_local) {
return Err(WebFetchError::SsrfBlocked {
host: host.to_string(),
ip,
@@ -112,7 +170,7 @@ pub(crate) async fn check_ssrf(url: &Url) -> Result<(), WebFetchError> {
addrs
.iter()
.find(|addr| is_blocked_ip(&addr.ip()))
.find(|addr| is_blocked_for_host(&addr.ip(), host, allow_local))
.map_or(Ok(()), |addr| {
Err(WebFetchError::SsrfBlocked {
host: host.to_string(),
@@ -127,86 +185,201 @@ mod tests {
#[test]
fn blocks_rfc1918_10x() {
assert!(is_blocked_ip(&"10.0.0.1".parse().unwrap()));
assert!(is_blocked_ip(&"10.255.255.255".parse().unwrap()));
assert!(is_non_public_ip(&"10.0.0.1".parse().unwrap()));
assert!(is_non_public_ip(&"10.255.255.255".parse().unwrap()));
}
#[test]
fn blocks_rfc1918_172x() {
assert!(is_blocked_ip(&"172.16.0.1".parse().unwrap()));
assert!(is_blocked_ip(&"172.31.255.255".parse().unwrap()));
assert!(!is_blocked_ip(&"172.15.0.1".parse().unwrap()));
assert!(!is_blocked_ip(&"172.32.0.1".parse().unwrap()));
assert!(is_non_public_ip(&"172.16.0.1".parse().unwrap()));
assert!(is_non_public_ip(&"172.31.255.255".parse().unwrap()));
assert!(!is_non_public_ip(&"172.15.0.1".parse().unwrap()));
assert!(!is_non_public_ip(&"172.32.0.1".parse().unwrap()));
}
#[test]
fn blocks_rfc1918_192168() {
assert!(is_blocked_ip(&"192.168.0.1".parse().unwrap()));
assert!(is_blocked_ip(&"192.168.255.255".parse().unwrap()));
assert!(is_non_public_ip(&"192.168.0.1".parse().unwrap()));
assert!(is_non_public_ip(&"192.168.255.255".parse().unwrap()));
}
#[test]
fn blocks_link_local() {
assert!(is_blocked_ip(&"169.254.0.1".parse().unwrap()));
assert!(is_blocked_ip(&"169.254.169.254".parse().unwrap()));
assert!(is_non_public_ip(&"169.254.0.1".parse().unwrap()));
assert!(is_non_public_ip(&"169.254.169.254".parse().unwrap()));
}
#[test]
fn blocks_cgnat_cloud_metadata() {
assert!(is_blocked_ip(&"100.64.0.1".parse().unwrap()));
assert!(is_blocked_ip(&"100.127.255.255".parse().unwrap()));
assert!(!is_blocked_ip(&"100.63.0.1".parse().unwrap()));
assert!(!is_blocked_ip(&"100.128.0.1".parse().unwrap()));
assert!(is_non_public_ip(&"100.64.0.1".parse().unwrap()));
assert!(is_non_public_ip(&"100.127.255.255".parse().unwrap()));
assert!(!is_non_public_ip(&"100.63.0.1".parse().unwrap()));
assert!(!is_non_public_ip(&"100.128.0.1".parse().unwrap()));
}
#[test]
fn blocks_unspecified() {
assert!(is_blocked_ip(&"0.0.0.0".parse().unwrap()));
assert!(is_blocked_ip(&"::".parse().unwrap()));
assert!(is_non_public_ip(&"0.0.0.0".parse().unwrap()));
assert!(is_non_public_ip(&"::".parse().unwrap()));
}
#[test]
fn allows_loopback() {
assert!(!is_blocked_ip(&"127.0.0.1".parse().unwrap()));
assert!(!is_blocked_ip(&"127.0.0.2".parse().unwrap()));
assert!(!is_blocked_ip(&"::1".parse().unwrap()));
fn blocks_loopback_by_default() {
for ip in ["127.0.0.1", "127.0.0.2", "::1", "::ffff:127.0.0.1"] {
let ip: IpAddr = ip.parse().unwrap();
assert!(is_non_public_ip(&ip), "{ip} must not be public");
assert!(
is_blocked_for_host(&ip, "localhost", false),
"{ip} must be blocked without allow_local"
);
}
}
#[test]
fn allow_local_opens_loopback_only_for_an_explicit_local_host() {
for (ip, host) in [
("127.0.0.1", "localhost"),
("127.0.0.1", "127.0.0.1"),
("::1", "::1"),
("::ffff:127.0.0.1", "localhost"),
] {
let ip: IpAddr = ip.parse().unwrap();
assert!(!is_blocked_for_host(&ip, host, true), "{ip} via {host}");
}
assert!(
is_blocked_for_host(&"127.0.0.1".parse().unwrap(), "evil.example.com", true),
"a public name resolving to loopback is DNS rebinding"
);
}
#[test]
fn allow_local_never_opens_private_or_link_local() {
for ip in ["10.0.0.1", "169.254.169.254", "192.168.1.1"] {
let ip: IpAddr = ip.parse().unwrap();
assert!(
is_blocked_for_host(&ip, "localhost", true),
"{ip} must stay blocked even with allow_local"
);
}
}
#[test]
fn blocks_test_net_and_reserved_ranges() {
for ip in [
"0.0.0.1",
"192.0.0.1",
"192.0.2.1",
"198.18.0.1",
"198.19.255.255",
"198.51.100.1",
"203.0.113.1",
"240.0.0.1",
] {
assert!(is_non_public_ip(&ip.parse().unwrap()), "{ip}");
}
// Neighbours of every range above must stay reachable.
for ip in [
"1.0.0.1",
"192.0.1.1",
"192.0.3.1",
"198.17.255.255",
"198.20.0.1",
"198.51.101.1",
"203.0.114.1",
"223.255.255.255",
] {
assert!(!is_non_public_ip(&ip.parse().unwrap()), "{ip}");
}
}
/// A v6 record can smuggle v4 through four wrapper prefixes.
#[test]
fn blocks_ipv4_smuggled_through_ipv6_wrappers() {
for ip in [
"64:ff9b::a9fe:a9fe",
"64:ff9b::7f00:1",
"2002:7f00:1::",
"::7f00:1",
"::a00:1",
"fec0::1",
"2001:db8::1",
] {
assert!(is_non_public_ip(&ip.parse().unwrap()), "{ip}");
}
for ip in ["64:ff9b::808:808", "2002:808:808::", "2001:db9::1"] {
assert!(!is_non_public_ip(&ip.parse().unwrap()), "{ip}");
}
}
#[test]
fn explicit_local_host_tolerates_brackets_dots_and_zone_ids() {
for host in [
"localhost",
"LOCALHOST.",
"127.0.0.1",
"127.1.2.3",
"::1",
"[::1]",
"::1%lo0",
] {
assert!(is_explicit_local_host(host), "{host}");
}
for host in [
"example.com",
"notlocalhost",
"localhost.evil.com",
"10.0.0.1",
] {
assert!(!is_explicit_local_host(host), "{host}");
}
}
#[test]
fn allows_public_ips() {
assert!(!is_blocked_ip(&"1.1.1.1".parse().unwrap()));
assert!(!is_blocked_ip(&"8.8.8.8".parse().unwrap()));
assert!(!is_blocked_ip(&"142.250.80.46".parse().unwrap()));
for ip in [
"1.1.1.1",
"8.8.8.8",
"142.250.80.46",
// Global unicast v6: guards the new masks against over-matching.
"2606:4700::1111",
"2001:4860:4860::8888",
] {
assert!(!is_non_public_ip(&ip.parse().unwrap()), "{ip}");
}
}
#[test]
fn blocks_ipv6_link_local() {
assert!(is_blocked_ip(&"fe80::1".parse().unwrap()));
assert!(is_non_public_ip(&"fe80::1".parse().unwrap()));
}
#[test]
fn blocks_ipv6_unique_local() {
assert!(is_blocked_ip(&"fc00::1".parse().unwrap()));
assert!(is_blocked_ip(&"fd00::1".parse().unwrap()));
assert!(is_non_public_ip(&"fc00::1".parse().unwrap()));
assert!(is_non_public_ip(&"fd00::1".parse().unwrap()));
}
#[test]
fn blocks_ipv4_mapped_ipv6_private() {
assert!(is_blocked_ip(&"::ffff:10.0.0.1".parse::<IpAddr>().unwrap()));
assert!(is_blocked_ip(
assert!(is_non_public_ip(
&"::ffff:10.0.0.1".parse::<IpAddr>().unwrap()
));
assert!(is_non_public_ip(
&"::ffff:192.168.1.1".parse::<IpAddr>().unwrap()
));
}
#[test]
fn allows_ipv4_mapped_ipv6_public() {
assert!(!is_blocked_ip(&"::ffff:8.8.8.8".parse::<IpAddr>().unwrap()));
assert!(!is_non_public_ip(
&"::ffff:8.8.8.8".parse::<IpAddr>().unwrap()
));
}
#[tokio::test]
async fn ssrf_blocks_ip_literal_private() {
let url = Url::parse("https://10.0.0.1/secret").unwrap();
let result = check_ssrf(&url).await;
let result = check_ssrf(&url, false).await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("private"));
}
@@ -214,7 +387,19 @@ mod tests {
#[tokio::test]
async fn ssrf_allows_ip_literal_public() {
let url = Url::parse("https://1.1.1.1/").unwrap();
let result = check_ssrf(&url).await;
let result = check_ssrf(&url, false).await;
assert!(result.is_ok());
}
#[tokio::test]
async fn ssrf_blocks_loopback_literal_by_default() {
let url = Url::parse("http://127.0.0.1:8080/").unwrap();
assert!(check_ssrf(&url, false).await.is_err());
}
#[tokio::test]
async fn ssrf_allows_loopback_literal_when_opted_in() {
let url = Url::parse("http://127.0.0.1:8080/").unwrap();
assert!(check_ssrf(&url, true).await.is_ok());
}
}
crates/codegen/kigi-tools/src/implementations/mod.rs
+3 -3
View File
@@ -15,9 +15,9 @@ pub mod use_tool;
pub mod web_search;
pub use kigi::bash::{BashError, BashToolInput};
pub use kigi::{
AskUserQuestionTool, BashTool, EnterPlanModeTool, ExitPlanModeTool, GrepTool, KillTaskTool,
ListDirTool, ReadFileTool, SearchReplaceTool, TaskOutputTool, TaskTool, TodoWriteTool,
WaitTasksTool, WebFetchTool, WebSearchTool,
AgentSwarmTool, AskUserQuestionTool, BashTool, EnterPlanModeTool, ExitPlanModeTool, GrepTool,
KillTaskTool, ListDirTool, ReadFileTool, SearchReplaceTool, TaskOutputTool, TaskTool,
TodoWriteTool, WaitTasksTool, WebFetchTool, WebSearchTool,
};
pub use memory::{MemoryGetImpl, MemorySearchImpl};
pub use opencode::{
crates/codegen/kigi-tools/src/normalization.rs
+1
View File
@@ -103,6 +103,7 @@ pub fn canonical_input(input: &ToolInput) -> Option<serde_json::Value> {
| ToolInput::WaitTasks(_)
| ToolInput::KillTask(_)
| ToolInput::Task(_)
| ToolInput::AgentSwarm(_)
| ToolInput::WebSearch(_)
| ToolInput::WebFetch(_)
| ToolInput::ApplyPatch(_)
crates/codegen/kigi-tools/src/registry/types.rs
+1
View File
@@ -650,6 +650,7 @@ impl ToolRegistryBuilder {
b.register::<kigi::GetTerminalCommandOutputTool>();
b.register::<kigi::WaitTasksTool>();
b.register::<kigi::TaskTool>();
b.register::<kigi::AgentSwarmTool>();
b.register::<kigi::WebSearchTool>();
b.register_with_params::<kigi::WebFetchTool, kigi::web_fetch::WebFetchParams>();
b.register::<kigi::LspTool>();
crates/codegen/kigi-tools/src/tool_taxonomy.rs
+2
View File
@@ -56,6 +56,7 @@ impl ToolKind {
ToolKind::MemorySearch => "Memory Search",
ToolKind::MemoryGet => "Memory Read",
ToolKind::Task => "Subagent",
ToolKind::AgentSwarm => "Subagent Swarm",
ToolKind::EnterPlan => "Enter Plan Mode",
ToolKind::ExitPlan => "Exit Plan Mode",
ToolKind::AskUser => "Ask User",
@@ -96,6 +97,7 @@ impl ToolKind {
| ToolKind::KillTaskAction
| ToolKind::Skill
| ToolKind::Task
| ToolKind::AgentSwarm
| ToolKind::DeployApp
| ToolKind::SearchTool
| ToolKind::UseTool
crates/codegen/kigi-tools/src/types/tool.rs
+5
View File
@@ -88,6 +88,11 @@ pub enum ToolKind {
MemorySearch,
MemoryGet,
Task,
/// Fan-out sibling of [`ToolKind::Task`]: its own kind because
/// `TemplateRenderer`'s `by_kind` map holds ONE tool name per kind, so
/// sharing `Task` would let the swarm win that slot and silently redirect
/// every `${{ tools.by_kind.task }}` reference in other tools' prompts.
AgentSwarm,
EnterPlan,
ExitPlan,
AskUser,
crates/codegen/kigi-tools/src/types/tool_io.rs
+2
View File
@@ -32,6 +32,7 @@ use crate::implementations::opencode::write::WriteInput;
use crate::implementations::search_tool::SearchToolInput;
use crate::implementations::skills::skill::SkillInput;
use crate::implementations::use_tool::UseToolInput;
use kigi_tool_types::AgentSwarmToolInput;
use kigi_tool_types::KillTaskToolInput;
use kigi_tool_types::TaskOutputToolInput;
use kigi_tool_types::TaskToolInput;
@@ -67,6 +68,7 @@ pub enum ToolInput {
WaitTasks(WaitTasksToolInput),
KillTask(KillTaskToolInput),
Task(TaskToolInput),
AgentSwarm(AgentSwarmToolInput),
WebSearch(WebSearchInput),
WebFetch(WebFetchInput),
Write(WriteInput),
crates/codegen/kigi-tui/docs/user-guide/05-configuration.md
+8
View File
@@ -225,8 +225,16 @@ timeout_secs = 1800 # seconds to wait when enabled (default:
[toolset.web_fetch]
proxy_endpoint = "https://proxy.example.com" # egress proxy URL
allowed_domains = ["docs.rs", "x.ai"] # override the built-in allowlist
allow_local = false # true = reach localhost / 127.0.0.0/8 / ::1
```
`allow_local` opens **loopback only**, and only when the URL names it
explicitly (`http://127.0.0.1:8080/`, `http://localhost:3000/`). A public
domain whose DNS record points at loopback stays blocked — that is DNS
rebinding, not local development. Private, link-local, CGNAT and cloud
metadata ranges are never reachable, with or without this flag. Precedence:
user config → `KIGI_WEB_FETCH_ALLOW_LOCAL` → off.
`[toolset.ask_user_question]` is honored across **requirements.toml**, **managed
config**, and **user `config.toml`**. Precedence: requirements → env
(`KIGI_ASK_USER_QUESTION_TIMEOUT_ENABLED` /
crates/codegen/kigi-tui/src/acp/tracker.rs
+8 -3
View File
@@ -2157,7 +2157,10 @@ fn extract_variant(tc: &acp::ToolCall) -> Option<&str> {
}
/// Twin without the optional-toolset spelling.
fn is_task_variant(variant: Option<&str>) -> bool {
matches!(variant, Some("Task"))
// The swarm belongs here for the same reason `Task` does: its members each
// raise their own SubagentBlock, and it must register the blocking wait or
// a call that holds the turn until every member finishes shows no spinner.
matches!(variant, Some("Task" | "AgentSwarm"))
}
/// Twin without the optional-toolset spelling.
fn is_write_variant(variant: Option<&str>) -> bool {
@@ -2190,8 +2193,10 @@ fn is_goal_tool(tc: &acp::ToolCall) -> bool {
/// SubagentSpawned notification) provides better visibility. Covers the
/// `task` / `Task` / `spawn_subagent` ids and Task-family variant tags.
fn is_task_tool(tc: &acp::ToolCall) -> bool {
matches!(tc.title.as_str(), "task" | "Task" | "spawn_subagent")
|| is_task_variant(extract_variant(tc))
matches!(
tc.title.as_str(),
"task" | "Task" | "spawn_subagent" | "agent_swarm"
) || is_task_variant(extract_variant(tc))
}
/// Check if a tool call is a scheduler tool (scheduler_create/delete/list).
///
crates/codegen/kigi-tui/src/app/dispatch/router.rs
+7 -1
View File
@@ -62,6 +62,7 @@ use super::session::lifecycle::{
clear_startup_actions, dispatch_agent_type_mismatch_answered, dispatch_exit_session,
dispatch_new_session, dispatch_new_session_inner, dispatch_new_session_with_id,
dispatch_new_worktree_session, dispatch_trust_folder, open_new_session_question,
skip_picker_and_create_session,
};
use super::session::load::{
dispatch_cycle_session_source_filter, dispatch_load_session, dispatch_pick_content_session,
@@ -749,7 +750,12 @@ pub(crate) fn dispatch(action: Action, app: &mut AppView) -> Vec<Effect> {
};
let Some(session_id) = agent.session.session_id.clone() else {
agent.session.deferred_model_switch = Some((model_id, effort));
return vec![];
// No session bound: with a create in flight this is a no-op
// and `SessionCreated` applies the stash; with none in flight
// (project question pending — only a plain prompt opens it) a
// switch would dangle forever, so start the session like the
// QueueCommand arm does for queued slash work.
return skip_picker_and_create_session(app, id);
};
agent.session.model_switch_pending = true;
vec![Effect::SwitchModel {
crates/codegen/kigi-tui/src/app/dispatch/settings/setters.rs
+14 -6
View File
@@ -1519,12 +1519,20 @@ pub(in crate::app::dispatch) fn set_default_model(
effort: None,
prev_model_id: prev_id.clone(),
});
} else if let Some(agent) = app.agents.get_mut(&aid) {
// No session id yet — stash for
// `EventLoop::on_session_created` to apply once the session
// id materialises. Mirrors the deferred-switch handling in
// `Action::SwitchModel`.
agent.session.deferred_model_switch = Some((new_id, None));
} else {
if let Some(agent) = app.agents.get_mut(&aid) {
// No session id yet — stash for
// `EventLoop::on_session_created` to apply once the session
// id materialises. Mirrors the deferred-switch handling in
// `Action::SwitchModel`.
agent.session.deferred_model_switch = Some((new_id, None));
}
// With no create in flight (project question pending), the stash
// would dangle forever — start the session it drains into. No-op
// when a create is already pending.
effects.extend(
crate::app::dispatch::session::lifecycle::skip_picker_and_create_session(app, aid),
);
}
effects
}
crates/codegen/kigi-tui/src/app/dispatch/tests/session/lifecycle.rs
+16 -3
View File
@@ -481,7 +481,7 @@ fn session_failed_clears_flag_no_fetches() {
assert!(!app.agents[&id].pending_extensions_fetch);
}
#[test]
fn switch_model_without_session_does_nothing() {
fn switch_model_without_session_starts_the_session() {
let mut app = test_app_with_agent();
let id = AgentId(0);
app.agents.get_mut(&id).unwrap().session.session_id = None;
@@ -493,7 +493,14 @@ fn switch_model_without_session_does_nothing() {
},
&mut app,
);
assert!(effects.is_empty());
// No create was in flight, so the switch starts the session its stash
// drains into (the stash alone dangled forever — the /model-in-Downloads
// silent no-op). `model_switch_pending` flips on SessionCreated.
assert!(
effects
.iter()
.any(|e| matches!(e, Effect::CreateSession { .. }))
);
assert!(!app.agents[&id].session.model_switch_pending);
}
#[test]
@@ -668,11 +675,17 @@ fn switch_model_deferred_when_no_session_id() {
},
&mut app,
);
assert!(effects.is_empty());
// Stashed for SessionCreated — and the session it drains into is started
// (no create was in flight; a bare stash never drained).
assert_eq!(
app.agents[&id].session.deferred_model_switch,
Some((model_id, None))
);
assert!(
effects
.iter()
.any(|e| matches!(e, Effect::CreateSession { .. }))
);
assert!(!app.agents[&id].session.model_switch_pending);
}
#[test]
crates/codegen/kigi-tui/src/app/dispatch/tests/task_result.rs
+71
View File
@@ -977,6 +977,77 @@ fn switch_model_pending_lifecycle() {
assert!(!app.agents[&id].session.model_switch_pending);
}
/// A model switch with no session AND no create in flight (the project-picker
/// question is pending; only a plain prompt can open it) must start the
/// deferred session itself, or the stashed switch dangles forever with zero
/// feedback — `/model X eff` in `~/Downloads` looked like "the model never
/// changes". Mirrors the `QueueCommand` arm: queued slash work bypasses the
/// picker and creates the session so the stash drains.
#[test]
fn switch_model_without_session_creates_the_deferred_session() {
let mut app = test_app_with_agent();
let id = AgentId(0);
app.agents.get_mut(&id).unwrap().session.session_id = None;
// Harness cwd is `/tmp` (a non-project dir); arm the picker gate the way
// startup leaves it (the harness pre-marks it shown for other tests).
app.project_picker_shown = false;
assert!(app.needs_project_picker());
let model_id = acp::ModelId::new(std::sync::Arc::from("kigi-4.5"));
let effects = dispatch(
Action::SwitchModel {
model_id: model_id.clone(),
effort: None,
},
&mut app,
);
assert_eq!(
app.agents[&id].session.deferred_model_switch,
Some((model_id, None)),
"switch must stay stashed for SessionCreated to apply"
);
assert!(
effects
.iter()
.any(|e| matches!(e, Effect::CreateSession { .. })),
"sessionless switch must start the session the stash drains into"
);
}
/// Same stash path while a create IS in flight (`mcp_init_progress` set):
/// no duplicate `CreateSession` — the pending create applies the stash.
#[test]
fn switch_model_with_create_in_flight_does_not_duplicate_create() {
let mut app = test_app_with_agent();
let id = AgentId(0);
{
let agent = app.agents.get_mut(&id).unwrap();
agent.session.session_id = None;
agent.mcp_init_progress = Some(crate::app::agent_view::McpInitProgress {
total: 0,
connected: 0,
started_at: std::time::Instant::now(),
});
}
let effects = dispatch(
Action::SwitchModel {
model_id: acp::ModelId::new(std::sync::Arc::from("kigi-4.5")),
effort: None,
},
&mut app,
);
assert!(app.agents[&id].session.deferred_model_switch.is_some());
assert!(
!effects
.iter()
.any(|e| matches!(e, Effect::CreateSession { .. })),
"an in-flight create must not be duplicated"
);
}
#[test]
fn no_deferred_switch_means_no_extra_effect() {
// When there is no deferred model switch, SessionCreated should
crates/codegen/kigi-tui/src/app/event_loop.rs
+11 -27
View File
@@ -2333,11 +2333,7 @@ async fn drain_and_process(
}
// On terminals without bracketed paste, try to capture more events
// that may still be in transit from the input reader thread. A batch
// 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.
// that may still be in transit from the input reader thread.
if should_extend_for_paste(&raw_events)
&& (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);
/// Timeout for subsequent rounds once paste has been detected. Must exceed
/// one Windows scheduler quantum (~15.6 ms): ConPTY delivers a paste as a
/// per-character key burst, and a mid-burst gap of one quantum is routine.
/// 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.
/// one Windows scheduler quantum (~15.6 ms), or a routine mid-burst gap
/// splits the paste and repaste-to-expand duplicates instead of expanding.
const PASTE_CONTINUE_TIMEOUT: Duration = Duration::from_millis(25);
/// Safety cap on events accumulated in one extension pass. On the
/// no-bracketed-paste path a paste is one key event per character, so the
/// cap must exceed any real paste (~200 KB of text); hitting it splits the
/// paste with the same duplicate-on-repaste effect as a timeout split. The
/// idle timeout above, not this cap, is the normal terminator.
/// Safety cap on events accumulated in one extension pass. Pastes arrive
/// one key event per character here, so the cap must exceed any real paste;
/// hitting it splits the paste like a timeout miss would.
const PASTE_EXTEND_MAX_EVENTS: usize = 200_000;
/// 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
/// paste burst in progress. Human typing and key auto-repeat deliver one
/// or two events per batch; only a paste chunk lands more at once.
/// paste burst; typing and auto-repeat deliver only a couple per batch.
const PASTE_BURST_BATCH_EVENTS: usize = 8;
/// True when the batch alone proves a paste is in progress, without
/// waiting on [`PASTE_DETECT_TIMEOUT`].
/// True when the batch alone proves a paste is in progress, so collection
/// can start without the [`PASTE_DETECT_TIMEOUT`] round-trip.
fn is_paste_burst(events: &[Event]) -> bool {
events.iter().filter(|e| is_pasteable_key_event(e)).count() >= PASTE_BURST_BATCH_EVENTS
}
@@ -3188,11 +3178,9 @@ mod tests {
#[test]
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'))];
assert!(!is_paste_burst(&typed));
// Release events carry no content and must not count toward it.
let releases: Vec<Event> = "releases only!"
.chars()
.map(|c| release(KeyCode::Char(c)))
@@ -3200,10 +3188,8 @@ mod tests {
assert!(!is_paste_burst(&releases));
}
/// A key-burst paste whose tail trickles in with sub-window gaps must
/// reassemble into ONE synthetic paste. A split here is what made
/// repaste-to-expand insert a duplicate chip on Windows: the fragment
/// could never byte-equal the chip content.
/// A burst tail trickling in with one-quantum gaps must reassemble into
/// ONE paste — a split is what duplicated chips on Windows.
#[tokio::test(start_paused = true)]
async fn burst_tail_with_scheduler_gaps_reassembles_into_one_paste() {
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();
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 {
for c in tail.chars() {
tokio::time::sleep(Duration::from_millis(15)).await;
crates/codegen/kigi-tui/src/slash/commands/effort_levels.rs
-1
View File
@@ -25,7 +25,6 @@ pub(crate) fn effort_description(level: ReasoningEffort) -> &'static str {
ReasoningEffort::High => "Heavy reasoning",
ReasoningEffort::Xhigh => "Extra-heavy reasoning",
ReasoningEffort::Max => "Maximum reasoning",
ReasoningEffort::Ultra => "Ultra reasoning",
}
}
crates/codegen/kigi-workspace/src/capability.rs
+4 -2
View File
@@ -91,6 +91,7 @@ pub(crate) const ALL_TOOL_KINDS: &[ToolKind] = &[
ToolKind::MemorySearch,
ToolKind::MemoryGet,
ToolKind::Task,
ToolKind::AgentSwarm,
ToolKind::EnterPlan,
ToolKind::ExitPlan,
ToolKind::AskUser,
@@ -144,8 +145,9 @@ pub(crate) fn kind_allowed(mode: CapabilityMode, kind: ToolKind) -> bool {
// Bash / shell.
Execute => matches!(mode, M::Execute),
// Process control (background tasks, monitors).
BackgroundTaskAction | WaitTasksAction | KillTaskAction | Task | Monitor => {
// Process control (background tasks, monitors). A swarm is a fan-out of
// subagent spawns, so it gates exactly as `Task` does.
BackgroundTaskAction | WaitTasksAction | KillTaskAction | Task | AgentSwarm | Monitor => {
matches!(mode, M::Execute)
}
crates/codegen/kigi-workspace/src/session/tool_config.rs
+3
View File
@@ -325,6 +325,9 @@ fn build_web_fetch_config() -> kigi_tools::implementations::kigi::web_fetch::Web
if let Ok(proxy) = std::env::var("KIGI_WEB_FETCH_PROXY") {
params.proxy_endpoint = Some(proxy);
}
if kigi_config::env_bool("KIGI_WEB_FETCH_ALLOW_LOCAL") == Some(true) {
params.allow_local = Some(true);
}
WebFetchConfig::Enabled { params }
}
#[cfg(any(test, feature = "test-support"))]
crates/common/kigi-tool-types/src/agent_swarm.rs
+116
View File
@@ -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()
}
}
crates/common/kigi-tool-types/src/lib.rs
+5
View File
@@ -1,10 +1,15 @@
//! Canonical, extensible tool types.
mod agent_swarm;
mod ext;
mod schema_utils;
pub mod serde_lenient;
mod task;
mod types;
pub use agent_swarm::{
AgentSwarmToolInput, MAX_AGENT_SWARM_MEMBERS, PROMPT_TEMPLATE_PLACEHOLDER, SwarmMemberOutcome,
SwarmMemberResult,
};
pub use ext::Extensions;
pub use schema_utils::parse_arguments_from_schema_lossy;
pub use serde_lenient::{