Compare commits

..
Author SHA1 Message Date
Gnanam b6897d98d1 chore: clear staticcheck nits surfaced by the v16 audit
Fixes the six minor staticcheck findings (no behavior change):
- S1011 ×2: loop+append → append(dst, src...) in agenteval/score.go and
  cli/agent_eval.go.
- S1016: convert rawResult to searchResult instead of a field-by-field struct
  literal in tools/web_search.go (identical fields).
- ST1005: lowercase the "unsafe redirect URL" error string in tools/web_fetch.go
  (+ update its test assertion).
- SA4006: drop the dead final assignment to next in a provider_wizard test
  (keep the type assertion).
- S1009: remove the redundant nil check before len() in a provideronboarding
  test.

staticcheck drops 32 → 26 (remaining: 15 U1000 unused + 11 SA1019 deprecated
mouse API — separate audit items). gofmt/vet/build/test (full) pass;
deadcode unchanged (100); govulncheck 0.
2026-06-15 16:20:50 +05:30
KRATOS 9ddef0e957 feat(swarm): multi-agent swarm over the specialist sub-agent (#207)
* feat(swarm): multi-agent swarm over the specialist sub-agent

Add internal/swarm so an orchestrator can spawn and coordinate MULTIPLE
specialist members concurrently under one team, communicate via per-agent
mailboxes, hand work off, and adopt orphaned tasks. Additive and opt-in: the
single Task tool is unchanged; swarm tools only act when invoked.

- Registry/Definition: agent roster (built-in teammate/subagent + user-defined).
- Mailbox: per-team/per-agent JSON inbox, owner-only 0600, file-locked with
  stale-break, size-capped, path-confined, and malformed-fails-closed.
- Coordinator: in-memory task registry with terminal-state guards and stable
  per-agent colors.
- Team: bounded concurrent members per team with a FIFO overflow queue that
  drains one-per-slot.
- Lifecycle: spawn / handoff / orphan-adoption with bounded relaunch on
  temporary failures; members run under the Swarm's base context so they
  outlive the spawn call.
- Tools: swarm_spawn, swarm_send, swarm_inbox, swarm_status, swarm_handoff,
  swarm_collect.
- Wiring: each member is launched through specialist.Executor and inherits the
  orchestrator's model + permission mode (never widened) and runs under the same
  sandbox/policy. Mailbox state lives under the workspace so writes fall within
  the sandbox write rules. Swarm lifetime is paired with the specialist runtime.

Scheduling (cron/wakeup) and daemon-pool-backed members are deferred follow-ups.

* swarm: address CodeRabbit review (7 findings)

- coordinator: SetStatus rejects unknown TaskStatus values (fail closed) so a
  task can never enter a state Summarize/lifecycle don't reason about.
- permission inheritance ENFORCED end-to-end (was resolved but discarded):
  thread PermissionMode through specialist.TaskRunOptions -> BuildArgs/
  BuildResumeArgs. A headless specialist child now runs "--auto high" (unsafe)
  ONLY when the parent is unsafe or no mode is given (preserves the Task tool's
  behavior); any other parent mode yields non-unsafe "--auto low", so a member
  never gains more authority than its orchestrator. Also fixed the swarm
  permission vocabulary to the real agent runtime modes (ask/spec-draft/auto/
  unsafe) instead of TUI display names, and the launcher substitutes "auto" for
  an empty mode so a member can't fall into the empty-means-high path.
- mailbox: symlink-aware confinement (EvalSymlinks) rejects an inbox dir that
  resolves outside BaseDir; dirs tightened to 0700 even if pre-existing with
  broad perms; lock release is ownership-token-aware so a stale-broken writer
  can't delete a newer writer's lock (split-brain).
- member: recover panics inside a member goroutine -> fails that member only,
  never crashes the orchestrator.
- tools: collapse() truncates on rune boundaries (no invalid UTF-8).

Tests added for every fix. Existing specialist behavior is unchanged (empty
PermissionMode keeps --auto high). Gates: gofmt/vet/build (host+linux)/test
-race/staticcheck(no new)/govulncheck(0)/deadcode(no new) all pass.

* swarm: address review round 2 (Windows lock + non-blocking notes)

Blocking (Vasanth): mailbox acquireLock now treats Windows ERROR_ACCESS_DENIED
/ ERROR_SHARING_VIOLATION as lock-contended (retry/wait) instead of a fatal
error, via a platform-aware isLockContended (lock_windows.go / lock_other.go).
A contended lock under Windows file-sharing semantics no longer fails a
send/read spuriously. TestMailboxConcurrentSends raised to 200 goroutines with a
10s lock timeout and now asserts zero send failures (regression guard).

Non-blocking notes addressed:
- Removed third-party reference-project names from doc comments (clean-room +
  provider-neutral); the Go code is original.
- Documented the inbox-file size bound's ×2 (pretty-print overhead) constant.
- Documented Send's O(N) rewrite cost + when to switch to append-only.
- permissionRank: spec-draft now ranks below ask (it never widens authority).
- Documented watch's per-member-id retry binding (Member.ID == MemberSpec.ID).
- handoff Safety.Reason now also mentions the inbox write.

Gates: gofmt · vet · build (host+linux+windows) · test -race · staticcheck
(no new) · govulncheck (0) · deadcode (no new) all pass.
2026-06-15 15:10:21 +05:30
KRATOS ba8fbb8b71 feat(sandbox): complete sandbox breadth — WSL fallback + opt-in MITM TLS inspection (#206)
* feat(sandbox): complete sandbox breadth — WSL fallback + opt-in MITM TLS inspection

Adds the two remaining sandbox features on top of the strict fail-closed core
(the other four — SOCKS5 egress, re-entrancy guard, fine-grained path lists,
sandbox-aware ripgrep — already shipped and are untouched). Both new features are
opt-in and the DEFAULT policy is unchanged (enforce + network deny, InspectTLS
off; AllowPolicyOnlyRunner keeps its existing default).

1) WSL detection + safe fail-closed fallback
   - parseWSL/detectWSL (wsl.go + wsl_linux.go reading /proc/version; wsl_other.go
     stub) classify WSL/WSL2.
   - SelectBackend: on Linux without bubblewrap, if running under WSL it selects a
     new BackendWSL instead of silently degrading to a plain policy-only runner.
   - BuildCommandPlan(BackendWSL): FAILS CLOSED unless Policy.AllowPolicyOnlyRunner
     is set (errWSLPolicyOnlyDisabled explains how to opt in). When allowed it runs
     the policy-only fallback that still routes egress through the filtering proxy
     (Backend.ProxyEgress => EnforcesScopedEgress, best-effort HTTP+SOCKS), sets
     ZERO_SANDBOXED=1 / ZERO_SANDBOX_BACKEND=wsl, and records an auditable
     least-privilege downgrade note (CommandPlan.Notes). Never runs unsandboxed
     silently.

2) Optional gated MITM TLS inspection (Policy.InspectTLS, default OFF)
   - Off => the egress proxy is a pure CONNECT passthrough, byte-for-byte
     unchanged (no CA generated).
   - On => the proxy terminates TLS toward the sandboxed client with an EPHEMERAL,
     locally generated ECDSA CA (mitm.go), minting per-host leaves on the fly, so
     it can re-check the DECRYPTED request Host through the SAME
     authorize()/domainAllowed() gate (deny-wins, empty allowlist => deny) — MITM
     widens visibility, never authority. The UPSTREAM dial is ALWAYS validated
     against the system roots (never InsecureSkipVerify; tests inject a known root
     pool to exercise validation). The CA's PUBLIC cert is written owner-only and
     surfaced via ZERO_SANDBOX_CA_CERT; MITM only runs on backends that already
     permit reading it (sandbox-exec allows file-read; WSL shares the host fs), so
     no new bind path is introduced. Trust implications documented loudly on the
     field. (No "raw passthrough only" flag exists in zero, so the mutual-exclusion
     case is N/A.)

Tests: parseWSL table; WSL fail-closed-without-opt-in and policy-only-with-proxy
(markers + ProxyEnvWithSocks) + deny-mode markers-no-proxy; MITM passthrough-by-
default (no CA), CA mints a chaining leaf, full TLS-terminate-and-forward of an
allowlisted host (upstream validated via an explicit root pool), and decrypted-
Host deny. Gates: gofmt, go vet (host + GOOS=linux), go build (host/linux/
windows), go test (incl -race), staticcheck (no new, host + linux), govulncheck
(0), deadcode -test=false under GOOS=linux (no new unreachable funcs).

* sandbox: address review on #206 (WSL + MITM)

- mitm: bound the upstream transport's TCP connect (DialContext 30s) so a
  blackholed upstream can't park a goroutine; disable upstream HTTP/2 and
  normalize the response to HTTP/1.1 before writing to the HTTP/1.1 client.
- mitm: mint leaves for the AUTHORIZED CONNECT host (not arbitrary client SNI)
  and cap the leaf cache (512) so SNI churn can't drive unbounded keygen/memory.
- mitm: re-check the decrypted Host with the FULL authorize gate
  (authorizeTarget) so a DomainPrompt-allowed host is honored consistently with
  the CONNECT check, not just domainAllowed.
- runner (WSL): preserve the caller's env in the direct-run WSL plan (it is not
  OS-wrapped), appending the sandbox markers last; the downgrade note says
  "with proxy egress" only when a proxy was actually started.
- tests: WSL caller-env preservation; deny-mode asserts no HTTP/HTTPS/ALL_PROXY
  leak + the downgrade note; MITM forward asserts HTTP 200; mitmConnect sets an
  I/O deadline.

Gates: gofmt, vet (host+linux), build (host/linux/windows), test (incl -race),
staticcheck (no new), govulncheck (0), deadcode -test=false (GOOS=linux) no new.
2026-06-15 11:38:14 +05:30
Vasanth TandClaude 265aa52ba7 Polish /permissions command output into a bounded card (#205)
* Convert /permissions output to bounded command card

Replace status: ok style with commandCard format used by /tools and /context.

Summary shows mode and grant count; grants render as bullet rows.

Co-Authored-By: Claude <[email protected]>

* Address CodeRabbit nitpick: add error path test for /permissions

Adds TestPermissionsCommandCardHandlesGrantListError and a grantLister

interface so tests can inject a stub store without a real file path.

Co-Authored-By: Claude <[email protected]>

* Fix gofmt formatting in permissions card

Tabs alignment only; no functional change.

Co-Authored-By: Claude <[email protected]>

---------

Co-authored-by: Claude <[email protected]>
2026-06-15 11:37:51 +05:30
KRATOS 135e3c5657 feat(daemon): zero daemon — worker pool + session supervisor over a local control socket (#203)
* feat(daemon): add `zero daemon` — worker pool + session supervisor over a local control socket

Adds internal/daemon + `zero daemon start|stop|status|run|attach`: a long-running
local service that supervises a pool of headless `zero exec` workers and routes
multiple agent sessions to them over an owner-only Unix-domain control socket.
Additive — interactive and one-shot exec are unchanged.

Reuses zero's building blocks:
- internal/background: ConfigureChildProcessGroup (process-group setup) +
  TerminateProcess (cross-platform terminate) for worker kill/drain. The pool
  does its own os/exec and uses background only for those helpers.
- internal/streamjson: the worker speaks stream-json on stdout; the daemon
  forwards those lines to clients.
- internal/cli exec: a worker is `zero exec --output-format stream-json`.

Pieces (reimplemented in Go from a reference daemon):
- protocol.go: framed control codec — 4-byte BE length + 1 kind byte, 1 MiB cap,
  JSON control messages, version negotiation.
- pool.go + backoff.go: bounded worker pool — at-least-once dispatch with
  exponential backoff+jitter retry capped by MaxAttempts, EXIT_TEMPFAIL(75)/
  EXIT_PERMANENT(76) handling, graceful drain + kill-timeout. zero's exec is
  one-shot, so workers are on-demand (one session per worker = the lease), not
  warm long-lived slots (a persistent worker mode is a follow-up).
- session.go: SessionManager — sessionID->worker routing, lease (one active
  request per worker), queue when the pool is full, per-session output buffer for
  attach + metrics.
- server.go + lock.go + socket.go + paths.go: single-instance PID-file lock with
  stale-lock (dead-PID) recovery, status file, owner-only socket, accept loop;
  cleans up lock/socket/status on exit.
- launcher.go: production worker launcher. client.go: control-socket client.

Security model:
- Control socket is LOCAL-ONLY (AF_UNIX loopback file, never a TCP port) and
  owner-only: socket 0600 + parent dir 0700 on POSIX (Windows relies on the
  per-user profile ACL; a dedicated pipe ACL is a follow-up).
- Workers stay sandboxed: the launcher SCRUBS the sandbox re-entrancy markers
  (ZERO_SANDBOXED / ZERO_SANDBOX_BACKEND) from each worker's env so the worker
  re-establishes its OWN sandbox instead of inheriting a pass-through plan, while
  propagating the rest of the policy config/env. The daemon never bypasses the
  sandbox.
- Control frames over the 1 MiB cap are rejected before allocation.

Robustness: a crashed worker never takes down the daemon (backoff+jitter retry,
MaxAttempts cap, EXIT_PERMANENT stops retry); drain gives a grace window then
force-kills stragglers; a stale lock from an unclean exit is reclaimed.

Out of scope (follow-up): remote session-linking; warm pre-spawned persistent
workers; Windows named-pipe ACL.

Tests: protocol round-trip + oversize/unknown-kind rejection; backoff bounds;
pool spawn/retry-backoff/permanent/cap/queue/drain; session lease/route/queue/
attach; lock single-instance + stale recovery; launcher env-scrub (incl. an
end-to-end check that the spawned worker process receives scrubbed env); server
e2e start->run->stream-json->status->attach->stop with a fake worker; CLI
usage/validation/not-running paths. Live-verified start/status/stop (0600 socket,
0700 dir, file cleanup). Platform files behind explicit //go:build tags.

* daemon: address review — help listing, extra-arg rejection, lock write error, bounded sessions

- cli/app.go: list `daemon` in `zero --help` so the command is discoverable.
- cli/daemon: `stop`/`status` reject extra arguments; `attach` rejects more than
  one session id (fail fast on typos instead of silently ignoring tokens).
- daemon/lock: check the PID fmt.Fprintf error — a partial/failed write would
  leave a malformed lock file that another process reads as stale and wrongly
  reclaims, breaking the single-instance guarantee; on write failure remove the
  file and return the error.
- daemon/session: bound the registry with MaxSessions (default 256) — evict the
  oldest FINISHED sessions once past the cap so a long-running daemon doesn't
  accumulate sessions without limit; running/queued sessions are never evicted, so
  attach-after-finish for recent sessions is preserved.

Tests: CLI extra-arg rejection; finished-over-cap eviction + running-never-evicted.
Gates: gofmt, vet, build, test (incl -race), staticcheck (no new), govulncheck (0),
deadcode (no new), GOOS=linux+windows cross-compile.
2026-06-14 22:16:37 +05:30
KRATOS 355b76cc78 sandbox: exempt first-party network tools from the Evaluate network-deny too (#202)
Follow-up to the EnforceToolNetwork change: web_search/web_fetch were still
hard-denied under NetworkDeny because the engine-level gate in Evaluate
(netMode==NetworkDeny && risk has "network") fired BEFORE the tool's
RunWithSandbox/NetworkHostAllowed ever ran. The earlier change only relaxed
NetworkHostAllowed, so the user-facing path (Evaluate -> registry) still blocked
the tool with "network access is blocked by sandbox policy".

Fix: Evaluate now skips the network-deny for a request that declares
SideEffectNetwork (the first-party in-process tools, web_search/web_fetch) when
EnforceToolNetwork is off — mirroring the NetworkHostAllowed exemption so the two
gates agree. A SHELL command merely classified as network (SideEffectShell) is
NOT exempt, so shell egress stays blocked under deny. EnforceToolNetwork restores
the strict deny for the tools.

Tests: TestEvaluateExemptsNetworkToolsFromDenyByDefault covers the full path
(default-exempt, granted->allow, EnforceToolNetwork->deny, shell stays denied);
the two existing gate tests now opt into EnforceToolNetwork to keep testing the
strict path. Gates: gofmt, vet, build, test (incl. -race), staticcheck (no new),
govulncheck (0), deadcode (no new unreachable funcs).
2026-06-14 20:36:51 +05:30
Anandan 6a9f44a716 Polish TUI transcript and reasoning display (#200)
* Polish TUI transcript and reasoning display

* Address provider and transcript review fixes

* Gate inline think tag parsing

* Cover disabled think tag parsing config
2026-06-14 18:55:18 +05:30
Vasanth T e1f1b111a1 Add command dashboard cards (#201)
* Add context command dashboard card

* Render tools command as dashboard card

* Render command cards as styled TUI cards

* Fix command card action rendering

* fix: harden command card edge cases
2026-06-14 18:34:46 +05:30
Vasanth T 284a948f00 feat: add doctor fix flow (#198)
* feat: add doctor fix flow

* fix: animate doctor connectivity status

* fix: make doctor report problem first

* fix: address doctor review feedback
2026-06-14 17:54:00 +05:30
KRATOS 671af594fb feat(sandbox): harden internal/sandbox — re-entrancy guard, SOCKS5 egress, sandbox-aware search, fine-grained path lists (#199)
* feat(sandbox): harden internal/sandbox with four opt-in capabilities

All four are off by default and fail closed; existing behavior is unchanged
when the new policy fields are empty / ZERO_SANDBOXED is unset.

1. Re-entrancy guard — ZERO_SANDBOXED=1 marks every wrapped command; a nested
   BuildCommandPlan returns a pass-through plan (no double bwrap/sandbox-exec,
   no second egress proxy). New: EnvSandboxed, IsAlreadySandboxed().

2. SOCKS5 egress — the scoped-egress proxy also serves SOCKS5 CONNECT on a
   second loopback port through the SAME authorize()/domainAllowed() gate.
   ProxyEnvWithSocks points ALL_PROXY at socks5://<addr>; HTTP(S)_PROXY stay on
   the HTTP listener; the sandbox-exec profile allows both proxy ports.
   Bubblewrap scoped still collapses to deny.

3. Sandbox-aware search — ReadExclusionGlobs(policy, scope) returns ripgrep
   --glob '!<pat>' for DenyRead subtrees; grep/glob skip read-denied subtrees
   at walk time via Engine.ReadPathExcluded/ReadDirExcluded when sandboxed.

4. Fine-grained path lists — Policy.AllowRead/DenyRead/AllowWrite/DenyWrite.
   Read = deny-then-allow (more-specific AllowRead re-includes); Write =
   DenyWrite > workspace/scope > AllowWrite > deny. AllowWrite is reflected in
   the OS write binds; DenyWrite is emitted as sandbox-exec deny rules. Each
   entry is home-expanded, made absolute, and symlink-resolved.

Tests: internal/sandbox/{reentrancy,egress_socks,pathlists}_test.go and
internal/tools/read_exclusions_test.go.
Gates: gofmt, go vet, go build, go test (incl. -race), staticcheck (no new
findings), govulncheck (0), deadcode -test=false (no new unreachable funcs).

* sandbox: address review — SOCKS5 method negotiation + cache read-exclusion roots

- egress_socks: respect SOCKS5 method negotiation (RFC 1928). The proxy now
  inspects the client's offered methods and replies 0xFF ("no acceptable
  methods") when no-auth was not offered, instead of forcing no-auth.
- pathlists/engine: resolve DenyRead/AllowRead roots ONCE per search walk via a
  new ReadExclusions matcher (Engine.ReadExclusions), replacing the per-path
  ReadPathExcluded/ReadDirExcluded that re-ran Abs/EvalSymlinks for every visited
  path during a grep/glob walk.

Re-entrancy guard left as-is by design: ZERO_SANDBOXED is an internal marker set
only on commands zero wraps; cross-process re-entrancy detection legitimately
relies on the inherited env, and a manually-exported flag is a self-opt-out, not
an external bypass.

Tests: SOCKS no-acceptable-methods rejection; nested-AllowRead descent via the
matcher. Gates: gofmt, vet, build, test (incl. -race), staticcheck (no new
findings), govulncheck (0), deadcode (no new unreachable funcs).

* sandbox: fix windows smoke test + symlink-resolve DirExcluded prefix

- pathlists_test: TestSandboxExecProfileEmitsDenyWriteRule now builds the
  expected deny clause via sandboxProfileString(secret), so it matches the
  profile's own escaping (Windows doubles backslashes). Fixes the
  windows-latest smoke failure; macOS/Linux unaffected (slash paths, no-op
  escaping).
- pathlists: ReadExclusions.DirExcluded now EvalSymlinks-resolves the candidate
  path before the nested-AllowRead prefix check, matching the resolution applied
  to allowRoots — otherwise a denied dir reached THROUGH a symlink fails to match
  a nested AllowRead root and is wrongly skipped, dropping a re-included subtree.
  Added TestDirExcludedResolvesSymlinkPrefix (skips where symlinks are
  unavailable).

Gates: gofmt, vet, build, test (incl. -race), staticcheck (no new findings),
govulncheck (0), deadcode (no new unreachable funcs).

* sandbox: strengthen re-entrancy guard with a corroborating backend marker

Addresses review feedback that the re-entrancy guard trusted a single ambient
env var. IsAlreadySandboxed now requires BOTH correlated markers that
sandboxEnvironment sets together on every wrapped command — ZERO_SANDBOXED=1 AND
a non-empty ZERO_SANDBOX_BACKEND (new EnvSandboxBackend const) — so a lone
stray/hand-exported ZERO_SANDBOXED=1 no longer forces an unsandboxed
pass-through. Pass-through (direct) plans set neither marker, so genuine nesting
still detected and double-wrapping still avoided.

Tests updated: TestIsAlreadySandboxed now asserts each marker alone is
insufficient and both together are required; the re-entrancy and wrapping plan
tests set/clear both markers explicitly.

Gates: gofmt, vet, build, test (incl. -race), staticcheck (no new findings),
govulncheck (0), deadcode (no new unreachable funcs).

* sandbox: don't subject first-party web tools to the network policy by default

The sandbox was reported as too strict: web_search (and web_fetch) were blocked
with "network access is disabled by the sandbox policy" because the default
policy is network=deny and the in-process web tools were gated by the same
NetworkHostAllowed check as the sandboxed shell's egress.

That network policy exists to confine the sandboxed SHELL's egress, which these
in-process tools don't use. New Policy.EnforceToolNetwork (off by default) makes
NetworkHostAllowed exempt web_search/web_fetch unless an operator opts in. The
sandboxed-shell egress decision (Evaluate / effectiveNetworkMode) is unchanged,
and the tools keep their own SSRF/private-IP, port, redirect, and redaction
safeguards. Set EnforceToolNetwork to hold the tools to the allow/scoped/deny
policy as before.

Tests: engine default-exempt + enforced-table cases; web_search default-allows
and enforced-when-flag; web_fetch enforced cases set the flag. Gates: gofmt, vet,
build, test (incl. -race), staticcheck (no new findings), govulncheck (0),
deadcode (no new unreachable funcs).

* sandbox: re-review fixes — disabled policy inert; path lists independent of EnforceWorkspace

- engine: ReadExclusions() / ReadExclusionGlobs() now return inert results under
  ModeDisabled, so a disabled policy never filters grep/glob (parity with
  Evaluate, which already allows everything when disabled).
- engine/pathlists: the fine-grained path lists (DenyRead/DenyWrite, with
  AllowRead/AllowWrite) now apply in Evaluate whenever the sandbox is enforcing,
  independent of EnforceWorkspace — matching the grep/glob path that already
  honors DenyRead directly. The workspace boundary (scope.validate) stays gated
  by EnforceWorkspace via a param threaded through validatePathWithPolicy /
  validateWritePath; DenyWrite wins regardless.
- reentrancy_test: TestSandboxEnvironmentMarksSandboxed now asserts BOTH
  correlated markers (ZERO_SANDBOXED=1 and ZERO_SANDBOX_BACKEND) so dropping
  either would fail the test.

Tests added: path lists enforced with EnforceWorkspace off; ReadExclusions inert
under ModeDisabled. Gates: gofmt, vet, build, test (incl. -race), staticcheck
(no new findings), govulncheck (0), deadcode (no new unreachable funcs).

* sandbox: enforce path lists even without a workspace root

Re-review follow-up: the Evaluate path-validation loop was still guarded by
request.WorkspaceRoot != "", so an engine built without a workspace root skipped
DenyRead/DenyWrite/AllowRead/AllowWrite even for absolute paths. The loop now
runs unconditionally; the workspace boundary (scope.validate) stays gated on
having a root (enforceWorkspace = EnforceWorkspace && WorkspaceRoot != ""), while
the path lists match absolute paths regardless. An unanchorable relative path
(no workspace root) fails closed when anything is configured to enforce, and is a
no-op otherwise (unchanged from prior behavior).

Test: TestEvaluatePathListsWithoutWorkspaceRoot. Gates: gofmt, vet, build, test
(incl. -race), staticcheck (no new findings), govulncheck (0), deadcode (no new
unreachable funcs).

* sandbox: base no-workspace fail-closed on normalized (enforceable) path lists

policyHasPathLists now resolves the lists (resolvePolicyPaths/
resolveWriteRootPaths) instead of counting raw config, matching how the rest of
the file normalizes them. A typo or non-existent entry (dropped during
resolution) therefore no longer makes validatePathWithPolicy fail-close every
relative request when there is no workspace root.

Gates: gofmt, vet, build, test (incl. -race), staticcheck (no new findings),
govulncheck (0), deadcode (no new unreachable funcs).
2026-06-14 17:37:40 +05:30
Vasanth T 7153ec7a7d feat: add TUI diagnostics center (#197) 2026-06-14 13:18:43 +05:30
KRATOS 73eb8d3376 Solidify sandbox: harden command analysis + wire/clean dormant policy surfaces (#196)
* Harden sandbox command risk by wiring the AST analyzer

The AST-based command analyzer (sandbox/analyzer.go) was built and tested but never
reached at runtime — the live risk classifier used only the regex detectors. Wire
AnalyzeCommand into classifyWithScope as a second opinion. It walks the parsed shell tree,
so it catches destructive/network programs the regexes miss (shred, fdisk, parted; telnet,
ftp, sftp; and commands hidden behind sudo/env wrappers or a sh -c launcher) and flags an
unparseable (obfuscated) script as elevated risk.

Purely additive — it only ADDs risk categories, so a benign, parseable command classifies
exactly as before, and a program name inside a quoted argument is not flagged (AST
precision). Revives the dormant analyzer (sandbox unreachable funcs 13 -> 6). New tests cover
the regex-missed destructive/network programs, nested launchers, the unparseable signal, the
quoted-name no-false-positive case, and benign-stays-clean. Full suite + -race green.

* Wire sandbox auto-allow-bash env opt-in; drop orphaned batch-autonomy

Two dormant sandbox surfaces resolved:

- WIRE: AutoAllowBashWhenSandboxed was enforced live (engine.go) but never set — no
  config field, and the ZERO_SANDBOX_AUTO_ALLOW_BASH env overlay (ApplyAutoAllowBashEnv)
  was unwired, so the documented opt-in did nothing. Wire it through the central policy
  builder (applyConfiguredSandboxPolicy, used by all engine constructions). Fail-safe: only
  turns ON when the env is truthy, only honored when the shell sandbox is actually active,
  default stays prompt, and an explicit config opt-in is preserved. Revives
  AutoAllowBashEnvEnabled + ApplyAutoAllowBashEnv; adds a policy-builder test.

- DELETE: RequiredAutonomyForBatch/riskToAutonomy (batch.go + test) had no consumer — the
  agent loop enforces autonomy per tool call (loop.go), so a batch-ceiling helper was
  speculative and redundant. Wiring it would have meant fabricating a batch model.

Sandbox now has no dormant funcs; the only remaining unreachable entries are the
cross-platform seccomp pair (unixSocketBlockFilter is live on Linux; ApplyUnixSocketBlock
is the intentional non-Linux no-op stub).

* Assert risk level (and add wrapper cases) in AST analyzer tests

Per review (Vasanth + CodeRabbit): the AST hardening tests checked only the risk category,
not the escalation level, so a regression that downgraded destructive/network shell commands
while preserving the category would slip through. Assert risk.Level too — Critical for
destructive/network, High for unparseable — matching the pattern used throughout the file.
Add sudo/env wrapper cases (sudo/env shred, sudo telnet) to cover effectiveProgram's wrapper
unwrapping.
2026-06-14 13:13:16 +05:30
Vasanth T 0444eb0e6f feat: add backend doctor diagnostics (#194) 2026-06-14 11:57:10 +05:30
122 changed files with 12990 additions and 352 deletions

No files matched your search

+1 -3
View File
@@ -125,9 +125,7 @@ func Score(suite Suite, input ScoreInput) Report {
if len(task.RequiredTraceEvents) > 0 {
report.Results = append(report.Results, scoreTraceEvents(task.RequiredTraceEvents, input))
}
for _, result := range unknownCommandResults(input.CommandResults, seenCommands, input) {
report.Results = append(report.Results, result)
}
report.Results = append(report.Results, unknownCommandResults(input.CommandResults, seenCommands, input)...)
report.finishSummary()
return report
}
+1 -3
View File
@@ -515,9 +515,7 @@ func agentEvalReportFromBenchmark(suite agenteval.Suite, report agenteval.Benchm
if task.Agent.Truncated {
converted.Truncated = true
}
for _, failure := range agentEvalFailuresFromTaskReport(task) {
converted.Failures = append(converted.Failures, failure)
}
converted.Failures = append(converted.Failures, agentEvalFailuresFromTaskReport(task)...)
}
return converted
}
+64 -19
View File
@@ -7,6 +7,7 @@ import (
"fmt"
"io"
"os"
"path/filepath"
"strings"
"time"
@@ -23,6 +24,7 @@ import (
"github.com/Gitlawb/zero/internal/sessions"
"github.com/Gitlawb/zero/internal/skills"
"github.com/Gitlawb/zero/internal/specialist"
"github.com/Gitlawb/zero/internal/swarm"
"github.com/Gitlawb/zero/internal/tools"
"github.com/Gitlawb/zero/internal/tui"
"github.com/Gitlawb/zero/internal/update"
@@ -240,6 +242,8 @@ func runWithDeps(args []string, stdout io.Writer, stderr io.Writer, deps appDeps
case "exec":
// Forward leading --add-dir occurrences so exec's own parser collects them.
return runExec(append(addDirFlagArgs(addDirs), args[1:]...), stdout, stderr, deps)
case "daemon":
return runDaemon(args[1:], stdout, stderr, deps)
case "config":
return runConfig(args[1:], stdout, stderr, deps)
case "models":
@@ -415,6 +419,12 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
if err != nil {
return writeAppError(stderr, "failed to resolve user config path: "+err.Error(), 1)
}
doctorUserConfigPath := ""
projectConfigPath := ""
if resolveOptions, optErr := config.DefaultResolveOptions(workspaceRoot); optErr == nil {
doctorUserConfigPath = resolveOptions.UserConfigPath
projectConfigPath = resolveOptions.ProjectConfigPath
}
needsSetup := setupRequired(resolved)
setupVisible := forceSetup || needsSetup
@@ -507,20 +517,24 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
})
lastKnownMCPConfig := mcpConfig
return deps.runTUI(context.Background(), tui.Options{
Cwd: workspaceRoot,
UserConfigPath: userConfigPath,
ProviderName: resolved.Provider.Name,
ModelName: resolved.Provider.Model,
ProviderProfile: resolved.Provider,
FavoriteModels: resolved.Preferences.FavoriteModels,
Provider: provider,
NewProvider: deps.newProvider,
Registry: registry,
SessionStore: deps.newSessionStore(),
SandboxStore: sandboxStore,
MCPConfig: mcpConfig,
MCPPermissionStore: mcpPermissionStore,
MCPTokenStore: mcpTokenStore,
Cwd: workspaceRoot,
UserConfigPath: userConfigPath,
DoctorUserConfigPath: doctorUserConfigPath,
ProjectConfigPath: projectConfigPath,
ProviderName: resolved.Provider.Name,
ModelName: resolved.Provider.Model,
ProviderProfile: resolved.Provider,
FavoriteModels: resolved.Preferences.FavoriteModels,
Provider: provider,
NewProvider: deps.newProvider,
ProbeProviderHealth: deps.probeProviderHealth,
UserAgent: userAgent(),
Registry: registry,
SessionStore: deps.newSessionStore(),
SandboxStore: sandboxStore,
MCPConfig: mcpConfig,
MCPPermissionStore: mcpPermissionStore,
MCPTokenStore: mcpTokenStore,
MCPCommand: func(ctx context.Context, args []string) tui.MCPCommandResult {
if ctx == nil {
ctx = context.Background()
@@ -582,12 +596,37 @@ func newCoreRegistryScoped(workspaceRoot string, scope tools.PathScope) *tools.R
return registry
}
func registerSpecialistTools(registry *tools.Registry, workspaceRoot string) (*specialist.Runtime, error) {
// agentToolRuntime bundles the specialist runtime with the swarm it backs so
// their lifetimes (and shutdown) stay paired: registerSpecialistTools brings up
// both, closeSpecialistRuntime tears down both.
type agentToolRuntime struct {
specialist *specialist.Runtime
swarm *swarm.Swarm
}
func registerSpecialistTools(registry *tools.Registry, workspaceRoot string) (*agentToolRuntime, error) {
paths, err := specialist.DefaultPaths(workspaceRoot)
if err != nil {
return nil, err
}
return specialist.RegisterTools(registry, specialist.Executor{Paths: paths})
executor := specialist.Executor{Paths: paths}
runtime, err := specialist.RegisterTools(registry, executor)
if err != nil {
return nil, err
}
// The swarm reuses the same specialist executor to launch each member, so
// every member runs under the orchestrator's sandbox + policy. Mailbox state
// lives under the workspace so its files fall within the sandbox write rules.
sw, err := swarm.New(swarm.Options{
BaseDir: filepath.Join(workspaceRoot, ".zero", "swarm"),
Launcher: swarm.NewSpecialistLauncher(executor),
})
if err != nil {
runtime.Close()
return nil, err
}
swarm.RegisterTools(registry, sw)
return &agentToolRuntime{specialist: runtime, swarm: sw}, nil
}
func shouldRegisterExecSpecialistTools(options execOptions) bool {
@@ -609,12 +648,17 @@ func closeMCPRuntime(stderr io.Writer, runtime mcpToolRuntime) {
}
}
func closeSpecialistRuntime(stderr io.Writer, runtime *specialist.Runtime) {
func closeSpecialistRuntime(stderr io.Writer, runtime *agentToolRuntime) {
if runtime == nil {
return
}
if err := runtime.Close(); err != nil {
_, _ = fmt.Fprintf(stderr, "[zero] specialist_cleanup_error: %s\n", err)
if runtime.swarm != nil {
runtime.swarm.Close()
}
if runtime.specialist != nil {
if err := runtime.specialist.Close(); err != nil {
_, _ = fmt.Fprintf(stderr, "[zero] specialist_cleanup_error: %s\n", err)
}
}
}
@@ -637,6 +681,7 @@ Usage:
Commands:
exec Run a one-shot prompt through the Go agent runtime
daemon Manage the local background worker daemon (start/stop/status/run/attach)
setup Guide first-run provider setup
config Inspect resolved Go configuration without leaking secrets
models List Zero model registry entries
+35
View File
@@ -8,6 +8,7 @@ import (
"io"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"time"
@@ -59,10 +60,38 @@ func TestRunPrintsHelp(t *testing.T) {
}
}
func setCLIUserConfigRoot(t *testing.T) string {
t.Helper()
root := t.TempDir()
switch runtime.GOOS {
case "windows":
t.Setenv("APPDATA", root)
case "darwin":
t.Setenv("HOME", root)
default:
t.Setenv("XDG_CONFIG_HOME", root)
}
configRoot, err := os.UserConfigDir()
if err != nil {
t.Fatalf("UserConfigDir() error = %v", err)
}
return configRoot
}
func TestRunNoArgsLaunchesSetupTUIWithNilProviderWhenNoProviderConfigured(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
cwd := t.TempDir()
setCLIUserConfigRoot(t)
projectConfigPath := filepath.Join(cwd, ".zero", "config.json")
if err := os.MkdirAll(filepath.Dir(projectConfigPath), 0o700); err != nil {
t.Fatalf("create project config parent: %v", err)
}
if err := os.WriteFile(projectConfigPath, []byte("{}"), 0o600); err != nil {
t.Fatalf("write project config: %v", err)
}
userConfigPath := filepath.Join(t.TempDir(), "zero", "config.json")
var launchedOptions tui.Options
@@ -104,6 +133,12 @@ func TestRunNoArgsLaunchesSetupTUIWithNilProviderWhenNoProviderConfigured(t *tes
if launchedOptions.UserConfigPath != userConfigPath {
t.Fatalf("UserConfigPath = %q, want %q", launchedOptions.UserConfigPath, userConfigPath)
}
if launchedOptions.DoctorUserConfigPath != "" {
t.Fatalf("DoctorUserConfigPath = %q, want empty for missing user config", launchedOptions.DoctorUserConfigPath)
}
if launchedOptions.ProjectConfigPath != projectConfigPath {
t.Fatalf("ProjectConfigPath = %q, want %q", launchedOptions.ProjectConfigPath, projectConfigPath)
}
if launchedOptions.Provider != nil {
t.Fatalf("Provider = %#v, want nil", launchedOptions.Provider)
}
+122
View File
@@ -17,6 +17,13 @@ type backendStatusOptions struct {
}
func runBackends(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int {
if len(args) > 0 {
switch args[0] {
case "doctor":
return runBackendsDoctor(args[1:], stdout, stderr, deps)
}
}
options, help, err := parseBackendsArgs(args)
if err != nil {
return writeExecUsageError(stderr, err.Error())
@@ -44,6 +51,44 @@ func runBackends(args []string, stdout io.Writer, stderr io.Writer, deps appDeps
return exitSuccess
}
type backendDoctorOptions struct {
json bool
}
func runBackendsDoctor(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int {
options, help, err := parseBackendsDoctorArgs(args)
if err != nil {
return writeExecUsageError(stderr, err.Error())
}
if help {
if err := writeBackendsDoctorHelp(stdout); err != nil {
return exitCrash
}
return exitSuccess
}
report, err := backendDoctorReport(deps)
if err != nil {
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
if options.json {
if err := writePrettyJSON(stdout, report); err != nil {
return exitCrash
}
if !report.OK {
return exitProvider
}
return exitSuccess
}
if _, err := fmt.Fprintln(stdout, redaction.RedactString(formatBackendDoctorReport(report), redaction.Options{})); err != nil {
return exitCrash
}
if !report.OK {
return exitProvider
}
return exitSuccess
}
func parseBackendsArgs(args []string) (backendStatusOptions, bool, error) {
options := backendStatusOptions{}
for _, arg := range args {
@@ -59,6 +104,21 @@ func parseBackendsArgs(args []string) (backendStatusOptions, bool, error) {
return options, false, nil
}
func parseBackendsDoctorArgs(args []string) (backendDoctorOptions, bool, error) {
options := backendDoctorOptions{}
for _, arg := range args {
switch arg {
case "-h", "--help", "help":
return options, true, nil
case "--json":
options.json = true
default:
return options, false, execUsageError{fmt.Sprintf("unknown backends doctor flag %q", arg)}
}
}
return options, false, nil
}
func backendLifecycleSnapshot(deps appDeps) (zerocommands.BackendLifecycleSnapshot, error) {
cwd, err := deps.getwd()
if err != nil {
@@ -86,6 +146,32 @@ func backendLifecycleSnapshot(deps appDeps) (zerocommands.BackendLifecycleSnapsh
return zerocommands.NewBackendLifecycleSnapshot(servers, hookResult.Config.Hooks, pluginResult.Plugins), nil
}
func backendDoctorReport(deps appDeps) (zerocommands.BackendDoctorReport, error) {
cwd, err := deps.getwd()
if err != nil {
return zerocommands.BackendDoctorReport{}, fmt.Errorf("failed to resolve workspace: %w", err)
}
cfg, err := deps.resolveMCPConfig(cwd)
if err != nil {
return zerocommands.BackendDoctorReport{}, err
}
hookResult, err := deps.loadHooks(hooks.LoadOptions{Cwd: cwd})
if err != nil {
return zerocommands.BackendDoctorReport{}, err
}
pluginResult, err := deps.loadPlugins(plugins.LoadOptions{Cwd: cwd})
if err != nil {
return zerocommands.BackendDoctorReport{}, err
}
return zerocommands.NewBackendDoctorReport(zerocommands.BackendDoctorInput{
MCP: cfg,
Hooks: hookResult,
Plugins: pluginResult,
}), nil
}
func formatBackendLifecycleSnapshot(snapshot zerocommands.BackendLifecycleSnapshot) string {
lines := []string{"Zero Backends:"}
lines = append(lines, fmt.Sprintf(" MCP servers: %d", len(snapshot.MCPServers)))
@@ -135,16 +221,52 @@ func formatBackendLifecycleSnapshot(snapshot zerocommands.BackendLifecycleSnapsh
return strings.Join(lines, "\n")
}
func formatBackendDoctorReport(report zerocommands.BackendDoctorReport) string {
lines := []string{
"Zero backend doctor",
"Overall: " + string(report.Status),
}
for _, check := range report.Checks {
target := ""
if check.Target != "" {
target = " " + check.Target
}
lines = append(lines, fmt.Sprintf("[%s] %s%s - %s", check.Status, check.ID, target, check.Message))
if check.Action != "" {
lines = append(lines, " action: "+check.Action)
}
}
return strings.Join(lines, "\n")
}
func writeBackendsHelp(w io.Writer) error {
_, err := fmt.Fprint(w, `Usage:
zero backends [flags]
zero backends doctor [flags]
Inspect MCP, hook, and plugin backend lifecycle state without connecting to
external MCP servers.
Commands:
doctor Diagnose MCP, hook, and plugin backend setup
Flags:
--json Print backend lifecycle data as JSON
-h, --help Show this help
`)
return err
}
func writeBackendsDoctorHelp(w io.Writer) error {
_, err := fmt.Fprint(w, `Usage:
zero backends doctor [flags]
Diagnose MCP, hook, and plugin backend setup without connecting to external MCP
servers or executing hooks/plugins.
Flags:
--json Print backend diagnostics as JSON
-h, --help Show this help
`)
return err
}
+222
View File
@@ -5,7 +5,12 @@ import (
"context"
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"sync/atomic"
"testing"
"github.com/Gitlawb/zero/internal/config"
@@ -153,3 +158,220 @@ func TestRunBackendsTextAndHelp(t *testing.T) {
}
}
}
func TestRunBackendsDoctorJSONAndTextWithoutConnectingMCP(t *testing.T) {
cwd := t.TempDir()
secret := "sk-proj-" + strings.Repeat("b", 24)
deps := appDeps{
getwd: func() (string, error) { return cwd, nil },
resolveMCPConfig: func(workspaceRoot string) (config.MCPConfig, error) {
if workspaceRoot != cwd {
t.Fatalf("workspaceRoot = %q, want %q", workspaceRoot, cwd)
}
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
"remote": {
Type: "http",
URL: "https://api.example.com/mcp?token=" + secret,
},
"broken": {
Type: "http",
},
}}, nil
},
loadHooks: func(options hooks.LoadOptions) (hooks.LoadResult, error) {
if options.Cwd != cwd {
t.Fatalf("hook Cwd = %q, want %q", options.Cwd, cwd)
}
return hooks.LoadResult{Diagnostics: []hooks.Diagnostic{{
Kind: hooks.DiagnosticSchema,
Message: "schema failed " + secret,
HookID: "zero.bad",
}}}, nil
},
loadPlugins: func(options plugins.LoadOptions) (plugins.LoadResult, error) {
if options.Cwd != cwd {
t.Fatalf("plugin Cwd = %q, want %q", options.Cwd, cwd)
}
return plugins.LoadResult{Plugins: []plugins.LoadedPlugin{{
ID: "zero.docs",
Name: "Docs",
Enabled: true,
Source: plugins.SourceProject,
}}}, nil
},
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
return nil, errors.New("zero backends doctor must not connect to MCP servers")
},
}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"backends", "doctor", "--json"}, &stdout, &stderr, deps)
if exitCode != exitProvider {
t.Fatalf("exitCode = %d, want %d stderr=%s stdout=%s", exitCode, exitProvider, stderr.String(), stdout.String())
}
if strings.Contains(stdout.String(), secret) || strings.Contains(stdout.String(), "sk-proj-") {
t.Fatalf("backend doctor JSON leaked secret material:\n%s", stdout.String())
}
var payload zerocommands.BackendDoctorReport
if err := json.Unmarshal(stdout.Bytes(), &payload); err != nil {
t.Fatalf("backend doctor JSON failed to decode: %v\n%s", err, stdout.String())
}
if payload.OK {
t.Fatalf("payload.OK = true, want false for invalid MCP/hook diagnostic: %#v", payload.Checks)
}
if payload.Status != zerocommands.BackendDoctorStatusFail {
t.Fatalf("payload.Status = %q, want %q", payload.Status, zerocommands.BackendDoctorStatusFail)
}
assertBackendDoctorPayloadCheck(t, payload, "backend.mcp.invalid", "broken")
assertBackendDoctorPayloadCheck(t, payload, "backend.hooks.diagnostic", "zero.bad")
stdout.Reset()
stderr.Reset()
exitCode = runWithDeps([]string{"backends", "doctor"}, &stdout, &stderr, deps)
if exitCode != exitProvider {
t.Fatalf("text exitCode = %d, want %d stderr=%s", exitCode, exitProvider, stderr.String())
}
for _, want := range []string{"Zero backend doctor", "[fail] backend.mcp.invalid", "zero mcp add broken", "[fail] backend.hooks.diagnostic"} {
if !strings.Contains(stdout.String(), want) {
t.Fatalf("backend doctor text missing %q:\n%s", want, stdout.String())
}
}
}
func TestRunBackendsDoctorHelp(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"backends", "doctor", "--help"}, &stdout, &stderr, appDeps{})
if exitCode != exitSuccess {
t.Fatalf("exitCode = %d stderr=%s", exitCode, stderr.String())
}
for _, want := range []string{"Usage:", "zero backends doctor", "--json"} {
if !strings.Contains(stdout.String(), want) {
t.Fatalf("backend doctor help missing %q:\n%s", want, stdout.String())
}
}
}
func TestRunBackendsDoctorDoesNotConnectOrExecuteConfiguredBackends(t *testing.T) {
cwd := t.TempDir()
hookSentinel := filepath.Join(cwd, "hook-ran")
pluginToolSentinel := filepath.Join(cwd, "plugin-tool-ran")
pluginHookSentinel := filepath.Join(cwd, "plugin-hook-ran")
helperCommand := os.Args[0]
helperArgs := func(path string) []string {
return []string{"-test.run=TestBackendDoctorHelperProcess", "--", "--zero-backend-doctor-sentinel", path}
}
var hits int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
atomic.AddInt32(&hits, 1)
w.WriteHeader(http.StatusOK)
}))
defer server.Close()
hookConfigPath := filepath.Join(cwd, ".zero", "hooks.json")
writeBackendDoctorJSON(t, hookConfigPath, map[string]any{
"enabled": true,
"hooks": []any{map[string]any{
"id": "zero.sentinel",
"event": "sessionStart",
"command": helperCommand,
"args": helperArgs(hookSentinel),
}},
})
pluginDir := filepath.Join(cwd, ".zero", "plugins", "sentinel")
writeBackendDoctorJSON(t, filepath.Join(pluginDir, "plugin.json"), map[string]any{
"schemaVersion": 1,
"id": "zero.sentinel",
"name": "Sentinel",
"version": "1.0.0",
"tools": []any{map[string]any{
"name": "sentinel_tool",
"command": helperCommand,
"args": helperArgs(pluginToolSentinel),
}},
"hooks": []any{map[string]any{
"name": "sentinel_hook",
"event": "sessionStart",
"command": helperCommand,
"args": helperArgs(pluginHookSentinel),
}},
})
deps := appDeps{
getwd: func() (string, error) { return cwd, nil },
resolveMCPConfig: func(string) (config.MCPConfig, error) {
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
"remote": {Type: "http", URL: server.URL + "/mcp"},
}}, nil
},
loadHooks: func(options hooks.LoadOptions) (hooks.LoadResult, error) {
return hooks.LoadConfig(hooks.LoadOptions{
Cwd: options.Cwd,
UserConfigPath: filepath.Join(cwd, "missing-user-hooks.json"),
ProjectConfigPath: hookConfigPath,
})
},
loadPlugins: func(options plugins.LoadOptions) (plugins.LoadResult, error) {
return plugins.Load(plugins.LoadOptions{
Roots: []plugins.Root{{Source: plugins.SourceProject, Path: filepath.Join(cwd, ".zero", "plugins")}},
Cwd: options.Cwd,
})
},
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
return nil, errors.New("zero backends doctor must not register MCP tools")
},
}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"backends", "doctor", "--json"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("exitCode = %d stderr=%s stdout=%s", exitCode, stderr.String(), stdout.String())
}
if got := atomic.LoadInt32(&hits); got != 0 {
t.Fatalf("backend doctor connected to MCP HTTP server %d time(s)", got)
}
for _, sentinel := range []string{hookSentinel, pluginToolSentinel, pluginHookSentinel} {
if _, err := os.Stat(sentinel); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("backend doctor executed configured backend command, sentinel=%s err=%v", sentinel, err)
}
}
}
func assertBackendDoctorPayloadCheck(t *testing.T, report zerocommands.BackendDoctorReport, id string, target string) {
t.Helper()
for _, check := range report.Checks {
if check.ID == id && check.Target == target {
return
}
}
t.Fatalf("check %s/%s not found in %#v", id, target, report.Checks)
}
func writeBackendDoctorJSON(t *testing.T, path string, value any) {
t.Helper()
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
t.Fatalf("create dir for %s: %v", path, err)
}
data, err := json.MarshalIndent(value, "", " ")
if err != nil {
t.Fatalf("marshal %s: %v", path, err)
}
if err := os.WriteFile(path, data, 0o600); err != nil {
t.Fatalf("write %s: %v", path, err)
}
}
func TestBackendDoctorHelperProcess(t *testing.T) {
for index, arg := range os.Args {
if arg != "--zero-backend-doctor-sentinel" || index+1 >= len(os.Args) {
continue
}
if err := os.WriteFile(os.Args[index+1], []byte("executed"), 0o600); err != nil {
os.Exit(2)
}
os.Exit(0)
}
}
+331
View File
@@ -0,0 +1,331 @@
package cli
import (
"fmt"
"io"
"os"
"os/exec"
"os/signal"
"path/filepath"
"strings"
"syscall"
"time"
"github.com/Gitlawb/zero/internal/background"
"github.com/Gitlawb/zero/internal/daemon"
)
// runDaemon dispatches the `zero daemon ...` subcommands. The daemon supervises a
// pool of headless `zero exec` workers and routes sessions to them over an
// owner-only local control socket. It is an ADDITIVE surface — interactive and
// one-shot exec are unchanged.
func runDaemon(args []string, stdout io.Writer, stderr io.Writer, _ appDeps) int {
if len(args) == 0 {
return writeDaemonUsage(stderr, exitUsage)
}
sub, rest := args[0], args[1:]
switch sub {
case "start":
return runDaemonStart(rest, stdout, stderr)
case "stop":
return runDaemonStop(rest, stdout, stderr)
case "status":
return runDaemonStatus(rest, stdout, stderr)
case "run":
return runDaemonRun(rest, stdout, stderr)
case "attach":
return runDaemonAttach(rest, stdout, stderr)
case "-h", "--help", "help":
return writeDaemonUsage(stdout, exitSuccess)
default:
fmt.Fprintf(stderr, "unknown daemon subcommand %q\n", sub)
return writeDaemonUsage(stderr, exitUsage)
}
}
func writeDaemonUsage(w io.Writer, code int) int {
fmt.Fprint(w, `Usage: zero daemon <command>
Commands:
start [--foreground] Start the daemon (background by default).
stop Gracefully stop the running daemon.
status Show daemon / pool / session status.
run --session <id> [--cwd <dir>] [--prompt <text>] [exec flags...]
Create/route a session and stream its output.
attach <session> Attach to a running session's stream.
`)
return code
}
// runDaemonStart starts the daemon. Without --foreground it spawns a detached
// background process running the foreground daemon and returns once it is up.
func runDaemonStart(args []string, stdout io.Writer, stderr io.Writer) int {
foreground := false
for _, a := range args {
switch a {
case "--foreground", "-f":
foreground = true
case "-h", "--help":
return writeDaemonUsage(stdout, exitSuccess)
default:
return writeExecUsageError(stderr, fmt.Sprintf("unknown flag %q for daemon start", a))
}
}
paths, err := daemon.DefaultPaths()
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
if foreground {
return runDaemonForeground(paths, stdout, stderr)
}
return runDaemonStartDetached(paths, stdout, stderr)
}
// runDaemonForeground runs the daemon in this process until SIGINT/SIGTERM. Each
// worker is a headless `zero exec` child with the sandbox re-entrancy markers
// scrubbed (NewExecLauncher) so it establishes its own sandbox.
func runDaemonForeground(paths daemon.Paths, stdout io.Writer, stderr io.Writer) int {
launcher, err := daemon.NewExecLauncher(daemon.ExecLauncherConfig{})
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
logf := func(line string) { fmt.Fprintln(stderr, "[daemon] "+line) }
pool, err := daemon.NewPool(daemon.PoolOptions{Launcher: launcher, Log: logf})
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
mgr, err := daemon.NewSessionManager(daemon.SessionManagerOptions{Pool: pool})
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
srv, err := daemon.NewServer(daemon.ServerOptions{Paths: paths, Manager: mgr, Pool: pool, Log: logf})
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
sigCh := make(chan os.Signal, 1)
signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
go func() {
<-sigCh
srv.Shutdown()
}()
if err := srv.Serve(); err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
return exitSuccess
}
// runDaemonStartDetached spawns the foreground daemon as a detached background
// process (its own process group, output to a log file) and waits for it to bind.
func runDaemonStartDetached(paths daemon.Paths, stdout io.Writer, stderr io.Writer) int {
if daemonReachable(paths) {
fmt.Fprintln(stdout, "zero daemon is already running")
return exitSuccess
}
exe, err := os.Executable()
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
if err := os.MkdirAll(filepath.Dir(paths.Socket), 0o700); err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
logPath := filepath.Join(filepath.Dir(paths.Socket), "daemon.log")
logFile, err := os.OpenFile(logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600)
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
defer logFile.Close()
cmd := exec.Command(exe, "daemon", "start", "--foreground")
cmd.Stdout = logFile
cmd.Stderr = logFile
background.ConfigureChildProcessGroup(cmd) // own process group: outlives this shell
if err := cmd.Start(); err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
_ = cmd.Process.Release()
deadline := time.Now().Add(5 * time.Second)
for time.Now().Before(deadline) {
if daemonReachable(paths) {
fmt.Fprintf(stdout, "zero daemon started (socket %s)\n", paths.Socket)
return exitSuccess
}
time.Sleep(25 * time.Millisecond)
}
return writeAppError(stderr, "daemon did not come up within timeout; see "+logPath, exitCrash)
}
func runDaemonStop(args []string, stdout io.Writer, stderr io.Writer) int {
if helpRequested(args) {
return writeDaemonUsage(stdout, exitSuccess)
}
if len(args) > 0 {
return writeExecUsageError(stderr, "daemon stop takes no arguments")
}
paths, err := daemon.DefaultPaths()
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
client, err := daemon.Dial(paths.Socket)
if err != nil {
fmt.Fprintln(stdout, "zero daemon is not running")
return exitSuccess
}
defer client.Close()
if err := client.Shutdown(); err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
fmt.Fprintln(stdout, "zero daemon stopped")
return exitSuccess
}
func runDaemonStatus(args []string, stdout io.Writer, stderr io.Writer) int {
if helpRequested(args) {
return writeDaemonUsage(stdout, exitSuccess)
}
if len(args) > 0 {
return writeExecUsageError(stderr, "daemon status takes no arguments")
}
paths, err := daemon.DefaultPaths()
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
client, err := daemon.Dial(paths.Socket)
if err != nil {
fmt.Fprintln(stdout, "zero daemon is not running")
return exitSuccess
}
defer client.Close()
report, err := client.Status()
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
fmt.Fprintf(stdout, "daemon pid=%d version=%d socket=%s\n", report.PID, report.Version, report.Socket)
fmt.Fprintf(stdout, "pool size=%d busy=%d queue=%d\n", report.PoolSize, len(report.Workers), report.QueueDepth)
for _, s := range report.Sessions {
fmt.Fprintf(stdout, " session %s state=%s lines=%d\n", s.ID, s.State, s.Lines)
}
return exitSuccess
}
func runDaemonRun(args []string, stdout io.Writer, stderr io.Writer) int {
session, cwd, prompt := "", "", ""
var forward []string
for i := 0; i < len(args); i++ {
a := args[i]
value := func() (string, bool) {
if i+1 >= len(args) {
return "", false
}
i++
return args[i], true
}
switch {
case a == "-h" || a == "--help":
return writeDaemonUsage(stdout, exitSuccess)
case a == "--session":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--session requires a value")
}
session = v
case strings.HasPrefix(a, "--session="):
session = strings.TrimPrefix(a, "--session=")
case a == "--cwd":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--cwd requires a value")
}
cwd = v
case strings.HasPrefix(a, "--cwd="):
cwd = strings.TrimPrefix(a, "--cwd=")
case a == "--prompt" || a == "-p":
v, ok := value()
if !ok {
return writeExecUsageError(stderr, "--prompt requires a value")
}
prompt = v
case strings.HasPrefix(a, "--prompt="):
prompt = strings.TrimPrefix(a, "--prompt=")
default:
// Forwarded verbatim to the worker `zero exec` (reuses exec's own flag
// parsing for per-session run options).
forward = append(forward, a)
}
}
if strings.TrimSpace(session) == "" {
return writeExecUsageError(stderr, "daemon run requires --session <id>")
}
if prompt == "" && len(forward) == 0 {
return writeExecUsageError(stderr, "daemon run requires --prompt <text> or exec args")
}
paths, err := daemon.DefaultPaths()
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
client, err := daemon.Dial(paths.Socket)
if err != nil {
return writeAppError(stderr, "zero daemon is not running (start it with `zero daemon start`)", exitCrash)
}
defer client.Close()
if err := client.Run(session, cwd, prompt, forward, func(line string) { fmt.Fprintln(stdout, line) }); err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
return exitSuccess
}
func runDaemonAttach(args []string, stdout io.Writer, stderr io.Writer) int {
session := ""
extra := 0
for _, a := range args {
if a == "-h" || a == "--help" {
return writeDaemonUsage(stdout, exitSuccess)
}
if session == "" {
session = a
continue
}
extra++
}
if strings.TrimSpace(session) == "" {
return writeExecUsageError(stderr, "daemon attach requires a <session> id")
}
if extra > 0 {
return writeExecUsageError(stderr, "daemon attach accepts exactly one <session> id")
}
paths, err := daemon.DefaultPaths()
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
client, err := daemon.Dial(paths.Socket)
if err != nil {
return writeAppError(stderr, "zero daemon is not running (start it with `zero daemon start`)", exitCrash)
}
defer client.Close()
if err := client.Attach(session, func(line string) { fmt.Fprintln(stdout, line) }); err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
}
return exitSuccess
}
// daemonReachable reports whether a daemon is accepting connections (a successful
// handshake). Used for the single-instance check and start-up wait.
func daemonReachable(paths daemon.Paths) bool {
client, err := daemon.Dial(paths.Socket)
if err != nil {
return false
}
_ = client.Close()
return true
}
func helpRequested(args []string) bool {
for _, a := range args {
if a == "-h" || a == "--help" {
return true
}
}
return false
}
+123
View File
@@ -0,0 +1,123 @@
package cli
import (
"bytes"
"strings"
"testing"
)
// isolateDaemonPaths points DefaultPaths at a temp dir so the test never touches
// a real daemon on the dev machine.
func isolateDaemonPaths(t *testing.T) {
t.Helper()
t.Setenv("XDG_RUNTIME_DIR", t.TempDir())
}
func runDaemonCLI(t *testing.T, args ...string) (int, string, string) {
t.Helper()
var out, errb bytes.Buffer
code := runDaemon(args, &out, &errb, appDeps{})
return code, out.String(), errb.String()
}
func TestDaemonUsage(t *testing.T) {
code, _, _ := runDaemonCLI(t)
if code != exitUsage {
t.Fatalf("no-args exit = %d, want exitUsage", code)
}
code, out, _ := runDaemonCLI(t, "--help")
if code != exitSuccess || !strings.Contains(out, "Usage: zero daemon") {
t.Fatalf("--help exit=%d out=%q", code, out)
}
}
func TestDaemonUnknownSubcommand(t *testing.T) {
code, _, errb := runDaemonCLI(t, "frobnicate")
if code != exitUsage {
t.Fatalf("unknown subcommand exit = %d, want exitUsage", code)
}
if !strings.Contains(errb, "unknown daemon subcommand") {
t.Fatalf("stderr = %q, want unknown-subcommand message", errb)
}
}
func TestDaemonRunRequiresSession(t *testing.T) {
isolateDaemonPaths(t)
code, _, errb := runDaemonCLI(t, "run", "--prompt", "hi")
if code != exitUsage {
t.Fatalf("run without --session exit = %d, want exitUsage", code)
}
if !strings.Contains(errb, "--session") {
t.Fatalf("stderr = %q, want a --session hint", errb)
}
}
func TestDaemonRunRequiresPromptOrArgs(t *testing.T) {
isolateDaemonPaths(t)
code, _, errb := runDaemonCLI(t, "run", "--session", "s1")
if code != exitUsage {
t.Fatalf("run without prompt/args exit = %d, want exitUsage", code)
}
if !strings.Contains(errb, "--prompt") {
t.Fatalf("stderr = %q, want a --prompt hint", errb)
}
}
func TestDaemonStopWhenNotRunning(t *testing.T) {
isolateDaemonPaths(t)
code, out, _ := runDaemonCLI(t, "stop")
if code != exitSuccess {
t.Fatalf("stop (not running) exit = %d, want exitSuccess", code)
}
if !strings.Contains(out, "not running") {
t.Fatalf("stop output = %q, want 'not running'", out)
}
}
func TestDaemonStatusWhenNotRunning(t *testing.T) {
isolateDaemonPaths(t)
code, out, _ := runDaemonCLI(t, "status")
if code != exitSuccess {
t.Fatalf("status (not running) exit = %d, want exitSuccess", code)
}
if !strings.Contains(out, "not running") {
t.Fatalf("status output = %q, want 'not running'", out)
}
}
func TestDaemonAttachRequiresSession(t *testing.T) {
isolateDaemonPaths(t)
code, _, errb := runDaemonCLI(t, "attach")
if code != exitUsage {
t.Fatalf("attach without session exit = %d, want exitUsage", code)
}
if !strings.Contains(errb, "session") {
t.Fatalf("stderr = %q, want a session hint", errb)
}
}
func TestDaemonRunWhenNotRunning(t *testing.T) {
isolateDaemonPaths(t)
code, _, errb := runDaemonCLI(t, "run", "--session", "s1", "--prompt", "hello")
if code != exitCrash {
t.Fatalf("run (no daemon) exit = %d, want exitCrash", code)
}
if !strings.Contains(errb, "not running") {
t.Fatalf("stderr = %q, want 'not running'", errb)
}
}
func TestDaemonSubcommandsRejectExtraArgs(t *testing.T) {
isolateDaemonPaths(t)
cases := [][]string{
{"stop", "oops"},
{"status", "oops"},
{"attach", "s1", "extra"},
}
for _, args := range cases {
code, _, errb := runDaemonCLI(t, args...)
if code != exitUsage {
t.Fatalf("%v exit = %d, want exitUsage (reject extra args); stderr=%q", args, code, errb)
}
}
}
+9 -3
View File
@@ -19,7 +19,6 @@ import (
"github.com/Gitlawb/zero/internal/providers"
"github.com/Gitlawb/zero/internal/sandbox"
"github.com/Gitlawb/zero/internal/sessions"
"github.com/Gitlawb/zero/internal/specialist"
"github.com/Gitlawb/zero/internal/specmode"
"github.com/Gitlawb/zero/internal/streamjson"
"github.com/Gitlawb/zero/internal/tools"
@@ -163,7 +162,7 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
if options.allowEscalation {
registry.Register(tools.NewEscalateModelTool())
}
var specialistRuntime *specialist.Runtime
var specialistRuntime *agentToolRuntime
if shouldRegisterExecSpecialistTools(options) {
var err error
specialistRuntime, err = registerSpecialistTools(registry, workspaceRoot)
@@ -704,7 +703,8 @@ func buildExecSandboxEngine(workspaceRoot string, resolved config.ResolvedConfig
// depth — it guarantees that even a direct/programmatic caller passing a typo
// can never WIDEN the ceiling back to the High default.
// applyConfiguredSandboxPolicy overlays every config-sourced sandbox knob onto
// the default policy: the autonomy ceiling and the network mode.
// the default policy (the autonomy ceiling and the network mode), plus the
// ZERO_SANDBOX_AUTO_ALLOW_BASH environment opt-in.
func applyConfiguredSandboxPolicy(policy sandbox.Policy, cfg config.SandboxConfig) sandbox.Policy {
policy = applyConfiguredAutonomyCeiling(policy, cfg.MaxAutonomy)
if network := strings.TrimSpace(cfg.Network); network != "" {
@@ -721,6 +721,12 @@ func applyConfiguredSandboxPolicy(policy sandbox.Policy, cfg config.SandboxConfi
if cfg.MonitorDenials {
policy.MonitorDenials = true
}
// ZERO_SANDBOX_AUTO_ALLOW_BASH opts into auto-allowing a *sandboxed* shell
// command without a prompt — the active sandbox is the safety boundary, and
// the engine only honors this when the shell sandbox is actually active. Like
// the flags above it only ever turns the opt-in ON; unset/blank stays off
// (prompt), and an explicit config opt-in is preserved.
policy = sandbox.ApplyAutoAllowBashEnv(policy)
return policy
}
+13
View File
@@ -662,6 +662,19 @@ func TestApplyConfiguredSandboxPolicyHardeningFlags(t *testing.T) {
}
}
func TestApplyConfiguredSandboxPolicyHonorsAutoAllowBashEnv(t *testing.T) {
// Unset/blank: the opt-in stays off so a sandboxed bash command still prompts.
t.Setenv(sandbox.EnvAutoAllowBash, "")
if got := applyConfiguredSandboxPolicy(sandbox.DefaultPolicy(), config.SandboxConfig{}); got.AutoAllowBashWhenSandboxed {
t.Fatal("auto-allow-bash must be off when ZERO_SANDBOX_AUTO_ALLOW_BASH is unset")
}
// Truthy: the env opt-in is wired through the central policy builder.
t.Setenv(sandbox.EnvAutoAllowBash, "true")
if got := applyConfiguredSandboxPolicy(sandbox.DefaultPolicy(), config.SandboxConfig{}); !got.AutoAllowBashWhenSandboxed {
t.Fatal("ZERO_SANDBOX_AUTO_ALLOW_BASH=true must enable AutoAllowBashWhenSandboxed")
}
}
func newSandboxTestStore(t *testing.T) *sandbox.GrantStore {
t.Helper()
store, err := sandbox.NewGrantStore(sandbox.StoreOptions{
+26 -1
View File
@@ -1094,7 +1094,8 @@ func TestResolveProviderProfileExtendedJSONAliases(t *testing.T) {
"auth_scheme": "raw",
"auth_header_value": "header-secret",
"custom_headers": {"X-Zero": "1"},
"model_id": "custom-model"
"model_id": "custom-model",
"parse_think_tags": true
}]
}`)
@@ -1119,6 +1120,30 @@ func TestResolveProviderProfileExtendedJSONAliases(t *testing.T) {
if profile.CustomHeaders["X-Zero"] != "1" {
t.Fatalf("CustomHeaders = %#v, want X-Zero header", profile.CustomHeaders)
}
if profile.ParseThinkTags == nil || !*profile.ParseThinkTags {
t.Fatalf("ParseThinkTags = %#v, want true", profile.ParseThinkTags)
}
}
func TestResolveProviderProfileParseThinkTagsFalseAlias(t *testing.T) {
path := writeConfig(t, `{
"activeProvider": "custom",
"providers": [{
"name": "custom",
"provider_kind": "openai-compatible",
"base_url": "https://custom.example/v1",
"model_id": "custom-model",
"parse_think_tags": false
}]
}`)
resolved, err := Resolve(ResolveOptions{UserConfigPath: path, Env: map[string]string{}})
if err != nil {
t.Fatalf("Resolve() error = %v", err)
}
if resolved.Provider.ParseThinkTags == nil || *resolved.Provider.ParseThinkTags {
t.Fatalf("ParseThinkTags = %#v, want explicit false", resolved.Provider.ParseThinkTags)
}
}
func TestResolveRejectsUnknownProviderCatalogID(t *testing.T) {
+14
View File
@@ -34,6 +34,7 @@ type ProviderProfile struct {
AuthHeaderValue string `json:"authHeaderValue,omitempty"`
CustomHeaders map[string]string `json:"customHeaders,omitempty"`
Model string `json:"model,omitempty"`
ParseThinkTags *bool `json:"parseThinkTags,omitempty"`
Description string `json:"description,omitempty"`
}
@@ -51,6 +52,7 @@ func HasProviderProfile(profile ProviderProfile) bool {
strings.TrimSpace(profile.AuthHeaderValue) != "" ||
profile.CustomHeaders != nil ||
strings.TrimSpace(profile.Model) != "" ||
profile.ParseThinkTags != nil ||
strings.TrimSpace(profile.Description) != ""
}
@@ -304,6 +306,8 @@ func (profile *ProviderProfile) UnmarshalJSON(data []byte) error {
CustomHeadersSnake map[string]string `json:"custom_headers"`
Model string `json:"model"`
ModelID string `json:"model_id"`
ParseThinkTags *bool `json:"parseThinkTags"`
ParseThinkTagsSnake *bool `json:"parse_think_tags"`
Description string `json:"description"`
}
@@ -325,6 +329,7 @@ func (profile *ProviderProfile) UnmarshalJSON(data []byte) error {
profile.AuthHeaderValue = firstNonEmpty(raw.AuthHeaderValue, raw.AuthHeaderValueSnake)
profile.CustomHeaders = firstNonNilMap(raw.CustomHeaders, raw.CustomHeadersSnake)
profile.Model = strings.TrimSpace(firstNonEmpty(raw.Model, raw.ModelID))
profile.ParseThinkTags = firstNonNilBool(raw.ParseThinkTags, raw.ParseThinkTagsSnake)
profile.Description = strings.TrimSpace(raw.Description)
return nil
}
@@ -348,6 +353,15 @@ func firstNonNilMap(values ...map[string]string) map[string]string {
return nil
}
func firstNonNilBool(values ...*bool) *bool {
for _, value := range values {
if value != nil {
return value
}
}
return nil
}
func copyStringMap(values map[string]string) map[string]string {
if values == nil {
return nil
+41
View File
@@ -0,0 +1,41 @@
package daemon
import (
"math/rand"
"time"
)
// Backoff bounds. Mirrors reference-daemon-code-agent-js/worker-manager.js:
// backoffMs(consecutive) = jitter(min(1000 * 2**consecutive, 300000)).
const (
backoffUnit = time.Second
backoffCap = 5 * time.Minute
)
// backoffBase returns the un-jittered exponential delay for a 1-based restart
// count: min(1s * 2^attempt, 5m). attempt<=0 yields the unit delay.
func backoffBase(attempt int) time.Duration {
if attempt <= 0 {
return backoffUnit
}
// Shift carefully: cap before overflow. 2^attempt * 1s exceeds the cap well
// before any int64 overflow (cap is 5m), so clamp once it would exceed it.
d := backoffUnit
for i := 0; i < attempt; i++ {
d *= 2
if d >= backoffCap {
return backoffCap
}
}
return d
}
// BackoffWithJitter returns backoffBase(attempt) scaled by a uniform random
// factor in [0.5, 1.5), matching worker-manager.js jitter(base) =
// base*(0.5 + random()). The jitter de-synchronizes simultaneous restarts so a
// fleet of workers does not thunder back at the same instant.
func BackoffWithJitter(attempt int) time.Duration {
base := backoffBase(attempt)
factor := 0.5 + rand.Float64() //nolint:gosec // jitter, not security-sensitive
return time.Duration(float64(base) * factor)
}
+38
View File
@@ -0,0 +1,38 @@
package daemon
import (
"testing"
"time"
)
func TestBackoffBaseExponentialAndCap(t *testing.T) {
cases := []struct {
attempt int
want time.Duration
}{
{0, 1 * time.Second},
{1, 2 * time.Second},
{2, 4 * time.Second},
{3, 8 * time.Second},
{20, backoffCap}, // far past the cap
}
for _, tc := range cases {
if got := backoffBase(tc.attempt); got != tc.want {
t.Fatalf("backoffBase(%d) = %s, want %s", tc.attempt, got, tc.want)
}
}
}
func TestBackoffWithJitterBounds(t *testing.T) {
for attempt := 1; attempt <= 6; attempt++ {
base := backoffBase(attempt)
lo := time.Duration(float64(base) * 0.5)
hi := time.Duration(float64(base) * 1.5)
for i := 0; i < 200; i++ {
got := BackoffWithJitter(attempt)
if got < lo || got >= hi {
t.Fatalf("BackoffWithJitter(%d) = %s, want in [%s,%s)", attempt, got, lo, hi)
}
}
}
}
+117
View File
@@ -0,0 +1,117 @@
package daemon
import (
"fmt"
"net"
"time"
)
// Client is a control-socket client used by the `zero daemon run|attach|status|
// stop` subcommands. It performs the version handshake on Dial.
type Client struct {
conn net.Conn
}
// Dial connects to the daemon control socket and completes the handshake.
func Dial(socketPath string) (*Client, error) {
conn, err := net.DialTimeout("unix", socketPath, 5*time.Second)
if err != nil {
return nil, err
}
c := &Client{conn: conn}
if err := c.handshake(); err != nil {
_ = conn.Close()
return nil, err
}
return c, nil
}
func (c *Client) handshake() error {
if err := WriteControl(c.conn, Ctrl{Type: CtrlHello, Version: ProtoVersion}); err != nil {
return err
}
reply, err := ReadControl(c.conn)
if err != nil {
return err
}
if reply.Type != CtrlHelloOK {
return fmt.Errorf("daemon: handshake rejected: %s", reply.Message)
}
return nil
}
// Close closes the connection.
func (c *Client) Close() error { return c.conn.Close() }
// Run starts/routes a session and streams its stream-json output lines to onLine
// until the session ends. It returns the session's terminal error, if any.
func (c *Client) Run(session, cwd, prompt string, args []string, onLine func(string)) error {
if err := WriteControl(c.conn, Ctrl{Type: CtrlRun, Session: session, Cwd: cwd, Prompt: prompt, Args: args}); err != nil {
return err
}
return c.streamLoop(onLine)
}
// Attach streams a running session's output (buffered history + live) to onLine.
func (c *Client) Attach(session string, onLine func(string)) error {
if err := WriteControl(c.conn, Ctrl{Type: CtrlAttach, Session: session}); err != nil {
return err
}
return c.streamLoop(onLine)
}
func (c *Client) streamLoop(onLine func(string)) error {
for {
msg, err := ReadControl(c.conn)
if err != nil {
return err
}
switch msg.Type {
case CtrlAck:
// dispatch acknowledged; keep reading
case CtrlData:
if onLine != nil {
onLine(msg.Line)
}
case CtrlEnd:
return nil
case CtrlError:
return fmt.Errorf("daemon: %s", msg.Message)
default:
return fmt.Errorf("daemon: unexpected message %q", msg.Type)
}
}
}
// Status requests the daemon/worker/session status report.
func (c *Client) Status() (*StatusReport, error) {
if err := WriteControl(c.conn, Ctrl{Type: CtrlStatus}); err != nil {
return nil, err
}
msg, err := ReadControl(c.conn)
if err != nil {
return nil, err
}
if msg.Type == CtrlError {
return nil, fmt.Errorf("daemon: %s", msg.Message)
}
if msg.Type != CtrlStatusResult || msg.Status == nil {
return nil, fmt.Errorf("daemon: unexpected status reply %q", msg.Type)
}
return msg.Status, nil
}
// Shutdown asks the daemon to drain and stop.
func (c *Client) Shutdown() error {
if err := WriteControl(c.conn, Ctrl{Type: CtrlShutdown}); err != nil {
return err
}
msg, err := ReadControl(c.conn)
if err != nil {
return err
}
if msg.Type == CtrlError {
return fmt.Errorf("daemon: %s", msg.Message)
}
return nil
}
+165
View File
@@ -0,0 +1,165 @@
package daemon
import (
"bufio"
"context"
"errors"
"fmt"
"io"
"os"
"os/exec"
"strings"
"github.com/Gitlawb/zero/internal/background"
"github.com/Gitlawb/zero/internal/sandbox"
)
// reentrancyMarkers are the env vars zero sets on a command it has ALREADY
// wrapped in a sandbox. The daemon MUST strip them from every worker's
// environment: a worker that inherited them would trip zero's re-entrancy guard
// (sandbox.IsAlreadySandboxed needs BOTH) and run with a pass-through plan —
// i.e. UNSANDBOXED — silently defeating the sandbox. Stripping them forces each
// worker to (re)establish its own sandbox. This is the daemon's #1 security
// invariant: it must NOT bypass the sandbox.
var reentrancyMarkers = []string{sandbox.EnvSandboxed, sandbox.EnvSandboxBackend}
// scrubWorkerEnv returns env with the sandbox re-entrancy markers removed. Every
// other variable is preserved — including the ZERO_SANDBOX_* policy config the
// worker reads to rebuild its own policy — so the sandbox policy is PROPAGATED,
// not bypassed. The caller's slice is never mutated.
func scrubWorkerEnv(env []string) []string {
out := make([]string, 0, len(env))
for _, kv := range env {
drop := false
for _, marker := range reentrancyMarkers {
if strings.HasPrefix(kv, marker+"=") {
drop = true
break
}
}
if !drop {
out = append(out, kv)
}
}
return out
}
// execWorker is a WorkerHandle backed by a `zero exec` child process speaking
// stream-json on stdout.
type execWorker struct {
cmd *exec.Cmd
stdout io.ReadCloser
lines *scannerLines
pid int
}
func (w *execWorker) Stdout() Lines { return w.lines }
func (w *execWorker) Pid() int { return w.pid }
func (w *execWorker) Wait() (int, error) {
err := w.cmd.Wait()
if err == nil {
return 0, nil
}
var exitErr *exec.ExitError
if errors.As(err, &exitErr) {
return exitErr.ExitCode(), nil
}
return -1, err
}
func (w *execWorker) Kill() error {
if w.cmd.Process == nil {
return nil
}
// background.TerminateProcess is the cross-platform terminate (kills the
// process group on POSIX, taskkill /T on Windows).
return background.TerminateProcess(w.cmd.Process.Pid)
}
// scannerLines adapts a bufio.Scanner to the Lines interface.
type scannerLines struct {
sc *bufio.Scanner
}
func (s *scannerLines) Next() (string, bool, error) {
if s.sc.Scan() {
return s.sc.Text(), true, nil
}
if err := s.sc.Err(); err != nil {
return "", false, err
}
return "", false, nil
}
// ExecLauncherConfig configures the production worker launcher.
type ExecLauncherConfig struct {
// Executable is the zero binary to spawn (defaults to os.Executable()).
Executable string
// BaseArgs are prepended before the per-session args (defaults to the headless
// stream-json exec invocation `exec --output-format stream-json`).
BaseArgs []string
// Env is the base environment for workers (defaults to os.Environ()). The
// re-entrancy markers are always scrubbed from it.
Env []string
}
// NewExecLauncher builds a Launcher that spawns `zero exec` workers which speak
// stream-json on stdout. Each worker:
// - runs with the re-entrancy markers SCRUBBED from its env (so it establishes
// its own sandbox; the daemon never bypasses the sandbox), while the rest of
// the policy config/env is propagated;
// - is placed in its own process group (background.ConfigureChildProcessGroup)
// so Kill / drain can terminate it and any children cleanly;
// - has its per-session working directory (spec.Cwd) and flags (spec.Args)
// applied. The pool does the os/exec itself and uses internal/background only
// for process-group setup + cross-platform terminate.
func NewExecLauncher(cfg ExecLauncherConfig) (Launcher, error) {
exe := cfg.Executable
if exe == "" {
resolved, err := os.Executable()
if err != nil {
return nil, fmt.Errorf("daemon: locate zero executable: %w", err)
}
exe = resolved
}
baseArgs := cfg.BaseArgs
if baseArgs == nil {
baseArgs = []string{"exec", "--output-format", "stream-json"}
}
baseEnv := cfg.Env
if baseEnv == nil {
baseEnv = os.Environ()
}
env := scrubWorkerEnv(baseEnv)
return func(ctx context.Context, spec WorkerSpec) (WorkerHandle, error) {
args := make([]string, 0, len(baseArgs)+len(spec.Args))
args = append(args, baseArgs...)
args = append(args, spec.Args...)
cmd := exec.CommandContext(ctx, exe, args...)
cmd.Dir = spec.Cwd
cmd.Env = env
cmd.Stdin = nil // the prompt is passed via flags; no streamed stdin input
background.ConfigureChildProcessGroup(cmd)
stdout, err := cmd.StdoutPipe()
if err != nil {
return nil, fmt.Errorf("daemon: worker stdout pipe: %w", err)
}
if err := cmd.Start(); err != nil {
return nil, fmt.Errorf("daemon: start worker: %w", err)
}
scanner := bufio.NewScanner(stdout)
// Allow long stream-json lines up to the control-frame cap.
scanner.Buffer(make([]byte, 0, 64*1024), MaxFrameSize)
return &execWorker{
cmd: cmd,
stdout: stdout,
lines: &scannerLines{sc: scanner},
pid: cmd.Process.Pid,
}, nil
}, nil
}
+102
View File
@@ -0,0 +1,102 @@
package daemon
import (
"context"
"runtime"
"strings"
"testing"
"github.com/Gitlawb/zero/internal/sandbox"
)
func TestScrubWorkerEnvRemovesReentrancyMarkers(t *testing.T) {
in := []string{
"PATH=/bin",
sandbox.EnvSandboxed + "=1",
"ZERO_SANDBOX_AUTO_ALLOW_BASH=1", // policy config — must be PRESERVED
sandbox.EnvSandboxBackend + "=bubblewrap",
"HOME=/home/u",
}
out := scrubWorkerEnv(in)
for _, kv := range out {
if strings.HasPrefix(kv, sandbox.EnvSandboxed+"=") || strings.HasPrefix(kv, sandbox.EnvSandboxBackend+"=") {
t.Fatalf("re-entrancy marker not scrubbed: %q", kv)
}
}
want := map[string]bool{"PATH=/bin": true, "ZERO_SANDBOX_AUTO_ALLOW_BASH=1": true, "HOME=/home/u": true}
got := map[string]bool{}
for _, kv := range out {
got[kv] = true
}
for w := range want {
if !got[w] {
t.Fatalf("scrub dropped a non-marker var: %q", w)
}
}
if len(out) != len(want) {
t.Fatalf("scrub result = %v, want exactly %d preserved vars", out, len(want))
}
if len(in) != 5 {
t.Fatal("scrubWorkerEnv must not mutate the caller's slice")
}
}
// TestExecLauncherScrubsEnvEndToEnd proves the spawned WORKER PROCESS actually
// receives an environment without the re-entrancy markers — the security
// invariant. It uses /bin/sh to dump the worker env, so it is POSIX-only.
func TestExecLauncherScrubsEnvEndToEnd(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses the POSIX env dumper /bin/sh -c env")
}
launcher, err := NewExecLauncher(ExecLauncherConfig{
Executable: "/bin/sh",
BaseArgs: []string{"-c", "env"},
Env: append([]string{"KEEP=yes"},
sandbox.EnvSandboxed+"=1",
sandbox.EnvSandboxBackend+"=bubblewrap"),
})
if err != nil {
t.Fatalf("NewExecLauncher: %v", err)
}
h, err := launcher(context.Background(), WorkerSpec{})
if err != nil {
t.Fatalf("launch: %v", err)
}
lines := h.Stdout()
var got []string
for {
line, ok, lerr := lines.Next()
if lerr != nil {
t.Fatalf("read worker stdout: %v", lerr)
}
if !ok {
break
}
got = append(got, line)
}
if _, err := h.Wait(); err != nil {
t.Fatalf("worker wait: %v", err)
}
keepSeen := false
for _, l := range got {
if strings.HasPrefix(l, sandbox.EnvSandboxed+"=") || strings.HasPrefix(l, sandbox.EnvSandboxBackend+"=") {
t.Fatalf("worker process still has a re-entrancy marker in its env: %q", l)
}
if l == "KEEP=yes" {
keepSeen = true
}
}
if !keepSeen {
t.Fatalf("worker env lost a non-marker var (KEEP=yes); got %v", got)
}
}
func TestNewExecLauncherDefaults(t *testing.T) {
// With no Executable it resolves the test binary's own path (os.Executable),
// and constructs without error.
if _, err := NewExecLauncher(ExecLauncherConfig{}); err != nil {
t.Fatalf("NewExecLauncher with defaults: %v", err)
}
}
+87
View File
@@ -0,0 +1,87 @@
package daemon
import (
"errors"
"fmt"
"io/fs"
"os"
"strconv"
"strings"
)
// Single-instance lock. Mirrors reference-daemon-code-agent-js/supervisor.js's
// lock file: a PID file created with O_EXCL. A second start fails; a STALE lock
// left by a dead daemon (the recorded PID is no longer alive) is reclaimed so the
// daemon recovers from an unclean shutdown without manual cleanup.
// ErrAlreadyRunning is returned when a live daemon already holds the lock.
var ErrAlreadyRunning = errors.New("daemon: another instance is already running")
// fileLock is an acquired single-instance lock.
type fileLock struct {
path string
}
// processAlive reports whether pid is a live process. Implemented per-platform
// (lock_posix.go / lock_windows.go). It is a package var so tests can stub it.
var processAlive = osProcessAlive
// acquireLock takes the single-instance lock at path, reclaiming a stale lock
// whose recorded PID is dead. isAlive overrides the liveness check (tests pass a
// stub); nil uses the real processAlive.
func acquireLock(path string, isAlive func(pid int) bool) (*fileLock, error) {
if isAlive == nil {
isAlive = processAlive
}
// At most two passes: create, or detect-stale-then-retry once.
for attempt := 0; attempt < 2; attempt++ {
f, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
if err == nil {
// A failed PID write would leave a malformed lock file that another
// process reads as stale (unparsable PID) and wrongly reclaims, breaking
// the single-instance guarantee — so on write failure, remove it and fail.
if _, werr := fmt.Fprintf(f, "%d\n", os.Getpid()); werr != nil {
_ = f.Close()
_ = os.Remove(path)
return nil, werr
}
if cerr := f.Close(); cerr != nil {
_ = os.Remove(path)
return nil, cerr
}
return &fileLock{path: path}, nil
}
if !errors.Is(err, fs.ErrExist) {
return nil, err
}
pid, perr := readPidFile(path)
if perr == nil && pid > 0 && isAlive(pid) {
return nil, fmt.Errorf("%w (pid %d)", ErrAlreadyRunning, pid)
}
// Stale lock (dead PID or unreadable) — remove and retry once.
if rmErr := os.Remove(path); rmErr != nil && !errors.Is(rmErr, fs.ErrNotExist) {
return nil, rmErr
}
}
return nil, ErrAlreadyRunning
}
// release removes the lock file. Safe to call once.
func (l *fileLock) release() error {
if l == nil || l.path == "" {
return nil
}
if err := os.Remove(l.path); err != nil && !errors.Is(err, fs.ErrNotExist) {
return err
}
return nil
}
// readPidFile reads and parses the PID recorded in a lock file.
func readPidFile(path string) (int, error) {
data, err := os.ReadFile(path)
if err != nil {
return 0, err
}
return strconv.Atoi(strings.TrimSpace(string(data)))
}
+22
View File
@@ -0,0 +1,22 @@
//go:build !windows
package daemon
import (
"errors"
"syscall"
)
// osProcessAlive reports whether pid is a live process on POSIX. Signal 0 does
// not deliver a signal; it only checks existence/permission. ESRCH means no such
// process (dead); EPERM means it exists but we may not signal it (alive).
func osProcessAlive(pid int) bool {
if pid <= 0 {
return false
}
err := syscall.Kill(pid, 0)
if err == nil {
return true
}
return errors.Is(err, syscall.EPERM)
}
+77
View File
@@ -0,0 +1,77 @@
package daemon
import (
"errors"
"os"
"path/filepath"
"strconv"
"strings"
"testing"
)
func TestLockSingleInstance(t *testing.T) {
path := filepath.Join(t.TempDir(), "daemon.lock")
alive := func(int) bool { return true }
l1, err := acquireLock(path, alive)
if err != nil {
t.Fatalf("first acquireLock: %v", err)
}
// Second acquire while the holder is "alive" must be refused.
if _, err := acquireLock(path, alive); !errors.Is(err, ErrAlreadyRunning) {
t.Fatalf("second acquireLock err = %v, want ErrAlreadyRunning", err)
}
// After release, a new acquire succeeds.
if err := l1.release(); err != nil {
t.Fatalf("release: %v", err)
}
l2, err := acquireLock(path, alive)
if err != nil {
t.Fatalf("acquire after release: %v", err)
}
_ = l2.release()
}
func TestLockStaleRecovery(t *testing.T) {
path := filepath.Join(t.TempDir(), "daemon.lock")
// Simulate a stale lock from a crashed daemon: a PID file whose process is
// dead.
if err := os.WriteFile(path, []byte("4242\n"), 0o600); err != nil {
t.Fatalf("write stale lock: %v", err)
}
dead := func(int) bool { return false }
l, err := acquireLock(path, dead)
if err != nil {
t.Fatalf("stale-lock recovery failed: %v", err)
}
// The lock now records OUR pid, not the stale one.
data, _ := os.ReadFile(path)
if strings.TrimSpace(string(data)) != strconv.Itoa(os.Getpid()) {
t.Fatalf("reclaimed lock pid = %q, want %d", strings.TrimSpace(string(data)), os.Getpid())
}
_ = l.release()
}
func TestLockReleaseRemovesFile(t *testing.T) {
path := filepath.Join(t.TempDir(), "daemon.lock")
l, err := acquireLock(path, func(int) bool { return true })
if err != nil {
t.Fatalf("acquireLock: %v", err)
}
if err := l.release(); err != nil {
t.Fatalf("release: %v", err)
}
if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("lock file still present after release: %v", err)
}
}
func TestProcessAliveSelfAndDead(t *testing.T) {
if !osProcessAlive(os.Getpid()) {
t.Fatal("osProcessAlive(self) = false, want true")
}
// PID 0 / negative are never live.
if osProcessAlive(0) || osProcessAlive(-1) {
t.Fatal("osProcessAlive must reject non-positive pids")
}
}
+27
View File
@@ -0,0 +1,27 @@
//go:build windows
package daemon
import "golang.org/x/sys/windows"
// osProcessAlive reports whether pid is a live process on Windows. It opens the
// process for limited query access; a successful open with a non-exited status
// means the PID is live. x/sys is already a module dependency.
func osProcessAlive(pid int) bool {
if pid <= 0 {
return false
}
handle, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION, false, uint32(pid))
if err != nil {
return false
}
defer windows.CloseHandle(handle)
var code uint32
if err := windows.GetExitCodeProcess(handle, &code); err != nil {
// Could not query; treat the open as proof of life (conservative: do not
// reclaim a lock we are unsure about).
return true
}
const stillActive = 259 // STILL_ACTIVE
return code == stillActive
}
+42
View File
@@ -0,0 +1,42 @@
package daemon
import (
"os"
"path/filepath"
"strings"
)
// DefaultDir returns the per-user directory for the daemon's runtime files:
// $XDG_RUNTIME_DIR/zero when set (a tmpfs owned by the user on Linux), otherwise
// ~/.zero. Mirrors supervisor.js / config.js choosing a per-user runtime dir.
func DefaultDir() (string, error) {
if rt := strings.TrimSpace(os.Getenv("XDG_RUNTIME_DIR")); rt != "" {
return filepath.Join(rt, "zero"), nil
}
home, err := os.UserHomeDir()
if err != nil {
return "", err
}
return filepath.Join(home, ".zero"), nil
}
// Paths bundles the daemon's socket, lock, and status file paths.
type Paths struct {
Socket string
Lock string
Status string
}
// DefaultPaths returns the daemon control-socket, lock, and status file paths
// under DefaultDir.
func DefaultPaths() (Paths, error) {
dir, err := DefaultDir()
if err != nil {
return Paths{}, err
}
return Paths{
Socket: filepath.Join(dir, "daemon.sock"),
Lock: filepath.Join(dir, "daemon.lock"),
Status: filepath.Join(dir, "daemon.status"),
}, nil
}
+355
View File
@@ -0,0 +1,355 @@
package daemon
import (
"context"
"errors"
"fmt"
"sync"
"time"
)
// Exit codes a worker may use to signal restart policy. Mirrors
// reference-daemon-code-agent-js/worker-manager.js EXIT_TEMPFAIL/EXIT_PERMANENT.
const (
// ExitTempfail asks for a retry after a fixed cool-off (transient failure).
ExitTempfail = 75
// ExitPermanent asks the daemon to STOP retrying this work (fatal).
ExitPermanent = 76
)
// defaultPoolSize, defaultMaxAttempts and defaultKillTimeout are used when a
// PoolOptions field is left zero.
const (
defaultPoolSize = 4
defaultMaxAttempts = 5
defaultKillTimeout = 5 * time.Second
)
// ErrPoolDraining is returned by Run once the pool is shutting down.
var ErrPoolDraining = errors.New("daemon: pool is draining")
// ErrPermanent is returned when a worker exits ExitPermanent or the attempt cap
// is exhausted; the request must not be retried.
var ErrPermanent = errors.New("daemon: worker failed permanently")
// WorkerSpec describes the worker to launch for one session attempt.
type WorkerSpec struct {
Session string
Cwd string
// Args are the per-session exec flags (e.g. --prompt ...), parsed by the CLI
// and appended after `exec` by the production launcher.
Args []string
}
// WorkerHandle is a launched worker process the pool supervises. The production
// launcher wraps an exec.Cmd running `zero exec -i/-o stream-json`; tests inject
// a fake. Stdout yields the worker's stream-json event lines.
type WorkerHandle interface {
Stdout() Lines
// Wait blocks until the worker exits and returns its process exit code.
Wait() (int, error)
// Kill force-terminates the worker (used on drain timeout).
Kill() error
Pid() int
}
// Lines is a stream of stream-json output lines from a worker. Next returns the
// next line and io-EOF-style ok=false when the stream ends.
type Lines interface {
Next() (line string, ok bool, err error)
}
// Launcher spawns a worker for a single attempt. It must apply the sandbox/env
// policy (the production launcher scrubs the re-entrancy markers so the worker
// re-establishes its own sandbox — see newExecLauncher).
type Launcher func(ctx context.Context, spec WorkerSpec) (WorkerHandle, error)
// PoolOptions configures a Pool. Zero fields take documented defaults.
type PoolOptions struct {
// Size is the max number of concurrent workers (lease slots). Sessions beyond
// this queue until a slot frees.
Size int
// Launcher spawns a worker. Required.
Launcher Launcher
// MaxAttempts caps the total attempts per request (first + retries). When
// exhausted, Run returns ErrPermanent (mirrors capping consecutive restarts).
MaxAttempts int
// KillTimeout bounds graceful termination before SIGKILL-equivalent on drain.
KillTimeout time.Duration
// Backoff returns the delay before retry attempt n (1-based crash count).
// Defaults to BackoffWithJitter; tests inject a deterministic function.
Backoff func(attempt int) time.Duration
// TempfailDelay is the cool-off after an ExitTempfail exit. Defaults to 30s.
TempfailDelay time.Duration
// Log, when set, receives one line per lifecycle event.
Log func(string)
}
// Pool is a bounded, self-healing worker pool. A crashed worker never takes down
// the pool: Run retries on a fresh worker with backoff up to MaxAttempts, and an
// ExitPermanent worker stops retries. Drain stops new work and kills stragglers.
type Pool struct {
opts PoolOptions
slots chan struct{}
mu sync.Mutex
draining bool
active map[int]WorkerHandle // pid -> handle, for drain/kill + status
nextID int
drainOnce sync.Once
drained chan struct{}
}
// workerStat tracks one in-flight request's restart count (local to Run).
type workerStat struct {
id int
restarts int
}
// NewPool builds a pool from opts, filling defaults.
func NewPool(opts PoolOptions) (*Pool, error) {
if opts.Launcher == nil {
return nil, errors.New("daemon: pool requires a Launcher")
}
if opts.Size <= 0 {
opts.Size = defaultPoolSize
}
if opts.MaxAttempts <= 0 {
opts.MaxAttempts = defaultMaxAttempts
}
if opts.KillTimeout <= 0 {
opts.KillTimeout = defaultKillTimeout
}
if opts.TempfailDelay <= 0 {
opts.TempfailDelay = 30 * time.Second
}
if opts.Backoff == nil {
opts.Backoff = BackoffWithJitter
}
return &Pool{
opts: opts,
slots: make(chan struct{}, opts.Size),
active: map[int]WorkerHandle{},
drained: make(chan struct{}),
}, nil
}
// Size returns the configured max concurrency.
func (p *Pool) Size() int { return p.opts.Size }
func (p *Pool) logf(format string, args ...any) {
if p.opts.Log != nil {
p.opts.Log(fmt.Sprintf(format, args...))
}
}
// acquire takes a lease slot, queuing until one frees or ctx is cancelled / the
// pool drains.
func (p *Pool) acquire(ctx context.Context) error {
p.mu.Lock()
draining := p.draining
p.mu.Unlock()
if draining {
return ErrPoolDraining
}
select {
case p.slots <- struct{}{}:
return nil
case <-ctx.Done():
return ctx.Err()
case <-p.drained:
return ErrPoolDraining
}
}
func (p *Pool) release() { <-p.slots }
// QueueDepth reports how many lease slots are currently in use.
func (p *Pool) QueueDepth() int { return len(p.slots) }
// Sink receives a worker's stream-json output lines for one request.
type Sink interface {
Line(line string)
}
// Run leases a worker slot and dispatches spec to a worker, streaming its
// stream-json lines to sink. It is at-least-once with bounded retries: a worker
// that crashes (non-zero, non-permanent) is retried on a fresh worker after a
// backoff; ExitPermanent or exhausting MaxAttempts returns ErrPermanent. Run
// queues when all slots are busy. The returned int is the final worker exit code.
func (p *Pool) Run(ctx context.Context, spec WorkerSpec, sink Sink) (int, error) {
if err := p.acquire(ctx); err != nil {
return 0, err
}
defer p.release()
// Lease hook: a sink that tracks queue->running state is notified the instant
// a worker slot is leased (mirrors lease.js acquiring a worker for a session).
if s, ok := sink.(interface{ Started() }); ok {
s.Started()
}
stat := p.newStat()
var lastErr error
for attempt := 1; attempt <= p.opts.MaxAttempts; attempt++ {
if p.isDraining() {
return 0, ErrPoolDraining
}
code, err := p.runOnce(ctx, stat.id, spec, sink)
switch {
case err != nil:
lastErr = err
if ctx.Err() != nil {
return 0, ctx.Err()
}
p.logf("worker %d launch/run error: %v", stat.id, err)
case code == 0:
return 0, nil // clean success
case code == ExitPermanent:
p.logf("worker %d exited permanently (code=%d) — not retrying", stat.id, code)
return code, ErrPermanent
case code == ExitTempfail:
lastErr = fmt.Errorf("worker %d tempfail (code=%d)", stat.id, code)
p.logf("worker %d tempfail — retry after %s", stat.id, p.opts.TempfailDelay)
if !p.sleep(ctx, p.opts.TempfailDelay) {
return 0, ctx.Err()
}
continue // tempfail retries do not count against the crash backoff
default:
lastErr = fmt.Errorf("worker %d exited code=%d", stat.id, code)
}
if attempt == p.opts.MaxAttempts {
break
}
stat.restarts++
delay := p.opts.Backoff(stat.restarts)
p.logf("worker %d restart %d after backoff %s", stat.id, stat.restarts, delay)
if !p.sleep(ctx, delay) {
return 0, ctx.Err()
}
}
if lastErr == nil {
lastErr = ErrPermanent
}
return ExitPermanent, fmt.Errorf("%w: %v", ErrPermanent, lastErr)
}
// runOnce launches a single worker, pumps its output to sink, and returns its
// exit code. The worker handle is tracked so Drain can kill it.
func (p *Pool) runOnce(ctx context.Context, id int, spec WorkerSpec, sink Sink) (int, error) {
handle, err := p.opts.Launcher(ctx, spec)
if err != nil {
return 0, err
}
p.track(handle)
defer p.untrack(handle)
// Pump stdout lines until the stream ends.
lines := handle.Stdout()
for {
line, ok, lerr := lines.Next()
if lerr != nil {
break
}
if !ok {
break
}
if sink != nil {
sink.Line(line)
}
}
return handle.Wait()
}
func (p *Pool) newStat() *workerStat {
p.mu.Lock()
defer p.mu.Unlock()
p.nextID++
return &workerStat{id: p.nextID}
}
func (p *Pool) track(h WorkerHandle) {
p.mu.Lock()
p.active[h.Pid()] = h
p.mu.Unlock()
}
func (p *Pool) untrack(h WorkerHandle) {
p.mu.Lock()
delete(p.active, h.Pid())
p.mu.Unlock()
}
func (p *Pool) isDraining() bool {
p.mu.Lock()
defer p.mu.Unlock()
return p.draining
}
// sleep waits d, returning false if ctx is cancelled first. A non-positive d
// returns immediately.
func (p *Pool) sleep(ctx context.Context, d time.Duration) bool {
if d <= 0 {
return ctx.Err() == nil
}
timer := time.NewTimer(d)
defer timer.Stop()
select {
case <-timer.C:
return true
case <-ctx.Done():
return false
}
}
// Drain stops accepting new work, gives in-flight workers a grace window
// (KillTimeout) to finish on their own, then force-kills any straggler. It
// returns as soon as the pool is idle (graceful) or the window elapses. Safe to
// call once; subsequent calls are no-ops.
func (p *Pool) Drain() {
p.drainOnce.Do(func() {
p.mu.Lock()
p.draining = true
p.mu.Unlock()
close(p.drained)
// Grace window: poll until idle or the deadline.
deadline := time.Now().Add(p.opts.KillTimeout)
for time.Now().Before(deadline) {
p.mu.Lock()
n := len(p.active)
p.mu.Unlock()
if n == 0 {
return // all workers drained gracefully
}
time.Sleep(5 * time.Millisecond)
}
// Timed out: force-kill remaining workers.
p.mu.Lock()
handles := make([]WorkerHandle, 0, len(p.active))
for _, h := range p.active {
handles = append(handles, h)
}
p.mu.Unlock()
for _, h := range handles {
p.logf("drain: killing straggler worker pid=%d", h.Pid())
_ = h.Kill()
}
})
}
// WorkerStats returns a snapshot of the currently-active (busy) workers, bounded
// by the pool size. The on-demand pool keeps no idle workers, so an empty result
// means the pool is idle.
func (p *Pool) WorkerStats() []WorkerStatus {
p.mu.Lock()
defer p.mu.Unlock()
out := make([]WorkerStatus, 0, len(p.active))
for pid := range p.active {
out = append(out, WorkerStatus{PID: pid, State: "busy"})
}
return out
}
+248
View File
@@ -0,0 +1,248 @@
package daemon
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
)
// --- test doubles ---------------------------------------------------------
type fakeLines struct {
lines []string
i int
}
func (f *fakeLines) Next() (string, bool, error) {
if f.i >= len(f.lines) {
return "", false, nil
}
l := f.lines[f.i]
f.i++
return l, true, nil
}
type fakeWorker struct {
pid int
out []string
exitCode int
killed int32
waitCh chan struct{} // when non-nil, Wait blocks until closed (drain tests)
}
func (w *fakeWorker) Stdout() Lines { return &fakeLines{lines: w.out} }
func (w *fakeWorker) Wait() (int, error) {
if w.waitCh != nil {
<-w.waitCh
}
return w.exitCode, nil
}
func (w *fakeWorker) Kill() error {
atomic.StoreInt32(&w.killed, 1)
if w.waitCh != nil {
select {
case <-w.waitCh:
default:
close(w.waitCh)
}
}
return nil
}
func (w *fakeWorker) Pid() int { return w.pid }
// seqLauncher hands out the given workers in order, then errors.
func seqLauncher(workers ...*fakeWorker) (Launcher, *int32) {
var calls int32
l := func(_ context.Context, _ WorkerSpec) (WorkerHandle, error) {
n := atomic.AddInt32(&calls, 1)
idx := int(n) - 1
if idx >= len(workers) {
return nil, errors.New("no more fake workers")
}
return workers[idx], nil
}
return l, &calls
}
type collectSink struct {
mu sync.Mutex
lines []string
}
func (c *collectSink) Line(line string) {
c.mu.Lock()
c.lines = append(c.lines, line)
c.mu.Unlock()
}
// --- tests ----------------------------------------------------------------
func TestPoolRunStreamsAndSucceeds(t *testing.T) {
launcher, calls := seqLauncher(&fakeWorker{pid: 1, out: []string{"a", "b"}, exitCode: 0})
pool, err := NewPool(PoolOptions{Size: 2, Launcher: launcher})
if err != nil {
t.Fatalf("NewPool: %v", err)
}
sink := &collectSink{}
code, err := pool.Run(context.Background(), WorkerSpec{Session: "s"}, sink)
if err != nil || code != 0 {
t.Fatalf("Run = (%d,%v), want (0,nil)", code, err)
}
if len(sink.lines) != 2 || sink.lines[0] != "a" || sink.lines[1] != "b" {
t.Fatalf("sink lines = %v, want [a b]", sink.lines)
}
if *calls != 1 {
t.Fatalf("launcher calls = %d, want 1", *calls)
}
}
func TestPoolRunRetriesWithBackoff(t *testing.T) {
launcher, calls := seqLauncher(
&fakeWorker{pid: 1, exitCode: 1}, // crash
&fakeWorker{pid: 2, exitCode: 1}, // crash
&fakeWorker{pid: 3, exitCode: 0, out: []string{"ok"}}, // success
)
var backoffArgs []int
pool, _ := NewPool(PoolOptions{
Size: 1, Launcher: launcher, MaxAttempts: 5,
Backoff: func(attempt int) time.Duration { backoffArgs = append(backoffArgs, attempt); return 0 },
})
sink := &collectSink{}
code, err := pool.Run(context.Background(), WorkerSpec{Session: "s"}, sink)
if err != nil || code != 0 {
t.Fatalf("Run = (%d,%v), want (0,nil)", code, err)
}
if *calls != 3 {
t.Fatalf("launcher calls = %d, want 3 (2 crashes + success)", *calls)
}
// Backoff invoked between attempts with an increasing 1-based restart count.
if len(backoffArgs) != 2 || backoffArgs[0] != 1 || backoffArgs[1] != 2 {
t.Fatalf("backoff attempts = %v, want [1 2]", backoffArgs)
}
}
func TestPoolRunPermanentStopsImmediately(t *testing.T) {
launcher, calls := seqLauncher(&fakeWorker{pid: 1, exitCode: ExitPermanent})
pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher, Backoff: func(int) time.Duration { return 0 }})
code, err := pool.Run(context.Background(), WorkerSpec{Session: "s"}, &collectSink{})
if !errors.Is(err, ErrPermanent) {
t.Fatalf("Run err = %v, want ErrPermanent", err)
}
if code != ExitPermanent {
t.Fatalf("Run code = %d, want %d", code, ExitPermanent)
}
if *calls != 1 {
t.Fatalf("permanent exit must not retry; launcher calls = %d, want 1", *calls)
}
}
func TestPoolRunCapExhausted(t *testing.T) {
launcher, calls := seqLauncher(
&fakeWorker{pid: 1, exitCode: 1},
&fakeWorker{pid: 2, exitCode: 1},
&fakeWorker{pid: 3, exitCode: 1},
)
pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher, MaxAttempts: 3, Backoff: func(int) time.Duration { return 0 }})
_, err := pool.Run(context.Background(), WorkerSpec{Session: "s"}, &collectSink{})
if !errors.Is(err, ErrPermanent) {
t.Fatalf("Run err = %v, want ErrPermanent after cap", err)
}
if *calls != 3 {
t.Fatalf("launcher calls = %d, want 3 (== MaxAttempts)", *calls)
}
}
func TestPoolRunTempfailRetries(t *testing.T) {
launcher, calls := seqLauncher(
&fakeWorker{pid: 1, exitCode: ExitTempfail},
&fakeWorker{pid: 2, exitCode: 0, out: []string{"ok"}},
)
pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher, TempfailDelay: time.Millisecond})
code, err := pool.Run(context.Background(), WorkerSpec{Session: "s"}, &collectSink{})
if err != nil || code != 0 {
t.Fatalf("Run = (%d,%v), want (0,nil)", code, err)
}
if *calls != 2 {
t.Fatalf("launcher calls = %d, want 2 (tempfail + success)", *calls)
}
}
func TestPoolQueuesWhenFull(t *testing.T) {
block := make(chan struct{})
first := &fakeWorker{pid: 1, waitCh: block}
second := &fakeWorker{pid: 2, exitCode: 0}
launcher, _ := seqLauncher(first, second)
pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})
started := make(chan struct{})
go func() {
close(started)
_, _ = pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{})
}()
<-started
// Wait until the first run holds the only slot.
waitFor(t, func() bool { return pool.QueueDepth() == 1 })
secondDone := make(chan struct{})
go func() {
_, _ = pool.Run(context.Background(), WorkerSpec{Session: "b"}, &collectSink{})
close(secondDone)
}()
// The second run must NOT complete while the slot is held.
select {
case <-secondDone:
t.Fatal("second run completed while the only slot was busy (no queueing)")
case <-time.After(50 * time.Millisecond):
}
// Free the slot; the queued run now proceeds.
close(block)
select {
case <-secondDone:
case <-time.After(2 * time.Second):
t.Fatal("queued run did not proceed after slot freed")
}
}
func TestPoolDrainKillsStraggler(t *testing.T) {
block := make(chan struct{})
straggler := &fakeWorker{pid: 1, waitCh: block}
launcher, _ := seqLauncher(straggler)
pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher, MaxAttempts: 1, KillTimeout: 10 * time.Millisecond})
runDone := make(chan struct{})
go func() {
_, _ = pool.Run(context.Background(), WorkerSpec{Session: "a"}, &collectSink{})
close(runDone)
}()
waitFor(t, func() bool { return pool.QueueDepth() == 1 })
pool.Drain() // KillTimeout elapses, straggler is force-killed
if atomic.LoadInt32(&straggler.killed) != 1 {
t.Fatal("Drain must kill a straggler still running after KillTimeout")
}
select {
case <-runDone:
case <-time.After(2 * time.Second):
t.Fatal("run did not finish after drain kill")
}
// A new Run after drain is refused.
if _, err := pool.Run(context.Background(), WorkerSpec{Session: "b"}, &collectSink{}); !errors.Is(err, ErrPoolDraining) {
t.Fatalf("Run after drain err = %v, want ErrPoolDraining", err)
}
}
func waitFor(t *testing.T, cond func() bool) {
t.Helper()
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
if cond() {
return
}
time.Sleep(time.Millisecond)
}
t.Fatal("condition not met within timeout")
}
+184
View File
@@ -0,0 +1,184 @@
// Package daemon implements zero's long-running daemon/server mode: a local
// control server that supervises a pool of headless `zero exec` worker processes
// and routes multiple agent sessions to them over an owner-only local socket.
//
// It reuses zero's existing building blocks rather than inventing new ones:
// - internal/background : child-process group setup + cross-platform terminate.
// - internal/streamjson : the line-based agent event protocol on worker stdio.
// - internal/cli (exec) : a worker is a `zero exec -i/-o stream-json` process.
//
// The control plane (this file) is a small framed codec mirrored from the
// reference daemon's protocol.js: a 4-byte big-endian length prefix, a 1-byte
// frame kind, then a JSON control payload, with a 1 MiB frame cap and a
// version-negotiation handshake. The agent event stream itself stays stream-json.
package daemon
import (
"encoding/binary"
"encoding/json"
"errors"
"fmt"
"io"
)
// Wire constants. Mirrors reference-daemon-code-agent-js/protocol.js
// (FRAME_DATA/FRAME_CTRL, HEADER_SIZE, MAX_FRAME_SIZE, PROTO_VERSION).
const (
// ProtoVersion is the control-protocol version advertised in the handshake.
ProtoVersion = 1
// MaxFrameSize caps a single control frame payload at 1 MiB. A larger
// advertised length is rejected before any allocation (fail closed).
MaxFrameSize = 1 << 20
// frameHeaderSize is 4 bytes big-endian length + 1 byte kind.
frameHeaderSize = 5
)
// FrameKind tags a frame's payload. Mirrors protocol.js FRAME_DATA/FRAME_CTRL.
type FrameKind uint8
const (
// KindData carries an opaque payload (e.g. a stream-json line). Reserved for
// future use; the control plane uses KindCtrl.
KindData FrameKind = 0
// KindCtrl carries a JSON-encoded control message (see Ctrl).
KindCtrl FrameKind = 1
)
// ErrFrameTooLarge is returned when a frame's advertised payload length exceeds
// MaxFrameSize. The reader rejects it without allocating the buffer.
var ErrFrameTooLarge = errors.New("daemon: control frame exceeds size cap")
// ErrUnknownFrameKind is returned for a frame whose kind byte is not recognized.
var ErrUnknownFrameKind = errors.New("daemon: unknown control frame kind")
// WriteFrame writes a single framed message: [uint32BE len][uint8 kind][payload].
// It rejects an over-cap payload before writing anything so a buggy/hostile caller
// cannot emit a frame a conforming reader would refuse.
func WriteFrame(w io.Writer, kind FrameKind, payload []byte) error {
if len(payload) > MaxFrameSize {
return fmt.Errorf("%w (%d > %d)", ErrFrameTooLarge, len(payload), MaxFrameSize)
}
header := make([]byte, frameHeaderSize)
binary.BigEndian.PutUint32(header[:4], uint32(len(payload)))
header[4] = byte(kind)
if _, err := w.Write(header); err != nil {
return err
}
if len(payload) == 0 {
return nil
}
_, err := w.Write(payload)
return err
}
// ReadFrame reads a single framed message. It validates the advertised length
// against MaxFrameSize BEFORE allocating, so an oversize or hostile frame is
// rejected (ErrFrameTooLarge) rather than triggering a huge allocation. An
// unrecognized kind byte yields ErrUnknownFrameKind after the payload is drained
// so the stream stays in sync.
func ReadFrame(r io.Reader) (FrameKind, []byte, error) {
header := make([]byte, frameHeaderSize)
if _, err := io.ReadFull(r, header); err != nil {
return 0, nil, err
}
length := binary.BigEndian.Uint32(header[:4])
if length > MaxFrameSize {
return 0, nil, fmt.Errorf("%w (%d > %d)", ErrFrameTooLarge, length, MaxFrameSize)
}
kind := FrameKind(header[4])
payload := make([]byte, length)
if length > 0 {
if _, err := io.ReadFull(r, payload); err != nil {
return 0, nil, err
}
}
if kind != KindData && kind != KindCtrl {
return kind, nil, ErrUnknownFrameKind
}
return kind, payload, nil
}
// CtrlType is the discriminator for a control message.
type CtrlType string
const (
// CtrlHello is the client's opening handshake; Version is the highest control
// version it speaks.
CtrlHello CtrlType = "hello"
// CtrlHelloOK is the daemon's handshake reply with the negotiated Version.
CtrlHelloOK CtrlType = "hello_ok"
// CtrlRun asks the daemon to create/route a session and stream its output.
CtrlRun CtrlType = "run"
// CtrlAttach asks the daemon to attach the client to a running session's stream.
CtrlAttach CtrlType = "attach"
// CtrlStatus requests daemon/worker/session status.
CtrlStatus CtrlType = "status"
// CtrlStatusResult carries the status payload (see StatusReport in Status).
CtrlStatusResult CtrlType = "status_result"
// CtrlData carries one stream-json line of agent output back to the client.
CtrlData CtrlType = "data"
// CtrlAck acknowledges receipt/dispatch of a request (at-least-once dispatch).
CtrlAck CtrlType = "ack"
// CtrlEnd marks the end of a session's stream.
CtrlEnd CtrlType = "end"
// CtrlShutdown asks the daemon to drain and stop.
CtrlShutdown CtrlType = "shutdown"
// CtrlError carries a human-readable error for the client.
CtrlError CtrlType = "error"
)
// Ctrl is a control message carried in a KindCtrl frame as JSON. Unused fields
// are omitted so frames stay small.
type Ctrl struct {
Type CtrlType `json:"type"`
Version int `json:"version,omitempty"`
Session string `json:"session,omitempty"`
Cwd string `json:"cwd,omitempty"`
Args []string `json:"args,omitempty"`
Prompt string `json:"prompt,omitempty"`
// Line is one stream-json event line (for CtrlData).
Line string `json:"line,omitempty"`
// Message is a human-readable detail (for CtrlError / CtrlAck).
Message string `json:"message,omitempty"`
// Status is populated on CtrlStatusResult.
Status *StatusReport `json:"status,omitempty"`
}
// WriteControl marshals a control message into a KindCtrl frame.
func WriteControl(w io.Writer, msg Ctrl) error {
payload, err := json.Marshal(msg)
if err != nil {
return err
}
return WriteFrame(w, KindCtrl, payload)
}
// ReadControl reads the next frame and decodes it as a control message. A
// non-control frame yields ErrUnknownFrameKind so the caller fails closed.
func ReadControl(r io.Reader) (Ctrl, error) {
kind, payload, err := ReadFrame(r)
if err != nil {
return Ctrl{}, err
}
if kind != KindCtrl {
return Ctrl{}, ErrUnknownFrameKind
}
var msg Ctrl
if err := json.Unmarshal(payload, &msg); err != nil {
return Ctrl{}, fmt.Errorf("daemon: decode control message: %w", err)
}
return msg, nil
}
// NegotiateVersion picks the protocol version both ends can speak: the lower of
// the client's advertised version and ProtoVersion. A non-positive client
// version is invalid and yields ok=false so the daemon can reject the connection.
func NegotiateVersion(clientVersion int) (int, bool) {
if clientVersion <= 0 {
return 0, false
}
if clientVersion < ProtoVersion {
return clientVersion, true
}
return ProtoVersion, true
}
+120
View File
@@ -0,0 +1,120 @@
package daemon
import (
"bytes"
"encoding/binary"
"errors"
"io"
"strings"
"testing"
)
func TestFrameRoundTrip(t *testing.T) {
var buf bytes.Buffer
payloads := [][]byte{[]byte("hello"), {}, []byte(strings.Repeat("x", 4096))}
for _, p := range payloads {
if err := WriteFrame(&buf, KindCtrl, p); err != nil {
t.Fatalf("WriteFrame: %v", err)
}
}
for i, want := range payloads {
kind, got, err := ReadFrame(&buf)
if err != nil {
t.Fatalf("ReadFrame[%d]: %v", i, err)
}
if kind != KindCtrl {
t.Fatalf("frame[%d] kind = %d, want KindCtrl", i, kind)
}
if !bytes.Equal(got, want) {
t.Fatalf("frame[%d] payload = %q, want %q", i, got, want)
}
}
}
func TestControlRoundTrip(t *testing.T) {
var buf bytes.Buffer
msg := Ctrl{Type: CtrlRun, Session: "s1", Cwd: "/tmp", Args: []string{"--prompt", "hi"}}
if err := WriteControl(&buf, msg); err != nil {
t.Fatalf("WriteControl: %v", err)
}
got, err := ReadControl(&buf)
if err != nil {
t.Fatalf("ReadControl: %v", err)
}
if got.Type != CtrlRun || got.Session != "s1" || got.Cwd != "/tmp" || len(got.Args) != 2 {
t.Fatalf("round-trip mismatch: %+v", got)
}
}
func TestWriteFrameRejectsOversize(t *testing.T) {
var buf bytes.Buffer
err := WriteFrame(&buf, KindCtrl, make([]byte, MaxFrameSize+1))
if !errors.Is(err, ErrFrameTooLarge) {
t.Fatalf("WriteFrame oversize err = %v, want ErrFrameTooLarge", err)
}
if buf.Len() != 0 {
t.Fatalf("oversize WriteFrame must not emit any bytes, wrote %d", buf.Len())
}
}
func TestReadFrameRejectsOversizeBeforeAlloc(t *testing.T) {
// A header advertising a 4 GiB payload must be rejected without reading or
// allocating that payload.
header := make([]byte, frameHeaderSize)
binary.BigEndian.PutUint32(header[:4], MaxFrameSize+1)
header[4] = byte(KindCtrl)
_, _, err := ReadFrame(bytes.NewReader(header))
if !errors.Is(err, ErrFrameTooLarge) {
t.Fatalf("ReadFrame oversize err = %v, want ErrFrameTooLarge", err)
}
}
func TestReadFrameUnknownKind(t *testing.T) {
var buf bytes.Buffer
// Hand-write a frame with an unknown kind byte (7) and a small payload.
header := make([]byte, frameHeaderSize)
binary.BigEndian.PutUint32(header[:4], 2)
header[4] = 7
buf.Write(header)
buf.Write([]byte("ab"))
_, _, err := ReadFrame(&buf)
if !errors.Is(err, ErrUnknownFrameKind) {
t.Fatalf("unknown kind err = %v, want ErrUnknownFrameKind", err)
}
}
func TestReadFrameTruncated(t *testing.T) {
// Header promises 10 bytes but only 3 follow → unexpected EOF.
var buf bytes.Buffer
header := make([]byte, frameHeaderSize)
binary.BigEndian.PutUint32(header[:4], 10)
header[4] = byte(KindCtrl)
buf.Write(header)
buf.Write([]byte("abc"))
_, _, err := ReadFrame(&buf)
if err == nil || errors.Is(err, ErrFrameTooLarge) {
t.Fatalf("truncated frame err = %v, want a read error", err)
}
if !errors.Is(err, io.ErrUnexpectedEOF) {
t.Fatalf("truncated frame err = %v, want io.ErrUnexpectedEOF", err)
}
}
func TestNegotiateVersion(t *testing.T) {
cases := []struct {
client int
want int
wantOK bool
}{
{client: 1, want: 1, wantOK: true},
{client: 99, want: ProtoVersion, wantOK: true}, // clamp down to ours
{client: 0, want: 0, wantOK: false},
{client: -3, want: 0, wantOK: false},
}
for _, tc := range cases {
got, ok := NegotiateVersion(tc.client)
if got != tc.want || ok != tc.wantOK {
t.Fatalf("NegotiateVersion(%d) = (%d,%v), want (%d,%v)", tc.client, got, ok, tc.want, tc.wantOK)
}
}
}
+294
View File
@@ -0,0 +1,294 @@
package daemon
import (
"context"
"encoding/json"
"errors"
"fmt"
"net"
"os"
"sync"
"time"
)
// Server is the daemon control plane. Mirrors reference-daemon-code-agent-js/
// supervisor.js (single-instance lock, status file, lifecycle) + the accept loop
// that routes framed control requests to the SessionManager/Pool. It listens on
// an owner-only local Unix socket and NEVER binds a TCP port.
type Server struct {
opts ServerOptions
startedAt time.Time
listener net.Listener
lock *fileLock
ctx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
shutdownOnce sync.Once
done chan struct{}
}
// ServerOptions configures a Server.
type ServerOptions struct {
Paths Paths
Manager *SessionManager
Pool *Pool
Version int
Now func() time.Time
Log func(string)
isAlive func(int) bool // test hook for the single-instance lock
}
// NewServer validates options and builds a Server.
func NewServer(opts ServerOptions) (*Server, error) {
if opts.Manager == nil || opts.Pool == nil {
return nil, errors.New("daemon: server requires a Pool and SessionManager")
}
if opts.Paths.Socket == "" || opts.Paths.Lock == "" || opts.Paths.Status == "" {
return nil, errors.New("daemon: server requires socket, lock, and status paths")
}
if opts.Version <= 0 {
opts.Version = ProtoVersion
}
if opts.Now == nil {
opts.Now = time.Now
}
ctx, cancel := context.WithCancel(context.Background())
return &Server{
opts: opts,
ctx: ctx,
cancel: cancel,
done: make(chan struct{}),
}, nil
}
func (s *Server) logf(format string, args ...any) {
if s.opts.Log != nil {
s.opts.Log(fmt.Sprintf(format, args...))
}
}
// Serve acquires the single-instance lock, binds the owner-only control socket,
// writes the status file, and serves connections until Shutdown. It blocks. On
// return it has released the lock and removed the socket/status files.
func (s *Server) Serve() error {
if err := checkSocketPathLength(s.opts.Paths.Socket); err != nil {
return err
}
if err := secureSocketParent(s.opts.Paths.Socket); err != nil {
return err
}
lock, err := acquireLock(s.opts.Paths.Lock, s.opts.isAlive)
if err != nil {
return err
}
s.lock = lock
defer s.cleanup()
// A leftover socket file from an unclean exit would make Listen fail with
// "address already in use"; we hold the lock, so any socket here is stale.
_ = os.Remove(s.opts.Paths.Socket)
listener, err := net.Listen("unix", s.opts.Paths.Socket)
if err != nil {
return fmt.Errorf("daemon: bind control socket: %w", err)
}
s.listener = listener
if err := hardenSocketFile(s.opts.Paths.Socket); err != nil {
return fmt.Errorf("daemon: harden control socket: %w", err)
}
s.startedAt = s.opts.Now()
if err := s.writeStatusFile(); err != nil {
return err
}
s.logf("daemon listening on %s (pid %d)", s.opts.Paths.Socket, os.Getpid())
for {
conn, err := listener.Accept()
if err != nil {
select {
case <-s.done:
s.wg.Wait()
return nil // clean shutdown
default:
// Transient accept error during normal operation.
s.logf("accept error: %v", err)
return fmt.Errorf("daemon: accept: %w", err)
}
}
s.wg.Add(1)
go func() {
defer s.wg.Done()
s.handleConn(conn)
}()
}
}
// Shutdown stops accepting connections, cancels in-flight runs, drains the pool,
// and removes the socket/lock/status files. Safe to call multiple times.
func (s *Server) Shutdown() {
s.shutdownOnce.Do(func() {
close(s.done)
s.cancel() // stop in-flight pool runs
if s.listener != nil {
_ = s.listener.Close()
}
s.opts.Pool.Drain()
})
}
func (s *Server) cleanup() {
if s.listener != nil {
_ = s.listener.Close()
}
_ = os.Remove(s.opts.Paths.Socket)
_ = os.Remove(s.opts.Paths.Status)
if s.lock != nil {
_ = s.lock.release()
}
}
func (s *Server) writeStatusFile() error {
status := StatusFile{
PID: os.Getpid(),
Socket: s.opts.Paths.Socket,
Version: s.opts.Version,
StartedAt: s.startedAt,
}
data, err := json.MarshalIndent(status, "", " ")
if err != nil {
return err
}
if err := os.WriteFile(s.opts.Paths.Status, data, 0o600); err != nil {
return fmt.Errorf("daemon: write status file: %w", err)
}
return nil
}
// handleConn performs the handshake then dispatches a single control command.
func (s *Server) handleConn(conn net.Conn) {
defer conn.Close()
hello, err := ReadControl(conn)
if err != nil {
return
}
if hello.Type != CtrlHello {
_ = WriteControl(conn, Ctrl{Type: CtrlError, Message: "expected hello"})
return
}
version, ok := NegotiateVersion(hello.Version)
if !ok {
_ = WriteControl(conn, Ctrl{Type: CtrlError, Message: "unsupported protocol version"})
return
}
if err := WriteControl(conn, Ctrl{Type: CtrlHelloOK, Version: version}); err != nil {
return
}
cmd, err := ReadControl(conn)
if err != nil {
return
}
switch cmd.Type {
case CtrlRun:
s.handleRun(conn, cmd)
case CtrlAttach:
s.handleAttach(conn, cmd)
case CtrlStatus:
s.handleStatus(conn)
case CtrlShutdown:
_ = WriteControl(conn, Ctrl{Type: CtrlAck, Message: "shutting down"})
s.Shutdown()
default:
_ = WriteControl(conn, Ctrl{Type: CtrlError, Message: fmt.Sprintf("unknown command %q", cmd.Type)})
}
}
func (s *Server) handleRun(conn net.Conn, cmd Ctrl) {
if cmd.Session == "" {
_ = WriteControl(conn, Ctrl{Type: CtrlError, Message: "run requires a session id"})
return
}
args := cmd.Args
if cmd.Prompt != "" {
args = append(args, "--prompt", cmd.Prompt)
}
sess, err := s.opts.Manager.Start(s.ctx, WorkerSpec{Session: cmd.Session, Cwd: cmd.Cwd, Args: args})
if err != nil {
_ = WriteControl(conn, Ctrl{Type: CtrlError, Message: err.Error()})
return
}
_ = WriteControl(conn, Ctrl{Type: CtrlAck, Session: sess.ID()})
buffered, live, cancel := sess.Subscribe()
defer cancel()
s.streamToClient(conn, sess, buffered, live)
}
func (s *Server) handleAttach(conn net.Conn, cmd Ctrl) {
if cmd.Session == "" {
_ = WriteControl(conn, Ctrl{Type: CtrlError, Message: "attach requires a session id"})
return
}
buffered, live, cancel, err := s.opts.Manager.Attach(cmd.Session)
if err != nil {
_ = WriteControl(conn, Ctrl{Type: CtrlError, Message: err.Error()})
return
}
defer cancel()
sess, _ := s.opts.Manager.Get(cmd.Session)
_ = WriteControl(conn, Ctrl{Type: CtrlAck, Session: cmd.Session})
s.streamToClient(conn, sess, buffered, live)
}
// streamToClient writes the buffered history then live lines as CtrlData frames,
// finishing with CtrlEnd (or CtrlError if the session failed). A write error
// (client disconnected) ends the stream without affecting the session.
func (s *Server) streamToClient(conn net.Conn, sess *Session, buffered []string, live <-chan string) {
for _, line := range buffered {
if err := WriteControl(conn, Ctrl{Type: CtrlData, Line: line}); err != nil {
return
}
}
for {
select {
case line, ok := <-live:
if !ok {
s.finishStream(conn, sess)
return
}
if err := WriteControl(conn, Ctrl{Type: CtrlData, Line: line}); err != nil {
return
}
case <-s.done:
_ = WriteControl(conn, Ctrl{Type: CtrlEnd, Message: "daemon shutting down"})
return
}
}
}
func (s *Server) finishStream(conn net.Conn, sess *Session) {
if sess != nil {
if err := sess.Err(); err != nil {
_ = WriteControl(conn, Ctrl{Type: CtrlError, Session: sess.ID(), Message: err.Error()})
return
}
}
_ = WriteControl(conn, Ctrl{Type: CtrlEnd})
}
func (s *Server) handleStatus(conn net.Conn) {
report := StatusReport{
PID: os.Getpid(),
Version: s.opts.Version,
Socket: s.opts.Paths.Socket,
StartedAt: s.startedAt,
PoolSize: s.opts.Pool.Size(),
Workers: s.opts.Pool.WorkerStats(),
Sessions: s.opts.Manager.Statuses(),
QueueDepth: s.opts.Pool.QueueDepth(),
}
_ = WriteControl(conn, Ctrl{Type: CtrlStatusResult, Status: &report})
}
+176
View File
@@ -0,0 +1,176 @@
package daemon
import (
"errors"
"os"
"path/filepath"
"testing"
"time"
)
func newTestServer(t *testing.T, launcher Launcher) (*Server, Paths) {
t.Helper()
dir := t.TempDir()
paths := Paths{
Socket: filepath.Join(dir, "d.sock"),
Lock: filepath.Join(dir, "d.lock"),
Status: filepath.Join(dir, "d.status"),
}
pool, err := NewPool(PoolOptions{Size: 2, Launcher: launcher, KillTimeout: 200 * time.Millisecond})
if err != nil {
t.Fatalf("NewPool: %v", err)
}
mgr, err := NewSessionManager(SessionManagerOptions{Pool: pool})
if err != nil {
t.Fatalf("NewSessionManager: %v", err)
}
srv, err := NewServer(ServerOptions{Paths: paths, Manager: mgr, Pool: pool})
if err != nil {
t.Fatalf("NewServer: %v", err)
}
return srv, paths
}
func waitForFile(t *testing.T, path string) {
t.Helper()
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
if _, err := os.Stat(path); err == nil {
return
}
time.Sleep(2 * time.Millisecond)
}
t.Fatalf("file %s did not appear within timeout", path)
}
func TestServerEndToEnd(t *testing.T) {
out := []string{`{"type":"event","seq":1}`, `{"type":"event","seq":2}`}
launcher, _ := seqLauncher(&fakeWorker{pid: 1, out: out, exitCode: 0})
srv, paths := newTestServer(t, launcher)
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve() }()
waitForFile(t, paths.Status)
// --- run a session and collect its stream-json output ---
runClient, err := Dial(paths.Socket)
if err != nil {
t.Fatalf("Dial(run): %v", err)
}
var got []string
if err := runClient.Run("s1", "", "hello", nil, func(line string) { got = append(got, line) }); err != nil {
t.Fatalf("Run: %v", err)
}
runClient.Close()
if len(got) != 2 || got[0] != out[0] || got[1] != out[1] {
t.Fatalf("streamed lines = %v, want %v", got, out)
}
// --- status reflects the session ---
statusClient, err := Dial(paths.Socket)
if err != nil {
t.Fatalf("Dial(status): %v", err)
}
report, err := statusClient.Status()
statusClient.Close()
if err != nil {
t.Fatalf("Status: %v", err)
}
if report.PoolSize != 2 {
t.Fatalf("status PoolSize = %d, want 2", report.PoolSize)
}
foundSession := false
for _, s := range report.Sessions {
if s.ID == "s1" {
foundSession = true
if s.State != string(SessionDone) {
t.Fatalf("session s1 state = %s, want done", s.State)
}
}
}
if !foundSession {
t.Fatalf("status missing session s1: %+v", report.Sessions)
}
// --- attach to the finished session sees the buffered history ---
attachClient, err := Dial(paths.Socket)
if err != nil {
t.Fatalf("Dial(attach): %v", err)
}
var attached []string
if err := attachClient.Attach("s1", func(line string) { attached = append(attached, line) }); err != nil {
t.Fatalf("Attach: %v", err)
}
attachClient.Close()
if len(attached) != 2 {
t.Fatalf("attach lines = %v, want 2 buffered", attached)
}
// --- graceful stop, daemon cleans up ---
stopClient, err := Dial(paths.Socket)
if err != nil {
t.Fatalf("Dial(stop): %v", err)
}
if err := stopClient.Shutdown(); err != nil {
t.Fatalf("Shutdown: %v", err)
}
stopClient.Close()
select {
case err := <-serveErr:
if err != nil {
t.Fatalf("Serve returned error: %v", err)
}
case <-time.After(3 * time.Second):
t.Fatal("Serve did not return after shutdown")
}
// Socket, status, and lock files are removed on exit.
for _, p := range []string{paths.Socket, paths.Status, paths.Lock} {
if _, err := os.Stat(p); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("file %s not cleaned up after shutdown: %v", p, err)
}
}
}
func TestServerSecondInstanceFails(t *testing.T) {
block := make(chan struct{})
defer close(block)
launcher, _ := seqLauncher(&fakeWorker{pid: 1, waitCh: block})
srv1, paths := newTestServer(t, launcher)
serveErr := make(chan error, 1)
go func() { serveErr <- srv1.Serve() }()
waitForFile(t, paths.Status)
// A second server on the same paths must refuse to start (single instance).
pool2, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})
mgr2, _ := NewSessionManager(SessionManagerOptions{Pool: pool2})
srv2, _ := NewServer(ServerOptions{Paths: paths, Manager: mgr2, Pool: pool2})
if err := srv2.Serve(); !errors.Is(err, ErrAlreadyRunning) {
t.Fatalf("second Serve err = %v, want ErrAlreadyRunning", err)
}
srv1.Shutdown()
<-serveErr
}
func TestServerRejectsUnknownCommand(t *testing.T) {
block := make(chan struct{})
defer close(block)
launcher, _ := seqLauncher(&fakeWorker{pid: 1, waitCh: block})
srv, paths := newTestServer(t, launcher)
go func() { _ = srv.Serve() }()
defer srv.Shutdown()
waitForFile(t, paths.Status)
client, err := Dial(paths.Socket)
if err != nil {
t.Fatalf("Dial: %v", err)
}
defer client.Close()
// Run with an empty session id is rejected with an error frame.
err = client.Run("", "", "hi", nil, nil)
if err == nil {
t.Fatal("run with empty session id must return an error")
}
}
+306
View File
@@ -0,0 +1,306 @@
package daemon
import (
"context"
"errors"
"sort"
"sync"
)
// Session/lease/roster model. Mirrors reference-daemon-code-agent-js/
// session-manager.js (sessionID -> worker + metrics), lease.js (one active
// request per worker) and roster.js (the live session registry). Because zero's
// `exec` worker is one-shot, each session is served by exactly one worker via
// Pool.Run, so "one active request per worker" holds inherently and the pool's
// bounded slots provide the queue-when-full behavior.
// SessionState is a session's lifecycle phase.
type SessionState string
const (
// SessionQueued: created, waiting for a free pool slot.
SessionQueued SessionState = "queued"
// SessionRunning: a worker is leased and producing output.
SessionRunning SessionState = "running"
// SessionDone: the worker finished cleanly.
SessionDone SessionState = "done"
// SessionFailed: the worker failed permanently / the run errored.
SessionFailed SessionState = "failed"
)
// ErrSessionExists is returned when starting a session ID already in use.
var ErrSessionExists = errors.New("daemon: session already exists")
// ErrSessionNotFound is returned by Get/Attach for an unknown session ID.
var ErrSessionNotFound = errors.New("daemon: session not found")
const defaultSessionBuffer = 1024
// Session is one routed agent run. It buffers recent output (so a late `attach`
// sees history) and fans live lines out to subscribers. It implements the pool's
// Sink (Line) and the lease "started" hook (Started).
type Session struct {
id string
cwd string
mu sync.Mutex
state SessionState
buffer []string
maxBuffer int
lines int
finished bool
err error
exitCode int
subscribers map[int]chan string
nextSub int
done chan struct{}
}
func newSession(id, cwd string, maxBuffer int) *Session {
if maxBuffer <= 0 {
maxBuffer = defaultSessionBuffer
}
return &Session{
id: id,
cwd: cwd,
state: SessionQueued,
maxBuffer: maxBuffer,
subscribers: map[int]chan string{},
done: make(chan struct{}),
}
}
// ID returns the session ID.
func (s *Session) ID() string { return s.id }
// Done is closed when the session finishes (success or failure).
func (s *Session) Done() <-chan struct{} { return s.done }
// Err returns the terminal error (nil on clean completion); valid after Done.
func (s *Session) Err() error {
s.mu.Lock()
defer s.mu.Unlock()
return s.err
}
// State returns the current lifecycle phase.
func (s *Session) State() SessionState {
s.mu.Lock()
defer s.mu.Unlock()
return s.state
}
// Started marks the session running once the pool leases a worker (lease hook).
func (s *Session) Started() {
s.mu.Lock()
if s.state == SessionQueued {
s.state = SessionRunning
}
s.mu.Unlock()
}
// Line records and broadcasts one stream-json output line.
func (s *Session) Line(line string) {
s.mu.Lock()
s.lines++
s.buffer = append(s.buffer, line)
if len(s.buffer) > s.maxBuffer {
// Drop the oldest line; the ring keeps only the most recent maxBuffer.
s.buffer = s.buffer[len(s.buffer)-s.maxBuffer:]
}
subs := make([]chan string, 0, len(s.subscribers))
for _, ch := range s.subscribers {
subs = append(subs, ch)
}
s.mu.Unlock()
for _, ch := range subs {
// Non-blocking: a slow/stalled subscriber must not stall the worker pump.
// Dropped live lines remain available via the buffer for a fresh attach.
select {
case ch <- line:
default:
}
}
}
func (s *Session) finish(exitCode int, err error) {
s.mu.Lock()
s.finished = true
s.exitCode = exitCode
s.err = err
if err != nil {
s.state = SessionFailed
} else {
s.state = SessionDone
}
subs := s.subscribers
s.subscribers = map[int]chan string{}
s.mu.Unlock()
for _, ch := range subs {
close(ch)
}
close(s.done)
}
// Subscribe returns the buffered history plus a channel of subsequent live lines
// and a cancel func. If the session has already finished, the live channel is
// returned already closed. Mirrors attach.js.
func (s *Session) Subscribe() (buffered []string, live <-chan string, cancel func()) {
s.mu.Lock()
defer s.mu.Unlock()
buffered = append([]string(nil), s.buffer...)
if s.finished {
ch := make(chan string)
close(ch)
return buffered, ch, func() {}
}
id := s.nextSub
s.nextSub++
ch := make(chan string, 256)
s.subscribers[id] = ch
cancel = func() {
s.mu.Lock()
if c, ok := s.subscribers[id]; ok {
delete(s.subscribers, id)
close(c)
}
s.mu.Unlock()
}
return buffered, ch, cancel
}
func (s *Session) status() SessionStatus {
s.mu.Lock()
defer s.mu.Unlock()
return SessionStatus{ID: s.id, State: string(s.state), Lines: s.lines}
}
func (s *Session) isFinished() bool {
s.mu.Lock()
defer s.mu.Unlock()
return s.finished
}
// defaultMaxSessions bounds the retained session registry so a long-running
// daemon does not accumulate finished sessions without limit.
const defaultMaxSessions = 256
// SessionManager owns the live session registry and routes each session to a
// pool worker. Mirrors session-manager.js + roster.js. Finished sessions are
// retained (so a late `attach` sees history) up to MaxSessions, past which the
// oldest FINISHED ones are evicted; running/queued sessions are never evicted.
type SessionManager struct {
pool *Pool
maxBuffer int
maxSessions int
mu sync.Mutex
sessions map[string]*Session
order []string // creation order, for FIFO eviction of finished sessions
}
// SessionManagerOptions configures a SessionManager.
type SessionManagerOptions struct {
Pool *Pool
MaxBuffer int // per-session output ring size; 0 => default
// MaxSessions caps the retained session registry; 0 => default. The oldest
// FINISHED sessions are evicted once the cap is exceeded.
MaxSessions int
}
// NewSessionManager builds a manager over pool.
func NewSessionManager(opts SessionManagerOptions) (*SessionManager, error) {
if opts.Pool == nil {
return nil, errors.New("daemon: session manager requires a Pool")
}
maxSessions := opts.MaxSessions
if maxSessions <= 0 {
maxSessions = defaultMaxSessions
}
return &SessionManager{
pool: opts.Pool,
maxBuffer: opts.MaxBuffer,
maxSessions: maxSessions,
sessions: map[string]*Session{},
}, nil
}
// Start creates a session for spec and dispatches it to the pool. It returns the
// Session immediately (non-blocking); the run proceeds in the background and the
// session transitions queued -> running -> done/failed. A duplicate ID is
// rejected with ErrSessionExists.
func (m *SessionManager) Start(ctx context.Context, spec WorkerSpec) (*Session, error) {
m.mu.Lock()
if _, exists := m.sessions[spec.Session]; exists {
m.mu.Unlock()
return nil, ErrSessionExists
}
sess := newSession(spec.Session, spec.Cwd, m.maxBuffer)
m.sessions[spec.Session] = sess
m.order = append(m.order, spec.Session)
m.pruneLocked()
m.mu.Unlock()
go func() {
code, err := m.pool.Run(ctx, spec, sess)
sess.finish(code, err)
}()
return sess, nil
}
// pruneLocked evicts the oldest FINISHED sessions while the registry exceeds
// maxSessions, so a long-running daemon does not retain finished sessions without
// bound. Running/queued sessions are never evicted (a transient burst can exceed
// the cap until those finish). Caller must hold m.mu.
func (m *SessionManager) pruneLocked() {
if len(m.sessions) <= m.maxSessions {
return
}
kept := m.order[:0]
for _, id := range m.order {
s, ok := m.sessions[id]
if !ok {
continue // already removed
}
if len(m.sessions) > m.maxSessions && s.isFinished() {
delete(m.sessions, id)
continue
}
kept = append(kept, id)
}
m.order = kept
}
// Get returns a session by ID.
func (m *SessionManager) Get(id string) (*Session, bool) {
m.mu.Lock()
defer m.mu.Unlock()
s, ok := m.sessions[id]
return s, ok
}
// Attach returns the buffered history + live stream for an existing session.
func (m *SessionManager) Attach(id string) (buffered []string, live <-chan string, cancel func(), err error) {
s, ok := m.Get(id)
if !ok {
return nil, nil, nil, ErrSessionNotFound
}
b, l, c := s.Subscribe()
return b, l, c, nil
}
// Statuses returns a snapshot of all sessions, sorted by ID for stable output.
func (m *SessionManager) Statuses() []SessionStatus {
m.mu.Lock()
sessions := make([]*Session, 0, len(m.sessions))
for _, s := range m.sessions {
sessions = append(sessions, s)
}
m.mu.Unlock()
out := make([]SessionStatus, 0, len(sessions))
for _, s := range sessions {
out = append(out, s.status())
}
sort.Slice(out, func(i, j int) bool { return out[i].ID < out[j].ID })
return out
}
+179
View File
@@ -0,0 +1,179 @@
package daemon
import (
"context"
"errors"
"sync/atomic"
"testing"
"time"
)
func drain(t *testing.T, buffered []string, live <-chan string) []string {
t.Helper()
got := append([]string(nil), buffered...)
timeout := time.After(2 * time.Second)
for {
select {
case line, ok := <-live:
if !ok {
return got
}
got = append(got, line)
case <-timeout:
t.Fatal("timed out draining session stream")
}
}
}
func TestSessionStartRoutesAndStreams(t *testing.T) {
launcher, _ := seqLauncher(&fakeWorker{pid: 1, out: []string{"e1", "e2", "e3"}, exitCode: 0})
pool, _ := NewPool(PoolOptions{Size: 2, Launcher: launcher})
mgr, err := NewSessionManager(SessionManagerOptions{Pool: pool})
if err != nil {
t.Fatalf("NewSessionManager: %v", err)
}
sess, err := mgr.Start(context.Background(), WorkerSpec{Session: "s1"})
if err != nil {
t.Fatalf("Start: %v", err)
}
buffered, live, cancel := sess.Subscribe()
defer cancel()
got := drain(t, buffered, live)
if len(got) != 3 || got[0] != "e1" || got[2] != "e3" {
t.Fatalf("streamed lines = %v, want [e1 e2 e3]", got)
}
<-sess.Done()
if sess.State() != SessionDone {
t.Fatalf("state = %s, want done", sess.State())
}
}
func TestSessionDuplicateIDRejected(t *testing.T) {
block := make(chan struct{})
defer close(block)
launcher, _ := seqLauncher(&fakeWorker{pid: 1, waitCh: block})
pool, _ := NewPool(PoolOptions{Size: 2, Launcher: launcher})
mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool})
if _, err := mgr.Start(context.Background(), WorkerSpec{Session: "dup"}); err != nil {
t.Fatalf("first Start: %v", err)
}
if _, err := mgr.Start(context.Background(), WorkerSpec{Session: "dup"}); !errors.Is(err, ErrSessionExists) {
t.Fatalf("second Start err = %v, want ErrSessionExists", err)
}
}
func TestSessionLeaseQueuesWhenPoolFull(t *testing.T) {
block := make(chan struct{})
a := &fakeWorker{pid: 1, waitCh: block}
b := &fakeWorker{pid: 2, out: []string{"b1"}, exitCode: 0}
launcher, calls := seqLauncher(a, b)
pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})
mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool})
sa, _ := mgr.Start(context.Background(), WorkerSpec{Session: "a"})
waitFor(t, func() bool { return sa.State() == SessionRunning })
sb, _ := mgr.Start(context.Background(), WorkerSpec{Session: "b"})
// b must stay queued (its worker not yet launched) while a holds the slot.
time.Sleep(50 * time.Millisecond)
if sb.State() != SessionQueued {
t.Fatalf("session b state = %s, want queued while pool full", sb.State())
}
if n := atomic.LoadInt32(calls); n != 1 {
t.Fatalf("launcher calls = %d, want 1 (b not launched yet)", n)
}
close(block) // a finishes, slot frees, b proceeds
<-sb.Done()
if sb.State() != SessionDone {
t.Fatalf("session b state = %s, want done", sb.State())
}
if n := atomic.LoadInt32(calls); n != 2 {
t.Fatalf("launcher calls = %d, want 2 after slot freed", n)
}
}
func TestSessionAttachAfterFinishSeesBuffer(t *testing.T) {
launcher, _ := seqLauncher(&fakeWorker{pid: 1, out: []string{"x", "y"}, exitCode: 0})
pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})
mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool})
sess, _ := mgr.Start(context.Background(), WorkerSpec{Session: "s"})
<-sess.Done()
buffered, live, cancel, err := mgr.Attach("s")
if err != nil {
t.Fatalf("Attach: %v", err)
}
defer cancel()
got := drain(t, buffered, live)
if len(got) != 2 || got[0] != "x" || got[1] != "y" {
t.Fatalf("attach buffered = %v, want [x y]", got)
}
if _, _, _, err := mgr.Attach("missing"); !errors.Is(err, ErrSessionNotFound) {
t.Fatalf("Attach(missing) err = %v, want ErrSessionNotFound", err)
}
}
func TestSessionManagerEvictsFinishedOverCap(t *testing.T) {
launcher, _ := seqLauncher(
&fakeWorker{pid: 1, out: []string{"a"}, exitCode: 0},
&fakeWorker{pid: 2, out: []string{"b"}, exitCode: 0},
&fakeWorker{pid: 3, out: []string{"c"}, exitCode: 0},
)
pool, _ := NewPool(PoolOptions{Size: 2, Launcher: launcher})
mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool, MaxSessions: 2})
for _, id := range []string{"s1", "s2", "s3"} {
s, err := mgr.Start(context.Background(), WorkerSpec{Session: id})
if err != nil {
t.Fatalf("Start(%s): %v", id, err)
}
<-s.Done() // finish before starting the next so prune can evict it
}
// Past the cap of 2, the oldest FINISHED session (s1) is evicted; the two
// newest are retained.
if _, ok := mgr.Get("s1"); ok {
t.Fatal("oldest finished session must be evicted past the cap")
}
if _, ok := mgr.Get("s2"); !ok {
t.Fatal("s2 must be retained")
}
if _, ok := mgr.Get("s3"); !ok {
t.Fatal("s3 (newest) must be retained")
}
}
func TestSessionManagerKeepsRunningOverCap(t *testing.T) {
// All sessions are running (workers block), so none are evicted even past cap.
block := make(chan struct{})
defer close(block)
launcher, _ := seqLauncher(
&fakeWorker{pid: 1, waitCh: block},
&fakeWorker{pid: 2, waitCh: block},
)
pool, _ := NewPool(PoolOptions{Size: 2, Launcher: launcher})
mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool, MaxSessions: 1})
s1, _ := mgr.Start(context.Background(), WorkerSpec{Session: "r1"})
s2, _ := mgr.Start(context.Background(), WorkerSpec{Session: "r2"})
waitFor(t, func() bool { return s1.State() == SessionRunning && s2.State() == SessionRunning })
if _, ok := mgr.Get("r1"); !ok {
t.Fatal("a running session must never be evicted, even past the cap")
}
if _, ok := mgr.Get("r2"); !ok {
t.Fatal("a running session must never be evicted, even past the cap")
}
}
func TestSessionStatuses(t *testing.T) {
launcher, _ := seqLauncher(&fakeWorker{pid: 1, out: []string{"l"}, exitCode: 0})
pool, _ := NewPool(PoolOptions{Size: 1, Launcher: launcher})
mgr, _ := NewSessionManager(SessionManagerOptions{Pool: pool})
sess, _ := mgr.Start(context.Background(), WorkerSpec{Session: "only"})
<-sess.Done()
st := mgr.Statuses()
if len(st) != 1 || st[0].ID != "only" || st[0].State != string(SessionDone) || st[0].Lines != 1 {
t.Fatalf("statuses = %+v, want one done session with 1 line", st)
}
}
+27
View File
@@ -0,0 +1,27 @@
package daemon
import (
"fmt"
"os"
"path/filepath"
)
// maxUnixSocketPath bounds the socket path to the smallest platform sun_path
// limit so we surface a clear error instead of a cryptic bind failure. macOS
// allows 104 bytes incl. the NUL (so 103 chars); Linux allows 108. 103 is the
// safe cross-platform ceiling.
const maxUnixSocketPath = 103
// secureSocketParent creates the socket's parent directory owner-only (0700 on
// POSIX; on Windows the per-user profile directory is already ACL-restricted).
func secureSocketParent(socketPath string) error {
return os.MkdirAll(filepath.Dir(socketPath), 0o700)
}
// checkSocketPathLength rejects an over-long unix socket path before bind.
func checkSocketPathLength(socketPath string) error {
if len(socketPath) > maxUnixSocketPath {
return fmt.Errorf("daemon: socket path too long (%d > %d): %s", len(socketPath), maxUnixSocketPath, socketPath)
}
return nil
}
+12
View File
@@ -0,0 +1,12 @@
//go:build !windows
package daemon
import "os"
// hardenSocketFile restricts the control socket to its owner (0600) so no other
// local user can connect. Combined with the 0700 parent directory this makes the
// control plane owner-only, as required by the security model.
func hardenSocketFile(path string) error {
return os.Chmod(path, 0o600)
}
+9
View File
@@ -0,0 +1,9 @@
//go:build windows
package daemon
// hardenSocketFile is a no-op on Windows: AF_UNIX socket files do not honor POSIX
// mode bits, and the daemon places the socket under the per-user profile
// directory, which is already ACL-restricted to the owner. A dedicated restricted
// ACL (or a named pipe with an explicit DACL) is a documented follow-up.
func hardenSocketFile(string) error { return nil }
+43
View File
@@ -0,0 +1,43 @@
package daemon
import "time"
// StatusReport is the daemon/worker/session snapshot returned by `zero daemon
// status` (CtrlStatusResult). Mirrors reference-daemon-code-agent-js/status.js.
type StatusReport struct {
PID int `json:"pid"`
Version int `json:"version"`
Socket string `json:"socket"`
StartedAt time.Time `json:"startedAt"`
PoolSize int `json:"poolSize"`
Workers []WorkerStatus `json:"workers"`
Sessions []SessionStatus `json:"sessions"`
QueueDepth int `json:"queueDepth"`
}
// WorkerStatus is one worker's lifecycle snapshot.
type WorkerStatus struct {
ID int `json:"id"`
PID int `json:"pid"`
State string `json:"state"`
ConsecutiveCrashes int `json:"consecutiveCrashes"`
Restarts int `json:"restarts"`
Session string `json:"session,omitempty"`
}
// SessionStatus is one session's routing/metrics snapshot.
type SessionStatus struct {
ID string `json:"id"`
State string `json:"state"`
Lines int `json:"lines"`
}
// StatusFile is the on-disk daemon status document (pid, socket, version,
// started-at), written next to the lock so `status`/`stop` work without a live
// connection. Mirrors supervisor.js's status file.
type StatusFile struct {
PID int `json:"pid"`
Socket string `json:"socket"`
Version int `json:"version"`
StartedAt time.Time `json:"startedAt"`
}
+4 -1
View File
@@ -158,6 +158,9 @@ func providerModelCheck(profile config.ProviderProfile) Check {
if emptyProviderProfile(profile) {
return check("provider.model", "Provider model", StatusWarn, "Model validity was skipped because provider config is unavailable.", nil)
}
if strings.TrimSpace(profile.Model) == "" {
return check("provider.model", "Provider model", StatusFail, "Provider model is required.", map[string]any{"provider": providerName(profile)})
}
registry, err := modelregistry.DefaultRegistry()
if err != nil {
return check("provider.model", "Provider model", StatusFail, "Model registry could not be loaded: "+err.Error(), nil)
@@ -165,7 +168,7 @@ func providerModelCheck(profile config.ProviderProfile) Check {
model, err := registry.Require(profile.Model)
if err != nil {
if profile.ProviderKind == config.ProviderKindOpenAICompatible || profile.ProviderKind == config.ProviderKindAnthropicCompat {
return check("provider.model", "Provider model", StatusWarn, fmt.Sprintf("Custom %s model was not found in the Zero registry; runtime will pass it through to the configured provider.", profile.ProviderKind), map[string]any{"model": profile.Model, "provider": providerName(profile)})
return check("provider.model", "Provider model", StatusWarn, fmt.Sprintf("Custom %s model was not found in the Zero registry; runtime will pass it through to the configured provider. Run `zero doctor --connectivity` to validate the endpoint and auth.", profile.ProviderKind), map[string]any{"model": profile.Model, "provider": providerName(profile)})
}
return check("provider.model", "Provider model", StatusFail, "Provider model is invalid: "+err.Error(), map[string]any{"model": profile.Model})
}
+90
View File
@@ -86,6 +86,96 @@ func TestRunReportWarnsForUnknownOpenAICompatibleModel(t *testing.T) {
}
}
func TestProviderModelUnknownOpenAICompatibleModelWarnsPassThrough(t *testing.T) {
report := Run(Options{
Now: fixedDoctorClock("2026-06-04T15:45:00Z"),
Runtime: "go",
Provider: config.ProviderProfile{
Name: "openrouter",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "https://openrouter.ai/api/v1",
Model: "vendor/new-dynamic-model",
},
})
if !report.OK {
t.Fatalf("unknown openai-compatible model should warn, not fail: %#v", report)
}
check := report.Check("provider.model")
if check == nil {
t.Fatalf("expected provider.model check: %#v", report.Checks)
}
if check.Status != StatusWarn {
t.Fatalf("provider.model status = %s, want warn: %#v", check.Status, check)
}
if !strings.Contains(check.Message, "runtime will pass it through") || !strings.Contains(check.Message, "doctor --connectivity") {
t.Fatalf("provider.model message should explain pass-through connectivity validation, got %q", check.Message)
}
}
func TestProviderModelMissingModelFails(t *testing.T) {
report := Run(Options{
Now: fixedDoctorClock("2026-06-04T15:45:00Z"),
Runtime: "go",
Provider: config.ProviderProfile{
Name: "ollama-cloud",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "https://ollama.com/v1",
Model: " ",
},
})
if report.OK {
t.Fatalf("missing model should fail: %#v", report)
}
check := report.Check("provider.model")
if check == nil || check.Status != StatusFail || !strings.Contains(check.Message, "required") {
t.Fatalf("expected missing model failure: %#v", report.Checks)
}
}
func TestProviderModelBuiltInUnknownModelFails(t *testing.T) {
report := Run(Options{
Now: fixedDoctorClock("2026-06-04T15:45:00Z"),
Runtime: "go",
Provider: config.ProviderProfile{
Name: "openai",
ProviderKind: config.ProviderKindOpenAI,
BaseURL: config.OpenAIBaseURL,
Model: "vendor/new-dynamic-model",
},
})
if report.OK {
t.Fatalf("built-in provider unknown model should fail: %#v", report)
}
check := report.Check("provider.model")
if check == nil || check.Status != StatusFail || !strings.Contains(check.Message, "unknown Zero model") {
t.Fatalf("expected built-in unknown model failure: %#v", report.Checks)
}
}
func TestProviderModelRegisteredProviderMismatchFails(t *testing.T) {
report := Run(Options{
Now: fixedDoctorClock("2026-06-04T15:45:00Z"),
Runtime: "go",
Provider: config.ProviderProfile{
Name: "openai",
ProviderKind: config.ProviderKindOpenAI,
BaseURL: config.OpenAIBaseURL,
Model: "gemini-2.5-pro",
},
})
if report.OK {
t.Fatalf("registered provider mismatch should fail: %#v", report)
}
check := report.Check("provider.model")
if check == nil || check.Status != StatusFail || !strings.Contains(check.Message, "not available for provider") {
t.Fatalf("expected registered provider mismatch failure: %#v", report.Checks)
}
}
func writeDoctorConfig(t *testing.T, body string) string {
t.Helper()
dir := t.TempDir()
@@ -62,7 +62,7 @@ func TestDetectLocalRuntimesReportsReachableRuntime(t *testing.T) {
if !detected[0].Reachable {
t.Fatalf("runtime should be reachable: %#v", detected[0])
}
if detected[0].Models == nil || len(detected[0].Models) == 0 || detected[0].Models[0] != "llama3.1" {
if len(detected[0].Models) == 0 || detected[0].Models[0] != "llama3.1" {
t.Fatalf("expected discovered model list, got %#v", detected[0].Models)
}
}
+35
View File
@@ -40,6 +40,7 @@ func New(profile config.ProviderProfile, options Options) (zeroruntime.Provider,
MaxTokens: resolved.maxOutputTokens,
HTTPClient: options.HTTPClient,
UserAgent: options.UserAgent,
ParseThinkTags: parseThinkTagsForProfile(profile, resolved),
})
case config.ProviderKindAnthropic, config.ProviderKindAnthropicCompat:
return anthropic.New(anthropic.Options{
@@ -72,6 +73,40 @@ func New(profile config.ProviderProfile, options Options) (zeroruntime.Provider,
}
}
func parseThinkTagsForProfile(profile config.ProviderProfile, resolved resolvedProfile) bool {
if profile.ParseThinkTags != nil {
return *profile.ParseThinkTags
}
if resolved.providerKind != config.ProviderKindOpenAICompatible {
return false
}
return modelMayEmitThinkTags(resolved.apiModel)
}
func modelMayEmitThinkTags(model string) bool {
model = strings.ToLower(strings.TrimSpace(model))
for _, marker := range []string{
"deepseek-r1",
"deepseek-reasoner",
"gpt-oss",
"glm-z1",
"kimi-k2-thinking",
"magistral",
"minimax-m3",
"nemotron",
"qwen3",
"qwq",
"reasoner",
"reasoning",
"thinking",
} {
if strings.Contains(model, marker) {
return true
}
}
return false
}
type resolvedProfile struct {
providerKind config.ProviderKind
apiModel string
+27
View File
@@ -108,6 +108,33 @@ func TestNewSupportsOpenAIProviderKind(t *testing.T) {
}
}
func TestParseThinkTagsForProfileUsesConservativeDefaultsAndOverride(t *testing.T) {
openAICompatible := resolvedProfile{providerKind: config.ProviderKindOpenAICompatible, apiModel: "qwen3-coder:480b"}
if !parseThinkTagsForProfile(config.ProviderProfile{}, openAICompatible) {
t.Fatal("qwen3 OpenAI-compatible model should parse inline think tags by default")
}
generic := resolvedProfile{providerKind: config.ProviderKindOpenAICompatible, apiModel: "factory-model"}
if parseThinkTagsForProfile(config.ProviderProfile{}, generic) {
t.Fatal("generic OpenAI-compatible model should preserve literal think tags by default")
}
official := resolvedProfile{providerKind: config.ProviderKindOpenAI, apiModel: "gpt-4.1"}
if parseThinkTagsForProfile(config.ProviderProfile{}, official) {
t.Fatal("official OpenAI model should preserve literal think tags by default")
}
enabled := true
if !parseThinkTagsForProfile(config.ProviderProfile{ParseThinkTags: &enabled}, generic) {
t.Fatal("explicit parseThinkTags=true should enable inline think parsing")
}
disabled := false
if parseThinkTagsForProfile(config.ProviderProfile{ParseThinkTags: &disabled}, openAICompatible) {
t.Fatal("explicit parseThinkTags=false should disable inline think parsing")
}
}
func TestNewResolvesKnownModelToAPIModelAndProvider(t *testing.T) {
transport := &captureTransport{
responseBody: "data: {\"type\":\"message_stop\"}\n\n",
+24 -8
View File
@@ -40,6 +40,9 @@ type Options struct {
// StreamIdleTimeout aborts the stream if no data arrives for this long.
// Zero uses defaultStreamIdleTimeout.
StreamIdleTimeout time.Duration
// ParseThinkTags converts streamed <think>...</think> content into reasoning
// events for OpenAI-compatible models known to emit that legacy format.
ParseThinkTags bool
}
// Provider streams completions from an OpenAI-compatible chat completions API.
@@ -55,6 +58,7 @@ type Provider struct {
httpClient *http.Client
userAgent string
streamIdleTimeout time.Duration
parseThinkTags bool
}
// New creates an OpenAI-compatible provider.
@@ -100,6 +104,7 @@ func New(options Options) (*Provider, error) {
httpClient: httpClient,
userAgent: options.UserAgent,
streamIdleTimeout: idleTimeout,
parseThinkTags: options.ParseThinkTags,
}, nil
}
@@ -162,7 +167,7 @@ func (provider *Provider) stream(ctx context.Context, body []byte, events chan<-
return
}
state := newToolState()
state := newToolState(provider.parseThinkTags)
// Use the shared SSE reader (also used by the Anthropic/Gemini providers) so
// multi-line "data:" continuation fields are joined into one payload, and the
// idle watchdog / context cancellation are handled uniformly.
@@ -170,7 +175,8 @@ func (provider *Provider) stream(ctx context.Context, body []byte, events chan<-
return provider.emitPayload(ctx, data, state, events)
})
if errors.Is(err, providerio.ErrStreamIdle) {
state.closeOpen(ctx, events)
state.flushBufferedContent(events)
state.closeBufferedOpen(events)
sendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventError,
Error: provider.redact(fmt.Sprintf("provider stream error: idle timeout after %s (upstream stopped sending data)", provider.streamIdleTimeout)),
@@ -178,16 +184,19 @@ func (provider *Provider) stream(ctx context.Context, body []byte, events chan<-
return
}
if err != nil {
state.closeOpen(ctx, events)
state.flushBufferedContent(events)
state.closeBufferedOpen(events)
sendEvent(ctx, events, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventError, Error: provider.redact("provider stream error: " + err.Error())})
return
}
if ctxErr := ctx.Err(); ctxErr != nil {
state.closeOpen(ctx, events)
state.flushBufferedContent(events)
state.closeBufferedOpen(events)
sendEvent(ctx, events, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventError, Error: provider.redact("provider stream error: " + ctxErr.Error())})
return
}
if !state.done {
state.flushContent(ctx, events)
state.closeOpen(ctx, events)
sendEvent(ctx, events, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone, FinishReason: state.finishReason})
}
@@ -199,6 +208,7 @@ func (provider *Provider) stream(ctx context.Context, body []byte, events chan<-
func (provider *Provider) emitPayload(ctx context.Context, data string, state *toolState, events chan<- zeroruntime.StreamEvent) bool {
var chunk streamChunk
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
state.flushContent(ctx, events)
state.closeOpen(ctx, events)
sendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventError,
@@ -208,6 +218,7 @@ func (provider *Provider) emitPayload(ctx context.Context, data string, state *t
return false
}
if chunk.Error != nil {
state.flushContent(ctx, events)
state.closeOpen(ctx, events)
sendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventError,
@@ -234,15 +245,13 @@ func (provider *Provider) emitChunk(
})
}
if choice.Delta.Content != "" {
sendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventText,
Content: choice.Delta.Content,
})
state.emitContent(ctx, events, choice.Delta.Content)
}
for _, toolCall := range choice.Delta.ToolCalls {
state.applyDelta(ctx, toolCall, events)
}
if choice.FinishReason == "tool_calls" {
state.flushContent(ctx, events)
state.closeOpen(ctx, events)
}
if reason := mapFinishReason(choice.FinishReason); reason != "" {
@@ -314,6 +323,13 @@ func sendEvent(ctx context.Context, events chan<- zeroruntime.StreamEvent, event
}
}
func sendBufferedEvent(events chan<- zeroruntime.StreamEvent, event zeroruntime.StreamEvent) {
select {
case events <- event:
default:
}
}
func (provider *Provider) openAIRequest(request zeroruntime.CompletionRequest) chatCompletionRequest {
messages := make([]chatMessage, 0, len(request.Messages))
for _, message := range request.Messages {
+114 -3
View File
@@ -244,6 +244,53 @@ func TestStreamCompletionEmitsReasoningBeforeRegularContent(t *testing.T) {
}
}
func TestStreamCompletionPreservesLiteralThinkTagsByDefault(t *testing.T) {
provider := newTestProvider(t, func(w http.ResponseWriter, r *http.Request) {
writeSSE(w, `{"choices":[{"delta":{"content":"show <think>literal</think> markup"}}]}`)
writeSSE(w, `[DONE]`)
})
events := collectProviderEvents(t, provider)
assertEvent(t, events[0], zeroruntime.StreamEventText, "show <think>literal</think> markup")
if reasoning := eventsOfType(events, zeroruntime.StreamEventReasoning); len(reasoning) != 0 {
t.Fatalf("literal think tags must not emit reasoning by default, got %#v", reasoning)
}
}
func TestStreamCompletionSplitsInlineThinkTagsFromContent(t *testing.T) {
provider := newTestProviderWithThinkTags(t, func(w http.ResponseWriter, r *http.Request) {
writeSSE(w, `{"choices":[{"delta":{"content":"<think>private reasoning</think>public answer"}}]}`)
writeSSE(w, `[DONE]`)
})
events := collectProviderEvents(t, provider)
assertEvent(t, events[0], zeroruntime.StreamEventReasoning, "private reasoning")
assertEvent(t, events[1], zeroruntime.StreamEventText, "public answer")
if events[2].Type != zeroruntime.StreamEventDone {
t.Fatalf("third event = %#v, want done", events[2])
}
}
func TestStreamCompletionSplitsInlineThinkTagsAcrossChunks(t *testing.T) {
provider := newTestProviderWithThinkTags(t, func(w http.ResponseWriter, r *http.Request) {
writeSSE(w, `{"choices":[{"delta":{"content":"<thi"}}]}`)
writeSSE(w, `{"choices":[{"delta":{"content":"nk>reason"}}]}`)
writeSSE(w, `{"choices":[{"delta":{"content":"ing</"}}]}`)
writeSSE(w, `{"choices":[{"delta":{"content":"think> answer"}}]}`)
writeSSE(w, `[DONE]`)
})
events := collectProviderEvents(t, provider)
reasoning := eventsOfType(events, zeroruntime.StreamEventReasoning)
if len(reasoning) != 2 || reasoning[0].Content != "reason" || reasoning[1].Content != "ing" {
t.Fatalf("reasoning events = %#v, want split reasoning content", reasoning)
}
text := eventsOfType(events, zeroruntime.StreamEventText)
if len(text) != 1 || text[0].Content != " answer" {
t.Fatalf("text events = %#v, want answer-only content", text)
}
}
func TestStreamCompletionBuffersToolArgsUntilIDAndNameArrive(t *testing.T) {
provider := newTestProvider(t, func(w http.ResponseWriter, r *http.Request) {
writeSSE(w, `{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"path\":"}}]}}]}`)
@@ -385,19 +432,83 @@ func TestStreamCompletionEmitsErrorWhenContextCancels(t *testing.T) {
}
}
func TestStreamCompletionFlushesBufferedContentWhenContextCancels(t *testing.T) {
release := make(chan struct{})
provider := newTestProviderWithThinkTags(t, func(w http.ResponseWriter, r *http.Request) {
writeSSE(w, `{"choices":[{"delta":{"content":"visible <thi"}}]}`)
select {
case <-r.Context().Done():
case <-release:
}
})
defer close(release)
ctx, cancel := context.WithCancel(context.Background())
stream, err := provider.StreamCompletion(ctx, zeroruntime.CompletionRequest{})
if err != nil {
t.Fatalf("StreamCompletion returned setup error: %v", err)
}
events := []zeroruntime.StreamEvent{}
select {
case event, ok := <-stream:
if !ok {
t.Fatal("stream closed before first text event")
}
events = append(events, event)
case <-time.After(2 * time.Second):
t.Fatal("timed out waiting for first text event")
}
cancel()
events = append(events, readAll(stream)...)
if len(events) != 3 ||
events[0].Type != zeroruntime.StreamEventText || events[0].Content != "visible " ||
events[1].Type != zeroruntime.StreamEventText || events[1].Content != "<thi" ||
events[2].Type != zeroruntime.StreamEventError {
t.Fatalf("events = %#v, want text, buffered text, then context error", events)
}
if done := eventsOfType(events, zeroruntime.StreamEventDone); len(done) != 0 {
t.Fatalf("events = %#v, want no done after context cancel", events)
}
}
func newTestProvider(t *testing.T, handler http.HandlerFunc) *Provider {
t.Helper()
return newTestProviderWithKey(t, "", handler)
}
func newTestProviderWithKey(t *testing.T, apiKey string, handler http.HandlerFunc) *Provider {
t.Helper()
return newTestProviderWithOptions(t, Options{APIKey: apiKey}, handler)
}
func newTestProviderWithThinkTags(t *testing.T, handler http.HandlerFunc) *Provider {
t.Helper()
return newTestProviderWithOptions(t, Options{ParseThinkTags: true}, handler)
}
func newTestProviderWithOptions(t *testing.T, options Options, handler http.HandlerFunc) *Provider {
t.Helper()
server := httptest.NewServer(handler)
t.Cleanup(server.Close)
options.BaseURL = server.URL
if strings.TrimSpace(options.Model) == "" {
options.Model = "gpt-test"
}
provider, err := New(Options{
APIKey: apiKey,
BaseURL: server.URL,
Model: "gpt-test",
APIKey: options.APIKey,
BaseURL: options.BaseURL,
Model: options.Model,
AuthHeader: options.AuthHeader,
AuthScheme: options.AuthScheme,
AuthHeaderValue: options.AuthHeaderValue,
CustomHeaders: options.CustomHeaders,
HTTPClient: options.HTTPClient,
UserAgent: options.UserAgent,
MaxTokens: options.MaxTokens,
StreamIdleTimeout: options.StreamIdleTimeout,
ParseThinkTags: options.ParseThinkTags,
})
if err != nil {
t.Fatalf("New returned error: %v", err)
+147 -7
View File
@@ -3,12 +3,20 @@ package openai
import (
"context"
"sort"
"strings"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
const (
thinkOpenTag = "<think>"
thinkCloseTag = "</think>"
)
type toolState struct {
calls map[int]*pendingToolCall
calls map[int]*pendingToolCall
think thinkTagSplitter
parseThinkTags bool
// finishReason holds the normalized terminal stop reason (zeroruntime
// FinishReason*) when the response ended abnormally (length/content_filter),
// so the provider can attach it to the done event. Empty for a normal finish.
@@ -26,8 +34,128 @@ type pendingToolCall struct {
ended bool
}
func newToolState() *toolState {
return &toolState{calls: make(map[int]*pendingToolCall)}
func newToolState(parseThinkTags bool) *toolState {
return &toolState{calls: make(map[int]*pendingToolCall), parseThinkTags: parseThinkTags}
}
type streamEventEmitter func(zeroruntime.StreamEvent)
type thinkTagSplitter struct {
pending string
inThinking bool
}
func (state *toolState) emitContent(ctx context.Context, events chan<- zeroruntime.StreamEvent, content string) {
state.emitContentWith(content, func(event zeroruntime.StreamEvent) {
sendEvent(ctx, events, event)
})
}
func (state *toolState) flushContent(ctx context.Context, events chan<- zeroruntime.StreamEvent) {
state.flushContentWith(func(event zeroruntime.StreamEvent) {
sendEvent(ctx, events, event)
})
}
func (state *toolState) flushBufferedContent(events chan<- zeroruntime.StreamEvent) {
state.flushContentWith(func(event zeroruntime.StreamEvent) {
sendBufferedEvent(events, event)
})
}
func (state *toolState) emitContentWith(content string, emit streamEventEmitter) {
if !state.parseThinkTags {
emit(zeroruntime.StreamEvent{Type: zeroruntime.StreamEventText, Content: content})
return
}
state.think.push(content, func(eventType zeroruntime.StreamEventType, text string) {
emit(zeroruntime.StreamEvent{Type: eventType, Content: text})
})
}
func (state *toolState) flushContentWith(emit streamEventEmitter) {
if !state.parseThinkTags {
return
}
state.think.flush(func(eventType zeroruntime.StreamEventType, text string) {
emit(zeroruntime.StreamEvent{Type: eventType, Content: text})
})
}
func (splitter *thinkTagSplitter) push(content string, emit func(zeroruntime.StreamEventType, string)) {
if content == "" {
return
}
splitter.pending += content
splitter.drain(false, emit)
}
func (splitter *thinkTagSplitter) flush(emit func(zeroruntime.StreamEventType, string)) {
splitter.drain(true, emit)
}
func (splitter *thinkTagSplitter) drain(final bool, emit func(zeroruntime.StreamEventType, string)) {
for {
tag := thinkOpenTag
if splitter.inThinking {
tag = thinkCloseTag
}
if index := indexFold(splitter.pending, tag); index >= 0 {
splitter.emitCurrent(emit, splitter.pending[:index])
splitter.pending = splitter.pending[index+len(tag):]
splitter.inThinking = !splitter.inThinking
continue
}
if final {
splitter.emitCurrent(emit, splitter.pending)
splitter.pending = ""
return
}
keep := trailingTagPrefixLen(splitter.pending, tag)
if keep == len(splitter.pending) {
return
}
emitText := splitter.pending[:len(splitter.pending)-keep]
splitter.emitCurrent(emit, emitText)
splitter.pending = splitter.pending[len(splitter.pending)-keep:]
return
}
}
func (splitter *thinkTagSplitter) emitCurrent(emit func(zeroruntime.StreamEventType, string), text string) {
if text == "" {
return
}
eventType := zeroruntime.StreamEventText
if splitter.inThinking {
eventType = zeroruntime.StreamEventReasoning
}
emit(eventType, text)
}
func trailingTagPrefixLen(text string, tag string) int {
maxLen := min(len(text), len(tag)-1)
for length := maxLen; length > 0; length-- {
if strings.EqualFold(text[len(text)-length:], tag[:length]) {
return length
}
}
return 0
}
func indexFold(text string, needle string) int {
if needle == "" {
return 0
}
if len(text) < len(needle) {
return -1
}
for index := 0; index <= len(text)-len(needle); index++ {
if strings.EqualFold(text[index:index+len(needle)], needle) {
return index
}
}
return -1
}
func (state *toolState) applyDelta(
@@ -78,6 +206,18 @@ func (state *toolState) applyDelta(
}
func (state *toolState) closeOpen(ctx context.Context, events chan<- zeroruntime.StreamEvent) {
state.closeOpenWith(func(event zeroruntime.StreamEvent) {
sendEvent(ctx, events, event)
})
}
func (state *toolState) closeBufferedOpen(events chan<- zeroruntime.StreamEvent) {
state.closeOpenWith(func(event zeroruntime.StreamEvent) {
sendBufferedEvent(events, event)
})
}
func (state *toolState) closeOpenWith(emit streamEventEmitter) {
indexes := make([]int, 0, len(state.calls))
for index := range state.calls {
indexes = append(indexes, index)
@@ -95,20 +235,20 @@ func (state *toolState) closeOpen(ctx context.Context, events chan<- zeroruntime
if call.id == "" || call.name == "" {
if call.id != "" || call.name != "" || call.arguments != "" {
call.ended = true
sendEvent(ctx, events, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventToolCallDropped})
emit(zeroruntime.StreamEvent{Type: zeroruntime.StreamEventToolCallDropped})
}
continue
}
if !call.started {
call.started = true
sendEvent(ctx, events, zeroruntime.StreamEvent{
emit(zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventToolCallStart,
ToolCallID: call.id,
ToolName: call.name,
})
}
if call.arguments != "" {
sendEvent(ctx, events, zeroruntime.StreamEvent{
emit(zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventToolCallDelta,
ToolCallID: call.id,
ArgumentsFragment: call.arguments,
@@ -116,7 +256,7 @@ func (state *toolState) closeOpen(ctx context.Context, events chan<- zeroruntime
call.arguments = ""
}
call.ended = true
sendEvent(ctx, events, zeroruntime.StreamEvent{
emit(zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventToolCallEnd,
ToolCallID: call.id,
})
+42 -5
View File
@@ -23,15 +23,26 @@ type Backend struct {
// under a seatbelt profile that restricts outbound to the proxy port);
// bubblewrap isolates the network namespace with no bridge to the host proxy,
// so scoped egress there collapses to deny until a real relay exists.
ScopedEgress bool `json:"scopedEgress,omitempty"`
Executable string `json:"executable,omitempty"`
Message string `json:"message,omitempty"`
ScopedEgress bool `json:"scopedEgress,omitempty"`
// ProxyEgress reports a backend that routes scoped egress through the local
// filtering proxy on a BEST-EFFORT basis (the child must honor the proxy env)
// rather than via OS-level network isolation. The WSL policy-only fallback uses
// this: there is no netns/seatbelt to enforce egress, but the proxy still
// applies the allow/deny gate to well-behaved clients.
ProxyEgress bool `json:"proxyEgress,omitempty"`
Executable string `json:"executable,omitempty"`
Message string `json:"message,omitempty"`
}
// EnforcesScopedEgress reports whether a populated NetworkScoped allowlist can be
// enforced by this backend. When false, scoped egress must fail closed (collapse
// to deny) rather than silently run with unrestricted networking.
// routed through the filtering proxy by this backend. When false, scoped egress
// must fail closed (collapse to deny) rather than silently run with unrestricted
// networking. A native backend needs an executable; the WSL fallback uses
// best-effort proxy egress (ProxyEgress) with no native isolation.
func (backend Backend) EnforcesScopedEgress() bool {
if backend.ProxyEgress {
return true
}
return backend.Available && backend.Executable != "" && backend.ScopedEgress
}
@@ -65,6 +76,12 @@ func SelectBackend(options BackendOptions) Backend {
if path, err := lookup("bwrap"); err == nil && path != "" {
return nativeBackend(goos, BackendBubblewrap, path, "bubblewrap sandbox available")
}
// Under WSL (no bubblewrap), do NOT silently degrade to a plain policy-only
// runner: select the WSL backend, which still routes egress through the
// filtering proxy and is gated on AllowPolicyOnlyRunner in the runner.
if info := detectWSL(); info.IsWSL {
return wslBackend(goos, info)
}
return policyOnlyBackend(goos, "policy-only fallback: bubblewrap is not installed")
case "darwin":
if path, err := lookup("sandbox-exec"); err == nil && path != "" {
@@ -94,6 +111,26 @@ func nativeBackend(goos string, name BackendName, executable string, message str
}
}
// wslBackend is the policy-only WSL fallback: no native isolation, but scoped
// egress is routed best-effort through the filtering proxy (ProxyEgress). The
// runner gates it on Policy.AllowPolicyOnlyRunner and records a downgrade note.
func wslBackend(goos string, info WSLInfo) Backend {
msg := "policy-only WSL fallback: bubblewrap unavailable under WSL; egress routed through the filtering proxy"
if info.IsWSL2 {
msg = "policy-only WSL2 fallback: bubblewrap unavailable/unreliable under WSL2; egress routed through the filtering proxy"
}
return Backend{
Name: BackendWSL,
Available: false,
Platform: goos,
Fallback: true,
CommandWrapping: false,
NativeIsolation: false,
ProxyEgress: true,
Message: msg,
}
}
func policyOnlyBackend(goos string, message string) Backend {
return Backend{
Name: BackendPolicyOnly,
-33
View File
@@ -1,33 +0,0 @@
package sandbox
// RequiredAutonomyForBatch returns the highest autonomy tier required to run a
// batch of operations without further prompting, given each operation's
// classified risk. RiskCritical and RiskHigh both map to AutonomyHigh because
// the autonomy ladder has three tiers and the risk ladder has four. Any risk
// level that is empty or unrecognised fails closed to AutonomyHigh so an
// unclassified operation can never silently lower the batch ceiling. An empty
// batch requires only AutonomyLow.
func RequiredAutonomyForBatch(risks []Risk) Autonomy {
required := AutonomyLow
for _, risk := range risks {
tier := riskToAutonomy(risk.Level)
if autonomyRank[tier] > autonomyRank[required] {
required = tier
}
}
return required
}
func riskToAutonomy(level RiskLevel) Autonomy {
switch level {
case RiskLow:
return AutonomyLow
case RiskMedium:
return AutonomyMedium
case RiskHigh, RiskCritical:
return AutonomyHigh
default:
// Unknown / empty risk is treated as the highest tier (fail-closed).
return AutonomyHigh
}
}
-26
View File
@@ -1,26 +0,0 @@
package sandbox
import "testing"
func TestRequiredAutonomyForBatch(t *testing.T) {
cases := []struct {
name string
risks []Risk
want Autonomy
}{
{"empty slice is low", nil, AutonomyLow},
{"single low", []Risk{{Level: RiskLow}}, AutonomyLow},
{"mixed low and medium is medium", []Risk{{Level: RiskLow}, {Level: RiskMedium}}, AutonomyMedium},
{"any high is high", []Risk{{Level: RiskLow}, {Level: RiskHigh}}, AutonomyHigh},
{"any critical is high", []Risk{{Level: RiskMedium}, {Level: RiskCritical}}, AutonomyHigh},
{"unknown risk fails closed to high", []Risk{{Level: RiskLevel("weird")}}, AutonomyHigh},
{"empty-level risk fails closed to high", []Risk{{Level: RiskLevel("")}}, AutonomyHigh},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if got := RequiredAutonomyForBatch(tc.risks); got != tc.want {
t.Fatalf("RequiredAutonomyForBatch(%v) = %q, want %q", tc.risks, got, tc.want)
}
})
}
}
+94 -4
View File
@@ -2,6 +2,8 @@ package sandbox
import (
"context"
"crypto/tls"
"crypto/x509"
"errors"
"fmt"
"io"
@@ -32,9 +34,21 @@ import (
type egressProxy struct {
listener net.Listener
server *http.Server
allowed []string
denied []string
log func(string)
// socksListener serves the SOCKS5 CONNECT front-end on a SEPARATE loopback
// port. It runs the same authorize()/domainAllowed() gate as the HTTP proxy, so
// a client configured with ALL_PROXY=socks5://<socksAddr> is filtered
// identically. Bound to 127.0.0.1:0 like the HTTP listener.
socksListener net.Listener
allowed []string
denied []string
log func(string)
// mitm, when non-nil (Policy.InspectTLS), makes handleConnect TERMINATE TLS
// using a locally generated CA so the decrypted Host can be re-checked against
// the allow/deny gate. mitmTransport dials upstreams with full validation
// (system roots unless upstreamRoots is set for tests) — never InsecureSkipVerify.
mitm *mitmCA
mitmTransport *http.Transport
// domainPrompt, when set, is asked to authorize an UNKNOWN host (one neither
// allowed nor explicitly denied) instead of failing closed. It must return
@@ -71,6 +85,14 @@ type egressOptions struct {
// fail-closed behavior.
DomainPrompt func(ctx context.Context, host string, port int) (bool, error)
PromptTimeout time.Duration
// InspectTLS enables TLS termination (MITM) so the proxy can enforce the
// allow/deny gate on the decrypted request Host. Off => pure CONNECT
// passthrough (the default, byte-for-byte unchanged).
InspectTLS bool
// upstreamRoots overrides the root CAs used to validate UPSTREAM TLS. nil =
// system roots (production). Tests set it to an httptest server's pool so
// validation is exercised WITHOUT InsecureSkipVerify.
upstreamRoots *x509.CertPool
}
// errEmptyAllowlist is returned when a scoped-egress proxy is requested with no
@@ -89,9 +111,18 @@ func newEgressProxy(options egressOptions) (*egressProxy, error) {
if err != nil {
return nil, fmt.Errorf("bind scoped-egress proxy: %w", err)
}
// Bind the SOCKS5 front-end on its own loopback port. Fail closed: if it
// cannot bind, tear down the HTTP listener and return an error so the caller
// treats scoped egress as a full deny rather than starting a partial proxy.
socksListener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
_ = listener.Close()
return nil, fmt.Errorf("bind scoped-egress socks proxy: %w", err)
}
proxy := &egressProxy{
listener: listener,
socksListener: socksListener,
allowed: allowed,
denied: denied,
log: options.Log,
@@ -100,6 +131,31 @@ func newEgressProxy(options egressOptions) (*egressProxy, error) {
decisionCache: map[string]bool{},
inflight: map[string]chan struct{}{},
}
if options.InspectTLS {
ca, err := newMITMCA()
if err != nil {
_ = listener.Close()
_ = socksListener.Close()
return nil, err
}
proxy.mitm = ca
dialer := &net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}
proxy.mitmTransport = &http.Transport{
// Bound the TCP connect phase so a blackholed upstream cannot park a
// handler+goroutine for the OS default (minutes) — mirrors the
// net.DialTimeout the passthrough CONNECT path uses.
DialContext: dialer.DialContext,
// Upstream TLS is validated normally; upstreamRoots is nil in production
// (system roots) and set only by tests. NEVER InsecureSkipVerify.
TLSClientConfig: &tls.Config{MinVersion: tls.VersionTLS12, RootCAs: options.upstreamRoots},
// Keep upstream on HTTP/1.1: the decrypted client side is HTTP/1.1, so an
// HTTP/2 upstream response would serialize an invalid status line back to
// the client. (handleConnectMITM also normalizes the response proto.)
ForceAttemptHTTP2: false,
TLSHandshakeTimeout: 30 * time.Second,
ResponseHeaderTimeout: 30 * time.Second,
}
}
proxy.server = &http.Server{
Handler: http.HandlerFunc(proxy.handle),
ReadHeaderTimeout: 30 * time.Second,
@@ -108,10 +164,11 @@ func newEgressProxy(options egressOptions) (*egressProxy, error) {
// Serve returns ErrServerClosed on Close; nothing else to do.
_ = proxy.server.Serve(listener)
}()
go proxy.serveSocks(socksListener)
return proxy, nil
}
// Addr returns the loopback host:port the proxy is listening on.
// Addr returns the loopback host:port the HTTP proxy is listening on.
func (proxy *egressProxy) Addr() string {
if proxy == nil || proxy.listener == nil {
return ""
@@ -119,6 +176,25 @@ func (proxy *egressProxy) Addr() string {
return proxy.listener.Addr().String()
}
// SocksAddr returns the loopback host:port the SOCKS5 front-end is listening on,
// or "" when no SOCKS listener is active.
func (proxy *egressProxy) SocksAddr() string {
if proxy == nil || proxy.socksListener == nil {
return ""
}
return proxy.socksListener.Addr().String()
}
// CACertPEM returns the MITM CA's public certificate (PEM) when TLS inspection
// is active, or nil for a plain passthrough proxy. The runner writes this to a
// file and surfaces it via ZERO_SANDBOX_CA_CERT.
func (proxy *egressProxy) CACertPEM() []byte {
if proxy == nil || proxy.mitm == nil {
return nil
}
return proxy.mitm.caPEM()
}
// Close stops the proxy and releases its listener. It is safe to call multiple
// times.
func (proxy *egressProxy) Close() error {
@@ -130,12 +206,26 @@ func (proxy *egressProxy) Close() error {
if proxy.server != nil {
err = proxy.server.Close()
}
// Closing the SOCKS listener unblocks serveSocks's Accept loop so the
// goroutine returns. Report its error only if the HTTP close succeeded.
if proxy.socksListener != nil {
if cerr := proxy.socksListener.Close(); cerr != nil && err == nil {
err = cerr
}
}
if proxy.mitmTransport != nil {
proxy.mitmTransport.CloseIdleConnections()
}
})
return err
}
func (proxy *egressProxy) handle(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodConnect {
if proxy.mitm != nil {
proxy.handleConnectMITM(w, r)
return
}
proxy.handleConnect(w, r)
return
}
+156
View File
@@ -0,0 +1,156 @@
package sandbox
import (
"encoding/binary"
"io"
"net"
"strconv"
"time"
)
// SOCKS5 wire constants (RFC 1928). The proxy speaks only the no-auth CONNECT
// subset: enough for a sandboxed process configured with ALL_PROXY=socks5://...
// to tunnel TCP through the same allow/deny gate the HTTP proxy enforces.
const (
socksVersion5 = 0x05
socksNoAuth = 0x00
socksNoAcceptable = 0xFF
socksCmdConnect = 0x01
socksAtypIPv4 = 0x01
socksAtypDomain = 0x03
socksAtypIPv6 = 0x04
socksRepSuccess = 0x00
socksRepGeneralFailure = 0x01
socksRepNotAllowed = 0x02
socksRepCmdUnsupported = 0x07
)
// socksHandshakeTimeout bounds the greeting+request exchange so a stalled or
// malformed client cannot pin a goroutine; it is cleared before the relay so an
// established tunnel is not torn down.
const socksHandshakeTimeout = 30 * time.Second
// serveSocks accepts SOCKS5 connections until the listener is closed.
func (proxy *egressProxy) serveSocks(listener net.Listener) {
for {
conn, err := listener.Accept()
if err != nil {
return // listener closed
}
go proxy.handleSocksConn(conn)
}
}
// handleSocksConn runs one SOCKS5 CONNECT through the SAME authorize() gate as
// the HTTP proxy: an allowed (host, port) is dialed and relayed; a denied target
// gets a SOCKS "connection not allowed" reply and is never dialed. Fail-closed:
// any protocol/read error just closes the connection.
func (proxy *egressProxy) handleSocksConn(conn net.Conn) {
defer conn.Close()
_ = conn.SetDeadline(time.Now().Add(socksHandshakeTimeout))
host, port, ok := socksNegotiate(conn)
if !ok {
return
}
targetLabel := net.JoinHostPort(host, strconv.Itoa(port))
if !proxy.authorize(host, port) {
proxy.logDecision("deny", "SOCKS", targetLabel)
_ = socksReply(conn, socksRepNotAllowed)
return
}
upstream, err := net.DialTimeout("tcp", targetLabel, 30*time.Second)
if err != nil {
proxy.logDecision("deny", "SOCKS", targetLabel)
_ = socksReply(conn, socksRepGeneralFailure)
return
}
defer upstream.Close()
if err := socksReply(conn, socksRepSuccess); err != nil {
return
}
// Clear the handshake deadline; the tunnel is long-lived.
_ = conn.SetDeadline(time.Time{})
proxy.logDecision("allow", "SOCKS", targetLabel)
relay(conn, upstream)
}
// socksNegotiate performs the no-auth greeting and a CONNECT request, returning
// the requested host and port. ok=false means negotiation failed (an error reply
// has already been written where the protocol calls for one) and the caller
// should close the connection.
func socksNegotiate(conn net.Conn) (host string, port int, ok bool) {
greeting := make([]byte, 2) // VER, NMETHODS
if _, err := io.ReadFull(conn, greeting); err != nil || greeting[0] != socksVersion5 {
return "", 0, false
}
methods := make([]byte, int(greeting[1]))
if _, err := io.ReadFull(conn, methods); err != nil {
return "", 0, false
}
// RFC 1928: select a method the client actually offered. We only support
// no-auth; if the client did not offer it, reply "no acceptable methods"
// (0xFF) and abort rather than forcing a method it never advertised.
supportsNoAuth := false
for _, method := range methods {
if method == socksNoAuth {
supportsNoAuth = true
break
}
}
if !supportsNoAuth {
_, _ = conn.Write([]byte{socksVersion5, socksNoAcceptable})
return "", 0, false
}
if _, err := conn.Write([]byte{socksVersion5, socksNoAuth}); err != nil {
return "", 0, false
}
header := make([]byte, 4) // VER, CMD, RSV, ATYP
if _, err := io.ReadFull(conn, header); err != nil || header[0] != socksVersion5 {
return "", 0, false
}
if header[1] != socksCmdConnect {
_ = socksReply(conn, socksRepCmdUnsupported)
return "", 0, false
}
switch header[3] {
case socksAtypIPv4:
addr := make([]byte, net.IPv4len)
if _, err := io.ReadFull(conn, addr); err != nil {
return "", 0, false
}
host = net.IP(addr).String()
case socksAtypIPv6:
addr := make([]byte, net.IPv6len)
if _, err := io.ReadFull(conn, addr); err != nil {
return "", 0, false
}
host = net.IP(addr).String()
case socksAtypDomain:
length := make([]byte, 1)
if _, err := io.ReadFull(conn, length); err != nil {
return "", 0, false
}
name := make([]byte, int(length[0]))
if _, err := io.ReadFull(conn, name); err != nil {
return "", 0, false
}
host = string(name)
default:
_ = socksReply(conn, socksRepGeneralFailure)
return "", 0, false
}
portBytes := make([]byte, 2)
if _, err := io.ReadFull(conn, portBytes); err != nil {
return "", 0, false
}
return host, int(binary.BigEndian.Uint16(portBytes)), true
}
// socksReply writes a minimal SOCKS5 reply with the given status and a zero IPv4
// bound address (clients ignore BND.ADDR/BND.PORT for CONNECT).
func socksReply(conn net.Conn, status byte) error {
_, err := conn.Write([]byte{socksVersion5, status, 0x00, socksAtypIPv4, 0, 0, 0, 0, 0, 0})
return err
}
+187
View File
@@ -0,0 +1,187 @@
package sandbox
import (
"io"
"net"
"strconv"
"testing"
"time"
)
// socks5Connect performs a SOCKS5 no-auth greeting and a CONNECT request for the
// given ATYP/address/port through proxyAddr, returning the open connection and
// the server's REP byte. The caller closes the connection.
func socks5Connect(t *testing.T, proxyAddr string, atyp byte, addr []byte, port int) (net.Conn, byte) {
t.Helper()
conn, err := net.DialTimeout("tcp", proxyAddr, 2*time.Second)
if err != nil {
t.Fatalf("dial socks proxy: %v", err)
}
_ = conn.SetDeadline(time.Now().Add(2 * time.Second))
// Greeting: VER=5, NMETHODS=1, METHOD=0 (no auth).
if _, err := conn.Write([]byte{socksVersion5, 0x01, socksNoAuth}); err != nil {
t.Fatalf("write greeting: %v", err)
}
method := make([]byte, 2)
if _, err := io.ReadFull(conn, method); err != nil {
t.Fatalf("read method selection: %v", err)
}
if method[0] != socksVersion5 || method[1] != socksNoAuth {
t.Fatalf("method selection = %v, want [5 0]", method)
}
// Request: VER=5, CMD=CONNECT, RSV=0, ATYP, ADDR, PORT(2, big-endian).
request := append([]byte{socksVersion5, socksCmdConnect, 0x00, atyp}, addr...)
request = append(request, byte(port>>8), byte(port&0xff))
if _, err := conn.Write(request); err != nil {
t.Fatalf("write request: %v", err)
}
reply := make([]byte, 10) // VER, REP, RSV, ATYP(IPv4), BND.ADDR(4), BND.PORT(2)
if _, err := io.ReadFull(conn, reply); err != nil {
t.Fatalf("read reply: %v", err)
}
return conn, reply[1]
}
func TestEgressProxySocksAddrBindsLoopback(t *testing.T) {
proxy, err := newEgressProxy(egressOptions{Allowed: []string{"github.com"}})
if err != nil {
t.Fatalf("newEgressProxy: %v", err)
}
defer proxy.Close()
if proxy.SocksAddr() == "" {
t.Fatal("SocksAddr must be non-empty while the proxy is running")
}
if proxy.SocksAddr() == proxy.Addr() {
t.Fatalf("SOCKS and HTTP listeners must be distinct ports, both %q", proxy.Addr())
}
host, _, err := net.SplitHostPort(proxy.SocksAddr())
if err != nil {
t.Fatalf("SplitHostPort(%q): %v", proxy.SocksAddr(), err)
}
if host != "127.0.0.1" {
t.Fatalf("SOCKS bound to %q, want loopback 127.0.0.1", host)
}
}
// TestEgressProxySocksRejectsUnsupportedAuth verifies the proxy honors SOCKS5
// method negotiation: a client offering only methods the proxy does not support
// (here user/pass 0x02, no no-auth) gets the "no acceptable methods" reply (0xFF)
// rather than the proxy forcing a method the client never advertised.
func TestEgressProxySocksRejectsUnsupportedAuth(t *testing.T) {
proxy, err := newEgressProxy(egressOptions{Allowed: []string{"github.com"}})
if err != nil {
t.Fatalf("newEgressProxy: %v", err)
}
defer proxy.Close()
conn, err := net.DialTimeout("tcp", proxy.SocksAddr(), 2*time.Second)
if err != nil {
t.Fatalf("dial socks proxy: %v", err)
}
defer conn.Close()
_ = conn.SetDeadline(time.Now().Add(2 * time.Second))
// Greeting: VER=5, NMETHODS=1, METHOD=0x02 (user/pass) — no-auth NOT offered.
if _, err := conn.Write([]byte{socksVersion5, 0x01, 0x02}); err != nil {
t.Fatalf("write greeting: %v", err)
}
resp := make([]byte, 2)
if _, err := io.ReadFull(conn, resp); err != nil {
t.Fatalf("read method selection: %v", err)
}
if resp[0] != socksVersion5 || resp[1] != socksNoAcceptable {
t.Fatalf("method selection = %v, want [5 0xFF] (no acceptable methods)", resp)
}
}
// TestEgressProxySocksConnectAllowed verifies an allowlisted SOCKS5 CONNECT is
// authorized through the SAME gate, dialed, and relayed to a fake upstream.
func TestEgressProxySocksConnectAllowed(t *testing.T) {
upstream, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen upstream: %v", err)
}
defer upstream.Close()
go func() {
c, aerr := upstream.Accept()
if aerr != nil {
return
}
defer c.Close()
buf := make([]byte, 4)
if _, rerr := io.ReadFull(c, buf); rerr != nil {
return
}
_, _ = c.Write([]byte("pong"))
}()
_, portStr, _ := net.SplitHostPort(upstream.Addr().String())
port, _ := strconv.Atoi(portStr)
// Allow the loopback IP so the IPv4 CONNECT is authorized.
proxy, err := newEgressProxy(egressOptions{Allowed: []string{"127.0.0.1"}})
if err != nil {
t.Fatalf("newEgressProxy: %v", err)
}
defer proxy.Close()
conn, rep := socks5Connect(t, proxy.SocksAddr(), socksAtypIPv4, net.IPv4(127, 0, 0, 1).To4(), port)
defer conn.Close()
if rep != socksRepSuccess {
t.Fatalf("SOCKS CONNECT to allowed host REP = %#x, want success (0)", rep)
}
if _, err := conn.Write([]byte("ping")); err != nil {
t.Fatalf("write through tunnel: %v", err)
}
got := make([]byte, 4)
if _, err := io.ReadFull(conn, got); err != nil {
t.Fatalf("read through tunnel: %v", err)
}
if string(got) != "pong" {
t.Fatalf("tunneled response = %q, want pong", string(got))
}
}
// TestEgressProxySocksConnectDenied verifies a non-allowlisted SOCKS5 CONNECT
// (domain form) is refused with the "connection not allowed" reply.
func TestEgressProxySocksConnectDenied(t *testing.T) {
proxy, err := newEgressProxy(egressOptions{Allowed: []string{"allowed.example"}})
if err != nil {
t.Fatalf("newEgressProxy: %v", err)
}
defer proxy.Close()
domain := "denied.example"
addr := append([]byte{byte(len(domain))}, []byte(domain)...)
conn, rep := socks5Connect(t, proxy.SocksAddr(), socksAtypDomain, addr, 443)
conn.Close()
if rep != socksRepNotAllowed {
t.Fatalf("SOCKS CONNECT to denied host REP = %#x, want not-allowed (2)", rep)
}
}
// TestEgressProxySocksDenyWinsOverAllow verifies the SOCKS gate honors
// deny-wins: a denied subdomain of an allowed domain is refused (REP 2), proving
// the SOCKS path runs the same authorize()/domainAllowed() logic as the HTTP
// proxy rather than a parallel, weaker check.
func TestEgressProxySocksDenyWinsOverAllow(t *testing.T) {
proxy, err := newEgressProxy(egressOptions{
Allowed: []string{"allowed.example"},
Denied: []string{"blocked.allowed.example"},
})
if err != nil {
t.Fatalf("newEgressProxy: %v", err)
}
defer proxy.Close()
denied := "blocked.allowed.example"
addr := append([]byte{byte(len(denied))}, []byte(denied)...)
conn, rep := socks5Connect(t, proxy.SocksAddr(), socksAtypDomain, addr, 443)
conn.Close()
if rep != socksRepNotAllowed {
t.Fatalf("denied subdomain REP = %#x, want not-allowed (2) (deny must win)", rep)
}
}
+65 -6
View File
@@ -65,6 +65,35 @@ func (engine *Engine) Scope() *Scope {
return engine.scope
}
// ReadExclusions returns the resolved DenyRead/AllowRead exclusion matcher for
// this engine's policy, resolving each policy entry ONCE. The search tools build
// it a single time per run and reuse it across the whole walk so the predicates
// don't re-run Abs/EvalSymlinks per visited path. Returns nil for a nil engine
// (the matcher's methods treat nil as "exclude nothing").
func (engine *Engine) ReadExclusions() *ReadExclusions {
// A disabled policy enforces nothing, so it must not filter search results
// either (Evaluate already allows every request under ModeDisabled).
if engine == nil || engine.policy.Mode == ModeDisabled {
return nil
}
return &ReadExclusions{
workspaceRoot: engine.workspaceRoot,
denyRoots: resolvePolicyPaths(engine.policy.DenyRead),
allowRoots: resolvePolicyPaths(engine.policy.AllowRead),
}
}
// ReadExclusionGlobs returns the ripgrep-style --glob exclusion args for this
// engine's policy + scope (see the package-level ReadExclusionGlobs). Empty when
// DenyRead is unset or the engine has no scope.
func (engine *Engine) ReadExclusionGlobs() []string {
// A disabled policy filters nothing (parity with ReadExclusions / Evaluate).
if engine == nil || engine.policy.Mode == ModeDisabled {
return nil
}
return ReadExclusionGlobs(engine.policy, engine.scope)
}
// effectiveNetworkMode is the single source of truth for the engine's active
// network mode: it collapses an empty-allowlist scoped policy to deny, and ALSO
// downgrades scoped to deny when the backend cannot actually route scoped egress
@@ -90,6 +119,13 @@ func (engine *Engine) effectiveNetworkMode(policy Policy) NetworkMode {
// their pre-sandbox behaviour); deny blocks everything; scoped allows only hosts
// in AllowedDomains minus DeniedDomains (an empty allowlist, or a backend that
// can't enforce scoped egress, collapses to deny).
//
// By default the first-party, in-process tools that consult this gate are NOT
// subject to the network policy (EnforceToolNetwork is off): that policy exists
// to confine the sandboxed SHELL's egress, which these tools don't use, and they
// retain their own SSRF/port/redirect safeguards. Set EnforceToolNetwork to also
// hold them to the allow/scoped/deny policy. The sandboxed-shell egress decision
// lives in Evaluate via effectiveNetworkMode and is unaffected by this flag.
func (engine *Engine) NetworkHostAllowed(host string) (bool, NetworkMode) {
if engine == nil {
return true, NetworkAllow
@@ -98,6 +134,10 @@ func (engine *Engine) NetworkHostAllowed(host string) (bool, NetworkMode) {
if policy.Mode == ModeDisabled {
return true, NetworkAllow
}
if !policy.EnforceToolNetwork {
// First-party in-process tools are exempt unless the operator opts in.
return true, NetworkAllow
}
switch mode := engine.effectiveNetworkMode(policy); mode {
case NetworkAllow:
return true, mode
@@ -109,6 +149,19 @@ func (engine *Engine) NetworkHostAllowed(host string) (bool, NetworkMode) {
}
}
// toolNetworkExempt reports whether a request is exempt from the engine-level
// network deny because it is a first-party, in-process network TOOL — one that
// declares SideEffectNetwork (web_search / web_fetch) — and the operator has not
// opted into EnforceToolNetwork. Such tools do not use the sandboxed shell's
// egress; they keep their own SSRF/host safeguards, and NetworkHostAllowed (also
// gated by EnforceToolNetwork) governs them at run time. A SHELL command merely
// classified as network (SideEffectShell) is NOT exempt, so shell egress stays
// blocked under deny. Mirrors the NetworkHostAllowed exemption so the Evaluate
// gate and the per-tool gate never diverge.
func (engine *Engine) toolNetworkExempt(policy Policy, request Request) bool {
return !policy.EnforceToolNetwork && request.SideEffect == SideEffectNetwork
}
// scopeFor returns the scope to validate request paths against. The engine's
// shared scope applies only when the request targets the engine's own
// workspace root; a per-request override root gets an ad-hoc single-root scope
@@ -221,11 +274,17 @@ func (engine *Engine) Evaluate(ctx context.Context, request Request) Decision {
if request.Permission == PermissionDeny {
return deny(request, risk, ViolationDeniedPermission, "", permissionReason(request), false)
}
if policy.EnforceWorkspace && request.WorkspaceRoot != "" {
for _, requested := range requestPaths(request) {
if violation := scope.validate(requested); violation != nil {
return deny(request, risk, violation.Code, violation.Path, violation.Reason, false)
}
// The fine-grained path lists (DenyRead/DenyWrite/AllowRead/AllowWrite) apply
// whenever the sandbox is enforcing, independent of EnforceWorkspace and even
// when there is no workspace root (absolute paths are still resolved and
// matched), so they are honored consistently with the grep/glob exclusion path
// and can't be bypassed by an engine built without a workspace root. The
// workspace boundary itself needs a root, so it is gated on having one. Mode is
// already known to be enforcing here (ModeDisabled returned above).
enforceWorkspace := policy.EnforceWorkspace && request.WorkspaceRoot != ""
for _, requested := range requestPaths(request) {
if violation := validatePathWithPolicy(scope, policy, request.SideEffect, enforceWorkspace, request.WorkspaceRoot, requested); violation != nil {
return deny(request, risk, violation.Code, violation.Path, violation.Reason, false)
}
}
// effectiveNetworkMode collapses an empty-allowlist scoped policy to deny, and
@@ -235,7 +294,7 @@ func (engine *Engine) Evaluate(ctx context.Context, request Request) Decision {
// than running with unrestricted networking. Shared with NetworkHostAllowed so
// the per-tool gate can't diverge from this decision. Allow is unchanged.
netMode := engine.effectiveNetworkMode(policy)
if netMode == NetworkDeny && HasRiskCategory(risk, "network") {
if netMode == NetworkDeny && HasRiskCategory(risk, "network") && !engine.toolNetworkExempt(policy, request) {
return deny(request, risk, ViolationNetwork, "", "network access is blocked by sandbox policy", false)
}
if policy.DenyDestructiveShell && HasRiskCategory(risk, "destructive") {
+75 -2
View File
@@ -9,6 +9,53 @@ import (
"time"
)
// TestEvaluateExemptsNetworkToolsFromDenyByDefault reproduces the web_search
// hard-deny: under NetworkDeny, the engine-level network gate must NOT deny a
// first-party network TOOL (SideEffect=network) by default, so it follows the
// normal permission flow instead of being blocked outright. EnforceToolNetwork
// restores the strict deny, and a network SHELL command stays denied either way.
func TestEvaluateExemptsNetworkToolsFromDenyByDefault(t *testing.T) {
base := Policy{Mode: ModeEnforce, Network: NetworkDeny, MaxAutonomy: AutonomyHigh}
// Default (EnforceToolNetwork off): web_search is NOT network-denied.
engine := NewEngine(EngineOptions{Policy: base})
d := engine.Evaluate(context.Background(), Request{
ToolName: "web_search", SideEffect: SideEffectNetwork, Permission: PermissionPrompt,
})
if d.Action == ActionDeny {
t.Fatalf("web_search must not be network-denied by default, got deny: %s", d.ErrorString())
}
// Granted permission → it actually runs (ActionAllow), not blocked.
allowed := engine.Evaluate(context.Background(), Request{
ToolName: "web_search", SideEffect: SideEffectNetwork, Permission: PermissionPrompt, PermissionGranted: true,
})
if allowed.Action != ActionAllow {
t.Fatalf("granted web_search must be allowed under deny by default, got %q (%s)", allowed.Action, allowed.ErrorString())
}
// EnforceToolNetwork on → strict: the network tool is denied at the gate.
strict := base
strict.EnforceToolNetwork = true
strictEngine := NewEngine(EngineOptions{Policy: strict})
ds := strictEngine.Evaluate(context.Background(), Request{
ToolName: "web_search", SideEffect: SideEffectNetwork, Permission: PermissionPrompt, PermissionGranted: true,
})
if ds.Action != ActionDeny || ds.Violation == nil || ds.Violation.Code != ViolationNetwork {
t.Fatalf("with EnforceToolNetwork, web_search must be network-denied, got %q (%+v)", ds.Action, ds.Violation)
}
// A network SHELL command (SideEffect=shell) stays denied under deny — shell
// egress is never exempt, even with EnforceToolNetwork off.
shell := engine.Evaluate(context.Background(), Request{
ToolName: "bash", SideEffect: SideEffectShell, PermissionGranted: true,
Args: map[string]any{"command": "curl https://evil.test"},
})
if shell.Action != ActionDeny || shell.Violation == nil || shell.Violation.Code != ViolationNetwork {
t.Fatalf("a network shell command must stay denied under deny, got %q (%+v)", shell.Action, shell.Violation)
}
}
func TestEngineEvaluatesReadPromptAndPersistentDecisions(t *testing.T) {
root := t.TempDir()
store, err := NewGrantStore(StoreOptions{
@@ -295,7 +342,11 @@ func TestEngineClassifiesNetworkAndDestructiveShellCommands(t *testing.T) {
}
func TestEngineDeniesNetworkSideEffectWhenPolicyBlocksNetwork(t *testing.T) {
engine := NewEngine(EngineOptions{WorkspaceRoot: t.TempDir(), Policy: DefaultPolicy()})
// EnforceToolNetwork opts the in-process network tools into the policy; by
// default they are exempt (see TestEvaluateExemptsNetworkToolsFromDenyByDefault).
policy := DefaultPolicy()
policy.EnforceToolNetwork = true
engine := NewEngine(EngineOptions{WorkspaceRoot: t.TempDir(), Policy: policy})
decision := engine.Evaluate(context.Background(), Request{
ToolName: "web_fetch",
@@ -736,7 +787,11 @@ func TestEngineNetworkHostAllowed(t *testing.T) {
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
engine := NewEngine(EngineOptions{Policy: tc.policy, Backend: tc.backend})
// This table exercises the ENFORCED tool-network gate; the default
// (exempt) posture is covered by TestEngineNetworkHostAllowedToolNetworkDefaultExempt.
policy := tc.policy
policy.EnforceToolNetwork = true
engine := NewEngine(EngineOptions{Policy: policy, Backend: tc.backend})
allowed, mode := engine.NetworkHostAllowed(tc.host)
if allowed != tc.allowed || mode != tc.mode {
t.Fatalf("NetworkHostAllowed(%q) = (%v, %q), want (%v, %q)", tc.host, allowed, mode, tc.allowed, tc.mode)
@@ -749,3 +804,21 @@ func TestEngineNetworkHostAllowed(t *testing.T) {
t.Fatalf("nil engine NetworkHostAllowed = (%v, %q), want (true, allow)", allowed, mode)
}
}
// TestEngineNetworkHostAllowedToolNetworkDefaultExempt verifies the default
// posture: with EnforceToolNetwork off, the in-process tool gate is exempt from
// the network policy, so even a deny / scoped-unlisted host is allowed (the
// sandboxed-shell egress decision is separate and unaffected).
func TestEngineNetworkHostAllowedToolNetworkDefaultExempt(t *testing.T) {
egressBackend := Backend{Name: BackendSandboxExec, Available: true, Executable: "/usr/bin/sandbox-exec", ScopedEgress: true}
policies := []Policy{
{Mode: ModeEnforce, Network: NetworkDeny},
{Mode: ModeEnforce, Network: NetworkScoped, AllowedDomains: []string{"allowed.test"}},
}
for _, policy := range policies {
engine := NewEngine(EngineOptions{Policy: policy, Backend: egressBackend})
if allowed, _ := engine.NetworkHostAllowed("evil.test"); !allowed {
t.Fatalf("default (EnforceToolNetwork off) must allow the tool host under network=%q", policy.Network)
}
}
}
+261
View File
@@ -0,0 +1,261 @@
package sandbox
import (
"bufio"
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/tls"
"crypto/x509"
"crypto/x509/pkix"
"encoding/pem"
"fmt"
"io"
"math/big"
"net"
"net/http"
"os"
"strings"
"sync"
"time"
)
// MITM TLS interception (opt-in via Policy.InspectTLS). Mirrors
// reference-sandbox-wsl-mitm-code-agent-js/sandbox-manager.js mitmCA/getMitmCA
// and network/proxy.js. The proxy terminates TLS toward the SANDBOXED CLIENT
// using an ephemeral, locally generated CA so it can read the decrypted request
// and apply the SAME allow/deny gate on the real Host — visibility, not new
// authority. The UPSTREAM connection is always validated against real roots
// (never InsecureSkipVerify).
// mitmCA is an ephemeral certificate authority that mints per-host leaf certs.
// It lives only for the proxy's lifetime; only its PUBLIC cert is ever exposed
// (so the sandboxed client can trust the minted leaves).
type mitmCA struct {
cert *x509.Certificate
key *ecdsa.PrivateKey
certDER []byte
mu sync.Mutex
leaves map[string]*tls.Certificate
}
// maxMITMLeafCacheEntries caps the per-host leaf cache so an attacker cannot grow
// it without bound. Combined with minting for the AUTHORIZED host (not arbitrary
// client SNI), this keeps the cache size to roughly the allowlist size.
const maxMITMLeafCacheEntries = 512
func randomSerial() (*big.Int, error) {
serial, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
if err != nil {
return nil, fmt.Errorf("mitm: serial: %w", err)
}
return serial, nil
}
// newMITMCA generates a fresh ECDSA CA. The key never leaves the process.
func newMITMCA() (*mitmCA, error) {
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if err != nil {
return nil, fmt.Errorf("mitm: generate CA key: %w", err)
}
serial, err := randomSerial()
if err != nil {
return nil, err
}
now := time.Now()
tmpl := &x509.Certificate{
SerialNumber: serial,
Subject: pkix.Name{CommonName: "zero sandbox MITM CA", Organization: []string{"zero-sandbox"}},
NotBefore: now.Add(-time.Hour),
NotAfter: now.Add(24 * time.Hour),
KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageCRLSign | x509.KeyUsageDigitalSignature,
BasicConstraintsValid: true,
IsCA: true,
MaxPathLenZero: true,
}
der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
if err != nil {
return nil, fmt.Errorf("mitm: create CA cert: %w", err)
}
cert, err := x509.ParseCertificate(der)
if err != nil {
return nil, fmt.Errorf("mitm: parse CA cert: %w", err)
}
return &mitmCA{cert: cert, key: key, certDER: der, leaves: map[string]*tls.Certificate{}}, nil
}
// caPEM returns the CA's PUBLIC certificate in PEM (safe to write/expose).
func (ca *mitmCA) caPEM() []byte {
return pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: ca.certDER})
}
// leafFor mints (and caches) a leaf certificate for host, signed by the CA, with
// the CA appended so the chain verifies against the CA alone.
func (ca *mitmCA) leafFor(host string) (*tls.Certificate, error) {
host = hostnameOnly(host)
if host == "" {
return nil, fmt.Errorf("mitm: empty host for leaf cert")
}
ca.mu.Lock()
defer ca.mu.Unlock()
if leaf, ok := ca.leaves[host]; ok {
return leaf, nil
}
// Bound the cache: past the cap, mint without caching so memory stays bounded
// (a small per-request cost rather than unbounded growth).
cacheable := len(ca.leaves) < maxMITMLeafCacheEntries
leafKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if err != nil {
return nil, err
}
serial, err := randomSerial()
if err != nil {
return nil, err
}
now := time.Now()
tmpl := &x509.Certificate{
SerialNumber: serial,
Subject: pkix.Name{CommonName: host},
NotBefore: now.Add(-time.Hour),
NotAfter: now.Add(24 * time.Hour),
KeyUsage: x509.KeyUsageDigitalSignature,
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
}
if ip := net.ParseIP(host); ip != nil {
tmpl.IPAddresses = []net.IP{ip}
} else {
tmpl.DNSNames = []string{host}
}
der, err := x509.CreateCertificate(rand.Reader, tmpl, ca.cert, &leafKey.PublicKey, ca.key)
if err != nil {
return nil, fmt.Errorf("mitm: mint leaf for %s: %w", host, err)
}
leaf := &tls.Certificate{Certificate: [][]byte{der, ca.certDER}, PrivateKey: leafKey}
if cacheable {
ca.leaves[host] = leaf
}
return leaf, nil
}
// writeMITMCAFile writes the CA's PUBLIC cert PEM to an owner-only temp file and
// returns its path, so the runner can surface it via ZERO_SANDBOX_CA_CERT.
func writeMITMCAFile(pemBytes []byte) (string, error) {
f, err := os.CreateTemp("", "zero-sandbox-ca-*.pem")
if err != nil {
return "", fmt.Errorf("mitm: create CA file: %w", err)
}
if err := f.Chmod(0o600); err != nil {
_ = f.Close()
_ = os.Remove(f.Name())
return "", err
}
if _, err := f.Write(pemBytes); err != nil {
_ = f.Close()
_ = os.Remove(f.Name())
return "", err
}
if err := f.Close(); err != nil {
_ = os.Remove(f.Name())
return "", err
}
return f.Name(), nil
}
// handleConnectMITM terminates TLS toward the client (minting a leaf for the
// requested host), then for each decrypted request re-checks the real Host
// through the SAME authorize gate and forwards it upstream with full system-root
// validation. A denied host gets a 403 over the decrypted channel and is never
// forwarded.
func (proxy *egressProxy) handleConnectMITM(w http.ResponseWriter, r *http.Request) {
target := r.Host
if target == "" {
target = r.URL.Host
}
if !proxy.authorizeTarget(target, 443) {
proxy.logDecision("deny", "CONNECT", target)
http.Error(w, "scoped egress: host not allowed", http.StatusForbidden)
return
}
hijacker, ok := w.(http.Hijacker)
if !ok {
http.Error(w, "scoped egress: hijack unsupported", http.StatusInternalServerError)
return
}
clientConn, _, err := hijacker.Hijack()
if err != nil {
return
}
defer clientConn.Close()
if _, err := io.WriteString(clientConn, "HTTP/1.1 200 Connection Established\r\n\r\n"); err != nil {
return
}
host := hostnameOnly(target)
tlsConn := tls.Server(clientConn, &tls.Config{
MinVersion: tls.VersionTLS12,
GetCertificate: func(_ *tls.ClientHelloInfo) (*tls.Certificate, error) {
// Mint for the AUTHORIZED CONNECT host, never the arbitrary client SNI:
// honoring SNI would let an attacker churn distinct names to force
// unbounded key generation and cache growth.
return proxy.mitm.leafFor(host)
},
})
if err := tlsConn.Handshake(); err != nil {
return
}
defer tlsConn.Close()
reader := bufio.NewReader(tlsConn)
for {
req, err := http.ReadRequest(reader)
if err != nil {
return // client closed or EOF
}
reqHost := req.Host
if reqHost == "" {
reqHost = host
}
// MITM widens VISIBILITY, never authority: the decrypted Host passes the
// SAME full authorize() gate as the CONNECT check (deny-wins, empty allowlist
// => deny, and an interactive DomainPrompt is honored consistently).
if !proxy.authorizeTarget(reqHost, 443) {
proxy.logDecision("deny", req.Method, reqHost)
writeMITMResponse(tlsConn, http.StatusForbidden, "scoped egress: host not allowed")
return
}
req.URL.Scheme = "https"
req.URL.Host = normalizeConnectTarget(target)
req.RequestURI = ""
resp, err := proxy.mitmTransport.RoundTrip(req)
if err != nil {
proxy.logDecision("deny", req.Method, reqHost)
writeMITMResponse(tlsConn, http.StatusBadGateway, "scoped egress: upstream request failed")
return
}
proxy.logDecision("allow", req.Method, reqHost)
// Re-frame the response as HTTP/1.1 for the decrypted (HTTP/1.1) client
// connection, regardless of how the upstream replied.
resp.Proto, resp.ProtoMajor, resp.ProtoMinor = "HTTP/1.1", 1, 1
writeErr := resp.Write(tlsConn)
_ = resp.Body.Close()
if writeErr != nil || req.Close || resp.Close {
return
}
}
}
// writeMITMResponse writes a minimal plaintext HTTP/1.1 response over the
// decrypted client connection.
func writeMITMResponse(conn net.Conn, code int, msg string) {
body := msg + "\n"
resp := &http.Response{
StatusCode: code,
ProtoMajor: 1,
ProtoMinor: 1,
Header: http.Header{"Content-Type": []string{"text/plain; charset=utf-8"}},
Body: io.NopCloser(strings.NewReader(body)),
ContentLength: int64(len(body)),
}
_ = resp.Write(conn)
}
+158
View File
@@ -0,0 +1,158 @@
package sandbox
import (
"bufio"
"crypto/tls"
"crypto/x509"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
)
func TestEgressProxyPassthroughNoCAByDefault(t *testing.T) {
proxy, err := newEgressProxy(egressOptions{Allowed: []string{"github.com"}})
if err != nil {
t.Fatalf("newEgressProxy: %v", err)
}
defer proxy.Close()
if proxy.mitm != nil {
t.Fatal("InspectTLS off must not create a MITM CA (pure passthrough)")
}
if proxy.CACertPEM() != nil {
t.Fatal("passthrough proxy must expose no CA cert")
}
}
func TestMITMCAMintsChainingLeaf(t *testing.T) {
ca, err := newMITMCA()
if err != nil {
t.Fatalf("newMITMCA: %v", err)
}
leaf, err := ca.leafFor("api.example.com")
if err != nil {
t.Fatalf("leafFor: %v", err)
}
leafCert, err := x509.ParseCertificate(leaf.Certificate[0])
if err != nil {
t.Fatalf("parse leaf: %v", err)
}
roots := x509.NewCertPool()
if !roots.AppendCertsFromPEM(ca.caPEM()) {
t.Fatal("append CA PEM")
}
if _, err := leafCert.Verify(x509.VerifyOptions{
Roots: roots,
DNSName: "api.example.com",
KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
}); err != nil {
t.Fatalf("minted leaf must chain to the CA: %v", err)
}
// Per-host caching.
if leaf2, _ := ca.leafFor("api.example.com"); leaf2 != leaf {
t.Fatal("leafFor must cache one leaf per host")
}
}
// mitmConnect performs a CONNECT through the proxy then opens a TLS client that
// trusts ONLY the MITM CA, returning the decrypted client connection.
func mitmConnect(t *testing.T, proxyAddr, target string, caPEM []byte, serverName string) *tls.Conn {
t.Helper()
conn, err := net.DialTimeout("tcp", proxyAddr, 2*time.Second)
if err != nil {
t.Fatalf("dial proxy: %v", err)
}
// Bound the whole exchange (CONNECT read + TLS handshake + request/response) so
// a regression cannot hang the test indefinitely.
if err := conn.SetDeadline(time.Now().Add(10 * time.Second)); err != nil {
t.Fatalf("set deadline: %v", err)
}
fmt.Fprintf(conn, "CONNECT %s HTTP/1.1\r\nHost: %s\r\n\r\n", target, target)
br := bufio.NewReader(conn)
resp, err := http.ReadResponse(br, &http.Request{Method: "CONNECT"})
if err != nil {
t.Fatalf("read CONNECT response: %v", err)
}
if resp.StatusCode != http.StatusOK {
t.Fatalf("CONNECT status = %d, want 200", resp.StatusCode)
}
if br.Buffered() != 0 {
t.Fatalf("proxy over-sent %d bytes after CONNECT 200", br.Buffered())
}
pool := x509.NewCertPool()
if !pool.AppendCertsFromPEM(caPEM) {
t.Fatal("append MITM CA PEM")
}
tlsConn := tls.Client(conn, &tls.Config{RootCAs: pool, ServerName: serverName})
if err := tlsConn.Handshake(); err != nil {
t.Fatalf("client TLS handshake trusting the MITM CA: %v", err)
}
return tlsConn
}
func TestEgressProxyMITMForwardsAllowlistedHost(t *testing.T) {
upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
fmt.Fprint(w, "upstream-ok")
}))
defer upstream.Close()
host := mustHost(t, upstream.URL)
target := mustHostPort(t, upstream.URL)
// Trust the upstream's real cert explicitly (NOT InsecureSkipVerify): proves
// the MITM still validates the upstream chain.
upstreamRoots := x509.NewCertPool()
upstreamRoots.AddCert(upstream.Certificate())
proxy, err := newEgressProxy(egressOptions{Allowed: []string{host}, InspectTLS: true, upstreamRoots: upstreamRoots})
if err != nil {
t.Fatalf("newEgressProxy: %v", err)
}
defer proxy.Close()
if proxy.CACertPEM() == nil {
t.Fatal("InspectTLS must generate a CA")
}
tlsConn := mitmConnect(t, proxy.Addr(), target, proxy.CACertPEM(), host)
defer tlsConn.Close()
fmt.Fprintf(tlsConn, "GET / HTTP/1.1\r\nHost: %s\r\nConnection: close\r\n\r\n", host)
resp, err := http.ReadResponse(bufio.NewReader(tlsConn), &http.Request{Method: http.MethodGet})
if err != nil {
t.Fatalf("read decrypted response: %v", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
t.Fatalf("allowlisted forwarded request status = %d, want 200", resp.StatusCode)
}
body, err := io.ReadAll(resp.Body)
if err != nil {
t.Fatalf("read decrypted body: %v", err)
}
if !strings.Contains(string(body), "upstream-ok") {
t.Fatalf("MITM did not forward the allowlisted request, got %q", string(body))
}
}
func TestEgressProxyMITMBlocksDeniedDecryptedHost(t *testing.T) {
// CONNECT to an allowed host, but send a Host header that is NOT allowlisted —
// the MITM re-checks the DECRYPTED Host and blocks it before forwarding.
proxy, err := newEgressProxy(egressOptions{Allowed: []string{"127.0.0.1"}, InspectTLS: true})
if err != nil {
t.Fatalf("newEgressProxy: %v", err)
}
defer proxy.Close()
tlsConn := mitmConnect(t, proxy.Addr(), "127.0.0.1:443", proxy.CACertPEM(), "127.0.0.1")
defer tlsConn.Close()
fmt.Fprint(tlsConn, "GET / HTTP/1.1\r\nHost: blocked.example\r\nConnection: close\r\n\r\n")
resp, err := http.ReadResponse(bufio.NewReader(tlsConn), nil)
if err != nil {
t.Fatalf("read decrypted response: %v", err)
}
if resp.StatusCode != http.StatusForbidden {
t.Fatalf("denied decrypted Host status = %d, want 403", resp.StatusCode)
}
}
+374
View File
@@ -0,0 +1,374 @@
package sandbox
import (
"os"
"path/filepath"
"strings"
)
// This file implements the fine-grained AllowRead/DenyRead/AllowWrite/DenyWrite
// path lists (Policy fields). They layer ON TOP of the workspace + Scope guards
// and never bypass them: AllowRead only re-includes inside a DenyRead carve-out,
// AllowWrite is consulted only after the workspace/Scope guard already denied a
// write, and every match is symlink-resolved so a symlink prefix cannot evade a
// deny or sneak past an allow. All lists default empty, so an unconfigured policy
// behaves exactly as before.
// resolvePolicyPath home-expands, makes absolute, and symlink-resolves a single
// policy path entry. ok is false for a blank entry or one that does not exist
// (EvalSymlinks requires existence) so a bogus entry is dropped — a non-existent
// deny protects nothing and a non-existent allow grants nothing.
func resolvePolicyPath(entry string) (string, bool) {
trimmed := strings.TrimSpace(entry)
if trimmed == "" {
return "", false
}
if trimmed == "~" || strings.HasPrefix(trimmed, "~/") || strings.HasPrefix(trimmed, "~"+string(filepath.Separator)) {
home, err := os.UserHomeDir()
if err != nil {
return "", false
}
trimmed = filepath.Join(home, strings.TrimPrefix(strings.TrimPrefix(trimmed[1:], "/"), string(filepath.Separator)))
}
absolute, err := filepath.Abs(trimmed)
if err != nil {
return "", false
}
resolved, err := filepath.EvalSymlinks(absolute)
if err != nil {
return "", false
}
return resolved, true
}
// resolvePolicyPaths resolves and de-duplicates a list of policy path entries,
// dropping blanks and non-existent entries. Files and directories are both kept
// (a DenyRead/DenyWrite entry may target a single sensitive file).
func resolvePolicyPaths(entries []string) []string {
if len(entries) == 0 {
return nil
}
seen := make(map[string]struct{}, len(entries))
out := make([]string, 0, len(entries))
for _, entry := range entries {
resolved, ok := resolvePolicyPath(entry)
if !ok {
continue
}
if _, dup := seen[resolved]; dup {
continue
}
seen[resolved] = struct{}{}
out = append(out, resolved)
}
return out
}
// resolveWriteRootPaths is resolvePolicyPaths restricted to existing directories
// that are not the filesystem root — the only valid targets for an OS write bind
// and for an AllowWrite grant root.
func resolveWriteRootPaths(entries []string) []string {
resolved := resolvePolicyPaths(entries)
if len(resolved) == 0 {
return nil
}
out := make([]string, 0, len(resolved))
for _, path := range resolved {
info, err := os.Stat(path)
if err != nil || !info.IsDir() {
continue
}
if filepath.Dir(path) == path {
continue // refuse the filesystem root as a write root
}
out = append(out, path)
}
return out
}
// pathUnderPolicyRoot reports whether requestedPath lies within root. A relative
// requestedPath is resolved against workspaceRoot; the portion of an absolute
// path outside root is symlink-resolved (via NormalizePrefixForRoot) so a
// symlink prefix cannot evade the match. root must be an already-resolved
// absolute path.
func pathUnderPolicyRoot(requestedPath, root, workspaceRoot string) bool {
if root == "" {
return false
}
abs := requestedPath
if !filepath.IsAbs(abs) {
if workspaceRoot == "" {
return false
}
abs = filepath.Join(workspaceRoot, abs)
}
normalized := NormalizePrefixForRoot(abs, root)
return pathWithinRoot(root, normalized)
}
// readDenied reports whether path is excluded by the DenyRead list with no
// more-specific AllowRead re-inclusion. "More specific" means an AllowRead entry
// nested inside the matched DenyRead entry — that subtree is read back in while
// the rest of the denied tree stays blocked. It resolves the policy entries on
// each call, so it suits a one-shot check (the Evaluate gate). A search walk that
// checks many paths should use ReadExclusions, which resolves the entries once.
func readDenied(policy Policy, workspaceRoot, path string) bool {
return readDeniedResolved(workspaceRoot, resolvePolicyPaths(policy.DenyRead), resolvePolicyPaths(policy.AllowRead), path)
}
// readDeniedResolved is readDenied operating on ALREADY-resolved deny/allow roots,
// so a caller that resolved the policy entries once can reuse them across many
// path checks instead of re-running Abs/EvalSymlinks per path.
func readDeniedResolved(workspaceRoot string, denyRoots, allowRoots []string, path string) bool {
if len(denyRoots) == 0 {
return false
}
for _, deny := range denyRoots {
if !pathUnderPolicyRoot(path, deny, workspaceRoot) {
continue
}
reincluded := false
for _, allow := range allowRoots {
// The allow entry must sit inside the deny entry to be "more specific",
// and the path must fall under that allow entry.
if pathWithinRoot(deny, allow) && pathUnderPolicyRoot(path, allow, workspaceRoot) {
reincluded = true
break
}
}
if !reincluded {
return true
}
}
return false
}
// ReadExclusions holds the resolved DenyRead/AllowRead roots for a policy so a
// search walk resolves each policy entry ONCE (Abs/EvalSymlinks) and reuses the
// result across every visited path, rather than re-resolving per path. Build it
// with Engine.ReadExclusions and reuse it for the whole grep/glob walk.
type ReadExclusions struct {
workspaceRoot string
denyRoots []string
allowRoots []string
}
// Active reports whether any DenyRead root is configured. When false the
// exclusions are a no-op and the search behaves exactly as before.
func (rx *ReadExclusions) Active() bool {
return rx != nil && len(rx.denyRoots) > 0
}
// PathExcluded reports whether reading path is excluded by DenyRead, honoring a
// more-specific AllowRead re-inclusion. It is the per-file predicate for a walk.
func (rx *ReadExclusions) PathExcluded(path string) bool {
if !rx.Active() {
return false
}
return readDeniedResolved(rx.workspaceRoot, rx.denyRoots, rx.allowRoots, path)
}
// DirExcluded reports whether a directory subtree can be skipped wholesale during
// a walk: it is read-denied AND contains no nested AllowRead root (descending is
// required to reach a re-included subtree). When it returns false on a denied dir
// because of a nested allow, PathExcluded still filters the denied siblings.
func (rx *ReadExclusions) DirExcluded(path string) bool {
if !rx.Active() {
return false
}
if !readDeniedResolved(rx.workspaceRoot, rx.denyRoots, rx.allowRoots, path) {
return false
}
abs := path
if !filepath.IsAbs(abs) && rx.workspaceRoot != "" {
abs = filepath.Join(rx.workspaceRoot, path)
}
// allowRoots are symlink-resolved (resolvePolicyPaths), so resolve abs the
// same way before the prefix comparison — otherwise a dir reached THROUGH a
// symlink would fail to match a nested AllowRead root and be wrongly skipped,
// dropping a re-included subtree. Best-effort: keep abs if it can't resolve
// (e.g. a non-existent path), matching the deny check's tolerant behavior.
if resolved, err := filepath.EvalSymlinks(abs); err == nil {
abs = resolved
}
return !hasNestedAllowReadResolved(rx.allowRoots, abs)
}
// allowWriteScope builds an ad-hoc Scope from the resolved AllowWrite roots so a
// write to an AllowWrite path is validated with the SAME symlink-traversal logic
// the workspace Scope uses. Returns nil when there are no usable AllowWrite roots.
func allowWriteScope(policy Policy) *Scope {
roots := resolveWriteRootPaths(policy.AllowWrite)
if len(roots) == 0 {
return nil
}
return &Scope{workspaceRoot: roots[0], extraRoots: roots[1:]}
}
// validateWritePath enforces the write precedence: DenyWrite wins, then (when
// workspace enforcement is on) a workspace/Scope-writable path is allowed, then
// an absolute path under an AllowWrite root is allowed, otherwise the base
// workspace/Scope violation stands. The DenyWrite list applies regardless of
// enforceWorkspace; the workspace boundary itself applies only when
// enforceWorkspace. It never bypasses the symlink/out-of-workspace guards.
func validateWritePath(scope *Scope, policy Policy, enforceWorkspace bool, workspaceRoot, path string) *pathViolation {
// DenyWrite wins regardless of workspace enforcement.
for _, deny := range resolvePolicyPaths(policy.DenyWrite) {
if pathUnderPolicyRoot(path, deny, workspaceRoot) {
return &pathViolation{
Code: ViolationPolicyDenied,
Path: path,
Reason: path + " is excluded by the sandbox DenyWrite policy",
}
}
}
if !enforceWorkspace {
// Workspace confinement is off: only the explicit DenyWrite list restricts.
return nil
}
base := scope.validate(path)
if base == nil {
return nil // writable under the workspace / Scope guard
}
// AllowWrite only extends ABSOLUTE paths: a relative path is inherently
// workspace-relative and already resolved by the base guard above.
if filepath.IsAbs(path) {
if allow := allowWriteScope(policy); allow != nil && allow.validate(path) == nil {
return nil
}
}
return base
}
// validatePathWithPolicy is the single entry point the engine uses to validate a
// request path. The fine-grained read/write lists (DenyRead/DenyWrite, with
// AllowRead/AllowWrite) apply whenever the sandbox is enforcing, INDEPENDENT of
// enforceWorkspace, so they match the grep/glob path that also honors DenyRead
// directly. The workspace boundary (scope.validate) applies only when
// enforceWorkspace. Behavior is unchanged when the lists are empty.
func validatePathWithPolicy(scope *Scope, policy Policy, sideEffect SideEffect, enforceWorkspace bool, workspaceRoot, path string) *pathViolation {
// A relative path cannot be anchored without a workspace root, so it cannot be
// checked against the (absolute) path lists or workspace boundary. Fail closed
// when there is anything to enforce; otherwise it is a no-op (unchanged from the
// pre-path-list behavior, where an empty workspace root skipped validation).
if workspaceRoot == "" && !filepath.IsAbs(path) {
if enforceWorkspace || policyHasPathLists(policy) {
return &pathViolation{
Code: ViolationOutsideWorkspace,
Path: path,
Reason: path + " cannot be validated without a workspace root",
}
}
return nil
}
switch sideEffect {
case SideEffectRead:
if readDenied(policy, workspaceRoot, path) {
return &pathViolation{
Code: ViolationPolicyDenied,
Path: path,
Reason: path + " is excluded by the sandbox DenyRead policy",
}
}
if enforceWorkspace {
return scope.validate(path)
}
return nil
case SideEffectWrite, SideEffectOutOfWorkspace:
return validateWritePath(scope, policy, enforceWorkspace, workspaceRoot, path)
default:
if enforceWorkspace {
return scope.validate(path)
}
return nil
}
}
// hasNestedAllowReadResolved reports whether any already-resolved AllowRead root
// sits strictly inside dir (an already-resolved absolute path). When true, a
// read-denied dir must still be descended during a walk so the re-included
// subtree is reachable.
func hasNestedAllowReadResolved(allowRoots []string, dir string) bool {
for _, allow := range allowRoots {
if allow != dir && pathWithinRoot(dir, allow) {
return true
}
}
return false
}
// workspaceRelGlob returns target as a clean, slash-separated path relative to
// workspaceRoot, or ok=false when target is the root itself or lies outside it
// (a workspace-rooted search never reaches such a path, so no glob is needed).
func workspaceRelGlob(workspaceRoot, target string) (string, bool) {
rel, err := filepath.Rel(workspaceRoot, target)
if err != nil {
return "", false
}
if rel == "." || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) || filepath.IsAbs(rel) {
return "", false
}
return filepath.ToSlash(rel), true
}
// ReadExclusionGlobs returns ripgrep-style --glob exclusion args for the policy's
// DenyRead subtrees that fall inside the scope's workspace root, so a
// ripgrep-based search never descends into a read-denied subtree. For each such
// entry it emits `--glob`, `!<rel>` and `--glob`, `!<rel>/**`. Mirrors the
// read-subtree exclusion globs used by comparable executor sandboxes.
//
// The projection is exclusions-only: a positive ripgrep glob would switch the
// search into whitelist mode and restrict it to only matching files, so AllowRead
// re-inclusion is NOT expressed here. The Go-native grep/glob tools honor
// AllowRead precisely via the cached predicate (Engine.ReadExclusions); this
// function is the coarser ripgrep-format export for an external rg-based
// consumer. Empty when DenyRead is unset (the default), so search behavior is
// unchanged.
func ReadExclusionGlobs(policy Policy, scope *Scope) []string {
denyRoots := resolvePolicyPaths(policy.DenyRead)
if len(denyRoots) == 0 || scope == nil {
return nil
}
workspaceRoot := scope.WorkspaceRoot()
if workspaceRoot == "" {
return nil
}
var globs []string
for _, deny := range denyRoots {
rel, ok := workspaceRelGlob(workspaceRoot, deny)
if !ok {
continue
}
globs = append(globs, "--glob", "!"+rel, "--glob", "!"+rel+"/**")
}
return globs
}
// policyHasPathLists reports whether any fine-grained path list has an
// ENFORCEABLE entry. It resolves the lists (matching how the rest of this file
// normalizes them) rather than counting raw config, so a typo or non-existent
// entry — which resolution drops — doesn't spuriously fail-close relative
// requests when there is no workspace root.
func policyHasPathLists(policy Policy) bool {
return len(resolvePolicyPaths(policy.DenyRead)) > 0 ||
len(resolvePolicyPaths(policy.AllowRead)) > 0 ||
len(resolvePolicyPaths(policy.DenyWrite)) > 0 ||
len(resolveWriteRootPaths(policy.AllowWrite)) > 0
}
// dedupeStrings returns xs with duplicates removed, preserving first-seen order.
func dedupeStrings(xs []string) []string {
if len(xs) <= 1 {
return xs
}
seen := make(map[string]struct{}, len(xs))
out := make([]string, 0, len(xs))
for _, x := range xs {
if _, dup := seen[x]; dup {
continue
}
seen[x] = struct{}{}
out = append(out, x)
}
return out
}
+426
View File
@@ -0,0 +1,426 @@
package sandbox
import (
"context"
"os"
"path/filepath"
"strings"
"testing"
)
// resolvedTempDir returns a symlink-resolved temp dir so paths built under it
// compare equal to the EvalSymlinks-normalized forms the path lists produce
// (macOS /var -> /private/var would otherwise diverge).
func resolvedTempDir(t *testing.T) string {
t.Helper()
resolved, err := filepath.EvalSymlinks(t.TempDir())
if err != nil {
t.Fatalf("EvalSymlinks(tempdir): %v", err)
}
return resolved
}
func mkdir(t *testing.T, path string) string {
t.Helper()
if err := os.MkdirAll(path, 0o755); err != nil {
t.Fatalf("MkdirAll(%q): %v", path, err)
}
return path
}
func TestReadDenyAllowPrecedence(t *testing.T) {
ws := resolvedTempDir(t)
mkdir(t, filepath.Join(ws, "src"))
secret := mkdir(t, filepath.Join(ws, "secret"))
public := mkdir(t, filepath.Join(ws, "secret", "public"))
cases := []struct {
name string
policy Policy
path string
wantDenied bool
}{
{
name: "no lists: nothing denied",
policy: Policy{},
path: filepath.Join(secret, "data"),
wantDenied: false,
},
{
name: "denyread blocks subtree",
policy: Policy{DenyRead: []string{secret}},
path: filepath.Join(secret, "data"),
wantDenied: true,
},
{
name: "denyread does not touch siblings",
policy: Policy{DenyRead: []string{secret}},
path: filepath.Join(ws, "src", "main.go"),
wantDenied: false,
},
{
name: "more-specific allowread re-includes",
policy: Policy{DenyRead: []string{secret}, AllowRead: []string{public}},
path: filepath.Join(public, "ok.txt"),
wantDenied: false,
},
{
name: "allowread only re-includes its own subtree",
policy: Policy{DenyRead: []string{secret}, AllowRead: []string{public}},
path: filepath.Join(secret, "other.txt"),
wantDenied: true,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if got := readDenied(tc.policy, ws, tc.path); got != tc.wantDenied {
t.Fatalf("readDenied(%q) = %v, want %v", tc.path, got, tc.wantDenied)
}
})
}
}
func TestWritePrecedenceMatrix(t *testing.T) {
ws := resolvedTempDir(t)
mkdir(t, filepath.Join(ws, "src"))
wsSecret := mkdir(t, filepath.Join(ws, "secret"))
ext := resolvedTempDir(t)
mkdir(t, filepath.Join(ext, "build"))
extProtected := mkdir(t, filepath.Join(ext, "build", "protected"))
outside := resolvedTempDir(t) // never allowed
scope, err := NewScope(ws, nil)
if err != nil {
t.Fatalf("NewScope: %v", err)
}
cases := []struct {
name string
policy Policy
path string
wantAllow bool
}{
{
name: "workspace path writable by default",
policy: Policy{},
path: filepath.Join(ws, "src", "main.go"),
wantAllow: true,
},
{
name: "outside path denied by default",
policy: Policy{},
path: filepath.Join(outside, "x"),
wantAllow: false,
},
{
name: "allowwrite extends to external root",
policy: Policy{AllowWrite: []string{filepath.Join(ext, "build")}},
path: filepath.Join(ext, "build", "out.o"),
wantAllow: true,
},
{
name: "denywrite wins over workspace",
policy: Policy{DenyWrite: []string{wsSecret}},
path: filepath.Join(wsSecret, "creds"),
wantAllow: false,
},
{
name: "denywrite wins over allowwrite",
policy: Policy{AllowWrite: []string{filepath.Join(ext, "build")}, DenyWrite: []string{extProtected}},
path: filepath.Join(extProtected, "x"),
wantAllow: false,
},
{
name: "allowwrite does not allow unlisted external path",
policy: Policy{AllowWrite: []string{filepath.Join(ext, "build")}},
path: filepath.Join(outside, "x"),
wantAllow: false,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
// This matrix exercises the workspace-enforced write precedence.
violation := validateWritePath(scope, tc.policy, true, ws, tc.path)
gotAllow := violation == nil
if gotAllow != tc.wantAllow {
t.Fatalf("validateWritePath(%q) allow=%v (violation=%v), want allow=%v", tc.path, gotAllow, violation, tc.wantAllow)
}
})
}
}
func TestReadExclusionGlobs(t *testing.T) {
ws := resolvedTempDir(t)
secret := mkdir(t, filepath.Join(ws, "secret"))
scope, err := NewScope(ws, nil)
if err != nil {
t.Fatalf("NewScope: %v", err)
}
// Empty DenyRead → no globs (default behavior unchanged).
if globs := ReadExclusionGlobs(Policy{}, scope); globs != nil {
t.Fatalf("empty DenyRead must yield no globs, got %v", globs)
}
// DenyRead inside the workspace → ripgrep exclusion globs for the rel path.
got := strings.Join(ReadExclusionGlobs(Policy{DenyRead: []string{secret}}, scope), " ")
want := "--glob !secret --glob !secret/**"
if got != want {
t.Fatalf("ReadExclusionGlobs = %q, want %q", got, want)
}
// DenyRead OUTSIDE the workspace → skipped (a rooted search never reaches it).
outside := resolvedTempDir(t)
if globs := ReadExclusionGlobs(Policy{DenyRead: []string{outside}}, scope); globs != nil {
t.Fatalf("out-of-workspace DenyRead must yield no globs, got %v", globs)
}
}
// TestReadDirExcludedDescendsForNestedAllow verifies a read-denied dir is NOT
// skipped wholesale when it contains a nested AllowRead root, so the walk can
// still reach the re-included subtree.
func TestReadDirExcludedDescendsForNestedAllow(t *testing.T) {
ws := resolvedTempDir(t)
secret := mkdir(t, filepath.Join(ws, "secret"))
public := mkdir(t, filepath.Join(ws, "secret", "public"))
engine := NewEngine(EngineOptions{
WorkspaceRoot: ws,
Policy: Policy{Mode: ModeEnforce, EnforceWorkspace: true, DenyRead: []string{secret}, AllowRead: []string{public}},
})
// The denied dir must be descended (nested allow), but the denied dir's own
// files are still excluded per-file.
rx := engine.ReadExclusions()
if rx.DirExcluded(secret) {
t.Fatalf("a denied dir with a nested AllowRead must not be skipped wholesale")
}
if !rx.PathExcluded(filepath.Join(secret, "creds")) {
t.Fatalf("a denied file outside the re-included subtree must be excluded")
}
if rx.PathExcluded(filepath.Join(public, "ok")) {
t.Fatalf("a re-included AllowRead file must not be excluded")
}
}
// TestDirExcludedResolvesSymlinkPrefix verifies DirExcluded resolves a symlink
// in the path prefix before checking for a nested AllowRead root, so a denied dir
// reached through a symlink is NOT skipped wholesale when it contains a
// re-included subtree (the resolved allowRoots would otherwise never match).
func TestDirExcludedResolvesSymlinkPrefix(t *testing.T) {
ws := resolvedTempDir(t)
real := mkdir(t, filepath.Join(ws, "real"))
secret := mkdir(t, filepath.Join(real, "secret"))
public := mkdir(t, filepath.Join(secret, "public"))
link := filepath.Join(ws, "link")
if err := os.Symlink(real, link); err != nil {
t.Skipf("symlinks unavailable on this platform: %v", err)
}
engine := NewEngine(EngineOptions{
WorkspaceRoot: ws,
Policy: Policy{Mode: ModeEnforce, EnforceWorkspace: true, DenyRead: []string{secret}, AllowRead: []string{public}},
})
rx := engine.ReadExclusions()
// secret reached THROUGH the symlink still contains the re-included AllowRead
// subtree once the prefix is resolved, so it must not be skipped wholesale.
if rx.DirExcluded(filepath.Join(link, "secret")) {
t.Fatalf("DirExcluded must resolve the symlink prefix and detect the nested AllowRead (must not skip)")
}
}
// TestPathListsApplyWithoutWorkspaceEnforcement verifies the explicit path lists
// (DenyWrite/DenyRead) are honored even when EnforceWorkspace is off, while the
// workspace boundary itself is not enforced in that mode.
func TestPathListsApplyWithoutWorkspaceEnforcement(t *testing.T) {
ws := resolvedTempDir(t)
secret := mkdir(t, filepath.Join(ws, "secret"))
outside := resolvedTempDir(t)
scope, err := NewScope(ws, nil)
if err != nil {
t.Fatalf("NewScope: %v", err)
}
// DenyWrite still blocks even with workspace enforcement off.
denyPolicy := Policy{DenyWrite: []string{secret}}
if v := validateWritePath(scope, denyPolicy, false, ws, filepath.Join(secret, "creds")); v == nil {
t.Fatal("DenyWrite must block even when EnforceWorkspace is off")
}
// With enforcement off and no DenyWrite match, an out-of-workspace write is
// NOT blocked (the workspace boundary is intentionally not applied).
if v := validateWritePath(scope, denyPolicy, false, ws, filepath.Join(outside, "x")); v != nil {
t.Fatalf("with EnforceWorkspace off, a non-denied path must be allowed, got %v", v)
}
// DenyRead still blocks reads even with workspace enforcement off.
if v := validatePathWithPolicy(scope, Policy{DenyRead: []string{secret}}, SideEffectRead, false, ws, filepath.Join(secret, "x")); v == nil {
t.Fatal("DenyRead must block reads even when EnforceWorkspace is off")
}
}
// TestEvaluatePathListsWithoutWorkspaceRoot verifies the path lists are enforced
// even on an engine built WITHOUT a workspace root: absolute paths are matched
// against DenyWrite, and an unanchorable relative path fails closed.
func TestEvaluatePathListsWithoutWorkspaceRoot(t *testing.T) {
secret := resolvedTempDir(t)
other := resolvedTempDir(t)
engine := NewEngine(EngineOptions{
Policy: Policy{Mode: ModeEnforce, MaxAutonomy: AutonomyHigh, DenyWrite: []string{secret}},
})
denied := engine.Evaluate(context.Background(), Request{
ToolName: "write_file", SideEffect: SideEffectWrite,
Args: map[string]any{"path": filepath.Join(secret, "x")},
})
if denied.Action != ActionDeny {
t.Fatalf("absolute DenyWrite path must be denied without a workspace root, got %q", denied.Action)
}
ok := engine.Evaluate(context.Background(), Request{
ToolName: "write_file", SideEffect: SideEffectWrite,
Args: map[string]any{"path": filepath.Join(other, "x")},
})
if ok.Action == ActionDeny {
t.Fatalf("a non-denied absolute path must not be denied without a workspace root, got deny: %s", ok.ErrorString())
}
rel := engine.Evaluate(context.Background(), Request{
ToolName: "write_file", SideEffect: SideEffectWrite,
Args: map[string]any{"path": "x"},
})
if rel.Action != ActionDeny {
t.Fatalf("an unanchorable relative path must fail closed when path lists exist, got %q", rel.Action)
}
}
// TestReadExclusionsInactiveWhenDisabled verifies a disabled policy filters
// nothing — ReadExclusions/ReadExclusionGlobs are inert under ModeDisabled even
// when DenyRead is configured.
func TestReadExclusionsInactiveWhenDisabled(t *testing.T) {
ws := resolvedTempDir(t)
secret := mkdir(t, filepath.Join(ws, "secret"))
engine := NewEngine(EngineOptions{
WorkspaceRoot: ws,
Policy: Policy{Mode: ModeDisabled, DenyRead: []string{secret}},
})
if rx := engine.ReadExclusions(); rx.Active() {
t.Fatal("ReadExclusions must be inactive under a disabled policy")
}
if globs := engine.ReadExclusionGlobs(); globs != nil {
t.Fatalf("ReadExclusionGlobs must be empty under a disabled policy, got %v", globs)
}
}
// TestEvaluateAppliesReadDeny verifies the lists are wired into Evaluate: a read
// under DenyRead is denied, while a sibling read is not.
func TestEvaluateAppliesReadDeny(t *testing.T) {
ws := resolvedTempDir(t)
mkdir(t, filepath.Join(ws, "src"))
secret := mkdir(t, filepath.Join(ws, "secret"))
engine := NewEngine(EngineOptions{
WorkspaceRoot: ws,
Policy: Policy{
Mode: ModeEnforce,
EnforceWorkspace: true,
MaxAutonomy: AutonomyHigh,
DenyRead: []string{secret},
},
})
denied := engine.Evaluate(context.Background(), Request{
ToolName: "read_file",
SideEffect: SideEffectRead,
Args: map[string]any{"path": filepath.Join(secret, "creds.txt")},
})
if denied.Action != ActionDeny {
t.Fatalf("read under DenyRead = %q, want deny", denied.Action)
}
ok := engine.Evaluate(context.Background(), Request{
ToolName: "read_file",
SideEffect: SideEffectRead,
Args: map[string]any{"path": filepath.Join(ws, "src", "main.go")},
})
if ok.Action == ActionDeny {
t.Fatalf("read of a non-denied sibling must not be denied, got %q (%s)", ok.Action, ok.ErrorString())
}
}
// TestEvaluateAppliesWriteAllow verifies an external write is denied by default
// but permitted once the path is on AllowWrite.
func TestEvaluateAppliesWriteAllow(t *testing.T) {
ws := resolvedTempDir(t)
ext := resolvedTempDir(t)
build := mkdir(t, filepath.Join(ext, "build"))
base := Policy{Mode: ModeEnforce, EnforceWorkspace: true, MaxAutonomy: AutonomyHigh}
target := filepath.Join(build, "out.o")
denyEngine := NewEngine(EngineOptions{WorkspaceRoot: ws, Policy: base})
denied := denyEngine.Evaluate(context.Background(), Request{
ToolName: "write_file",
SideEffect: SideEffectWrite,
Args: map[string]any{"path": target},
})
if denied.Action != ActionDeny {
t.Fatalf("external write without AllowWrite = %q, want deny", denied.Action)
}
allowPolicy := base
allowPolicy.AllowWrite = []string{build}
allowEngine := NewEngine(EngineOptions{WorkspaceRoot: ws, Policy: allowPolicy})
allowed := allowEngine.Evaluate(context.Background(), Request{
ToolName: "write_file",
SideEffect: SideEffectWrite,
Args: map[string]any{"path": target},
})
if allowed.Action == ActionDeny {
t.Fatalf("external write under AllowWrite must not be denied, got deny: %s", allowed.ErrorString())
}
}
// TestWriteRootsReflectAllowWrite verifies AllowWrite roots flow into the OS
// backend write binds (so a sandboxed shell can write there).
func TestWriteRootsReflectAllowWrite(t *testing.T) {
ws := resolvedTempDir(t)
ext := resolvedTempDir(t)
build := mkdir(t, filepath.Join(ext, "build"))
engine := NewEngine(EngineOptions{
WorkspaceRoot: ws,
Policy: Policy{Mode: ModeEnforce, EnforceWorkspace: true, AllowWrite: []string{build}},
})
roots := engine.writeRoots(ws)
found := false
for _, r := range roots {
if r == build {
found = true
}
}
if !found {
t.Fatalf("writeRoots must include AllowWrite root %q, got %v", build, roots)
}
}
// TestSandboxExecProfileEmitsDenyWriteRule verifies a DenyWrite directory becomes
// an explicit seatbelt deny clause AFTER the write-allow (last-match-wins).
func TestSandboxExecProfileEmitsDenyWriteRule(t *testing.T) {
ws := resolvedTempDir(t)
secret := mkdir(t, filepath.Join(ws, "secret"))
policy := Policy{Mode: ModeEnforce, EnforceWorkspace: true, DenyWrite: []string{secret}}
profile := sandboxExecProfile([]string{ws}, policy, "", "", "")
// The path is escaped the same way the profile escapes it (e.g. Windows
// backslashes are doubled), so the assertion holds on every platform.
wantDeny := `(deny file-write* (subpath "` + sandboxProfileString(secret) + `"))`
if !strings.Contains(profile, wantDeny) {
t.Fatalf("profile missing DenyWrite rule %q:\n%s", wantDeny, profile)
}
// The deny clause must appear AFTER the write-allow so seatbelt's
// last-match-wins lets it override the bind.
allowIdx := strings.Index(profile, "(allow file-write*")
denyIdx := strings.Index(profile, wantDeny)
if allowIdx < 0 || denyIdx < 0 || denyIdx < allowIdx {
t.Fatalf("DenyWrite rule must follow the write-allow (allow@%d deny@%d):\n%s", allowIdx, denyIdx, profile)
}
}
+90
View File
@@ -0,0 +1,90 @@
package sandbox
import "testing"
func TestIsAlreadySandboxed(t *testing.T) {
// A lone ZERO_SANDBOXED=1 (no corroborating backend marker) must NOT count as
// sandboxed — stronger provenance than a single ambient flag.
t.Setenv(EnvSandboxBackend, "")
t.Setenv(EnvSandboxed, "1")
if IsAlreadySandboxed() {
t.Fatalf("IsAlreadySandboxed must be false when only %s=1 is set without %s", EnvSandboxed, EnvSandboxBackend)
}
// The backend marker alone (no ZERO_SANDBOXED=1) must not count either.
t.Setenv(EnvSandboxBackend, string(BackendBubblewrap))
t.Setenv(EnvSandboxed, "")
if IsAlreadySandboxed() {
t.Fatalf("IsAlreadySandboxed must be false when only %s is set", EnvSandboxBackend)
}
// Both correlated markers present (as sandboxEnvironment sets them) → sandboxed.
t.Setenv(EnvSandboxed, "1")
if !IsAlreadySandboxed() {
t.Fatalf("IsAlreadySandboxed must be true when both %s=1 and %s are set", EnvSandboxed, EnvSandboxBackend)
}
t.Setenv(EnvSandboxed, "0")
if IsAlreadySandboxed() {
t.Fatalf("IsAlreadySandboxed must be false when %s=0", EnvSandboxed)
}
}
func TestSandboxEnvironmentMarksSandboxed(t *testing.T) {
env := sandboxEnvironment(DefaultPolicy(), BackendBubblewrap, "/home/agent")
// Both correlated markers must be set: IsAlreadySandboxed requires BOTH, so a
// regression that drops either one would silently break re-entrancy detection.
wantSandboxed := EnvSandboxed + "=1"
wantBackend := EnvSandboxBackend + "=" + string(BackendBubblewrap)
var gotSandboxed, gotBackend bool
for _, entry := range env {
switch entry {
case wantSandboxed:
gotSandboxed = true
case wantBackend:
gotBackend = true
}
}
if !gotSandboxed || !gotBackend {
t.Fatalf("sandboxEnvironment must set %q and %q so a wrapped command is detectable; got %#v", wantSandboxed, wantBackend, env)
}
}
func TestBuildCommandPlanWrapsWhenNotAlreadySandboxed(t *testing.T) {
// Ensure neither re-entrancy marker is set so we are NOT seen as sandboxed.
t.Setenv(EnvSandboxed, "")
t.Setenv(EnvSandboxBackend, "")
root := t.TempDir()
engine := NewEngine(EngineOptions{
WorkspaceRoot: root,
Policy: DefaultPolicy(),
Backend: Backend{Name: BackendBubblewrap, Available: true, Executable: "/usr/bin/bwrap"},
})
plan, err := engine.BuildCommandPlan(CommandSpec{Name: "/bin/sh", Args: []string{"-c", "pwd"}, Dir: root})
if err != nil {
t.Fatalf("BuildCommandPlan: %v", err)
}
if !plan.Wrapped || plan.Name != "/usr/bin/bwrap" {
t.Fatalf("expected a wrapped bubblewrap plan, got wrapped=%v name=%q", plan.Wrapped, plan.Name)
}
}
func TestBuildCommandPlanReEntrancyGuardReturnsPassThrough(t *testing.T) {
// Simulate running inside an already-wrapped process: BOTH correlated markers
// that sandboxEnvironment sets together must be present for the guard to fire.
t.Setenv(EnvSandboxed, "1")
t.Setenv(EnvSandboxBackend, string(BackendBubblewrap))
root := t.TempDir()
engine := NewEngine(EngineOptions{
WorkspaceRoot: root,
Policy: DefaultPolicy(),
Backend: Backend{Name: BackendBubblewrap, Available: true, Executable: "/usr/bin/bwrap"},
})
plan, err := engine.BuildCommandPlan(CommandSpec{Name: "/bin/sh", Args: []string{"-c", "pwd"}, Dir: root})
if err != nil {
t.Fatalf("BuildCommandPlan: %v", err)
}
if plan.Wrapped {
t.Fatalf("re-entrancy guard must return an unwrapped pass-through plan, got wrapped=%v name=%q args=%v", plan.Wrapped, plan.Name, plan.Args)
}
if plan.Name != "/bin/sh" {
t.Fatalf("pass-through plan must run the command directly, got name=%q", plan.Name)
}
}
+16
View File
@@ -106,6 +106,22 @@ func classifyWithScope(request Request, scope *Scope) Risk {
if pipedInstallerPattern.MatchString(command) {
add("piped_installer", RiskCritical)
}
// AST second opinion (analyzer.go): walks the parsed shell tree, so it
// catches destructive/network programs the regexes miss — e.g. shred,
// fdisk, parted, and commands hidden behind sudo/env wrappers or a
// `sh -c <payload>` launcher — and flags an unparseable (obfuscated)
// script as elevated risk. It only ADDS categories, so a benign,
// parseable command is classified exactly as before.
analysis := AnalyzeCommand(command)
if analysis.Network {
add("network", RiskCritical)
}
if analysis.Destructive {
add("destructive", RiskCritical)
}
if analysis.TooComplex {
add("unparseable_command", RiskHigh)
}
}
for _, path := range requestPaths(request) {
+66
View File
@@ -229,3 +229,69 @@ func TestClassifyNoneSideEffectIsLowRisk(t *testing.T) {
t.Fatalf("control-only tool must not classify as out_of_workspace: %#v", risk)
}
}
// The following tests cover the AST analyzer wired into classifyWithScope as a
// second opinion to the regex detectors.
func TestClassifyASTCatchesDestructiveProgramsRegexMisses(t *testing.T) {
// shred/fdisk/parted are irrecoverably destructive but absent from the regex
// pattern; the AST analyzer flags them — including behind a sh -c launcher or
// a sudo/env wrapper (effectiveProgram resolves the real program). The escalated
// level (Critical) is part of the contract, so assert it alongside the category.
for _, command := range []string{
"shred -u secret.txt",
"fdisk /dev/sda",
"parted /dev/sda mklabel gpt",
"bash -c 'shred /etc/passwd'",
"sudo shred -u secret.txt",
"env shred -u secret.txt",
} {
risk := classifyCommand(command)
if risk.Level != RiskCritical || !HasRiskCategory(risk, "destructive") {
t.Fatalf("Classify(%q) = level %s, categories %v; want critical destructive", command, risk.Level, risk.Categories)
}
}
}
func TestClassifyASTCatchesNetworkProgramsRegexMisses(t *testing.T) {
for _, command := range []string{
"telnet example.com 23",
"ftp ftp.example.com",
"sftp user@host",
"sudo telnet example.com 23",
} {
risk := classifyCommand(command)
if risk.Level != RiskCritical || !HasRiskCategory(risk, "network") {
t.Fatalf("Classify(%q) = level %s, categories %v; want critical network", command, risk.Level, risk.Categories)
}
}
}
func TestClassifyFlagsUnparseableCommand(t *testing.T) {
// An unparseable (e.g. obfuscated) script can't be analyzed statically; the
// AST analyzer reports TooComplex so the classifier elevates it to High.
risk := classifyCommand(`echo "unterminated`)
if risk.Level != RiskHigh || !HasRiskCategory(risk, "unparseable_command") {
t.Fatalf("Classify of unparseable command = level %s, categories %v; want high unparseable_command", risk.Level, risk.Categories)
}
}
func TestClassifyASTDoesNotFlagQuotedProgramName(t *testing.T) {
// A program name inside a quoted argument is not a command, so the AST
// analyzer must not flag it (documenting a destructive command in an echo).
risk := classifyCommand(`echo "shred wipes files"`)
if HasRiskCategory(risk, "destructive") {
t.Fatalf("Classify(%q) wrongly flagged destructive: %v", `echo "shred wipes files"`, risk.Categories)
}
}
func TestClassifyBenignCommandStaysClean(t *testing.T) {
for _, command := range []string{"echo hello", "ls -la", "go build ./..."} {
risk := classifyCommand(command)
for _, category := range []string{"destructive", "network", "unparseable_command"} {
if HasRiskCategory(risk, category) {
t.Fatalf("Classify(%q) wrongly flagged %s: %v", command, category, risk.Categories)
}
}
}
}
+220 -24
View File
@@ -17,6 +17,13 @@ const bubblewrapWorkspace = "/workspace"
var errPolicyOnlyRunnerDisabled = errors.New("policy-only sandbox runner is disabled")
// errWSLPolicyOnlyDisabled is returned when running under WSL without bubblewrap
// and the policy has NOT opted into the policy-only runner — fail closed rather
// than run with no native isolation.
var errWSLPolicyOnlyDisabled = errors.New("bubblewrap is unavailable under WSL and Policy.AllowPolicyOnlyRunner is disabled: " +
"enable AllowPolicyOnlyRunner to run under the policy-only WSL fallback (engine policy + proxy egress, no native isolation), " +
"or install/enable bubblewrap")
type CommandSpec struct {
Name string
Args []string
@@ -39,6 +46,10 @@ type CommandPlan struct {
// StartDenialMonitor to capture what the sandbox blocked. Empty unless
// Policy.MonitorDenials is set on a macOS sandbox-exec plan.
MonitorTag string `json:"monitorTag,omitempty"`
// Notes records auditable least-privilege downgrade notes — e.g. the WSL
// policy-only fallback noting that native isolation was unavailable. Surfaced to
// the operator; never affects enforcement.
Notes []string `json:"notes,omitempty"`
// cleanup releases resources tied to the plan's lifetime — currently the
// scoped-egress proxy, which must outlive the command run and be shut down
// afterwards. It is never serialized; callers invoke it via Cleanup() once the
@@ -99,6 +110,29 @@ func ProxyEnv(addr string) []string {
}
}
// ProxyEnvWithSocks is the SOCKS-aware form of ProxyEnv built ON TOP of it:
// HTTP_PROXY/HTTPS_PROXY stay on the HTTP listener (httpAddr) for HTTP(S)-aware
// clients, while ALL_PROXY/all_proxy are re-pointed at the SOCKS5 front-end
// (socksAddr) so a client that honors ALL_PROXY tunnels arbitrary TCP through the
// same allow/deny gate. When socksAddr is empty it returns exactly ProxyEnv's
// output, so the SOCKS upgrade is purely additive and never weakens the default.
func ProxyEnvWithSocks(httpAddr, socksAddr string) []string {
env := ProxyEnv(httpAddr)
if socksAddr == "" {
return env
}
socksURL := "socks5://" + socksAddr
for i, kv := range env {
switch {
case strings.HasPrefix(kv, "ALL_PROXY="):
env[i] = "ALL_PROXY=" + socksURL
case strings.HasPrefix(kv, "all_proxy="):
env[i] = "all_proxy=" + socksURL
}
}
return env
}
func (engine *Engine) CommandContext(ctx context.Context, spec CommandSpec) (*exec.Cmd, CommandPlan, error) {
if ctx == nil {
ctx = context.Background()
@@ -118,10 +152,20 @@ func (engine *Engine) CommandContext(ctx context.Context, spec CommandSpec) (*ex
// only applies to engines built without a workspace root (NewEngine always
// builds a scope otherwise); it is kept as defense in depth.
func (engine *Engine) writeRoots(workspaceRoot string) []string {
var roots []string
if engine.scope != nil {
return engine.scope.Roots()
roots = engine.scope.Roots()
} else {
roots = []string{workspaceRoot}
}
return []string{workspaceRoot}
// Reflect the policy's AllowWrite roots in the OS backend write binds so a
// sandboxed shell command may write where the policy grants writes. DenyWrite
// is enforced at the policy gate, and on sandbox-exec additionally as an
// explicit deny rule (see sandboxExecProfile).
if extra := resolveWriteRootPaths(engine.policy.AllowWrite); len(extra) > 0 {
roots = dedupeStrings(append(roots, extra...))
}
return roots
}
func (engine *Engine) BuildCommandPlan(spec CommandSpec) (CommandPlan, error) {
@@ -146,6 +190,15 @@ func (engine *Engine) BuildCommandPlan(spec CommandSpec) (CommandPlan, error) {
if backend.Name == "" {
backend = Backend{Name: BackendPolicyOnly, Message: "policy-only fallback: sandbox backend was not selected"}
}
// Re-entrancy guard: a command spawned by a process we already wrapped (both
// ZERO_SANDBOXED=1 and ZERO_SANDBOX_BACKEND set in its env — see
// IsAlreadySandboxed) must not be wrapped again — nested bwrap / sandbox-exec
// fails and a second egress proxy would be redundant. Return a pass-through
// plan. Mirrors the already-sandboxed guard used by comparable executor
// sandboxes.
if IsAlreadySandboxed() {
return directCommandPlan(spec, backend, policy, workspaceRoot), nil
}
if policy.Mode == ModeDisabled {
return directCommandPlan(spec, backend, policy, workspaceRoot), nil
}
@@ -166,6 +219,18 @@ func (engine *Engine) BuildCommandPlan(spec CommandSpec) (CommandPlan, error) {
}
return sandboxExecCommandPlan(spec, workspaceRoot, engine.writeRoots(workspaceRoot), policy, backend, egress), nil
}
case BackendWSL:
// WSL fallback: no native isolation available. Fail closed unless the policy
// opted into the policy-only runner; otherwise run policy-only with the
// command's egress routed through the filtering proxy and an auditable note.
if !policy.AllowPolicyOnlyRunner {
return CommandPlan{}, errWSLPolicyOnlyDisabled
}
egress, err := startScopedEgress(policy, backend)
if err != nil {
return CommandPlan{}, err
}
return wslCommandPlan(spec, workspaceRoot, policy, backend, egress), nil
}
if !policy.AllowPolicyOnlyRunner {
return CommandPlan{}, errPolicyOnlyRunnerDisabled
@@ -173,12 +238,79 @@ func (engine *Engine) BuildCommandPlan(spec CommandSpec) (CommandPlan, error) {
return directCommandPlan(spec, backend, policy, workspaceRoot), nil
}
// wslCommandPlan builds the WSL policy-only fallback plan: the command runs
// DIRECTLY (no native wrap, since bubblewrap is unavailable under WSL) but with
// the sandbox env markers (ZERO_SANDBOXED=1, ZERO_SANDBOX_BACKEND=wsl) and, when
// a scoped-egress proxy was started, the proxy env so well-behaved clients route
// through the allow/deny gate. It carries a least-privilege downgrade note.
func wslCommandPlan(spec CommandSpec, workspaceRoot string, policy Policy, backend Backend, egress *scopedEgress) CommandPlan {
// The WSL fallback runs the command DIRECTLY (no native wrap), so it inherits
// the caller's environment like the generic policy-only runner; the sandbox
// markers are appended last so they cannot be shadowed by the caller's env.
var env []string
if spec.Env != nil {
env = cloneStrings(spec.Env)
} else {
env = append(env, os.Environ()...)
}
env = append(env,
EnvSandboxBackend+"="+string(BackendWSL),
"ZERO_SANDBOX_NETWORK="+string(policy.Network),
EnvSandboxed+"=1",
)
note := "WSL: native process isolation (bubblewrap) is unavailable; downgraded to the policy-only runner " +
"(least privilege). The engine still enforces the full policy."
if egress != nil {
env = append(env, ProxyEnvWithSocks(egress.addr, egress.socksAddr)...)
env = appendCACertEnv(env, egress)
note = "WSL: native process isolation (bubblewrap) is unavailable; downgraded to the policy-only runner " +
"with proxy egress (least privilege). The engine still enforces the full policy."
}
plan := CommandPlan{
Backend: backend,
WorkspaceRoot: workspaceRoot,
Policy: policy,
Wrapped: false,
Name: spec.Name,
Args: cloneStrings(spec.Args),
Dir: spec.Dir,
Env: env,
Notes: []string{note},
}
if egress != nil {
plan.cleanup = egress.cleanup
}
return plan
}
// scopedEgress holds the address of a started scoped-egress proxy and the
// cleanup that shuts it down. A nil *scopedEgress means scoped egress is not in
// effect for this command (the network mode is allow or deny-equivalent).
type scopedEgress struct {
addr string
cleanup func()
addr string
// socksAddr is the loopback address of the proxy's SOCKS5 front-end, used to
// set ALL_PROXY=socks5://<socksAddr>. Empty when the proxy exposes no SOCKS
// listener, in which case ALL_PROXY falls back to the HTTP proxy.
socksAddr string
// caCertPath is the path to the MITM CA's public cert (PEM), set only when
// Policy.InspectTLS started a TLS-terminating proxy. Surfaced to the sandboxed
// client via ZERO_SANDBOX_CA_CERT so it can trust the minted leaves.
caCertPath string
cleanup func()
}
// EnvCACert is the env var pointing the sandboxed client at the MITM CA's public
// cert so it can trust the proxy's minted leaf certificates (only set when
// Policy.InspectTLS is on).
const EnvCACert = "ZERO_SANDBOX_CA_CERT"
// appendCACertEnv adds ZERO_SANDBOX_CA_CERT when the egress proxy is terminating
// TLS (InspectTLS). It is a no-op for a plain passthrough proxy.
func appendCACertEnv(env []string, egress *scopedEgress) []string {
if egress == nil || egress.caCertPath == "" {
return env
}
return append(env, EnvCACert+"="+egress.caCertPath)
}
// startScopedEgress starts the local filtering proxy when the policy's effective
@@ -196,13 +328,32 @@ func startScopedEgress(policy Policy, backend Backend) (*scopedEgress, error) {
return nil, nil
}
proxy, err := startEgressProxy(egressOptions{
Allowed: policy.AllowedDomains,
Denied: policy.DeniedDomains,
Allowed: policy.AllowedDomains,
Denied: policy.DeniedDomains,
InspectTLS: policy.InspectTLS,
})
if err != nil {
return nil, fmt.Errorf("scoped egress unavailable, denying network: %w", err)
}
return &scopedEgress{addr: proxy.Addr(), cleanup: func() { _ = proxy.Close() }}, nil
se := &scopedEgress{addr: proxy.Addr(), socksAddr: proxy.SocksAddr(), cleanup: func() { _ = proxy.Close() }}
if policy.InspectTLS {
// Write the MITM CA's PUBLIC cert so the sandboxed client can trust the
// minted leaves; surface it via ZERO_SANDBOX_CA_CERT and clean it up with
// the proxy. Reuses no new bind path: sandbox-exec already allows file-read,
// and the WSL policy-only runner shares the host filesystem.
caPath, werr := writeMITMCAFile(proxy.CACertPEM())
if werr != nil {
_ = proxy.Close()
return nil, werr
}
se.caCertPath = caPath
base := se.cleanup
se.cleanup = func() {
base()
_ = os.Remove(caPath)
}
}
return se, nil
}
func directCommandPlan(spec CommandSpec, backend Backend, policy Policy, workspaceRoot string) CommandPlan {
@@ -374,21 +525,25 @@ func findSeccompHelper() string {
}
func sandboxExecCommandPlan(spec CommandSpec, workspaceRoot string, writeRoots []string, policy Policy, backend Backend, egress *scopedEgress) CommandPlan {
var proxyPort string
var proxyPort, socksPort string
if egress != nil {
if _, port, err := net.SplitHostPort(egress.addr); err == nil {
proxyPort = port
}
if _, port, err := net.SplitHostPort(egress.socksAddr); err == nil {
socksPort = port
}
}
denialTag := ""
if policy.MonitorDenials {
denialTag = nextSandboxDenialTag()
}
args := []string{"-p", sandboxExecProfile(writeRoots, policy, proxyPort, denialTag), spec.Name}
args := []string{"-p", sandboxExecProfile(writeRoots, policy, proxyPort, socksPort, denialTag), spec.Name}
args = append(args, spec.Args...)
env := sandboxEnvironment(policy, BackendSandboxExec, workspaceRoot)
if egress != nil {
env = append(env, ProxyEnv(egress.addr)...)
env = append(env, ProxyEnvWithSocks(egress.addr, egress.socksAddr)...)
env = appendCACertEnv(env, egress)
}
plan := CommandPlan{
Backend: backend,
@@ -426,8 +581,9 @@ func sandboxEnvironment(policy Policy, backend BackendName, home string) []strin
"HOME=" + home,
"PATH=" + firstEnv("PATH", defaultPath()),
"TERM=" + firstEnv("TERM", "dumb"),
"ZERO_SANDBOX_BACKEND=" + string(backend),
EnvSandboxBackend + "=" + string(backend),
"ZERO_SANDBOX_NETWORK=" + string(policy.Network),
EnvSandboxed + "=1",
}
if runtime.GOOS == "windows" {
env = append(env, "COMSPEC="+firstEnv("COMSPEC", "cmd.exe"))
@@ -540,8 +696,8 @@ func sandboxMachLookupRule() string {
return "(allow mach-lookup\n " + strings.Join(filters, "\n ") + ")"
}
func sandboxExecProfile(writeRoots []string, policy Policy, proxyPort string, denialTag string) string {
networkRule := networkRuleFor(policy, proxyPort)
func sandboxExecProfile(writeRoots []string, policy Policy, proxyPort string, socksPort string, denialTag string) string {
networkRule := networkRuleFor(policy, proxyPort, socksPort)
writeRule := "(allow file-write*)"
if policy.EnforceWorkspace {
// The granted write roots are the only writable *project* locations. Temp
@@ -566,7 +722,7 @@ func sandboxExecProfile(writeRoots []string, policy Policy, proxyPort string, de
// filters `log stream` for this exact (per-plan) tag.
denyDefault = `(deny default (with message "` + sandboxProfileString(denialTag) + `"))`
}
return strings.Join([]string{
rules := []string{
"(version 1)",
denyDefault,
"(allow process*)",
@@ -579,27 +735,67 @@ func sandboxExecProfile(writeRoots []string, policy Policy, proxyPort string, de
sandboxMachLookupRule(),
"(allow file-read*)",
writeRule,
networkRule,
}, "\n")
}
// DenyWrite entries are emitted AFTER the write allow so seatbelt's
// last-match-wins makes them override the workspace/AllowWrite binds — so
// "DenyWrite wins" holds for shell commands too, not just the policy gate.
rules = append(rules, denyWriteRules(policy)...)
rules = append(rules, networkRule)
return strings.Join(rules, "\n")
}
// denyWriteRules returns seatbelt deny clauses for the policy's resolved
// DenyWrite paths: a (subpath ...) clause for a directory, a (literal ...) clause
// for a single file. Empty when DenyWrite is unset.
func denyWriteRules(policy Policy) []string {
resolved := resolvePolicyPaths(policy.DenyWrite)
if len(resolved) == 0 {
return nil
}
out := make([]string, 0, len(resolved))
for _, path := range resolved {
filter := `(subpath "` + sandboxProfileString(path) + `")`
if info, err := os.Stat(path); err == nil && !info.IsDir() {
filter = `(literal "` + sandboxProfileString(path) + `")`
}
out = append(out, "(deny file-write* "+filter+")")
}
return out
}
// networkRuleFor returns the seatbelt network clause for a policy. allow opens
// all network; deny (and an empty-allowlist scoped policy, which effectiveNetwork
// collapses to deny) blocks all network; scoped denies general network but
// permits only outbound to the local proxy port on localhost, so traffic must
// flow through the filtering proxy. A scoped policy with no resolvable proxy port
// falls back to a full deny (fail closed).
func networkRuleFor(policy Policy, proxyPort string) string {
// permits only outbound to the local proxy ports on localhost, so traffic must
// flow through the filtering proxy. Both the HTTP proxy port and the SOCKS5
// front-end port are allowed (a sandboxed process configured with
// ALL_PROXY=socks5://... must reach the SOCKS listener). A scoped policy with no
// resolvable proxy port falls back to a full deny (fail closed).
func networkRuleFor(policy Policy, proxyPort string, socksPort string) string {
switch effectiveNetwork(policy) {
case NetworkAllow:
return "(allow network*)"
case NetworkScoped:
if strings.TrimSpace(proxyPort) == "" {
rules := []string{"(deny network*)"}
seen := map[string]struct{}{}
// Deny by default, then allow only outbound to each distinct proxy port on
// loopback. Duplicate/empty ports are skipped so the clause stays minimal.
for _, port := range []string{proxyPort, socksPort} {
port = strings.TrimSpace(port)
if port == "" {
continue
}
if _, ok := seen[port]; ok {
continue
}
seen[port] = struct{}{}
rules = append(rules, `(allow network-outbound (remote ip "localhost:`+sandboxProfileString(port)+`"))`)
}
if len(rules) == 1 {
// No resolvable proxy port: fail closed.
return "(deny network*)"
}
// Deny by default, then allow only outbound to the proxy on loopback.
return "(deny network*)\n" +
`(allow network-outbound (remote ip "localhost:` + sandboxProfileString(proxyPort) + `"))`
return strings.Join(rules, "\n")
default:
return "(deny network*)"
}
+12 -3
View File
@@ -23,8 +23,17 @@ func TestScopedNetworkGateInEvaluate(t *testing.T) {
}
sandboxExec := Backend{Name: BackendSandboxExec, Available: true, Executable: "/usr/sbin/sandbox-exec", ScopedEgress: true}
// This test exercises the network GATE for a network tool, so it opts the
// in-process tools into the network policy (EnforceToolNetwork); by default
// those tools are exempt and would not be gated here.
enforcedScoped := func(allowed, denied []string) Policy {
p := scopedPolicy(allowed, denied)
p.EnforceToolNetwork = true
return p
}
// An enforcing backend lets a populated scoped policy permit a network tool.
enforcing := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: scopedPolicy([]string{"github.com"}, nil), Backend: sandboxExec})
enforcing := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: enforcedScoped([]string{"github.com"}, nil), Backend: sandboxExec})
if decision := enforcing.Evaluate(context.Background(), networkRequest); decision.Action == ActionDeny {
t.Fatalf("enforcing backend + scoped policy denied a network tool: %#v", decision)
}
@@ -35,7 +44,7 @@ func TestScopedNetworkGateInEvaluate(t *testing.T) {
"policy-only": {},
}
for name, backend := range unenforceable {
engine := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: scopedPolicy([]string{"github.com"}, nil), Backend: backend})
engine := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: enforcedScoped([]string{"github.com"}, nil), Backend: backend})
decision := engine.Evaluate(context.Background(), networkRequest)
if decision.Action != ActionDeny || decision.Violation == nil || decision.Violation.Code != ViolationNetwork {
t.Fatalf("%s backend must deny scoped network it cannot enforce, got %#v", name, decision)
@@ -43,7 +52,7 @@ func TestScopedNetworkGateInEvaluate(t *testing.T) {
}
// Empty allowlist still fails closed regardless of backend.
empty := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: scopedPolicy(nil, nil), Backend: sandboxExec})
empty := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: enforcedScoped(nil, nil), Backend: sandboxExec})
decision := empty.Evaluate(context.Background(), networkRequest)
if decision.Action != ActionDeny || decision.Violation == nil || decision.Violation.Code != ViolationNetwork {
t.Fatalf("empty scoped policy must deny network like NetworkDeny, got %#v", decision)
+48 -4
View File
@@ -233,7 +233,7 @@ func resolvedTestPath(t *testing.T, path string) string {
}
func TestSandboxExecProfileIncludesExtraWriteRoots(t *testing.T) {
profile := sandboxExecProfile([]string{"/ws", "/extra root"}, Policy{Mode: ModeEnforce, EnforceWorkspace: true}, "", "")
profile := sandboxExecProfile([]string{"/ws", "/extra root"}, Policy{Mode: ModeEnforce, EnforceWorkspace: true}, "", "", "")
if !strings.Contains(profile, "(allow file-write*") {
t.Fatalf("profile missing file-write rule:\n%s", profile)
}
@@ -251,7 +251,7 @@ func TestSandboxExecProfileIncludesExtraWriteRoots(t *testing.T) {
}
func TestSandboxExecProfileTagsDenialsWhenMonitoring(t *testing.T) {
off := sandboxExecProfile([]string{"/ws"}, Policy{Mode: ModeEnforce, EnforceWorkspace: true}, "", "")
off := sandboxExecProfile([]string{"/ws"}, Policy{Mode: ModeEnforce, EnforceWorkspace: true}, "", "", "")
if strings.Contains(off, "with message") {
t.Fatalf("denials must not be tagged when monitoring is off:\n%s", off)
}
@@ -259,7 +259,7 @@ func TestSandboxExecProfileTagsDenialsWhenMonitoring(t *testing.T) {
t.Fatalf("profile missing the plain default-deny:\n%s", off)
}
on := sandboxExecProfile([]string{"/ws"}, Policy{Mode: ModeEnforce, EnforceWorkspace: true, MonitorDenials: true}, "", "run-tag-123")
on := sandboxExecProfile([]string{"/ws"}, Policy{Mode: ModeEnforce, EnforceWorkspace: true, MonitorDenials: true}, "", "", "run-tag-123")
if !strings.Contains(on, `(deny default (with message "run-tag-123"))`) {
t.Fatalf("denials must be tagged when monitoring is on:\n%s", on)
}
@@ -291,7 +291,7 @@ func TestSandboxExecCommandPlanUsesUniquePerPlanDenialTag(t *testing.T) {
}
func TestSandboxExecProfileGrantsSignalAndMachLookup(t *testing.T) {
profile := sandboxExecProfile([]string{"/ws"}, Policy{Mode: ModeEnforce, EnforceWorkspace: true}, "", "")
profile := sandboxExecProfile([]string{"/ws"}, Policy{Mode: ModeEnforce, EnforceWorkspace: true}, "", "", "")
// Signalling own process group lets a sandboxed script kill the children it
// spawns; without it seatbelt denies kill() with "Operation not permitted".
@@ -465,3 +465,47 @@ func TestProxyEnv(t *testing.T) {
t.Fatalf("ProxyEnv NO_PROXY must exclude loopback, got %q/%q", env["NO_PROXY"], env["no_proxy"])
}
}
func TestProxyEnvWithSocks(t *testing.T) {
env := map[string]string{}
for _, kv := range ProxyEnvWithSocks("127.0.0.1:8899", "127.0.0.1:1080") {
if i := strings.IndexByte(kv, '='); i >= 0 {
env[kv[:i]] = kv[i+1:]
}
}
// HTTP_PROXY/HTTPS_PROXY stay on the HTTP listener...
const httpURL = "http://127.0.0.1:8899"
for _, key := range []string{"HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy"} {
if env[key] != httpURL {
t.Fatalf("ProxyEnvWithSocks[%s] = %q, want %q", key, env[key], httpURL)
}
}
// ...while ALL_PROXY/all_proxy point at the SOCKS5 front-end.
const socksURL = "socks5://127.0.0.1:1080"
for _, key := range []string{"ALL_PROXY", "all_proxy"} {
if env[key] != socksURL {
t.Fatalf("ProxyEnvWithSocks[%s] = %q, want %q", key, env[key], socksURL)
}
}
if !strings.Contains(env["NO_PROXY"], "127.0.0.1") {
t.Fatalf("ProxyEnvWithSocks NO_PROXY must exclude loopback, got %q", env["NO_PROXY"])
}
}
// TestProxyEnvWithSocksEmptyFallsBackToHTTP verifies the SOCKS upgrade is purely
// additive: with no SOCKS address, ALL_PROXY falls back to the HTTP proxy exactly
// as the legacy ProxyEnv produced, never weakening the default.
func TestProxyEnvWithSocksEmptyFallsBackToHTTP(t *testing.T) {
env := map[string]string{}
for _, kv := range ProxyEnvWithSocks("127.0.0.1:8899", "") {
if i := strings.IndexByte(kv, '='); i >= 0 {
env[kv[:i]] = kv[i+1:]
}
}
const httpURL = "http://127.0.0.1:8899"
for _, key := range []string{"ALL_PROXY", "all_proxy"} {
if env[key] != httpURL {
t.Fatalf("ProxyEnvWithSocks(empty socks)[%s] = %q, want HTTP fallback %q", key, env[key], httpURL)
}
}
}
+72
View File
@@ -12,6 +12,31 @@ import (
// only an explicit truthy value enables it.
const EnvAutoAllowBash = "ZERO_SANDBOX_AUTO_ALLOW_BASH"
// EnvSandboxed marks a process that zero has already wrapped in a sandbox: every
// wrapped command carries ZERO_SANDBOXED=1 in its environment. When such a
// process spawns another command through the engine, the re-entrancy guard
// returns a pass-through plan instead of double-wrapping it — nested bwrap /
// sandbox-exec fails, and a second egress proxy would be redundant. Mirrors the
// already-sandboxed guard used by comparable executor sandboxes. Unset by default.
const EnvSandboxed = "ZERO_SANDBOXED"
// EnvSandboxBackend records which backend wrapped the command. sandboxEnvironment
// always sets it alongside EnvSandboxed, so it serves as a corroborating marker:
// the re-entrancy guard requires BOTH, raising the provenance bar above a single
// ambient flag (a stray or hand-exported ZERO_SANDBOXED=1 with no backend marker
// no longer forces an unsandboxed pass-through).
const EnvSandboxBackend = "ZERO_SANDBOX_BACKEND"
// IsAlreadySandboxed reports whether the current process is already running
// inside a zero-created sandbox. It requires BOTH correlated markers that
// sandboxEnvironment sets together — EnvSandboxed == "1" AND a non-empty
// EnvSandboxBackend — so a single user-set/inherited ZERO_SANDBOXED=1 cannot by
// itself disable wrapping. zero sets both only on genuinely wrapped commands;
// pass-through (direct) plans set neither.
func IsAlreadySandboxed() bool {
return os.Getenv(EnvSandboxed) == "1" && strings.TrimSpace(os.Getenv(EnvSandboxBackend)) != ""
}
type SideEffect string
type Permission string
type PermissionMode string
@@ -106,6 +131,11 @@ const (
BackendBubblewrap BackendName = "bubblewrap"
BackendSandboxExec BackendName = "sandbox-exec"
BackendPolicyOnly BackendName = "policy-only"
// BackendWSL is the policy-only fallback used under WSL when bubblewrap is
// unavailable/unreliable: there is no native OS isolation, but network egress
// is still routed through the local filtering proxy and the command runs under
// the full policy engine. It is surfaced via ZERO_SANDBOX_BACKEND=wsl.
BackendWSL BackendName = "wsl"
)
const (
@@ -133,6 +163,14 @@ type Policy struct {
// NetworkDeny so existing policies keep their exact behaviour.
AllowedDomains []string `json:"allowedDomains,omitempty"`
DeniedDomains []string `json:"deniedDomains,omitempty"`
// EnforceToolNetwork, when true, subjects the first-party in-process network
// tools (web_search / web_fetch) to the sandbox network policy via
// NetworkHostAllowed. It is OFF by default: that policy exists to confine the
// sandboxed SHELL's egress, which these tools do not use, so by default they
// are allowed to run (keeping their own SSRF/port/redirect/redaction
// safeguards). The sandboxed-shell egress decision is independent of this flag.
// Turn it on to also hold web_search/web_fetch to the allow/scoped/deny policy.
EnforceToolNetwork bool `json:"enforceToolNetwork,omitempty"`
// AutoAllowBashWhenSandboxed, when true, auto-allows the bash tool WITHOUT a
// permission prompt — but only when the sandbox is actually active (a
// native-isolation backend wraps the command). The sandbox is then the safety
@@ -153,6 +191,40 @@ type Policy struct {
// (runs without the filter) when the helper binary is not found. Ignored on
// non-bubblewrap backends.
BlockUnixSockets bool `json:"blockUnixSockets,omitempty"`
// AllowRead/DenyRead/AllowWrite/DenyWrite are fine-grained path lists layered
// ON TOP of the workspace + Scope guards; they never bypass the symlink /
// out-of-workspace protections. Each entry is home-expanded, made absolute, and
// symlink-resolved (an entry that does not exist is dropped). All default empty,
// so an unconfigured policy behaves exactly as before. Semantics:
//
// - Read: a path readable under the base workspace/Scope guard is denied if it
// falls under a DenyRead entry, UNLESS a more-specific AllowRead entry (one
// nested inside that DenyRead entry) re-includes it. AllowRead only
// re-includes within a DenyRead carve-out; it never extends reads beyond the
// workspace.
// - Write: DenyWrite wins over everything; otherwise a path writable under the
// workspace/Scope guard is allowed; otherwise an absolute path under an
// AllowWrite root is allowed; otherwise it is denied. AllowWrite roots are
// also reflected in the OS backend write binds, and on sandbox-exec DenyWrite
// entries are emitted as explicit deny rules so the precedence holds for
// shell commands too.
AllowRead []string `json:"allowRead,omitempty"`
DenyRead []string `json:"denyRead,omitempty"`
AllowWrite []string `json:"allowWrite,omitempty"`
DenyWrite []string `json:"denyWrite,omitempty"`
// InspectTLS, when true, lets the in-process scoped-egress proxy TERMINATE TLS
// (using a locally generated, ephemeral CA) so it can enforce the per-host
// allow/deny rules on the DECRYPTED request Host — today a CONNECT tunnel only
// reveals the host once. OFF by default: the proxy stays a pure CONNECT
// passthrough, byte-for-byte unchanged. SECURITY/TRUST: enabling this means the
// sandbox can read the plaintext of the sandboxed process's TLS traffic. The
// MITM only re-signs toward the sandboxed client; the upstream connection is
// ALWAYS validated against the system roots (never InsecureSkipVerify), and the
// decrypted Host still passes the SAME authorize()/domainAllowed() gate — MITM
// widens visibility, never authority. The generated CA's public cert is written
// to the sandbox runtime dir and surfaced via ZERO_SANDBOX_CA_CERT so the
// sandboxed client can trust it.
InspectTLS bool `json:"inspectTLS,omitempty"`
}
type Request struct {
+26
View File
@@ -0,0 +1,26 @@
package sandbox
import "strings"
// WSLInfo describes the Windows Subsystem for Linux environment, if any. It is
// derived from /proc/version. Mirrors reference wsl/index.js isWSL/isWSL2 and
// sandbox-manager.js isWsl (reads /proc/version, matches "microsoft").
type WSLInfo struct {
IsWSL bool
IsWSL2 bool
Kernel string
}
// parseWSL classifies a /proc/version string. IsWSL is set when it contains
// "microsoft" or "wsl"; IsWSL2 additionally requires "microsoft" and the absence
// of the legacy "wsl1" marker. Factored out (pure) so it is table-testable
// without a real /proc on any GOOS.
func parseWSL(procVersion string) WSLInfo {
lower := strings.ToLower(procVersion)
isWSL := strings.Contains(lower, "microsoft") || strings.Contains(lower, "wsl")
info := WSLInfo{IsWSL: isWSL, Kernel: strings.TrimSpace(procVersion)}
if isWSL {
info.IsWSL2 = strings.Contains(lower, "microsoft") && !strings.Contains(lower, "wsl1")
}
return info
}
+15
View File
@@ -0,0 +1,15 @@
//go:build linux
package sandbox
import "os"
// detectWSL reads /proc/version once and classifies the WSL environment. A read
// error (no /proc) yields the zero WSLInfo (not WSL).
func detectWSL() WSLInfo {
data, err := os.ReadFile("/proc/version")
if err != nil {
return WSLInfo{}
}
return parseWSL(string(data))
}
+8
View File
@@ -0,0 +1,8 @@
//go:build !linux
package sandbox
// detectWSL is a no-op on non-Linux platforms: WSL is a Linux-under-Windows
// environment, so the host GOOS is never WSL. Keeps the package building on all
// platforms.
func detectWSL() WSLInfo { return WSLInfo{} }
+139
View File
@@ -0,0 +1,139 @@
package sandbox
import (
"errors"
"strings"
"testing"
)
func TestParseWSL(t *testing.T) {
cases := []struct {
name string
input string
wantWSL bool
wantWSL2 bool
}{
{
name: "wsl2 microsoft kernel",
input: "Linux version 5.15.90.1-microsoft-standard-WSL2 (...)",
wantWSL: true,
wantWSL2: true,
},
{
name: "wsl1 legacy marker",
input: "Linux version 4.4.0-19041-Microsoft (wsl1 build) ...",
wantWSL: true,
wantWSL2: false,
},
{
name: "plain linux",
input: "Linux version 6.8.0-generic (gcc ...) #1 SMP",
wantWSL: false,
wantWSL2: false,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
got := parseWSL(tc.input)
if got.IsWSL != tc.wantWSL || got.IsWSL2 != tc.wantWSL2 {
t.Fatalf("parseWSL(%q) = {IsWSL:%v IsWSL2:%v}, want {%v %v}", tc.input, got.IsWSL, got.IsWSL2, tc.wantWSL, tc.wantWSL2)
}
if got.IsWSL && got.Kernel == "" {
t.Fatalf("parseWSL should record the kernel string")
}
})
}
}
func wslBackendForTest() Backend {
return Backend{Name: BackendWSL, Platform: "linux", Fallback: true, ProxyEgress: true}
}
func TestWSLPlanFailsClosedWithoutPolicyOnly(t *testing.T) {
root := t.TempDir()
policy := scopedPolicy([]string{"github.com"}, nil)
policy.AllowPolicyOnlyRunner = false // explicitly refuse the degraded runner
engine := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: policy, Backend: wslBackendForTest()})
_, err := engine.BuildCommandPlan(CommandSpec{Name: "/bin/sh", Args: []string{"-c", "pwd"}, Dir: root})
if !errors.Is(err, errWSLPolicyOnlyDisabled) {
t.Fatalf("WSL without AllowPolicyOnlyRunner err = %v, want errWSLPolicyOnlyDisabled", err)
}
}
func TestWSLPolicyOnlyPlanCarriesProxyAndMarkers(t *testing.T) {
root := t.TempDir()
// scopedPolicy uses DefaultPolicy (AllowPolicyOnlyRunner=true) + an allowlist,
// so the WSL fallback starts the filtering proxy and runs policy-only.
engine := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: scopedPolicy([]string{"github.com"}, nil), Backend: wslBackendForTest()})
plan, err := engine.BuildCommandPlan(CommandSpec{Name: "/bin/sh", Args: []string{"-c", "pwd"}, Dir: root, Env: []string{}})
if err != nil {
t.Fatalf("BuildCommandPlan: %v", err)
}
defer plan.Cleanup()
// Policy-only: the command runs directly, not OS-wrapped.
if plan.Wrapped || plan.Name != "/bin/sh" {
t.Fatalf("WSL plan must run the command directly, got wrapped=%v name=%q", plan.Wrapped, plan.Name)
}
env := strings.Join(plan.Env, "\n")
if !strings.Contains(env, EnvSandboxed+"=1") {
t.Fatalf("WSL plan env missing %s=1:\n%s", EnvSandboxed, env)
}
if !strings.Contains(env, EnvSandboxBackend+"=wsl") {
t.Fatalf("WSL plan env missing %s=wsl:\n%s", EnvSandboxBackend, env)
}
if !strings.Contains(env, "HTTP_PROXY=") || !strings.Contains(env, "ALL_PROXY=") {
t.Fatalf("WSL plan env missing proxy env (ProxyEnvWithSocks):\n%s", env)
}
if len(plan.Notes) == 0 || !strings.Contains(strings.Join(plan.Notes, " "), "policy-only") {
t.Fatalf("WSL plan must record a least-privilege downgrade note, got %v", plan.Notes)
}
}
func TestWSLPlanPreservesCallerEnv(t *testing.T) {
root := t.TempDir()
engine := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: DefaultPolicy(), Backend: wslBackendForTest()})
plan, err := engine.BuildCommandPlan(CommandSpec{
Name: "/bin/sh", Args: []string{"-c", "true"}, Dir: root,
Env: []string{"FOO=bar", "PATH=/custom/bin"},
})
if err != nil {
t.Fatalf("BuildCommandPlan: %v", err)
}
defer plan.Cleanup()
env := strings.Join(plan.Env, "\n")
// The WSL direct-run inherits the caller's env (it is not OS-wrapped)...
if !strings.Contains(env, "FOO=bar") || !strings.Contains(env, "PATH=/custom/bin") {
t.Fatalf("WSL plan must preserve caller env, got:\n%s", env)
}
// ...with the sandbox markers appended (so they win over caller env).
if !strings.Contains(env, EnvSandboxed+"=1") || !strings.Contains(env, EnvSandboxBackend+"=wsl") {
t.Fatalf("WSL plan must append sandbox markers:\n%s", env)
}
}
func TestWSLPolicyOnlyDenyPlanHasMarkersNoProxy(t *testing.T) {
root := t.TempDir()
// Default policy (NetworkDeny) → no proxy started, but the WSL markers + note
// still apply and the plan still runs (AllowPolicyOnlyRunner default true).
engine := NewEngine(EngineOptions{WorkspaceRoot: root, Policy: DefaultPolicy(), Backend: wslBackendForTest()})
// Explicit empty env so the assertion below is not confused by a dev/CI
// HTTP_PROXY inherited via os.Environ().
plan, err := engine.BuildCommandPlan(CommandSpec{Name: "/bin/sh", Args: []string{"-c", "pwd"}, Dir: root, Env: []string{}})
if err != nil {
t.Fatalf("BuildCommandPlan: %v", err)
}
defer plan.Cleanup()
env := strings.Join(plan.Env, "\n")
if !strings.Contains(env, EnvSandboxBackend+"=wsl") || !strings.Contains(env, EnvSandboxed+"=1") {
t.Fatalf("WSL deny plan must still carry sandbox markers:\n%s", env)
}
if len(plan.Notes) == 0 || !strings.Contains(strings.Join(plan.Notes, " "), "policy-only") {
t.Fatalf("WSL deny plan must record a least-privilege downgrade note, got %v", plan.Notes)
}
if strings.Contains(env, "HTTP_PROXY=") || strings.Contains(env, "HTTPS_PROXY=") || strings.Contains(env, "ALL_PROXY=") {
t.Fatalf("WSL deny plan must NOT export any proxy env vars (network denied):\n%s", env)
}
}
+44 -12
View File
@@ -58,14 +58,20 @@ type BuildArgsInput struct {
CurrentDepth int
Description string
Cwd string
// PermissionMode is the parent's resolved permission mode. Empty preserves
// the historical "--auto high" (the Task tool omits it); a caller that
// passes a non-unsafe mode (e.g. the swarm propagating a non-unsafe
// orchestrator) gets a non-unsafe child so authority is never widened.
PermissionMode string
}
type BuildResumeArgsInput struct {
SessionID string
Prompt string
CurrentDepth int
Manifest Manifest
Cwd string
SessionID string
Prompt string
CurrentDepth int
Manifest Manifest
Cwd string
PermissionMode string
}
type BuildArgsResult struct {
@@ -90,8 +96,32 @@ type TaskRunOptions struct {
ParentReasoningEffort string
CurrentDepth int
Cwd string
// PermissionMode propagates the parent's resolved permission mode to the
// child. Empty keeps the historical "--auto high"; a non-unsafe value runs
// the child non-unsafe so it never gains more authority than the parent.
PermissionMode string
}
// specialistAutonomy maps a parent permission mode to the child's "--auto"
// level. A headless specialist child cannot answer interactive prompts, so it
// runs autonomously — but only at "high" (unsafe) when the parent is itself
// unsafe (or when no mode is provided, preserving the Task tool's behavior).
// Any other parent mode yields the non-unsafe "low" level, so the child's
// authority never exceeds the parent's.
func specialistAutonomy(permissionMode string) string {
switch strings.TrimSpace(permissionMode) {
case "", string(permissionModeUnsafe):
return "high"
default:
return "low"
}
}
// permissionModeUnsafe mirrors agent.PermissionModeUnsafe without importing the
// agent package (which would create an import cycle): exec resolves "--auto high"
// to this mode.
const permissionModeUnsafe = "unsafe"
type ExecResult struct {
Result tools.Result
SessionID string
@@ -147,7 +177,7 @@ func (executor Executor) BuildArgs(input BuildArgsInput) (BuildArgsResult, error
args := []string{"exec", "--init-session-id", sessionID}
args = append(args, promptArgs...)
args = appendModelArgs(args, input.Manifest, input.ParentModel, input.ParentReasoningEffort)
args = append(args, "--auto", "high", "--output-format", "stream-json")
args = append(args, "--auto", specialistAutonomy(input.PermissionMode), "--output-format", "stream-json")
toolAllowlist, err := resolvedToolAllowlist(input.Manifest)
if err != nil {
return BuildArgsResult{}, err
@@ -199,7 +229,7 @@ func (executor Executor) BuildResumeArgs(input BuildResumeArgsInput) (BuildArgsR
}
args := []string{"exec", "--resume", sessionID}
args = append(args, promptArgs...)
args = append(args, "--auto", "high", "--output-format", "stream-json")
args = append(args, "--auto", specialistAutonomy(input.PermissionMode), "--output-format", "stream-json")
toolAllowlist, err := resolvedToolAllowlist(input.Manifest)
if err != nil {
return BuildArgsResult{}, err
@@ -230,6 +260,7 @@ func (executor Executor) runFresh(ctx context.Context, params TaskParameters, op
CurrentDepth: options.CurrentDepth,
Description: params.Description,
Cwd: options.Cwd,
PermissionMode: options.PermissionMode,
})
if err != nil {
return ExecResult{}, err
@@ -257,11 +288,12 @@ func (executor Executor) runResume(ctx context.Context, params TaskParameters, o
return ExecResult{}, err
}
built, err := executor.BuildResumeArgs(BuildResumeArgsInput{
SessionID: params.Resume,
Prompt: params.Prompt,
CurrentDepth: options.CurrentDepth,
Manifest: manifest,
Cwd: options.Cwd,
SessionID: params.Resume,
Prompt: params.Prompt,
CurrentDepth: options.CurrentDepth,
Manifest: manifest,
Cwd: options.Cwd,
PermissionMode: options.PermissionMode,
})
if err != nil {
return ExecResult{}, err
+59
View File
@@ -64,6 +64,65 @@ func TestBuildArgsCreatesFreshSpecialistExecInvocation(t *testing.T) {
}
}
func TestSpecialistAutonomyByPermissionMode(t *testing.T) {
cases := map[string]string{
"": "high", // back-compat: the Task tool omits the mode
"unsafe": "high",
"auto": "low",
"ask": "low",
"spec-draft": "low",
"whatever": "low",
}
for mode, want := range cases {
if got := specialistAutonomy(mode); got != want {
t.Errorf("specialistAutonomy(%q) = %q, want %q", mode, got, want)
}
}
}
func TestBuildArgsAutonomyHonorsPermissionMode(t *testing.T) {
executor := Executor{NewSessionID: func() (string, error) { return "child", nil }}
manifest := Manifest{Metadata: Metadata{Name: "reviewer"}, SystemPrompt: "x", ResolvedTools: []string{"read_file"}}
// A non-unsafe parent mode yields a non-unsafe child (--auto low), so a
// swarm member never gains more authority than a non-unsafe orchestrator.
res, err := executor.BuildArgs(BuildArgsInput{Manifest: manifest, Prompt: "p", PermissionMode: "auto"})
if err != nil {
t.Fatalf("BuildArgs: %v", err)
}
if !containsSequence(res.Args, []string{"--auto", "low"}) {
t.Fatalf("non-unsafe parent must yield --auto low, got %v", res.Args)
}
if containsSequence(res.Args, []string{"--auto", "high"}) {
t.Fatalf("non-unsafe parent must NOT yield --auto high: %v", res.Args)
}
// An unsafe parent (and the empty back-compat case) keep --auto high.
for _, mode := range []string{"unsafe", ""} {
out, err := executor.BuildArgs(BuildArgsInput{Manifest: manifest, Prompt: "p", PermissionMode: mode})
if err != nil {
t.Fatalf("BuildArgs(%q): %v", mode, err)
}
if !containsSequence(out.Args, []string{"--auto", "high"}) {
t.Fatalf("mode %q must yield --auto high, got %v", mode, out.Args)
}
}
}
func TestBuildResumeArgsAutonomyHonorsPermissionMode(t *testing.T) {
executor := Executor{}
manifest := Manifest{Metadata: Metadata{Name: "reviewer"}, ResolvedTools: []string{"read_file"}}
res, err := executor.BuildResumeArgs(BuildResumeArgsInput{
SessionID: "child_session", Prompt: "p", Manifest: manifest, PermissionMode: "auto",
})
if err != nil {
t.Fatalf("BuildResumeArgs: %v", err)
}
if !containsSequence(res.Args, []string{"--auto", "low"}) || containsSequence(res.Args, []string{"--auto", "high"}) {
t.Fatalf("resume non-unsafe parent must yield --auto low, got %v", res.Args)
}
}
func TestBuildArgsSessionTitleFallsBackToNameWithoutDescription(t *testing.T) {
executor := Executor{NewSessionID: func() (string, error) { return "child", nil }}
manifest := Manifest{
+259
View File
@@ -0,0 +1,259 @@
package swarm
import (
"errors"
"fmt"
"sort"
"sync"
"time"
)
// TaskStatus is the lifecycle state of a swarm task.
type TaskStatus string
const (
// StatusPending: registered but the member has not started.
StatusPending TaskStatus = "pending"
// StatusRunning: the member is executing.
StatusRunning TaskStatus = "running"
// StatusDone: the member finished successfully.
StatusDone TaskStatus = "done"
// StatusFailed: the member exited with an error.
StatusFailed TaskStatus = "failed"
// StatusHandedOff: ownership was transferred to another agent.
StatusHandedOff TaskStatus = "handed-off"
)
// terminal reports whether a status is final (no further transitions expected).
func (s TaskStatus) terminal() bool {
return s == StatusDone || s == StatusFailed || s == StatusHandedOff
}
// valid reports whether a status is one of the known lifecycle states. Unknown
// statuses are rejected (fail closed) so a task can never enter an undefined
// state that Summarize and the lifecycle logic cannot reason about.
func (s TaskStatus) valid() bool {
switch s {
case StatusPending, StatusRunning, StatusDone, StatusFailed, StatusHandedOff:
return true
default:
return false
}
}
// Task is one unit of work tracked by the coordinator. It is returned by value
// from snapshots so callers cannot mutate coordinator state without the lock.
type Task struct {
ID string
AgentID string
Team string
Description string
Status TaskStatus
Result string
Err string
CreatedAt time.Time
UpdatedAt time.Time
}
// agentColors palette for stable per-agent coloring in the TUI/status. Provider
// -neutral ANSI-ish names; the renderer maps them to styles.
var agentColors = []string{"cyan", "magenta", "green", "yellow", "blue", "red"}
// Coordinator is the in-memory task registry + team color assigner shared by an
// orchestrator and its members. It is safe for concurrent use.
type Coordinator struct {
mu sync.RWMutex
tasks map[string]*Task
colors map[string]string // agentID -> color
colorIndex int
now func() time.Time // injectable clock for tests
}
// NewCoordinator returns an empty coordinator using the wall clock.
func NewCoordinator() *Coordinator {
return &Coordinator{
tasks: map[string]*Task{},
colors: map[string]string{},
now: time.Now,
}
}
// ErrTaskExists is returned when registering a duplicate task ID.
var ErrTaskExists = errors.New("swarm: task already registered")
// ErrUnknownTask is returned when updating/removing a missing task ID.
var ErrUnknownTask = errors.New("swarm: unknown task")
// Register adds a new pending task. The ID must be unique and non-empty.
func (c *Coordinator) Register(id, agentID, team, description string) (Task, error) {
if id == "" {
return Task{}, errors.New("swarm: task id is required")
}
c.mu.Lock()
defer c.mu.Unlock()
if _, ok := c.tasks[id]; ok {
return Task{}, fmt.Errorf("%w: %s", ErrTaskExists, id)
}
now := c.now()
t := &Task{
ID: id,
AgentID: agentID,
Team: team,
Description: description,
Status: StatusPending,
CreatedAt: now,
UpdatedAt: now,
}
c.tasks[id] = t
c.assignColorLocked(agentID)
return *t, nil
}
// SetStatus transitions a task. Transitions out of a terminal state are
// rejected (fail closed) so a late member update cannot resurrect a finished
// task.
func (c *Coordinator) SetStatus(id string, status TaskStatus) error {
if !status.valid() {
return fmt.Errorf("swarm: invalid task status %q", status)
}
c.mu.Lock()
defer c.mu.Unlock()
t, ok := c.tasks[id]
if !ok {
return fmt.Errorf("%w: %s", ErrUnknownTask, id)
}
if t.Status.terminal() && t.Status != status {
return fmt.Errorf("swarm: task %s is %s (terminal); cannot move to %s", id, t.Status, status)
}
t.Status = status
t.UpdatedAt = c.now()
return nil
}
// Complete marks a task done with its result.
func (c *Coordinator) Complete(id, result string) error {
return c.finish(id, StatusDone, result, "")
}
// Fail marks a task failed with an error message.
func (c *Coordinator) Fail(id, errMsg string) error {
return c.finish(id, StatusFailed, "", errMsg)
}
func (c *Coordinator) finish(id string, status TaskStatus, result, errMsg string) error {
c.mu.Lock()
defer c.mu.Unlock()
t, ok := c.tasks[id]
if !ok {
return fmt.Errorf("%w: %s", ErrUnknownTask, id)
}
if t.Status.terminal() {
return fmt.Errorf("swarm: task %s already %s", id, t.Status)
}
t.Status = status
t.Result = result
t.Err = errMsg
t.UpdatedAt = c.now()
return nil
}
// Reassign transfers a task to a new agent (handoff/orphan-adoption). The task
// returns to pending under the new owner unless it is already terminal.
func (c *Coordinator) Reassign(id, newAgentID string) error {
c.mu.Lock()
defer c.mu.Unlock()
t, ok := c.tasks[id]
if !ok {
return fmt.Errorf("%w: %s", ErrUnknownTask, id)
}
if t.Status.terminal() {
return fmt.Errorf("swarm: task %s already %s; cannot reassign", id, t.Status)
}
t.AgentID = newAgentID
t.Status = StatusPending
t.UpdatedAt = c.now()
c.assignColorLocked(newAgentID)
return nil
}
// Get returns a snapshot of one task.
func (c *Coordinator) Get(id string) (Task, bool) {
c.mu.RLock()
defer c.mu.RUnlock()
t, ok := c.tasks[id]
if !ok {
return Task{}, false
}
return *t, true
}
// List returns snapshots of all tasks, ordered by creation time then ID for
// stable output.
func (c *Coordinator) List() []Task {
c.mu.RLock()
out := make([]Task, 0, len(c.tasks))
for _, t := range c.tasks {
out = append(out, *t)
}
c.mu.RUnlock()
sort.Slice(out, func(i, j int) bool {
if out[i].CreatedAt.Equal(out[j].CreatedAt) {
return out[i].ID < out[j].ID
}
return out[i].CreatedAt.Before(out[j].CreatedAt)
})
return out
}
// Color returns the stable color assigned to an agent (assigning one on first
// use), so status output colors each member consistently.
func (c *Coordinator) Color(agentID string) string {
c.mu.Lock()
defer c.mu.Unlock()
return c.assignColorLocked(agentID)
}
func (c *Coordinator) assignColorLocked(agentID string) string {
if agentID == "" {
return ""
}
if color, ok := c.colors[agentID]; ok {
return color
}
color := agentColors[c.colorIndex%len(agentColors)]
c.colors[agentID] = color
c.colorIndex++
return color
}
// Summary is an aggregate count of tasks by status for the status tool.
type Summary struct {
Total int
Pending int
Running int
Done int
Failed int
HandedOff int
}
// Summarize aggregates current task counts.
func (c *Coordinator) Summarize() Summary {
c.mu.RLock()
defer c.mu.RUnlock()
s := Summary{Total: len(c.tasks)}
for _, t := range c.tasks {
switch t.Status {
case StatusPending:
s.Pending++
case StatusRunning:
s.Running++
case StatusDone:
s.Done++
case StatusFailed:
s.Failed++
case StatusHandedOff:
s.HandedOff++
}
}
return s
}
+151
View File
@@ -0,0 +1,151 @@
package swarm
import (
"errors"
"sync"
"testing"
)
func TestCoordinatorRegisterAndDuplicate(t *testing.T) {
c := NewCoordinator()
if _, err := c.Register("t1", "a1", "team", "desc"); err != nil {
t.Fatalf("Register: %v", err)
}
task, ok := c.Get("t1")
if !ok || task.Status != StatusPending {
t.Fatalf("Get after register: ok=%v status=%v", ok, task.Status)
}
if _, err := c.Register("t1", "a1", "team", "desc"); !errors.Is(err, ErrTaskExists) {
t.Fatalf("duplicate Register err = %v, want ErrTaskExists", err)
}
if _, err := c.Register("", "a", "team", "d"); err == nil {
t.Fatal("empty id must error")
}
}
func TestCoordinatorLifecycleTransitions(t *testing.T) {
c := NewCoordinator()
_, _ = c.Register("t1", "a1", "team", "desc")
if err := c.SetStatus("t1", StatusRunning); err != nil {
t.Fatalf("SetStatus running: %v", err)
}
if err := c.Complete("t1", "result-data"); err != nil {
t.Fatalf("Complete: %v", err)
}
task, _ := c.Get("t1")
if task.Status != StatusDone || task.Result != "result-data" {
t.Fatalf("after Complete: status=%v result=%q", task.Status, task.Result)
}
// Terminal guard: cannot move done -> running, nor complete twice.
if err := c.SetStatus("t1", StatusRunning); err == nil {
t.Fatal("moving a done task to running must fail")
}
if err := c.Complete("t1", "x"); err == nil {
t.Fatal("completing a done task must fail")
}
if err := c.Fail("t1", "boom"); err == nil {
t.Fatal("failing a done task must fail")
}
}
func TestCoordinatorRejectsInvalidStatus(t *testing.T) {
c := NewCoordinator()
_, _ = c.Register("t1", "a1", "team", "desc")
if err := c.SetStatus("t1", TaskStatus("bogus")); err == nil {
t.Fatal("SetStatus must reject an unknown status")
}
// The task stays in its prior (valid) state.
if task, _ := c.Get("t1"); task.Status != StatusPending {
t.Fatalf("rejected status must not mutate the task, got %v", task.Status)
}
}
func TestCoordinatorUnknownTask(t *testing.T) {
c := NewCoordinator()
if err := c.SetStatus("missing", StatusRunning); !errors.Is(err, ErrUnknownTask) {
t.Fatalf("SetStatus(missing) = %v, want ErrUnknownTask", err)
}
if err := c.Complete("missing", "x"); !errors.Is(err, ErrUnknownTask) {
t.Fatalf("Complete(missing) = %v, want ErrUnknownTask", err)
}
}
func TestCoordinatorReassign(t *testing.T) {
c := NewCoordinator()
_, _ = c.Register("t1", "a1", "team", "desc")
_ = c.SetStatus("t1", StatusRunning)
if err := c.Reassign("t1", "a2"); err != nil {
t.Fatalf("Reassign: %v", err)
}
task, _ := c.Get("t1")
if task.AgentID != "a2" || task.Status != StatusPending {
t.Fatalf("after Reassign: agent=%q status=%v", task.AgentID, task.Status)
}
// Cannot reassign a terminal task.
_ = c.Complete("t1", "done")
if err := c.Reassign("t1", "a3"); err == nil {
t.Fatal("reassigning a terminal task must fail")
}
}
func TestCoordinatorColorStability(t *testing.T) {
c := NewCoordinator()
first := c.Color("a1")
if first == "" {
t.Fatal("Color should assign a non-empty color")
}
if again := c.Color("a1"); again != first {
t.Fatalf("Color not stable: %q vs %q", first, again)
}
if other := c.Color("a2"); other == "" {
t.Fatal("second agent should also get a color")
}
if c.Color("") != "" {
t.Fatal("empty agent id should get no color")
}
}
func TestCoordinatorSummarize(t *testing.T) {
c := NewCoordinator()
_, _ = c.Register("p", "a", "team", "")
_, _ = c.Register("r", "b", "team", "")
_ = c.SetStatus("r", StatusRunning)
_, _ = c.Register("d", "c", "team", "")
_ = c.Complete("d", "ok")
s := c.Summarize()
if s.Total != 3 || s.Pending != 1 || s.Running != 1 || s.Done != 1 {
t.Fatalf("Summarize = %+v", s)
}
}
func TestCoordinatorConcurrentRegister(t *testing.T) {
c := NewCoordinator()
var wg sync.WaitGroup
for i := 0; i < 50; i++ {
wg.Add(1)
go func(n int) {
defer wg.Done()
id := "task-" + string(rune('a'+n%26)) + "-" + itoa(n)
_, _ = c.Register(id, id, "team", "desc")
_ = c.SetStatus(id, StatusRunning)
_ = c.Complete(id, "done")
}(i)
}
wg.Wait()
if got := len(c.List()); got != 50 {
t.Fatalf("List len = %d, want 50", got)
}
}
// itoa avoids importing strconv just for the concurrency test id.
func itoa(n int) string {
if n == 0 {
return "0"
}
var b []byte
for n > 0 {
b = append([]byte{byte('0' + n%10)}, b...)
n /= 10
}
return string(b)
}
+133
View File
@@ -0,0 +1,133 @@
// Package swarm adds a multi-agent SWARM on top of zero's single sub-agent
// mechanism (internal/specialist): an orchestrator can spawn and coordinate
// MULTIPLE specialist members that run concurrently, communicate via per-agent
// mailboxes, and hand work off. It composes internal/specialist (to launch each
// member), internal/daemon (pooled workers when a daemon is running) and
// internal/background. Everything is additive and opt-in — the existing single
// "task" tool is unchanged; swarm tools are only active when invoked.
//
// Every member runs under the SAME sandbox + risk/autonomy policy as the
// orchestrator: the swarm never grants a member more authority than its parent.
package swarm
import (
"errors"
"sort"
"strings"
"sync"
)
// ErrUnknownAgentType is returned when a definition lookup misses.
var ErrUnknownAgentType = errors.New("swarm: unknown agent type")
// modelInherit is the sentinel meaning "use the orchestrator's model" — members
// never silently pick a different (e.g. more capable) model.
const modelInherit = "inherit"
// Definition is one entry in the agent roster: how a member of a given type
// behaves (agent type, when to use it, model — "inherit" by default —,
// permission mode, and a system-prompt builder).
type Definition struct {
AgentType string
WhenToUse string
Model string // "inherit" => use the orchestrator's model
PermissionMode string
// SystemPrompt returns the member's system prompt for the given task context.
// It is a func so a definition can fold the task briefing in.
SystemPrompt func(ctx PromptContext) string
}
// PromptContext is the minimal context handed to a definition's SystemPrompt.
type PromptContext struct {
Team string
Task string
}
// Registry is the agent roster: definitions looked up by agent type. Built-ins
// are seeded at construction; user-defined definitions extend or override them.
type Registry struct {
mu sync.RWMutex
defs map[string]Definition
}
// NewRegistry builds a registry seeded with the built-in roster.
func NewRegistry() *Registry {
r := &Registry{defs: map[string]Definition{}}
for _, def := range builtinDefinitions() {
r.defs[def.AgentType] = def
}
return r
}
// Register adds or overrides a definition (user-defined agents extend the
// built-ins). An empty AgentType is rejected.
func (r *Registry) Register(def Definition) error {
agentType := strings.TrimSpace(def.AgentType)
if agentType == "" {
return errors.New("swarm: definition requires an agentType")
}
if def.Model == "" {
def.Model = modelInherit
}
def.AgentType = agentType
r.mu.Lock()
r.defs[agentType] = def
r.mu.Unlock()
return nil
}
// Lookup returns the definition for agentType, or ErrUnknownAgentType.
func (r *Registry) Lookup(agentType string) (Definition, error) {
r.mu.RLock()
def, ok := r.defs[strings.TrimSpace(agentType)]
r.mu.RUnlock()
if !ok {
return Definition{}, ErrUnknownAgentType
}
return def, nil
}
// AgentTypes returns the registered agent types, sorted, for status/help.
func (r *Registry) AgentTypes() []string {
r.mu.RLock()
types := make([]string, 0, len(r.defs))
for t := range r.defs {
types = append(types, t)
}
r.mu.RUnlock()
sort.Strings(types)
return types
}
// builtinDefinitions returns the seeded roster (teammate, subagent) — kept
// minimal; user-defined agents extend it via Register.
func builtinDefinitions() []Definition {
return []Definition{
{
AgentType: "teammate",
WhenToUse: "In-process teammate for parallel task execution; delegate work to run alongside the orchestrator.",
Model: modelInherit,
// Empty PermissionMode => inherit the orchestrator's mode (never widened).
SystemPrompt: func(ctx PromptContext) string {
return "You are a teammate agent collaborating with an orchestrator on team " +
displayTeam(ctx.Team) + ". Complete your assigned task and report results.\n\nTask: " + ctx.Task
},
},
{
AgentType: "subagent",
WhenToUse: "General-purpose subagent for an isolated, delegated task; starts with zero prior context.",
Model: modelInherit,
SystemPrompt: func(ctx PromptContext) string {
return "You are a subagent spawned to complete a specific task. You start with zero context — " +
"the briefing below is all you know.\n\nTask: " + ctx.Task
},
},
}
}
func displayTeam(team string) string {
if strings.TrimSpace(team) == "" {
return "default"
}
return team
}
+74
View File
@@ -0,0 +1,74 @@
package swarm
import (
"errors"
"testing"
)
func TestRegistryBuiltins(t *testing.T) {
r := NewRegistry()
for _, want := range []string{"teammate", "subagent"} {
def, err := r.Lookup(want)
if err != nil {
t.Fatalf("builtin %q missing: %v", want, err)
}
if def.Model != modelInherit {
t.Fatalf("builtin %q model = %q, want inherit", want, def.Model)
}
if def.SystemPrompt == nil {
t.Fatalf("builtin %q must have a SystemPrompt", want)
}
if got := def.SystemPrompt(PromptContext{Team: "t", Task: "do x"}); got == "" {
t.Fatalf("builtin %q SystemPrompt returned empty", want)
}
}
}
func TestRegistryLookupUnknown(t *testing.T) {
r := NewRegistry()
if _, err := r.Lookup("nope"); !errors.Is(err, ErrUnknownAgentType) {
t.Fatalf("Lookup(nope) err = %v, want ErrUnknownAgentType", err)
}
}
func TestRegistryRegisterOverrideAndDefaultModel(t *testing.T) {
r := NewRegistry()
if err := r.Register(Definition{AgentType: " custom "}); err != nil {
t.Fatalf("Register: %v", err)
}
def, err := r.Lookup("custom")
if err != nil {
t.Fatalf("Lookup(custom): %v", err)
}
if def.AgentType != "custom" {
t.Fatalf("agentType not trimmed: %q", def.AgentType)
}
if def.Model != modelInherit {
t.Fatalf("default model = %q, want inherit", def.Model)
}
// Override a builtin.
if err := r.Register(Definition{AgentType: "teammate", Model: "fixed-model"}); err != nil {
t.Fatalf("override Register: %v", err)
}
if got, _ := r.Lookup("teammate"); got.Model != "fixed-model" {
t.Fatalf("override model = %q, want fixed-model", got.Model)
}
}
func TestRegistryRegisterRejectsEmpty(t *testing.T) {
r := NewRegistry()
if err := r.Register(Definition{AgentType: " "}); err == nil {
t.Fatal("Register with blank agentType must error")
}
}
func TestRegistryAgentTypesSorted(t *testing.T) {
r := NewRegistry()
_ = r.Register(Definition{AgentType: "alpha"})
types := r.AgentTypes()
for i := 1; i < len(types); i++ {
if types[i-1] > types[i] {
t.Fatalf("AgentTypes not sorted: %v", types)
}
}
}
+47
View File
@@ -0,0 +1,47 @@
package swarm
import (
"context"
"github.com/Gitlawb/zero/internal/specialist"
)
// NewSpecialistLauncher adapts internal/specialist.Executor into a
// MemberLauncher: each swarm member runs as a specialist sub-agent. This is the
// production "direct" launch path (used when no daemon pool is wired). The
// member inherits the orchestrator's model via spec.Model and runs in spec.Cwd,
// so it stays under the same sandbox + policy as its parent — the swarm never
// grants a member more authority.
//
// agentType maps to the specialist agent name; the specialist's own definition
// for that name governs the member's system prompt and tools. RunInBackground is
// deliberately left false: the swarm provides concurrency itself (one goroutine
// per member), so each Run executes to completion inside its goroutine.
//
// spec.PermissionMode (resolved + clamped by buildSpec, never above the
// orchestrator's) is propagated so the child only runs unsafe/high autonomy when
// the orchestrator is itself unsafe; otherwise it runs non-unsafe. An empty mode
// is substituted with the non-unsafe "auto" default so a swarm member can never
// fall into the historical empty-means-high path — that is the Task tool's
// behavior, not the swarm's.
func NewSpecialistLauncher(executor specialist.Executor) MemberLauncher {
return FuncLauncher{Run: func(ctx context.Context, spec MemberSpec) (MemberResult, error) {
permissionMode := spec.PermissionMode
if permissionMode == "" {
permissionMode = permissionModeAuto
}
res, err := executor.Run(ctx, specialist.TaskParameters{
Name: spec.AgentType,
Prompt: spec.Task,
Description: spec.AgentType + " · team " + spec.Team,
}, specialist.TaskRunOptions{
ParentModel: spec.Model,
Cwd: spec.Cwd,
PermissionMode: permissionMode,
})
if err != nil {
return MemberResult{}, err
}
return MemberResult{Result: res.Result.Output, SessionID: res.SessionID}, nil
}}
}
+176
View File
@@ -0,0 +1,176 @@
package swarm
import (
"fmt"
)
// Spawn registers a task and launches a member of agentType to run it under the
// given team, inheriting the orchestrator's policy. It returns the task id
// immediately; the member runs concurrently (under the Swarm's base context, so
// it outlives this call) and reports its result through the coordinator. If the
// team is at its slot cap the member is queued and launches when a slot frees
// (it stays pending in the coordinator until then).
func (s *Swarm) Spawn(pol Policy, teamName, agentType, task, cwd string) (string, error) {
def, err := s.registry.Lookup(agentType)
if err != nil {
return "", err
}
team := sanitizeName(teamName)
id := s.nextID(agentType)
if _, err := s.coord.Register(id, id, team, task); err != nil {
return "", err
}
s.rememberCwd(id, cwd)
spec := s.buildSpec(pol, id, id, team, def, task, cwd)
s.dispatch(spec)
return id, nil
}
// dispatch admits a spec to its team (launching now or queuing for a slot).
func (s *Swarm) dispatch(spec MemberSpec) {
t := s.team(spec.Team)
if t.admit(spec) {
s.launch(t, spec)
}
// Otherwise the spec is queued; the coordinator task stays pending until a
// slot frees and afterExit launches it.
}
// launch starts a member for spec and supervises it. A synchronous launch
// failure fails the task and frees the slot.
func (s *Swarm) launch(t *Team, spec MemberSpec) {
handle, err := s.launcher.Launch(s.baseCtx, spec)
if err != nil {
_ = s.coord.Fail(spec.TaskID, "launch: "+err.Error())
s.afterExit(t)
return
}
m := &Member{ID: spec.ID, AgentType: spec.AgentType, TaskID: spec.TaskID, handle: handle}
t.addMember(m)
_ = s.coord.SetStatus(spec.TaskID, StatusRunning)
go s.watch(t, m, spec)
}
// watch awaits a member, applies bounded relaunch on temporary failures, records
// the terminal outcome, then frees the slot and drains the queue.
//
// The relaunch counter is per-member (m.restarts) and the same spec is reused on
// retry. This is sound because a Member is bound 1:1 to its spec for its whole
// life (Member.ID == MemberSpec.ID); a retry never reuses the struct for a
// different spec.
func (s *Swarm) watch(t *Team, m *Member, spec MemberSpec) {
res, err := m.handle.Wait()
if err != nil {
if isRetryable(err) && m.restarts < maxMemberRestarts {
if nh, relErr := s.launcher.Launch(s.baseCtx, spec); relErr == nil {
m.restarts++
m.handle = nh
go s.watch(t, m, spec)
return
}
// fall through to record the original error if relaunch fails
}
_ = s.coord.Fail(m.TaskID, memberError(err))
} else {
_ = s.coord.Complete(m.TaskID, res.Result)
}
t.removeMember(m.ID)
s.afterExit(t)
}
// afterExit releases the just-vacated slot and launches the next queued spec, if
// any. Each exit drains at most one queued member; that member's own exit drains
// the next, so the queue empties one-per-slot without unbounded recursion.
func (s *Swarm) afterExit(t *Team) {
next, ok := t.onExit()
if !ok {
return
}
s.launch(t, next)
}
// Handoff transfers a task to a fresh member of toAgentType, delivering a note to
// the new member's inbox and marking the original task handed-off. It returns the
// new task id. A handoff of an already-terminal task is rejected (fail closed).
func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string) (string, error) {
task, ok := s.coord.Get(taskID)
if !ok {
return "", fmt.Errorf("%w: %s", ErrUnknownTask, taskID)
}
if task.Status.terminal() {
return "", fmt.Errorf("swarm: task %s already %s; cannot hand off", taskID, task.Status)
}
def, err := s.registry.Lookup(toAgentType)
if err != nil {
return "", err
}
team := sanitizeName(teamName)
newID := s.nextID(toAgentType)
handoffTask := task.Description
if note != "" {
handoffTask += "\n\nHandoff note: " + note
}
if _, err := s.coord.Register(newID, newID, team, handoffTask); err != nil {
return "", err
}
cwd := s.cwdFor(taskID)
s.rememberCwd(newID, cwd)
// Deliver the handoff note to the new member's inbox. A mailbox failure must
// not silently strand the handoff, so it is surfaced to the caller.
if note != "" {
if mbErr := s.mailbox.Send(team, newID, Message{
From: task.AgentID, Subject: "handoff", Body: note, Type: "handoff", Time: nowRFC3339(),
}); mbErr != nil {
return "", fmt.Errorf("swarm: deliver handoff note: %w", mbErr)
}
}
// Retire the original task (it has been re-delegated).
_ = s.coord.SetStatus(taskID, StatusHandedOff)
spec := s.buildSpec(pol, newID, newID, team, def, handoffTask, cwd)
s.dispatch(spec)
return newID, nil
}
// AdoptOrphans re-parents tasks in a team whose owning member is no longer live
// (e.g. a crashed worker) onto fresh members of toAgentType, returning the
// adopted task ids. Terminal tasks and tasks with a live owner are left alone.
func (s *Swarm) AdoptOrphans(pol Policy, teamName, toAgentType string) ([]string, error) {
def, err := s.registry.Lookup(toAgentType)
if err != nil {
return nil, err
}
team := sanitizeName(teamName)
t := s.team(team)
live := t.liveAgents()
var adopted []string
for _, task := range s.coord.List() {
if task.Team != team || task.Status.terminal() {
continue
}
if _, ok := live[task.AgentID]; ok {
continue // still has a live member
}
newAgent := s.nextID(toAgentType)
if err := s.coord.Reassign(task.ID, newAgent); err != nil {
continue // raced to terminal; skip
}
cwd := s.cwdFor(task.ID)
s.rememberCwd(task.ID, cwd)
spec := s.buildSpec(pol, newAgent, task.ID, team, def, task.Description, cwd)
s.dispatch(spec)
adopted = append(adopted, task.ID)
}
return adopted, nil
}
// Collect returns snapshots of every task in a team, for swarm_collect/status.
func (s *Swarm) Collect(teamName string) []Task {
team := sanitizeName(teamName)
var out []Task
for _, task := range s.coord.List() {
if task.Team == team {
out = append(out, task)
}
}
return out
}
+345
View File
@@ -0,0 +1,345 @@
package swarm
import (
"context"
"strings"
"sync"
"testing"
"time"
)
// controllableLauncher records every launched spec and lets a test control each
// member's outcome and (optionally) gate completion until a channel is closed.
type controllableLauncher struct {
mu sync.Mutex
specs []MemberSpec
attempts map[string]int
gate chan struct{} // nil => members return immediately
result func(spec MemberSpec, attempt int) (MemberResult, error)
launchErr error
}
func newLauncher(result func(MemberSpec, int) (MemberResult, error)) *controllableLauncher {
return &controllableLauncher{attempts: map[string]int{}, result: result}
}
func (l *controllableLauncher) Launch(ctx context.Context, spec MemberSpec) (MemberHandle, error) {
l.mu.Lock()
if l.launchErr != nil {
err := l.launchErr
l.mu.Unlock()
return nil, err
}
l.specs = append(l.specs, spec)
l.attempts[spec.ID]++
attempt := l.attempts[spec.ID]
gate := l.gate
result := l.result
l.mu.Unlock()
h := &funcHandle{id: spec.ID, done: make(chan struct{})}
go func() {
defer close(h.done)
if gate != nil {
select {
case <-gate:
case <-ctx.Done():
h.err = ctx.Err()
return
}
}
if result != nil {
h.res, h.err = result(spec, attempt)
}
}()
return h, nil
}
func (l *controllableLauncher) recorded() []MemberSpec {
l.mu.Lock()
defer l.mu.Unlock()
out := make([]MemberSpec, len(l.specs))
copy(out, l.specs)
return out
}
func (l *controllableLauncher) attemptCount(id string) int {
l.mu.Lock()
defer l.mu.Unlock()
return l.attempts[id]
}
func newSwarmFor(t *testing.T, l MemberLauncher) *Swarm {
t.Helper()
sw, err := New(Options{BaseDir: t.TempDir(), Launcher: l, MaxTeamSize: 2})
if err != nil {
t.Fatalf("New: %v", err)
}
t.Cleanup(sw.Close)
return sw
}
func waitFor(t *testing.T, what string, cond func() bool) {
t.Helper()
deadline := time.Now().Add(3 * time.Second)
for time.Now().Before(deadline) {
if cond() {
return
}
time.Sleep(5 * time.Millisecond)
}
t.Fatalf("timed out waiting for %s", what)
}
func okFor(spec MemberSpec, _ int) (MemberResult, error) {
return MemberResult{Result: "ok:" + spec.Task, SessionID: "sess-" + spec.ID}, nil
}
func TestSpawnCompletes(t *testing.T) {
l := newLauncher(okFor)
sw := newSwarmFor(t, l)
id, err := sw.Spawn(Policy{Model: "m"}, "team", "teammate", "build widget", "/work")
if err != nil {
t.Fatalf("Spawn: %v", err)
}
waitFor(t, "task done", func() bool {
task, ok := sw.Coordinator().Get(id)
return ok && task.Status == StatusDone
})
task, _ := sw.Coordinator().Get(id)
if task.Result != "ok:build widget" {
t.Fatalf("result = %q", task.Result)
}
}
func TestSpawnInheritsPolicy(t *testing.T) {
l := newLauncher(okFor)
sw := newSwarmFor(t, l)
_, err := sw.Spawn(Policy{Model: "orch-model", PermissionMode: permissionModeAuto}, "team", "teammate", "task", "/cwd")
if err != nil {
t.Fatalf("Spawn: %v", err)
}
waitFor(t, "spec recorded", func() bool { return len(l.recorded()) == 1 })
spec := l.recorded()[0]
if spec.Model != "orch-model" {
t.Fatalf("member model = %q, want inherited orch-model", spec.Model)
}
if spec.PermissionMode != permissionModeAuto {
t.Fatalf("member permission mode = %q, want inherited auto", spec.PermissionMode)
}
if spec.Cwd != "/cwd" {
t.Fatalf("member cwd = %q, want /cwd", spec.Cwd)
}
if spec.SystemPrompt == "" {
t.Fatal("member should carry a resolved system prompt")
}
}
func TestConcurrencyCapAndQueueDrains(t *testing.T) {
gate := make(chan struct{})
l := newLauncher(okFor)
l.gate = gate
sw := newSwarmFor(t, l) // MaxTeamSize 2
pol := Policy{Model: "m"}
for i := 0; i < 5; i++ {
if _, err := sw.Spawn(pol, "team", "teammate", "task", ""); err != nil {
t.Fatalf("Spawn %d: %v", i, err)
}
}
team := sw.team("team")
if team.Running() != 2 || team.QueueDepth() != 3 {
t.Fatalf("cap not enforced: running=%d queue=%d, want 2/3", team.Running(), team.QueueDepth())
}
// Release everyone; the queue should drain one-per-slot until all are done.
close(gate)
waitFor(t, "all tasks done", func() bool { return sw.Coordinator().Summarize().Done == 5 })
if team.Running() != 0 || team.QueueDepth() != 0 {
t.Fatalf("after drain running=%d queue=%d, want 0/0", team.Running(), team.QueueDepth())
}
if got := len(l.recorded()); got != 5 {
t.Fatalf("launched %d members, want 5", got)
}
}
func TestRetryOnTemporaryError(t *testing.T) {
l := newLauncher(func(spec MemberSpec, attempt int) (MemberResult, error) {
if attempt < 3 {
return MemberResult{}, ErrMemberTemporary
}
return MemberResult{Result: "recovered"}, nil
})
sw := newSwarmFor(t, l)
id, _ := sw.Spawn(Policy{}, "team", "teammate", "task", "")
waitFor(t, "task recovered", func() bool {
task, ok := sw.Coordinator().Get(id)
return ok && task.Status == StatusDone
})
if got := l.attemptCount(id); got != 3 {
t.Fatalf("attempts = %d, want 3 (initial + 2 retries)", got)
}
task, _ := sw.Coordinator().Get(id)
if task.Result != "recovered" {
t.Fatalf("result = %q, want recovered", task.Result)
}
}
func TestRetryExhaustionFails(t *testing.T) {
l := newLauncher(func(MemberSpec, int) (MemberResult, error) {
return MemberResult{}, ErrMemberTemporary
})
sw := newSwarmFor(t, l)
id, _ := sw.Spawn(Policy{}, "team", "teammate", "task", "")
waitFor(t, "task failed", func() bool {
task, ok := sw.Coordinator().Get(id)
return ok && task.Status == StatusFailed
})
if got := l.attemptCount(id); got != maxMemberRestarts+1 {
t.Fatalf("attempts = %d, want %d", got, maxMemberRestarts+1)
}
}
func TestPermanentErrorNoRetry(t *testing.T) {
l := newLauncher(func(MemberSpec, int) (MemberResult, error) {
return MemberResult{}, errPlain("hard failure")
})
sw := newSwarmFor(t, l)
id, _ := sw.Spawn(Policy{}, "team", "teammate", "task", "")
waitFor(t, "task failed", func() bool {
task, ok := sw.Coordinator().Get(id)
return ok && task.Status == StatusFailed
})
if got := l.attemptCount(id); got != 1 {
t.Fatalf("attempts = %d, want 1 (no retry on permanent error)", got)
}
task, _ := sw.Coordinator().Get(id)
if task.Err == "" {
t.Fatal("failed task should record an error message")
}
}
func TestHandoffDeliversNoteAndRetiresOriginal(t *testing.T) {
gate := make(chan struct{})
l := newLauncher(okFor)
l.gate = gate // keep the original member running so it is non-terminal
sw := newSwarmFor(t, l)
pol := Policy{Model: "m"}
origID, _ := sw.Spawn(pol, "team", "teammate", "original task", "/w")
waitFor(t, "original running", func() bool {
task, ok := sw.Coordinator().Get(origID)
return ok && task.Status == StatusRunning
})
newID, err := sw.Handoff(pol, "team", origID, "subagent", "please continue")
if err != nil {
t.Fatalf("Handoff: %v", err)
}
// Original retired.
orig, _ := sw.Coordinator().Get(origID)
if orig.Status != StatusHandedOff {
t.Fatalf("original status = %v, want handed-off", orig.Status)
}
// Note delivered to the new member's inbox.
msgs, err := sw.Mailbox().ReadAndConsume("team", newID)
if err != nil {
t.Fatalf("read new inbox: %v", err)
}
if len(msgs) != 1 || msgs[0].Body != "please continue" || msgs[0].Type != "handoff" {
t.Fatalf("handoff note not delivered: %+v", msgs)
}
// The new member carries the handoff note in its task and preserves cwd.
close(gate)
waitFor(t, "spec for new member", func() bool {
for _, s := range l.recorded() {
if s.ID == newID {
return true
}
}
return false
})
for _, s := range l.recorded() {
if s.ID == newID {
if s.Cwd != "/w" {
t.Fatalf("handoff lost cwd: %q", s.Cwd)
}
}
}
// A handoff of an already-terminal task is rejected.
waitFor(t, "new task done", func() bool {
task, ok := sw.Coordinator().Get(newID)
return ok && task.Status == StatusDone
})
if _, err := sw.Handoff(pol, "team", newID, "teammate", "again"); err == nil {
t.Fatal("handoff of a terminal task must fail")
}
}
func TestAdoptOrphans(t *testing.T) {
l := newLauncher(okFor)
sw := newSwarmFor(t, l)
// Simulate a crashed member: a running task in the coordinator whose owning
// agent has no live member in the team.
if _, err := sw.Coordinator().Register("orphan-1", "ghost", "team", "stranded work"); err != nil {
t.Fatalf("Register: %v", err)
}
_ = sw.Coordinator().SetStatus("orphan-1", StatusRunning)
adopted, err := sw.AdoptOrphans(Policy{Model: "m"}, "team", "subagent")
if err != nil {
t.Fatalf("AdoptOrphans: %v", err)
}
if len(adopted) != 1 || adopted[0] != "orphan-1" {
t.Fatalf("adopted = %v, want [orphan-1]", adopted)
}
// The orphan is relaunched under a fresh agent and completes.
waitFor(t, "orphan completed", func() bool {
task, ok := sw.Coordinator().Get("orphan-1")
return ok && task.Status == StatusDone
})
task, _ := sw.Coordinator().Get("orphan-1")
if task.AgentID == "ghost" {
t.Fatal("orphan should be reassigned to a new agent")
}
// A second adoption finds nothing (the task now has a live/owned outcome).
again, _ := sw.AdoptOrphans(Policy{}, "team", "subagent")
if len(again) != 0 {
t.Fatalf("second adoption = %v, want none", again)
}
}
func TestCollectScopesToTeam(t *testing.T) {
l := newLauncher(okFor)
sw := newSwarmFor(t, l)
a, _ := sw.Spawn(Policy{}, "alpha", "teammate", "ta", "")
_, _ = sw.Spawn(Policy{}, "beta", "teammate", "tb", "")
waitFor(t, "alpha done", func() bool {
task, ok := sw.Coordinator().Get(a)
return ok && task.Status == StatusDone
})
collected := sw.Collect("alpha")
if len(collected) != 1 || collected[0].Team != "alpha" {
t.Fatalf("Collect(alpha) = %+v, want one alpha task", collected)
}
}
func TestFuncLauncherRecoversPanic(t *testing.T) {
// A panic inside a member's Run must surface as that member's error, never
// escape the goroutine and crash the orchestrator.
l := FuncLauncher{Run: func(context.Context, MemberSpec) (MemberResult, error) {
panic("boom in member")
}}
h, err := l.Launch(context.Background(), MemberSpec{ID: "m1"})
if err != nil {
t.Fatalf("Launch: %v", err)
}
_, err = h.Wait()
if err == nil || !strings.Contains(err.Error(), "panicked") {
t.Fatalf("panic must surface as a member error, got %v", err)
}
}
func TestSpawnUnknownAgentType(t *testing.T) {
sw := newSwarmFor(t, newLauncher(okFor))
if _, err := sw.Spawn(Policy{}, "team", "does-not-exist", "task", ""); err == nil {
t.Fatal("Spawn with unknown agent type must error")
}
}
+15
View File
@@ -0,0 +1,15 @@
//go:build !windows
package swarm
import (
"errors"
"os"
)
// isLockContended reports whether an O_EXCL lock-create error means the lock is
// currently held (so the caller should retry/wait) rather than a hard failure.
// On POSIX an existing lock surfaces only as os.ErrExist.
func isLockContended(err error) bool {
return errors.Is(err, os.ErrExist)
}
+30
View File
@@ -0,0 +1,30 @@
//go:build windows
package swarm
import (
"errors"
"os"
"syscall"
)
// Windows error codes that an O_EXCL create (or a racing delete) can surface
// when another writer holds the lock file open, instead of ERROR_FILE_EXISTS.
// Treating them as "contended" makes the lock retry/wait rather than fail the
// whole send/read spuriously under contention.
const (
errorAccessDenied = syscall.Errno(5) // ERROR_ACCESS_DENIED
errorSharingViolation = syscall.Errno(32) // ERROR_SHARING_VIOLATION
)
// isLockContended reports whether an O_EXCL lock-create error means the lock is
// currently held (so the caller should retry/wait) rather than a hard failure.
// On Windows a held lock can appear as ERROR_FILE_EXISTS (mapped to os.ErrExist)
// but also as ERROR_ACCESS_DENIED / ERROR_SHARING_VIOLATION when the file is
// concurrently open or pending deletion.
func isLockContended(err error) bool {
if errors.Is(err, os.ErrExist) {
return true
}
return errors.Is(err, errorAccessDenied) || errors.Is(err, errorSharingViolation)
}
+383
View File
@@ -0,0 +1,383 @@
package swarm
import (
"encoding/json"
"errors"
"fmt"
"os"
"path/filepath"
"regexp"
"strings"
"sync/atomic"
"time"
)
// Mailbox is a per-agent, per-team message inbox persisted as a JSON array on
// disk at <baseDir>/<team>/inboxes/<agent>.json, written atomically under a
// lock file. It is hardened for untrusted input:
//
// - inbox files and their parent dirs are owner-only (0600/0700);
// - every write takes an exclusive lock file (bounded retry, stale-break);
// - message bodies and whole inboxes are size-capped (fail closed on oversize);
// - team/agent names are sanitized so a name can never escape the base dir;
// - malformed inbox JSON is rejected rather than silently reset (fail closed).
type Mailbox struct {
// BaseDir is the root under which <team>/inboxes/<agent>.json live. It must
// be set (NewMailbox enforces this).
BaseDir string
// MaxMessageBytes caps a single serialized message. 0 => defaultMaxMessageBytes.
MaxMessageBytes int
// MaxMessages caps the number of messages retained in one inbox. 0 =>
// defaultMaxMessages. Sends past the cap fail closed.
MaxMessages int
// LockTimeout bounds how long Send/MarkRead wait for the inbox lock.
LockTimeout time.Duration
}
const (
defaultMaxMessageBytes = 64 * 1024 // 64 KiB per message
defaultMaxMessages = 1000 // per inbox
defaultLockTimeout = 5 * time.Second
lockStaleAfter = 30 * time.Second
inboxFileMode = 0o600
inboxDirMode = 0o700
)
// ErrMailboxFull is returned when an inbox is at MaxMessages.
var ErrMailboxFull = errors.New("swarm: mailbox is full")
// ErrMessageTooLarge is returned when a single message exceeds MaxMessageBytes.
var ErrMessageTooLarge = errors.New("swarm: message exceeds size cap")
// Message is one inbox entry. Time is RFC3339 and set on Send if empty.
type Message struct {
From string `json:"from"`
To string `json:"to"`
Subject string `json:"subject,omitempty"`
Body string `json:"body"`
Type string `json:"type"`
Read bool `json:"read"`
Time string `json:"time,omitempty"`
}
// NewMailbox validates baseDir and returns a Mailbox with defaults applied.
func NewMailbox(baseDir string) (*Mailbox, error) {
if strings.TrimSpace(baseDir) == "" {
return nil, errors.New("swarm: mailbox base dir is required")
}
abs, err := filepath.Abs(baseDir)
if err != nil {
return nil, fmt.Errorf("swarm: resolve mailbox base dir: %w", err)
}
return &Mailbox{
BaseDir: abs,
MaxMessageBytes: defaultMaxMessageBytes,
MaxMessages: defaultMaxMessages,
LockTimeout: defaultLockTimeout,
}, nil
}
var nameUnsafe = regexp.MustCompile(`[^a-zA-Z0-9\-_]`)
// sanitizeName replaces any char outside [A-Za-z0-9-_] with '-', defaulting
// empty input to "default" and an all-unsafe result to "unknown". This
// guarantees the name is a single safe path segment (no '/', '\\', '.', '..').
func sanitizeName(name string) string {
trimmed := strings.TrimSpace(name)
if trimmed == "" {
trimmed = "default"
}
cleaned := nameUnsafe.ReplaceAllString(trimmed, "-")
if cleaned == "" {
return "unknown"
}
return cleaned
}
func (m *Mailbox) maxMessageBytes() int {
if m.MaxMessageBytes > 0 {
return m.MaxMessageBytes
}
return defaultMaxMessageBytes
}
func (m *Mailbox) maxMessages() int {
if m.MaxMessages > 0 {
return m.MaxMessages
}
return defaultMaxMessages
}
func (m *Mailbox) lockTimeout() time.Duration {
if m.LockTimeout > 0 {
return m.LockTimeout
}
return defaultLockTimeout
}
// inboxPath returns the sanitized, base-confined inbox path for (team, agent)
// and verifies it cannot escape BaseDir lexically (defense in depth on top of
// sanitize). Symlink-aware confinement happens in ensureInboxDir.
func (m *Mailbox) inboxPath(team, agent string) (string, error) {
dir := filepath.Join(m.BaseDir, sanitizeName(team), "inboxes")
path := filepath.Join(dir, sanitizeName(agent)+".json")
if within, err := pathWithin(m.BaseDir, path); err != nil || !within {
return "", fmt.Errorf("swarm: inbox path escapes base dir (team=%q agent=%q)", team, agent)
}
return path, nil
}
// pathWithin reports whether target is base or lives under it (lexically).
func pathWithin(base, target string) (bool, error) {
rel, err := filepath.Rel(base, target)
if err != nil {
return false, err
}
if rel == ".." || strings.HasPrefix(rel, ".."+string(os.PathSeparator)) {
return false, nil
}
return true, nil
}
// ensureInboxDir creates the inbox's parent directories owner-only, then
// verifies — after resolving symlinks — that the real directory stays within the
// real base dir. This defends against a symlink planted under BaseDir (e.g. a
// pre-existing "inboxes" -> /etc) that the lexical inboxPath check cannot catch.
// Only after confinement is confirmed are the directories tightened to 0700, so
// a chmod can never land outside the base via a symlink.
func (m *Mailbox) ensureInboxDir(inboxPath string) error {
dir := filepath.Dir(inboxPath) // <base>/<team>/inboxes
if err := os.MkdirAll(dir, inboxDirMode); err != nil {
return fmt.Errorf("swarm: create inbox dir: %w", err)
}
realBase, err := filepath.EvalSymlinks(m.BaseDir)
if err != nil {
return fmt.Errorf("swarm: resolve base dir: %w", err)
}
realDir, err := filepath.EvalSymlinks(dir)
if err != nil {
return fmt.Errorf("swarm: resolve inbox dir: %w", err)
}
if within, err := pathWithin(realBase, realDir); err != nil || !within {
return fmt.Errorf("swarm: inbox dir escapes base dir via symlink")
}
// MkdirAll leaves an existing dir's permissions untouched; tighten the team
// and inboxes dirs to owner-only so mailbox contents are never group/world
// readable even if a dir pre-existed with broad perms.
_ = os.Chmod(dir, inboxDirMode)
_ = os.Chmod(filepath.Dir(dir), inboxDirMode) // <base>/<team>
return nil
}
// Send appends msg to the recipient's inbox under the team, taking the inbox
// lock. It fails closed on oversize messages or a full inbox.
//
// Each send rewrites the whole inbox (read all → append one → write all), so it
// is O(N) in the inbox size. That is intentional given the small bounds
// (MaxMessages × MaxMessageBytes); if those caps are raised substantially,
// switch to an append-only log with periodic compaction.
func (m *Mailbox) Send(team, recipient string, msg Message) error {
path, err := m.inboxPath(team, recipient)
if err != nil {
return err
}
if msg.Type == "" {
msg.Type = "message"
}
if msg.To == "" {
msg.To = sanitizeName(recipient)
}
if msg.Time == "" {
msg.Time = time.Now().UTC().Format(time.RFC3339)
}
encoded, err := json.Marshal(msg)
if err != nil {
return fmt.Errorf("swarm: encode message: %w", err)
}
if len(encoded) > m.maxMessageBytes() {
return fmt.Errorf("%w: %d > %d bytes", ErrMessageTooLarge, len(encoded), m.maxMessageBytes())
}
if err := m.ensureInboxDir(path); err != nil {
return err
}
release, err := acquireLock(path+".lock", m.lockTimeout())
if err != nil {
return err
}
defer release()
messages, err := m.readLocked(path)
if err != nil {
return err
}
if len(messages) >= m.maxMessages() {
return fmt.Errorf("%w: %d messages", ErrMailboxFull, len(messages))
}
messages = append(messages, msg)
return atomicWriteJSON(path, messages)
}
// ReadAndConsume reads the recipient's inbox and marks every previously-unread
// message read under one lock, returning the messages with their read flags as
// they were BEFORE this read (so the caller can tell which were new). The inbox
// is empty (no error) if it does not yet exist; malformed inbox JSON is reported
// as an error (fail closed) and nothing is written.
func (m *Mailbox) ReadAndConsume(team, recipient string) ([]Message, error) {
path, err := m.inboxPath(team, recipient)
if err != nil {
return nil, err
}
// The lock file lives beside the inbox; create the dir so it can be taken
// even on first read of a not-yet-existing inbox.
if err := m.ensureInboxDir(path); err != nil {
return nil, err
}
release, err := acquireLock(path+".lock", m.lockTimeout())
if err != nil {
return nil, err
}
defer release()
messages, err := m.readLocked(path)
if err != nil {
return nil, err
}
if len(messages) == 0 {
return messages, nil
}
// Snapshot pre-consume state to return, then flip unread -> read on disk.
snapshot := make([]Message, len(messages))
copy(snapshot, messages)
changed := false
for i := range messages {
if !messages[i].Read {
messages[i].Read = true
changed = true
}
}
if changed {
if err := atomicWriteJSON(path, messages); err != nil {
return nil, err
}
}
return snapshot, nil
}
// readLocked reads and decodes an inbox file. Callers that mutate must hold the
// lock; pure readers may call it directly (a torn read surfaces as a decode
// error and is retried by the caller, never silently treated as empty).
func (m *Mailbox) readLocked(path string) ([]Message, error) {
raw, err := os.ReadFile(path)
if err != nil {
if errors.Is(err, os.ErrNotExist) {
return nil, nil
}
return nil, fmt.Errorf("swarm: read inbox: %w", err)
}
if len(raw) == 0 {
return nil, nil
}
// Reject an oversized inbox file outright rather than decoding untrusted,
// unbounded JSON into memory. The bound is the cap on total message bytes
// (maxMessages × maxMessageBytes) doubled to allow for pretty-printed JSON
// (indentation, field names, brackets) overhead versus the compact bytes the
// per-message cap measures.
if maxFile := m.maxMessages() * m.maxMessageBytes() * 2; len(raw) > maxFile {
return nil, fmt.Errorf("swarm: inbox file too large (%d bytes)", len(raw))
}
var messages []Message
if err := json.Unmarshal(raw, &messages); err != nil {
return nil, fmt.Errorf("swarm: malformed inbox %s: %w", filepath.Base(path), err)
}
for i := range messages {
if messages[i].Type == "" {
messages[i].Type = "message"
}
}
return messages, nil
}
// atomicWriteJSON writes data as pretty JSON to a sibling temp file (0600) then
// renames it over path, so a reader never observes a partial write.
func atomicWriteJSON(path string, data any) error {
encoded, err := json.MarshalIndent(data, "", " ")
if err != nil {
return fmt.Errorf("swarm: encode inbox: %w", err)
}
dir := filepath.Dir(path)
tmp, err := os.CreateTemp(dir, ".inbox-*.tmp")
if err != nil {
return fmt.Errorf("swarm: create temp inbox: %w", err)
}
tmpName := tmp.Name()
defer os.Remove(tmpName) // no-op after a successful rename
if err := tmp.Chmod(inboxFileMode); err != nil {
tmp.Close()
return fmt.Errorf("swarm: chmod temp inbox: %w", err)
}
if _, err := tmp.Write(encoded); err != nil {
tmp.Close()
return fmt.Errorf("swarm: write temp inbox: %w", err)
}
if err := tmp.Close(); err != nil {
return fmt.Errorf("swarm: close temp inbox: %w", err)
}
if err := os.Rename(tmpName, path); err != nil {
return fmt.Errorf("swarm: commit inbox: %w", err)
}
return nil
}
// lockSeq makes each lock token unique within the process even when two
// goroutines acquire in the same nanosecond.
var lockSeq atomic.Uint64
// acquireLock takes an exclusive lock by creating lockPath with O_EXCL. It
// retries with a short backoff until timeout, breaking a lock whose file is
// older than lockStaleAfter (so a crashed holder cannot deadlock the swarm).
//
// Each holder writes a unique token into the lock file and release is
// OWNERSHIP-AWARE: it removes the lock only if it still holds our token. After a
// stale-break another writer may legitimately own the lock; unconditionally
// removing it would let a third writer in and create split-brain writers that
// can corrupt or drop mailbox updates.
func acquireLock(lockPath string, timeout time.Duration) (func(), error) {
token := fmt.Sprintf("%d-%d-%d", os.Getpid(), time.Now().UnixNano(), lockSeq.Add(1))
deadline := time.Now().Add(timeout)
for {
f, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, inboxFileMode)
if err == nil {
// Record our ownership token; the exclusive create already holds the
// lock, so a write error only weakens diagnostics, not safety.
_, _ = f.WriteString(token)
f.Close()
var released bool
return func() {
if released {
return
}
released = true
if data, rerr := os.ReadFile(lockPath); rerr == nil && string(data) == token {
os.Remove(lockPath)
}
}, nil
}
if !isLockContended(err) {
return nil, fmt.Errorf("swarm: acquire lock: %w", err)
}
// Lock held: break it if stale, otherwise wait. The break is conditional
// on the file's content being unchanged since we observed it, so we do not
// delete a lock another writer rotated in between.
if info, statErr := os.Stat(lockPath); statErr == nil && time.Since(info.ModTime()) > lockStaleAfter {
stale, _ := os.ReadFile(lockPath)
if data, rerr := os.ReadFile(lockPath); rerr == nil && string(data) == string(stale) {
os.Remove(lockPath)
}
continue
}
if time.Now().After(deadline) {
return nil, fmt.Errorf("swarm: timed out acquiring lock %s", filepath.Base(lockPath))
}
time.Sleep(20 * time.Millisecond)
}
}
+276
View File
@@ -0,0 +1,276 @@
package swarm
import (
"errors"
"os"
"path/filepath"
"runtime"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
)
func newTestMailbox(t *testing.T) *Mailbox {
t.Helper()
mb, err := NewMailbox(t.TempDir())
if err != nil {
t.Fatalf("NewMailbox: %v", err)
}
return mb
}
func TestSanitizeName(t *testing.T) {
cases := map[string]string{
"": "default",
" ": "default",
"team-1_ok": "team-1_ok",
"../../etc": "------etc",
"a/b\\c": "a-b-c",
"team name!": "team-name-",
"!!!": "---",
"plain": "plain",
}
for in, want := range cases {
if got := sanitizeName(in); got != want {
t.Errorf("sanitizeName(%q) = %q, want %q", in, got, want)
}
}
}
func TestMailboxSendAndConsume(t *testing.T) {
mb := newTestMailbox(t)
if err := mb.Send("team", "bob", Message{From: "alice", Body: "hello"}); err != nil {
t.Fatalf("Send: %v", err)
}
if err := mb.Send("team", "bob", Message{From: "alice", Body: "again"}); err != nil {
t.Fatalf("Send 2: %v", err)
}
msgs, err := mb.ReadAndConsume("team", "bob")
if err != nil {
t.Fatalf("ReadAndConsume: %v", err)
}
if len(msgs) != 2 {
t.Fatalf("got %d messages, want 2", len(msgs))
}
// First read returns them as unread (pre-consume snapshot), with defaults set.
if msgs[0].Read || msgs[0].Type != "message" || msgs[0].To != "bob" || msgs[0].Time == "" {
t.Fatalf("message defaults wrong: %+v", msgs[0])
}
// Second read shows them already consumed (read).
msgs2, err := mb.ReadAndConsume("team", "bob")
if err != nil {
t.Fatalf("ReadAndConsume 2: %v", err)
}
for _, m := range msgs2 {
if !m.Read {
t.Fatalf("message should be read after consume: %+v", m)
}
}
}
func TestMailboxReadMissingIsEmpty(t *testing.T) {
mb := newTestMailbox(t)
msgs, err := mb.ReadAndConsume("team", "nobody")
if err != nil {
t.Fatalf("ReadAndConsume missing: %v", err)
}
if len(msgs) != 0 {
t.Fatalf("missing inbox should be empty, got %d", len(msgs))
}
}
func TestMailboxOversizeMessageRejected(t *testing.T) {
mb := newTestMailbox(t)
mb.MaxMessageBytes = 64
err := mb.Send("team", "bob", Message{From: "a", Body: strings.Repeat("x", 1000)})
if !errors.Is(err, ErrMessageTooLarge) {
t.Fatalf("oversize Send err = %v, want ErrMessageTooLarge", err)
}
// Nothing should have been written.
msgs, _ := mb.ReadAndConsume("team", "bob")
if len(msgs) != 0 {
t.Fatalf("oversize Send must not persist, got %d messages", len(msgs))
}
}
func TestMailboxFullRejected(t *testing.T) {
mb := newTestMailbox(t)
mb.MaxMessages = 2
for i := 0; i < 2; i++ {
if err := mb.Send("team", "bob", Message{From: "a", Body: "ok"}); err != nil {
t.Fatalf("Send %d: %v", i, err)
}
}
if err := mb.Send("team", "bob", Message{From: "a", Body: "overflow"}); !errors.Is(err, ErrMailboxFull) {
t.Fatalf("full Send err = %v, want ErrMailboxFull", err)
}
}
func TestMailboxMalformedFailsClosed(t *testing.T) {
mb := newTestMailbox(t)
// Write a valid message so the inbox path exists, then corrupt the file.
if err := mb.Send("team", "bob", Message{From: "a", Body: "ok"}); err != nil {
t.Fatalf("Send: %v", err)
}
path, err := mb.inboxPath("team", "bob")
if err != nil {
t.Fatalf("inboxPath: %v", err)
}
if err := os.WriteFile(path, []byte("{not json"), 0o600); err != nil {
t.Fatalf("corrupt inbox: %v", err)
}
if _, err := mb.ReadAndConsume("team", "bob"); err == nil {
t.Fatal("malformed inbox must fail closed (error), not be treated as empty")
}
// A subsequent Send also fails closed rather than clobbering unknown data.
if err := mb.Send("team", "bob", Message{From: "a", Body: "x"}); err == nil {
t.Fatal("Send into a malformed inbox must fail closed")
}
}
func TestMailboxPathConfinement(t *testing.T) {
mb := newTestMailbox(t)
// A traversal-style name is sanitized into a single safe segment; the inbox
// stays under BaseDir.
path, err := mb.inboxPath("../../evil", "../escape")
if err != nil {
t.Fatalf("inboxPath: %v", err)
}
rel, err := filepath.Rel(mb.BaseDir, path)
if err != nil || strings.HasPrefix(rel, "..") {
t.Fatalf("inbox path escaped base dir: base=%q path=%q rel=%q", mb.BaseDir, path, rel)
}
}
func TestMailboxFileMode(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("unix file modes")
}
mb := newTestMailbox(t)
if err := mb.Send("team", "bob", Message{From: "a", Body: "ok"}); err != nil {
t.Fatalf("Send: %v", err)
}
path, _ := mb.inboxPath("team", "bob")
info, err := os.Stat(path)
if err != nil {
t.Fatalf("stat: %v", err)
}
if perm := info.Mode().Perm(); perm != 0o600 {
t.Fatalf("inbox file mode = %o, want 600", perm)
}
}
func TestMailboxTightensDirPerms(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("unix file modes")
}
base := t.TempDir()
// Pre-create the team + inboxes dirs with broad perms; ensureInboxDir must
// tighten them to 0700 despite MkdirAll leaving existing dirs untouched.
broad := filepath.Join(base, "team", "inboxes")
if err := os.MkdirAll(broad, 0o755); err != nil {
t.Fatalf("pre-create: %v", err)
}
mb, err := NewMailbox(base)
if err != nil {
t.Fatalf("NewMailbox: %v", err)
}
if err := mb.Send("team", "bob", Message{From: "a", Body: "ok"}); err != nil {
t.Fatalf("Send: %v", err)
}
for _, dir := range []string{broad, filepath.Join(base, "team")} {
info, err := os.Stat(dir)
if err != nil {
t.Fatalf("stat %s: %v", dir, err)
}
if perm := info.Mode().Perm(); perm != 0o700 {
t.Fatalf("dir %s mode = %o, want 700", dir, perm)
}
}
}
func TestMailboxRejectsSymlinkedInboxDir(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("symlink semantics differ")
}
base := t.TempDir()
outside := t.TempDir() // a dir OUTSIDE base
// Plant a symlink: <base>/team/inboxes -> <outside>. A lexical check passes,
// but the symlink-aware confinement must reject the write.
if err := os.MkdirAll(filepath.Join(base, "team"), 0o700); err != nil {
t.Fatalf("mkdir team: %v", err)
}
if err := os.Symlink(outside, filepath.Join(base, "team", "inboxes")); err != nil {
t.Skipf("symlink unsupported: %v", err)
}
mb, err := NewMailbox(base)
if err != nil {
t.Fatalf("NewMailbox: %v", err)
}
if err := mb.Send("team", "bob", Message{From: "a", Body: "escape"}); err == nil {
t.Fatal("Send into a symlinked-out inbox dir must fail closed")
}
// And nothing was written outside the base.
if entries, _ := os.ReadDir(outside); len(entries) != 0 {
t.Fatalf("write escaped base via symlink: %d entries in %s", len(entries), outside)
}
}
func TestMailboxLockReleaseIsOwnershipAware(t *testing.T) {
mb := newTestMailbox(t)
path, _ := mb.inboxPath("team", "bob")
if err := mb.ensureInboxDir(path); err != nil {
t.Fatalf("ensureInboxDir: %v", err)
}
lockPath := path + ".lock"
// Writer A acquires the lock.
releaseA, err := acquireLock(lockPath, time.Second)
if err != nil {
t.Fatalf("acquire A: %v", err)
}
// Simulate a stale-break + takeover by writer B: overwrite the lock content
// with B's token (as a fresh acquire after a break would).
if err := os.WriteFile(lockPath, []byte("writer-B-token"), 0o600); err != nil {
t.Fatalf("simulate B takeover: %v", err)
}
// A's release must NOT delete B's lock (ownership-aware).
releaseA()
if _, err := os.Stat(lockPath); err != nil {
t.Fatalf("A's release deleted B's lock (split-brain): %v", err)
}
// B's own release removes it.
os.Remove(lockPath)
}
func TestMailboxConcurrentSends(t *testing.T) {
mb := newTestMailbox(t)
// High fan-out + a generous lock timeout so heavy contention is exercised
// (and a lock regression — e.g. a Windows sharing-violation treated as fatal —
// would surface as lost/failed messages here rather than flaking).
mb.LockTimeout = 10 * time.Second
const n = 200
var wg sync.WaitGroup
var failures atomic.Int32
for i := 0; i < n; i++ {
wg.Add(1)
go func() {
defer wg.Done()
if err := mb.Send("team", "bob", Message{From: "a", Body: "concurrent"}); err != nil {
failures.Add(1)
}
}()
}
wg.Wait()
if f := failures.Load(); f != 0 {
t.Fatalf("%d concurrent Sends failed (lock contention not handled)", f)
}
msgs, err := mb.ReadAndConsume("team", "bob")
if err != nil {
t.Fatalf("ReadAndConsume: %v", err)
}
if len(msgs) != n {
t.Fatalf("concurrent sends lost messages: got %d, want %d", len(msgs), n)
}
}
+97
View File
@@ -0,0 +1,97 @@
package swarm
import (
"context"
"errors"
"fmt"
)
// MemberSpec is the fully-resolved instruction to launch one member. The swarm
// builds it from a Definition + task; a MemberLauncher turns it into a running
// agent. PermissionMode/Model carry the orchestrator's policy so the launcher
// can never grant a member more authority than its parent.
type MemberSpec struct {
// ID identifies the running member (its agent id). For a freshly spawned
// member this equals TaskID; after an orphan adoption a new agent (new ID)
// takes over the original TaskID.
ID string
TaskID string
AgentType string
Team string
Task string
Cwd string
Model string // resolved (never the "inherit" sentinel)
PermissionMode string
SystemPrompt string
}
// MemberResult is what a finished member returns.
type MemberResult struct {
Result string
SessionID string
}
// MemberHandle is a running member. Wait blocks until the member exits and
// returns its result or error; ID identifies the member.
type MemberHandle interface {
ID() string
Wait() (MemberResult, error)
}
// MemberLauncher turns a MemberSpec into a running member. Production wires two
// implementations behind this seam — one over internal/specialist.Executor
// (direct) and one over internal/daemon.Pool (when a daemon is running) — while
// tests inject a fake. Keeping the seam here means the swarm core has no
// compile-time dependency on either heavy subsystem.
type MemberLauncher interface {
Launch(ctx context.Context, spec MemberSpec) (MemberHandle, error)
}
// RunFunc executes a member to completion. Both production launchers and tests
// supply one; FuncLauncher adapts it to the MemberLauncher interface by running
// it in a goroutine.
type RunFunc func(ctx context.Context, spec MemberSpec) (MemberResult, error)
// FuncLauncher is the single MemberLauncher implementation: it runs a RunFunc in
// a goroutine and exposes a handle. Production builds one whose RunFunc calls
// specialist.Executor.Run (or submits to daemon.Pool); tests build one whose
// RunFunc is a stub.
type FuncLauncher struct {
Run RunFunc
}
// Launch starts the RunFunc for spec and returns a handle to await it.
func (l FuncLauncher) Launch(ctx context.Context, spec MemberSpec) (MemberHandle, error) {
if l.Run == nil {
return nil, errors.New("swarm: FuncLauncher requires a Run func")
}
h := &funcHandle{id: spec.ID, done: make(chan struct{})}
go func() {
defer close(h.done)
// A panic in a member's Run must fail only that member, never crash the
// orchestrator process. Recover and surface it as the member's error.
defer func() {
if r := recover(); r != nil {
h.res = MemberResult{}
h.err = fmt.Errorf("swarm: member %s panicked: %v", spec.ID, r)
}
}()
h.res, h.err = l.Run(ctx, spec)
}()
return h, nil
}
type funcHandle struct {
id string
done chan struct{}
res MemberResult
err error
}
func (h *funcHandle) ID() string { return h.id }
// Wait blocks until the member exits and returns its result/error.
func (h *funcHandle) Wait() (MemberResult, error) {
<-h.done
return h.res, h.err
}
+360
View File
@@ -0,0 +1,360 @@
package swarm
import (
"context"
"errors"
"fmt"
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
)
// defaultMaxTeamSize bounds concurrently-running members per team. Spawns past
// the cap queue and launch as slots free up (mirrors daemon.Pool's bounded,
// on-demand worker model).
const defaultMaxTeamSize = 8
// maxMemberRestarts bounds automatic relaunches of a member that exits with a
// temporary error, mirroring daemon.Pool's bounded backoff retries.
const maxMemberRestarts = 2
// Options configures a Swarm.
type Options struct {
// BaseDir is the mailbox root (required).
BaseDir string
// Launcher turns specs into running members (required).
Launcher MemberLauncher
// Registry is the agent roster; nil => a fresh built-in registry.
Registry *Registry
// Coordinator tracks tasks; nil => a fresh coordinator.
Coordinator *Coordinator
// MaxTeamSize caps concurrent members per team; 0 => defaultMaxTeamSize.
MaxTeamSize int
// Definitions are user-defined agents to add to the roster (extending or
// overriding the built-ins). Empty agent types are rejected.
Definitions []Definition
// Context is the parent context for all members; members outlive the
// per-call tool context (a spawn returns immediately). nil => Background.
// Close cancels the derived context, stopping every member.
Context context.Context
}
// Policy is the orchestrator's live model + permission mode at spawn time. Every
// member inherits it: a member never runs on a different model than requested,
// nor with a more permissive mode than its parent. It is passed per call (not
// stored on the Swarm) because the orchestrator's model/mode can change
// mid-session (e.g. an escalate-model).
type Policy struct {
Model string
PermissionMode string
}
// Swarm is the façade the swarm tools call. It owns the agent roster, the task
// coordinator, the mailbox, and the live teams, and enforces that every member
// inherits the orchestrator's model + permission mode.
type Swarm struct {
registry *Registry
coord *Coordinator
mailbox *Mailbox
launcher MemberLauncher
maxTeamSize int
baseCtx context.Context
cancel context.CancelFunc
mu sync.Mutex
teams map[string]*Team
taskCwd map[string]string // taskID -> cwd, for handoff/adoption relaunch
idSeq atomic.Uint64
}
// Team is a named set of concurrently-running members with a bounded slot count
// and a FIFO queue for overflow spawns.
type Team struct {
Name string
mu sync.Mutex
members map[string]*Member
queue []MemberSpec
running int
maxSize int
}
// Member is one running agent within a team.
type Member struct {
ID string
AgentType string
TaskID string
handle MemberHandle
restarts int
}
// New validates options and returns a Swarm.
func New(opts Options) (*Swarm, error) {
if opts.Launcher == nil {
return nil, errors.New("swarm: a MemberLauncher is required")
}
mb, err := NewMailbox(opts.BaseDir)
if err != nil {
return nil, err
}
registry := opts.Registry
if registry == nil {
registry = NewRegistry()
}
for _, def := range opts.Definitions {
if err := registry.Register(def); err != nil {
return nil, err
}
}
coord := opts.Coordinator
if coord == nil {
coord = NewCoordinator()
}
maxTeam := opts.MaxTeamSize
if maxTeam <= 0 {
maxTeam = defaultMaxTeamSize
}
parent := opts.Context
if parent == nil {
parent = context.Background()
}
ctx, cancel := context.WithCancel(parent)
return &Swarm{
registry: registry,
coord: coord,
mailbox: mb,
launcher: opts.Launcher,
maxTeamSize: maxTeam,
baseCtx: ctx,
cancel: cancel,
teams: map[string]*Team{},
taskCwd: map[string]string{},
}, nil
}
// Close cancels every running member's context and releases resources. It is
// safe to call more than once.
func (s *Swarm) Close() {
if s.cancel != nil {
s.cancel()
}
}
// rememberCwd records a task's working dir so a handoff/adoption relaunch keeps it.
func (s *Swarm) rememberCwd(taskID, cwd string) {
s.mu.Lock()
s.taskCwd[taskID] = cwd
s.mu.Unlock()
}
// cwdFor returns a task's recorded working dir ("" if none).
func (s *Swarm) cwdFor(taskID string) string {
s.mu.Lock()
defer s.mu.Unlock()
return s.taskCwd[taskID]
}
// Registry exposes the roster (for tool listing / user-defined registration).
func (s *Swarm) Registry() *Registry { return s.registry }
// Coordinator exposes the task registry (for status/collect).
func (s *Swarm) Coordinator() *Coordinator { return s.coord }
// Mailbox exposes the message store (for send/inbox tools).
func (s *Swarm) Mailbox() *Mailbox { return s.mailbox }
// team returns the named team, creating it on first use.
func (s *Swarm) team(name string) *Team {
key := sanitizeName(name)
s.mu.Lock()
defer s.mu.Unlock()
t, ok := s.teams[key]
if !ok {
t = &Team{Name: key, members: map[string]*Member{}, maxSize: s.maxTeamSize}
s.teams[key] = t
}
return t
}
// nextID mints a stable, unique member/task id within the process.
func (s *Swarm) nextID(agentType string) string {
n := s.idSeq.Add(1)
return sanitizeName(agentType) + "-" + strconv.FormatUint(n, 10)
}
// resolveModel folds the "inherit" sentinel into the orchestrator's model so a
// launched member never silently runs on a different model.
func resolveModel(pol Policy, def Definition) string {
if def.Model == "" || def.Model == modelInherit {
return pol.Model
}
return def.Model
}
// resolvePermissionMode never widens authority: a member uses its definition's
// mode only if it is not more permissive than the orchestrator's. Unknown modes
// fall back to the orchestrator's mode (fail closed).
func resolvePermissionMode(pol Policy, def Definition) string {
if def.PermissionMode == "" {
return pol.PermissionMode
}
if permissionRank(def.PermissionMode) > permissionRank(pol.PermissionMode) {
// Definition asks for more than the parent has — clamp to the parent.
return pol.PermissionMode
}
return def.PermissionMode
}
// Runtime permission modes (mirror internal/agent.PermissionMode without
// importing it, to avoid an import cycle). These are the actual values that flow
// through tools.RunOptions.PermissionMode — NOT the TUI display names.
const (
permissionModeAsk = "ask" // prompts for every tool (most restrictive)
permissionModeSpecDraft = "spec-draft" // spec-drafting only
permissionModeAuto = "auto" // auto-approve low-risk
permissionModeUnsafe = "unsafe" // approve everything (most permissive)
)
// permissionRank orders permission modes from least to most permissive so the
// swarm can clamp a member to no more than its parent. Unknown/empty modes rank
// as the strictest (0) so they never accidentally widen access. spec-draft
// (which scopes a child to spec tooling) ranks below ask: it never grants more
// authority than prompting-for-everything, so a spec-draft parent clamps members
// hardest among the known modes.
func permissionRank(mode string) int {
switch strings.TrimSpace(mode) {
case permissionModeSpecDraft:
return 1
case permissionModeAsk:
return 2
case permissionModeAuto:
return 3
case permissionModeUnsafe:
return 4
default:
return 0
}
}
// admit reserves a running slot for a spec, or queues it. It reports whether the
// caller should launch immediately.
func (t *Team) admit(spec MemberSpec) (launchNow bool) {
t.mu.Lock()
defer t.mu.Unlock()
if t.running < t.maxSize {
t.running++
return true
}
t.queue = append(t.queue, spec)
return false
}
// onExit releases a slot after a member exits and returns the next queued spec
// (already re-reserving its slot) if one is waiting.
func (t *Team) onExit() (MemberSpec, bool) {
t.mu.Lock()
defer t.mu.Unlock()
if t.running > 0 {
t.running--
}
if len(t.queue) > 0 && t.running < t.maxSize {
next := t.queue[0]
t.queue = t.queue[1:]
t.running++
return next, true
}
return MemberSpec{}, false
}
func (t *Team) addMember(m *Member) {
t.mu.Lock()
t.members[m.ID] = m
t.mu.Unlock()
}
func (t *Team) removeMember(id string) {
t.mu.Lock()
delete(t.members, id)
t.mu.Unlock()
}
// liveAgents returns the set of agent ids that currently have a live member, so
// orphan detection can tell which tasks have lost their owner.
func (t *Team) liveAgents() map[string]struct{} {
t.mu.Lock()
defer t.mu.Unlock()
live := make(map[string]struct{}, len(t.members))
for _, m := range t.members {
live[m.ID] = struct{}{}
}
return live
}
// QueueDepth reports how many specs are waiting for a slot (status/tests).
func (t *Team) QueueDepth() int {
t.mu.Lock()
defer t.mu.Unlock()
return len(t.queue)
}
// Running reports how many members are currently occupying a slot.
func (t *Team) Running() int {
t.mu.Lock()
defer t.mu.Unlock()
return t.running
}
// ErrMemberTemporary marks a member failure as transient (worker died, transient
// I/O) so the swarm relaunches it within maxMemberRestarts, mirroring daemon
// EXIT_TEMPFAIL semantics. Permanent failures are not retried.
var ErrMemberTemporary = errors.New("swarm: temporary member failure")
// isRetryable reports whether a member error should trigger a bounded relaunch.
func isRetryable(err error) bool {
if err == nil {
return false
}
if errors.Is(err, ErrMemberTemporary) {
return true
}
var t interface{ Temporary() bool }
if errors.As(err, &t) {
return t.Temporary()
}
return false
}
// buildSpec resolves a definition + task into a launchable spec under the
// orchestrator's policy. memberID identifies the running agent; taskID is the
// coordinator task it serves (equal for a fresh spawn, distinct after adoption).
func (s *Swarm) buildSpec(pol Policy, memberID, taskID, team string, def Definition, task, cwd string) MemberSpec {
prompt := ""
if def.SystemPrompt != nil {
prompt = def.SystemPrompt(PromptContext{Team: team, Task: task})
}
return MemberSpec{
ID: memberID,
TaskID: taskID,
AgentType: def.AgentType,
Team: team,
Task: task,
Cwd: cwd,
Model: resolveModel(pol, def),
PermissionMode: resolvePermissionMode(pol, def),
SystemPrompt: prompt,
}
}
// memberError wraps a member's failure for coordinator recording.
func memberError(err error) string {
if err == nil {
return ""
}
return fmt.Sprintf("%v", err)
}
// nowRFC3339 is a tiny helper for mailbox handoff notes.
func nowRFC3339() string { return time.Now().UTC().Format(time.RFC3339) }
+91
View File
@@ -0,0 +1,91 @@
package swarm
import "testing"
func TestResolveModelInherit(t *testing.T) {
pol := Policy{Model: "orchestrator-model"}
if got := resolveModel(pol, Definition{Model: modelInherit}); got != "orchestrator-model" {
t.Fatalf("inherit model = %q, want orchestrator-model", got)
}
if got := resolveModel(pol, Definition{Model: ""}); got != "orchestrator-model" {
t.Fatalf("empty model = %q, want orchestrator-model", got)
}
if got := resolveModel(pol, Definition{Model: "explicit"}); got != "explicit" {
t.Fatalf("explicit model = %q, want explicit", got)
}
}
func TestResolvePermissionModeNeverWidens(t *testing.T) {
parent := Policy{PermissionMode: permissionModeAuto} // "auto"
// A definition asking for MORE than the parent (unsafe > auto) is clamped down.
if got := resolvePermissionMode(parent, Definition{PermissionMode: permissionModeUnsafe}); got != permissionModeAuto {
t.Fatalf("clamp = %q, want auto (never widen)", got)
}
// A definition asking for LESS keeps its stricter mode (ask < auto).
if got := resolvePermissionMode(parent, Definition{PermissionMode: permissionModeAsk}); got != permissionModeAsk {
t.Fatalf("stricter = %q, want ask", got)
}
// Empty inherits the parent.
if got := resolvePermissionMode(parent, Definition{PermissionMode: ""}); got != permissionModeAuto {
t.Fatalf("empty = %q, want auto", got)
}
// An unsafe parent still honors a stricter definition (auto < unsafe).
if got := resolvePermissionMode(Policy{PermissionMode: permissionModeUnsafe}, Definition{PermissionMode: permissionModeAuto}); got != permissionModeAuto {
t.Fatalf("unsafe parent + auto def = %q, want auto (stricter honored)", got)
}
}
func TestTeamAdmitAndQueue(t *testing.T) {
tm := &Team{Name: "t", members: map[string]*Member{}, maxSize: 2}
if !tm.admit(MemberSpec{ID: "a"}) || !tm.admit(MemberSpec{ID: "b"}) {
t.Fatal("first two admits should launch now")
}
if tm.admit(MemberSpec{ID: "c"}) {
t.Fatal("third admit should queue, not launch")
}
if tm.Running() != 2 || tm.QueueDepth() != 1 {
t.Fatalf("running=%d queue=%d, want 2/1", tm.Running(), tm.QueueDepth())
}
// One exit drains the single queued spec (re-reserving its slot).
next, ok := tm.onExit()
if !ok || next.ID != "c" {
t.Fatalf("onExit = %+v ok=%v, want spec c", next, ok)
}
if tm.Running() != 2 || tm.QueueDepth() != 0 {
t.Fatalf("after drain running=%d queue=%d, want 2/0", tm.Running(), tm.QueueDepth())
}
// Two more exits empty the team (no queue left).
if _, ok := tm.onExit(); ok {
t.Fatal("onExit with empty queue should not return a spec")
}
if _, ok := tm.onExit(); ok {
t.Fatal("onExit with empty queue should not return a spec")
}
if tm.Running() != 0 {
t.Fatalf("running = %d, want 0", tm.Running())
}
}
func TestIsRetryable(t *testing.T) {
if isRetryable(nil) {
t.Fatal("nil is not retryable")
}
if !isRetryable(ErrMemberTemporary) {
t.Fatal("ErrMemberTemporary must be retryable")
}
if isRetryable(errPlain("permanent")) {
t.Fatal("a plain error is not retryable")
}
if !isRetryable(tempErr{}) {
t.Fatal("an error with Temporary()==true must be retryable")
}
}
type errPlain string
func (e errPlain) Error() string { return string(e) }
type tempErr struct{}
func (tempErr) Error() string { return "temp" }
func (tempErr) Temporary() bool { return true }
+402
View File
@@ -0,0 +1,402 @@
package swarm
import (
"context"
"errors"
"fmt"
"sort"
"strings"
"github.com/Gitlawb/zero/internal/tools"
)
// Tool names. All swarm tools are additive and only act when invoked; the
// existing single "Task" tool is unchanged.
const (
SpawnToolName = "swarm_spawn"
SendToolName = "swarm_send"
InboxToolName = "swarm_inbox"
StatusToolName = "swarm_status"
HandoffToolName = "swarm_handoff"
CollectToolName = "swarm_collect"
)
// RegisterTools registers the six swarm tools backed by one shared Swarm. The
// caller owns the Swarm's lifetime (call Swarm.Close on shutdown).
func RegisterTools(registry *tools.Registry, sw *Swarm) {
registry.Register(&spawnTool{sw: sw})
registry.Register(&sendTool{sw: sw})
registry.Register(&inboxTool{sw: sw})
registry.Register(&statusTool{sw: sw})
registry.Register(&handoffTool{sw: sw})
registry.Register(&collectTool{sw: sw})
}
// policyFrom derives the member-inheritance policy from the live tool options so
// each spawned member runs on the orchestrator's current model + permission mode.
func policyFrom(options tools.RunOptions) Policy {
return Policy{Model: options.Model, PermissionMode: options.PermissionMode}
}
func swarmStr(args map[string]any, key string) string {
if args == nil {
return ""
}
if v, ok := args[key]; ok {
if s, ok := v.(string); ok {
return strings.TrimSpace(s)
}
}
return ""
}
func okResult(output, kind, summary string) tools.Result {
return tools.Result{Status: tools.StatusOK, Output: output, Display: tools.Display{Kind: kind, Summary: summary}}
}
func errResult(format string, a ...any) tools.Result {
return tools.Result{Status: tools.StatusError, Output: "Error: " + fmt.Sprintf(format, a...)}
}
// ---- swarm_spawn -----------------------------------------------------------
type spawnTool struct{ sw *Swarm }
func (t *spawnTool) Name() string { return SpawnToolName }
func (t *spawnTool) Description() string {
return "Spawn a swarm member of the given agent type to run a task concurrently under a team. Returns a task id immediately; results are reported through swarm_status/swarm_collect."
}
func (t *spawnTool) Parameters() tools.Schema {
return tools.Schema{
Type: "object",
Properties: map[string]tools.PropertySchema{
"agent_type": {Type: "string", Description: "Roster agent type to spawn (e.g. teammate, subagent)."},
"task": {Type: "string", Description: "The focused task/briefing for the member."},
"team": {Type: "string", Description: "Team to spawn into. Defaults to \"default\"."},
},
Required: []string{"agent_type", "task"},
AdditionalProperties: false,
}
}
func (t *spawnTool) Safety() tools.Safety {
return tools.Safety{
SideEffect: tools.SideEffectShell,
Permission: tools.PermissionPrompt,
Reason: "Spawns a swarm member sub-agent under the orchestrator's sandbox and policy.",
AdvertiseInAuto: true,
}
}
func (t *spawnTool) Run(ctx context.Context, args map[string]any) tools.Result {
return t.RunWithOptions(ctx, args, tools.RunOptions{})
}
func (t *spawnTool) RunWithOptions(_ context.Context, args map[string]any, options tools.RunOptions) tools.Result {
agentType := swarmStr(args, "agent_type")
task := swarmStr(args, "task")
team := swarmStr(args, "team")
if agentType == "" {
return errResult("swarm_spawn requires agent_type")
}
if task == "" {
return errResult("swarm_spawn requires task")
}
id, err := t.sw.Spawn(policyFrom(options), team, agentType, task, options.Cwd)
if err != nil {
if errors.Is(err, ErrUnknownAgentType) {
return errResult("%v; available agent types: %s", err, strings.Join(t.sw.Registry().AgentTypes(), ", "))
}
return errResult("%v", err)
}
out := fmt.Sprintf("Spawned %s as task %s on team %s.", agentType, id, displayTeam(team))
res := okResult(out, "swarm", out)
res.Meta = map[string]string{"task_id": id, "team": sanitizeName(team), "agent_type": agentType}
return res
}
// ---- swarm_send ------------------------------------------------------------
type sendTool struct{ sw *Swarm }
func (t *sendTool) Name() string { return SendToolName }
func (t *sendTool) Description() string {
return "Send a message to another swarm member's inbox on a team."
}
func (t *sendTool) Parameters() tools.Schema {
return tools.Schema{
Type: "object",
Properties: map[string]tools.PropertySchema{
"to": {Type: "string", Description: "Recipient member/agent id."},
"body": {Type: "string", Description: "Message body."},
"subject": {Type: "string", Description: "Optional subject."},
"from": {Type: "string", Description: "Optional sender id (the orchestrator by default)."},
"team": {Type: "string", Description: "Team to send within. Defaults to \"default\"."},
},
Required: []string{"to", "body"},
AdditionalProperties: false,
}
}
func (t *sendTool) Safety() tools.Safety {
return tools.Safety{
SideEffect: tools.SideEffectWrite,
Permission: tools.PermissionAllow,
Reason: "Writes a message to an owner-only swarm inbox file.",
AdvertiseInAuto: true,
}
}
func (t *sendTool) Run(ctx context.Context, args map[string]any) tools.Result {
return t.RunWithOptions(ctx, args, tools.RunOptions{})
}
func (t *sendTool) RunWithOptions(_ context.Context, args map[string]any, _ tools.RunOptions) tools.Result {
to := swarmStr(args, "to")
body := swarmStr(args, "body")
if to == "" {
return errResult("swarm_send requires to")
}
if body == "" {
return errResult("swarm_send requires body")
}
team := swarmStr(args, "team")
from := swarmStr(args, "from")
if from == "" {
from = "orchestrator"
}
msg := Message{From: from, Subject: swarmStr(args, "subject"), Body: body}
if err := t.sw.Mailbox().Send(team, to, msg); err != nil {
return errResult("%v", err)
}
out := fmt.Sprintf("Delivered message to %s on team %s.", sanitizeName(to), displayTeam(team))
return okResult(out, "swarm", out)
}
// ---- swarm_inbox -----------------------------------------------------------
type inboxTool struct{ sw *Swarm }
func (t *inboxTool) Name() string { return InboxToolName }
func (t *inboxTool) Description() string {
return "Read a member's inbox on a team. Returns all messages and marks previously-unread ones as read."
}
func (t *inboxTool) Parameters() tools.Schema {
return tools.Schema{
Type: "object",
Properties: map[string]tools.PropertySchema{
"agent": {Type: "string", Description: "Member/agent id whose inbox to read."},
"team": {Type: "string", Description: "Team the inbox belongs to. Defaults to \"default\"."},
},
Required: []string{"agent"},
AdditionalProperties: false,
}
}
func (t *inboxTool) Safety() tools.Safety {
return tools.Safety{
// Reading an inbox marks messages read (a benign metadata flip in an
// owner-only file), so it is classified as a write but auto-allowed.
SideEffect: tools.SideEffectWrite,
Permission: tools.PermissionAllow,
Reason: "Reads and marks-read an owner-only swarm inbox file.",
AdvertiseInAuto: true,
}
}
func (t *inboxTool) Run(ctx context.Context, args map[string]any) tools.Result {
return t.RunWithOptions(ctx, args, tools.RunOptions{})
}
func (t *inboxTool) RunWithOptions(_ context.Context, args map[string]any, _ tools.RunOptions) tools.Result {
agent := swarmStr(args, "agent")
if agent == "" {
return errResult("swarm_inbox requires agent")
}
team := swarmStr(args, "team")
messages, err := t.sw.Mailbox().ReadAndConsume(team, agent)
if err != nil {
return errResult("%v", err)
}
var b strings.Builder
fmt.Fprintf(&b, "Inbox for %s on team %s: %d message(s)\n", sanitizeName(agent), displayTeam(team), len(messages))
for i, msg := range messages {
state := "read"
if !msg.Read {
state = "unread"
}
subject := msg.Subject
if subject != "" {
subject = " [" + subject + "]"
}
fmt.Fprintf(&b, " [%d] (%s) from %s%s: %s\n", i, state, msg.From, subject, msg.Body)
}
out := strings.TrimRight(b.String(), "\n")
return okResult(out, "swarm", fmt.Sprintf("%d message(s)", len(messages)))
}
// ---- swarm_status ----------------------------------------------------------
type statusTool struct{ sw *Swarm }
func (t *statusTool) Name() string { return StatusToolName }
func (t *statusTool) Description() string {
return "Report swarm task status: counts by state and per-task lines. Optionally scoped to one team."
}
func (t *statusTool) Parameters() tools.Schema {
return tools.Schema{
Type: "object",
Properties: map[string]tools.PropertySchema{
"team": {Type: "string", Description: "Optional team to scope to."},
},
AdditionalProperties: false,
}
}
func (t *statusTool) Safety() tools.Safety {
return tools.Safety{
SideEffect: tools.SideEffectRead,
Permission: tools.PermissionAllow,
Reason: "Reports in-memory swarm task status.",
AdvertiseInAuto: true,
}
}
func (t *statusTool) Run(ctx context.Context, args map[string]any) tools.Result {
return t.RunWithOptions(ctx, args, tools.RunOptions{})
}
func (t *statusTool) RunWithOptions(_ context.Context, args map[string]any, _ tools.RunOptions) tools.Result {
team := swarmStr(args, "team")
tasks := t.sw.Coordinator().List()
if team != "" {
want := sanitizeName(team)
filtered := tasks[:0:0]
for _, task := range tasks {
if task.Team == want {
filtered = append(filtered, task)
}
}
tasks = filtered
}
return okResult(renderTasks(t.sw.Coordinator(), tasks, team), "swarm", fmt.Sprintf("%d task(s)", len(tasks)))
}
// ---- swarm_handoff ---------------------------------------------------------
type handoffTool struct{ sw *Swarm }
func (t *handoffTool) Name() string { return HandoffToolName }
func (t *handoffTool) Description() string {
return "Hand a task off to a fresh member of another agent type, delivering an optional note to the new member's inbox."
}
func (t *handoffTool) Parameters() tools.Schema {
return tools.Schema{
Type: "object",
Properties: map[string]tools.PropertySchema{
"task_id": {Type: "string", Description: "Task id to hand off."},
"to_agent_type": {Type: "string", Description: "Roster agent type to take over."},
"note": {Type: "string", Description: "Optional handoff note delivered to the new member's inbox."},
"team": {Type: "string", Description: "Team the task belongs to. Defaults to \"default\"."},
},
Required: []string{"task_id", "to_agent_type"},
AdditionalProperties: false,
}
}
func (t *handoffTool) Safety() tools.Safety {
return tools.Safety{
SideEffect: tools.SideEffectShell,
Permission: tools.PermissionPrompt,
Reason: "Spawns a replacement swarm member to take over a task, and writes the handoff note to the new member's inbox.",
AdvertiseInAuto: true,
}
}
func (t *handoffTool) Run(ctx context.Context, args map[string]any) tools.Result {
return t.RunWithOptions(ctx, args, tools.RunOptions{})
}
func (t *handoffTool) RunWithOptions(_ context.Context, args map[string]any, options tools.RunOptions) tools.Result {
taskID := swarmStr(args, "task_id")
to := swarmStr(args, "to_agent_type")
if taskID == "" {
return errResult("swarm_handoff requires task_id")
}
if to == "" {
return errResult("swarm_handoff requires to_agent_type")
}
team := swarmStr(args, "team")
newID, err := t.sw.Handoff(policyFrom(options), team, taskID, to, swarmStr(args, "note"))
if err != nil {
return errResult("%v", err)
}
out := fmt.Sprintf("Handed off task %s to %s as task %s on team %s.", taskID, to, newID, displayTeam(team))
res := okResult(out, "swarm", out)
res.Meta = map[string]string{"task_id": newID, "previous_task_id": taskID}
return res
}
// ---- swarm_collect ---------------------------------------------------------
type collectTool struct{ sw *Swarm }
func (t *collectTool) Name() string { return CollectToolName }
func (t *collectTool) Description() string {
return "Collect the results of a team's tasks (status, result, and any error per task)."
}
func (t *collectTool) Parameters() tools.Schema {
return tools.Schema{
Type: "object",
Properties: map[string]tools.PropertySchema{
"team": {Type: "string", Description: "Team to collect from. Defaults to \"default\"."},
},
AdditionalProperties: false,
}
}
func (t *collectTool) Safety() tools.Safety {
return tools.Safety{
SideEffect: tools.SideEffectRead,
Permission: tools.PermissionAllow,
Reason: "Reports in-memory swarm task results.",
AdvertiseInAuto: true,
}
}
func (t *collectTool) Run(ctx context.Context, args map[string]any) tools.Result {
return t.RunWithOptions(ctx, args, tools.RunOptions{})
}
func (t *collectTool) RunWithOptions(_ context.Context, args map[string]any, _ tools.RunOptions) tools.Result {
team := swarmStr(args, "team")
tasks := t.sw.Collect(team)
var b strings.Builder
fmt.Fprintf(&b, "Results for team %s: %d task(s)\n", displayTeam(team), len(tasks))
for _, task := range tasks {
line := fmt.Sprintf(" - %s [%s] %s", task.ID, task.Status, task.Description)
if task.Result != "" {
line += "\n result: " + collapse(task.Result)
}
if task.Err != "" {
line += "\n error: " + collapse(task.Err)
}
b.WriteString(line + "\n")
}
out := strings.TrimRight(b.String(), "\n")
return okResult(out, "swarm", fmt.Sprintf("%d task(s)", len(tasks)))
}
// renderTasks formats a status summary + per-task lines, colored per agent.
func renderTasks(coord *Coordinator, tasks []Task, team string) string {
s := coord.Summarize()
var b strings.Builder
scope := "all teams"
if strings.TrimSpace(team) != "" {
scope = "team " + sanitizeName(team)
}
fmt.Fprintf(&b, "Swarm status (%s): %d task(s) — %d running, %d pending, %d done, %d failed, %d handed-off\n",
scope, len(tasks), s.Running, s.Pending, s.Done, s.Failed, s.HandedOff)
sorted := make([]Task, len(tasks))
copy(sorted, tasks)
sort.Slice(sorted, func(i, j int) bool { return sorted[i].ID < sorted[j].ID })
for _, task := range sorted {
fmt.Fprintf(&b, " - %s [%s] (%s) %s\n", task.ID, task.Status, coord.Color(task.AgentID), collapse(task.Description))
}
return strings.TrimRight(b.String(), "\n")
}
// collapse trims a possibly-multiline string to a single compact line for
// status/collect output.
func collapse(s string) string {
s = strings.ReplaceAll(s, "\n", " ")
s = strings.Join(strings.Fields(s), " ")
const max = 200 // runes, not bytes — never slice mid-character
runes := []rune(s)
if len(runes) > max {
return string(runes[:max]) + "…"
}
return s
}
+202
View File
@@ -0,0 +1,202 @@
package swarm
import (
"context"
"strings"
"testing"
"unicode/utf8"
"github.com/Gitlawb/zero/internal/tools"
)
func TestCollapseRuneSafeTruncation(t *testing.T) {
// 300 multi-byte runes: truncation must not split a rune (no invalid UTF-8).
got := collapse(strings.Repeat("🚀", 300))
if !utf8.ValidString(got) {
t.Fatalf("collapse produced invalid UTF-8: %q", got)
}
if n := utf8.RuneCountInString(got); n != 201 { // 200 runes + ellipsis
t.Fatalf("collapse rune count = %d, want 201", n)
}
}
func newToolSwarm(t *testing.T, l MemberLauncher) (*tools.Registry, *Swarm) {
t.Helper()
sw := newSwarmFor(t, l)
reg := tools.NewRegistry()
RegisterTools(reg, sw)
return reg, sw
}
func TestRegisterToolsRegistersAll(t *testing.T) {
reg, _ := newToolSwarm(t, newLauncher(okFor))
for _, name := range []string{SpawnToolName, SendToolName, InboxToolName, StatusToolName, HandoffToolName, CollectToolName} {
if _, ok := reg.Get(name); !ok {
t.Fatalf("tool %q not registered", name)
}
}
}
func TestSpawnToolThroughRegistry(t *testing.T) {
l := newLauncher(okFor)
reg, sw := newToolSwarm(t, l)
res := reg.RunWithOptions(context.Background(), SpawnToolName, map[string]any{
"agent_type": "teammate",
"task": "do the thing",
"team": "alpha",
}, tools.RunOptions{PermissionGranted: true, Model: "m1", Cwd: "/work"})
if res.Status != tools.StatusOK {
t.Fatalf("spawn result status = %v, output=%q", res.Status, res.Output)
}
id := res.Meta["task_id"]
if id == "" {
t.Fatal("spawn must return a task_id in Meta")
}
waitFor(t, "spec recorded", func() bool { return len(l.recorded()) == 1 })
spec := l.recorded()[0]
if spec.Model != "m1" || spec.Cwd != "/work" {
t.Fatalf("policy/cwd not threaded into member: model=%q cwd=%q", spec.Model, spec.Cwd)
}
if _, ok := sw.Coordinator().Get(id); !ok {
t.Fatalf("task %q not tracked by coordinator", id)
}
}
func TestSpawnToolRequiresPermission(t *testing.T) {
reg, _ := newToolSwarm(t, newLauncher(okFor))
// swarm_spawn is a prompt tool: without a grant the registry refuses it.
res := reg.RunWithOptions(context.Background(), SpawnToolName, map[string]any{
"agent_type": "teammate", "task": "x",
}, tools.RunOptions{})
if res.Status != tools.StatusError || !strings.Contains(res.Output, "Permission required") {
t.Fatalf("ungranted spawn should be refused, got status=%v output=%q", res.Status, res.Output)
}
}
func TestSpawnToolValidation(t *testing.T) {
_, sw := newToolSwarm(t, newLauncher(okFor))
tool := &spawnTool{sw: sw}
if res := tool.Run(context.Background(), map[string]any{"task": "x"}); res.Status != tools.StatusError {
t.Fatal("missing agent_type must error")
}
if res := tool.Run(context.Background(), map[string]any{"agent_type": "teammate"}); res.Status != tools.StatusError {
t.Fatal("missing task must error")
}
}
func TestSendAndInboxTools(t *testing.T) {
reg, _ := newToolSwarm(t, newLauncher(okFor))
send := reg.RunWithOptions(context.Background(), SendToolName, map[string]any{
"to": "bob", "body": "ping", "team": "alpha", "subject": "hi",
}, tools.RunOptions{})
if send.Status != tools.StatusOK {
t.Fatalf("send status = %v output=%q", send.Status, send.Output)
}
inbox := reg.RunWithOptions(context.Background(), InboxToolName, map[string]any{
"agent": "bob", "team": "alpha",
}, tools.RunOptions{})
if inbox.Status != tools.StatusOK {
t.Fatalf("inbox status = %v", inbox.Status)
}
if !strings.Contains(inbox.Output, "ping") || !strings.Contains(inbox.Output, "[hi]") {
t.Fatalf("inbox output missing message: %q", inbox.Output)
}
if !strings.Contains(inbox.Output, "unread") {
t.Fatalf("first inbox read should show unread: %q", inbox.Output)
}
}
func TestSendToolValidation(t *testing.T) {
_, sw := newToolSwarm(t, newLauncher(okFor))
tool := &sendTool{sw: sw}
if res := tool.Run(context.Background(), map[string]any{"body": "x"}); res.Status != tools.StatusError {
t.Fatal("missing to must error")
}
if res := tool.Run(context.Background(), map[string]any{"to": "bob"}); res.Status != tools.StatusError {
t.Fatal("missing body must error")
}
}
func TestStatusAndCollectTools(t *testing.T) {
reg, sw := newToolSwarm(t, newLauncher(okFor))
spawn := reg.RunWithOptions(context.Background(), SpawnToolName, map[string]any{
"agent_type": "teammate", "task": "compute", "team": "alpha",
}, tools.RunOptions{PermissionGranted: true, Model: "m"})
id := spawn.Meta["task_id"]
waitFor(t, "task done", func() bool {
task, ok := sw.Coordinator().Get(id)
return ok && task.Status == StatusDone
})
status := reg.RunWithOptions(context.Background(), StatusToolName, map[string]any{"team": "alpha"}, tools.RunOptions{})
if status.Status != tools.StatusOK || !strings.Contains(status.Output, id) {
t.Fatalf("status output missing task: %q", status.Output)
}
collect := reg.RunWithOptions(context.Background(), CollectToolName, map[string]any{"team": "alpha"}, tools.RunOptions{})
if collect.Status != tools.StatusOK || !strings.Contains(collect.Output, "ok:compute") {
t.Fatalf("collect output missing result: %q", collect.Output)
}
}
func TestSpawnToolUnknownTypeListsAvailable(t *testing.T) {
_, sw := newToolSwarm(t, newLauncher(okFor))
tool := &spawnTool{sw: sw}
res := tool.RunWithOptions(context.Background(), map[string]any{
"agent_type": "ghost", "task": "x",
}, tools.RunOptions{})
if res.Status != tools.StatusError || !strings.Contains(res.Output, "available agent types") {
t.Fatalf("unknown type should list available agents, got %q", res.Output)
}
if !strings.Contains(res.Output, "teammate") {
t.Fatalf("available agents should include builtins, got %q", res.Output)
}
}
func TestNewWithUserDefinitions(t *testing.T) {
used := false
sw, err := New(Options{
BaseDir: t.TempDir(),
Launcher: newLauncher(okFor),
Definitions: []Definition{{
AgentType: "researcher",
SystemPrompt: func(PromptContext) string { used = true; return "research" },
}},
})
if err != nil {
t.Fatalf("New with definitions: %v", err)
}
t.Cleanup(sw.Close)
if _, err := sw.Spawn(Policy{Model: "m"}, "team", "researcher", "find things", ""); err != nil {
t.Fatalf("Spawn user-defined agent: %v", err)
}
waitFor(t, "researcher used", func() bool { return used })
// A bad definition makes New fail closed.
if _, err := New(Options{BaseDir: t.TempDir(), Launcher: newLauncher(okFor), Definitions: []Definition{{AgentType: " "}}}); err == nil {
t.Fatal("New with a blank-agentType definition must error")
}
}
func TestHandoffToolThroughRegistry(t *testing.T) {
gate := make(chan struct{})
l := newLauncher(okFor)
l.gate = gate
reg, sw := newToolSwarm(t, l)
spawn := reg.RunWithOptions(context.Background(), SpawnToolName, map[string]any{
"agent_type": "teammate", "task": "long", "team": "alpha",
}, tools.RunOptions{PermissionGranted: true, Model: "m"})
origID := spawn.Meta["task_id"]
waitFor(t, "running", func() bool {
task, ok := sw.Coordinator().Get(origID)
return ok && task.Status == StatusRunning
})
res := reg.RunWithOptions(context.Background(), HandoffToolName, map[string]any{
"task_id": origID, "to_agent_type": "subagent", "note": "carry on", "team": "alpha",
}, tools.RunOptions{PermissionGranted: true, Model: "m"})
if res.Status != tools.StatusOK {
t.Fatalf("handoff status = %v output=%q", res.Status, res.Output)
}
if res.Meta["previous_task_id"] != origID || res.Meta["task_id"] == "" {
t.Fatalf("handoff meta wrong: %+v", res.Meta)
}
close(gate)
}
+24 -2
View File
@@ -8,6 +8,8 @@ import (
"regexp"
"sort"
"strings"
"github.com/Gitlawb/zero/internal/sandbox"
)
type globTool struct {
@@ -44,6 +46,17 @@ func NewScopedGlobTool(workspaceRoot string, scope PathScope) Tool {
}
func (tool globTool) Run(_ context.Context, args map[string]any) Result {
return tool.runWith(args, readExcluder{})
}
// RunWithSandbox runs glob while skipping subtrees the sandbox policy denies
// reads to (DenyRead). With no DenyRead configured the excluder is a no-op and
// behavior is unchanged.
func (tool globTool) RunWithSandbox(_ context.Context, args map[string]any, engine *sandbox.Engine) Result {
return tool.runWith(args, sandboxReadExcluder(engine))
}
func (tool globTool) runWith(args map[string]any, exclude readExcluder) Result {
pattern, err := aliasedStringArg(args, []string{"pattern", "glob", "match", "query", "expression"}, "", true, false)
if err != nil {
return errorResult("Error: Invalid arguments for glob: " + err.Error())
@@ -75,7 +88,7 @@ func (tool globTool) Run(_ context.Context, args map[string]any) Result {
return errorResult("Error running glob " + fmt.Sprintf("%q", pattern) + ": " + err.Error())
}
matches, err := scanGlob(root, displayRoot, matcher, includeDirs)
matches, err := scanGlob(root, displayRoot, matcher, includeDirs, exclude)
if err != nil {
return errorResult("Error running glob " + fmt.Sprintf("%q", pattern) + ": " + err.Error())
}
@@ -99,7 +112,7 @@ func (tool globTool) Run(_ context.Context, args map[string]any) Result {
}
}
func scanGlob(root string, displayRoot string, matcher *regexp.Regexp, includeDirs bool) ([]string, error) {
func scanGlob(root string, displayRoot string, matcher *regexp.Regexp, includeDirs bool, exclude readExcluder) ([]string, error) {
matches := []string{}
err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, walkErr error) error {
if walkErr != nil {
@@ -117,6 +130,10 @@ func scanGlob(root string, displayRoot string, matcher *regexp.Regexp, includeDi
if entry.IsDir() && shouldSkipDirectory(entry.Name()) {
return filepath.SkipDir
}
if entry.IsDir() && exclude.dirExcluded(path) {
// Skip a read-denied subtree (sandbox DenyRead) wholesale.
return filepath.SkipDir
}
if entry.IsDir() && !includeDirs {
return nil
}
@@ -129,6 +146,11 @@ func scanGlob(root string, displayRoot string, matcher *regexp.Regexp, includeDi
if !entry.IsDir() && shouldSkipWorkspaceFile(normalized) {
return nil
}
// A read-denied path (even an includeDirs directory with a nested allow we
// still descended into) must not appear in results.
if exclude.fileExcluded(path) {
return nil
}
if matcher.MatchString(normalized) {
match := normalized
// When cwd resolved to an extra (non-workspace) root, resolveScopedPath
+26 -2
View File
@@ -8,6 +8,8 @@ import (
"regexp"
"sort"
"strings"
"github.com/Gitlawb/zero/internal/sandbox"
)
type grepTool struct {
@@ -54,6 +56,18 @@ func NewScopedGrepTool(workspaceRoot string, scope PathScope) Tool {
}
func (tool grepTool) Run(_ context.Context, args map[string]any) Result {
return tool.runWith(args, readExcluder{})
}
// RunWithSandbox runs the search while skipping subtrees the sandbox policy
// denies reads to (DenyRead), so grep never surfaces content from a read-denied
// path. With no DenyRead configured the excluder is a no-op and behavior is
// unchanged.
func (tool grepTool) RunWithSandbox(_ context.Context, args map[string]any, engine *sandbox.Engine) Result {
return tool.runWith(args, sandboxReadExcluder(engine))
}
func (tool grepTool) runWith(args map[string]any, exclude readExcluder) Result {
pattern, err := aliasedStringArg(args, []string{"pattern", "query", "regex", "search", "expression"}, "", true, false)
if err != nil {
return errorResult("Error: Invalid arguments for grep: " + err.Error())
@@ -124,7 +138,7 @@ func (tool grepTool) Run(_ context.Context, args map[string]any) Result {
}
}
files, err := grepFiles(resolvedRoot, target, globMatcher)
files, err := grepFiles(resolvedRoot, target, globMatcher, exclude)
if err != nil {
return errorResult("Error running grep: " + err.Error())
}
@@ -225,7 +239,7 @@ func confineGrepFile(resolvedRoot string, path string) (string, string, bool) {
return filepath.ToSlash(relative), resolved, true
}
func grepFiles(resolvedRoot string, target string, globMatcher *regexp.Regexp) ([]string, error) {
func grepFiles(resolvedRoot string, target string, globMatcher *regexp.Regexp, exclude readExcluder) ([]string, error) {
info, err := os.Stat(target)
if err != nil {
return nil, err
@@ -239,6 +253,9 @@ func grepFiles(resolvedRoot string, target string, globMatcher *regexp.Regexp) (
if shouldSkipWorkspaceFile(relative) {
return []string{}, nil
}
if exclude.fileExcluded(target) {
return []string{}, nil
}
// A single explicit file is matched by its base name so a pattern like
// "*.go" applies regardless of how deep the file sits under the workspace.
if globMatcher == nil || globMatcher.MatchString(filepath.Base(target)) {
@@ -265,6 +282,13 @@ func grepFiles(resolvedRoot string, target string, globMatcher *regexp.Regexp) (
return filepath.SkipDir
}
if entry.IsDir() {
// Skip a read-denied subtree (sandbox DenyRead) wholesale.
if exclude.dirExcluded(path) {
return filepath.SkipDir
}
return nil
}
if exclude.fileExcluded(path) {
return nil
}
// Confine each candidate through symlinks: a symlink inside the workspace
+31
View File
@@ -0,0 +1,31 @@
package tools
import "github.com/Gitlawb/zero/internal/sandbox"
// readExcluder skips read-denied paths (the sandbox DenyRead policy) during a
// search walk. The zero value (both funcs nil) excludes nothing, so a
// non-sandboxed search behaves exactly as before — the exclusions are opt-in and
// only ever REMOVE results, never add them.
type readExcluder struct {
file func(string) bool
dir func(string) bool
}
func (e readExcluder) fileExcluded(path string) bool { return e.file != nil && e.file(path) }
func (e readExcluder) dirExcluded(path string) bool { return e.dir != nil && e.dir(path) }
// sandboxReadExcluder builds a readExcluder from a sandbox engine's DenyRead
// policy. The engine's read-exclusion matcher resolves the policy paths ONCE here
// (not per visited path), and the closures reuse it for the whole walk. A nil
// engine or an inactive (no DenyRead) policy yields the no-op excluder, so the
// search tools keep their pre-sandbox behavior.
func sandboxReadExcluder(engine *sandbox.Engine) readExcluder {
if engine == nil {
return readExcluder{}
}
rx := engine.ReadExclusions()
if !rx.Active() {
return readExcluder{}
}
return readExcluder{file: rx.PathExcluded, dir: rx.DirExcluded}
}
+99
View File
@@ -0,0 +1,99 @@
package tools
import (
"context"
"os"
"path/filepath"
"strings"
"testing"
"github.com/Gitlawb/zero/internal/sandbox"
)
// denyReadFixture builds a workspace with an ordinary file and a secret subtree,
// returning the resolved workspace root and a sandbox engine whose DenyRead
// covers the secret dir.
func denyReadFixture(t *testing.T) (string, *sandbox.Engine) {
t.Helper()
ws, err := filepath.EvalSymlinks(t.TempDir())
if err != nil {
t.Fatalf("EvalSymlinks: %v", err)
}
if err := os.WriteFile(filepath.Join(ws, "main.go"), []byte("package main\n"), 0o600); err != nil {
t.Fatalf("write main.go: %v", err)
}
secret := filepath.Join(ws, "secret")
if err := os.MkdirAll(secret, 0o755); err != nil {
t.Fatalf("mkdir secret: %v", err)
}
if err := os.WriteFile(filepath.Join(secret, "creds.go"), []byte("package main\n"), 0o600); err != nil {
t.Fatalf("write creds.go: %v", err)
}
scope, err := sandbox.NewScope(ws, nil)
if err != nil {
t.Fatalf("NewScope: %v", err)
}
engine := sandbox.NewEngine(sandbox.EngineOptions{
WorkspaceRoot: ws,
Policy: sandbox.Policy{
Mode: sandbox.ModeEnforce,
EnforceWorkspace: true,
DenyRead: []string{secret},
},
Scope: scope,
})
return ws, engine
}
func TestGrepSkipsDenyReadSubtree(t *testing.T) {
ws, engine := denyReadFixture(t)
tool, ok := NewGrepTool(ws).(sandboxAwareTool)
if !ok {
t.Fatal("grep tool must be sandbox-aware")
}
args := map[string]any{"pattern": "package main", "output_mode": "files_with_matches"}
// Sandboxed: the DenyRead subtree must be excluded from results.
sandboxed := tool.RunWithSandbox(context.Background(), args, engine)
if sandboxed.Status != StatusOK {
t.Fatalf("grep failed: %s", sandboxed.Output)
}
if !strings.Contains(sandboxed.Output, "main.go") {
t.Fatalf("grep must still match the non-denied file, got:\n%s", sandboxed.Output)
}
if strings.Contains(sandboxed.Output, "creds.go") {
t.Fatalf("grep must NOT surface a DenyRead file, got:\n%s", sandboxed.Output)
}
// Without a sandbox, the same search includes the secret file (default behavior).
plain := NewGrepTool(ws).Run(context.Background(), args)
if !strings.Contains(plain.Output, "creds.go") {
t.Fatalf("non-sandboxed grep should include the secret file, got:\n%s", plain.Output)
}
}
func TestGlobSkipsDenyReadSubtree(t *testing.T) {
ws, engine := denyReadFixture(t)
tool, ok := NewGlobTool(ws).(sandboxAwareTool)
if !ok {
t.Fatal("glob tool must be sandbox-aware")
}
args := map[string]any{"pattern": "**/*.go"}
sandboxed := tool.RunWithSandbox(context.Background(), args, engine)
if sandboxed.Status != StatusOK {
t.Fatalf("glob failed: %s", sandboxed.Output)
}
if !strings.Contains(sandboxed.Output, "main.go") {
t.Fatalf("glob must still match the non-denied file, got:\n%s", sandboxed.Output)
}
if strings.Contains(sandboxed.Output, "creds.go") {
t.Fatalf("glob must NOT surface a DenyRead file, got:\n%s", sandboxed.Output)
}
plain := NewGlobTool(ws).Run(context.Background(), args)
if !strings.Contains(plain.Output, "creds.go") {
t.Fatalf("non-sandboxed glob should include the secret file, got:\n%s", plain.Output)
}
}
+2
View File
@@ -199,6 +199,7 @@ func CoreReadOnlyToolsScoped(workspaceRoot string, scope PathScope) []Tool {
}
}
func CoreWriteTools(workspaceRoot string) []Tool { return CoreWriteToolsScoped(workspaceRoot, nil) }
func CoreWriteToolsScoped(workspaceRoot string, scope PathScope) []Tool {
return []Tool{
NewScopedWriteFileTool(workspaceRoot, scope),
@@ -208,6 +209,7 @@ func CoreWriteToolsScoped(workspaceRoot string, scope PathScope) []Tool {
}
}
func CoreShellTools(workspaceRoot string) []Tool { return CoreShellToolsScoped(workspaceRoot, nil) }
func CoreShellToolsScoped(workspaceRoot string, scope PathScope) []Tool {
return []Tool{
NewScopedBashTool(workspaceRoot, scope),
+1 -1
View File
@@ -302,7 +302,7 @@ func webFetchRedirectPolicy(previous func(*http.Request, []*http.Request) error,
return fmt.Errorf("too many redirects: maximum is %d", webFetchRedirectLimit)
}
if err := validateParsedWebFetchURL(request.Context(), request.URL, resolver, gate); err != nil {
return fmt.Errorf("Unsafe redirect URL: %w", err)
return fmt.Errorf("unsafe redirect URL: %w", err)
}
if previous != nil {
return previous(request, via)
+6 -4
View File
@@ -285,7 +285,7 @@ func TestWebFetchToolRejectsUnsafeRedirects(t *testing.T) {
if result.Status != StatusError {
t.Fatalf("expected redirect safety error, got %s: %s", result.Status, result.Output)
}
if !strings.Contains(result.Output, "Unsafe redirect URL") {
if !strings.Contains(result.Output, "unsafe redirect URL") {
t.Fatalf("expected unsafe redirect message, got %q", result.Output)
}
})
@@ -386,7 +386,9 @@ func TestWebFetchRunWithSandboxBlocksDeniedHost(t *testing.T) {
}
engine := zeroSandbox.NewEngine(zeroSandbox.EngineOptions{
Policy: zeroSandbox.Policy{Mode: zeroSandbox.ModeEnforce, Network: zeroSandbox.NetworkDeny},
// EnforceToolNetwork opts the in-process tool into the network policy; off
// (the default) would let web_fetch run despite deny.
Policy: zeroSandbox.Policy{Mode: zeroSandbox.ModeEnforce, Network: zeroSandbox.NetworkDeny, EnforceToolNetwork: true},
})
result := tool.RunWithSandbox(context.Background(), map[string]any{"url": "https://example.com"}, engine)
if result.Status != StatusError {
@@ -405,7 +407,7 @@ func TestWebFetchRunWithSandboxScopedBlocksUnlistedHost(t *testing.T) {
tool := newWebFetchToolWithClient(failingClient).(webFetchTool)
engine := zeroSandbox.NewEngine(zeroSandbox.EngineOptions{
Policy: zeroSandbox.Policy{Mode: zeroSandbox.ModeEnforce, Network: zeroSandbox.NetworkScoped, AllowedDomains: []string{"allowed.test"}},
Policy: zeroSandbox.Policy{Mode: zeroSandbox.ModeEnforce, Network: zeroSandbox.NetworkScoped, AllowedDomains: []string{"allowed.test"}, EnforceToolNetwork: true},
Backend: zeroSandbox.Backend{
Name: zeroSandbox.BackendSandboxExec, Available: true,
Executable: "/usr/bin/sandbox-exec", ScopedEgress: true,
@@ -433,7 +435,7 @@ func TestRegistryRoutesWebFetchThroughSandboxScopedPolicy(t *testing.T) {
registry.Register(newWebFetchToolWithClient(failingClient))
engine := zeroSandbox.NewEngine(zeroSandbox.EngineOptions{
Policy: zeroSandbox.Policy{Mode: zeroSandbox.ModeEnforce, Network: zeroSandbox.NetworkScoped, AllowedDomains: []string{"allowed.test"}},
Policy: zeroSandbox.Policy{Mode: zeroSandbox.ModeEnforce, Network: zeroSandbox.NetworkScoped, AllowedDomains: []string{"allowed.test"}, EnforceToolNetwork: true},
Backend: zeroSandbox.Backend{
Name: zeroSandbox.BackendSandboxExec, Available: true,
Executable: "/usr/bin/sandbox-exec", ScopedEgress: true,
+1 -1
View File
@@ -276,7 +276,7 @@ func parseSearchResults(body []byte) ([]searchResult, error) {
convert := func(raw []rawResult) []searchResult {
out := make([]searchResult, 0, len(raw))
for _, item := range raw {
out = append(out, searchResult{Title: item.Title, URL: item.URL, Snippet: item.Snippet})
out = append(out, searchResult(item))
}
return out
}
+23 -5
View File
@@ -160,18 +160,36 @@ func TestHTTPSearchBackendEndpointHost(t *testing.T) {
}
}
func TestWebSearchRunWithSandboxBlocksUnderDeny(t *testing.T) {
backend := &fakeHostedBackend{host: "search.example"}
// TestWebSearchRunWithSandboxAllowsByDefaultUnderDeny verifies the default
// posture: a first-party in-process tool is NOT subject to the sandbox network
// policy unless EnforceToolNetwork is set, so web_search runs even under deny.
func TestWebSearchRunWithSandboxAllowsByDefaultUnderDeny(t *testing.T) {
backend := &fakeHostedBackend{host: "search.example", results: []searchResult{{Title: "T", URL: "https://x.test"}}}
tool := newWebSearchToolWithBackend(backend).(webSearchTool)
engine := zeroSandbox.NewEngine(zeroSandbox.EngineOptions{
Policy: zeroSandbox.Policy{Mode: zeroSandbox.ModeEnforce, Network: zeroSandbox.NetworkDeny},
})
res := tool.RunWithSandbox(context.Background(), map[string]any{"query": "hi"}, engine)
if res.Status != StatusOK {
t.Fatalf("web_search must run by default even under deny, got %q: %s", res.Status, res.Output)
}
if !backend.called {
t.Fatal("search backend must be called when tool-network enforcement is off (default)")
}
}
func TestWebSearchRunWithSandboxBlocksUnderDenyWhenEnforced(t *testing.T) {
backend := &fakeHostedBackend{host: "search.example"}
tool := newWebSearchToolWithBackend(backend).(webSearchTool)
engine := zeroSandbox.NewEngine(zeroSandbox.EngineOptions{
Policy: zeroSandbox.Policy{Mode: zeroSandbox.ModeEnforce, Network: zeroSandbox.NetworkDeny, EnforceToolNetwork: true},
})
res := tool.RunWithSandbox(context.Background(), map[string]any{"query": "hi"}, engine)
if res.Status != StatusError || !strings.Contains(res.Output, "disabled") {
t.Fatalf("expected deny block, got %q: %s", res.Status, res.Output)
}
if backend.called {
t.Fatal("search backend must not be called under network deny")
t.Fatal("search backend must not be called under network deny when enforced")
}
}
@@ -179,7 +197,7 @@ func TestWebSearchRunWithSandboxScopedBlocksUnlistedHost(t *testing.T) {
backend := &fakeHostedBackend{host: "search.example"}
tool := newWebSearchToolWithBackend(backend).(webSearchTool)
engine := zeroSandbox.NewEngine(zeroSandbox.EngineOptions{
Policy: zeroSandbox.Policy{Mode: zeroSandbox.ModeEnforce, Network: zeroSandbox.NetworkScoped, AllowedDomains: []string{"allowed.test"}},
Policy: zeroSandbox.Policy{Mode: zeroSandbox.ModeEnforce, Network: zeroSandbox.NetworkScoped, AllowedDomains: []string{"allowed.test"}, EnforceToolNetwork: true},
Backend: zeroSandbox.Backend{
Name: zeroSandbox.BackendSandboxExec, Available: true,
Executable: "/usr/bin/sandbox-exec", ScopedEgress: true,
@@ -198,7 +216,7 @@ func TestWebSearchRunWithSandboxScopedAllowsListedHost(t *testing.T) {
backend := &fakeHostedBackend{host: "search.example", results: []searchResult{{Title: "T", URL: "https://x.test"}}}
tool := newWebSearchToolWithBackend(backend).(webSearchTool)
engine := zeroSandbox.NewEngine(zeroSandbox.EngineOptions{
Policy: zeroSandbox.Policy{Mode: zeroSandbox.ModeEnforce, Network: zeroSandbox.NetworkScoped, AllowedDomains: []string{"search.example"}},
Policy: zeroSandbox.Policy{Mode: zeroSandbox.ModeEnforce, Network: zeroSandbox.NetworkScoped, AllowedDomains: []string{"search.example"}, EnforceToolNetwork: true},
Backend: zeroSandbox.Backend{
Name: zeroSandbox.BackendSandboxExec, Available: true,
Executable: "/usr/bin/sandbox-exec", ScopedEgress: true,
+195 -6
View File
@@ -6,6 +6,8 @@ import (
"strings"
"time"
tea "github.com/charmbracelet/bubbletea"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/doctor"
"github.com/Gitlawb/zero/internal/modelregistry"
@@ -14,13 +16,200 @@ import (
zsearch "github.com/Gitlawb/zero/internal/search"
)
func (m model) doctorText() string {
report := doctor.Run(doctor.Options{
Now: m.now,
Runtime: "go",
Provider: m.providerProfile,
const doctorStatusRowID = "doctor/status"
func (m model) startDoctorCommand(args string) (model, tea.Cmd) {
connectivity, fix, help, err := parseDoctorCommandArgs(args)
if err != nil {
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: doctorUsageText(commandStatusBlocked, err.Error())})
return m, nil
}
if help {
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: doctorUsageText(commandStatusInfo, "Show local diagnostics for provider, model, sandbox, LSP, and backend setup.")})
return m, nil
}
if fix {
return m.startDoctorFixCommand()
}
if !connectivity {
m = m.setDoctorStatusRow(m.doctorText(false))
return m, nil
}
m.doctorCommandSeq++
id := m.doctorCommandSeq
snapshot := m
m.doctorInFlight = true
m.doctorFrame = 0
m = m.setDoctorStatusRow(m.doctorConnectivityRunningText())
return m, tea.Batch(func() tea.Msg {
return doctorCommandResultMsg{id: id, text: snapshot.doctorText(true)}
}, m.spinner.Tick)
}
func (m model) startDoctorFixCommand() (model, tea.Cmd) {
report := doctor.Run(m.doctorOptions(false))
if doctorReportNeedsProviderSetup(report) {
if m.pending {
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: doctorFixBusyText()})
return m, nil
}
m.providerWizard = m.newProviderWizard()
m.clearSuggestions()
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: doctorFixProviderSetupText()})
return m, nil
}
if doctorReportCanProbeConnectivity(report) {
return m.startDoctorCommand("--connectivity")
}
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: doctorFixPlanText(report)})
return m, nil
}
func (m model) doctorText(connectivity bool) string {
report := doctor.Run(m.doctorOptions(connectivity))
return renderCommandOutput(doctorCommandOutput(report, nil))
}
func parseDoctorCommandArgs(args string) (connectivity bool, fix bool, help bool, err error) {
for _, field := range strings.Fields(args) {
switch strings.ToLower(field) {
case "--connectivity", "connectivity":
connectivity = true
case "--fix", "fix":
fix = true
case "-h", "--help", "help":
help = true
default:
return false, false, false, fmt.Errorf("unknown doctor flag %q", field)
}
}
if connectivity && fix {
return false, false, false, fmt.Errorf("choose either %q or %q, not both", "fix", "--connectivity")
}
return connectivity, fix, help, nil
}
func doctorReportNeedsProviderSetup(report doctor.Report) bool {
for _, id := range []string{"provider.config", "provider.model"} {
if check := report.Check(id); check != nil && check.Status == doctor.StatusFail {
return true
}
}
return false
}
func doctorReportCanProbeConnectivity(report doctor.Report) bool {
check := report.Check("provider.connectivity")
return check != nil && check.Status != doctor.StatusPass
}
func doctorUsageText(status commandStatus, message string) string {
return renderCommandOutput(commandOutput{
Title: "Diagnostics",
Status: status,
Sections: []commandSection{{
Title: "Usage",
Lines: []string{
message,
"/doctor",
"/doctor fix",
"/doctor --connectivity",
"/health",
},
}},
})
return doctor.Format(report)
}
func doctorFixProviderSetupText() string {
return renderCommandOutput(commandOutput{
Title: "Diagnostics fix",
Status: commandStatusInfo,
Sections: []commandSection{{
Title: "Provider",
Lines: []string{"Opening provider setup. Choose a provider, add credentials, then select a model."},
}},
Hints: []string{"Esc closes setup"},
})
}
func doctorFixBusyText() string {
return renderCommandOutput(commandOutput{
Title: "Diagnostics fix",
Status: commandStatusWarning,
Sections: []commandSection{{
Title: "Provider",
Lines: []string{"Cannot open provider setup while a run is active."},
}},
Hints: []string{"stop the current run, then retry /doctor fix"},
})
}
func doctorFixPlanText(report doctor.Report) string {
return renderCommandOutput(commandOutput{
Title: "Diagnostics fix",
Status: doctorCommandStatus(report),
Sections: []commandSection{{
Title: "Next actions",
Lines: doctorFixLines(report),
}},
Hints: []string{"run /doctor --connectivity to recheck provider health"},
})
}
func doctorFixLines(report doctor.Report) []string {
lines := []string{}
hasIssue := false
for _, check := range report.Checks {
if check.Status == doctor.StatusPass {
continue
}
hasIssue = true
switch check.ID {
case "sandbox.backend":
lines = append(lines, "native sandbox: use WSL2 or a Linux container on Windows")
case "lsp.servers":
lines = append(lines, "language servers: install missing LSP binaries on PATH")
case "provider.connectivity":
lines = append(lines, "provider connectivity: run /doctor --connectivity")
case "config.files", "config.validation":
lines = append(lines, "config: run /provider to create or repair provider config")
}
}
if len(lines) == 0 {
if hasIssue {
return []string{"No automatic fixes are available for the detected diagnostics."}
}
return []string{"No automatic fixes are available because diagnostics are already clean."}
}
return lines
}
func (m model) doctorConnectivityRunningText() string {
return strings.Join([]string{
"Checking provider",
"Zero is probing the active endpoint. Keep typing; messages will queue until the check finishes.",
m.doctorAnimationLine(),
"provider: " + displayValue(m.providerName, displayValue(m.providerProfile.Name, "unknown")),
"model: " + displayValue(m.modelName, displayValue(m.providerProfile.Model, "unknown")),
}, "\n")
}
func (m model) doctorAnimationLine() string {
frame := compactFrames[m.doctorFrame%len(compactFrames)]
return frame + " checking provider connectivity..."
}
func (m model) setDoctorStatusRow(text string) model {
row := transcriptRow{kind: rowSystem, id: doctorStatusRowID, tool: "doctor", text: text}
for i := len(m.transcript) - 1; i >= 0; i-- {
if m.transcript[i].id == doctorStatusRowID {
m.transcript[i] = row
return m
}
}
m.transcript = appendTranscriptRow(m.transcript, row)
return m
}
func (m model) searchText(query string) string {
+99
View File
@@ -39,6 +39,105 @@ type commandRow struct {
Text string
}
type commandCard struct {
Title string
Summary []string
Sections []commandCardSection
Actions []string
}
type commandCardSection struct {
Title string
Fields []commandField
Lines []string
Rows []commandRow
}
const commandCardTranscriptPrefix = "\x00command-card\x00"
func renderCommandCardTranscript(card commandCard) string {
return commandCardTranscriptPrefix + renderCommandCard(card)
}
func commandCardTranscriptPayload(text string) (string, bool) {
if !strings.HasPrefix(text, commandCardTranscriptPrefix) {
return text, false
}
return strings.TrimPrefix(text, commandCardTranscriptPrefix), true
}
func renderCommandCard(card commandCard) string {
lines := []string{}
title := compactCommandOutputText(card.Title)
if title == "" {
title = "Zero"
}
lines = append(lines, title)
if summary := compactCommandCardList(card.Summary); len(summary) > 0 {
lines = append(lines, strings.Join(summary, " | "))
}
for _, section := range card.Sections {
lines = append(lines, formatCommandCardSection(section)...)
}
if actions := compactCommandCardList(card.Actions); len(actions) > 0 {
lines = append(lines, "actions: "+strings.Join(actions, " | "))
}
return strings.Join(lines, "\n")
}
func formatCommandCardSection(section commandCardSection) []string {
lines := []string{}
if title := compactCommandOutputText(section.Title); title != "" {
lines = append(lines, title)
}
keyWidth := 0
fields := make([]commandField, 0, len(section.Fields))
for _, field := range section.Fields {
key := compactCommandOutputText(field.Key)
value := compactCommandOutputText(field.Value)
if key == "" || value == "" {
continue
}
if len(key) > keyWidth {
keyWidth = len(key)
}
fields = append(fields, commandField{Key: key, Value: value})
}
for _, field := range fields {
lines = append(lines, " "+field.Key+strings.Repeat(" ", keyWidth-len(field.Key)+2)+field.Value)
}
for _, line := range section.Lines {
if text := compactCommandOutputText(line); text != "" {
lines = append(lines, " "+text)
}
}
for _, row := range section.Rows {
if text := compactCommandOutputText(row.Text); text != "" {
lines = append(lines, " - "+text)
}
}
return lines
}
func compactCommandCardList(values []string) []string {
compacted := make([]string, 0, len(values))
for _, value := range values {
if text := compactCommandOutputText(value); text != "" {
compacted = append(compacted, text)
}
}
return compacted
}
func formatCommandOutput(output commandOutput) string {
lines := []string{}
+75
View File
@@ -49,6 +49,81 @@ func TestFormatCommandOutputRendersSectionsFieldsRowsAndHints(t *testing.T) {
}
}
func TestFormatCommandCardRendersBoundedDashboard(t *testing.T) {
got := renderCommandCard(commandCard{
Title: "Context",
Summary: []string{"go runtime", "ask permissions", "1 tool"},
Sections: []commandCardSection{
{
Title: "Runtime",
Fields: []commandField{
{Key: "cwd", Value: `D:\repo`},
{Key: "provider", Value: "openai"},
},
},
{
Title: "Tools",
Fields: []commandField{
{Key: "registered", Value: "1"},
},
},
},
Actions: []string{"/permissions manage access", "/tools inspect catalog"},
})
want := strings.Join([]string{
"Context",
"go runtime | ask permissions | 1 tool",
"Runtime",
" cwd D:\\repo",
" provider openai",
"Tools",
" registered 1",
"actions: /permissions manage access | /tools inspect catalog",
}, "\n")
if got != want {
t.Fatalf("unexpected command card:\nwant:\n%s\n\ngot:\n%s", want, got)
}
}
func TestFormatCommandCardRedactsTokenLikeText(t *testing.T) {
got := renderCommandCard(commandCard{
Title: "Context",
Summary: []string{"provider sk-proj-summary-secret-value"},
Sections: []commandCardSection{
{
Title: "Runtime",
Fields: []commandField{
{Key: "api key", Value: "sk-ant-api03-abcdefghijklmnopqrstuvwxyz"},
},
Lines: []string{
"google token AIza1234567890abcdef",
},
Rows: []commandRow{
{Text: "shell used sk-proj-row-secret-value"},
},
},
},
Actions: []string{"retry with sk-proj-action-secret-value"},
})
for _, secret := range []string{
"sk-proj-summary-secret-value",
"sk-ant-api03-abcdefghijklmnopqrstuvwxyz",
"AIza1234567890abcdef",
"sk-proj-row-secret-value",
"sk-proj-action-secret-value",
} {
if strings.Contains(got, secret) {
t.Fatalf("expected token-like text to be redacted, got:\n%s", got)
}
}
if !strings.Contains(got, "[REDACTED]") {
t.Fatalf("expected redacted marker in command card, got:\n%s", got)
}
}
func TestFormatCommandOutputSupportsAllStatuses(t *testing.T) {
tests := []struct {
name string
+229 -2
View File
@@ -2,6 +2,7 @@ package tui
import (
"context"
"errors"
"path/filepath"
"strings"
"testing"
@@ -107,6 +108,139 @@ func TestProviderCommandRedactsCredentialBearingBaseURL(t *testing.T) {
assertContains(t, text, "base url: https://proxy.local/v1?api_key=[REDACTED]")
}
func TestToolsCommandRendersCommandCard(t *testing.T) {
m := newModel(context.Background(), Options{
Registry: tools.NewRegistry(),
})
m.input.SetValue("/tools")
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
next := updated.(model)
if cmd != nil {
t.Fatal("expected /tools to be handled without starting an agent run")
}
emptyText := transcriptText(next.transcript)
for _, want := range []string{
"Tools",
"0 registered | no tools available",
"Registry",
"registered 0",
"actions: /mcp manage servers | /permissions manage access",
} {
assertContains(t, emptyText, want)
}
assertNotContains(t, emptyText, "status: warning")
assertNotContains(t, emptyText, "registered tools:")
registry := tools.NewRegistry()
registry.Register(tools.NewReadFileTool("."))
m = newModel(context.Background(), Options{
Registry: registry,
})
m.input.SetValue("/tools")
updated, cmd = m.Update(tea.KeyMsg{Type: tea.KeyEnter})
next = updated.(model)
if cmd != nil {
t.Fatal("expected /tools to be handled without starting an agent run")
}
toolsText := transcriptText(next.transcript)
for _, want := range []string{
"Tools",
"1 registered | registered catalog",
"Registry",
"registered 1",
"Available",
"- read_file",
"actions: /mcp manage servers | /permissions manage access",
} {
assertContains(t, toolsText, want)
}
assertNotContains(t, toolsText, "status: ok")
assertNotContains(t, toolsText, "registered tools:")
}
func TestToolsCommandCardHandlesNilRegistry(t *testing.T) {
text := model{}.toolsText()
for _, want := range []string{
"Tools",
"0 registered | no tools available",
"registered 0",
} {
assertContains(t, text, want)
}
}
func TestToolsCommandShowsFullSortedCatalog(t *testing.T) {
registry := tools.NewRegistry()
for _, name := range []string{
"write_file",
"read_file",
"grep",
"glob",
"edit_file",
"apply_patch",
"bash",
"web_search",
"web_fetch",
} {
registry.Register(commandTestMCPTool{name: name, description: name + " tool"})
}
m := newModel(context.Background(), Options{
Registry: registry,
})
m.input.SetValue("/tools")
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
next := updated.(model)
if cmd != nil {
t.Fatal("expected /tools to be handled without starting an agent run")
}
text := transcriptText(next.transcript)
for _, want := range []string{
"9 registered | registered catalog",
"- apply_patch",
"- bash",
"- edit_file",
"- glob",
"- grep",
"- read_file",
"- web_fetch",
"- web_search",
"- write_file",
} {
assertContains(t, text, want)
}
assertNotContains(t, text, "... 1 more")
previous := -1
for _, want := range []string{
"- apply_patch",
"- bash",
"- edit_file",
"- glob",
"- grep",
"- read_file",
"- web_fetch",
"- web_search",
"- write_file",
} {
current := strings.Index(text, want)
if current < 0 {
t.Fatalf("expected tools output to contain %q, got:\n%s", want, text)
}
if current <= previous {
t.Fatalf("expected tools output to keep sorted order at %q, got:\n%s", want, text)
}
previous = current
}
}
func TestContextAndPermissionsCommandsRenderProductState(t *testing.T) {
registry := tools.NewRegistry()
registry.Register(tools.NewReadFileTool("."))
@@ -140,9 +274,24 @@ func TestContextAndPermissionsCommandsRenderProductState(t *testing.T) {
t.Fatal("expected /context to be handled without starting an agent run")
}
contextText := transcriptText(next.transcript)
for _, want := range []string{"Context", "status: ok", "Runtime", "Session", "Tools", "cwd: D:\\codings\\Opensource\\Zero", "registered tools: 1"} {
for _, want := range []string{
"Context",
"go runtime | ask permissions | 1 tool",
"Runtime",
"cwd D:\\codings\\Opensource\\Zero",
"provider openai",
"model gpt-4.1",
"Session",
"active none",
"compaction idle",
"Tools",
"registered 1",
"actions: /permissions manage access | /tools inspect catalog",
} {
assertContains(t, contextText, want)
}
assertNotContains(t, contextText, "status: ok")
assertNotContains(t, contextText, "permission mode:")
next.input.SetValue("/permissions")
updated, cmd = next.Update(tea.KeyMsg{Type: tea.KeyEnter})
@@ -151,10 +300,88 @@ func TestContextAndPermissionsCommandsRenderProductState(t *testing.T) {
t.Fatal("expected /permissions to be handled without starting an agent run")
}
permissionText := transcriptText(next.transcript)
for _, want := range []string{"Permissions", "status: ok", "Permission mode: ask", "persistent grants: 1", "bash [allow/high]", "[REDACTED]"} {
for _, want := range []string{
"Permissions",
"ask permissions",
"1 persistent grant",
"State",
"mode ask",
"Grants",
"bash [allow/high]",
"[REDACTED]",
} {
assertContains(t, permissionText, want)
}
assertNotContains(t, permissionText, "sk-proj-sensitive")
assertNotContains(t, permissionText, "status: ok")
assertNotContains(t, permissionText, "Permission mode:")
}
// stubGrantStore is a minimal test double for sandbox.GrantStore. It only
// implements List, which is enough to exercise the error path of
// permissionsText().
type stubGrantStore struct {
listErr error
}
func (store *stubGrantStore) List() ([]sandbox.Grant, error) {
return nil, store.listErr
}
func TestPermissionsCommandCardHandlesNilStoreAndEmptyGrants(t *testing.T) {
nilStore := model{permissionMode: agent.PermissionModeAuto}.permissionsText()
for _, want := range []string{
"Permissions",
"auto permissions",
"grants unavailable",
"mode auto",
"persistent grants: unavailable",
} {
assertContains(t, nilStore, want)
}
assertNotContains(t, nilStore, "status: warning")
store, err := sandbox.NewGrantStore(sandbox.StoreOptions{FilePath: filepath.Join(t.TempDir(), "empty-grants.json")})
if err != nil {
t.Fatalf("NewGrantStore returned error: %v", err)
}
emptyText := model{permissionMode: agent.PermissionModeAsk, sandboxStore: store}.permissionsText()
for _, want := range []string{
"Permissions",
"ask permissions",
"no persistent grants",
"mode ask",
"none",
} {
assertContains(t, emptyText, want)
}
assertNotContains(t, emptyText, "status: info")
errStore := &stubGrantStore{listErr: errors.New("storage failure")}
errText := model{permissionMode: agent.PermissionModeAsk}.permissionsTextWithStore(errStore)
for _, want := range []string{
"Permissions",
"ask permissions",
"grants error",
"mode ask",
"error: storage failure",
} {
assertContains(t, errText, want)
}
assertNotContains(t, errText, "status: blocked")
}
func TestContextCommandCardHandlesNilRegistryAndStableStyle(t *testing.T) {
text := model{}.contextText()
for _, want := range []string{
"Context",
"0 tools",
"style balanced",
"root unknown",
} {
assertContains(t, text, want)
}
}
func TestCompactCommandAvoidsShellOnlyPlaceholder(t *testing.T) {
+154 -59
View File
@@ -10,40 +10,54 @@ import (
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/providercatalog"
"github.com/Gitlawb/zero/internal/sandbox"
"github.com/Gitlawb/zero/internal/tools"
"github.com/Gitlawb/zero/internal/zerocommands"
)
func (m model) toolsText() string {
registered := m.registry.All()
registered := m.registeredTools()
count := len(registered)
if len(registered) == 0 {
return renderCommandOutput(commandOutput{
Title: "Tools",
Status: commandStatusWarning,
Sections: []commandSection{{
return renderCommandCardTranscript(commandCard{
Title: "Tools",
Summary: []string{"0 registered", "no tools available"},
Sections: []commandCardSection{{
Title: "Registry",
Lines: []string{"registered tools: 0"},
Fields: []commandField{
{Key: "registered", Value: "0"},
},
}},
Actions: []string{"/mcp manage servers", "/permissions manage access"},
})
}
names := make([]string, 0, len(registered))
names := make([]string, 0, count)
for _, tool := range registered {
names = append(names, commandBullet(tool.Name()))
names = append(names, tool.Name())
}
sort.Strings(names)
return renderCommandOutput(commandOutput{
Title: "Tools",
Status: commandStatusOK,
Sections: []commandSection{
available := make([]string, 0, len(names))
for _, name := range names {
available = append(available, commandBullet(name))
}
return renderCommandCardTranscript(commandCard{
Title: "Tools",
Summary: []string{fmt.Sprintf("%d registered", count), "registered catalog"},
Sections: []commandCardSection{
{
Title: "Registry",
Lines: []string{fmt.Sprintf("registered tools: %d", len(names))},
Fields: []commandField{
{Key: "registered", Value: fmt.Sprint(count)},
},
},
{
Title: "Available",
Lines: names,
Lines: available,
},
},
Actions: []string{"/mcp manage servers", "/permissions manage access"},
})
}
@@ -247,35 +261,57 @@ func splitMCPCommandArgs(args string) ([]string, error) {
}
func (m model) permissionsText() string {
stateLines := []string{
"Permission mode: " + string(m.permissionMode),
}
if m.sandboxStore == nil {
return renderCommandOutput(commandOutput{
Title: "Permissions",
Status: commandStatusWarning,
Sections: []commandSection{
{Title: "State", Lines: stateLines},
{Title: "Sandbox grants:", Lines: []string{"persistent grants: unavailable"}},
return m.permissionsTextWithStore(nil)
}
return m.permissionsTextWithStore(m.sandboxStore)
}
func (m model) permissionsTextWithStore(store grantLister) string {
mode := string(m.permissionMode)
if store == nil {
return renderCommandCardTranscript(commandCard{
Title: "Permissions",
Summary: []string{mode + " permissions", "grants unavailable"},
Sections: []commandCardSection{
{
Title: "State",
Fields: []commandField{
{Key: "mode", Value: mode},
},
},
{
Title: "Grants",
Lines: []string{"persistent grants: unavailable"},
},
},
})
}
grants, err := m.sandboxStore.List()
grants, err := store.List()
if err != nil {
return renderCommandOutput(commandOutput{
Title: "Permissions",
Status: commandStatusBlocked,
Sections: []commandSection{
{Title: "State", Lines: stateLines},
{Title: "Sandbox grants:", Lines: []string{"error: " + err.Error()}},
return renderCommandCardTranscript(commandCard{
Title: "Permissions",
Summary: []string{mode + " permissions", "grants error"},
Sections: []commandCardSection{
{
Title: "State",
Fields: []commandField{
{Key: "mode", Value: mode},
},
},
{
Title: "Grants",
Lines: []string{"error: " + err.Error()},
},
},
})
}
snapshots := zerocommands.SandboxGrantSnapshots(grants)
grantLines := []string{fmt.Sprintf("persistent grants: %d", len(snapshots))}
grantRows := []commandRow{}
if len(snapshots) == 0 {
grantLines = append(grantLines, "none")
grantRows = append(grantRows, commandRow{Text: "none"})
} else {
for _, grant := range snapshots {
line := fmt.Sprintf("%s [%s/%s]", grant.ToolName, grant.Decision, grant.MaxAutonomy)
@@ -285,19 +321,45 @@ func (m model) permissionsText() string {
if grant.Reason != "" {
line += " - " + grant.Reason
}
grantLines = append(grantLines, commandBullet(line))
grantRows = append(grantRows, commandRow{Text: line})
}
}
return renderCommandOutput(commandOutput{
Title: "Permissions",
Status: commandStatusOK,
Sections: []commandSection{
{Title: "State", Lines: stateLines},
{Title: "Sandbox grants:", Lines: grantLines},
return renderCommandCardTranscript(commandCard{
Title: "Permissions",
Summary: []string{mode + " permissions", formatGrantCount(len(snapshots))},
Sections: []commandCardSection{
{
Title: "State",
Fields: []commandField{
{Key: "mode", Value: mode},
},
},
{
Title: "Grants",
Rows: grantRows,
},
},
})
}
// grantLister is the subset of sandbox.GrantStore used by permissionsText().
// It exists to let tests inject error-stub stores without reaching for a real
// filesystem path.
type grantLister interface {
List() ([]sandbox.Grant, error)
}
func formatGrantCount(count int) string {
if count == 0 {
return "no persistent grants"
}
if count == 1 {
return "1 persistent grant"
}
return fmt.Sprintf("%d persistent grants", count)
}
func (m model) providerText() string {
profileLines := []string{
"provider: " + displayValue(m.providerName, "none"),
@@ -461,42 +523,75 @@ func (m model) modelText(args string) string {
}
func (m model) contextText() string {
toolCount := len(m.registry.All())
return renderCommandOutput(commandOutput{
Title: "Context",
Status: commandStatusOK,
Sections: []commandSection{
toolCount := len(m.registeredTools())
return renderCommandCardTranscript(commandCard{
Title: "Context",
Summary: []string{
"go runtime",
string(m.permissionMode) + " permissions",
pluralizeCount(toolCount, "tool", "tools"),
},
Sections: []commandCardSection{
{
Title: "Runtime",
Lines: []string{
"cwd: " + displayValue(m.cwd, "unknown"),
"provider: " + displayValue(m.providerName, "none"),
"model: " + displayValue(m.modelName, "none"),
"permission mode: " + string(m.permissionMode),
"effort: " + m.effortDisplay(),
"style: " + m.responseStyle,
"usage: " + m.usageSummaryText(),
fmt.Sprintf("max turns: %d", m.agentOptions.MaxTurns),
Fields: []commandField{
{Key: "cwd", Value: displayValue(m.cwd, "unknown")},
{Key: "provider", Value: displayValue(m.providerName, "none")},
{Key: "model", Value: displayValue(m.modelName, "none")},
{Key: "effort", Value: m.effortDisplay()},
{Key: "style", Value: displayValue(m.responseStyle, defaultResponseStyle)},
{Key: "usage", Value: m.usageSummaryText()},
{Key: "max turns", Value: fmt.Sprint(m.agentOptions.MaxTurns)},
},
},
{
Title: "Session",
Lines: []string{
"active session: " + displayValue(m.activeSession.SessionID, "none"),
"session root: " + displayValue(m.sessionStore.RootDir, "unknown"),
"compaction: " + m.compactionStatus(),
Fields: []commandField{
{Key: "active", Value: displayValue(m.activeSession.SessionID, "none")},
{Key: "root", Value: displayValue(m.sessionRootDir(), "unknown")},
{Key: "compaction", Value: contextCompactionStatus(m.compactionStatus())},
},
},
{
Title: "Tools",
Lines: []string{
fmt.Sprintf("registered tools: %d", toolCount),
Fields: []commandField{
{Key: "registered", Value: fmt.Sprint(toolCount)},
},
},
},
Actions: []string{"/permissions manage access", "/tools inspect catalog"},
})
}
func (m model) registeredTools() []tools.Tool {
if m.registry == nil {
return nil
}
return m.registry.All()
}
func (m model) sessionRootDir() string {
if m.sessionStore == nil {
return ""
}
return m.sessionStore.RootDir
}
func pluralizeCount(count int, singular string, plural string) string {
label := plural
if count == 1 {
label = singular
}
return fmt.Sprintf("%d %s", count, label)
}
func contextCompactionStatus(status string) string {
if status == "not compacted" {
return "idle"
}
return status
}
func (m model) configText() string {
return renderCommandOutput(commandOutput{
Title: "Config",
+10 -9
View File
@@ -210,11 +210,19 @@ var commandDefinitions = []commandDefinition{
description: "Show or set post-edit self-correction depth (LSP-only default; on/tests/full add the project test plan; off/lsp disable tests, LSP-only).",
kind: commandSelfCorrect,
},
{
name: "/help",
usage: "/help",
group: commandGroupMeta,
description: "Show available commands.",
kind: commandHelp,
},
{
name: "/doctor",
usage: "/doctor",
aliases: []string{"/health"},
usage: "/doctor [fix|--connectivity]",
group: commandGroupRuntime,
description: "Show local diagnostics.",
description: "Show diagnostics, open safe fixes, or run provider connectivity checks.",
kind: commandDoctor,
},
{
@@ -246,13 +254,6 @@ var commandDefinitions = []commandDefinition{
description: "Show input style state.",
kind: commandInputStyle,
},
{
name: "/help",
usage: "/help",
group: commandGroupMeta,
description: "Show available commands.",
kind: commandHelp,
},
{
name: "/exit",
aliases: []string{"/quit"},
+345
View File
@@ -0,0 +1,345 @@
package tui
import (
"context"
"reflect"
"strings"
"testing"
tea "github.com/charmbracelet/bubbletea"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/doctor"
"github.com/Gitlawb/zero/internal/providerhealth"
)
func TestParseDoctorCommandArgsRejectsFixWithConnectivity(t *testing.T) {
_, _, _, err := parseDoctorCommandArgs("fix --connectivity")
if err == nil {
t.Fatal("parseDoctorCommandArgs accepted both fix and --connectivity")
}
if !strings.Contains(err.Error(), "choose either") {
t.Fatalf("error = %q, want mutually exclusive flag guidance", err)
}
}
func TestDoctorOptionsIncludeConfigPaths(t *testing.T) {
m := newModel(context.Background(), Options{
UserConfigPath: "C:/zero/user.json",
ProjectConfigPath: "C:/repo/.zero/config.json",
ProviderProfile: config.ProviderProfile{
Name: "openai",
ProviderKind: config.ProviderKindOpenAI,
Model: "gpt-4.1",
},
})
options := m.doctorOptions(false)
if options.UserConfig != "C:/zero/user.json" {
t.Fatalf("UserConfig = %q, want %q", options.UserConfig, "C:/zero/user.json")
}
if options.ProjectConfig != "C:/repo/.zero/config.json" {
t.Fatalf("ProjectConfig = %q, want %q", options.ProjectConfig, "C:/repo/.zero/config.json")
}
if options.Connectivity {
t.Fatal("Connectivity = true, want false")
}
if options.ProviderHealth != nil {
t.Fatalf("ProviderHealth = %#v, want nil without connectivity", options.ProviderHealth)
}
}
func TestDoctorOptionsConnectivityInvokesConfiguredProviderHealthProbe(t *testing.T) {
profile := config.ProviderProfile{
Name: "custom",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "https://api.example.com/v1",
Model: "custom-model",
}
var called int
var gotOptions providerhealth.Options
m := newModel(context.Background(), Options{
ProviderProfile: profile,
UserAgent: "zero-test",
ProbeProviderHealth: func(_ context.Context, options providerhealth.Options) providerhealth.Result {
called++
gotOptions = options
return providerhealth.Result{
Status: providerhealth.StatusPass,
Checks: []providerhealth.Check{{
ID: "provider.connectivity",
Status: providerhealth.StatusPass,
Message: "reachable",
}},
}
},
})
options := m.doctorOptions(true)
report := doctor.Run(options)
if called != 1 {
t.Fatalf("probe called %d time(s), want 1", called)
}
if !gotOptions.Connectivity {
t.Fatal("probe Connectivity = false, want true")
}
if gotOptions.UserAgent != "zero-test" {
t.Fatalf("probe UserAgent = %q, want %q", gotOptions.UserAgent, "zero-test")
}
if !reflect.DeepEqual(gotOptions.Profile, profile) {
t.Fatalf("probe Profile = %#v, want %#v", gotOptions.Profile, profile)
}
check := report.Check("provider.connectivity")
if options.ProviderHealth == nil || check == nil || check.Status != doctor.StatusPass {
t.Fatalf("connectivity check = %#v, ProviderHealth = %#v; want passing injected health", check, options.ProviderHealth)
}
}
func TestDoctorConnectivityCommandRunsProbeAsynchronously(t *testing.T) {
profile := config.ProviderProfile{
Name: "custom",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "https://api.example.com/v1",
Model: "custom-model",
}
called := false
m := newModel(context.Background(), Options{
ProviderProfile: profile,
ProbeProviderHealth: func(context.Context, providerhealth.Options) providerhealth.Result {
called = true
return providerhealth.Result{
Status: providerhealth.StatusPass,
Checks: []providerhealth.Check{{
ID: "provider.connectivity",
Status: providerhealth.StatusPass,
Message: "reachable",
}},
}
},
})
m.input.SetValue("/doctor --connectivity")
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
next := updated.(model)
if cmd == nil {
t.Fatal("expected /doctor --connectivity to return an async command")
}
if called {
t.Fatal("provider probe ran synchronously before the returned command executed")
}
if !transcriptContains(next.transcript, "checking provider connectivity") {
t.Fatalf("expected running doctor status, got %#v", next.transcript)
}
msg := execCmd(cmd)
if !called {
t.Fatal("provider probe did not run when the async command executed")
}
updated, _ = next.Update(msg)
final := updated.(model)
for _, want := range []string{"Diagnostics", "[pass] provider.connectivity", "reachable"} {
if !transcriptContains(final.transcript, want) {
t.Fatalf("expected final doctor transcript to contain %q, got %#v", want, final.transcript)
}
}
}
func TestDoctorCommandUsesDiagnosticCenterRow(t *testing.T) {
m := newModel(context.Background(), Options{
ProviderProfile: config.ProviderProfile{
Name: "custom",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "https://api.example.com/v1",
Model: "custom-model",
},
})
m.input.SetValue("/doctor")
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
next := updated.(model)
if cmd != nil {
t.Fatal("expected plain /doctor to render synchronously")
}
row := newestDoctorStatusRow(next.transcript)
if row == nil {
t.Fatalf("expected /doctor to render a doctor status row, got %#v", next.transcript)
}
if row.id != doctorStatusRowID {
t.Fatalf("expected /doctor row id %q, got %q", doctorStatusRowID, row.id)
}
text := row.text
for _, want := range []string{"Diagnostics", "checks need attention", "Actions"} {
if !strings.Contains(text, want) {
t.Fatalf("doctor row missing %q:\n%s", want, text)
}
}
for _, unwanted := range []string{"Generated", "Checks", "[pass]"} {
if strings.Contains(text, unwanted) {
t.Fatalf("doctor row should hide %q:\n%s", unwanted, text)
}
}
}
func TestDoctorConnectivityCommandAnimatesAndReplacesStatusRow(t *testing.T) {
profile := config.ProviderProfile{
Name: "custom",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "https://api.example.com/v1",
Model: "custom-model",
}
m := newModel(context.Background(), Options{
ProviderProfile: profile,
ProbeProviderHealth: func(context.Context, providerhealth.Options) providerhealth.Result {
return providerhealth.Result{
Status: providerhealth.StatusPass,
Checks: []providerhealth.Check{{
ID: "provider.connectivity",
Status: providerhealth.StatusPass,
Message: "reachable",
}},
}
},
})
m.input.SetValue("/doctor --connectivity")
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
next := updated.(model)
if cmd == nil {
t.Fatal("expected /doctor --connectivity to start an async command")
}
beforeRows := len(next.transcript)
before := newestDoctorStatusText(next.transcript)
updated, _ = next.Update(next.spinner.Tick())
ticked := updated.(model)
after := newestDoctorStatusText(ticked.transcript)
if before == "" || after == "" {
t.Fatalf("expected doctor status row before=%q after=%q", before, after)
}
if before == after {
t.Fatalf("expected doctor status to animate on tick, still %q", after)
}
updated, _ = ticked.Update(execCmd(cmd))
final := updated.(model)
if len(final.transcript) != beforeRows {
t.Fatalf("expected final doctor report to replace the running row, rows before=%d after=%d: %#v", beforeRows, len(final.transcript), final.transcript)
}
if got := countDoctorDiagnosticRows(final.transcript); got != 1 {
t.Fatalf("expected one doctor diagnostic row after completion, got %d: %#v", got, final.transcript)
}
if transcriptContains(final.transcript, "checking provider connectivity") {
t.Fatalf("expected stale running copy to be replaced, got %#v", final.transcript)
}
if !transcriptContains(final.transcript, "[pass] provider.connectivity") {
t.Fatalf("expected completed diagnostics result, got %#v", final.transcript)
}
}
func TestDoctorFixOpensProviderWizardWhenProviderMissing(t *testing.T) {
m := newModel(context.Background(), Options{})
m.input.SetValue("/doctor fix")
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
next := updated.(model)
if cmd != nil {
t.Fatal("expected /doctor fix provider setup path to be handled synchronously")
}
if next.providerWizard == nil {
t.Fatalf("expected /doctor fix to open provider setup wizard, transcript %#v", next.transcript)
}
if !transcriptContains(next.transcript, "Opening provider setup") {
t.Fatalf("expected doctor fix transcript to explain provider setup, got %#v", next.transcript)
}
}
func TestDoctorFixLinesShowNoAutomaticFixForUnmappedFailure(t *testing.T) {
lines := doctorFixLines(doctor.Report{Checks: []doctor.Check{{
ID: "provider.auth",
Status: doctor.StatusFail,
Message: "API key rejected.",
}}})
text := strings.Join(lines, "\n")
if strings.Contains(text, "already clean") {
t.Fatalf("doctorFixLines reported clean diagnostics for an unmapped failure:\n%s", text)
}
if !strings.Contains(text, "No automatic fixes are available") {
t.Fatalf("doctorFixLines = %q, want no automatic fix guidance", text)
}
}
func newestDoctorStatusText(rows []transcriptRow) string {
if row := newestDoctorStatusRow(rows); row != nil {
return row.text
}
return ""
}
func newestDoctorStatusRow(rows []transcriptRow) *transcriptRow {
for i := len(rows) - 1; i >= 0; i-- {
row := rows[i]
if row.kind == rowSystem && row.id == doctorStatusRowID {
return &rows[i]
}
}
return nil
}
func countDoctorDiagnosticRows(rows []transcriptRow) int {
count := 0
for _, row := range rows {
if row.kind == rowSystem && row.id == doctorStatusRowID {
count++
}
}
return count
}
func TestDoctorFixRunsConnectivityWhenProviderConfigured(t *testing.T) {
profile := config.ProviderProfile{
Name: "custom",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "https://api.example.com/v1",
Model: "custom-model",
}
called := false
m := newModel(context.Background(), Options{
ProviderProfile: profile,
ProbeProviderHealth: func(context.Context, providerhealth.Options) providerhealth.Result {
called = true
return providerhealth.Result{
Status: providerhealth.StatusPass,
Checks: []providerhealth.Check{{
ID: "provider.connectivity",
Status: providerhealth.StatusPass,
Message: "reachable",
}},
}
},
})
m.input.SetValue("/doctor fix")
updated, cmd := m.Update(tea.KeyMsg{Type: tea.KeyEnter})
next := updated.(model)
if cmd == nil {
t.Fatal("expected /doctor fix to run provider connectivity asynchronously")
}
if called {
t.Fatal("provider probe ran synchronously before the returned command executed")
}
if !transcriptContains(next.transcript, "checking provider connectivity") {
t.Fatalf("expected running doctor status, got %#v", next.transcript)
}
updated, _ = next.Update(execCmd(cmd))
final := updated.(model)
if !called {
t.Fatal("provider probe did not run when /doctor fix command executed")
}
if !transcriptContains(final.transcript, "[pass] provider.connectivity") {
t.Fatalf("expected connectivity result in transcript, got %#v", final.transcript)
}
}
+198
View File
@@ -0,0 +1,198 @@
package tui
import (
"fmt"
"strings"
"github.com/Gitlawb/zero/internal/doctor"
"github.com/Gitlawb/zero/internal/zerocommands"
)
func doctorCommandOutput(report doctor.Report, backend *zerocommands.BackendLifecycleSnapshot) commandOutput {
sections := []commandSection{{
Title: "Summary",
Lines: doctorSummaryLines(report.Checks),
}}
provider, platform, other := doctorCheckSections(report.Checks)
sections = appendNonEmptyDoctorSection(sections, "Provider", provider)
sections = appendNonEmptyDoctorSection(sections, "Platform", platform)
sections = appendNonEmptyDoctorSection(sections, "Other", other)
if backend != nil {
sections = append(sections, doctorBackendSection(*backend))
}
if actions := doctorActions(report.Checks, backend); len(actions) > 0 {
sections = append(sections, commandSection{
Title: "Actions",
Lines: actions,
})
}
return commandOutput{
Title: "Diagnostics",
Status: doctorCommandStatus(report),
Sections: sections,
}
}
func doctorCommandStatus(report doctor.Report) commandStatus {
hasWarning := false
for _, check := range report.Checks {
switch check.Status {
case doctor.StatusFail:
return commandStatusBlocked
case doctor.StatusWarn:
hasWarning = true
}
}
if hasWarning {
return commandStatusWarning
}
return commandStatusOK
}
func doctorCheckSections(checks []doctor.Check) (provider []commandRow, platform []commandRow, other []commandRow) {
for _, check := range checks {
if check.Status == doctor.StatusPass && check.ID != "provider.connectivity" {
continue
}
row := commandRow{Text: doctorCheckRow(check)}
switch doctorCheckGroup(check.ID) {
case "provider":
provider = append(provider, row)
case "platform":
platform = append(platform, row)
default:
other = append(other, row)
}
}
return provider, platform, other
}
func doctorCheckGroup(id string) string {
switch id {
case "provider.config", "provider.model", "provider.connectivity", "provider.auth", "provider.runtime":
return "provider"
case "sandbox.backend", "lsp.servers", "runtime.go", "config.files", "config.validation":
return "platform"
default:
return ""
}
}
func doctorSummaryLines(checks []doctor.Check) []string {
pass, warn, fail := doctorStatusCounts(checks)
attention := warn + fail
if attention == 0 {
return []string{
fmt.Sprintf("%s healthy", doctorCheckCountLabel(pass)),
"All core systems are ready.",
}
}
return []string{
fmt.Sprintf("%s %s attention", doctorCheckCountLabel(attention), doctorAttentionVerb(attention)),
fmt.Sprintf("%s healthy", doctorCheckCountLabel(pass)),
}
}
func doctorCheckCountLabel(count int) string {
if count == 1 {
return "1 check"
}
return fmt.Sprintf("%d checks", count)
}
func doctorAttentionVerb(count int) string {
if count == 1 {
return "needs"
}
return "need"
}
func doctorStatusCounts(checks []doctor.Check) (pass int, warn int, fail int) {
for _, check := range checks {
switch check.Status {
case doctor.StatusPass:
pass++
case doctor.StatusWarn:
warn++
case doctor.StatusFail:
fail++
}
}
return pass, warn, fail
}
func doctorCheckRow(check doctor.Check) string {
parts := []string{fmt.Sprintf("[%s]", check.Status)}
if check.ID != "" {
parts = append(parts, check.ID)
}
if check.Message != "" {
parts = append(parts, "-", check.Message)
}
return strings.Join(parts, " ")
}
func appendNonEmptyDoctorSection(sections []commandSection, title string, rows []commandRow) []commandSection {
if len(rows) == 0 {
return sections
}
return append(sections, commandSection{
Title: title,
Rows: rows,
})
}
func doctorBackendSection(backend zerocommands.BackendLifecycleSnapshot) commandSection {
return commandSection{
Title: "Backend",
Fields: []commandField{
{Key: "MCP servers", Value: fmt.Sprintf("%d", len(backend.MCPServers))},
{Key: "Hooks", Value: fmt.Sprintf("%d", len(backend.Hooks))},
{Key: "Plugins", Value: fmt.Sprintf("%d", len(backend.Plugins))},
},
}
}
func doctorActions(checks []doctor.Check, backend *zerocommands.BackendLifecycleSnapshot) []string {
seen := map[string]bool{}
actions := []string{}
add := func(action string) {
action = strings.TrimSpace(action)
if action == "" || seen[action] {
return
}
seen[action] = true
actions = append(actions, action)
}
for _, check := range checks {
if check.Status == doctor.StatusPass {
continue
}
message := strings.ToLower(check.Message)
switch check.ID {
case "provider.config", "provider.model":
add("/provider - configure the active provider and model")
add("/doctor fix - open guided diagnostics repair")
case "provider.connectivity":
add("/doctor --connectivity - probe the provider endpoint")
case "sandbox.backend":
if strings.Contains(message, "windows") || strings.Contains(message, "policy-only") {
add("WSL2 or Linux container - enable native sandbox isolation on Windows")
}
case "lsp.servers":
if strings.Contains(message, "missing") {
add("install missing language servers - restore code intelligence on PATH")
}
}
}
if backend != nil {
add("/mcp - inspect MCP servers")
add("/hooks - inspect hooks")
add("/plugins - inspect plugins")
}
return actions
}
Loaded 100 of 122 files, more files were not shown because too many files have changed in this diff. Show more