Compare commits

...
Author SHA1 Message Date
e034a93815 perf(tui): bounded file read, LRU rendering cache, and size caps for full-file view (#1203)
* perf(tui): bounded file read, LRU rendering cache, and size caps for full-file view

- Bound full-file view reads (4000 lines, 1 MiB total, 4 KiB per line) with async loading and an LRU render cache capped at 64 entries and 8 MiB.
- Refresh both BTW surfaces when a tool result (such as write_stdin) carries both exact changed paths and an unknown-scope or aggregated directory summary outside Git.
- Gate Run details FILES mouse hit-testing on Run details being the visible foreground overlay so higher-priority modals (speech API-key prompt, leader help) block clicks behind them.

(cherry picked from commit 100d847aaf)

* fix(tui): gate full-view read errors on snapshot validity

renderFileViewFull returned renderedContent whenever hasError was set,
before the snapshot-validity check that governs every other completed
snapshot. A /btw fork of a full view that had settled on a read error
invalidates the inherited snapshot and clears renderedContent, but the
error latch survived, so the side pane rendered an empty body instead of
the loading placeholder until its own load landed. Reloads of an errored
view (resize, refresh, theme) likewise kept showing the stale error.

Serve the error snapshot through the same validity predicate as content
snapshots, so a pending or invalidated load falls through to "Loading…".
Error displays still do not count as cache hits.

Co-Authored-By: Claude Opus 5.5 <[email protected]>
(cherry picked from commit a1887f4b03)

* fix(tui): keep file-view revisions monotonic and revalidate on git sweeps

Shared path revisions reset to zero on every cache purge, while each surface
kept its own required revision. Three sites could also raise that revision
past the shared one by incrementing it locally. After a purge, an exact-path
edit made on the other BTW surface could produce a shared revision no higher
than this surface's own. recoverInvalidatedFileView then saw nothing new, and
the view stayed stale. Revisions now come from one cache-wide clock that
survives purges. A purge raises the epoch floor to that clock instead of
resetting it, and each surface adopts the revision the cache returns.

Every git sweep invalidated each swept path and forced a full refresh. The
swept list is cumulative since the baseline, and each live sweep is launched
beside the tool-result, turn-end or BTW-return handling that already reloads
or rechecks the open view. The sweep that followed therefore reread and
re-highlighted an unchanged file and flashed "Loading…". A sweep now
revalidates the open view instead. It waits for an in-flight load to land,
then rereads the bounded source in the background and compares it with the
displayed snapshot's identity (hash and truncation flags). The view reloads
only if the bytes changed. A mutation the swept list omits is still caught.

Regression tests (all fail on 311290c0 and pass here):
- TestFileViewCrossSurfaceExactEditAfterPurgeReloadsOtherSurface
- TestFileViewGitSweepRevalidatesWithoutDuplicateReload
- TestFileViewGitSweepDuringLoadDefersToThatLoad
- TestFileViewStaleRevalidationDoesNotReloadNewerSnapshot
- TestFileViewCachePurgeResetsPathRevisions (now asserts monotonicity)

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

* refactor(tui): return source identity from loadAndRender directly

1ce4f96e added loadAndRenderSource and left loadAndRender as a two-value
wrapper. After that change only tests called the wrapper, so make deadcode
newly reported it. Fold the wrapper away: loadAndRender now returns the
rendered text, the source identity, and the error, and its callers take the
value they need. Behavior is unchanged.

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

* fix(tui): close three file-view staleness gaps from review

- fileViewSourceHash joined lines with separators, so an empty file (no
  lines) and a file holding one newline (one empty line) hashed alike.
  Sweep revalidation compares identities without the file size, so adding
  or removing that blank line went undetected. Every line is now
  terminated, which keeps the encoding injective.
- A load cancelled inside the bounded read returned errFileViewSuperseded
  with no lines. loadAndRender then treated it as a read error and evicted
  the path, discarding an entry a newer load had already committed. The
  sentinel now returns before the eviction branch.
- A cancelled run keeps going until it finishes, and agentRowMsg ignores its
  late tool results. Its completion only flushed session events, so a write
  that finished during cancellation left an open full-file view stale. The
  flush completion now runs the unknown-scope refresh.

Regression tests (each fails on ca1cdaf3 and passes here):
- TestFileViewGitSweepDetectsBlankLineAddedToEmptyFile
- TestFileViewSupersededReadKeepsNewerCacheEntry
- TestFileViewCancelledRunCompletionRefreshesOpenView

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

* fix(tui): gate Run details Enter on the visible overlay

The Run details Enter drill-in keyed off the raw runDetailsOpen flag,
while mouse FILES picks use runDetailsVisible(). The flag outlives the
overlay: clicking a specialist card beside Run details enters the
subchat, which hides the overlay but leaves runDetailsOpen set, and
Enter then opened the parent's file view from inside the subchat.

Use runDetailsVisible() and the shared runDetailsInnerWidth() geometry
on the keyboard path, matching the mouse path. The branch's existing
Esc/Ctrl+B and key-swallow behavior is unchanged.

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

---------

Co-authored-by: hazyhaar <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
2026-11-07 10:28:29 -08:00
2ca84cc127 fix(hooks): dispatch specialistStart/Stop lifecycle hooks (#1204)
* fix(hooks): dispatch specialistStart/Stop and allow matchers

These events were documented and accepted by hooks add, but never fired.
Wire advisory dispatch from specialist accounting and allow specialist
name matchers so EXTENDING.md matches runtime behavior.

* fix(hooks): unlock before specialistStop hooks and fire once

Release accountingMu before running lifecycle hooks, and dispatch
specialistStop only when the stop event is newly recorded so concurrent
finishers cannot re-run arbitrary hook commands.

* fix(hooks): fire specialistStop without relying on session append

Decouple stop-hook dispatch from parent-session event append so swarm
members (empty parent) and TaskOutput still get exactly one hook via a
shared LifecycleHooks claim, and wire that bridge into TaskOutput.

* fix(hooks): claim specialistStop durably and gate hooks on real launches

Address the review findings on specialist lifecycle dispatch:

- specialistStop is now claimed in the parent session's event log via a
  record-once specialist_stop_hook event, sharing the accounting run-ID
  equivalence (empty runId is a catch-all). A later reaper in another
  process (TaskOutput poll, recovered TaskStop, restart) no longer re-runs
  the hook. The in-process claim remains only as fallback for swarm
  members with no parent session and for session-store errors.
- Recovered TaskStop records the terminal transition after Kill, so an
  explicitly stopped recovered task emits specialistStop exactly once.
- specialistStart fires only after the child process actually started
  (onStarted after Start()); a launch that fails before exec records
  attempt accounting but no public start/stop hooks. Reapers derive the
  same eligibility from the durable stop event's childStarted flag.

Adds a CLI startup wire-up test asserting both real entry points
(interactive TUI, headless exec) attach the lifecycle bridge to the same
LifecycleHooks the TaskOutput tool holds, fixes the stale import path on
the rebased branch, and documents the once-per-run semantics.

* fix(hooks): gate background exit accounting behind the start decision

A background child that exits immediately can run its launcher exit
callback before Executor.LaunchBackground returns, so the terminal
accounting could claim and dispatch specialistStop before runBackground
dispatched specialistStart, and the recorded stop missed the pid that
SetPID had not stored yet.

The exit callback now waits on a settled channel closed when runBackground
returns (after the start hook and SetPID), from its own goroutine so a
synchronous injected launcher cannot deadlock waiting for a signal only
the launch return can close.

Regression test: TestRunBackgroundDefersExitUntilStartDecisionCompletes
drives a synchronous fast-exit launcher and asserts start-before-stop hook
order, the recorded stop pid, and no deadlock. It fails on the previous
code with hooks "specialistStop,specialistStart".

* fix(hooks): record specialistStop only after the orphan kill succeeds

When SetPID fails after a background child has started, runBackground kills
the unrecorded child and records the stop. If TerminateProcess fails, the
child may still be running, but the branch still marked the task errored,
recorded the stop, took the durable specialist_stop_hook claim and fired
specialistStop. The hook could then run while the sub-agent was alive, and
the one-shot claim suppressed the hook on the child's real exit. Marking the
task non-running also let a TaskOutput or TaskStop reap claim the stop
before the launcher's Wait observed the exit.

On a failed kill, surface the joined error and return without the status
update or the stop. The task stays running, so reapers remain no-ops, and the
launcher's exit callback marks it exited and records the single stop when the
child ends. A successful kill keeps the immediate stop as before.

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

* fix(hooks): keep the prompt on a failed orphan kill and share one stop claim

When SetPID fails and terminating the unrecorded child also fails, the child
may still be running and the task stays running. That branch still deleted
the --file prompt, so a child that had not yet read it failed at startup.
The launcher's exit callback already removes the prompt once the child
really exits, and Runtime.Close removes tracked prompts at shutdown, so the
failed-kill branch no longer removes it.

dispatchSpecialistStopHook made two independent dispatch decisions: the
durable specialist_stop_hook claim, and the in-process claim used when the
session store errors. A won durable claim never recorded the in-process
claim, so on one bridge a transient store error before or after a durable
claim dispatched the same stop twice, breaking the documented at-most-once
per process. A won durable claim now also takes the in-process claim
(including for an ineligible stop), and the hook runs only if both are free.

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

---------

Co-authored-by: Ayush7614 <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
2026-11-07 10:19:56 -08:00
05c1b5e160 feat(plugins): add CLI enable/disable for plugin.json (#1202)
* feat(plugins): add CLI enable/disable for plugin.json

Add zero plugins enable|disable so installed plugins can be toggled
without removing them, matching hooks/MCP parity and EXTENDING docs.

* fix(plugins): serialize toggles with installs and preserve manifest bytes

Address review findings on the enable/disable CLI:

- Take the same cross-process installtxn lock plugins add/remove take on
  the plugin's install root before the final lookup and read, and publish
  plugin.json with installtxn.WriteFileAtomically while holding it.
- Re-resolve the plugin under the lock so a concurrent install can never
  be overwritten with manifest bytes read before it committed, retrying
  when precedence moves to the other root; a plugin removed under the
  lock reports ErrNotInstalled instead of being recreated.
- Edit the enabled value in place (a targeted token walk) so field
  order, formatting, and numeric literals survive; the old
  decode/re-encode dropped precision (9007199254740993 became
  9007199254740992) and sorted the keys.
- Fail closed on manifest stat errors, refresh a clean plugins.lock hash
  after a toggle, and leave pre-existing drift in place instead of
  absorbing it into the lockfile.
- Redact the text-mode output, reject an empty plugin id as a usage
  error, and document project-over-user precedence in the help text.

* fix(plugins): recover interrupted install transactions before toggling

SetEnabledByID takes the plugin root's installtxn lock but went straight
to reading and rewriting the tree. installtxn.Recover requires every
holder of that lock to resolve interrupted transactions first and to
abort on its error, as Install and Remove already do.

Without it, a toggle after an install killed between its tree swap and
its lockfile publish edited the unpublished tree, and the next install
or removal restored the retained backup over it, silently undoing the
enable/disable. A transaction recovery could not resolve also failed to
stop the toggle.

Call installtxn.Recover with the plugins reconciler right after taking
the lock, and add regression tests for every interrupted state that
leaves the plugin discoverable, the unresolved-transaction abort, and a
plugin removed while the toggle waits on the lock.

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

* docs(plugins): say when a plugin toggle takes effect

Zero reads plugin.json "enabled" only when it activates plugins at
startup (runExec and the interactive TUI bootstrap), and nothing reloads
them afterwards. A running interactive session therefore keeps the
tools, hooks, and skills of a plugin disabled from another terminal.

State that in the plugins docs and in the enable/disable help so the
documented activation rule matches what the command does.

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

* fix(plugins): pin toggled manifests, redact toggle errors, honor last enabled key

Address the review findings on the plugin enable/disable toggle.

Pin the plugin directory before editing its manifest. The install lock
only serializes cooperating installers, so swapping the plugin directory
for a link after discovery redirected the read and the atomic publish
through it, overwriting another manifest. The toggle now opens the
directory under its install root with pathjail, refuses a link or
reparse point, checks that the handle is the directory entry itself,
and reads, stats, creates the temporary file, and renames through that
handle.

Redact every stderr message the toggle writes. Parser and not-installed
usage errors echoed the raw argument, so a mistyped credential reached
stderr unredacted while the toggle's other output was redacted.

Validate only the last "enabled" value, as the loader does. The loader
keeps the last duplicate key, but the toggle rejected an earlier,
overwritten non-boolean value, so a plugin discovered as enabled could
not be toggled.

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

* fix(plugins): keep the toggle on the pinned plugin directory throughout

The previous fix pinned the manifest read and publish to a handle on the
plugin directory but still hashed the tree for plugins.lock by path, and
reported success without checking that the path still named the pinned
directory. If the directory was replaced after it was opened, the edit
landed in the moved-away original while the lockfile check and refresh
read the replacement, and the command reported a state the installed
plugin did not have.

Hash the tree before and after the edit through the pinned handle, and
confirm that the plugin's directory entry is still the pinned directory
before refreshing the lockfile or reporting any result, including a
no-op. hashTree now delegates to the same rooted walker, so path and
handle callers hash identical bytes.

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

* fix(plugins): hash plugin trees through native os.Root calls

20b484f moved hashTree onto an os.Root handle but walked it through
handle.FS(). That adapter accepts only fs.ValidPath names, which must be
UTF-8, so a plugin holding a Linux-valid name such as byte 0xff failed
with "lstat ...: invalid argument". copyTree and the previous walker
accept those names, so Install, the toggle's lock refresh, and Info's
drift check all regressed for such plugins.

Walk the handle with its own Open/Lstat methods instead. The walk stays
pinned to the directory, and the hash bytes match the merge-base
implementation (checked on a fixture with nested and non-UTF-8 names,
an exec bit, a symlink, .git and a FIFO).

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

---------

Co-authored-by: Ayush7614 <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
2026-11-07 09:56:16 -08:00
16ccd78765 feat(mcp): add AI/ML API as an opt-in built-in server (#1193)
* feat(aimlapi): shared service key and attribution for aimlapi.com services

Adds the pieces the Toolbox and MCP integrations share:

- internal/aimlapi: X-AIMLAPI-Source (agent/zero), the fixed Toolbox and
  MCP URLs, and an exact-hostname check (https, default port, no
  userinfo) for api/tools/mcp.aimlapi.com before the key or attribution
  headers are attached. The partner id still comes from ResolvePartnerID.
- internal/config: AIMLAPIServiceKey picks the key from the aimlapi
  profile (or AIMLAPI_API_KEY without one), reading the credential
  store, and withholds it when any aimlapi-keyed profile points at a
  proxy base URL or AIMLAPI_INFERENCE_URL is overridden. Keys that are
  not a single token are refused.

* feat(mcp): AI/ML API MCP server as an opt-in built-in default

Seeds `aimlapi` (https://mcp.aimlapi.com/mcp, streamable HTTP) next to
Exa as a built-in MCP server that ships disabled: its tool runs and
generations are billed to the user's AI/ML API balance. The disable is
marked as explicitly set, so a project config cannot lift it; the user
turns it on with `zero mcp enable aimlapi`, which now writes
`"disabled": false` for a default that ships off (and `zero mcp add
aimlapi` does the same). `zero doctor` reports the untouched default as
a pass rather than a warning.

The MCP client attaches the AI/ML API key and the X-AIMLAPI-Source /
X-AIMLAPI-Partner-ID headers in one place (internal/mcp/aimlapi.go),
called from the three header-application points. Both apply only to a
request whose own URL is exactly https://mcp.aimlapi.com on the default
port, whatever the server is named. The key is resolved for every
request through config.LoadAIMLAPIServiceKey, so switching the aimlapi
provider to a proxy, or a key that is not a single token, means no
Authorization header on the next request. It is never put in config,
the server identity used for OAuth token binding, list output or errors.
A user Authorization header, `"auth": "oauth"`, or a project-touched
server goes without it. A 401 from that host explains how to add a key
or sign in with `zero mcp oauth login`.

* fix(aimlapi): validate whitespace and preserve store errors

* fix(mcp): tailor unauthorized hints to credential policy

* docs(mcp): clarify automatic credential eligibility

* test(mcp): cover edited disabled built-in defaults

* fix(aimlapi): report store reads after proxy eligibility checks

* docs(aimlapi): document MCP key precedence and clarify 401 hints

Operator copy now matches normalizeProvider resolution (inline key,
apiKeyEnv, then credential store) and warns that AIMLAPI_API_KEY can
bill MCP before a stored profile key. SECURITY and service-key godoc
stay aligned. 401 hints no longer imply a proxy-only failure mode.

Adds a regression test that resolved profile keys win over the store.

Co-authored-by: Cursor <[email protected]>

* fix(mcp): keep the aimlapi key off redirects and error text, persist enable

Redirect hops: net/http copies the first request's headers onto every
redirect hop without calling the header hooks, so the automatic AI/ML API
key and attribution followed a same-origin hop that fails the exact-URL
rule (userinfo) and an eligible hop reused the first hop's key. The
redirect guard now restores the server's own values for those headers and
applies the per-request rule again.

Error text: httpStatusError copied up to 1 KiB of the response body into
the error, so a body echoing the injected key (or its prefix, when cut by
the read limit) reached CLI output. When the request carried the automatic
key, the body is now omitted instead of matched.

Enable persistence: MCPServerConfig had no serializer and Disabled is
omitempty, so any read-modify-write of config.json (SetTheme,
SetActiveProvider, ...) dropped `"disabled": false` and the shipped
disable turned aimlapi off again. MarshalJSON now writes the flag whenever
it was explicitly set; `zero mcp add` keeps writing only a disable.

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

---------

Co-authored-by: Stan <[email protected]>
Co-authored-by: Hugo <[email protected]>
Co-authored-by: Cursor <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
2026-11-04 20:13:53 -08:00
c9f179482a feat(web-fetch): add an opt-in AI/ML API Toolbox scraping backend (#1192)
* feat(web-fetch): AI/ML API Toolbox as an opt-in web_fetch backend

With ZERO_WEBFETCH_PROVIDER=aimlapi, web_fetch fetches the page as Markdown
through the Toolbox tool search1api/crawl instead of a direct GET. A key
alone never selects it; unset or unknown values keep the direct fetch.

- The URL checks (pre-permission and resolved host) run before the URL is
  sent, so loopback, private and metadata URLs never reach the crawler.
- The key is resolved on every call through LoadAIMLAPIServiceKey; a
  missing, proxied or invalid key is an error with no request and no
  fallback. A failed run (is_error, failed status, no content) is also an
  error.
- Requests go only to aimlapi.ToolboxBaseURL, checked with
  IsAttributedURL before the key and attribution headers are attached;
  redirects are refused. The client uses the default transport, so the
  proxy environment applies and the target-site SSRF pin does not.
- Every run carries max_cost_usd 0.005, wait_ms 30000 and a fresh
  idempotency key. 202 runs are polled (hint clamped to 250 ms..5 s,
  120 s overall). 429/502/503/504 are retried with the identical body,
  honouring Retry-After (capped at 10 s); max_cost_exceeded is not.
- Output keeps the result shape with a Backend line, respects max_bytes,
  refuses format raw, and scrubs the key and redacts all Toolbox text.

* fix(web-fetch): honor poll retries and reject raw before permission

* docs(web-fetch): explain the bounded run response

* test(web-fetch): rename the Toolbox fake that collides with web_search tests

#1189 added a fakeToolbox type with a recorded method to the tools test
package. The web_fetch Toolbox tests declared their own fakeToolbox with a
different shape, so the package's tests no longer compiled once both were on
main. Rename the web_fetch one to webFetchFakeToolbox, matching the file's
webFetch prefix. Test-only; no behavior change.

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

* fix(web-fetch): describe the aimlapi backend in the tool schema

With ZERO_WEBFETCH_PROVIDER=aimlapi the max_bytes and format descriptions
still promised direct-fetch behavior: raw body bytes before conversion and
HTML conversion with a raw option. The Toolbox backend returns crawled
Markdown, applies max_bytes to that Markdown and refuses format "raw", as
the README says. Say so in both descriptions so the model sees the same
contract.

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

* fix(web-fetch): keep the AI/ML API key out of Toolbox errors and output

The Toolbox backend scrubbed the key with a literal string replace. A key
echoed back split by an invisible character (NUL, ZWSP) passed it on every
path the model sees: error bodies, run errors, page content and the run_id
meta field. scrubWebFetchKey now uses scrubAIMLAPIResult, the same
split-aware redaction the web_search backend uses.

A failed Toolbox request also returned the *url.Error from client.Do. That
error quotes the request URL, and a polling URL carries the run id
percent-escaped, a form no key scrubbing can match. send now reports the
cause without the URL, as web_search does.

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

* docs(web-fetch): state only established billing for Toolbox runs

The README said every fetch is a billed Toolbox run, while the next
sentence says a run over max_cost_usd is refused and not charged, and a
code comment said a 404 comes back as an unbilled is_error run. The
AI/ML API Toolbox documentation says you pay for runs but not whether a
failed run is billed. Both now say only that Toolbox runs are billed.

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

---------

Co-authored-by: Stan <[email protected]>
Co-authored-by: Hugo <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
2026-11-04 19:20:17 -08:00
98736457bd feat(web-search): add AI/ML API Toolbox as an opt-in backend (#1189)
* feat(aimlapi): shared service key and attribution for aimlapi.com services

Adds the pieces the Toolbox and MCP integrations share:

- internal/aimlapi: X-AIMLAPI-Source (agent/zero), the fixed Toolbox and
  MCP URLs, and an exact-hostname check (https, default port, no
  userinfo) for api/tools/mcp.aimlapi.com before the key or attribution
  headers are attached. The partner id still comes from ResolvePartnerID.
- internal/config: AIMLAPIServiceKey picks the key from the aimlapi
  profile (or AIMLAPI_API_KEY without one), reading the credential
  store, and withholds it when any aimlapi-keyed profile points at a
  proxy base URL or AIMLAPI_INFERENCE_URL is overridden. Keys that are
  not a single token are refused.

* feat(web-search): AI/ML API Toolbox as an opt-in web_search backend

ZERO_WEBSEARCH_PROVIDER=aimlapi runs the Toolbox tool web/search at the
fixed https://tools.aimlapi.com/v1/tools URL. It is selected only by that
setting; an AI/ML API key alone never enables it or registers web_search.

- The key comes from config.LoadAIMLAPIServiceKey on every search; a
  missing, proxy-profile or invalid key fails before any request.
  ZERO_WEBSEARCH_BASE_URL and ZERO_WEBSEARCH_API_KEY are ignored.
- The key and attribution headers are attached only after
  aimlapi.IsAttributedURL accepts the URL, and redirects are refused.
- Every run carries max_cost_usd 0.025 and a fresh idempotency key;
  limit maps to max_results and domains to include_domains, with the
  local domain filter kept.
- A queued run (202) is polled until it ends, honouring
  next_poll_hint_ms clamped to 250 ms..5 s. This backend gets a 30 s
  budget with wait_ms 20 s; the generic backends keep 10 s.
- Non-2xx responses, is_error runs, failed/stopped/timed_out runs and
  malformed bodies are errors carrying the error code and request id,
  with the key scrubbed.

* fix(aimlapi): preserve store errors and isolate Windows config tests

* fix(web-search): redact results and preserve domain errors

* docs(web-search): clarify minimum search charge

* fix(aimlapi): reject Unicode whitespace in service keys

* fix(aimlapi): report store reads after proxy eligibility checks

* fix(web-search): keep plain no-results for unscoped backends

Only a backend that applied the domain allowlist itself (the AI/ML API
Toolbox) should turn an empty result into the "no results matched
domains" error. A generic backend never saw the allowlist, so an empty
result from it is an ordinary "No results for query", as before this
branch.

Also clear ZERO_WEBSEARCH_PROVIDER in the unconfigured registry test:
ZERO_WEBSEARCH_PROVIDER=aimlapi alone now registers web_search, so an
exported value flipped the test's oracle.

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

* fix(aimlapi): resolve the Toolbox key without an active provider

A user config that lists several providers but marks none active (or
names a removed one) makes Resolve return ErrNoActiveProvider together
with the normalized profiles. The app keeps running on that partial
config, but LoadAIMLAPIServiceKey stopped on the error, so web_search
failed even with an aimlapi profile or AIMLAPI_API_KEY set. Look the
key up in the partial config instead; the proxy check still scans every
resolved profile, and all other resolve errors are still returned.

Also state the key rules in the README as implemented: user config
only (project .zero providers are not used), which aimlapi profile is
used, and that any aimlapi profile or any profile reading
AIMLAPI_API_KEY on a non-canonical base URL withholds the key.

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

* fix(aimlapi): withhold the Toolbox key for dropped proxy profiles

Resolve drops a non-active profile it cannot normalize (an unknown or
unsupported catalog entry, a missing base URL, ...), so the proxy check
never saw it. A dropped profile that reads AIMLAPI_API_KEY from a proxy
still shows the key belongs to that proxy, yet the key was sent to
tools.aimlapi.com. LoadAIMLAPIServiceKey now also checks the configured
user profiles, merged the same way Resolve merges them, and withholds
the key when a dropped profile carries it and is not explicitly on the
canonical base URL.

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

* test(web-search): pin the generic backend route in BASE_URL registration tests

defaultSearchBackend now checks ZERO_WEBSEARCH_PROVIDER=aimlapi before
ZERO_WEBSEARCH_BASE_URL. TestWebSearchRegisteredInCoreNetworkTools and
TestCoreNetworkToolsExposeSafetyMetadata set only the base URL and inherit
the provider from the process environment, so with an exported
ZERO_WEBSEARCH_PROVIDER=aimlapi they took the Toolbox route and still passed
even if generic base-URL registration broke. Clear the provider so they
exercise the generic route they describe, matching
TestCoreNetworkToolsOmitWebSearchWhenUnconfigured.

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

* fix(redaction): redact caller-supplied secrets split by invisible characters

RedactString replaced Options.ExtraSecretValues only as contiguous text,
and the gap-tolerant pass from #1067 covers only the built-in credential
shapes. An opaque key such as the AI/ML API service key, echoed by a
service with a NUL or zero-width space inside it, matched neither pass.
Both fragments then reached Toolbox web_search errors (before clipping)
and successful results sent to the model and transcript.

Match each caller-supplied secret with runs of the same invisible
separators allowed between its characters, right after the contiguous
replacement. Separators inside the match are covered; those beside it
stay. Text without a separator takes the old path unchanged.

Regressions: split explicit values in the redaction package, and split
keys in the Toolbox HTTP error, run error and successful-result paths.
Raised by Copilot on #1189.

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

* fix(redaction): claim each caller secret whole or split before shorter ones

RedactString replaced every ExtraSecretValues entry contiguously, longest
first, and only then ran the split match. A shorter secret inside a
longer one that was split by an invisible character was replaced first,
so "abc\x00def" with secrets abcdef and abc became "[REDACTED]\x00def"
and the rest of the longer secret stayed visible.

Run the contiguous and split match together per value, longest first,
so a longer value is claimed in either form before any shorter one.
Regressions cover a shorter secret at the start, end and middle of a
split longer one, and alongside it. Raised by CodeRabbit on #1189.

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

* fix(redaction): merge overlapping caller-secret occurrences before replacing

ExtraSecretValues were replaced one occurrence at a time, so each
replacement could consume bytes another occurrence needed. With secret
abcabc, "abc\x00abcabc" became "abc\x00[REDACTED]". A key whose first and
last characters match lost all but one character the same way
("xY7q​[REDACTED]"), and two partly overlapping values left the
rest of one visible. Reordering the passes cannot fix this: the
contiguous and regexp replacements are both non-overlapping.

Find every occurrence of every value, whole or split by invisible
separators, in one separator-free copy of the text (linear KMP, all
overlapping starts), map the spans back, merge overlapping ones, and
replace each merged span once. Occurrences that only touch stay
separate replacements; separators beside a match stay; tab, newline
and carriage return are still not gaps. This honours the existing
TestRedactString_OverlappingExtraSecrets contract for both forms.
Raised by CodeRabbit on #1189.

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

* fix(web-search): keep the request URL out of Toolbox transport errors

When http.Client.Do failed, the backend scrubbed err.Error() from the
*url.Error, which quotes the request URL. A polling URL carries the run
id the Toolbox returned, percent-escaped, so a key echoed in that id
appeared as "Lp9%00Xr4" or "Lp9%E2%80%8BXr4". The key scrubber matches
the key whole or split by raw invisible separators, not percent-escapes,
so the key stayed recoverable in the web_search error.

Report only the cause inside the *url.Error. The message prefix, the
redirect refusal text and key scrubbing are unchanged. The regression
fails a polling GET after a queued run whose id carries a NUL- or
ZWSP-split key and checks the output has neither the polling URL nor
the key. Raised by CodeRabbit on #1189.

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

* fix(redaction): always redact a caller secret as written

Once the text held any invisible separator, redactExtraSecretValues
matched each caller secret only with separators removed from both. That
replaced the contiguous match instead of adding to it: a secret made
only of separators compacted to nothing and was skipped, and separators
at either end of a secret fell outside its match. RedactString("a\x00b")
with secret "\x00" returned the text unchanged, where main replaced it.

In the separator path, also record every occurrence of the secret as
written, then merge spans as before. Text without a separator takes the
unchanged path. Raised by CodeRabbit on #1189.

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

---------

Co-authored-by: Stan <[email protected]>
Co-authored-by: Hugo <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
2026-11-04 18:24:09 -08:00
ShryukGrandhi e59632406b test(agent): isolate user home, config and cache dirs for the agent tests (#1184)
Sandboxed tool runs in the agent tests keep runtime state and lease files
under os.UserCacheDir, so go test ./internal/agent/ left
zero/runtime/v1/*.lease files in the developer's real cache. TestMain
now points every per-user root at its temp dir via
testutil.IsolateUserDirs, as internal/tui does, and a guard test pins it.

Refs #1103

Signed-off-by: ShryukGrandhi <[email protected]>
2026-11-04 12:27:48 +08:00
github-actions[bot] 7c55a09c07 chore(main): release 0.12.0 (#1150)
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
2026-11-03 08:10:03 +08:00
Octopusandocto-patch d99f92b81a feat(providers): add MiniMax-M3.1-Flash to catalog (#1180)
* feat(providers): add MiniMax-M3.1-Flash to catalog

* fix(providers): forward MiniMax adaptive effort

---------

Co-authored-by: octo-patch <[email protected]>
2026-11-02 18:28:17 +08:00
zzucs 5aa54fe1d2 fix(agent): mark headless guardrail stops incomplete (#1182)
* fix(agent): mark headless guardrail stops incomplete

The empty-turn and repeated-tool-failure guards returned the run as a
successful final answer even with RequireCompletionSignal enabled, so
zero exec reported exit 0 and a success run_end for work it had just
abandoned.

Set Result.Incomplete and a reason at both stops when the headless gate
is on. Thresholds, stop answers, retry behavior, interactive runs, and
--no-completion-gate are unchanged.

Fixes #1154

* fix(agent): count consecutive empty turns in stop reason

The headless no-output stop reason used result.Turns, which counts every
turn in the run, so a run with an earlier tool-call turn claimed "stopped
after 4 turns with no output" for a 3-turn empty streak. Use the guard's
emptyTurns counter and say "consecutive".

Pinned by a case with a successful tool turn before the empty streak.
2026-11-02 14:25:40 +08:00
oniisan f912f80d00 fix(chatgpt): refresh subscription defaults and compatibility version (#1179)
* fix(chatgpt): refresh subscription defaults and compatibility version

Replace retired preset suggestions with vendor-documented Sol, Terra, Astra, and Luna IDs. Update the Codex discovery compatibility floor for Terra and verify that saved model selections are retained.

Refs Twigpine/zero#1177 (defaults portion)

* test(chatgpt): assert discovered model in compatibility probe

Verify the live response contains exactly gpt-future-coder in addition to checking the Codex compatibility query value.

Refs Twigpine/zero#1177

* docs(chatgpt): align Terra description with vendor guidance
2026-11-02 14:25:00 +08:00
ShryukGrandhi 9c4e3693e7 test(providerio): widen the idle window in the content-reset test (#1172)
TestScanSSEDataWithContextContentResetsOnData wrote a data line every 30ms against a 100ms idle timeout, so a ~70ms scheduling stall on a shared runner tripped the idle watchdog. Use a 500ms idle timeout (600ms content window) and 25 writes (~750ms), which keeps every gap far under the idle timeout while the run still outlasts the content window.

Signed-off-by: ShryukGrandhi <[email protected]>
2026-11-02 14:24:24 +08:00
zzucs 63962081cf fix(sessions): capture freeform apply_patch checkpoints (#1183)
* fix(sessions): capture freeform apply_patch checkpoints

Freeform apply_patch arguments are raw patch text, but both checkpoint
paths passed them to json.Unmarshal and dropped the error, so they found
no mutation targets and recorded no checkpoint. The tool still executed,
leaving rewind unable to restore the edit.

Resolve targets through a shared agent.CheckpointTargets helper that
prepares freeform input on the registered built-in apply_patch, exactly
as executeToolCall does. This keeps absolute patch headers
workspace-relative (rewriting them and setting cwd) instead of handing
absolute paths to the checkpoint layer. Non-freeform calls now use the
same tolerant decoder as execution, so the two cannot disagree.

Covers the CLI, spec-draft, and TUI capture paths.

Fixes #1162

* test(sessions): canonicalize freeform absolute-header fixtures

t.TempDir returns the symlinked /var spelling on macOS and short names on
Windows, while the apply_patch tool resolves the physical root. The
absolute-header cases must build the header from the resolved root, as the
repo's cross-platform test guidance requires.
2026-11-02 14:23:58 +08:00
Vasanth T 2e10a4b60b test(swarm): wait for every slot to be released before checking the drain (#1181)
A finishing member's task reads Done before its slot is freed: the
supervisor completes the task, removes the member and only then frees
the slot in afterExit. The test waited for Done == 5 and checked the
running count straight after, so a busy runner could see one slot still
held. It now waits for the drain itself.
2026-10-29 06:30:41 -07:00
Vasanth T 4247a07aba test(tools): stop the tty session before removing its root (#1175)
TestExecCommandTTYSkipsInteractiveBlock runs ssh to a bogus host under
tty with a 250ms yield. When ssh is still resolving the host after the
yield, the session keeps running in the test's directory and nothing
stops it, so on Windows t.TempDir's cleanup fails with a sharing
violation. It failed Smoke (windows-latest) on a README-only PR.

Stop the session in a cleanup, and build the tool on execTestRoot,
which retries the removal while Windows still holds the directory right
after the stop, as the other exec tests that leave a session do.
2026-10-27 06:25:41 -07:00
Vasanth T b837448dd4 docs: list Vercel under Partners (#1176)
Adds Vercel beside GMI Cloud in the Partners table of both READMEs,
with Vercel's official logotype in light and dark variants, taken
unchanged from its brand page apart from a trailing newline.
2026-10-27 15:29:59 +08:00
Vasanth T a0b036a3c7 docs: add Vercel OSS Program badge (#1174)
* docs: add Vercel OSS Program badge

* docs: add the Vercel OSS Program badge to README_ZH too

The two READMEs carry the same badge row, so the Chinese one gets the
badge as well. Both now list alt before src, like the other badges in
the row.
2026-10-27 12:30:50 +05:30
Vasanth T 164150633a feat(skills): install a skill's bundled files and point the skill tool at them (#1167)
* feat(skills): install a skill's bundled files and point the skill tool at them

A skill was only its SKILL.md: install copied that one file, and the skill
tool returned only the markdown body, so a skill whose instructions run its
own scripts was broken by installing it, and a hand-placed skill directory
gave the model no way to find the files beside it.

Install now copies the whole skill directory, leaving out hidden entries
(.git, .env) and symlinks, which it neither follows nor copies but reports.
A script stays executable; everything else is installed 0644. A tree over
1,000 files or 32 MiB is refused. skills.lock records a treeHash beside the
SKILL.md hash; skill info reports drift in bundled files, and an entry from
before this is still checked on SKILL.md alone. A reinstall that changed
only a bundled file is reported as an update.

Both skill tools end their output with the skill's directory and a list of
up to 50 bundled files when there are any; a SKILL.md-only skill is
returned unchanged.

Closes #584

* fix(skills): compare a reinstall over an old lock entry with SKILL.md alone

A lock entry written before tree hashes has no treeHash, so the reinstall
compared the new tree hash with an empty string and always counted as an
update. Re-adding an unchanged skill that is only a SKILL.md printed
"Updated" and "the files bundled with it changed" when there were none.

That entry comes from an install that copied SKILL.md alone, so compare
with that: it is an update when SKILL.md changed or the reinstall now
brings bundled files. InstallResult gains previousTreeHash, and the CLI
uses it to say the bundled files are now installed rather than changed
on the first reinstall over such an entry.

* fix(skills): match the tree hash when recovering an interrupted install

lockReconciler matched both trees of an interrupted workspace on the
SKILL.md hash alone. A reinstall that changes only a bundled file keeps
SKILL.md, so a kill between the tree swap and the lockfile publish left
two trees carrying the recorded hash, and recovery kept the new one and
retired the backup the lockfile still described.

A tree now matches an entry with a tree hash only when both hashes
match. An entry from before tree hashes is still matched on SKILL.md
alone. A tree hash that cannot be computed is returned as an error, the
way an unreadable SKILL.md already is, so recovery decides nothing about
a tree it could not read.

Also bring the package comment up to date: install copies the skill
directory, not SKILL.md alone.
2026-10-25 07:24:59 -07:00
110 changed files with 16472 additions and 341 deletions

No files matched your search

+1 -1
View File
@@ -1 +1 @@
{".":"0.11.0"}
{".":"0.12.0"}
+24
View File
@@ -6,6 +6,30 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
aims to follow [Semantic Versioning](https://semver.org/spec/v2.0.0.html) once the first release is
tagged. Until then, source builds report the version `dev`.
## [0.12.0](https://github.com/Twigpine/zero/compare/v0.11.0...v0.12.0) (2026-11-02)
### Features
* **providers:** add MiniMax-M3.1-Flash to catalog ([#1180](https://github.com/Twigpine/zero/issues/1180)) ([d99f92b](https://github.com/Twigpine/zero/commit/d99f92b81a671f186cede81614c89b8e8a7f9451))
* **sandbox:** support self-exec pattern to eliminate separate helper binary ([#946](https://github.com/Twigpine/zero/issues/946)) ([8c4e759](https://github.com/Twigpine/zero/commit/8c4e7597989ef9b3ca67b7703dfd15d92372aed2))
* **skills:** install a skill's bundled files and point the skill tool at them ([#1167](https://github.com/Twigpine/zero/issues/1167)) ([1641506](https://github.com/Twigpine/zero/commit/164150633a8b2e6ae864ae02fd7105b201eeb68e))
### Bug Fixes
* **agent:** mark headless guardrail stops incomplete ([#1182](https://github.com/Twigpine/zero/issues/1182)) ([5aa54fe](https://github.com/Twigpine/zero/commit/5aa54fe1d2b419c9bbef73605a77919c9007cf7b))
* bump Go to 1.26.10 for GO-2026-6809 ([#1168](https://github.com/Twigpine/zero/issues/1168)) ([6393a90](https://github.com/Twigpine/zero/commit/6393a90ed618a3536e09c02852f4ec829b97545b))
* **chatgpt:** refresh subscription defaults and compatibility version ([#1179](https://github.com/Twigpine/zero/issues/1179)) ([f912f80](https://github.com/Twigpine/zero/commit/f912f80d002602e11ba0ef8548ecee58442874aa))
* **deps:** refresh the lockfile to hono 4.13.13 ([#1171](https://github.com/Twigpine/zero/issues/1171)) ([de1443b](https://github.com/Twigpine/zero/commit/de1443b2edb44fd3747cc77200d54dc8cffcd0a7))
* **oauth:** reuse the crash-safe secret store and sync encrypted publications ([#1165](https://github.com/Twigpine/zero/issues/1165)) ([84c1b15](https://github.com/Twigpine/zero/commit/84c1b1554076679a342e9cac61f8f8cdbaf3a429))
* **redaction:** redact credentials whose trigger a greedy span swallowed ([#1158](https://github.com/Twigpine/zero/issues/1158)) ([a97db88](https://github.com/Twigpine/zero/commit/a97db8832fdba077f2d31bdf03ecef8d7c774fe1))
* **sessions:** capture freeform apply_patch checkpoints ([#1183](https://github.com/Twigpine/zero/issues/1183)) ([6396208](https://github.com/Twigpine/zero/commit/63962081cf130c2171407936ce3beb7021d238b8))
* **sessions:** preserve image attachment summaries on resume ([#1152](https://github.com/Twigpine/zero/issues/1152)) ([1e84849](https://github.com/Twigpine/zero/commit/1e848492e5c1919d187fb5e1eb5e7ef9bc6140dc)), closes [#1135](https://github.com/Twigpine/zero/issues/1135)
* **tools:** grep stops where a file's line breaks stop ([#1160](https://github.com/Twigpine/zero/issues/1160)) ([aa9ba84](https://github.com/Twigpine/zero/commit/aa9ba8412eb8917daa2ca5cb3fdbd0a39276fa88))
* **workspaceindex:** import gitsafe under the Twigpine module path ([#1153](https://github.com/Twigpine/zero/issues/1153)) ([ef4da13](https://github.com/Twigpine/zero/commit/ef4da1323f1097863ae278224b5f304c2fff909a))
* **workspaceindex:** run Scan's git calls through gitsafe ([#1149](https://github.com/Twigpine/zero/issues/1149)) ([2ffe183](https://github.com/Twigpine/zero/commit/2ffe1831b0950b1463ec46d487d03a5fcb679cc6))
## [0.11.0](https://github.com/Twigpine/zero/compare/v0.10.0...v0.11.0) (2026-10-15)
+80
View File
@@ -15,6 +15,7 @@
<img alt="Go 1.26.10+" src="https://img.shields.io/badge/Go-1.26.10+-00ADD8?logo=go&logoColor=white">
<img alt="25+ providers" src="https://img.shields.io/badge/providers-25+-34E2EA">
<a href="https://discord.gg/CaQDS6wdFn"><img alt="Discord" src="https://img.shields.io/badge/Discord-join-5865F2?logo=discord&logoColor=white"></a>
<a href="https://vercel.com/open-source-program"><img alt="Vercel OSS Program" src="https://vercel.com/oss/program-badge-2026.svg"></a>
<br>
<strong>English</strong> | <a href="README_ZH.md">中文</a>
</p>
@@ -87,9 +88,18 @@ zero exec --output-format stream-json < turns.jsonl
</picture>
</a>
</td>
<td align="center" width="150" height="80">
<a href="https://vercel.com/">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="docs/assets/vercel-logo-dark.svg">
<img src="docs/assets/vercel-logo.svg" alt="Vercel logo" width="120">
</picture>
</a>
</td>
</tr>
<tr>
<td align="center"><a href="https://www.gmicloud.ai/"><strong>GMI Cloud</strong></a></td>
<td align="center"><a href="https://vercel.com/"><strong>Vercel</strong></a></td>
</tr>
</table>
@@ -387,6 +397,76 @@ and can rewrite every request through it, so this is the trust you accept by
configuring one. Leave the variables unset for the checks to be enforced end to
end.
The built-in `web_search` tool is registered only when a search backend is
configured: `ZERO_WEBSEARCH_BASE_URL` (with an optional `ZERO_WEBSEARCH_API_KEY`)
for a generic JSON endpoint, `ZERO_WEBSEARCH_PROVIDER=searxng` with a SearXNG
base URL, or the AI/ML API Toolbox:
```bash
export ZERO_WEBSEARCH_PROVIDER=aimlapi
```
The AI/ML API backend runs the Toolbox tool `web/search`, and every search is
billed to your AI/ML API balance: about $0.0013 per result, with a documented
minimum of $0.0091 per search. Each search is sent with a $0.025 spending cap,
which the Toolbox enforces. It is used only when
`ZERO_WEBSEARCH_PROVIDER` is set to `aimlapi`; an AI/ML API key alone never
enables it. The key is read from your user configuration only; provider
profiles in a project `.zero` config are not used. It is the aimlapi profile's
key (the active profile when it is aimlapi, otherwise the first aimlapi
profile), including a key in the credential store, or `AIMLAPI_API_KEY` when no
aimlapi profile exists. It is resolved again on every search;
`ZERO_WEBSEARCH_API_KEY` is not used. When any aimlapi profile, or any profile
that reads `AIMLAPI_API_KEY`, points at a base URL other than
`https://api.aimlapi.com/v1`, or `AIMLAPI_INFERENCE_URL` is overridden, the key
may belong to a proxy, so it is not sent and the search fails without a request.
Requests go only to
`https://tools.aimlapi.com` (`ZERO_WEBSEARCH_BASE_URL` is ignored and redirects
are refused) and carry `X-AIMLAPI-Source: agent/zero` and an
`X-AIMLAPI-Partner-ID` header. The `domains` argument is sent to the Toolbox as
`include_domains`, and results are still filtered locally.
### AI/ML API Toolbox backend for `web_fetch`
`web_fetch` can fetch pages through the AI/ML API Toolbox tool
`search1api/crawl` instead of a direct HTTP GET. It is off unless you choose it:
```bash
export ZERO_WEBFETCH_PROVIDER=aimlapi
```
Having an AI/ML API key never switches it on by itself, and any other value (or
none) keeps the direct fetch. The variable is read on every fetch.
- **Billing:** Toolbox runs are billed to your AI/ML API balance, about
$0.0013 per run. Each run is sent with `max_cost_usd` 0.005, so a run
priced above that is refused before it starts and not charged. Rate-limited
or gateway-failed requests are retried with the same idempotency key, so a
retry is never billed twice.
- **Key:** the key comes from your `aimlapi` provider profile (including a key
in the credential store), or from `AIMLAPI_API_KEY` when there is no such
profile. It is withheld, and `web_fetch` returns an error, when any `aimlapi`
profile or any profile using `AIMLAPI_API_KEY` points at a base URL other than
`https://api.aimlapi.com/v1`, or when `AIMLAPI_INFERENCE_URL` is overridden:
the key may then belong to a proxy.
- **Destination:** requests go only to the fixed `https://tools.aimlapi.com`
endpoint, with `Authorization: Bearer <key>` and the attribution headers
`X-AIMLAPI-Source: agent/zero` and `X-AIMLAPI-Partner-ID`. Redirects are
refused. `HTTP_PROXY`/`HTTPS_PROXY` apply to these requests as to Zero's other
aimlapi.com calls.
- **No fallback:** without a usable key, or when the run fails, `web_fetch`
returns an error and does not fetch the page directly. Unset the variable to
go back to direct fetches.
- **Remote fetch:** the page is fetched by a remote service, which may serve a
cached copy, and the full URL (query string included) is sent to it. The URL
checks above still run first, so loopback, private and metadata URLs are
never handed to it. The result is Markdown with a
`Backend: aimlapi (search1api/crawl)` line; `max_bytes` cuts the Markdown,
and `format: "raw"` is refused because the original response is not
available. The complete Toolbox run response is limited to 8 MiB before JSON
decoding; larger responses return an error even if `max_bytes` is smaller.
The remote run may already have been billed when this limit is detected.
For local dev servers, use shell commands such as `curl` through `exec_command`
so the normal sandbox and permission policy applies. Long-running commands stay
attached to a background terminal session and can be listed or stopped from the
+10
View File
@@ -15,6 +15,7 @@
<img alt="Go 1.26.10+" src="https://img.shields.io/badge/Go-1.26.10+-00ADD8?logo=go&logoColor=white">
<img alt="25+ providers" src="https://img.shields.io/badge/providers-25+-34E2EA">
<a href="https://discord.gg/CaQDS6wdFn"><img alt="Discord" src="https://img.shields.io/badge/Discord-join-5865F2?logo=discord&logoColor=white"></a>
<a href="https://vercel.com/open-source-program"><img alt="Vercel OSS Program" src="https://vercel.com/oss/program-badge-2026.svg"></a>
<br>
<a href="README.md">English</a> | <strong>中文</strong>
</p>
@@ -85,9 +86,18 @@ zero exec --output-format stream-json < turns.jsonl
</picture>
</a>
</td>
<td align="center" width="150" height="80">
<a href="https://vercel.com/">
<picture>
<source media="(prefers-color-scheme: dark)" srcset="docs/assets/vercel-logo-dark.svg">
<img src="docs/assets/vercel-logo.svg" alt="Vercel logo" width="120">
</picture>
</a>
</td>
</tr>
<tr>
<td align="center"><a href="https://www.gmicloud.ai/"><strong>GMI Cloud</strong></a></td>
<td align="center"><a href="https://vercel.com/"><strong>Vercel</strong></a></td>
</tr>
</table>
+13 -1
View File
@@ -33,4 +33,16 @@ Please include, where possible:
depend on a user explicitly disabling the sandbox (`--skip-permissions-unsafe`), are generally
out of scope — but please report anything you are unsure about.
- Zero sends no telemetry. Network calls go to the model/tool providers you configure (and, for
keyless web search, to Exa's hosted MCP endpoint at `https://mcp.exa.ai/mcp`).
keyless web search, to Exa's hosted MCP endpoint at `https://mcp.exa.ai/mcp`). The built-in
AI/ML API MCP server (`https://mcp.aimlapi.com/mcp`) is off until you run
`zero mcp enable aimlapi`. Once enabled, the Bearer key follows the aimlapi
profile resolution order (inline key, then the profile's `apiKeyEnv` — default
`AIMLAPI_API_KEY` — then the credential store), or `AIMLAPI_API_KEY` alone when
no aimlapi profile exists; a non-empty `AIMLAPI_API_KEY` can override a stored
profile key for MCP billing. The automatic key is withheld if an aimlapi profile
(or a profile reading `AIMLAPI_API_KEY`) uses a base URL other than
`https://api.aimlapi.com/v1`, `AIMLAPI_INFERENCE_URL` points elsewhere, a
non-empty user Authorization header is configured, OAuth is selected, project
config has touched the server, or the key is missing or not a single token
(including Unicode whitespace). See
[docs/EXTENDING.md](docs/EXTENDING.md#built-in-servers).
+65 -2
View File
@@ -110,6 +110,8 @@ Discovery roots (earlier wins on name collisions):
A missing directory is fine — Zero just omits it. Management commands (`zero skill add` / `zero skill remove`) always write to the primary Zero dir only; `list` / `info` search primary + `~/.agents/skills`. `skill info` reports lock source/hash when present and whether the on-disk `SKILL.md` has drifted from the recorded hash.
For a skill installed since Zero started installing bundled files, `skill info` also reports drift in any file installed with it.
```text
~/.local/share/zero/skills/ # primary (install/remove target)
run-benchmarks/
@@ -140,6 +142,15 @@ Only `name` and `description` are recognized in the frontmatter today. The `name
The `skill` tool lets the agent load any discovered skill by name (primary, agents, or plugin).
A skill can ship files beside its `SKILL.md`, in the [agentskills.io](https://agentskills.io/specification) layout: `scripts/`, `references/`, `assets/` or anything else the instructions refer to. `zero skill add` installs the whole skill directory, leaving out hidden entries such as `.git` and not following or copying symlinks (it lists any it skipped). A skill larger than 1,000 files or 32 MiB is refused. When the `skill` tool loads a skill that has such files, its output ends with the skill's directory and a list of the files, so instructions like "run `scripts/fill.py`" can be followed. Running a bundled script goes through the normal tool permissions; installing or loading a skill never runs anything.
```text
pdf/
SKILL.md
scripts/fill.py
references/form-fields.md
```
## 4. Hooks
Hooks fire shell commands on lifecycle events. Configure them in JSON:
@@ -192,7 +203,11 @@ Events the agent emits (in dispatch order):
| `specialistStart` | A sub-agent is spawned | yes (specialist name) |
| `specialistStop` | A sub-agent ends | yes (specialist name) |
A hook's exit code decides what happens next: `0` continues, non-zero blocks the tool call (`beforeTool`) or surfaces an error (`afterTool`). Hook execution is recorded in the audit log; the audit is reachable from the agent's view of past actions, not from a dedicated `zero doctor` check.
For swarm members the matcher subject is the **agent type** (`subagent`, `teammate`, …), not a registered specialist profile name — swarm members run from an inline manifest keyed by agent type.
`specialistStart` fires once the child process has actually started; a launch that fails before exec (missing binary, spawn error) records no public lifecycle hooks. `specialistStop` fires once per terminal stop: completion, error, explicit stop of a recovered task, and a background task reaped by a later `TaskOutput` all share one durable claim, so a hook that notifies or cleans up does not run twice for the same sub-agent run.
A hook's exit code decides what happens next: `0` continues, non-zero blocks the tool call (`beforeTool`) or surfaces an error (`afterTool`). `sessionStart` / `sessionEnd` / `specialistStart` / `specialistStop` hooks are advisory (failures are audited but do not interrupt the run). Hook execution is recorded in the audit log; the audit is reachable from the agent's view of past actions, not from a dedicated `zero doctor` check.
> **Roadmap.** An in-UI hooks manager is on the backlog. Today you edit the JSON directly.
@@ -246,6 +261,52 @@ or use a new server name. OAuth tokens are bound to the resolved server identity
so changing a server target can require `zero mcp oauth login <server>` again.
OAuth server names ending in `.<32 hex chars>` are reserved for token storage.
#### Built-in servers
Zero seeds two MCP servers before your config is merged. Both show up in
`zero mcp list`, and any field can be overridden in your user config.
- `exa` (`https://mcp.exa.ai/mcp`) is on by default for keyless web search.
`zero mcp disable exa` turns it off.
- `aimlapi` (`https://mcp.aimlapi.com/mcp`, the AI/ML API Toolbox, model
discovery, image/video/audio generation, chat, embeddings and balance tools)
is **off by default**, because its tool runs and generations are billed to
your AI/ML API balance. Turn it on with `zero mcp enable aimlapi`, which
writes only `"disabled": false` to your user config, and off again with
`zero mcp disable aimlapi`. A project config cannot turn it on.
When `aimlapi` connects, Zero sends your AI/ML API key as
`Authorization: Bearer <key>`. The key is resolved again on every request and
is never written to `config.json`. With an aimlapi.com provider profile
(`zero providers setup aimlapi`), the key is chosen in this order:
1. An inline key on the profile (if set in config).
2. Otherwise the environment variable named by the profile's `apiKeyEnv`
(the catalog default is `AIMLAPI_API_KEY`).
3. Otherwise the key held in the credential store when the profile uses
`apiKeyStored`.
When there is no aimlapi profile, Zero uses `AIMLAPI_API_KEY` alone. If both
`AIMLAPI_API_KEY` and a stored profile key exist, step 2 wins over the store
whenever the inline key is empty — MCP requests may bill the environment key's
account, not the stored key.
No automatic key is sent when an aimlapi profile, or any profile reading
`AIMLAPI_API_KEY`, uses a base URL other than `https://api.aimlapi.com/v1`, or
when `AIMLAPI_INFERENCE_URL` is overridden, because that credential may belong
to a proxy. No key is sent when the resolved value is not a single token (for
example a comma- or newline-separated pool), when you set your own
non-empty `Authorization` header, when the server uses `"auth": "oauth"`, or
when project config touched the server. Without a key the server answers 401.
To sign in with your aimlapi.com account instead, set `"auth": "oauth"` on the
server and run `zero mcp oauth login aimlapi`.
Every request to exactly `https://mcp.aimlapi.com` (https, default port,
whatever the server is named, including one you add by hand) also carries the
attribution headers `X-AIMLAPI-Source: agent/zero` and `X-AIMLAPI-Partner-ID`,
and gets the key under the rules above. Other MCP servers never receive these
headers or the key.
### As a server — expose Zero's tools to another agent
```bash
@@ -290,10 +351,12 @@ Install and manage:
zero plugins add ./github-pr-review # copy into ~/.config/zero/plugins/ or ./.zero/plugins/
zero plugins list
zero plugins info github-pr-review # show manifest path, enabled state, lock metadata
zero plugins disable github-pr-review # sets "enabled": false in plugin.json
zero plugins enable github-pr-review # sets "enabled": true in plugin.json
zero plugins remove github-pr-review # alias: rm
```
A plugin is enabled by being present in the plugins directory and disabled by removing it (or by the user setting `"enabled": false` in its `plugin.json`). Plugins are not enabled or disabled by a CLI subcommand today.
A plugin is active when it is present under a plugins directory and its `plugin.json` has `"enabled": true` (the default when the field is omitted). Use `zero plugins disable <id>` / `zero plugins enable <id>` to toggle that field without removing the plugin. Pass `--user` to target only user plugins and ignore `./.zero/plugins`; without it, a project plugin with the same id takes precedence. The toggle edits the `enabled` value in place — the rest of the manifest keeps its formatting, field order, and numeric literals — and serializes with `plugins add` / `plugins remove` on the same plugins root. Zero reads plugins when it starts, so the change applies to the next `zero exec` run or interactive session; an interactive session that is already running keeps the plugin tools, hooks, and skills it started with until you restart it.
Plugin commands run with the plugin directory as their working directory. Use relative paths; the loader resolves them at activation time.
+4
View File
@@ -634,6 +634,10 @@ Core tool groups include:
`HTTP_PROXY`/`HTTPS_PROXY` set, the proxy is dialed instead and resolves the
target itself, so destination enforcement ends at the proxy; see the note in
the README.
With `ZERO_WEBFETCH_PROVIDER=aimlapi` (opt-in; a key alone does not select
it), the same URL checks run and the page is then fetched as Markdown by the
billed AI/ML API Toolbox tool `search1api/crawl`, with no fallback to a
direct fetch; see the README.
- **Local-control tools**: config-gated wrappers for browser, terminal, desktop,
and artifact helpers backed by packages such as `internal/localcontrol` and
`internal/browser`.
+1
View File
@@ -0,0 +1 @@
<svg xmlns="http://www.w3.org/2000/svg" width="261" height="52" fill="none" viewBox="0 0 261 52"><path fill="white" d="M59.8 52H0L29.9 0zm67.82-38.45q4.9 0 8.81 2.13a15.5 15.5 0 0 1 6.22 6.32q2.3 4.2 2.38 10.26v2.06h-26.35q.27 4.4 2.58 6.92 2.38 2.47 6.36 2.47a8.4 8.4 0 0 0 7.76-4.93l9.16.67q-1.68 4.98-6.29 7.99t-10.63 3q-5.53 0-9.64-2.27a16 16 0 0 1-6.43-6.46 20 20 0 0 1-2.31-9.72q0-5.52 2.3-9.72a16 16 0 0 1 6.44-6.46 20 20 0 0 1 9.64-2.26m62.55 0q4.47 0 8.18 1.66a15.3 15.3 0 0 1 6.15 4.6q2.38 3 2.87 7.05l-9.23.47a8 8 0 0 0-2.8-5 7.7 7.7 0 0 0-5.17-1.86q-4.33 0-6.7 3-2.4 3-2.39 8.52t2.38 8.52q2.37 3 6.71 3 3.15 0 5.39-1.87 2.24-1.92 2.72-5.46l9.3.4a14.7 14.7 0 0 1-2.87 7.33 16 16 0 0 1-6.15 4.86 21 21 0 0 1-8.39 1.66q-5.53 0-9.64-2.26a16 16 0 0 1-6.44-6.46 20 20 0 0 1-2.3-9.72q0-5.52 2.3-9.72a16 16 0 0 1 6.44-6.46 20 20 0 0 1 9.64-2.26m38.66 0q4.9 0 8.8 2.13a15.5 15.5 0 0 1 6.22 6.32q2.31 4.2 2.38 10.26v2.06h-26.35q.28 4.4 2.58 6.92 2.38 2.47 6.36 2.47a8.4 8.4 0 0 0 7.77-4.93l9.15.67q-1.68 5-6.29 7.99t-10.62 3q-5.53 0-9.65-2.27a16 16 0 0 1-6.43-6.46 20 20 0 0 1-2.31-9.72q0-5.52 2.3-9.72a16 16 0 0 1 6.44-6.46 20 20 0 0 1 9.64-2.26M86.9 36.69l17.24-34.33h10.8L89.96 49.63h-6.12L58.85 2.36h10.81zm71.62-15.55a11 11 0 0 1 2.47-4.48q2.28-2.31 6.38-2.31h3.4v7.26h-3.47q-2.91 0-4.79.8a5.8 5.8 0 0 0-2.77 2.5q-.9 1.72-.9 4.37v20.35h-8.89V14.35h8.33zm101.73 28.5h-8.95V2.35h8.95zM127.62 20.26q-3.7 0-6 2.2-2.32 2.2-2.87 6.2h17.05q-.48-4.34-2.72-6.33a7.8 7.8 0 0 0-5.46-2.07m101.2 0q-3.7 0-6 2.2-2.31 2.2-2.87 6.2H237q-.5-4.34-2.73-6.33a7.8 7.8 0 0 0-5.46-2.07" style="fill:white;fill-opacity:1;"/></svg>

After

Width:  |  Height:  |  Size: 1.6 KiB

+1
View File
@@ -0,0 +1 @@
<svg xmlns="http://www.w3.org/2000/svg" width="261" height="52" fill="none" viewBox="0 0 261 52"><path fill="black" d="M59.8 52H0L29.9 0zm67.82-38.45q4.9 0 8.81 2.13a15.5 15.5 0 0 1 6.22 6.32q2.3 4.2 2.38 10.26v2.06h-26.35q.27 4.4 2.58 6.92 2.38 2.47 6.36 2.47a8.4 8.4 0 0 0 7.76-4.93l9.16.67q-1.68 4.98-6.29 7.99t-10.63 3q-5.53 0-9.64-2.27a16 16 0 0 1-6.43-6.46 20 20 0 0 1-2.31-9.72q0-5.52 2.3-9.72a16 16 0 0 1 6.44-6.46 20 20 0 0 1 9.64-2.26m62.55 0q4.47 0 8.18 1.66a15.3 15.3 0 0 1 6.15 4.6q2.38 3 2.87 7.05l-9.23.47a8 8 0 0 0-2.8-5 7.7 7.7 0 0 0-5.17-1.86q-4.33 0-6.7 3-2.4 3-2.39 8.52t2.38 8.52q2.37 3 6.71 3 3.15 0 5.39-1.87 2.24-1.92 2.72-5.46l9.3.4a14.7 14.7 0 0 1-2.87 7.33 16 16 0 0 1-6.15 4.86 21 21 0 0 1-8.39 1.66q-5.53 0-9.64-2.26a16 16 0 0 1-6.44-6.46 20 20 0 0 1-2.3-9.72q0-5.52 2.3-9.72a16 16 0 0 1 6.44-6.46 20 20 0 0 1 9.64-2.26m38.66 0q4.9 0 8.8 2.13a15.5 15.5 0 0 1 6.22 6.32q2.31 4.2 2.38 10.26v2.06h-26.35q.28 4.4 2.58 6.92 2.38 2.47 6.36 2.47a8.4 8.4 0 0 0 7.77-4.93l9.15.67q-1.68 5-6.29 7.99t-10.62 3q-5.53 0-9.65-2.27a16 16 0 0 1-6.43-6.46 20 20 0 0 1-2.31-9.72q0-5.52 2.3-9.72a16 16 0 0 1 6.44-6.46 20 20 0 0 1 9.64-2.26M86.9 36.69l17.24-34.33h10.8L89.96 49.63h-6.12L58.85 2.36h10.81zm71.62-15.55a11 11 0 0 1 2.47-4.48q2.28-2.31 6.38-2.31h3.4v7.26h-3.47q-2.91 0-4.79.8a5.8 5.8 0 0 0-2.77 2.5q-.9 1.72-.9 4.37v20.35h-8.89V14.35h8.33zm101.73 28.5h-8.95V2.35h8.95zM127.62 20.26q-3.7 0-6 2.2-2.32 2.2-2.87 6.2h17.05q-.48-4.34-2.72-6.33a7.8 7.8 0 0 0-5.46-2.07m101.2 0q-3.7 0-6 2.2-2.31 2.2-2.87 6.2H237q-.5-4.34-2.73-6.33a7.8 7.8 0 0 0-5.46-2.07" style="fill:black;fill-opacity:1;"/></svg>

After

Width:  |  Height:  |  Size: 1.6 KiB

+34
View File
@@ -0,0 +1,34 @@
package agent
import "github.com/Twigpine/zero/internal/tools"
// CheckpointTargets returns the workspace-relative paths a tool call will mutate,
// resolving its arguments exactly as executeToolCall does. Freeform apply_patch
// input is normalized through the registered built-in tool, so absolute headers
// are rewritten relative to the selected write root and cwd is set: checkpoint
// capture and execution therefore agree on the files at stake.
//
// Best-effort: nil when the call has no capture targets or its arguments cannot
// be resolved. Checkpoint capture never fails a run.
func CheckpointTargets(workspaceRoot string, registry *tools.Registry, call ToolCall) []string {
if registry == nil {
return nil
}
args := map[string]any{}
if call.Freeform {
tool, ok := registry.Get(call.Name)
if !ok || !tools.IsBuiltInApplyPatch(tool) {
return nil
}
prepared, err := tools.PrepareFreeformApplyPatchArguments(tool, call.Arguments)
if err != nil {
return nil
}
args = prepared
} else if call.Arguments != "" {
if err := decodeToolArguments(call.Arguments, &args); err != nil {
return nil
}
}
return tools.MutationTargets(workspaceRoot, call.Name, args)
}
+99
View File
@@ -0,0 +1,99 @@
package agent
import (
"encoding/json"
"os"
"path/filepath"
"slices"
"testing"
"github.com/Twigpine/zero/internal/tools"
)
func TestCheckpointTargets(t *testing.T) {
root := t.TempDir()
// Absolute headers must be compared against the physical root: t.TempDir
// differs from it on macOS (/var -> /private/var) and via Windows short names.
if resolved, err := filepath.EvalSymlinks(root); err == nil {
root = resolved
}
if err := os.WriteFile(filepath.Join(root, "main.txt"), []byte("before\n"), 0o644); err != nil {
t.Fatal(err)
}
relativePatch := "*** Begin Patch\n*** Update File: main.txt\n@@\n-before\n+after\n*** End Patch\n"
absolutePatch := "*** Begin Patch\n*** Update File: " + filepath.Join(root, "main.txt") + "\n@@\n-before\n+after\n*** End Patch\n"
multiFilePatch := "*** Begin Patch\n*** Update File: main.txt\n@@\n-before\n+after\n*** Add File: notes.txt\n+new\n*** End Patch\n"
jsonArgs, err := json.Marshal(map[string]any{"patch": relativePatch})
if err != nil {
t.Fatal(err)
}
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(root, nil))
unbuiltinRegistry := tools.NewRegistry()
unbuiltinRegistry.Register(&customApplyPatchTool{})
for _, tc := range []struct {
name string
registry *tools.Registry
call ToolCall
want []string
}{
{
name: "freeform relative header",
call: ToolCall{Name: "apply_patch", Arguments: relativePatch, Freeform: true},
want: []string{"main.txt"},
},
{
name: "freeform absolute header stays workspace-relative",
call: ToolCall{Name: "apply_patch", Arguments: absolutePatch, Freeform: true},
want: []string{"main.txt"},
},
{
name: "freeform multi-file patch",
call: ToolCall{Name: "apply_patch", Arguments: multiFilePatch, Freeform: true},
want: []string{"main.txt", "notes.txt"},
},
{
name: "json arguments",
call: ToolCall{Name: "apply_patch", Arguments: string(jsonArgs)},
want: []string{"main.txt"},
},
{
name: "freeform traversal is rejected",
call: ToolCall{Name: "apply_patch", Arguments: "*** Begin Patch\n*** Update File: ../escape.txt\n@@\n-x\n+y\n*** End Patch\n", Freeform: true},
want: nil,
},
{
name: "freeform custom apply_patch is not prepared",
registry: unbuiltinRegistry,
call: ToolCall{Name: "apply_patch", Arguments: relativePatch, Freeform: true},
want: nil,
},
{
name: "unknown freeform tool",
call: ToolCall{Name: "read_file", Arguments: "raw", Freeform: true},
want: nil,
},
{
name: "malformed json arguments",
call: ToolCall{Name: "apply_patch", Arguments: "{"},
want: nil,
},
} {
t.Run(tc.name, func(t *testing.T) {
reg := registry
if tc.registry != nil {
reg = tc.registry
}
got := CheckpointTargets(root, reg, tc.call)
if !slices.Equal(got, tc.want) {
t.Fatalf("CheckpointTargets() = %#v, want %#v", got, tc.want)
}
})
}
if got := CheckpointTargets(root, nil, ToolCall{Name: "apply_patch", Arguments: relativePatch}); got != nil {
t.Fatalf("CheckpointTargets(nil registry) = %#v, want nil", got)
}
}
+64
View File
@@ -80,6 +80,9 @@ func TestRunStopsAfterConsecutiveEmptyTurns(t *testing.T) {
if result.FinalAnswer == maxTurnsAnswer {
t.Fatalf("no-output guard must stop before reaching maxTurns, got max-turns answer")
}
if result.Incomplete {
t.Fatalf("interactive run must not be marked Incomplete; reason=%q", result.IncompleteReason)
}
}
func TestRunResetsEmptyTurnCounterOnVisibleOutput(t *testing.T) {
@@ -513,6 +516,22 @@ func repeatedFlakyTurns(n int) [][]zeroruntime.StreamEvent {
return turns
}
// alwaysSucceedingTool is the successful counterpart to alwaysFailingTool, used
// to put a productive tool turn before an empty-turn streak.
type alwaysSucceedingTool struct{}
func (alwaysSucceedingTool) Name() string { return "ok" }
func (alwaysSucceedingTool) Description() string { return "always succeeds for testing" }
func (alwaysSucceedingTool) Parameters() tools.Schema {
return tools.Schema{Type: "object", AdditionalProperties: false}
}
func (alwaysSucceedingTool) Safety() tools.Safety {
return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow}
}
func (alwaysSucceedingTool) Run(context.Context, map[string]any) tools.Result {
return tools.Result{Status: tools.StatusOK, Output: "ok"}
}
func TestRunStopsAfterRepeatedToolFailures(t *testing.T) {
registry := tools.NewRegistry()
registry.Register(alwaysFailingTool{})
@@ -529,6 +548,51 @@ func TestRunStopsAfterRepeatedToolFailures(t *testing.T) {
if len(provider.requests) != toolFailureStopAt {
t.Fatalf("expected stop at %d failures, made %d requests", toolFailureStopAt, len(provider.requests))
}
if result.Incomplete {
t.Fatalf("interactive run must not be marked Incomplete; reason=%q", result.IncompleteReason)
}
}
// Under the headless completion gate both guardrail stops are unfinished runs,
// so they must set Result.Incomplete and a reason; exec maps that to exit 4.
// Thresholds and stop answers stay pinned by the two tests above.
func TestHeadlessGuardStopsAreIncomplete(t *testing.T) {
for _, tc := range []struct {
name string
turns [][]zeroruntime.StreamEvent
wantTurns int
wantReason string
}{
{"empty_turns", [][]zeroruntime.StreamEvent{emptyTurn(), emptyTurn(), emptyTurn()}, maxEmptyTurns, "3 consecutive turns with no output"},
{"empty_turns_after_progress", [][]zeroruntime.StreamEvent{
toolTurn("ok-1", "ok", "{}"),
emptyTurn(), emptyTurn(), emptyTurn(),
}, maxEmptyTurns + 1, "3 consecutive turns with no output"},
{"repeated_tool_failures", repeatedFlakyTurns(10), toolFailureStopAt, "flaky"},
} {
t.Run(tc.name, func(t *testing.T) {
registry := tools.NewRegistry()
registry.Register(alwaysFailingTool{})
registry.Register(alwaysSucceedingTool{})
result, err := Run(context.Background(), "go", &mockProvider{turns: tc.turns}, Options{
Registry: registry,
MaxTurns: 12,
RequireCompletionSignal: true,
})
if err != nil {
t.Fatal(err)
}
if result.Turns != tc.wantTurns {
t.Fatalf("guard did not trigger: turns=%d want=%d", result.Turns, tc.wantTurns)
}
if !result.Incomplete {
t.Fatalf("guardrail stop under the headless gate must be Incomplete; final=%q", result.FinalAnswer)
}
if !strings.Contains(result.IncompleteReason, tc.wantReason) {
t.Fatalf("IncompleteReason = %q, want it to mention %q", result.IncompleteReason, tc.wantReason)
}
})
}
}
func TestRunInjectsToolFailureHintWithSchema(t *testing.T) {
+15
View File
@@ -587,6 +587,14 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
// counted toward the runaway cap so we stop before burning maxTurns.
if guards.observeTurn(collected) {
result.FinalAnswer = noOutputStopAnswer(result.Turns)
// Under the headless completion gate this stop is unfinished work,
// not a successful final answer: the guard exists precisely because
// the model made no progress. Interactive runs keep the gate off, so
// their behavior is unchanged.
if options.RequireCompletionSignal {
result.Incomplete = true
result.IncompleteReason = "stopped after " + strconv.Itoa(guards.emptyTurns) + " consecutive turns with no output"
}
result.Messages = copyMessages(messages)
return result, nil
}
@@ -808,6 +816,13 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
messages = appendAbortedToolResults(messages, collected.ToolCalls[index+1:])
messages = append(messages, buildToolImageMessages()...)
result.FinalAnswer = toolFailureStopAnswer(call.Name, outcome.Count)
// Same headless contract as the empty-turn guard: halting on a
// repeated failure is an unfinished run, not success. Gate-off
// (interactive / --no-completion-gate) keeps the old behavior.
if options.RequireCompletionSignal {
result.Incomplete = true
result.IncompleteReason = "the " + call.Name + " tool failed " + strconv.Itoa(outcome.Count) + " times in a row with the same error"
}
result.Messages = copyMessages(messages)
return result, nil
}
+45
View File
@@ -1,14 +1,20 @@
package agent
import (
"fmt"
"os"
"path/filepath"
"strings"
"testing"
"github.com/Twigpine/zero/internal/config"
"github.com/Twigpine/zero/internal/sandbox"
"github.com/Twigpine/zero/internal/testutil"
)
// isolatedUserDirsRoot is where TestMain points every per-user directory.
var isolatedUserDirsRoot string
// TestMain points userConfigDirForPrompt at an empty, package-wide temp
// directory for every test in this package. Without this, buildSystemPrompt
// falls back to the real config.UserConfigDir, so any developer with a
@@ -17,17 +23,56 @@ import (
// across contributor machines. Tests that specifically exercise user
// guidelines stub userConfigDirForPrompt themselves via
// withSystemPromptTestUserConfigDir(Func) and restore this default on cleanup.
//
// It also points every per-user directory root at the same temp directory:
// sandboxed tool runs keep their runtime state and lease files under the user
// cache dir, so without this they land in the developer's real cache.
func TestMain(m *testing.M) {
dir, err := os.MkdirTemp("", "zero-agent-tests-*")
if err != nil {
panic(err)
}
isolatedUserDirsRoot = dir
restore, err := testutil.IsolateUserDirs(dir)
if err != nil {
os.RemoveAll(dir)
panic(err)
}
userConfigDirForPrompt = func() (string, error) { return dir, nil }
code := m.Run()
if err := restore(); err != nil {
fmt.Fprintln(os.Stderr, "agent tests: restore user dirs:", err)
code = 1
}
os.RemoveAll(dir)
os.Exit(code)
}
// Guards TestMain: the per-user directories the package can read or write
// must resolve inside the isolated root on every OS.
func TestUserDirsAreIsolatedFromTheDeveloperHome(t *testing.T) {
if isolatedUserDirsRoot == "" {
t.Fatal("TestMain did not record an isolated root")
}
home, err := os.UserHomeDir()
if err != nil {
t.Fatal(err)
}
userConfig, err := config.UserConfigDir()
if err != nil {
t.Fatal(err)
}
cache, err := os.UserCacheDir()
if err != nil {
t.Fatal(err)
}
for name, dir := range map[string]string{"home": home, "config": userConfig, "cache": cache} {
if !testutil.PathWithin(isolatedUserDirsRoot, dir) {
t.Errorf("%s dir %q is outside the isolated root %q", name, dir, isolatedUserDirsRoot)
}
}
}
func TestCoreSystemPromptIncludesCodingQualityRules(t *testing.T) {
prompt := strings.Join(strings.Fields(strings.ToLower(buildSystemPrompt(Options{}))), " ")
if strings.Contains(prompt, "clickable file paths") {
+57
View File
@@ -0,0 +1,57 @@
package aimlapi
import (
"net/http"
"net/url"
"strings"
)
const (
// SourceHeaderName and SourceValue identify Zero as the calling agent on
// requests to aimlapi.com services (the Toolbox and the MCP server).
SourceHeaderName = "X-AIMLAPI-Source"
SourceValue = "agent/zero"
// CanonicalInferenceBaseURL is the catalog inference endpoint. A provider
// profile whose base URL differs is treated as a proxy: its key is never sent
// to other aimlapi.com services.
CanonicalInferenceBaseURL = "https://api.aimlapi.com/v1"
// ToolboxBaseURL is the fixed AI/ML API Toolbox REST base. It is not
// configurable, so the key and attribution can only reach tools.aimlapi.com.
ToolboxBaseURL = "https://tools.aimlapi.com/v1/tools"
// MCPServerURL is the AI/ML API MCP server (streamable HTTP).
MCPServerURL = "https://mcp.aimlapi.com/mcp"
)
// attributedHosts are the only hostnames that receive the AI/ML API key or the
// attribution headers. Matching is exact; a suffix or substring match would
// also accept hosts such as tools.aimlapi.com.example.net.
var attributedHosts = [...]string{"api.aimlapi.com", "tools.aimlapi.com", "mcp.aimlapi.com"}
// IsAttributedURL reports whether target is an https URL on one of the
// aimlapi.com service hosts, on the default port and without userinfo.
func IsAttributedURL(target *url.URL) bool {
if target == nil || !strings.EqualFold(target.Scheme, "https") || target.User != nil {
return false
}
if port := target.Port(); port != "" && port != "443" {
return false
}
host := strings.ToLower(target.Hostname())
for _, candidate := range attributedHosts {
if host == candidate {
return true
}
}
return false
}
// ApplyAttributionHeaders sets the source and partner headers on header only
// when target is an aimlapi.com service URL; any other target is left as is.
func ApplyAttributionHeaders(header http.Header, target *url.URL) {
if header == nil || !IsAttributedURL(target) {
return
}
header.Set(SourceHeaderName, SourceValue)
header.Set(PartnerHeaderName, ResolvePartnerID(""))
}
+54
View File
@@ -0,0 +1,54 @@
package aimlapi
import (
"net/http"
"net/url"
"testing"
)
func TestIsAttributedURLMatchesExactServiceHostsOnly(t *testing.T) {
cases := map[string]bool{
"https://tools.aimlapi.com/v1/tools/run": true,
"https://mcp.aimlapi.com/mcp": true,
"https://api.aimlapi.com/v1": true,
"https://TOOLS.aimlapi.com:443/v1/tools": true,
"http://tools.aimlapi.com/v1/tools/run": false,
"https://tools.aimlapi.com:8443/v1/tools": false,
"https://[email protected]/v1/tools": false,
"https://tools.aimlapi.com.example.net/run": false,
"https://eviltools.aimlapi.com/run": false,
"https://aimlapi.com/app": false,
"https://auth.aimlapi.com/": false,
"https://example.com/?h=tools.aimlapi.com": false,
}
for raw, want := range cases {
target, err := url.Parse(raw)
if err != nil {
t.Fatalf("parse %q: %v", raw, err)
}
if got := IsAttributedURL(target); got != want {
t.Errorf("IsAttributedURL(%q) = %v, want %v", raw, got, want)
}
}
if IsAttributedURL(nil) {
t.Fatal("IsAttributedURL(nil) = true")
}
}
func TestApplyAttributionHeadersSetsSourceAndPartnerOnServiceHosts(t *testing.T) {
t.Setenv("AIMLAPI_PARTNER_ID", "")
header := http.Header{}
ApplyAttributionHeaders(header, &url.URL{Scheme: "https", Host: "tools.aimlapi.com", Path: "/v1/tools/run"})
if got := header.Get(SourceHeaderName); got != "agent/zero" {
t.Fatalf("%s = %q, want agent/zero", SourceHeaderName, got)
}
if got := header.Get(PartnerHeaderName); got != DefaultPartnerID {
t.Fatalf("%s = %q, want %q", PartnerHeaderName, got, DefaultPartnerID)
}
other := http.Header{}
ApplyAttributionHeaders(other, &url.URL{Scheme: "https", Host: "tools.aimlapi.com.example.net"})
if len(other) != 0 {
t.Fatalf("attribution reached a non-aimlapi host: %#v", other)
}
}
+25 -2
View File
@@ -964,6 +964,7 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
// report can be combined with the plugin activation's, and emit at most one
// notice when project hooks/plugins were dropped for an untrusted workspace.
hookDispatcher, hookSkip := newHookDispatcherWithExtra(workspaceRoot, pluginActivation.hooks, trustRoot, executionRunner)
specialistRuntime.attachLifecycleHooks(hookDispatcher)
emitTrustNotice(stderr, hookSkip, pluginActivation.trustSkip, mcpSkip)
// From the engine the session's commands run through, so the notice
// describes the sandbox they actually get.
@@ -1221,6 +1222,9 @@ type agentToolRuntime struct {
// the orchestrator's system-prompt delegation section. Populated alongside the
// tools, so it is present exactly when the Task tool is.
specialists []agent.SpecialistInfo
// lifecycleHooks is shared with the specialist executor so the hooks
// dispatcher (built later in startup) can be attached before any specialist runs.
lifecycleHooks *specialist.LifecycleHooks
}
// specialistInfos returns the runtime's specialist summaries, nil-safe so the
@@ -1232,12 +1236,26 @@ func (r *agentToolRuntime) specialistInfos() []agent.SpecialistInfo {
return r.specialists
}
func (r *agentToolRuntime) attachLifecycleHooks(dispatcher *hooks.Dispatcher) {
if r == nil || r.lifecycleHooks == nil || dispatcher == nil {
return
}
r.lifecycleHooks.Dispatch = func(ctx context.Context, event string, specialistName string, payload map[string]any) {
dispatcher.Dispatch(ctx, hooks.DispatchInput{
Event: hooks.Event(event),
ToolName: specialistName,
Payload: payload,
})
}
}
func registerSpecialistTools(registry *tools.Registry, workspaceRoot string, maxTeamSize int) (*agentToolRuntime, error) {
paths, err := specialist.DefaultPaths(workspaceRoot)
if err != nil {
return nil, err
}
executor := specialist.Executor{Paths: paths}
lifecycleHooks := &specialist.LifecycleHooks{}
executor := specialist.Executor{Paths: paths, LifecycleHooks: lifecycleHooks}
runtime, err := specialist.RegisterTools(registry, executor)
if err != nil {
return nil, err
@@ -1256,7 +1274,12 @@ func registerSpecialistTools(registry *tools.Registry, workspaceRoot string, max
return nil, err
}
swarm.RegisterTools(registry, sw)
return &agentToolRuntime{specialist: runtime, swarm: sw, specialists: specialistSummaries(paths)}, nil
return &agentToolRuntime{
specialist: runtime,
swarm: sw,
specialists: specialistSummaries(paths),
lifecycleHooks: lifecycleHooks,
}, nil
}
// specialistSummaries loads the available specialists (built-ins + user/project
+13 -1
View File
@@ -92,13 +92,25 @@ func runSkillAdd(args []string, dir string, stdout io.Writer, stderr io.Writer)
lines := []string{}
if result.Updated {
lines = append(lines, fmt.Sprintf("Updated skill %q.", result.Name))
lines = append(lines, fmt.Sprintf(" hash: %s -> %s", shortHash(result.PreviousHash), shortHash(result.Hash)))
switch {
case result.PreviousHash != result.Hash:
lines = append(lines, fmt.Sprintf(" hash: %s -> %s", shortHash(result.PreviousHash), shortHash(result.Hash)))
case result.PreviousTreeHash == "":
// The install this replaced predates tree hashes and copied SKILL.md
// alone, so the bundled files are new rather than changed.
lines = append(lines, " SKILL.md is unchanged; the files bundled with it are now installed too")
default:
lines = append(lines, " SKILL.md is unchanged; the files bundled with it changed")
}
} else {
lines = append(lines, fmt.Sprintf("Installed skill %q.", result.Name))
lines = append(lines, " hash: "+result.Hash)
}
lines = append(lines, " source: "+result.Source)
lines = append(lines, " path: "+result.Path)
if len(result.SkippedLinks) > 0 {
lines = append(lines, " not copied (symlinks): "+strings.Join(result.SkippedLinks, ", "))
}
if _, err := fmt.Fprintln(stdout, redaction.RedactString(strings.Join(lines, "\n"), redaction.Options{})); err != nil {
return exitCrash
}
+132
View File
@@ -7,6 +7,8 @@ import (
"path/filepath"
"strings"
"testing"
"github.com/Twigpine/zero/internal/skills"
)
func writeSourceSkillDir(t *testing.T, dir string, content string) string {
@@ -391,3 +393,133 @@ func TestRunDistributionCommandsDoNotLaunchTUI(t *testing.T) {
}
}
}
// A reinstall that only changed a bundled script must still read as an update
// (#584): the SKILL.md hash is the same, so "hash: X -> X" would say nothing.
func TestRunSkillAddReportsAnUpdateToBundledFilesOnly(t *testing.T) {
skillsDir := t.TempDir()
src := writeSourceSkillDir(t, filepath.Join(t.TempDir(), "src"),
"---\nname: pdf\ndescription: Fill forms.\n---\nRun scripts/fill.py.\n")
script := filepath.Join(src, "scripts", "fill.py")
if err := os.MkdirAll(filepath.Dir(script), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(script, []byte("print('v1')\n"), 0o755); err != nil {
t.Fatal(err)
}
deps := appDeps{skillsDir: func() string { return skillsDir }}
var stdout, stderr bytes.Buffer
if exit := runWithDeps([]string{"skill", "add", src}, &stdout, &stderr, deps); exit != 0 {
t.Fatalf("first skill add exit = %d, stderr = %s", exit, stderr.String())
}
if err := os.WriteFile(script, []byte("print('v2')\n"), 0o755); err != nil {
t.Fatal(err)
}
stdout.Reset()
stderr.Reset()
if exit := runWithDeps([]string{"skill", "add", src}, &stdout, &stderr, deps); exit != 0 {
t.Fatalf("second skill add exit = %d, stderr = %s", exit, stderr.String())
}
for _, want := range []string{`Updated skill "pdf".`, "SKILL.md is unchanged; the files bundled with it changed"} {
if !strings.Contains(stdout.String(), want) {
t.Fatalf("second add output does not say %q:\n%s", want, stdout.String())
}
}
}
// An install from before tree hashes copied SKILL.md alone and recorded no tree
// hash. Adding the skill again over that entry must not claim bundled files
// changed: there are none, or they are installed for the first time.
func TestRunSkillAddOverAnEntryWithoutATreeHash(t *testing.T) {
for _, tc := range []struct {
name string
script bool
want []string
notWant []string
}{
{
name: "SKILL.md only",
want: []string{`Installed skill "pdf".`},
notWant: []string{"Updated skill", "files bundled with it"},
},
{
name: "with a bundled script",
script: true,
want: []string{`Updated skill "pdf".`, "SKILL.md is unchanged; the files bundled with it are now installed too"},
notWant: []string{"files bundled with it changed"},
},
} {
t.Run(tc.name, func(t *testing.T) {
skillsDir := t.TempDir()
src := writeSourceSkillDir(t, filepath.Join(t.TempDir(), "src"),
"---\nname: pdf\ndescription: Fill forms.\n---\nRun scripts/fill.py.\n")
if tc.script {
script := filepath.Join(src, "scripts", "fill.py")
if err := os.MkdirAll(filepath.Dir(script), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(script, []byte("print('fill')\n"), 0o755); err != nil {
t.Fatal(err)
}
}
deps := appDeps{skillsDir: func() string { return skillsDir }}
var stdout, stderr bytes.Buffer
if exit := runWithDeps([]string{"skill", "add", src}, &stdout, &stderr, deps); exit != 0 {
t.Fatalf("first skill add exit = %d, stderr = %s", exit, stderr.String())
}
lock, err := skills.ReadLock(skillsDir)
if err != nil {
t.Fatal(err)
}
entry := lock["pdf"]
entry.TreeHash = ""
lock["pdf"] = entry
data, err := json.Marshal(lock)
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(skillsDir, skills.LockFileName), data, 0o644); err != nil {
t.Fatal(err)
}
stdout.Reset()
stderr.Reset()
if exit := runWithDeps([]string{"skill", "add", src}, &stdout, &stderr, deps); exit != 0 {
t.Fatalf("second skill add exit = %d, stderr = %s", exit, stderr.String())
}
for _, want := range tc.want {
if !strings.Contains(stdout.String(), want) {
t.Fatalf("second add output does not say %q:\n%s", want, stdout.String())
}
}
for _, notWant := range tc.notWant {
if strings.Contains(stdout.String(), notWant) {
t.Fatalf("second add output says %q:\n%s", notWant, stdout.String())
}
}
})
}
}
func TestRunSkillAddListsTheSymlinksItDidNotCopy(t *testing.T) {
skillsDir := t.TempDir()
src := writeSourceSkillDir(t, filepath.Join(t.TempDir(), "src"),
"---\nname: pdf\ndescription: Fill forms.\n---\nbody\n")
outside := filepath.Join(t.TempDir(), "notes.md")
if err := os.WriteFile(outside, []byte("outside\n"), 0o644); err != nil {
t.Fatal(err)
}
if err := os.Symlink(outside, filepath.Join(src, "notes.md")); err != nil {
t.Skipf("symlinks unavailable: %v", err)
}
var stdout, stderr bytes.Buffer
deps := appDeps{skillsDir: func() string { return skillsDir }}
if exit := runWithDeps([]string{"skill", "add", src}, &stdout, &stderr, deps); exit != 0 {
t.Fatalf("skill add exit = %d, stderr = %s", exit, stderr.String())
}
if !strings.Contains(stdout.String(), "not copied (symlinks): notes.md") {
t.Fatalf("add output does not list the skipped symlink:\n%s", stdout.String())
}
}
+2 -1
View File
@@ -655,6 +655,7 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
// can be combined with the plugin activation's, and emit at most one notice when
// project hooks/plugins were dropped for an untrusted workspace.
hookDispatcher, hookSkip := newHookDispatcherWithExtra(workspaceRoot, pluginActivation.hooks, trustRoot, executionRunner)
specialistRuntime.attachLifecycleHooks(hookDispatcher)
emitTrustNotice(stderr, hookSkip, pluginActivation.trustSkip, mcpSkip)
result, err := agent.Run(runCtx, agentPrompt, provider, agent.Options{
MaxTurns: resolved.MaxTurns,
@@ -712,7 +713,7 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
"arguments": call.Arguments,
})
// Snapshot before-state of files this call will mutate (safe rewind).
if checkpoint, ok := sessionRecorder.captureCheckpoint(workspaceRoot, call); ok {
if checkpoint, ok := sessionRecorder.captureCheckpoint(workspaceRoot, registry, call); ok {
writer.checkpoint(checkpoint)
}
},
@@ -0,0 +1,206 @@
package cli
import (
"bytes"
"context"
"encoding/json"
"os"
"path/filepath"
"testing"
"github.com/Twigpine/zero/internal/agent"
"github.com/Twigpine/zero/internal/config"
"github.com/Twigpine/zero/internal/sessions"
"github.com/Twigpine/zero/internal/tools"
"github.com/Twigpine/zero/internal/zeroruntime"
)
func freeformApplyPatchArguments(root string, absolute bool) string {
target := "main.txt"
if absolute {
target = filepath.Join(root, "main.txt")
}
return "*** Begin Patch\n*** Update File: " + target + "\n@@\n-before\n+after\n*** End Patch\n"
}
// captureCheckpoint must resolve a freeform apply_patch call through the same
// preparation execution uses, so the before-state is captured and rewind can
// restore it. The JSON control pins the unchanged path.
func TestCaptureCheckpointFreeformApplyPatch(t *testing.T) {
for _, tc := range []struct {
name string
freeform bool
absolute bool
jsonArgs bool
}{
{name: "freeform", freeform: true},
{name: "freeform absolute header", freeform: true, absolute: true},
{name: "json arguments", jsonArgs: true},
} {
t.Run(tc.name, func(t *testing.T) {
t.Setenv("XDG_DATA_HOME", t.TempDir())
t.Setenv("ZERO_CHECKPOINTS", "on")
root := t.TempDir()
// The absolute-header case needs the physical root spelling the tool
// resolves (macOS /var vs /private/var, Windows short names).
if resolved, err := filepath.EvalSymlinks(root); err == nil {
root = resolved
}
if err := os.WriteFile(filepath.Join(root, "main.txt"), []byte("before\n"), 0o644); err != nil {
t.Fatal(err)
}
store := sessions.NewStore(sessions.StoreOptions{})
session, err := store.Create(sessions.CreateInput{Title: "freeform checkpoint"})
if err != nil {
t.Fatal(err)
}
if _, err := store.AppendEvent(session.SessionID, sessions.AppendEventInput{
Type: sessions.EventMessage,
Payload: map[string]string{"role": "user", "content": "apply the patch"},
}); err != nil {
t.Fatal(err)
}
recorder := &execSessionRecorder{prepared: sessions.PreparedExec{Store: store, Session: session}}
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(root, nil))
call := agent.ToolCall{ID: "call", Name: "apply_patch", Freeform: tc.freeform}
if tc.jsonArgs {
encoded, err := json.Marshal(map[string]string{"patch": freeformApplyPatchArguments(root, false)})
if err != nil {
t.Fatal(err)
}
call.Arguments = string(encoded)
} else {
call.Arguments = freeformApplyPatchArguments(root, tc.absolute)
}
event, ok := recorder.captureCheckpoint(root, registry, call)
if !ok || event.Type != sessions.EventSessionCheckpoint {
t.Fatalf("captureCheckpoint() = %#v, %t; want a checkpoint event", event, ok)
}
var payload sessions.CheckpointPayload
if err := json.Unmarshal(event.Payload, &payload); err != nil {
t.Fatalf("decode checkpoint payload: %v", err)
}
if len(payload.Files) != 1 || payload.Files[0].Path != "main.txt" || payload.Files[0].Blob == "" {
t.Fatalf("checkpoint files = %#v, want main.txt with a blob", payload.Files)
}
// Simulate the tool's mutation, then rewind past it.
if err := os.WriteFile(filepath.Join(root, "main.txt"), []byte("after\n"), 0o644); err != nil {
t.Fatal(err)
}
report, err := store.ApplyRewind(session.SessionID, root, 1)
if err != nil {
t.Fatalf("ApplyRewind: %v", err)
}
if report.FilesRestored != 1 {
t.Fatalf("restore report = %#v, want one restored file", report)
}
content, err := os.ReadFile(filepath.Join(root, "main.txt"))
if err != nil {
t.Fatal(err)
}
if string(content) != "before\n" {
t.Fatalf("restored content = %q, want before newline", content)
}
})
}
}
// freeformPatchExecProvider emits one native freeform apply_patch call, then a
// final text turn.
type freeformPatchExecProvider struct {
patch string
calls int
}
func (provider *freeformPatchExecProvider) StreamCompletion(ctx context.Context, request zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) {
ch := make(chan zeroruntime.StreamEvent, 5)
calls := provider.calls
provider.calls++
if calls == 0 {
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call_patch", ToolName: "apply_patch", ToolCallFreeform: true}
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call_patch", ArgumentsFragment: provider.patch}
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call_patch"}
} else {
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventText, Content: "patched"}
}
ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}
close(ch)
return ch, nil
}
// End-to-end: exec captures the freeform checkpoint and `sessions rewind`
// restores the pre-patch content.
func TestRunExecFreeformApplyPatchCheckpointsForRewind(t *testing.T) {
isolateCLIUserState(t)
t.Setenv("ZERO_CHECKPOINTS", "on")
root := t.TempDir()
if err := os.WriteFile(filepath.Join(root, "main.txt"), []byte("before\n"), 0o644); err != nil {
t.Fatal(err)
}
provider := &freeformPatchExecProvider{patch: freeformApplyPatchArguments(root, false)}
var stdout, stderr bytes.Buffer
exitCode := runWithDeps([]string{
"exec", "--init-session-id", "freeform_rewind", "--skip-permissions-unsafe",
"--output-format", "stream-json", "apply the patch",
}, &stdout, &stderr, appDeps{
getwd: func() (string, error) { return root, nil },
resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) {
return execResolvedConfig(), nil
},
newProvider: func(config.ProviderProfile) (zeroruntime.Provider, error) {
return provider, nil
},
})
if exitCode != exitSuccess {
t.Fatalf("exit = %d, want %d; stderr=%q stdout=%q", exitCode, exitSuccess, stderr.String(), stdout.String())
}
content, err := os.ReadFile(filepath.Join(root, "main.txt"))
if err != nil {
t.Fatal(err)
}
if string(content) != "after\n" {
t.Fatalf("patched content = %q, want after newline", content)
}
store := sessions.NewStore(sessions.StoreOptions{})
events, err := store.ReadEvents("freeform_rewind")
if err != nil {
t.Fatalf("ReadEvents: %v", err)
}
checkpoints := 0
for _, event := range events {
if event.Type == sessions.EventSessionCheckpoint {
checkpoints++
}
}
if checkpoints != 1 {
t.Fatalf("checkpoint events = %d, want 1; events=%v", checkpoints, sessionEventTypes(events))
}
report, err := store.ApplyRewind("freeform_rewind", root, 1)
if err != nil {
t.Fatalf("ApplyRewind: %v", err)
}
if report.FilesRestored != 1 {
t.Fatalf("restore report = %#v, want one restored file", report)
}
content, err = os.ReadFile(filepath.Join(root, "main.txt"))
if err != nil {
t.Fatal(err)
}
if string(content) != "before\n" {
t.Fatalf("rewound content = %q, want before newline", content)
}
}
func sessionEventTypes(events []sessions.Event) []string {
types := make([]string, 0, len(events))
for _, event := range events {
types = append(types, string(event.Type))
}
return types
}
@@ -0,0 +1,90 @@
package cli
import (
"bytes"
"context"
"slices"
"strings"
"testing"
"github.com/Twigpine/zero/internal/config"
"github.com/Twigpine/zero/internal/zeroruntime"
)
// guardStopExecProvider emits turns with no visible text and no tool calls, so
// the empty-turn guardrail stops the run after its consecutive-empty-turn cap.
type guardStopExecProvider struct{}
func (guardStopExecProvider) StreamCompletion(ctx context.Context, request zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) {
ch := make(chan zeroruntime.StreamEvent, 1)
select {
case <-ctx.Done():
close(ch)
return ch, ctx.Err()
case ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}:
}
close(ch)
return ch, nil
}
// A guardrail stop in headless exec must use the existing INCOMPLETE contract
// (exit 4, error event, run_end status), and --no-completion-gate must keep
// reporting success. isolateCLIUserState also turns off the models.dev refresh
// that runExec starts in the background (see the #1152/#1161 flake).
func TestRunExecGuardrailStopIncompleteContract(t *testing.T) {
for _, tc := range []struct {
name string
args []string
wantExit int
wantStatus string
}{
{"gate on", []string{"exec", "--output-format", "stream-json", "complete the task"}, exitIncomplete, "incomplete"},
{"gate off", []string{"exec", "--no-completion-gate", "--output-format", "stream-json", "complete the task"}, exitSuccess, "success"},
} {
t.Run(tc.name, func(t *testing.T) {
isolateCLIUserState(t)
cwd := t.TempDir()
var stdout, stderr bytes.Buffer
exitCode := runWithDeps(tc.args, &stdout, &stderr, appDeps{
getwd: func() (string, error) { return cwd, nil },
resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) {
return execResolvedConfig(), nil
},
newProvider: func(config.ProviderProfile) (zeroruntime.Provider, error) {
return guardStopExecProvider{}, nil
},
})
if exitCode != tc.wantExit {
t.Fatalf("exit code = %d, want %d; stderr=%q", exitCode, tc.wantExit, stderr.String())
}
events := decodeJSONLines(t, stdout.String())
if len(events) == 0 || events[len(events)-1]["type"] != "run_end" {
t.Fatalf("expected run_end as the last event, got %#v", events)
}
runEnd := events[len(events)-1]
if runEnd["status"] != tc.wantStatus {
t.Fatalf("run_end status = %v, want %q", runEnd["status"], tc.wantStatus)
}
if runEnd["exitCode"] != float64(tc.wantExit) {
t.Fatalf("run_end exitCode = %v, want %d", runEnd["exitCode"], tc.wantExit)
}
finalEvent := findJSONEvent(t, events, "final")
if text, _ := finalEvent["text"].(string); !strings.Contains(text, "no output") {
t.Fatalf("final text = %q, want the guard stop answer", text)
}
if tc.wantStatus == "incomplete" {
errorEvent := findJSONEvent(t, events, "error")
if errorEvent["code"] != "incomplete" {
t.Fatalf("error code = %v, want incomplete", errorEvent["code"])
}
if message, _ := errorEvent["message"].(string); !strings.Contains(message, "no output") {
t.Fatalf("error message = %q, want the guard reason", message)
}
} else if slices.Contains(jsonEventTypes(events), "error") {
t.Fatalf("gate-off run must not emit an error event: %#v", events)
}
})
}
}
+2 -7
View File
@@ -1,7 +1,6 @@
package cli
import (
"encoding/json"
"fmt"
"io"
"strings"
@@ -19,15 +18,11 @@ type execSessionRecorder struct {
// captureCheckpoint snapshots the before-state of files a tool call will mutate.
// Best-effort: failures never affect the run. Returns the checkpoint event and
// true when one was recorded.
func (recorder *execSessionRecorder) captureCheckpoint(workspaceRoot string, call agent.ToolCall) (sessions.Event, bool) {
func (recorder *execSessionRecorder) captureCheckpoint(workspaceRoot string, registry *tools.Registry, call agent.ToolCall) (sessions.Event, bool) {
if recorder.prepared.Store == nil || recorder.prepared.Session.SessionID == "" {
return sessions.Event{}, false
}
var args map[string]any
if call.Arguments != "" {
_ = json.Unmarshal([]byte(call.Arguments), &args)
}
targets := tools.MutationTargets(workspaceRoot, call.Name, args)
targets := agent.CheckpointTargets(workspaceRoot, registry, call)
if len(targets) == 0 {
return sessions.Event{}, false
}
+1 -1
View File
@@ -148,7 +148,7 @@ func runExecSpecDraft(run execSpecDraftRun) int {
"name": call.Name,
"arguments": call.Arguments,
})
if checkpoint, ok := sessionRecorder.captureCheckpoint(run.workspaceRoot, call); ok {
if checkpoint, ok := sessionRecorder.captureCheckpoint(run.workspaceRoot, run.registry, call); ok {
writer.checkpoint(checkpoint)
}
},
+116
View File
@@ -2,6 +2,7 @@ package cli
import (
"context"
"errors"
"fmt"
"io"
"strings"
@@ -82,11 +83,105 @@ func runPlugins(args []string, stdout io.Writer, stderr io.Writer, deps appDeps)
return runPluginInfo(args[1:], stdout, stderr, deps)
case "remove", "rm":
return runPluginRemove(args[1:], deps.pluginsDir(), stdout, stderr)
case "enable":
return runPluginsToggle(args[1:], stdout, stderr, deps, false)
case "disable":
return runPluginsToggle(args[1:], stdout, stderr, deps, true)
default:
return writeExecUsageError(stderr, fmt.Sprintf("unknown plugins subcommand %q", args[0]))
}
}
type pluginToggleOptions struct {
json bool
user bool
}
// runPluginsToggle implements `zero plugins enable|disable <id>`.
//
// Every message it writes, including usage errors that echo the caller's
// arguments, passes through the redaction pipeline: an argument is user input
// and can be a mistyped credential.
func runPluginsToggle(args []string, stdout io.Writer, stderr io.Writer, deps appDeps, disabled bool) int {
commandName := "enable"
if disabled {
commandName = "disable"
}
usageError := func(message string) int {
return writeExecUsageError(stderr, redaction.RedactString(message, redaction.Options{}))
}
options, positional, help, err := parsePluginToggleArgs(args, commandName)
if err != nil {
return usageError(err.Error())
}
if help {
if err := writePluginsToggleHelp(stdout, commandName); err != nil {
return exitCrash
}
return exitSuccess
}
if len(positional) != 1 {
return usageError(fmt.Sprintf("usage: zero plugins %s <id> [--user] [--json]", commandName))
}
pluginID := strings.TrimSpace(positional[0])
if pluginID == "" {
return usageError(fmt.Sprintf("usage: zero plugins %s <id> [--user] [--json]", commandName))
}
cwd, err := deps.getwd()
if err != nil {
return writeAppError(stderr, "failed to resolve workspace: "+redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
loadOptions := plugins.LoadOptions{Cwd: cwd, ExcludeProject: options.user}
result, err := plugins.SetEnabledByID(loadOptions, pluginID, !disabled)
if err != nil {
if errors.Is(err, plugins.ErrNotInstalled) {
return usageError(err.Error())
}
return writeAppError(stderr, redaction.ErrorMessage(err, redaction.Options{}), exitCrash)
}
state := "enabled"
if disabled {
state = "disabled"
}
if options.json {
if err := writePrettyJSON(stdout, redaction.RedactValue(result, redaction.Options{})); err != nil {
return exitCrash
}
return exitSuccess
}
message := fmt.Sprintf("Plugin %s is now %s in %s.\n", result.ID, state, result.ManifestPath)
if !result.Changed {
message = fmt.Sprintf("Plugin %s was already %s in %s.\n", result.ID, state, result.ManifestPath)
}
if _, err := fmt.Fprint(stdout, redaction.RedactString(message, redaction.Options{})); err != nil {
return exitCrash
}
return exitSuccess
}
func parsePluginToggleArgs(args []string, command string) (pluginToggleOptions, []string, bool, error) {
options := pluginToggleOptions{}
positional := []string{}
for _, arg := range args {
switch arg {
case "-h", "--help", "help":
return options, positional, true, nil
case "--json":
options.json = true
case "--user":
options.user = true
default:
if strings.HasPrefix(arg, "-") {
return options, positional, false, execUsageError{fmt.Sprintf("unknown plugins %s flag %q", command, arg)}
}
positional = append(positional, arg)
}
}
return options, positional, false, nil
}
func runHooks(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int {
if len(args) == 0 {
return writeExecUsageError(stderr, "hooks subcommand required. Use `zero hooks list`.")
@@ -566,10 +661,31 @@ Commands:
info <id> Show plugin details and lockfile metadata
add <git-url|path> Install a plugin (manifest-validated, pinned in plugins.lock)
remove <id> Remove an installed plugin and its lockfile entry
enable <id> Enable a plugin via its plugin.json
disable <id> Disable a plugin via its plugin.json
`)
return err
}
func writePluginsToggleHelp(w io.Writer, command string) error {
_, err := fmt.Fprintf(w, `Usage:
zero plugins %s <id> [flags]
Flags:
--user Target only user plugins (ignore ./.zero/plugins)
--json Print command result as JSON
-h, --help Show this help
Without --user, a project plugin with the same id takes precedence over a user
plugin, matching the loader's discovery order.
Zero reads plugins when it starts, so the change applies to the next zero exec
run or interactive session. A running interactive session keeps the plugins it
started with until you restart it.
`, command)
return err
}
func writePluginsListHelp(w io.Writer) error {
_, err := fmt.Fprint(w, `Usage:
zero plugins list [flags]
+232
View File
@@ -5,6 +5,7 @@ import (
"context"
"encoding/json"
"errors"
"os"
"path/filepath"
"strings"
"testing"
@@ -333,3 +334,234 @@ func TestRunMCPPermissionsHelpDoesNotOpenStore(t *testing.T) {
})
}
}
func TestRunPluginsEnableDisable(t *testing.T) {
cwd := t.TempDir()
home := t.TempDir()
t.Setenv("XDG_CONFIG_HOME", home)
t.Setenv("HOME", home)
pluginDir := filepath.Join(cwd, ".zero", "plugins", "demo")
if err := os.MkdirAll(pluginDir, 0o700); err != nil {
t.Fatal(err)
}
manifestPath := filepath.Join(pluginDir, "plugin.json")
if err := os.WriteFile(manifestPath, []byte(`{
"schemaVersion": 1,
"id": "zero.demo",
"name": "Demo",
"version": "1.0.0"
}
`), 0o600); err != nil {
t.Fatal(err)
}
deps := appDeps{getwd: func() (string, error) { return cwd, nil }}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"plugins", "disable", "zero.demo", "--json"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("disable exit=%d stderr=%s", exitCode, stderr.String())
}
var disablePayload plugins.SetEnabledResult
if err := json.Unmarshal(stdout.Bytes(), &disablePayload); err != nil {
t.Fatalf("decode disable JSON: %v\n%s", err, stdout.String())
}
if disablePayload.ID != "zero.demo" || disablePayload.Enabled || !disablePayload.Changed {
t.Fatalf("unexpected disable payload: %#v", disablePayload)
}
raw, err := os.ReadFile(manifestPath)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(string(raw), `"enabled": false`) {
t.Fatalf("manifest not disabled: %s", raw)
}
stdout.Reset()
stderr.Reset()
exitCode = runWithDeps([]string{"plugins", "disable", "zero.demo"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("noop disable exit=%d stderr=%s", exitCode, stderr.String())
}
if !strings.Contains(stdout.String(), "was already disabled") {
t.Fatalf("noop disable output: %s", stdout.String())
}
stdout.Reset()
stderr.Reset()
exitCode = runWithDeps([]string{"plugins", "enable", "zero.demo"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("enable exit=%d stderr=%s", exitCode, stderr.String())
}
if !strings.Contains(stdout.String(), "is now enabled") {
t.Fatalf("enable output: %s", stdout.String())
}
stdout.Reset()
stderr.Reset()
exitCode = runWithDeps([]string{"plugins", "enable", "missing-plugin"}, &stdout, &stderr, deps)
if exitCode != exitUsage {
t.Fatalf("missing plugin exit=%d want %d stderr=%s", exitCode, exitUsage, stderr.String())
}
}
// The text-mode toggle output runs through the redaction pipeline, so a
// manifest path that contains something secret-shaped is masked.
func TestRunPluginsToggleTextOutputIsRedacted(t *testing.T) {
cwd := t.TempDir()
home := t.TempDir()
t.Setenv("XDG_CONFIG_HOME", home)
t.Setenv("HOME", home)
secretDir := "sk-proj-abcdefghijklmnopqrstuvwxyz"
pluginDir := filepath.Join(cwd, ".zero", "plugins", secretDir)
if err := os.MkdirAll(pluginDir, 0o700); err != nil {
t.Fatal(err)
}
manifestPath := filepath.Join(pluginDir, "plugin.json")
if err := os.WriteFile(manifestPath, []byte(`{
"schemaVersion": 1,
"id": "zero.demo",
"name": "Demo",
"version": "1.0.0"
}
`), 0o600); err != nil {
t.Fatal(err)
}
deps := appDeps{getwd: func() (string, error) { return cwd, nil }}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"plugins", "disable", "zero.demo"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("disable exit=%d stderr=%s", exitCode, stderr.String())
}
if strings.Contains(stdout.String(), secretDir) {
t.Fatalf("toggle text leaked a secret-shaped manifest path: %q", stdout.String())
}
if !strings.Contains(stdout.String(), "[REDACTED]") {
t.Fatalf("toggle text was not redacted: %q", stdout.String())
}
stdout.Reset()
stderr.Reset()
exitCode = runWithDeps([]string{"plugins", "disable", "zero.demo"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("noop disable exit=%d stderr=%s", exitCode, stderr.String())
}
if !strings.Contains(stdout.String(), "was already disabled") {
t.Fatalf("noop toggle output: %q", stdout.String())
}
if strings.Contains(stdout.String(), secretDir) {
t.Fatalf("noop toggle text leaked a secret-shaped manifest path: %q", stdout.String())
}
}
func TestRunPluginsToggleRejectsEmptyID(t *testing.T) {
deps := appDeps{getwd: func() (string, error) { return t.TempDir(), nil }}
for _, subcommand := range []string{"enable", "disable"} {
t.Run(subcommand, func(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"plugins", subcommand, ""}, &stdout, &stderr, deps)
if exitCode != exitUsage {
t.Fatalf("exit=%d want %d stderr=%s", exitCode, exitUsage, stderr.String())
}
if !strings.Contains(stderr.String(), "usage: zero plugins "+subcommand) {
t.Fatalf("stderr = %q", stderr.String())
}
})
}
}
func TestRunPluginsDisableUserIgnoresProject(t *testing.T) {
home := t.TempDir()
cwd := t.TempDir()
t.Setenv("XDG_CONFIG_HOME", home)
t.Setenv("HOME", home)
userManifest := filepath.Join(home, "zero", "plugins", "demo", "plugin.json")
projectManifest := filepath.Join(cwd, ".zero", "plugins", "demo", "plugin.json")
for _, path := range []string{userManifest, projectManifest} {
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte(`{
"schemaVersion": 1,
"id": "zero.demo",
"name": "Demo",
"version": "1.0.0",
"enabled": true
}
`), 0o600); err != nil {
t.Fatal(err)
}
}
deps := appDeps{getwd: func() (string, error) { return cwd, nil }}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"plugins", "disable", "zero.demo", "--user"}, &stdout, &stderr, deps)
if exitCode != exitSuccess {
t.Fatalf("disable --user exit=%d stderr=%s", exitCode, stderr.String())
}
if !strings.Contains(stdout.String(), "is now disabled") {
t.Fatalf("disable --user output: %s", stdout.String())
}
userRaw, err := os.ReadFile(userManifest)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(string(userRaw), `"enabled": false`) {
t.Fatalf("user manifest not disabled: %s", userRaw)
}
projectRaw, err := os.ReadFile(projectManifest)
if err != nil {
t.Fatal(err)
}
if strings.Contains(string(projectRaw), `"enabled": false`) {
t.Fatalf("project manifest should stay enabled: %s", projectRaw)
}
}
// The toggle's usage errors echo the caller's arguments, which can be a
// mistyped credential, so they pass through redaction like its other output.
func TestRunPluginsToggleUsageErrorsAreRedacted(t *testing.T) {
home := t.TempDir()
t.Setenv("XDG_CONFIG_HOME", home)
t.Setenv("HOME", home)
deps := appDeps{getwd: func() (string, error) { return t.TempDir(), nil }}
const secret = "sk-proj-abcdefghijklmnopqrstuvwxyz"
cases := []struct {
name string
args []string
want string
}{
{"unknown flag", []string{"zero.demo", "--token=" + secret}, "unknown plugins"},
{"missing plugin", []string{secret}, "plugin is not installed"},
}
for _, subcommand := range []string{"enable", "disable"} {
for _, tc := range cases {
t.Run(subcommand+" "+tc.name, func(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
args := append([]string{"plugins", subcommand}, tc.args...)
exitCode := runWithDeps(args, &stdout, &stderr, deps)
if exitCode != exitUsage {
t.Fatalf("exit=%d want %d stderr=%s", exitCode, exitUsage, stderr.String())
}
if strings.Contains(stderr.String(), secret) {
t.Fatalf("usage error leaked the argument: %q", stderr.String())
}
if !strings.Contains(stderr.String(), tc.want) || !strings.Contains(stderr.String(), "[REDACTED]") {
t.Fatalf("stderr = %q, want a redacted %q error", stderr.String(), tc.want)
}
})
}
}
}
+1 -1
View File
@@ -319,7 +319,7 @@ Flags:
--command <cmd> Command to run (required)
--name <name> Human-readable hook name
--description <text> Hook description
--matcher <pattern> Tool matcher (beforeTool/afterTool only)
--matcher <pattern> Tool/specialist matcher (beforeTool/afterTool/specialistStart/specialistStop)
--arg <value> Command argument (repeatable)
--user Write to user config instead of the project
--json Print command result as JSON
+32 -3
View File
@@ -741,6 +741,13 @@ func (cfg *mcpWritableConfig) upsertServer(name string, server config.MCPServerC
if err != nil {
return false, err
}
if !mergedServer.Disabled && config.IsDefaultMCPServerDisabled(name) {
// Adding a server under the name of a default that ships disabled means
// the user wants it on; the omitted flag would otherwise keep it off.
if data, err = setMCPServerRawDisabledFalse(data); err != nil {
return false, err
}
}
cfg.serverRaw[name] = data
if _, err := cfg.removeLegacyServer(name); err != nil {
return false, err
@@ -800,11 +807,29 @@ func mergeMCPServerRaw(existing json.RawMessage, server config.MCPServerConfig)
if key == "oauth" && preserveRawOAuth {
continue
}
// `mcp add` replaces the entry: it writes only a disable, never a stale
// explicit false carried over from the old entry (upsertServer adds the
// false a default that ships disabled needs).
if key == "disabled" && string(value) == "false" {
continue
}
merged[key] = value
}
return json.Marshal(merged)
}
func setMCPServerRawDisabledFalse(data json.RawMessage) (json.RawMessage, error) {
var server map[string]json.RawMessage
if err := json.Unmarshal(data, &server); err != nil {
return nil, err
}
if server == nil {
server = map[string]json.RawMessage{}
}
server["disabled"] = json.RawMessage("false")
return json.Marshal(server)
}
func (cfg *mcpWritableConfig) removeServer(name string) (bool, error) {
cfg.ensureRaw()
removed := false
@@ -856,15 +881,19 @@ func (cfg *mcpWritableConfig) setServerDisabled(name string, disabled bool) (boo
server = map[string]json.RawMessage{}
}
current := false
// A default that ships disabled (aimlapi) stays off until an entry says
// otherwise, so an absent flag reads as disabled and enabling it must be
// written out explicitly: dropping the flag would leave it off.
shippedDisabled := config.IsDefaultMCPServerDisabled(name)
current := shippedDisabled
if rawDisabled, ok := server["disabled"]; ok && len(rawDisabled) > 0 && string(rawDisabled) != "null" {
if err := json.Unmarshal(rawDisabled, &current); err != nil {
return false, false, err
}
}
changed := current != disabled
if disabled {
data, err := json.Marshal(true)
if disabled || shippedDisabled {
data, err := json.Marshal(disabled)
if err != nil {
return false, false, err
}
+260
View File
@@ -2,10 +2,18 @@ package cli
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"strings"
"sync"
"testing"
"github.com/Twigpine/zero/internal/aimlapi"
"github.com/Twigpine/zero/internal/config"
)
@@ -44,3 +52,255 @@ func TestRunMCPDisableSeededExaDefault(t *testing.T) {
t.Fatal("expected `mcp disable exa` to turn the seeded default off in the resolved config")
}
}
// aimlapiWire stands in for https://mcp.aimlapi.com/mcp at the transport layer:
// requests keep their real URL, nothing is dialed, and each one is recorded.
type aimlapiWire struct {
mu sync.Mutex
requests []*http.Request
}
func (wire *aimlapiWire) RoundTrip(request *http.Request) (*http.Response, error) {
var body []byte
if request.Body != nil {
body, _ = io.ReadAll(request.Body)
_ = request.Body.Close()
}
wire.mu.Lock()
wire.requests = append(wire.requests, request.Clone(context.Background()))
wire.mu.Unlock()
var message struct {
ID json.RawMessage `json:"id"`
Method string `json:"method"`
}
_ = json.Unmarshal(body, &message)
respond := func(status int, payload string) *http.Response {
return &http.Response{StatusCode: status, Header: http.Header{"Content-Type": {"application/json"}}, Body: io.NopCloser(strings.NewReader(payload)), Request: request}
}
if request.URL.String() != aimlapi.MCPServerURL {
return respond(http.StatusBadGateway, `{}`), nil
}
if len(message.ID) == 0 {
return respond(http.StatusAccepted, ""), nil
}
result := `{"protocolVersion":"2024-11-05","capabilities":{},"serverInfo":{"name":"fake","version":"1"}}`
if message.Method == "tools/list" {
result = `{"tools":[{"name":"tools_search","description":"Search the Toolbox","inputSchema":{"type":"object"}}]}`
}
return respond(http.StatusOK, fmt.Sprintf(`{"jsonrpc":"2.0","id":%s,"result":%s}`, message.ID, result)), nil
}
func (wire *aimlapiWire) recorded() []*http.Request {
wire.mu.Lock()
defer wire.mu.Unlock()
return append([]*http.Request(nil), wire.requests...)
}
// isolateAIMLAPIMCPCommand redirects every root the real resolve path reads to
// a temporary home, sets an AI/ML API key in the environment, and routes the
// default HTTP transport to a recorder. It returns the user config path.
func isolateAIMLAPIMCPCommand(t *testing.T, key string) (string, *aimlapiWire, appDeps) {
t.Helper()
home := t.TempDir()
t.Setenv("HOME", home)
t.Setenv("USERPROFILE", home)
t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, ".config"))
t.Setenv("XDG_STATE_HOME", filepath.Join(home, ".state"))
t.Setenv("XDG_CACHE_HOME", filepath.Join(home, ".cache"))
t.Setenv("APPDATA", filepath.Join(home, "AppData", "Roaming"))
t.Setenv("LOCALAPPDATA", filepath.Join(home, "AppData", "Local"))
t.Setenv("ZERO_CRED_STORAGE", "file")
t.Setenv("ZERO_PROVIDER", "")
t.Setenv("AIMLAPI_INFERENCE_URL", "")
t.Setenv("AIMLAPI_PARTNER_ID", "")
t.Setenv("AIMLAPI_API_KEY", key)
configPath, err := config.DefaultUserConfigPath()
if err != nil {
t.Fatal(err)
}
writeMCPCommandRawConfig(t, configPath, `{}`)
wire := &aimlapiWire{}
previous := http.DefaultTransport
http.DefaultTransport = wire
t.Cleanup(func() { http.DefaultTransport = previous })
workspace := t.TempDir()
return configPath, wire, appDeps{getwd: func() (string, error) { return workspace, nil }}
}
func TestRunMCPEnableAIMLAPIDefaultWritesOnlyTheToggle(t *testing.T) {
configPath, wire, deps := isolateAIMLAPIMCPCommand(t, "aiml-secret-key")
var out, errBuf bytes.Buffer
if code := runWithDeps([]string{"mcp", "disable", "aimlapi", "--json"}, &out, &errBuf, deps); code != exitSuccess {
t.Fatalf("disable exit=%d stderr=%s", code, errBuf.String())
}
if !strings.Contains(out.String(), `"changed": false`) {
t.Fatalf("disabling the off-by-default server should be a no-op: %s", out.String())
}
out.Reset()
errBuf.Reset()
if code := runWithDeps([]string{"mcp", "enable", "aimlapi", "--json"}, &out, &errBuf, deps); code != exitSuccess {
t.Fatalf("enable exit=%d stderr=%s", code, errBuf.String())
}
if !strings.Contains(out.String(), `"changed": true`) {
t.Fatalf("enable should report a change: %s", out.String())
}
data, err := os.ReadFile(configPath)
if err != nil {
t.Fatal(err)
}
var written struct {
MCP struct {
Servers map[string]map[string]any `json:"servers"`
} `json:"mcp"`
}
if err := json.Unmarshal(data, &written); err != nil {
t.Fatal(err)
}
if entry := written.MCP.Servers["aimlapi"]; len(entry) != 1 || entry["disabled"] != false {
t.Fatalf("enable must write only the toggle, got %#v in %s", entry, data)
}
if strings.Contains(string(data), "aiml-secret-key") || strings.Contains(string(data), "Authorization") {
t.Fatalf("config holds credential material: %s", data)
}
cfg, err := config.ResolveMCP(config.ResolveOptions{UserConfigPath: configPath})
if err != nil {
t.Fatal(err)
}
if server := cfg.Servers["aimlapi"]; server.Disabled || config.IsUnconfiguredDefault("aimlapi", server) {
t.Fatalf("enabled aimlapi = %#v, want enabled and user-configured", server)
}
// An unrelated save rewrites the whole file through FileConfig; the
// explicit enable has to survive it or the shipped disable comes back.
if _, err := config.SetTheme(configPath, "dark"); err != nil {
t.Fatal(err)
}
cfg, err = config.ResolveMCP(config.ResolveOptions{UserConfigPath: configPath})
if err != nil {
t.Fatal(err)
}
if cfg.Servers["aimlapi"].Disabled {
data, _ := os.ReadFile(configPath)
t.Fatalf("an unrelated config save turned aimlapi off again: %s", data)
}
if len(wire.recorded()) != 0 {
t.Fatal("enable must not connect")
}
out.Reset()
if code := runWithDeps([]string{"mcp", "disable", "aimlapi"}, &out, &errBuf, deps); code != exitSuccess {
t.Fatalf("disable exit=%d stderr=%s", code, errBuf.String())
}
cfg, err = config.ResolveMCP(config.ResolveOptions{UserConfigPath: configPath})
if err != nil {
t.Fatal(err)
}
if !cfg.Servers["aimlapi"].Disabled {
t.Fatal("disable after enable must turn aimlapi off again")
}
}
func TestRunMCPAddUnderAIMLAPIDefaultNameEnablesIt(t *testing.T) {
configPath, _, deps := isolateAIMLAPIMCPCommand(t, "")
var out, errBuf bytes.Buffer
if code := runWithDeps([]string{"mcp", "add", "aimlapi", "--type", "http", "--url", aimlapi.MCPServerURL, "--auth", "oauth"}, &out, &errBuf, deps); code != exitSuccess {
t.Fatalf("add exit=%d stderr=%s", code, errBuf.String())
}
cfg, err := config.ResolveMCP(config.ResolveOptions{UserConfigPath: configPath})
if err != nil {
t.Fatal(err)
}
if server := cfg.Servers["aimlapi"]; server.Disabled || server.Auth != "oauth" {
t.Fatalf("added aimlapi = %#v, want enabled with oauth", server)
}
if code := runWithDeps([]string{"mcp", "add", "aimlapi", "--type", "http", "--url", aimlapi.MCPServerURL, "--disabled"}, &out, &errBuf, deps); code != exitSuccess {
t.Fatalf("add --disabled exit=%d stderr=%s", code, errBuf.String())
}
if cfg, err = config.ResolveMCP(config.ResolveOptions{UserConfigPath: configPath}); err != nil || !cfg.Servers["aimlapi"].Disabled {
t.Fatalf("add --disabled left aimlapi enabled: %#v, %v", cfg.Servers["aimlapi"], err)
}
}
// Through the real resolve path (no injected resolveMCPConfig): the default is
// listed as disabled, never connects, and once enabled it reaches session
// startup and `zero mcp check`, which connects with the key on the wire.
func TestAIMLAPIMCPDefaultThroughRealResolvePath(t *testing.T) {
_, wire, deps := isolateAIMLAPIMCPCommand(t, "aiml-secret-key")
var out, errBuf bytes.Buffer
if code := runWithDeps([]string{"mcp", "list"}, &out, &errBuf, deps); code != exitSuccess {
t.Fatalf("list exit=%d stderr=%s", code, errBuf.String())
}
if !strings.Contains(out.String(), "aimlapi [http] disabled https://mcp.aimlapi.com/mcp") || strings.Contains(out.String(), "aiml-secret-key") {
t.Fatalf("list output should show the disabled aimlapi default and no key:\n%s", out.String())
}
out.Reset()
if code := runWithDeps([]string{"mcp", "list", "--json"}, &out, &errBuf, deps); code != exitSuccess {
t.Fatalf("list --json exit=%d stderr=%s", code, errBuf.String())
}
if !strings.Contains(out.String(), `"aimlapi"`) || strings.Contains(out.String(), "aiml-secret-key") {
t.Fatalf("list --json = %s", out.String())
}
out.Reset()
errBuf.Reset()
if code := runWithDeps([]string{"mcp", "check", "aimlapi"}, &out, &errBuf, deps); code == exitSuccess || !strings.Contains(errBuf.String(), "is disabled") {
t.Fatalf("check of the disabled default: exit=%d stderr=%s", code, errBuf.String())
}
resolved, err := fillAppDeps(deps).resolveMCPConfig(t.TempDir(), true)
if err != nil {
t.Fatal(err)
}
critical, optional := splitMCPStartupConfig(resolved)
if _, ok := critical.Servers["aimlapi"]; ok {
t.Fatal("the disabled default must not start with a session")
}
if _, ok := optional.Servers["aimlapi"]; ok {
t.Fatal("the disabled default must not start with a session")
}
if len(wire.recorded()) != 0 {
t.Fatalf("the disabled default connected: %d requests", len(wire.recorded()))
}
if code := runWithDeps([]string{"mcp", "enable", "aimlapi"}, &out, &errBuf, deps); code != exitSuccess {
t.Fatalf("enable exit=%d stderr=%s", code, errBuf.String())
}
resolved, err = fillAppDeps(deps).resolveMCPConfig(t.TempDir(), true)
if err != nil {
t.Fatal(err)
}
critical, _ = splitMCPStartupConfig(resolved)
if _, ok := critical.Servers["aimlapi"]; !ok {
t.Fatal("an enabled aimlapi must start with the session as a user-configured server")
}
out.Reset()
errBuf.Reset()
if code := runWithDeps([]string{"mcp", "check", "aimlapi", "--json"}, &out, &errBuf, deps); code != exitSuccess {
t.Fatalf("check exit=%d stdout=%s stderr=%s", code, out.String(), errBuf.String())
}
if !strings.Contains(out.String(), `"status": "ok"`) || !strings.Contains(out.String(), `"toolCount": 1`) {
t.Fatalf("check output = %s", out.String())
}
if strings.Contains(out.String()+errBuf.String(), "aiml-secret-key") {
t.Fatalf("check output leaks the key: %s %s", out.String(), errBuf.String())
}
requests := wire.recorded()
if len(requests) == 0 {
t.Fatal("check did not connect")
}
for _, request := range requests {
if request.URL.String() != aimlapi.MCPServerURL {
t.Fatalf("request went to %s", request.URL)
}
if got := request.Header.Get("Authorization"); got != "Bearer aiml-secret-key" {
t.Fatalf("Authorization = %q, want the AI/ML API key", got)
}
if request.Header.Get(aimlapi.SourceHeaderName) != aimlapi.SourceValue || request.Header.Get(aimlapi.PartnerHeaderName) == "" {
t.Fatalf("attribution headers missing: %v", request.Header)
}
}
}
@@ -0,0 +1,145 @@
package cli
import (
"bytes"
"context"
"io"
"testing"
"github.com/Twigpine/zero/internal/config"
"github.com/Twigpine/zero/internal/hooks"
"github.com/Twigpine/zero/internal/mcp"
"github.com/Twigpine/zero/internal/specialist"
"github.com/Twigpine/zero/internal/tools"
"github.com/Twigpine/zero/internal/tui"
"github.com/Twigpine/zero/internal/zeroruntime"
)
func TestRegisterSpecialistToolsWiresLifecycleHooksIntoTaskOutput(t *testing.T) {
registry := tools.NewRegistry()
runtime, err := registerSpecialistTools(registry, t.TempDir(), 0)
if err != nil {
t.Fatalf("registerSpecialistTools returned error: %v", err)
}
t.Cleanup(func() { closeSpecialistRuntime(io.Discard, runtime) })
if runtime.lifecycleHooks == nil {
t.Fatal("lifecycleHooks is nil after registerSpecialistTools")
}
raw, ok := registry.Get("TaskOutput")
if !ok {
t.Fatal("TaskOutput tool not registered")
}
outputTool, ok := raw.(*specialist.OutputTool)
if !ok {
t.Fatalf("TaskOutput type = %T, want *specialist.OutputTool", raw)
}
if outputTool.LifecycleHooks != runtime.lifecycleHooks {
t.Fatal("TaskOutput LifecycleHooks is not the shared registerSpecialistTools bridge")
}
dispatcher := hooks.NewDispatcher(hooks.DispatcherOptions{})
runtime.attachLifecycleHooks(dispatcher)
if runtime.lifecycleHooks.Dispatch == nil {
t.Fatal("attachLifecycleHooks left Dispatch nil")
}
if outputTool.LifecycleHooks.Dispatch == nil {
t.Fatal("TaskOutput did not observe Dispatch after attachLifecycleHooks")
}
}
// TestRealEntryPointsAttachLifecycleHooks closes the gap the registration-level
// test above cannot see: wiring the bridge when tools register is not enough
// unless the real startup paths actually call attachLifecycleHooks before any
// run. Both surfaces that can launch specialists must do so — the interactive
// TUI and headless exec. attachLifecycleHooks mutates the same
// *specialist.LifecycleHooks the TaskOutput tool holds, so a nil Dispatch on
// that tool after a full startup proves the call never ran and the
// specialistStart/specialistStop hooks would silently never fire.
func TestRealEntryPointsAttachLifecycleHooks(t *testing.T) {
requireTaskOutputDispatchAttached := func(t *testing.T, registry *tools.Registry) {
t.Helper()
if registry == nil {
t.Fatal("startup did not reach the run registry")
}
raw, ok := registry.Get("TaskOutput")
if !ok {
t.Fatal("TaskOutput tool not registered by startup")
}
outputTool, ok := raw.(*specialist.OutputTool)
if !ok {
t.Fatalf("TaskOutput type = %T, want *specialist.OutputTool", raw)
}
if outputTool.LifecycleHooks == nil || outputTool.LifecycleHooks.Dispatch == nil {
t.Fatal("startup left the specialist lifecycle bridge unattached: specialistStart/specialistStop hooks would never fire")
}
}
t.Run("interactive TUI", func(t *testing.T) {
isolateCLIUserState(t)
cwd := t.TempDir()
var registry *tools.Registry
// Empty args route runWithDeps to the interactive TUI path; the injected
// runTUI receives the same registry specialist tools registered into.
exitCode := runWithDeps([]string{}, io.Discard, io.Discard, appDeps{
getwd: func() (string, error) { return cwd, nil },
resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) {
return config.ResolvedConfig{
Provider: config.ProviderProfile{
Name: "p",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "http://127.0.0.1/v1",
Model: "m",
},
}, nil
},
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) { return config.MCPConfig{}, nil },
runTUI: func(_ context.Context, options tui.Options) int {
registry = options.Registry
return exitSuccess
},
})
if exitCode != exitSuccess {
t.Fatalf("runWithDeps exit = %d, want %d", exitCode, exitSuccess)
}
requireTaskOutputDispatchAttached(t, registry)
})
t.Run("headless exec", func(t *testing.T) {
isolateCLIUserState(t)
cwd := t.TempDir()
var registry *tools.Registry
var stdout, stderr bytes.Buffer
// --auto high is what registers specialist tools on the exec surface.
// MCP registration is the one startup chokepoint that sees the run
// registry before attachLifecycleHooks runs, so capture it there and
// assert after the run completes.
exitCode := runWithDeps([]string{
"exec", "--auto", "high", "--cwd", cwd, "--init-session-id", "hook_wireup", "hello",
}, &stdout, &stderr, appDeps{
getwd: func() (string, error) { return cwd, nil },
resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) {
return execResolvedConfig(), nil
},
newProvider: func(config.ProviderProfile) (zeroruntime.Provider, error) {
return echoExecProvider{}, nil
},
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) {
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
"wireup": {Type: "stdio", Command: "wireup-mcp"},
}}, nil
},
newMCPStore: func() (*mcp.PermissionStore, error) { return nil, nil },
registerMCPTools: func(_ context.Context, captured *tools.Registry, _ config.MCPConfig, _ mcp.RegisterOptions) (mcpToolRuntime, error) {
registry = captured
return closeFunc(func() error { return nil }), nil
},
})
if exitCode != exitSuccess {
t.Fatalf("runWithDeps exit = %d, want %d (stderr=%q)", exitCode, exitSuccess, stderr.String())
}
requireTaskOutputDispatchAttached(t, registry)
})
}
+168
View File
@@ -0,0 +1,168 @@
package config
import (
"errors"
"fmt"
"os"
"path/filepath"
"strings"
"unicode"
"github.com/Twigpine/zero/internal/aimlapi"
"github.com/Twigpine/zero/internal/providercatalog"
)
const aimlapiKeyEnv = "AIMLAPI_API_KEY"
var (
// ErrAIMLAPIKeyMissing means no AI/ML API key is configured.
ErrAIMLAPIKeyMissing = errors.New("no AI/ML API key is configured; set up aimlapi.com with /provider or set AIMLAPI_API_KEY")
// ErrAIMLAPIProxyProfile means the aimlapi profile points at a proxy, so its
// key belongs to that proxy and is withheld from aimlapi.com services.
ErrAIMLAPIProxyProfile = errors.New("the aimlapi provider points at a proxy base URL; its key is not sent to aimlapi.com services")
// ErrAIMLAPIKeyInvalid means the key cannot be sent as a single Bearer token.
ErrAIMLAPIKeyInvalid = errors.New("the AI/ML API key contains whitespace, a comma or a control character")
)
// AIMLAPIServiceKey returns the AI/ML API key that may be sent to the
// aimlapi.com Toolbox and MCP server. It is resolved again on every call, so a
// provider change is seen by the next request.
//
// The key comes from the aimlapi provider profile (the active one when it is
// aimlapi, otherwise the first), including a key held in the credential store,
// or from AIMLAPI_API_KEY when no aimlapi profile exists. The key is withheld
// when that profile, any aimlapi profile, or any profile reading
// AIMLAPI_API_KEY points at a base URL other than the canonical
// https://api.aimlapi.com/v1, and when AIMLAPI_INFERENCE_URL is overridden:
// in those cases the credential may belong to a proxy.
func AIMLAPIServiceKey(resolved ResolvedConfig, store APIKeyGetter, getenv func(string) string) (string, error) {
if getenv == nil {
getenv = os.Getenv
}
if override := strings.TrimSpace(getenv("AIMLAPI_INFERENCE_URL")); override != "" && !sameBaseURL(override, aimlapi.CanonicalInferenceBaseURL) {
return "", ErrAIMLAPIProxyProfile
}
for _, profile := range resolved.Providers {
if aimlapiKeyedProfile(profile) && !sameBaseURL(profile.BaseURL, aimlapi.CanonicalInferenceBaseURL) {
return "", ErrAIMLAPIProxyProfile
}
}
if profile, ok := aimlapiServiceProfile(resolved); ok {
if !sameBaseURL(profile.BaseURL, aimlapi.CanonicalInferenceBaseURL) {
return "", ErrAIMLAPIProxyProfile
}
if store != nil && profile.APIKeyStored && strings.TrimSpace(profile.APIKey) == "" && strings.TrimSpace(profile.Name) != "" {
key, found, err := store.Get(strings.TrimSpace(profile.Name))
if err != nil {
return "", fmt.Errorf("read AI/ML API credential store: %w", err)
}
if found {
profile.APIKey = key
}
}
return validAIMLAPIServiceKey(profile.APIKey)
}
return validAIMLAPIServiceKey(getenv(aimlapiKeyEnv))
}
// LoadAIMLAPIServiceKey resolves the user configuration, the environment and
// the user credential store, then applies AIMLAPIServiceKey. Project config is
// not read: provider keys are user-scoped. A config with no active provider
// still resolves its profiles (ErrNoActiveProvider), so the key is looked up in
// that partial config; every other resolve error is returned. Resolve drops a
// non-active profile it cannot normalize, so the configured profiles are
// checked as well: such a profile still withholds the key when it carries it.
func LoadAIMLAPIServiceKey() (string, error) {
userConfigPath, err := DefaultUserConfigPath()
if err != nil {
return "", err
}
existing, err := existingConfigFile(userConfigPath)
if err != nil {
return "", err
}
resolved, err := Resolve(ResolveOptions{UserConfigPath: existing, ExcludeProject: true})
if err != nil && !errors.Is(err, ErrNoActiveProvider) {
return "", err
}
if existing != "" {
fileConfig, err := loadConfigFile(existing)
if err != nil {
return "", err
}
var configured FileConfig
mergeConfig(&configured, fileConfig)
if droppedAIMLAPIProxyProfile(configured.Providers, resolved.Providers) {
return "", ErrAIMLAPIProxyProfile
}
}
// Check proxy eligibility and already-resolved keys before opening a store.
key, keyErr := AIMLAPIServiceKey(resolved, nil, os.Getenv)
if !errors.Is(keyErr, ErrAIMLAPIKeyMissing) {
return key, keyErr
}
profile, ok := aimlapiServiceProfile(resolved)
if !ok || !profile.APIKeyStored {
return key, keyErr
}
store, err := ProviderKeyStoreAt(filepath.Dir(userConfigPath))
if err != nil {
return "", fmt.Errorf("open AI/ML API credential store: %w", err)
}
return AIMLAPIServiceKey(resolved, store, os.Getenv)
}
func isAIMLAPIProfile(profile ProviderProfile) bool {
return strings.EqualFold(providercatalog.NormalizeID(profile.CatalogID), providercatalog.AIMLAPIID)
}
// aimlapiKeyedProfile reports whether profile carries an AI/ML API credential:
// an aimlapi catalog profile, or any profile that reads AIMLAPI_API_KEY.
func aimlapiKeyedProfile(profile ProviderProfile) bool {
return isAIMLAPIProfile(profile) || strings.EqualFold(strings.TrimSpace(profile.APIKeyEnv), aimlapiKeyEnv)
}
// droppedAIMLAPIProxyProfile reports whether a configured profile missing from
// the resolved list (Resolve could not normalize it) carries an AI/ML API
// credential and is not explicitly on the canonical base URL. Its endpoint was
// never resolved, so an empty base URL counts as a proxy.
func droppedAIMLAPIProxyProfile(configured, resolved []ProviderProfile) bool {
kept := make(map[string]bool, len(resolved))
for _, profile := range resolved {
kept[profile.Name] = true
}
for _, profile := range configured {
if kept[strings.TrimSpace(profile.Name)] || !aimlapiKeyedProfile(profile) {
continue
}
if !sameBaseURL(profile.BaseURL, aimlapi.CanonicalInferenceBaseURL) {
return true
}
}
return false
}
func aimlapiServiceProfile(resolved ResolvedConfig) (ProviderProfile, bool) {
if isAIMLAPIProfile(resolved.Provider) {
return resolved.Provider, true
}
for _, profile := range resolved.Providers {
if isAIMLAPIProfile(profile) {
return profile, true
}
}
return ProviderProfile{}, false
}
func validAIMLAPIServiceKey(raw string) (string, error) {
key := strings.TrimSpace(raw)
if key == "" {
return "", ErrAIMLAPIKeyMissing
}
for _, r := range key {
if unicode.IsSpace(r) || r <= ' ' || r == 0x7f || r == ',' || (r >= 0x80 && r <= 0x9f) {
return "", ErrAIMLAPIKeyInvalid
}
}
return key, nil
}
+282
View File
@@ -0,0 +1,282 @@
package config
import (
"errors"
"os"
"path/filepath"
"strings"
"testing"
)
func aimlapiServiceKeyEnv(values map[string]string) func(string) string {
return func(name string) string { return values[name] }
}
func TestAIMLAPIServiceKeyUsesCanonicalProfileKey(t *testing.T) {
resolved := ResolvedConfig{Providers: []ProviderProfile{
{Name: "openai", BaseURL: "https://api.openai.com/v1", APIKey: "sk-openai"},
{Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://api.aimlapi.com/v1/", APIKey: " aiml-key "},
}}
key, err := AIMLAPIServiceKey(resolved, nil, aimlapiServiceKeyEnv(nil))
if err != nil || key != "aiml-key" {
t.Fatalf("AIMLAPIServiceKey = %q, %v; want aiml-key", key, err)
}
}
func TestAIMLAPIServiceKeyReadsCredentialStore(t *testing.T) {
resolved := ResolvedConfig{Provider: ProviderProfile{Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://api.aimlapi.com/v1", APIKeyStored: true}}
store := fakeKeyGetter{keys: map[string]string{"aimlapi": "stored-key"}}
key, err := AIMLAPIServiceKey(resolved, store, aimlapiServiceKeyEnv(nil))
if err != nil || key != "stored-key" {
t.Fatalf("AIMLAPIServiceKey = %q, %v; want stored-key", key, err)
}
}
func TestAIMLAPIServiceKeyPrefersResolvedProfileKeyOverCredentialStore(t *testing.T) {
resolved := ResolvedConfig{Provider: ProviderProfile{
Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://api.aimlapi.com/v1",
APIKey: "env-key", APIKeyStored: true,
}}
store := fakeKeyGetter{keys: map[string]string{"aimlapi": "stored-key"}}
key, err := AIMLAPIServiceKey(resolved, store, aimlapiServiceKeyEnv(map[string]string{"AIMLAPI_API_KEY": "env-key"}))
if err != nil || key != "env-key" {
t.Fatalf("AIMLAPIServiceKey = %q, %v; want env-key from resolved profile, not stored-key", key, err)
}
}
func TestAIMLAPIServiceKeyFallsBackToEnvironmentWithoutProfile(t *testing.T) {
key, err := AIMLAPIServiceKey(ResolvedConfig{}, nil, aimlapiServiceKeyEnv(map[string]string{"AIMLAPI_API_KEY": "env-key"}))
if err != nil || key != "env-key" {
t.Fatalf("AIMLAPIServiceKey = %q, %v; want env-key", key, err)
}
if _, err := AIMLAPIServiceKey(ResolvedConfig{}, nil, aimlapiServiceKeyEnv(nil)); !errors.Is(err, ErrAIMLAPIKeyMissing) {
t.Fatalf("missing key error = %v, want ErrAIMLAPIKeyMissing", err)
}
}
func TestAIMLAPIServiceKeyWithholdsProxyCredentials(t *testing.T) {
env := aimlapiServiceKeyEnv(map[string]string{"AIMLAPI_API_KEY": "proxy-key"})
cases := map[string]ResolvedConfig{
"active aimlapi profile on a proxy": {Provider: ProviderProfile{Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://proxy.example/v1", APIKey: "proxy-key"}},
"a second aimlapi profile on a proxy": {Providers: []ProviderProfile{
{Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://api.aimlapi.com/v1", APIKey: "proxy-key"},
{Name: "aimlapi-proxy", CatalogID: "aimlapi", BaseURL: "https://proxy.example/v1", APIKey: "proxy-key"},
}},
"a custom profile reading AIMLAPI_API_KEY": {Providers: []ProviderProfile{
{Name: "gateway", BaseURL: "https://proxy.example/v1", APIKeyEnv: "AIMLAPI_API_KEY", APIKey: "proxy-key"},
}},
}
for name, resolved := range cases {
key, err := AIMLAPIServiceKey(resolved, nil, env)
if !errors.Is(err, ErrAIMLAPIProxyProfile) || key != "" {
t.Errorf("%s: AIMLAPIServiceKey = %q, %v; want ErrAIMLAPIProxyProfile", name, key, err)
}
}
overridden := aimlapiServiceKeyEnv(map[string]string{"AIMLAPI_API_KEY": "proxy-key", "AIMLAPI_INFERENCE_URL": "https://staging.example/v1"})
if key, err := AIMLAPIServiceKey(ResolvedConfig{}, nil, overridden); !errors.Is(err, ErrAIMLAPIProxyProfile) || key != "" {
t.Fatalf("inference URL override: AIMLAPIServiceKey = %q, %v; want ErrAIMLAPIProxyProfile", key, err)
}
}
func TestAIMLAPIServiceKeyRejectsKeysThatAreNotOneToken(t *testing.T) {
for _, raw := range []string{"first\u00a0second", "first\u2003second", "first,second", "first\nsecond", "first\r\nsecond", "a b", "key\x00tail", "key\u0085tail"} {
resolved := ResolvedConfig{Provider: ProviderProfile{Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://api.aimlapi.com/v1", APIKey: raw}}
if key, err := AIMLAPIServiceKey(resolved, nil, aimlapiServiceKeyEnv(nil)); !errors.Is(err, ErrAIMLAPIKeyInvalid) || key != "" {
t.Errorf("key %q: AIMLAPIServiceKey = %q, %v; want ErrAIMLAPIKeyInvalid", raw, key, err)
}
}
}
func TestLoadAIMLAPIServiceKeyResolvesUserConfig(t *testing.T) {
home := t.TempDir()
t.Setenv("HOME", home)
t.Setenv("APPDATA", filepath.Join(home, "AppData"))
t.Setenv("USERPROFILE", home)
t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, ".config"))
t.Setenv("ZERO_CRED_STORAGE", "file")
t.Setenv("ZERO_PROVIDER", "")
t.Setenv("AIMLAPI_INFERENCE_URL", "")
t.Setenv("AIMLAPI_API_KEY", "env-key")
if key, err := LoadAIMLAPIServiceKey(); err != nil || key != "env-key" {
t.Fatalf("without config: LoadAIMLAPIServiceKey = %q, %v; want env-key", key, err)
}
path, err := DefaultUserConfigPath()
if err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
t.Fatal(err)
}
proxy := `{"activeProvider":"aimlapi","providers":[{"name":"aimlapi","catalogID":"aimlapi","baseURL":"https://proxy.example/v1","model":"m"}]}`
if err := os.WriteFile(path, []byte(proxy), 0o600); err != nil {
t.Fatal(err)
}
if key, err := LoadAIMLAPIServiceKey(); !errors.Is(err, ErrAIMLAPIProxyProfile) || key != "" {
t.Fatalf("proxy profile: LoadAIMLAPIServiceKey = %q, %v; want ErrAIMLAPIProxyProfile", key, err)
}
canonical := `{"activeProvider":"aimlapi","providers":[{"name":"aimlapi","catalogID":"aimlapi"}]}`
if err := os.WriteFile(path, []byte(canonical), 0o600); err != nil {
t.Fatal(err)
}
if key, err := LoadAIMLAPIServiceKey(); err != nil || key != "env-key" {
t.Fatalf("canonical profile: LoadAIMLAPIServiceKey = %q, %v; want env-key", key, err)
}
}
// Several profiles and no activeProvider make Resolve return
// ErrNoActiveProvider with the normalized profiles; the interactive app keeps
// running on that partial config, so the service key must still resolve.
func TestLoadAIMLAPIServiceKeyWithoutActiveProvider(t *testing.T) {
home := t.TempDir()
t.Setenv("HOME", home)
t.Setenv("APPDATA", filepath.Join(home, "AppData"))
t.Setenv("USERPROFILE", home)
t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, ".config"))
t.Setenv("ZERO_CRED_STORAGE", "file")
t.Setenv("ZERO_PROVIDER", "")
t.Setenv("AIMLAPI_INFERENCE_URL", "")
t.Setenv("AIMLAPI_API_KEY", "env-key")
path, err := DefaultUserConfigPath()
if err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
t.Fatal(err)
}
other := `{"name":"other","provider":"openai","model":"m","apiKey":"other-key"}`
for _, tc := range []struct {
name string
active string
profile string
want string
wantErr error
}{
{name: "environment key", profile: `{"name":"second","provider":"openai","model":"m","apiKey":"second-key"}`, want: "env-key"},
{name: "stale active provider", active: "gone", profile: `{"name":"second","provider":"openai","model":"m","apiKey":"second-key"}`, want: "env-key"},
{name: "aimlapi profile key", profile: `{"name":"aimlapi","catalogID":"aimlapi","model":"m","apiKey":"profile-key"}`, want: "profile-key"},
{name: "proxy profile", profile: `{"name":"aimlapi","catalogID":"aimlapi","baseURL":"https://proxy.example/v1","model":"m"}`, wantErr: ErrAIMLAPIProxyProfile},
} {
data := `{"activeProvider":"` + tc.active + `","providers":[` + other + `,` + tc.profile + `]}`
if err := os.WriteFile(path, []byte(data), 0o600); err != nil {
t.Fatal(err)
}
key, err := LoadAIMLAPIServiceKey()
if tc.wantErr != nil {
if !errors.Is(err, tc.wantErr) || key != "" {
t.Fatalf("%s: LoadAIMLAPIServiceKey = %q, %v; want %v", tc.name, key, err, tc.wantErr)
}
continue
}
if err != nil || key != tc.want {
t.Fatalf("%s: LoadAIMLAPIServiceKey = %q, %v; want %q", tc.name, key, err, tc.want)
}
}
}
// Resolve drops a non-active profile it cannot normalize (here an unknown
// catalog entry). Such a profile still shows that AIMLAPI_API_KEY belongs to a
// proxy, so the key is withheld whether or not a provider is active.
func TestLoadAIMLAPIServiceKeyWithholdsForDroppedProxyProfile(t *testing.T) {
home := t.TempDir()
t.Setenv("HOME", home)
t.Setenv("APPDATA", filepath.Join(home, "AppData"))
t.Setenv("USERPROFILE", home)
t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, ".config"))
t.Setenv("ZERO_CRED_STORAGE", "file")
t.Setenv("ZERO_PROVIDER", "")
t.Setenv("AIMLAPI_INFERENCE_URL", "")
t.Setenv("AIMLAPI_API_KEY", "env-key")
path, err := DefaultUserConfigPath()
if err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
t.Fatal(err)
}
canonical := `{"name":"aimlapi","catalogID":"aimlapi","model":"m","apiKey":"profile-key"}`
other := `{"name":"other","provider":"openai","model":"m","apiKey":"other-key"}`
dropped := `{"name":"future","catalogID":"no-such-catalog","baseURL":"https://proxy.example/v1","apiKeyEnv":"AIMLAPI_API_KEY","model":"m"}`
unrelated := `{"name":"future","catalogID":"no-such-catalog","baseURL":"https://proxy.example/v1","apiKeyEnv":"OTHER_API_KEY","model":"m"}`
for _, tc := range []struct {
name string
data string
want string
wantErr error
}{
{name: "active aimlapi profile", data: `{"activeProvider":"aimlapi","providers":[` + canonical + `,` + dropped + `]}`, wantErr: ErrAIMLAPIProxyProfile},
{name: "no active provider", data: `{"providers":[` + other + `,` + dropped + `]}`, wantErr: ErrAIMLAPIProxyProfile},
{name: "dropped profile without the key", data: `{"activeProvider":"aimlapi","providers":[` + canonical + `,` + unrelated + `]}`, want: "profile-key"},
} {
if err := os.WriteFile(path, []byte(tc.data), 0o600); err != nil {
t.Fatal(err)
}
key, err := LoadAIMLAPIServiceKey()
if tc.wantErr != nil {
if !errors.Is(err, tc.wantErr) || key != "" {
t.Fatalf("%s: LoadAIMLAPIServiceKey = %q, %v; want %v", tc.name, key, err, tc.wantErr)
}
continue
}
if err != nil || key != tc.want {
t.Fatalf("%s: LoadAIMLAPIServiceKey = %q, %v; want %q", tc.name, key, err, tc.want)
}
}
}
func TestLoadAIMLAPIServiceKeyStoreOpenFailure(t *testing.T) {
home := t.TempDir()
t.Setenv("HOME", home)
t.Setenv("APPDATA", filepath.Join(home, "AppData"))
t.Setenv("USERPROFILE", home)
t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, ".config"))
t.Setenv("ZERO_PROVIDER", "")
t.Setenv("AIMLAPI_INFERENCE_URL", "")
t.Setenv("AIMLAPI_API_KEY", "env-key")
t.Setenv("ZERO_CRED_STORAGE", "invalid-test-storage")
if key, err := LoadAIMLAPIServiceKey(); err != nil || key != "env-key" {
t.Fatalf("environment key must not require a store: %q, %v", key, err)
}
path, err := DefaultUserConfigPath()
if err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
t.Fatal(err)
}
t.Setenv("AIMLAPI_API_KEY", "")
data := `{"activeProvider":"aimlapi","providers":[{"name":"aimlapi","catalogID":"aimlapi","model":"m","apiKeyStored":true}]}`
if err := os.WriteFile(path, []byte(data), 0o600); err != nil {
t.Fatal(err)
}
if key, err := LoadAIMLAPIServiceKey(); key != "" || err == nil || !strings.Contains(err.Error(), "unknown storage") {
t.Fatalf("stored key must report store failure: %q, %v", key, err)
}
t.Setenv("AIMLAPI_INFERENCE_URL", "https://proxy.example/v1")
if key, err := LoadAIMLAPIServiceKey(); key != "" || !errors.Is(err, ErrAIMLAPIProxyProfile) {
t.Fatalf("proxy rejection must precede store open: %q, %v", key, err)
}
t.Setenv("AIMLAPI_INFERENCE_URL", "")
data = `{"activeProvider":"aimlapi","providers":[{"name":"aimlapi","catalogID":"aimlapi","model":"m","apiKeyStored":true,"apiKey":"inline-key"}]}`
if err := os.WriteFile(path, []byte(data), 0o600); err != nil {
t.Fatal(err)
}
if key, err := LoadAIMLAPIServiceKey(); err != nil || key != "inline-key" {
t.Fatalf("inline key must not require a store: %q, %v", key, err)
}
}
func TestAIMLAPIServiceKeyReportsStoreReadFailure(t *testing.T) {
want := errors.New("test credential read failure")
resolved := ResolvedConfig{Provider: ProviderProfile{Name: "aimlapi", CatalogID: "aimlapi", BaseURL: "https://api.aimlapi.com/v1", APIKeyStored: true}}
if key, err := AIMLAPIServiceKey(resolved, fakeKeyGetter{err: want}, aimlapiServiceKeyEnv(nil)); key != "" || !errors.Is(err, want) {
t.Fatalf("read failure was hidden: %q, %v", key, err)
}
resolved.Provider.BaseURL = "https://proxy.example/v1"
if _, err := AIMLAPIServiceKey(resolved, fakeKeyGetter{err: want}, aimlapiServiceKeyEnv(nil)); !errors.Is(err, ErrAIMLAPIProxyProfile) {
t.Fatalf("proxy check must precede store read: %v", err)
}
}
+51
View File
@@ -0,0 +1,51 @@
package config
import (
"encoding/json"
"os"
"path/filepath"
"testing"
)
func TestEnsureChatGPTDefaultPreservesSavedSelection(t *testing.T) {
for _, saved := range []string{"", "gpt-explicit-saved"} {
name := "new profile"
if saved != "" {
name = "saved selection"
}
t.Run(name, func(t *testing.T) {
path := filepath.Join(t.TempDir(), "config.json")
if saved != "" {
body, err := json.Marshal(FileConfig{ActiveProvider: "codex", Providers: []ProviderProfile{{Name: "codex", CatalogID: "chatgpt", Model: saved}}})
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, body, 0600); err != nil {
t.Fatal(err)
}
}
ensured, err := EnsureCatalogProvider(path, "chatgpt")
if err != nil {
t.Fatal(err)
}
body, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
var cfg FileConfig
if err := json.Unmarshal(body, &cfg); err != nil {
t.Fatal(err)
}
want := "gpt-6.1-sol"
if saved != "" {
want = saved
}
if len(cfg.Providers) != 1 || cfg.Providers[0].Model != want {
t.Fatalf("ChatGPT default/selection = %#v, want model %q", cfg.Providers, want)
}
if saved != "" && (ensured.Name != "codex" || ensured.Created) {
t.Fatalf("saved subscription profile was replaced: %#v", ensured)
}
})
}
}
+27 -5
View File
@@ -3,26 +3,48 @@ package config
import (
"reflect"
"strings"
"github.com/Twigpine/zero/internal/aimlapi"
)
// DefaultMCPServers returns the MCP servers Zero ships ENABLED by default so
// web search and page fetching work out of the box with no setup and no API
// key. They are seeded before user/project config is merged (see ResolveMCP),
// so a user can override any field — for example add an API-key header to lift
// Exa's anonymous rate limit — or disable it entirely with
// DefaultMCPServers returns the MCP servers Zero ships built in. Exa is ENABLED
// by default so web search and page fetching work out of the box with no setup
// and no API key. They are seeded before user/project config is merged (see
// ResolveMCP), so a user can override any field — for example add an API-key
// header to lift Exa's anonymous rate limit — or disable it entirely with
// `zero mcp disable <name>` (which writes `"disabled": true`).
//
// Exa's hosted MCP server works anonymously with rate limits. Users can add an
// Exa API key for higher limits.
//
// The AI/ML API server ships DISABLED: its tool runs and generations are billed
// to the user's AI/ML API balance, so it only connects after
// `zero mcp enable aimlapi` (which writes `"disabled": false`). The disable is
// marked as explicitly set so the lower-trust project layer cannot lift it
// (see mergeMCPServer); only the user and CLI-override scopes can. The key is
// attached per request by the MCP client and is never part of this config.
func DefaultMCPServers() map[string]MCPServerConfig {
return map[string]MCPServerConfig{
"exa": {
Type: "http",
URL: "https://mcp.exa.ai/mcp",
},
"aimlapi": {
Type: "http",
URL: aimlapi.MCPServerURL,
Disabled: true,
disabledSet: true,
},
}
}
// IsDefaultMCPServerDisabled reports whether name is a built-in default that
// ships disabled, so enabling it has to be written out as `"disabled": false`
// rather than by dropping the flag.
func IsDefaultMCPServerDisabled(name string) bool {
return DefaultMCPServers()[strings.TrimSpace(name)].Disabled
}
// IsDefaultMCPServer reports whether name is one of Zero's built-in default MCP
// servers. The config commands use it so a default can be disabled/enabled even
// though it is not written to the user's config file until overridden.
+89
View File
@@ -291,3 +291,92 @@ func TestResolveMCPDisabledRetiredEntryStaysReEnableable(t *testing.T) {
t.Fatal("the disable must still carry to the successor")
}
}
func writeMCPDefaultsTestConfig(t *testing.T, data string) string {
t.Helper()
path := filepath.Join(t.TempDir(), "config.json")
if err := os.WriteFile(path, []byte(data), 0o600); err != nil {
t.Fatal(err)
}
return path
}
func TestResolveMCPSeedsDisabledAIMLAPIDefault(t *testing.T) {
if !IsDefaultMCPServer("aimlapi") || !IsDefaultMCPServerDisabled("aimlapi") || IsDefaultMCPServerDisabled("exa") {
t.Fatal("aimlapi must be a built-in default that ships disabled, exa one that ships enabled")
}
// No user config at all, and a user config that never mentions aimlapi.
for name, options := range map[string]ResolveOptions{
"no config": {},
"unrelated user file": {UserConfigPath: writeMCPDefaultsTestConfig(t, `{"mcp":{"servers":{"exa":{"disabled":true}}}}`)},
"user entry with only headers": {UserConfigPath: writeMCPDefaultsTestConfig(t,
`{"mcp":{"servers":{"aimlapi":{"headers":{"X-Trace":"1"}}}}}`)},
} {
cfg, err := ResolveMCP(options)
if err != nil {
t.Fatalf("%s: ResolveMCP: %v", name, err)
}
server, ok := cfg.Servers["aimlapi"]
if !ok {
t.Fatalf("%s: the aimlapi default must be visible in the resolved config", name)
}
if server.Type != "http" || server.URL != "https://mcp.aimlapi.com/mcp" || !server.Disabled {
t.Fatalf("%s: aimlapi default = %#v, want a disabled http server at the MCP URL", name, server)
}
}
cfg, err := ResolveMCP(ResolveOptions{})
if err != nil {
t.Fatal(err)
}
if !IsUnconfiguredDefault("aimlapi", cfg.Servers["aimlapi"]) {
t.Fatal("the untouched aimlapi default should be reported as unconfigured")
}
if exa := cfg.Servers["exa"]; exa.Disabled || !IsUnconfiguredDefault("exa", exa) {
t.Fatalf("exa default changed: %#v", exa)
}
}
func TestResolveMCPUserEnablesAIMLAPIDefault(t *testing.T) {
path := writeMCPDefaultsTestConfig(t, `{"mcp":{"servers":{"aimlapi":{"disabled":false}}}}`)
cfg, err := ResolveMCP(ResolveOptions{UserConfigPath: path})
if err != nil {
t.Fatalf("ResolveMCP: %v", err)
}
server := cfg.Servers["aimlapi"]
if server.Disabled || server.URL != "https://mcp.aimlapi.com/mcp" || server.Type != "http" {
t.Fatalf("user enable did not turn the default on: %#v", server)
}
if IsUnconfiguredDefault("aimlapi", server) {
t.Fatal("an enabled aimlapi is user-configured, so its failures must be reported")
}
enable := MCPServerConfig{Disabled: false, disabledSet: true}
cfg, err = ResolveMCP(ResolveOptions{Overrides: Overrides{MCP: MCPConfig{Servers: map[string]MCPServerConfig{"aimlapi": enable}}}})
if err != nil {
t.Fatal(err)
}
if cfg.Servers["aimlapi"].Disabled {
t.Fatal("the CLI override scope must be able to enable the aimlapi default")
}
}
func TestResolveMCPProjectCannotEnableAIMLAPIDefault(t *testing.T) {
project := writeMCPDefaultsTestConfig(t, `{"mcp":{"servers":{"aimlapi":{"disabled":false}}}}`)
cfg, err := ResolveMCP(ResolveOptions{ProjectConfigPath: project})
if err != nil {
t.Fatalf("ResolveMCP: %v", err)
}
if !cfg.Servers["aimlapi"].Disabled {
t.Fatal("a project config must not turn on the billed aimlapi default")
}
user := writeMCPDefaultsTestConfig(t, `{"mcp":{"servers":{"aimlapi":{"disabled":false}}}}`)
projectDisable := writeMCPDefaultsTestConfig(t, `{"mcp":{"servers":{"aimlapi":{"disabled":true}}}}`)
cfg, err = ResolveMCP(ResolveOptions{UserConfigPath: user, ProjectConfigPath: projectDisable})
if err != nil {
t.Fatalf("ResolveMCP: %v", err)
}
if cfg.Servers["aimlapi"].Disabled {
t.Fatal("the user's explicit enable wins over the project layer, as for any server")
}
}
+33
View File
@@ -724,6 +724,39 @@ func (cfg LocalControlDriverConfig) MarshalJSON() ([]byte, error) {
return json.Marshal(raw)
}
// MarshalJSON writes "disabled" whenever it was explicitly set, not only when
// true. A read-modify-write of config.json (SetTheme, SetActiveProvider, ...)
// would otherwise drop `"disabled": false`, which turns a default that ships
// disabled (aimlapi) back off and lets project config override a user's
// explicit enable (see mergeMCPServer).
func (server MCPServerConfig) MarshalJSON() ([]byte, error) {
type rawServer struct {
Type string `json:"type,omitempty"`
Command string `json:"command,omitempty"`
Args []string `json:"args,omitempty"`
Env map[string]string `json:"env,omitempty"`
URL string `json:"url,omitempty"`
Headers map[string]string `json:"headers,omitempty"`
Auth string `json:"auth,omitempty"`
OAuth *MCPOAuthConfig `json:"oauth,omitempty"`
Disabled *bool `json:"disabled,omitempty"`
}
raw := rawServer{
Type: server.Type,
Command: server.Command,
Args: server.Args,
Env: server.Env,
URL: server.URL,
Headers: server.Headers,
Auth: server.Auth,
OAuth: server.OAuth,
}
if server.disabledSet || server.Disabled {
raw.Disabled = &server.Disabled
}
return json.Marshal(raw)
}
func (server *MCPServerConfig) UnmarshalJSON(data []byte) error {
type rawServer struct {
Type string `json:"type"`
+2 -2
View File
@@ -779,7 +779,7 @@ func normalizeDefinition(raw any, field string) (Definition, error) {
return Definition{}, err
}
if matcher != "" && !eventSupportsMatcher(event) {
return Definition{}, manifestError{fieldPath: field + ".matcher", message: "matcher is only supported for beforeTool and afterTool hooks."}
return Definition{}, manifestError{fieldPath: field + ".matcher", message: "matcher is only supported for beforeTool, afterTool, specialistStart, and specialistStop hooks."}
}
command, err := requiredString(obj, field+".command")
if err != nil {
@@ -843,7 +843,7 @@ func matchesHookMatcher(matcher string, toolName string) bool {
}
func eventSupportsMatcher(event Event) bool {
return event == EventBeforeTool || event == EventAfterTool
return event == EventBeforeTool || event == EventAfterTool || event == EventSpecialistStart || event == EventSpecialistStop
}
func toDiagnostic(err error, source ConfigSource, path string) Diagnostic {
+41 -1
View File
@@ -180,7 +180,7 @@ func TestLoadConfigPreservesUserDisabledStateWhenProjectOmitsEnabled(t *testing.
}
func TestLoadConfigRejectsMatchersOnLifecycleHooks(t *testing.T) {
for _, event := range []string{"sessionStart", "specialistStart"} {
for _, event := range []string{"sessionStart", "sessionEnd"} {
t.Run(event, func(t *testing.T) {
dir := t.TempDir()
projectConfigPath := filepath.Join(dir, "hooks.json")
@@ -210,6 +210,46 @@ func TestLoadConfigRejectsMatchersOnLifecycleHooks(t *testing.T) {
}
}
func TestLoadConfigAcceptsMatchersOnSpecialistHooks(t *testing.T) {
for _, tc := range []struct {
event Event
id string
}{
{event: EventSpecialistStart, id: "zero.explorer-start"},
{event: EventSpecialistStop, id: "zero.explorer-stop"},
} {
t.Run(string(tc.event), func(t *testing.T) {
dir := t.TempDir()
projectConfigPath := filepath.Join(dir, "hooks.json")
writeHookJSON(t, projectConfigPath, map[string]any{
"hooks": []any{map[string]any{
"id": tc.id,
"event": string(tc.event),
"matcher": "explorer",
"command": "node",
}},
})
result, err := LoadConfig(LoadOptions{
UserConfigPath: filepath.Join(dir, "missing-user-hooks.json"),
ProjectConfigPath: projectConfigPath,
})
if err != nil {
t.Fatalf("LoadConfig returned error: %v", err)
}
if len(result.Config.Hooks) != 1 || result.Config.Hooks[0].Matcher != "explorer" {
t.Fatalf("specialist matcher hook = %#v", result.Config.Hooks)
}
if got := hookIDs(Select(result.Config, SelectInput{Event: tc.event, ToolName: "explorer"})); !reflect.DeepEqual(got, []string{tc.id}) {
t.Fatalf("matched selection = %#v", got)
}
if got := Select(result.Config, SelectInput{Event: tc.event, ToolName: "worker"}); len(got) != 0 {
t.Fatalf("unmatched specialist should be empty, got %#v", got)
}
})
}
}
func TestConfigStorePersistsUpdates(t *testing.T) {
configPath := filepath.Join(t.TempDir(), "hooks.json")
store, err := NewConfigStore(StoreOptions{ConfigPath: configPath})
+119
View File
@@ -0,0 +1,119 @@
package mcp
import (
"fmt"
"net/http"
"net/url"
"strings"
"github.com/Twigpine/zero/internal/aimlapi"
"github.com/Twigpine/zero/internal/config"
)
// The AI/ML API MCP server authenticates with the user's AI/ML API key.
//
// Invariant: the key and the attribution headers are attached only to an
// outgoing request whose own URL is https://mcp.aimlapi.com on the default port
// with no userinfo (exact hostname, see aimlapi.IsAttributedURL), whatever the
// server is named. The key is resolved again for every request, redirect hops
// included (see reapplyAIMLAPIRedirectHeaders), through
// config.LoadAIMLAPIServiceKey, which in the same call decides whether it may
// be sent (no proxy profile, a single token), so a provider change applies to
// the next request and a withheld or invalid key means no Authorization
// header. It lives only in that request's headers: never in Server, its
// Identity (OAuth token binding), config.json, list output or errors (an
// error response to such a request is reported without its body). A
// user-configured Authorization header, OAuth, or a project-touched server
// always goes without it.
var loadAIMLAPIServiceKey = config.LoadAIMLAPIServiceKey
var aimlapiMCPHost = func() string {
parsed, err := url.Parse(aimlapi.MCPServerURL)
if err != nil {
panic(err)
}
return parsed.Hostname()
}()
func isAIMLAPIMCPURL(target *url.URL) bool {
return aimlapi.IsAttributedURL(target) && strings.EqualFold(target.Hostname(), aimlapiMCPHost)
}
// aimlapiAutomaticHeaders are the headers applyAIMLAPIRequestHeaders may add.
var aimlapiAutomaticHeaders = [...]string{"Authorization", aimlapi.SourceHeaderName, aimlapi.PartnerHeaderName}
// applyAIMLAPIRequestHeaders runs after the server's own headers are set.
func applyAIMLAPIRequestHeaders(server Server, request *http.Request) {
if request == nil || !isAIMLAPIMCPURL(request.URL) {
return
}
aimlapi.ApplyAttributionHeaders(request.Header, request.URL)
if request.Header.Get("Authorization") != "" || strings.EqualFold(server.Auth, ServerAuthOAuth) || server.ProjectConfigured {
return
}
key, err := loadAIMLAPIServiceKey()
if err != nil || key == "" {
return
}
request.Header.Set("Authorization", "Bearer "+key)
}
// reapplyAIMLAPIRedirectHeaders gives a redirect hop the automatic headers its
// own URL earns. net/http copies the first request's headers onto every hop
// without calling the header hooks, so without this a key or attribution added
// for the first URL would follow a same-origin hop that fails isAIMLAPIMCPURL
// (userinfo, say), and an eligible hop would reuse the first hop's key instead
// of resolving it again. The hop gets back the server's own values for these
// headers, then the same rule as any other request.
func reapplyAIMLAPIRedirectHeaders(server Server, request *http.Request, first *http.Request) {
if request == nil || first == nil || !isAIMLAPIMCPURL(first.URL) {
return
}
for _, name := range aimlapiAutomaticHeaders {
request.Header.Del(name)
}
for name, value := range server.Headers {
for _, automatic := range aimlapiAutomaticHeaders {
if strings.EqualFold(name, automatic) {
request.Header.Set(name, value)
}
}
}
applyAIMLAPIRequestHeaders(server, request)
}
// sentAIMLAPIServiceKey reports whether request went out carrying the key
// applyAIMLAPIRequestHeaders injects, as opposed to the server's own
// Authorization header or OAuth.
func sentAIMLAPIServiceKey(server Server, request *http.Request) bool {
if request == nil || !isAIMLAPIMCPURL(request.URL) || strings.EqualFold(server.Auth, ServerAuthOAuth) || server.ProjectConfigured {
return false
}
return request.Header.Get("Authorization") != "" && !hasUserAuthorizationHeader(server)
}
func hasUserAuthorizationHeader(server Server) bool {
for name, value := range server.Headers {
if strings.EqualFold(name, "Authorization") && value != "" {
return true
}
}
return false
}
// aimlapiUnauthorizedHint explains how to authenticate when the AI/ML API MCP
// server rejects a request without OAuth. It names no key material.
func aimlapiUnauthorizedHint(server Server, status int) string {
if status != http.StatusUnauthorized || strings.EqualFold(server.Auth, ServerAuthOAuth) {
return ""
}
target, err := url.Parse(server.URL)
if err != nil || !isAIMLAPIMCPURL(target) {
return ""
}
oauthHint := fmt.Sprintf("To sign in instead, set \"auth\": \"oauth\" on this server and run `zero mcp oauth login %s`", server.Name)
if server.ProjectConfigured || hasUserAuthorizationHeader(server) {
return " (" + oauthHint + ")"
}
return fmt.Sprintf(" (set up the aimlapi.com provider with `zero providers setup aimlapi` or set AIMLAPI_API_KEY; automatic injection is also disabled for proxy profiles, OAuth, project config, and non-empty Authorization headers — see docs/EXTENDING.md#built-in-servers. To sign in instead, set \"auth\": \"oauth\" on this server and run `zero mcp oauth login %s`)", server.Name)
}
+537
View File
@@ -0,0 +1,537 @@
package mcp
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"strings"
"sync"
"testing"
"github.com/Twigpine/zero/internal/aimlapi"
"github.com/Twigpine/zero/internal/config"
)
// aimlapiRecorder stands in for mcp.aimlapi.com: it records every request it
// is handed (the URL stays https://mcp.aimlapi.com/mcp, nothing is dialed) and
// answers JSON-RPC, or 401 with the server's OAuth challenge when unauthorized.
type aimlapiRecorder struct {
mu sync.Mutex
headers []http.Header
urls []string
unauthorized bool
// redirectFirstTo, when set, answers the first request with a 307 to it.
redirectFirstTo string
// failStatus, when set, answers every request with it and failBody.
failStatus int
failBody string
}
func (recorder *aimlapiRecorder) RoundTrip(request *http.Request) (*http.Response, error) {
var body []byte
if request.Body != nil {
body, _ = io.ReadAll(request.Body)
_ = request.Body.Close()
}
recorder.mu.Lock()
recorder.headers = append(recorder.headers, request.Header.Clone())
recorder.urls = append(recorder.urls, request.URL.String())
unauthorized := recorder.unauthorized
redirect := recorder.redirectFirstTo != "" && len(recorder.headers) == 1
recorder.mu.Unlock()
if redirect {
response := fakeHTTPResponse(request, http.StatusTemporaryRedirect, "")
response.Header.Set("Location", recorder.redirectFirstTo)
return response, nil
}
if recorder.failStatus != 0 {
return fakeHTTPResponse(request, recorder.failStatus, recorder.failBody), nil
}
if unauthorized {
response := fakeHTTPResponse(request, http.StatusUnauthorized, `{"error":"invalid_token"}`)
response.Header.Set("WWW-Authenticate", `Bearer resource_metadata="https://mcp.aimlapi.com/.well-known/oauth-protected-resource/mcp"`)
return response, nil
}
var message struct {
ID json.RawMessage `json:"id"`
Method string `json:"method"`
}
_ = json.Unmarshal(body, &message)
if len(message.ID) == 0 {
return fakeHTTPResponse(request, http.StatusAccepted, ""), nil
}
result := `{}`
switch message.Method {
case "initialize":
result = `{"protocolVersion":"2024-11-05","capabilities":{},"serverInfo":{"name":"fake","version":"1"}}`
case "tools/list":
result = `{"tools":[{"name":"tools_search","inputSchema":{"type":"object"}}]}`
}
return fakeHTTPResponse(request, http.StatusOK, fmt.Sprintf(`{"jsonrpc":"2.0","id":%s,"result":%s}`, message.ID, result)), nil
}
func (recorder *aimlapiRecorder) recorded() []http.Header {
recorder.mu.Lock()
defer recorder.mu.Unlock()
return append([]http.Header(nil), recorder.headers...)
}
func fakeHTTPResponse(request *http.Request, status int, body string) *http.Response {
header := http.Header{}
if body != "" {
header.Set("Content-Type", "application/json")
}
return &http.Response{
StatusCode: status,
Header: header,
Body: io.NopCloser(bytes.NewBufferString(body)),
Request: request,
}
}
// isolateAIMLAPIKeyLookup points the real key resolver at an empty temporary
// home and returns the user config path it reads.
func isolateAIMLAPIKeyLookup(t *testing.T, envKey string) string {
t.Helper()
home := t.TempDir()
t.Setenv("HOME", home)
t.Setenv("USERPROFILE", home)
t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, ".config"))
t.Setenv("APPDATA", filepath.Join(home, "AppData"))
t.Setenv("ZERO_CRED_STORAGE", "file")
t.Setenv("ZERO_PROVIDER", "")
t.Setenv("AIMLAPI_INFERENCE_URL", "")
t.Setenv("AIMLAPI_PARTNER_ID", "")
t.Setenv("AIMLAPI_API_KEY", envKey)
path, err := config.DefaultUserConfigPath()
if err != nil {
t.Fatal(err)
}
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
t.Fatal(err)
}
return path
}
func writeAIMLAPIProviderConfig(t *testing.T, path string, baseURL string, apiKey string) {
t.Helper()
profile := map[string]any{"name": "aimlapi", "catalogID": "aimlapi", "model": "m", "apiKey": apiKey}
if baseURL != "" {
profile["baseURL"] = baseURL
}
data, err := json.Marshal(map[string]any{"activeProvider": "aimlapi", "providers": []any{profile}})
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, data, 0o600); err != nil {
t.Fatal(err)
}
}
func newRecordedNetworkClient(server Server, recorder *aimlapiRecorder) *networkClient {
return &networkClient{server: server, client: mcpHTTPClient(server, recorder), nextID: 1}
}
func aimlapiTestServer(name string, rawURL string) Server {
return Server{Name: name, Type: ServerTypeHTTP, URL: rawURL}
}
func assertAIMLAPIAttribution(t *testing.T, header http.Header) {
t.Helper()
if got := header.Get(aimlapi.SourceHeaderName); got != aimlapi.SourceValue {
t.Fatalf("%s = %q, want %q", aimlapi.SourceHeaderName, got, aimlapi.SourceValue)
}
if got, want := header.Get(aimlapi.PartnerHeaderName), aimlapi.ResolvePartnerID(""); got == "" || got != want {
t.Fatalf("%s = %q, want %q", aimlapi.PartnerHeaderName, got, want)
}
}
func TestAIMLAPIMCPRequestsCarryServiceKeyAndAttribution(t *testing.T) {
isolateAIMLAPIKeyLookup(t, "aiml-secret")
// A hand-added server on the exact host is treated like the built-in one.
for _, name := range []string{"aimlapi", "media"} {
recorder := &aimlapiRecorder{}
client := newRecordedNetworkClient(aimlapiTestServer(name, aimlapi.MCPServerURL), recorder)
if err := client.initialize(context.Background()); err != nil {
t.Fatalf("%s: initialize: %v", name, err)
}
if _, err := client.ListTools(context.Background()); err != nil {
t.Fatalf("%s: ListTools: %v", name, err)
}
requests := recorder.recorded()
if len(requests) != 3 {
t.Fatalf("%s: recorded %d requests, want initialize, initialized and tools/list", name, len(requests))
}
for _, header := range requests {
if got := header.Get("Authorization"); got != "Bearer aiml-secret" {
t.Fatalf("%s: Authorization = %q, want the AI/ML API key", name, got)
}
assertAIMLAPIAttribution(t, header)
}
}
}
func TestAIMLAPIMCPHeadersOnlyForExactHost(t *testing.T) {
isolateAIMLAPIKeyLookup(t, "aiml-secret")
for _, rawURL := range []string{
"https://mcp.aimlapi.com.example.net/mcp",
"http://mcp.aimlapi.com/mcp",
"https://mcp.aimlapi.com:8443/mcp",
"https://[email protected]/mcp",
"https://api.aimlapi.com/mcp",
"https://evilmcp.aimlapi.com/mcp",
"https://mcp.exa.ai/mcp",
} {
recorder := &aimlapiRecorder{}
server := aimlapiTestServer("other", rawURL)
server.Headers = map[string]string{"X-Custom": "kept"}
client := newRecordedNetworkClient(server, recorder)
if err := client.initialize(context.Background()); err != nil {
t.Fatalf("%s: initialize: %v", rawURL, err)
}
for _, header := range recorder.recorded() {
if strings.Contains(rawURL, "user@") && strings.HasPrefix(header.Get("Authorization"), "Basic ") {
// net/http turns URL userinfo into Basic auth on its own.
header.Del("Authorization")
}
// Byte-for-byte what any other MCP server gets.
want := http.Header{
"X-Custom": {"kept"},
"Content-Type": {"application/json"},
"Accept": {"application/json, text/event-stream"},
}
if fmt.Sprint(header) != fmt.Sprint(want) {
t.Fatalf("%s: headers = %v, want %v", rawURL, header, want)
}
}
}
// Port 443 spelled out and a different host case are still the exact host.
for _, rawURL := range []string{"https://mcp.aimlapi.com:443/mcp", "https://MCP.AIMLAPI.COM/mcp"} {
recorder := &aimlapiRecorder{}
client := newRecordedNetworkClient(aimlapiTestServer("aimlapi", rawURL), recorder)
if err := client.initialize(context.Background()); err != nil {
t.Fatalf("%s: initialize: %v", rawURL, err)
}
header := recorder.recorded()[0]
if header.Get("Authorization") != "Bearer aiml-secret" {
t.Fatalf("%s: Authorization = %q, want the key", rawURL, header.Get("Authorization"))
}
assertAIMLAPIAttribution(t, header)
}
}
// The key and its eligibility are resolved for every request, so switching the
// aimlapi profile to a proxy keeps the proxy's key off mcp.aimlapi.com on the
// very next request of an already-connected client and on new connections.
func TestAIMLAPIMCPKeyIsResolvedPerRequestAcrossProviderSwitch(t *testing.T) {
path := isolateAIMLAPIKeyLookup(t, "")
writeAIMLAPIProviderConfig(t, path, "", "canonical-key")
recorder := &aimlapiRecorder{}
server := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL)
client := newRecordedNetworkClient(server, recorder)
if err := client.initialize(context.Background()); err != nil {
t.Fatalf("initialize: %v", err)
}
if got := recorder.recorded()[0].Get("Authorization"); got != "Bearer canonical-key" {
t.Fatalf("canonical profile: Authorization = %q, want Bearer canonical-key", got)
}
writeAIMLAPIProviderConfig(t, path, "https://proxy.example/v1", "proxy-key")
if _, err := client.ListTools(context.Background()); err != nil {
t.Fatalf("ListTools after switch: %v", err)
}
fresh := newRecordedNetworkClient(server, recorder)
if err := fresh.initialize(context.Background()); err != nil {
t.Fatalf("initialize after switch: %v", err)
}
requests := recorder.recorded()
for _, header := range requests[2:] {
if got := header.Get("Authorization"); got != "" {
t.Fatalf("proxy profile: Authorization = %q, want none", got)
}
assertAIMLAPIAttribution(t, header)
}
for _, header := range requests {
if strings.Contains(fmt.Sprint(header), "proxy-key") {
t.Fatalf("proxy key reached mcp.aimlapi.com: %v", header)
}
}
}
func TestAIMLAPIMCPUnusableKeySendsNoAuthorization(t *testing.T) {
cases := map[string]string{
"multi-line key pool": "first-key\nsecond-key",
"comma key pool": "first-key,second-key",
"no key": "",
}
for name, apiKey := range cases {
t.Run(name, func(t *testing.T) {
path := isolateAIMLAPIKeyLookup(t, "")
writeAIMLAPIProviderConfig(t, path, "", apiKey)
recorder := &aimlapiRecorder{unauthorized: true}
client := newRecordedNetworkClient(aimlapiTestServer("aimlapi", aimlapi.MCPServerURL), recorder)
err := client.initialize(context.Background())
if err == nil {
t.Fatal("initialize succeeded against a 401")
}
requests := recorder.recorded()
if len(requests) != 1 {
t.Fatalf("recorded %d requests, want one intact request", len(requests))
}
header := requests[0]
if got := header.Get("Authorization"); got != "" {
t.Fatalf("Authorization = %q, want none", got)
}
if header.Get("Content-Type") != "application/json" {
t.Fatalf("request headers broken: %v", header)
}
assertAIMLAPIAttribution(t, header)
if !strings.Contains(err.Error(), "returned HTTP 401") || !strings.Contains(err.Error(), "zero mcp oauth login aimlapi") || !strings.Contains(err.Error(), "AIMLAPI_API_KEY") {
t.Fatalf("401 error lacks the authentication hint: %v", err)
}
if strings.Contains(err.Error(), "first-key") {
t.Fatalf("error leaks key material: %v", err)
}
})
}
}
func TestAIMLAPIMCPUnauthorizedHintOnlyForAIMLAPIWithoutOAuth(t *testing.T) {
if hint := aimlapiUnauthorizedHint(aimlapiTestServer("docs", "https://mcp.example.com/mcp"), http.StatusUnauthorized); hint != "" {
t.Fatalf("hint for another server = %q", hint)
}
oauth := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL)
oauth.Auth = ServerAuthOAuth
if hint := aimlapiUnauthorizedHint(oauth, http.StatusUnauthorized); hint != "" {
t.Fatalf("hint for an OAuth server = %q", hint)
}
if hint := aimlapiUnauthorizedHint(aimlapiTestServer("aimlapi", aimlapi.MCPServerURL), http.StatusForbidden); hint != "" {
t.Fatalf("hint for 403 = %q", hint)
}
}
func TestAIMLAPIMCPKeyDefersToUserAuthorizationOAuthAndProjectConfig(t *testing.T) {
isolateAIMLAPIKeyLookup(t, "aiml-secret")
userHeader := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL)
userHeader.Headers = map[string]string{"authorization": "Bearer user-token"}
oauth := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL)
oauth.Auth = ServerAuthOAuth
project := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL)
project.ProjectConfigured = true
cases := []struct {
name string
server Server
want string
}{
{"user Authorization header", userHeader, "Bearer user-token"},
{"OAuth server", oauth, ""},
{"project-configured server", project, ""},
}
for _, tc := range cases {
recorder := &aimlapiRecorder{}
client := newRecordedNetworkClient(tc.server, recorder)
if err := client.initialize(context.Background()); err != nil {
t.Fatalf("%s: initialize: %v", tc.name, err)
}
header := recorder.recorded()[0]
if got := header.Get("Authorization"); got != tc.want {
t.Fatalf("%s: Authorization = %q, want %q", tc.name, got, tc.want)
}
assertAIMLAPIAttribution(t, header)
}
}
func TestAIMLAPIMCPSSERequestsCarryKeyAndAttribution(t *testing.T) {
isolateAIMLAPIKeyLookup(t, "aiml-secret")
server := Server{Name: "aimlapi", Type: ServerTypeSSE, URL: aimlapi.MCPServerURL}
client := &remoteSSEClient{server: server}
apply := map[string]func(*http.Request){
"stream": client.applyStreamHeaders,
"post": client.applyPostHeaders,
}
for name, applyHeaders := range apply {
request, err := http.NewRequest(http.MethodPost, aimlapi.MCPServerURL, nil)
if err != nil {
t.Fatal(err)
}
applyHeaders(request)
if got := request.Header.Get("Authorization"); got != "Bearer aiml-secret" {
t.Fatalf("%s: Authorization = %q, want the key", name, got)
}
assertAIMLAPIAttribution(t, request.Header)
other, err := http.NewRequest(http.MethodPost, "https://mcp.aimlapi.com.example.net/mcp", nil)
if err != nil {
t.Fatal(err)
}
applyHeaders(other)
if other.Header.Get("Authorization") != "" || other.Header.Get(aimlapi.SourceHeaderName) != "" || other.Header.Get(aimlapi.PartnerHeaderName) != "" {
t.Fatalf("%s: lookalike host got AI/ML API headers: %v", name, other.Header)
}
}
}
func TestAIMLAPIMCPKeyStaysOutOfServerConfigAndIdentity(t *testing.T) {
enabled := config.MCPConfig{Servers: map[string]config.MCPServerConfig{
"aimlapi": {Type: "http", URL: aimlapi.MCPServerURL},
}}
isolateAIMLAPIKeyLookup(t, "")
withoutKey, err := NormalizeConfig(enabled)
if err != nil {
t.Fatal(err)
}
isolateAIMLAPIKeyLookup(t, "aiml-secret")
withKey, err := NormalizeConfig(enabled)
if err != nil {
t.Fatal(err)
}
if len(withKey) != 1 || withKey[0].Identity != withoutKey[0].Identity {
t.Fatalf("server identity depends on the key: %#v vs %#v", withKey, withoutKey)
}
if strings.Contains(fmt.Sprintf("%#v", withKey), "aiml-secret") {
t.Fatalf("normalized server carries the key: %#v", withKey)
}
// The shipped default is off, so it is not even normalized for a session.
defaults, err := NormalizeConfig(config.MCPConfig{Servers: config.DefaultMCPServers()})
if err != nil {
t.Fatal(err)
}
for _, server := range defaults {
if server.Name == "aimlapi" {
t.Fatalf("disabled aimlapi default was normalized for connection: %#v", server)
}
}
}
func TestAIMLAPIMCPUnauthorizedHintRespectsCredentialEligibility(t *testing.T) {
for _, tc := range []struct {
name string
headers map[string]string
project bool
wantSetup bool
}{
{"default", nil, false, true},
{"project", nil, true, false},
{"custom bearer", map[string]string{"Authorization": "Bearer manual"}, false, false},
{"mixed case header", map[string]string{"aUtHoRiZaTiOn": "Bearer manual"}, false, false},
{"empty header", map[string]string{"authorization": ""}, false, true},
} {
t.Run(tc.name, func(t *testing.T) {
server := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL)
server.Headers = tc.headers
server.ProjectConfigured = tc.project
hint := aimlapiUnauthorizedHint(server, http.StatusUnauthorized)
if strings.Contains(hint, "zero providers setup aimlapi") != tc.wantSetup || !strings.Contains(hint, "zero mcp oauth login aimlapi") {
t.Fatalf("incorrect hint: %s", hint)
}
})
}
}
// net/http copies the first request's headers onto every redirect hop without
// calling the header hooks, so the redirect guard has to apply the per-request
// rule again: a same-origin hop that fails the exact-URL rule (userinfo) gets
// neither the key nor attribution, an eligible hop gets a freshly resolved key,
// and a user's own Authorization header still follows as before.
func TestAIMLAPIMCPRedirectHopsReapplyAutomaticHeaders(t *testing.T) {
isolateAIMLAPIKeyLookup(t, "")
userHeader := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL)
userHeader.Headers = map[string]string{"authorization": "Bearer user-token"}
for _, tc := range []struct {
name string
server Server
location string
wantFirst string
wantHop string
wantAttribution bool
}{
{"userinfo hop", aimlapiTestServer("aimlapi", aimlapi.MCPServerURL), "https://[email protected]/mcp", "Bearer aiml-first", "", false},
{"eligible hop", aimlapiTestServer("aimlapi", aimlapi.MCPServerURL), "https://mcp.aimlapi.com/mcp/", "Bearer aiml-first", "Bearer aiml-second", true},
{"user header hop", userHeader, "https://[email protected]/mcp", "Bearer user-token", "Bearer user-token", false},
} {
t.Run(tc.name, func(t *testing.T) {
keys := []string{"aiml-first", "aiml-second"}
previous := loadAIMLAPIServiceKey
loadAIMLAPIServiceKey = func() (string, error) {
key := keys[0]
if len(keys) > 1 {
keys = keys[1:]
}
return key, nil
}
t.Cleanup(func() { loadAIMLAPIServiceKey = previous })
recorder := &aimlapiRecorder{redirectFirstTo: tc.location}
client := newRecordedNetworkClient(tc.server, recorder)
if err := client.initialize(context.Background()); err != nil {
t.Fatalf("initialize: %v", err)
}
requests := recorder.recorded()
if len(requests) < 2 || recorder.urls[1] != tc.location {
t.Fatalf("redirect was not followed: %v", recorder.urls)
}
if got := requests[0].Get("Authorization"); got != tc.wantFirst {
t.Fatalf("first hop Authorization = %q, want %q", got, tc.wantFirst)
}
hop := requests[1]
// With no key, net/http fills Basic auth from the hop's own userinfo;
// only the automatic key must be absent.
if got := hop.Get("Authorization"); (tc.wantHop == "" && strings.Contains(got, "aiml-")) || (tc.wantHop != "" && got != tc.wantHop) {
t.Fatalf("redirect hop Authorization = %q, want %q", got, tc.wantHop)
}
if tc.wantAttribution {
assertAIMLAPIAttribution(t, hop)
} else if hop.Get(aimlapi.SourceHeaderName) != "" || hop.Get(aimlapi.PartnerHeaderName) != "" {
t.Fatalf("redirect hop kept attribution headers: %v", hop)
}
})
}
}
// An error body can echo the key Zero injected; the body is left out of the
// error rather than matched, so a key cut by the 1024-byte read limit cannot
// leak its prefix either. Requests without the automatic key keep the detail.
func TestAIMLAPIMCPStatusErrorOmitsBodyWhenKeyWasInjected(t *testing.T) {
const key = "sk-echo-0123456789abcdef"
isolateAIMLAPIKeyLookup(t, key)
for _, tc := range []struct {
name string
body string
}{
{"echoed key", `{"error":"invalid credential ` + key + `"}`},
{"key across the read limit", strings.Repeat("x", 1020) + key},
} {
t.Run(tc.name, func(t *testing.T) {
recorder := &aimlapiRecorder{failStatus: http.StatusPaymentRequired, failBody: tc.body}
err := newRecordedNetworkClient(aimlapiTestServer("aimlapi", aimlapi.MCPServerURL), recorder).initialize(context.Background())
if err == nil {
t.Fatal("initialize succeeded against a failing server")
}
if got := recorder.recorded()[0].Get("Authorization"); got != "Bearer "+key {
t.Fatalf("request did not carry the injected key: %q", got)
}
if strings.Contains(err.Error(), key[:4]) || !strings.Contains(err.Error(), "HTTP 402") {
t.Fatalf("error = %q, want HTTP 402 without the echoed key", err)
}
})
}
userHeader := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL)
userHeader.Headers = map[string]string{"Authorization": "Bearer user-token"}
recorder := &aimlapiRecorder{failStatus: http.StatusPaymentRequired, failBody: `{"error":"insufficient balance"}`}
err := newRecordedNetworkClient(userHeader, recorder).initialize(context.Background())
if err == nil || !strings.Contains(err.Error(), "insufficient balance") {
t.Fatalf("error without the automatic key = %v, want the response detail", err)
}
}
+11 -2
View File
@@ -427,6 +427,7 @@ func (client *networkClient) applyHeaders(request *http.Request) {
if client.sessionID != "" {
request.Header.Set("Mcp-Session-Id", client.sessionID)
}
applyAIMLAPIRequestHeaders(client.server, request)
}
func (client *remoteSSEClient) applyStreamHeaders(request *http.Request) {
@@ -434,6 +435,7 @@ func (client *remoteSSEClient) applyStreamHeaders(request *http.Request) {
request.Header.Set(key, value)
}
request.Header.Set("Accept", "text/event-stream")
applyAIMLAPIRequestHeaders(client.server, request)
}
func (client *remoteSSEClient) applyPostHeaders(request *http.Request) {
@@ -442,6 +444,7 @@ func (client *remoteSSEClient) applyPostHeaders(request *http.Request) {
}
request.Header.Set("Content-Type", "application/json")
request.Header.Set("Accept", "application/json")
applyAIMLAPIRequestHeaders(client.server, request)
}
func closeResponseBody(errp *error, server Server, body io.Closer) {
@@ -607,12 +610,18 @@ func (client *remoteSSEClient) failPending(err error) {
}
func httpStatusError(server Server, response *http.Response) error {
hint := aimlapiUnauthorizedHint(server, response.StatusCode)
if sentAIMLAPIServiceKey(server, response.Request) {
// The body is server text that may echo the key Zero injected, whole or
// cut by the read limit, so it is left out rather than matched.
return fmt.Errorf("MCP %s server %s returned HTTP %d (response body omitted: the request carried your AI/ML API key)%s", server.Type, server.Name, response.StatusCode, hint)
}
body, _ := io.ReadAll(io.LimitReader(response.Body, 1024))
detail := strings.TrimSpace(string(body))
if detail == "" {
return fmt.Errorf("MCP %s server %s returned HTTP %d", server.Type, server.Name, response.StatusCode)
return fmt.Errorf("MCP %s server %s returned HTTP %d%s", server.Type, server.Name, response.StatusCode, hint)
}
return fmt.Errorf("MCP %s server %s returned HTTP %d: %s", server.Type, server.Name, response.StatusCode, detail)
return fmt.Errorf("MCP %s server %s returned HTTP %d: %s%s", server.Type, server.Name, response.StatusCode, detail, hint)
}
func decodeSSERPCMessage(reader io.Reader) (rpcMessage, error) {
+1
View File
@@ -28,6 +28,7 @@ func checkMCPRedirect(server Server) func(*http.Request, []*http.Request) error
if !sameMCPOrigin(via[0].URL, request.URL) {
return fmt.Errorf("MCP %s server %s refused cross-origin redirect to %s", server.Type, server.Name, mcpOrigin(request.URL))
}
reapplyAIMLAPIRedirectHeaders(server, request, via[0])
return nil
}
}
+1 -1
View File
@@ -374,7 +374,7 @@ func (tool skillTool) Run(_ context.Context, args map[string]any) tools.Result {
available := make([]string, 0, len(merged))
for _, skill := range merged {
if skill.Name == name {
return tools.Result{Status: tools.StatusOK, Output: skill.Content}
return tools.Result{Status: tools.StatusOK, Output: skills.ToolOutput(skill)}
}
available = append(available, skill.Name)
}
+25
View File
@@ -673,3 +673,28 @@ func TestExpandPluginRootPathRelative(t *testing.T) {
}
}
}
// The multi-root skill tool is the one the CLI registers, so it has to point at
// a skill's bundled files the same way the core tool does (#584).
func TestNewSkillToolPointsAtTheSkillsBundledFiles(t *testing.T) {
home := t.TempDir()
t.Setenv("HOME", home)
t.Setenv("USERPROFILE", home)
dir := t.TempDir()
writeTestSkill(t, dir, "pdf", "Run scripts/fill.py.")
script := filepath.Join(dir, "pdf", "scripts", "fill.py")
if err := os.MkdirAll(filepath.Dir(script), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(script, []byte("print('fill')\n"), 0o755); err != nil {
t.Fatal(err)
}
got := NewSkillTool(dir, nil).Run(context.Background(), map[string]any{"name": "pdf"})
if got.Status != tools.StatusOK {
t.Fatalf("Status = %s (output: %s)", got.Status, got.Output)
}
if !strings.Contains(got.Output, "- scripts/fill.py") || !strings.Contains(got.Output, filepath.Join(dir, "pdf")) {
t.Fatalf("output does not point at the bundled script:\n%s", got.Output)
}
}
+552
View File
@@ -0,0 +1,552 @@
package plugins
import (
"bytes"
"encoding/json"
"errors"
"fmt"
"io"
"io/fs"
"os"
"path/filepath"
"strings"
"github.com/Twigpine/zero/internal/installtxn"
"github.com/Twigpine/zero/internal/pathjail"
)
// SetEnabledResult describes a successful enable/disable of one plugin manifest.
type SetEnabledResult struct {
ID string `json:"id"`
Enabled bool `json:"enabled"`
Changed bool `json:"changed"`
Source Source `json:"source"`
ManifestPath string `json:"manifestPath"`
}
// maxEnabledRelockAttempts bounds how many times SetEnabledByID re-resolves the
// plugin after it moves to another install root while waiting for the lock.
const maxEnabledRelockAttempts = 4
// FindByID returns the loaded plugin with the given id after normal discovery
// and precedence (project overrides user). A missing id returns ok=false.
func FindByID(options LoadOptions, id string) (LoadedPlugin, bool, error) {
id = strings.TrimSpace(id)
if id == "" {
return LoadedPlugin{}, false, errors.New("plugin id is required")
}
result, err := Load(options)
if err != nil {
return LoadedPlugin{}, false, err
}
for _, plugin := range result.Plugins {
if plugin.ID == id {
return plugin, true, nil
}
}
return LoadedPlugin{}, false, nil
}
// SetEnabledByID finds a plugin by id (respecting LoadOptions) and toggles the
// top-level "enabled" field of its plugin.json.
//
// The read-modify-write is serialized with plugin installs and removals: it
// takes the same cross-process installtxn lock they take on the plugin's
// install root, runs the same interrupted-transaction recovery they run under
// it, re-resolves the plugin under that lock, and publishes the
// manifest with an atomic rename while still holding it, so a concurrent
// install can never be interleaved between the read and the write. Only the
// bytes of the "enabled" value change, so field order, formatting, and numeric
// literals in the rest of the manifest survive the toggle.
func SetEnabledByID(options LoadOptions, id string, enabled bool) (SetEnabledResult, error) {
for attempt := 0; ; attempt++ {
plugin, ok, err := FindByID(options, id)
if err != nil {
return SetEnabledResult{}, err
}
if !ok {
return SetEnabledResult{}, fmt.Errorf("%w: %q", ErrNotInstalled, id)
}
// Take the install root lock before the final lookup and read so an
// install or remove cannot swap the tree between them.
unlock, err := installtxn.Lock(plugin.Root)
if err != nil {
return SetEnabledResult{}, err
}
// Put back anything an earlier install or removal was killed mid-commit,
// and abort if a transaction could not be resolved: editing a tree that is
// still owed a restore loses the edit when the next install or removal
// puts the backup back over it. See installtxn.Recover.
if err := installtxn.Recover(plugin.Root, reconcileInterrupted(plugin.Root)); err != nil {
unlock()
return SetEnabledResult{}, err
}
final, ok, err := FindByID(options, id)
if err != nil {
unlock()
return SetEnabledResult{}, err
}
if !ok {
unlock()
return SetEnabledResult{}, fmt.Errorf("%w: %q", ErrNotInstalled, id)
}
if final.Root != plugin.Root {
// Precedence moved to another root while this lock was pending;
// retry against the root that wins now so the edit still serializes
// with that root's installs.
unlock()
if attempt >= maxEnabledRelockAttempts {
return SetEnabledResult{}, fmt.Errorf("plugin %q kept changing roots while waiting for the install lock", id)
}
continue
}
changed, err := setEnabledLocked(final, enabled)
unlock()
if err != nil {
return SetEnabledResult{}, err
}
return SetEnabledResult{
ID: final.ID,
Enabled: enabled,
Changed: changed,
Source: final.Source,
ManifestPath: final.ManifestPath,
}, nil
}
}
// setEnabledLocked performs the manifest read-modify-write while the caller
// holds the installtxn lock for the plugin's install root.
//
// The lock only serializes cooperating installers; it does not stop the plugin
// directory from being swapped after discovery found it. So the whole update
// works on one directory, pinned by openPluginDir: the manifest is read and
// published, and the tree hashed for the lockfile, through that handle, never
// by re-resolving the path.
func setEnabledLocked(plugin LoadedPlugin, enabled bool) (bool, error) {
pinned, err := openPluginDir(plugin)
if err != nil {
return false, err
}
defer pinned.Close()
return setEnabledPinned(plugin, pinned, enabled)
}
// setEnabledPinned is setEnabledLocked on an already pinned plugin directory.
//
// The edit lands in the pinned directory even if the path now names another
// one, so before refreshing the lockfile or reporting any result it confirms
// that the directory entry still is the pinned directory. Otherwise the
// command would report a state that the installed plugin does not have.
func setEnabledPinned(plugin LoadedPlugin, pinned *pinnedPluginDir, enabled bool) (bool, error) {
manifestName := filepath.Base(plugin.ManifestPath)
raw, perm, err := readManifest(pinned.dir, manifestName)
if err != nil {
return false, err
}
data, changed, err := toggleEnabledManifest(raw, enabled)
if err != nil {
return false, err
}
if !changed {
if err := pinned.verifyLive(); err != nil {
return false, err
}
return false, nil
}
// Read the lockfile before the edit so the refresh below only runs when the
// installed tree still matched its recorded hash; a toggle must never mask
// pre-existing drift.
lock, entry, hashClean := cleanLockEntry(plugin, pinned.dir)
if err := publishManifest(pinned.dir, manifestName, data, perm); err != nil {
return false, err
}
if !hashClean {
if err := pinned.verifyLive(); err != nil {
return false, err
}
return true, nil
}
after, err := hashRootedTree(pinned.dir)
if err != nil {
return true, fmt.Errorf("plugin manifest updated but %s hash refresh failed: %w", LockFileName, err)
}
if err := pinned.verifyLive(); err != nil {
return false, err
}
entry.Hash = after
lock[plugin.ID] = entry
if err := writeLock(plugin.Root, lock); err != nil {
return true, fmt.Errorf("plugin manifest updated but %s hash refresh failed: %w", LockFileName, err)
}
return true, nil
}
// errPluginDirReplaced reports that the plugin's directory entry no longer
// names the directory the toggle pinned and edited.
var errPluginDirReplaced = errors.New("plugin directory was replaced while it was being updated; the plugin installed there now was not changed")
// pinnedPluginDir is the plugin directory discovery found, held open for the
// length of one toggle.
type pinnedPluginDir struct {
root *os.Root // the plugin's install root
name string // the plugin directory's entry name in root
path string // the plugin directory's path, for messages
dir *os.Root // the plugin directory itself
}
// Close releases both handles.
func (p *pinnedPluginDir) Close() {
_ = p.dir.Close()
_ = p.root.Close()
}
// verifyLive confirms that the plugin's entry in its install root is still a
// real directory, not a symlink, junction, or other reparse point, and is the
// very directory p.dir holds open.
func (p *pinnedPluginDir) verifyLive() error {
opened, err := p.dir.Stat(".")
if err != nil {
return fmt.Errorf("inspect plugin directory: %w", err)
}
entry, err := p.root.Lstat(p.name)
if errors.Is(err, fs.ErrNotExist) {
return fmt.Errorf("%w: %s", errPluginDirReplaced, p.path)
}
if err != nil {
return fmt.Errorf("inspect plugin directory: %w", err)
}
if entry.Mode()&(os.ModeSymlink|os.ModeIrregular) != 0 {
return fmt.Errorf("%w: plugin directory %s", pathjail.ErrReparse, p.path)
}
if !os.SameFile(entry, opened) {
return fmt.Errorf("%w: %s", errPluginDirReplaced, p.path)
}
return nil
}
// openPluginDir pins the plugin directory discovery found, confined to the
// plugin's install root.
//
// A link swapped in before the open is refused here. One swapped in after it
// cannot redirect anything, because every later read, write, and hash is
// relative to the returned handle rather than to a path; verifyLive then
// detects that the path no longer names the edited directory.
func openPluginDir(plugin LoadedPlugin) (*pinnedPluginDir, error) {
root, name, err := pathjail.Open(plugin.Root, plugin.PluginDir)
if err != nil {
return nil, fmt.Errorf("open plugin directory: %w", err)
}
if name == "." || filepath.Base(name) != name {
_ = root.Close()
return nil, fmt.Errorf("%w: plugin directory %s is not an entry of %s", pathjail.ErrEscapes, plugin.PluginDir, plugin.Root)
}
if err := pathjail.RefuseReparse(root, name); err != nil {
_ = root.Close()
return nil, fmt.Errorf("open plugin directory: %w", err)
}
dir, err := root.OpenRoot(name)
if err != nil {
_ = root.Close()
return nil, fmt.Errorf("open plugin directory: %w", err)
}
pinned := &pinnedPluginDir{root: root, name: name, path: plugin.PluginDir, dir: dir}
if err := pinned.verifyLive(); err != nil {
pinned.Close()
return nil, err
}
return pinned, nil
}
// readManifest reads name from dir and returns its bytes and permission bits.
// Failing on a stat error keeps the publish from widening the permissions.
func readManifest(dir *os.Root, name string) ([]byte, os.FileMode, error) {
file, err := dir.Open(name)
if err != nil {
return nil, 0, fmt.Errorf("read plugin manifest: %w", err)
}
defer file.Close()
info, err := file.Stat()
if err != nil {
return nil, 0, fmt.Errorf("inspect plugin manifest: %w", err)
}
raw, err := io.ReadAll(file)
if err != nil {
return nil, 0, fmt.Errorf("read plugin manifest: %w", err)
}
return raw, info.Mode().Perm(), nil
}
// publishManifest replaces name in dir with data by renaming a complete
// temporary sibling over it, so a concurrent reader sees either the old
// manifest or the new one, never a partial write.
func publishManifest(dir *os.Root, name string, data []byte, perm os.FileMode) error {
temp, tempName, err := pathjail.CreateTemp(dir, ".", ".zero-manifest", "")
if err != nil {
return fmt.Errorf("write plugin manifest: %w", err)
}
defer func() { _ = dir.Remove(tempName) }()
if err := temp.Chmod(perm); err != nil {
_ = temp.Close()
return fmt.Errorf("write plugin manifest: %w", err)
}
if _, err := temp.Write(data); err != nil {
_ = temp.Close()
return fmt.Errorf("write plugin manifest: %w", err)
}
if err := temp.Sync(); err != nil {
_ = temp.Close()
return fmt.Errorf("write plugin manifest: %w", err)
}
if err := temp.Close(); err != nil {
return fmt.Errorf("write plugin manifest: %w", err)
}
if err := dir.Rename(tempName, name); err != nil {
return fmt.Errorf("write plugin manifest: %w", err)
}
return nil
}
// cleanLockEntry reports whether the pinned plugin directory dir still matches
// the lockfile entry recorded for the plugin and returns both for a refresh.
func cleanLockEntry(plugin LoadedPlugin, dir *os.Root) (map[string]LockEntry, LockEntry, bool) {
lock, err := ReadLock(plugin.Root)
if err != nil {
return nil, LockEntry{}, false
}
entry, ok := lock[plugin.ID]
if !ok {
return nil, LockEntry{}, false
}
current, err := hashRootedTree(dir)
if err != nil || current != entry.Hash {
return nil, LockEntry{}, false
}
return lock, entry, true
}
// toggleEnabledManifest returns raw with the top-level "enabled" member set to
// enabled. A missing or null value counts as enabled, matching the loader's
// optionalBool fallback. changed reports whether the result differs from raw.
func toggleEnabledManifest(raw []byte, enabled bool) (out []byte, changed bool, err error) {
dec := json.NewDecoder(bytes.NewReader(raw))
openTok, err := dec.Token()
if err != nil {
return nil, false, fmt.Errorf("parse plugin manifest: %w", err)
}
if delim, ok := openTok.(json.Delim); !ok || delim != '{' {
return nil, false, errors.New("plugin manifest must be a JSON object")
}
openEnd := dec.InputOffset()
current := true
lastInvalid := false
var valueStart, valueEnd int64 = -1, -1
for dec.More() {
keyTok, err := dec.Token()
if err != nil {
return nil, false, fmt.Errorf("parse plugin manifest: %w", err)
}
key, ok := keyTok.(string)
if !ok {
return nil, false, errors.New("plugin manifest must be a JSON object")
}
keyEnd := dec.InputOffset()
if key != "enabled" {
if err := skipManifestValue(dec); err != nil {
return nil, false, fmt.Errorf("parse plugin manifest: %w", err)
}
continue
}
valueTok, err := dec.Token()
if err != nil {
return nil, false, fmt.Errorf("parse plugin manifest: %w", err)
}
var literal string
switch value := valueTok.(type) {
case bool:
current = value
literal = "false"
if value {
literal = "true"
}
case nil:
current = true
literal = "null"
default:
// Only the last occurrence of a duplicate key counts, as in the
// loader's json.Unmarshal, so an earlier non-boolean value that a
// later one overwrites is not an error. Remember it and decide after
// the last member.
if err := skipManifestComposite(dec, valueTok); err != nil {
return nil, false, fmt.Errorf("parse plugin manifest: %w", err)
}
lastInvalid = true
continue
}
lastInvalid = false
start, err := enabledValueOffset(raw, keyEnd)
if err != nil {
return nil, false, err
}
if start+int64(len(literal)) > int64(len(raw)) ||
!bytes.Equal(raw[start:start+int64(len(literal))], []byte(literal)) {
return nil, false, errors.New("parse plugin manifest: cannot locate the enabled value")
}
// A duplicate key keeps the last occurrence, matching json.Unmarshal.
valueStart, valueEnd = start, start+int64(len(literal))
}
closingTok, err := dec.Token()
if err != nil {
return nil, false, fmt.Errorf("parse plugin manifest: %w", err)
}
if delim, ok := closingTok.(json.Delim); !ok || delim != '}' {
return nil, false, errors.New("plugin manifest must be a JSON object")
}
closingEnd := dec.InputOffset()
if _, err := dec.Token(); err != io.EOF {
return nil, false, errors.New("parse plugin manifest: unexpected data after the top-level object")
}
if lastInvalid {
return nil, false, ManifestError{FieldPath: "enabled", Message: "Expected a boolean."}
}
if current == enabled {
return nil, false, nil
}
newLiteral := "false"
if enabled {
newLiteral = "true"
}
if valueStart >= 0 {
out := make([]byte, 0, len(raw)-int(valueEnd-valueStart)+len(newLiteral))
out = append(out, raw[:valueStart]...)
out = append(out, newLiteral...)
out = append(out, raw[valueEnd:]...)
return out, true, nil
}
return insertEnabledMember(raw, openEnd, closingEnd-1, newLiteral), true, nil
}
// enabledValueOffset returns the byte offset of the value that follows the
// colon after an "enabled" key token that ended at keyEnd.
func enabledValueOffset(raw []byte, keyEnd int64) (int64, error) {
i := keyEnd
for i < int64(len(raw)) && isJSONSpace(raw[i]) {
i++
}
if i >= int64(len(raw)) || raw[i] != ':' {
return 0, errors.New("parse plugin manifest: missing colon after key")
}
i++
for i < int64(len(raw)) && isJSONSpace(raw[i]) {
i++
}
return i, nil
}
// insertEnabledMember adds a top-level "enabled" member to a manifest that has
// none, mirroring the document's own whitespace style where it can. openEnd is
// the byte offset just after the opening brace and bracePos is the byte index
// of the closing brace.
func insertEnabledMember(raw []byte, openEnd, bracePos int64, literal string) []byte {
insertPos := bracePos
for insertPos > openEnd && isJSONSpace(raw[insertPos-1]) {
insertPos--
}
newline := "\n"
if bytes.Contains(raw, []byte("\r\n")) {
newline = "\r\n"
}
first := openEnd
for first < bracePos && isJSONSpace(raw[first]) {
first++
}
empty := first >= bracePos
multiline := bytes.IndexByte(raw[openEnd:bracePos], '\n') >= 0
member := `"enabled": ` + literal
var insert string
switch {
case empty && !multiline:
insert = member
case empty && multiline:
insert = newline + closingLineIndent(raw, openEnd, bracePos) + " " + member
case !empty && !multiline:
insert = "," + member
default:
insert = "," + newline + memberLineIndent(raw, openEnd, first) + member
}
out := make([]byte, 0, len(raw)+len(insert))
out = append(out, raw[:insertPos]...)
out = append(out, insert...)
out = append(out, raw[insertPos:]...)
return out
}
// memberLineIndent returns the indentation of the line the first member of a
// JSON object starts on. first is the byte index of that member's first token.
func memberLineIndent(raw []byte, openEnd, first int64) string {
lineStart := first
for lineStart > openEnd && raw[lineStart-1] != '\n' {
lineStart--
}
end := lineStart
for end < first && (raw[end] == ' ' || raw[end] == '\t') {
end++
}
return string(raw[lineStart:end])
}
// closingLineIndent returns the indentation before an object's closing brace.
func closingLineIndent(raw []byte, openEnd, bracePos int64) string {
start := bracePos
for start > openEnd && (raw[start-1] == ' ' || raw[start-1] == '\t') {
start--
}
return string(raw[start:bracePos])
}
// skipManifestValue consumes the next JSON value from dec.
func skipManifestValue(dec *json.Decoder) error {
tok, err := dec.Token()
if err != nil {
return err
}
return skipManifestComposite(dec, tok)
}
// skipManifestComposite consumes the rest of the value whose first token tok
// was just read from dec. A scalar is already complete; an object or array is
// consumed through its matching closing delimiter.
func skipManifestComposite(dec *json.Decoder, tok json.Token) error {
delim, ok := tok.(json.Delim)
if !ok || (delim != '{' && delim != '[') {
return nil
}
depth := 1
for depth > 0 {
tok, err := dec.Token()
if err != nil {
return err
}
if inner, ok := tok.(json.Delim); ok {
switch inner {
case '{', '[':
depth++
case '}', ']':
depth--
}
}
}
return nil
}
func isJSONSpace(b byte) bool {
return b == ' ' || b == '\t' || b == '\n' || b == '\r'
}
+999
View File
@@ -0,0 +1,999 @@
package plugins
import (
"errors"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"time"
"github.com/Twigpine/zero/internal/installtxn"
"github.com/Twigpine/zero/internal/pathjail"
)
func enableTestOptions(root string) LoadOptions {
return LoadOptions{Roots: []Root{{Source: SourceUser, Path: root}}}
}
func writeRawManifest(t *testing.T, root, dir, raw string) string {
t.Helper()
pluginDir := filepath.Join(root, dir)
if err := os.MkdirAll(pluginDir, 0o700); err != nil {
t.Fatal(err)
}
manifestPath := filepath.Join(pluginDir, "plugin.json")
if err := os.WriteFile(manifestPath, []byte(raw), 0o600); err != nil {
t.Fatal(err)
}
return manifestPath
}
func demoManifest(enabled bool) string {
state := "false"
if enabled {
state = "true"
}
return `{
"schemaVersion": 1,
"id": "zero.demo",
"name": "Demo",
"version": "1.0.0",
"enabled": ` + state + `
}
`
}
func TestSetEnabledByIDPreservesManifestBytes(t *testing.T) {
root := t.TempDir()
raw := `{
"schemaVersion": 1,
"id": "zero.demo",
"name": "Demo",
"version": "1.0.0",
"limits": {
"maxBytes": 9007199254740993
},
"description": "keep me",
"enabled": true,
"tools": []
}
`
manifestPath := writeRawManifest(t, root, "demo", raw)
result, err := SetEnabledByID(enableTestOptions(root), "zero.demo", false)
if err != nil {
t.Fatalf("SetEnabledByID(false): %v", err)
}
if !result.Changed || result.Enabled {
t.Fatalf("unexpected result: %#v", result)
}
disabled, err := os.ReadFile(manifestPath)
if err != nil {
t.Fatal(err)
}
want := strings.Replace(raw, `"enabled": true`, `"enabled": false`, 1)
if string(disabled) != want {
t.Fatalf("toggle changed more than the enabled value:\n got %q\nwant %q", disabled, want)
}
if !strings.Contains(string(disabled), "9007199254740993") {
t.Fatal("large integer literal was rounded by the toggle")
}
plugin, ok, err := FindByID(enableTestOptions(root), "zero.demo")
if err != nil || !ok {
t.Fatalf("FindByID after disable: ok=%v err=%v", ok, err)
}
if plugin.Enabled {
t.Fatal("loader still reports the plugin enabled after disable")
}
result, err = SetEnabledByID(enableTestOptions(root), "zero.demo", true)
if err != nil {
t.Fatalf("SetEnabledByID(true): %v", err)
}
if !result.Changed || !result.Enabled {
t.Fatalf("unexpected enable result: %#v", result)
}
enabled, err := os.ReadFile(manifestPath)
if err != nil {
t.Fatal(err)
}
if string(enabled) != raw {
t.Fatalf("round trip did not restore the original bytes:\n got %q\nwant %q", enabled, raw)
}
}
func TestSetEnabledByIDTogglesManifestAndIsIdempotent(t *testing.T) {
root := t.TempDir()
writePluginManifest(t, filepath.Join(root, "demo"), map[string]any{
"schemaVersion": 1,
"id": "zero.demo",
"name": "Demo",
"version": "1.0.0",
"description": "keep me",
"tools": []any{},
})
manifestPath := filepath.Join(root, "demo", "plugin.json")
result, err := SetEnabledByID(enableTestOptions(root), "zero.demo", false)
if err != nil {
t.Fatalf("SetEnabledByID(false): %v", err)
}
if !result.Changed {
t.Fatal("expected first disable to change the manifest")
}
raw, err := os.ReadFile(manifestPath)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(string(raw), `"enabled": false`) {
t.Fatalf("manifest not disabled: %s", raw)
}
if !strings.Contains(string(raw), `"keep me"`) {
t.Fatalf("description was not preserved: %s", raw)
}
result, err = SetEnabledByID(enableTestOptions(root), "zero.demo", false)
if err != nil {
t.Fatalf("SetEnabledByID(false) again: %v", err)
}
if result.Changed {
t.Fatal("expected second disable to be a no-op")
}
result, err = SetEnabledByID(enableTestOptions(root), "zero.demo", true)
if err != nil {
t.Fatalf("SetEnabledByID(true): %v", err)
}
if !result.Changed {
t.Fatal("expected enable to change the manifest")
}
}
func TestSetEnabledByIDInsertsMissingField(t *testing.T) {
t.Run("pretty manifest", func(t *testing.T) {
root := t.TempDir()
raw := `{
"schemaVersion": 1,
"id": "zero.demo",
"name": "Demo",
"version": "1.0.0"
}
`
manifestPath := writeRawManifest(t, root, "demo", raw)
if _, err := SetEnabledByID(enableTestOptions(root), "zero.demo", false); err != nil {
t.Fatalf("SetEnabledByID(false): %v", err)
}
disabled, err := os.ReadFile(manifestPath)
if err != nil {
t.Fatal(err)
}
wantDisabled := `{
"schemaVersion": 1,
"id": "zero.demo",
"name": "Demo",
"version": "1.0.0",
"enabled": false
}
`
if string(disabled) != wantDisabled {
t.Fatalf("after insert:\n got %q\nwant %q", disabled, wantDisabled)
}
if _, err := SetEnabledByID(enableTestOptions(root), "zero.demo", true); err != nil {
t.Fatalf("SetEnabledByID(true): %v", err)
}
enabled, err := os.ReadFile(manifestPath)
if err != nil {
t.Fatal(err)
}
wantEnabled := strings.Replace(wantDisabled, `"enabled": false`, `"enabled": true`, 1)
if string(enabled) != wantEnabled {
t.Fatalf("after replace:\n got %q\nwant %q", enabled, wantEnabled)
}
})
t.Run("compact manifest", func(t *testing.T) {
root := t.TempDir()
raw := `{"schemaVersion":1,"id":"zero.demo","name":"Demo","version":"1.0.0"}`
manifestPath := writeRawManifest(t, root, "demo", raw)
if _, err := SetEnabledByID(enableTestOptions(root), "zero.demo", false); err != nil {
t.Fatalf("SetEnabledByID(false): %v", err)
}
disabled, err := os.ReadFile(manifestPath)
if err != nil {
t.Fatal(err)
}
want := `{"schemaVersion":1,"id":"zero.demo","name":"Demo","version":"1.0.0","enabled": false}`
if string(disabled) != want {
t.Fatalf("compact insert:\n got %q\nwant %q", disabled, want)
}
})
t.Run("null value", func(t *testing.T) {
root := t.TempDir()
raw := `{
"schemaVersion": 1,
"id": "zero.demo",
"name": "Demo",
"version": "1.0.0",
"enabled": null
}
`
manifestPath := writeRawManifest(t, root, "demo", raw)
result, err := SetEnabledByID(enableTestOptions(root), "zero.demo", true)
if err != nil {
t.Fatalf("SetEnabledByID(true): %v", err)
}
if result.Changed {
t.Fatal("null counts as enabled, so enabling must be a no-op")
}
unchanged, err := os.ReadFile(manifestPath)
if err != nil {
t.Fatal(err)
}
if string(unchanged) != raw {
t.Fatalf("no-op toggle rewrote the manifest: %q", unchanged)
}
if _, err := SetEnabledByID(enableTestOptions(root), "zero.demo", false); err != nil {
t.Fatalf("SetEnabledByID(false): %v", err)
}
disabled, err := os.ReadFile(manifestPath)
if err != nil {
t.Fatal(err)
}
want := strings.Replace(raw, `"enabled": null`, `"enabled": false`, 1)
if string(disabled) != want {
t.Fatalf("null replacement:\n got %q\nwant %q", disabled, want)
}
})
}
func TestToggleEnabledManifestFormattingEdges(t *testing.T) {
cases := []struct {
name string
in string
enabled bool
want string
changed bool
}{
{"empty object", `{}`, false, `{"enabled": false}`, true},
{"empty object with space", `{ }`, false, `{"enabled": false }`, true},
{"empty multiline", "{\n}", false, "{\n \"enabled\": false\n}", true},
{"compact members", `{"a":1}`, false, `{"a":1,"enabled": false}`, true},
{"crlf members", "{\r\n \"a\": 1\r\n}", false, "{\r\n \"a\": 1,\r\n \"enabled\": false\r\n}", true},
{"nested enabled untouched", `{"nested":{"enabled":true},"a":2}`, false, `{"nested":{"enabled":true},"a":2,"enabled": false}`, true},
{"replace with true", `{"enabled": false}`, true, `{"enabled": true}`, true},
{"replace with false", `{"enabled":true}`, false, `{"enabled":false}`, true},
{"null to false", `{"enabled": null}`, false, `{"enabled": false}`, true},
{"null to true is a no-op", `{"enabled": null}`, true, `{"enabled": null}`, false},
{"already enabled is a no-op", `{"enabled": true}`, true, `{"enabled": true}`, false},
// The loader keeps the last duplicate, so an overwritten non-boolean
// value is not an error and only the last value is rewritten.
{"overwritten string", `{"enabled":"obsolete","enabled":true}`, false, `{"enabled":"obsolete","enabled":false}`, true},
{"overwritten composite", `{"enabled":{"a":[1,{"enabled":2}]},"enabled":null}`, false, `{"enabled":{"a":[1,{"enabled":2}]},"enabled":false}`, true},
{"overwritten number is a no-op", `{"enabled":1,"enabled":false}`, false, `{"enabled":1,"enabled":false}`, false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
out, changed, err := toggleEnabledManifest([]byte(tc.in), tc.enabled)
if err != nil {
t.Fatalf("toggleEnabledManifest: %v", err)
}
if changed != tc.changed {
t.Fatalf("changed = %v, want %v", changed, tc.changed)
}
if !tc.changed {
if out != nil {
t.Fatalf("no-op returned bytes: %q", out)
}
return
}
if string(out) != tc.want {
t.Fatalf("got %q, want %q", out, tc.want)
}
})
}
}
func TestToggleEnabledManifestRejectsMalformedInput(t *testing.T) {
for _, raw := range []string{`{"enabled": "yes"}`, `{"enabled": 1}`, `{"enabled": []}`, `{"enabled": {}}`, `{"enabled": true, "enabled": "yes"}`} {
t.Run("value "+raw, func(t *testing.T) {
_, _, err := toggleEnabledManifest([]byte(raw), false)
if err == nil {
t.Fatal("expected error")
}
var manifestErr ManifestError
if !errors.As(err, &manifestErr) {
t.Fatalf("error = %v, want ManifestError", err)
}
if !strings.Contains(err.Error(), "enabled") {
t.Fatalf("error = %v, want it to name the enabled field", err)
}
})
}
for _, raw := range []string{``, `{`, `[]`, `null`, `"str"`, `{} {}`, `{}{}`} {
t.Run("document "+raw, func(t *testing.T) {
if _, _, err := toggleEnabledManifest([]byte(raw), false); err == nil {
t.Fatalf("toggleEnabledManifest(%q) succeeded, want error", raw)
}
})
}
}
// TestSetEnabledByIDHonorsTheLastDuplicateEnabledValue keeps the toggle in step
// with the loader on a manifest whose earlier "enabled" value is overwritten:
// the loader discovers the plugin from the last value, so the toggle must be
// able to change it too.
func TestSetEnabledByIDHonorsTheLastDuplicateEnabledValue(t *testing.T) {
root := t.TempDir()
writeRawManifest(t, root, "demo", `{
"schemaVersion": 1,
"id": "zero.demo",
"name": "Demo",
"version": "1.0.0",
"enabled": "obsolete",
"enabled": true
}
`)
loadedEnabled := func() bool {
t.Helper()
plugin, ok, err := FindByID(enableTestOptions(root), "zero.demo")
if err != nil || !ok {
t.Fatalf("FindByID: ok=%v err=%v", ok, err)
}
return plugin.Enabled
}
if !loadedEnabled() {
t.Fatal("the loader should discover the plugin as enabled from the last value")
}
for _, enabled := range []bool{false, true} {
result, err := SetEnabledByID(enableTestOptions(root), "zero.demo", enabled)
if err != nil {
t.Fatalf("SetEnabledByID(%v): %v", enabled, err)
}
if !result.Changed {
t.Fatalf("SetEnabledByID(%v) reported no change", enabled)
}
if got := loadedEnabled(); got != enabled {
t.Fatalf("after SetEnabledByID(%v) the loader sees enabled=%v", enabled, got)
}
}
}
// TestSetEnabledRefusesASwappedPluginDirectory swaps the plugin directory that
// discovery found for a link between the lookup and the edit. The install lock
// does not cover that window, because it only serializes cooperating
// installers. The edit must fail without following the link: neither the
// link's target nor the moved-away original may change.
func TestSetEnabledRefusesASwappedPluginDirectory(t *testing.T) {
cases := []struct {
name string
// link returns the symlink target to plant and the manifest it leads to.
link func(t *testing.T, root string) (target string, manifest string)
}{
{"to a directory outside the plugins root", func(t *testing.T, root string) (string, string) {
manifest := writeRawManifest(t, t.TempDir(), "elsewhere", demoManifest(true))
return filepath.Dir(manifest), manifest
}},
{"to another plugin in the same root", func(t *testing.T, root string) (string, string) {
// A relative link stays inside the root, so confinement to the
// root alone would follow it.
manifest := writeRawManifest(t, root, "other", strings.Replace(demoManifest(true), "zero.demo", "zero.other", 1))
return "other", manifest
}},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
root := t.TempDir()
raw := demoManifest(true)
writeRawManifest(t, root, "demo", raw)
plugin, ok, err := FindByID(enableTestOptions(root), "zero.demo")
if err != nil || !ok {
t.Fatalf("FindByID: ok=%v err=%v", ok, err)
}
target, targetManifest := tc.link(t, root)
targetBefore, err := os.ReadFile(targetManifest)
if err != nil {
t.Fatal(err)
}
moved := filepath.Join(t.TempDir(), "moved-away")
if err := os.Rename(plugin.PluginDir, moved); err != nil {
t.Fatal(err)
}
if err := os.Symlink(target, plugin.PluginDir); err != nil {
t.Skipf("cannot create a directory symlink here: %v", err)
}
if _, err := setEnabledLocked(plugin, false); !errors.Is(err, pathjail.ErrReparse) {
t.Fatalf("setEnabledLocked through a swapped directory: err = %v, want %v", err, pathjail.ErrReparse)
}
if after, err := os.ReadFile(targetManifest); err != nil || string(after) != string(targetBefore) {
t.Fatalf("the edit followed the link and changed %s (err %v):\n got %q\nwant %q", targetManifest, err, after, targetBefore)
}
if after, err := os.ReadFile(filepath.Join(moved, "plugin.json")); err != nil || string(after) != raw {
t.Fatalf("the moved-away manifest changed (err %v): %q", err, after)
}
})
}
}
// TestSetEnabledDetectsAPluginDirectoryReplacedAfterOpening replaces the live
// plugin directory after the toggle pinned it. The edit then lands in the
// moved-away directory, so the toggle must not refresh the lockfile from either
// tree and must not report success: the plugin installed at the path now is
// unchanged.
func TestSetEnabledDetectsAPluginDirectoryReplacedAfterOpening(t *testing.T) {
for _, tc := range []struct {
name string
pinned bool // the pinned manifest's enabled value
changes bool // whether disabling it is a real edit
}{
{"edit", true, true},
{"no-op", false, false},
} {
t.Run(tc.name, func(t *testing.T) {
root := t.TempDir()
manifestPath := writeRawManifest(t, root, "demo", demoManifest(tc.pinned))
pluginDir := filepath.Dir(manifestPath)
hash, err := hashTree(pluginDir)
if err != nil {
t.Fatal(err)
}
if err := writeLock(root, map[string]LockEntry{"zero.demo": {Source: "./demo", Hash: hash}}); err != nil {
t.Fatal(err)
}
lockPath := filepath.Join(root, LockFileName)
lockBefore, err := os.ReadFile(lockPath)
if err != nil {
t.Fatal(err)
}
plugin, ok, err := FindByID(enableTestOptions(root), "zero.demo")
if err != nil || !ok {
t.Fatalf("FindByID: ok=%v err=%v", ok, err)
}
pinned, err := openPluginDir(plugin)
if err != nil {
t.Fatalf("openPluginDir: %v", err)
}
defer pinned.Close()
moved := filepath.Join(t.TempDir(), "moved-away")
if err := os.Rename(pluginDir, moved); err != nil {
t.Fatal(err)
}
replacement := strings.Replace(demoManifest(true), `"version": "1.0.0"`, `"version": "2.0.0"`, 1)
writeRawManifest(t, root, "demo", replacement)
// The lockfile check hashes the pinned tree, which still matches the
// recorded hash, not the replacement now at the path.
if _, _, clean := cleanLockEntry(plugin, pinned.dir); !clean {
t.Error("cleanLockEntry hashed the replacement at the path instead of the pinned directory")
}
if _, err := setEnabledPinned(plugin, pinned, false); !errors.Is(err, errPluginDirReplaced) {
t.Fatalf("setEnabledPinned after the directory was replaced: err = %v, want %v", err, errPluginDirReplaced)
}
if after, err := os.ReadFile(manifestPath); err != nil || string(after) != replacement {
t.Fatalf("the replacement manifest changed (err %v): %q", err, after)
}
if after, err := os.ReadFile(lockPath); err != nil || string(after) != string(lockBefore) {
t.Fatalf("the lockfile changed (err %v):\n got %s\nwant %s", err, after, lockBefore)
}
movedRaw, err := os.ReadFile(filepath.Join(moved, "plugin.json"))
if err != nil {
t.Fatal(err)
}
if tc.changes != (string(movedRaw) != demoManifest(tc.pinned)) {
t.Fatalf("pinned manifest edit = %v, want %v: %q", !tc.changes, tc.changes, movedRaw)
}
})
}
}
func TestSetEnabledByIDRespectsUserOnlyFilter(t *testing.T) {
dir := t.TempDir()
userRoot := filepath.Join(dir, "user")
projectRoot := filepath.Join(dir, "project")
writePluginManifest(t, filepath.Join(userRoot, "demo"), map[string]any{
"schemaVersion": 1,
"id": "zero.demo",
"name": "User Demo",
"version": "0.1.0",
"enabled": true,
})
writePluginManifest(t, filepath.Join(projectRoot, "demo"), map[string]any{
"schemaVersion": 1,
"id": "zero.demo",
"name": "Project Demo",
"version": "0.2.0",
"enabled": true,
})
result, err := SetEnabledByID(LoadOptions{
Roots: []Root{
{Source: SourceUser, Path: userRoot},
{Source: SourceProject, Path: projectRoot},
},
}, "zero.demo", false)
if err != nil {
t.Fatalf("SetEnabledByID: %v", err)
}
if result.Source != SourceProject {
t.Fatalf("source = %s, want project (precedence)", result.Source)
}
if !result.Changed || result.Enabled {
t.Fatalf("unexpected result: %#v", result)
}
userOnly, err := SetEnabledByID(LoadOptions{
Roots: []Root{
{Source: SourceUser, Path: userRoot},
{Source: SourceProject, Path: projectRoot},
},
ExcludeProject: true,
}, "zero.demo", false)
if err != nil {
t.Fatalf("SetEnabledByID user-only: %v", err)
}
if userOnly.Source != SourceUser {
t.Fatalf("user-only source = %s, want user", userOnly.Source)
}
}
func TestSetEnabledPreservesManifestMode(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("Windows does not preserve chmod bits in os.Stat mode")
}
root := t.TempDir()
writePluginManifest(t, filepath.Join(root, "demo"), map[string]any{
"schemaVersion": 1,
"id": "zero.demo",
"name": "Demo",
"version": "1.0.0",
})
manifestPath := filepath.Join(root, "demo", "plugin.json")
if err := os.Chmod(manifestPath, 0o644); err != nil {
t.Fatal(err)
}
if _, err := SetEnabledByID(enableTestOptions(root), "zero.demo", false); err != nil {
t.Fatalf("SetEnabledByID: %v", err)
}
info, err := os.Stat(manifestPath)
if err != nil {
t.Fatal(err)
}
if got := info.Mode().Perm(); got != 0o644 {
t.Fatalf("mode = %04o, want 0644", got)
}
}
func TestSetEnabledByIDMissingPlugin(t *testing.T) {
_, err := SetEnabledByID(LoadOptions{Roots: []Root{{Source: SourceUser, Path: t.TempDir()}}}, "missing", false)
if err == nil {
t.Fatal("expected missing plugin error")
}
if !errors.Is(err, ErrNotInstalled) {
t.Fatalf("error = %v, want ErrNotInstalled", err)
}
}
func TestSetEnabledPublishFailureLeavesManifestIntact(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("Windows directory permissions do not block writes the way this test needs")
}
if os.Geteuid() == 0 {
t.Skip("root bypasses directory permission checks")
}
root := t.TempDir()
raw := demoManifest(true)
manifestPath := writeRawManifest(t, root, "demo", raw)
pluginDir := filepath.Dir(manifestPath)
if err := os.Chmod(pluginDir, 0o555); err != nil {
t.Fatal(err)
}
defer func() { _ = os.Chmod(pluginDir, 0o755) }()
if _, err := SetEnabledByID(enableTestOptions(root), "zero.demo", false); err == nil {
t.Fatal("expected the publish to fail in a read-only plugin directory")
}
after, err := os.ReadFile(manifestPath)
if err != nil {
t.Fatal(err)
}
if string(after) != raw {
t.Fatalf("failed publish changed the manifest:\n got %q\nwant %q", after, raw)
}
}
// TestSetEnabledByIDWaitsForInstallRootLock proves the toggle takes the same
// cross-process lock Install and Remove take on the plugin's install root: it
// must not complete while an install holds that lock.
func TestSetEnabledByIDWaitsForInstallRootLock(t *testing.T) {
root := t.TempDir()
manifestPath := writeRawManifest(t, root, "demo", demoManifest(true))
unlock, err := installtxn.Lock(root)
if err != nil {
t.Fatal(err)
}
results := make(chan SetEnabledResult, 1)
failures := make(chan error, 1)
go func() {
result, err := SetEnabledByID(enableTestOptions(root), "zero.demo", false)
if err != nil {
failures <- err
return
}
results <- result
}()
select {
case result := <-results:
unlock()
t.Fatalf("toggle completed while the install lock was held: %#v", result)
case err := <-failures:
unlock()
t.Fatalf("toggle failed while the install lock was held: %v", err)
case <-time.After(200 * time.Millisecond):
}
unlock()
select {
case result := <-results:
if !result.Changed {
t.Fatalf("toggle did not report a change: %#v", result)
}
case err := <-failures:
t.Fatalf("toggle failed after the lock was released: %v", err)
case <-time.After(10 * time.Second):
t.Fatal("toggle did not finish after the install lock was released")
}
raw, err := os.ReadFile(manifestPath)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(string(raw), `"enabled": false`) {
t.Fatalf("manifest not disabled: %s", raw)
}
}
// TestSetEnabledByIDSynchronizesWithStagedInstallCommit pauses an install
// between staging and commit, lets the toggle try to run, then commits the
// install. The toggle must apply to the newly installed manifest instead of
// overwriting it with bytes read before the install, and the lockfile entry
// must describe the tree actually on disk.
func TestSetEnabledByIDSynchronizesWithStagedInstallCommit(t *testing.T) {
root := t.TempDir()
oldRaw := demoManifest(true)
pluginDir := filepath.Dir(writeRawManifest(t, root, "zero.demo", oldRaw))
// Build the replacement tree with the same staging primitives Install
// uses and leave it staged but uncommitted.
src := t.TempDir()
newRaw := `{
"schemaVersion": 1,
"id": "zero.demo",
"name": "Demo",
"version": "2.0.0",
"description": "reinstalled",
"enabled": true
}
`
srcPlugin := filepath.Join(src, "zero.demo")
if err := os.MkdirAll(srcPlugin, 0o700); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(srcPlugin, "plugin.json"), []byte(newRaw), 0o600); err != nil {
t.Fatal(err)
}
staged, cleanup, err := installtxn.StageDir(root)
if err != nil {
t.Fatal(err)
}
defer cleanup()
if err := copyTree(srcPlugin, staged); err != nil {
t.Fatal(err)
}
stagedHash, err := hashTree(staged)
if err != nil {
t.Fatal(err)
}
// An install holds the root lock from before its commit until after it.
unlock, err := installtxn.Lock(root)
if err != nil {
t.Fatal(err)
}
results := make(chan SetEnabledResult, 1)
failures := make(chan error, 1)
go func() {
result, err := SetEnabledByID(enableTestOptions(root), "zero.demo", false)
if err != nil {
failures <- err
return
}
results <- result
}()
select {
case result := <-results:
unlock()
t.Fatalf("toggle completed while the install lock was held: %#v", result)
case err := <-failures:
unlock()
t.Fatalf("toggle failed while the install lock was held: %v", err)
case <-time.After(200 * time.Millisecond):
}
// Commit the staged install exactly like Install does: swap the tree and
// publish the lockfile entry while still holding the lock.
lock, err := ReadLock(root)
if err != nil {
unlock()
t.Fatal(err)
}
lock["zero.demo"] = LockEntry{Source: string(SourceUser), Hash: stagedHash}
if err := installtxn.CommitDir(pluginDir, staged, func() error { return writeLock(root, lock) }); err != nil {
unlock()
t.Fatal(err)
}
unlock()
var result SetEnabledResult
select {
case result = <-results:
case err := <-failures:
t.Fatalf("toggle failed after the install committed: %v", err)
case <-time.After(10 * time.Second):
t.Fatal("toggle did not finish after the install committed")
}
if !result.Changed {
t.Fatalf("toggle result = %#v, want changed", result)
}
// The toggle must have applied to the newly installed manifest, not to the
// bytes it first read.
final, err := os.ReadFile(filepath.Join(pluginDir, "plugin.json"))
if err != nil {
t.Fatal(err)
}
want := strings.Replace(newRaw, `"enabled": true`, `"enabled": false`, 1)
if string(final) != want {
t.Fatalf("manifest after concurrent install + disable:\n got %q\nwant %q", final, want)
}
// The lockfile entry must describe the tree that is actually on disk.
entries, err := ReadLock(root)
if err != nil {
t.Fatal(err)
}
current, err := hashTree(pluginDir)
if err != nil {
t.Fatal(err)
}
if entries["zero.demo"].Hash != current {
t.Fatalf("plugins.lock hash = %q, want current tree hash %q", entries["zero.demo"].Hash, current)
}
}
func TestSetEnabledByIDRefreshesCleanLockEntry(t *testing.T) {
root := t.TempDir()
manifestPath := writeRawManifest(t, root, "zero.demo", demoManifest(true))
pluginDir := filepath.Dir(manifestPath)
current, err := hashTree(pluginDir)
if err != nil {
t.Fatal(err)
}
if err := writeLock(root, map[string]LockEntry{"zero.demo": {Source: string(SourceUser), Hash: current}}); err != nil {
t.Fatal(err)
}
if _, err := SetEnabledByID(enableTestOptions(root), "zero.demo", false); err != nil {
t.Fatalf("SetEnabledByID: %v", err)
}
entries, err := ReadLock(root)
if err != nil {
t.Fatal(err)
}
afterToggle, err := hashTree(pluginDir)
if err != nil {
t.Fatal(err)
}
if entries["zero.demo"].Hash == current {
t.Fatal("lock hash was not refreshed for a clean install")
}
if entries["zero.demo"].Hash != afterToggle {
t.Fatalf("lock hash = %q, want refreshed tree hash %q", entries["zero.demo"].Hash, afterToggle)
}
info, err := Info(InfoOptions{LoadOptions: enableTestOptions(root)}, "zero.demo")
if err != nil {
t.Fatal(err)
}
if info.HashDrift {
t.Fatal("toggle left hash drift behind for a clean install")
}
// Pre-existing drift must survive the next toggle instead of being
// absorbed into the lockfile.
if err := os.WriteFile(filepath.Join(pluginDir, "extra.txt"), []byte("edited\n"), 0o600); err != nil {
t.Fatal(err)
}
if _, err := SetEnabledByID(enableTestOptions(root), "zero.demo", true); err != nil {
t.Fatalf("SetEnabledByID(true): %v", err)
}
entries, err = ReadLock(root)
if err != nil {
t.Fatal(err)
}
if entries["zero.demo"].Hash != afterToggle {
t.Fatalf("toggle absorbed pre-existing drift: hash = %q, want %q", entries["zero.demo"].Hash, afterToggle)
}
}
// TestSetEnabledByIDReportsAPluginRemovedWhileWaiting removes the plugin while
// the toggle is blocked on the install-root lock, exactly as Remove does while
// it holds that lock. The toggle must report the plugin as not installed
// instead of editing, or recreating, a tree that no longer exists.
func TestSetEnabledByIDReportsAPluginRemovedWhileWaiting(t *testing.T) {
root := t.TempDir()
pluginDir := filepath.Dir(writeRawManifest(t, root, "zero.demo", demoManifest(true)))
unlock, err := installtxn.Lock(root)
if err != nil {
t.Fatal(err)
}
failures := make(chan error, 1)
go func() {
_, err := SetEnabledByID(enableTestOptions(root), "zero.demo", false)
failures <- err
}()
select {
case err := <-failures:
unlock()
t.Fatalf("toggle returned while the install lock was held: %v", err)
case <-time.After(200 * time.Millisecond):
}
if err := installtxn.RemoveDir(pluginDir, func() error { return nil }); err != nil {
unlock()
t.Fatal(err)
}
unlock()
select {
case err := <-failures:
if !errors.Is(err, ErrNotInstalled) {
t.Fatalf("toggle after a concurrent removal: error = %v, want ErrNotInstalled", err)
}
case <-time.After(10 * time.Second):
t.Fatal("toggle did not finish after the removal released the lock")
}
if _, err := os.Stat(pluginDir); !os.IsNotExist(err) {
t.Fatalf("toggle recreated the removed plugin: %v", err)
}
}
// TestSetEnabledByIDResolvesInterruptedUpdatesFirst covers every point an
// update can be killed at that leaves the plugin discoverable. The toggle
// holds the install-root lock, so it has to
// run the same recovery pass Install and Remove run before it reads the tree:
// otherwise it edits, or fails to find, a tree that is still owed a restore,
// and the next install or removal puts the backup back over the user's edit.
func TestSetEnabledByIDResolvesInterruptedUpdatesFirst(t *testing.T) {
for _, state := range interruptedStates() {
t.Run(state.name, func(t *testing.T) {
u := seedInterruptedUpdate(t)
workspace := state.plant(t, u)
want := state.want(u)
if _, err := os.Lstat(u.target()); os.IsNotExist(err) {
// No root lists a plugin whose only copy is in a workspace, so the
// toggle reports it not installed without taking any lock, just as
// list and info do; the next install or removal restores it.
t.Skip("nothing is at the target, so the toggle never takes the lock")
}
result, err := SetEnabledByID(enableTestOptions(u.dir), "zero.demo", false)
if err != nil {
t.Fatalf("toggle after the interrupted update: %v", err)
}
if !result.Changed || result.Enabled {
t.Fatalf("toggle result = %#v, want changed and disabled", result)
}
assertDisabledAfterRecovery(t, u, workspace, want)
// The next lock holder runs recovery again. The disable has to survive
// it rather than be replaced by a backup the toggle never saw.
if err := installOther(t, u.dir, "zero.other"); err != nil {
t.Fatalf("install after the toggle: %v", err)
}
assertDisabledAfterRecovery(t, u, workspace, want)
})
}
}
func assertDisabledAfterRecovery(t *testing.T, u interruptedUpdate, workspace string, want recoveryWant) {
t.Helper()
data, err := os.ReadFile(filepath.Join(u.target(), "generation.txt"))
if err != nil {
t.Fatalf("reading the live tree: %v", err)
}
if string(data) != want.generation {
t.Errorf("the live tree is the %q generation, want %q", data, want.generation)
}
plugin, ok, err := FindByID(enableTestOptions(u.dir), "zero.demo")
if err != nil {
t.Fatal(err)
}
if !ok {
t.Fatal("Load does not see the toggled plugin")
}
if plugin.Version != want.version {
t.Errorf("Load sees version %q, want %q", plugin.Version, want.version)
}
if plugin.Enabled {
t.Errorf("the disable did not stick: Load sees the plugin enabled")
}
_, err = os.Stat(workspace)
if want.workspace && err != nil {
t.Errorf("the workspace should be left alone, got %v", err)
}
if !want.workspace && !os.IsNotExist(err) {
t.Errorf("the resolved workspace should be cleared, got %v", err)
}
}
// TestSetEnabledByIDAbortsOnAnUnresolvedInterruptedUpdate leaves a workspace
// whose trees match neither lockfile hash. Recovery cannot pick one, so the
// toggle must fail without editing the live manifest or the lockfile.
func TestSetEnabledByIDAbortsOnAnUnresolvedInterruptedUpdate(t *testing.T) {
u := seedInterruptedUpdate(t)
workspace := u.plantWorkspace(t)
u.mustRename(t, u.target(), filepath.Join(workspace, "previous"))
u.mustRename(t, filepath.Join(workspace, "staged"), u.target())
unmatched := LockEntry{Source: u.oldSource, Hash: "sha256:0000000000000000000000000000000000000000000000000000000000000000"}
lock, err := ReadLock(u.dir)
if err != nil {
t.Fatal(err)
}
lock["zero.demo"] = unmatched
if err := writeLock(u.dir, lock); err != nil {
t.Fatal(err)
}
manifestPath := filepath.Join(u.target(), "plugin.json")
before, err := os.ReadFile(manifestPath)
if err != nil {
t.Fatal(err)
}
if _, err := SetEnabledByID(enableTestOptions(u.dir), "zero.demo", false); err == nil {
t.Fatal("toggle did not report the unresolved transaction")
}
after, err := os.ReadFile(manifestPath)
if err != nil {
t.Fatal(err)
}
if string(after) != string(before) {
t.Errorf("the aborted toggle edited the manifest:\n got %q\nwant %q", after, before)
}
assertUntouched(t, u, "new", unmatched)
if _, err := os.Stat(filepath.Join(workspace, "previous")); err != nil {
t.Errorf("the backup should still be there for a hand rescue, got %v", err)
}
}
+2 -1
View File
@@ -7,7 +7,8 @@ import (
"strings"
)
// ErrNotInstalled is returned when Info cannot find a plugin with the requested id.
// ErrNotInstalled is returned when a plugin id cannot be found under the given
// LoadOptions (for example by Info or SetEnabledByID).
var ErrNotInstalled = errors.New("plugin is not installed")
// PluginInfo combines a loaded plugin with optional lockfile metadata.
+37 -11
View File
@@ -567,15 +567,32 @@ func validInstallID(id string) bool {
// sorted order. Each file contributes its plugin-relative path, executable bit,
// and bytes, so renames, mode flips, and content edits all change the hash.
func hashTree(root string) (string, error) {
handle, err := os.OpenRoot(root)
if err != nil {
return "", err
}
defer handle.Close()
return hashRootedTree(handle)
}
// hashRootedTree is hashTree over a directory that is already open. Walking
// relative to the handle hashes the directory it names even if its path is
// replaced meanwhile, and keeps the walk inside it.
//
// The walk calls the handle's own methods rather than handle.FS(): the fs.FS
// adapter only accepts fs.ValidPath names, which must be UTF-8, so it would
// refuse file names that the filesystem and copyTree accept.
func hashRootedTree(handle *os.Root) (string, error) {
hasher := sha256.New()
if err := hashTreeInto(hasher, root, root); err != nil {
if err := hashTreeInto(hasher, handle, "."); err != nil {
return "", err
}
return "sha256:" + hex.EncodeToString(hasher.Sum(nil)), nil
}
func hashTreeInto(hasher io.Writer, root string, dir string) error {
entries, err := os.ReadDir(dir)
// hashTreeInto hashes dir, a path relative to handle, into hasher.
func hashTreeInto(hasher io.Writer, handle *os.Root, dir string) error {
entries, err := readRootedDir(handle, dir)
if err != nil {
return err
}
@@ -588,8 +605,8 @@ func hashTreeInto(hasher io.Writer, root string, dir string) error {
if name == ".git" {
continue
}
path := filepath.Join(dir, name)
info, err := os.Lstat(path)
rel := filepath.Join(dir, name)
info, err := handle.Lstat(rel)
if err != nil {
return err
}
@@ -598,14 +615,10 @@ func hashTreeInto(hasher io.Writer, root string, dir string) error {
// Skipped by copyTree, so excluded from the hash too.
continue
case info.IsDir():
if err := hashTreeInto(hasher, root, path); err != nil {
if err := hashTreeInto(hasher, handle, rel); err != nil {
return err
}
case info.Mode().IsRegular():
rel, err := filepath.Rel(root, path)
if err != nil {
return err
}
executable := 0
if info.Mode().Perm()&0o111 != 0 {
executable = 1
@@ -616,7 +629,7 @@ func hashTreeInto(hasher io.Writer, root string, dir string) error {
if _, err := io.WriteString(hasher, header); err != nil {
return err
}
file, err := os.Open(path)
file, err := handle.Open(rel)
if err != nil {
return err
}
@@ -634,3 +647,16 @@ func hashTreeInto(hasher io.Writer, root string, dir string) error {
}
return nil
}
// readRootedDir lists dir, a path relative to handle, like os.ReadDir.
func readRootedDir(handle *os.Root, dir string) ([]os.DirEntry, error) {
file, err := handle.Open(dir)
if err != nil {
return nil, err
}
entries, err := file.ReadDir(-1)
if closeErr := file.Close(); err == nil {
err = closeErr
}
return entries, err
}
+59
View File
@@ -1181,3 +1181,62 @@ func TestRecoveryLeavesAPrefixCollidingDirectoryAlone(t *testing.T) {
t.Errorf("recovery published content it does not own: %v", err)
}
}
// TestNonUTF8FileNamesAreInstalledHashedAndToggled covers a plugin whose tree
// holds a file and a directory named with bytes that are not valid UTF-8, which
// Linux allows. copyTree copies them, so hashTree must hash them too: Install
// records the hash, the toggle refreshes it, and Info sees an edit to that file
// as drift. A walk that rejects such names breaks all three.
func TestNonUTF8FileNamesAreInstalledHashedAndToggled(t *testing.T) {
src := writeSourcePlugin(t, filepath.Join(t.TempDir(), "src"), validManifest())
nested := filepath.Join("assets\xff", "data\xfe.bin")
if err := os.Mkdir(filepath.Join(src, "assets\xff"), 0o755); err != nil {
t.Skipf("this filesystem rejects non-UTF-8 names: %v", err)
}
if err := os.WriteFile(filepath.Join(src, nested), []byte("original\n"), 0o644); err != nil {
t.Skipf("this filesystem rejects non-UTF-8 names: %v", err)
}
destDir := t.TempDir()
if _, err := Install(context.Background(), InstallOptions{Source: src, Dir: destDir}); err != nil {
t.Fatalf("Install a plugin with a non-UTF-8 file name: %v", err)
}
pluginDir := filepath.Join(destDir, "zero.demo")
if data, err := os.ReadFile(filepath.Join(pluginDir, nested)); err != nil || string(data) != "original\n" {
t.Fatalf("installed copy of the non-UTF-8 file = %q, %v", data, err)
}
options := LoadOptions{Roots: []Root{{Source: SourceUser, Path: destDir}}}
before, err := ReadLock(destDir)
if err != nil {
t.Fatal(err)
}
if _, err := SetEnabledByID(options, "zero.demo", false); err != nil {
t.Fatalf("SetEnabledByID on a plugin with a non-UTF-8 file name: %v", err)
}
after, err := ReadLock(destDir)
if err != nil {
t.Fatal(err)
}
if after["zero.demo"].Hash == before["zero.demo"].Hash {
t.Fatal("the toggle did not refresh the lock hash")
}
info, err := Info(InfoOptions{LoadOptions: options}, "zero.demo")
if err != nil {
t.Fatal(err)
}
if info.HashDrift {
t.Fatal("the toggle left hash drift behind")
}
if err := os.WriteFile(filepath.Join(pluginDir, nested), []byte("edited\n"), 0o644); err != nil {
t.Fatal(err)
}
info, err = Info(InfoOptions{LoadOptions: options}, "zero.demo")
if err != nil {
t.Fatal(err)
}
if !info.HashDrift {
t.Fatal("an edit to the non-UTF-8 file was not reported as drift")
}
}
+3 -1
View File
@@ -132,7 +132,9 @@ var descriptors = []Descriptor{
// way in; RequiresAuth is forced true so catalog consumers know the provider
// needs an interactive login before a request will succeed.
func() Descriptor {
d := openAICompat("chatgpt", "ChatGPT", "https://chatgpt.com/backend-api/codex", "gpt-5.5", nil)
// New ChatGPT profiles follow the vendor's GPT-5.5 replacement guidance:
// https://learn.chatgpt.com/docs/models#gpt-55-retirement
d := openAICompat("chatgpt", "ChatGPT", "https://chatgpt.com/backend-api/codex", "gpt-6.1-sol", nil)
d.RequiresAuth = true
return oauthProvider(d, false, false)
}(),
+20 -4
View File
@@ -44,6 +44,20 @@ var minimaxCuratedModels = []Model{
OutputModalities: []string{"text"},
Source: minimaxModelSource,
},
{
ID: "MiniMax-M3.1-Flash",
Description: "fast reasoning model",
ContextWindow: 1_000_000,
ToolCall: true,
Reasoning: true,
ReasoningEfforts: []string{"low", "medium", "high", "xhigh", "max"},
DefaultReasoningEffort: "max",
InputModalities: []string{"text", "image", "video"},
OutputModalities: []string{"text"},
InputCost: 0.15,
OutputCost: 0.60,
Source: minimaxModelSource,
},
{
ID: "MiniMax-M2.7",
Description: "agentic coding model",
@@ -113,10 +127,12 @@ var curatedModels = map[string][]Model{
{ID: "deepseek/deepseek-chat", Description: "coding model"},
},
"chatgpt": {
{ID: "gpt-5.5", Description: "recommended Codex model"},
{ID: "gpt-5.4", Description: "strong Codex model"},
{ID: "gpt-5.4-mini", Description: "fast Codex model"},
{ID: "gpt-5.3-codex-spark", Description: "research preview fast Codex model"},
// Suggestions only: availability comes from the authenticated Codex catalog.
// https://learn.chatgpt.com/docs/models
{ID: "gpt-6.1-sol", Description: "balanced Codex model"},
{ID: "gpt-6.1-terra", Description: "fast model for writing, visual understanding, and computer use"},
{ID: "gpt-6-astra", Description: "advanced reasoning model"},
{ID: "gpt-6-luna", Description: "efficient coding model"},
},
"deepseek": {
{ID: "deepseek-chat", Description: "catalog default"},
+22 -6
View File
@@ -35,8 +35,8 @@ func TestModelsAreProviderScoped(t *testing.T) {
},
{
provider: "chatgpt",
want: []string{"gpt-5.5", "gpt-5.4", "gpt-5.4-mini", "gpt-5.3-codex-spark"},
notWant: []string{"gpt-5", "gpt-4.1", "openai/gpt-4.1"},
want: []string{"gpt-6.1-sol", "gpt-6.1-terra", "gpt-6-astra", "gpt-6-luna"},
notWant: []string{"gpt-5.5", "gpt-5.4", "gpt-5.4-mini", "gpt-5.3-codex-spark", "openai/gpt-4.1"},
},
{
provider: "aimlapi",
@@ -65,7 +65,7 @@ func TestModelsAreProviderScoped(t *testing.T) {
},
{
provider: "minimaxi-cn",
want: []string{"MiniMax-M3", "MiniMax-M2.7", "MiniMax-M2.1"},
want: []string{"MiniMax-M3", "MiniMax-M3.1-Flash", "MiniMax-M2.7", "MiniMax-M2.1"},
notWant: []string{"gpt-4.1", "claude-sonnet-4.5"},
},
{
@@ -165,8 +165,8 @@ func TestMiniMaxModelsExposeCurrentCapabilities(t *testing.T) {
t.Fatal("provider minimax missing from catalog")
}
models := Models(descriptor)
if len(models) < 2 || models[0].ID != "MiniMax-M3" || models[1].ID != "MiniMax-M2.7" {
t.Fatalf("MiniMax current models = %#v, want MiniMax-M3 followed by MiniMax-M2.7", modelIDs(models))
if len(models) < 3 || models[0].ID != "MiniMax-M3" || models[1].ID != "MiniMax-M3.1-Flash" || models[2].ID != "MiniMax-M2.7" {
t.Fatalf("MiniMax current models = %#v, want MiniMax-M3, MiniMax-M3.1-Flash, then MiniMax-M2.7", modelIDs(models))
}
m3 := models[0]
@@ -178,7 +178,23 @@ func TestMiniMaxModelsExposeCurrentCapabilities(t *testing.T) {
t.Fatalf("MiniMax-M3 modalities = input:%#v output:%#v", m3.InputModalities, m3.OutputModalities)
}
m27 := models[1]
m31Flash := models[1]
if m31Flash.ContextWindow != 1_000_000 || !m31Flash.ToolCall || !m31Flash.Reasoning {
t.Fatalf("MiniMax-M3.1-Flash capabilities = %#v, want 1M context with tools and reasoning", m31Flash)
}
if !reflect.DeepEqual(m31Flash.ReasoningEfforts, []string{"low", "medium", "high", "xhigh", "max"}) ||
m31Flash.DefaultReasoningEffort != "max" {
t.Fatalf("MiniMax-M3.1-Flash reasoning efforts = %#v default %q", m31Flash.ReasoningEfforts, m31Flash.DefaultReasoningEffort)
}
if !reflect.DeepEqual(m31Flash.InputModalities, []string{"text", "image", "video"}) ||
!reflect.DeepEqual(m31Flash.OutputModalities, []string{"text"}) {
t.Fatalf("MiniMax-M3.1-Flash modalities = input:%#v output:%#v", m31Flash.InputModalities, m31Flash.OutputModalities)
}
if m31Flash.InputCost != 0.15 || m31Flash.OutputCost != 0.60 {
t.Fatalf("MiniMax-M3.1-Flash costs = %f/%f, want 0.15/0.60", m31Flash.InputCost, m31Flash.OutputCost)
}
m27 := models[2]
if m27.ContextWindow != 204_800 || !m27.ToolCall || !m27.Reasoning {
t.Fatalf("MiniMax-M2.7 capabilities = %#v, want 204.8K context with tools and reasoning", m27)
}
@@ -0,0 +1,31 @@
package providermodeldiscovery
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/Twigpine/zero/internal/config"
)
func TestChatGPTDiscoveryUsesCurrentCompatibilityVersion(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if got := r.URL.Query().Get("client_version"); got != "0.159.0" {
t.Errorf("client_version = %q, want documented GPT-6.1 Terra compatibility floor 0.159.0", got)
}
_, _ = w.Write([]byte(`{"models":[{"slug":"gpt-future-coder"}]}`))
}))
defer server.Close()
profile := config.ProviderProfile{
CatalogID: "chatgpt", ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: server.URL + "/backend-api/codex", APIKey: "test-bearer",
}
models, err := Discover(context.Background(), profile, Options{HTTPClient: server.Client()})
if err != nil {
t.Fatal(err)
}
if len(models) != 1 || models[0].ID != "gpt-future-coder" {
t.Fatalf("models = %#v, want only gpt-future-coder", models)
}
}
+3 -2
View File
@@ -20,8 +20,9 @@ import (
)
const (
anthropicVersion = "2023-06-01"
chatGPTModelsProtocolVersion = "0.146.0"
anthropicVersion = "2023-06-01"
// GPT-6.1 Terra's documented Codex client floor: https://learn.chatgpt.com/docs/models.
chatGPTModelsProtocolVersion = "0.159.0"
)
type Model struct {
+22 -5
View File
@@ -49,6 +49,19 @@ func thinkingBudgetForEffort(effort string) int {
}
}
func miniMaxM31FlashEffort(model, effort string) string {
if !strings.EqualFold(strings.TrimSpace(model), "MiniMax-M3.1-Flash") {
return ""
}
effort = strings.ToLower(strings.TrimSpace(effort))
switch effort {
case "low", "medium", "high", "xhigh", "max":
return effort
default:
return ""
}
}
// resolveThinking returns the thinking budget and the max_tokens to send. When a
// budget is requested, max_tokens is raised if needed so the budget plus a
// minimum response both fit (Anthropic rejects budget >= max_tokens). enabled is
@@ -377,11 +390,15 @@ func (provider *Provider) anthropicRequest(request zeroruntime.CompletionRequest
Messages: messages,
Stream: true,
}
// Extended thinking: enable when a reasoning effort was requested. The budget
// is counted against max_tokens, so raise max_tokens to leave room for the
// response. Temperature is intentionally left unset (Anthropic requires the
// default when thinking is on).
if budget, effectiveMax, enabled := resolveThinking(request.ReasoningEffort, provider.maxTokens); enabled {
// MiniMax-M3.1-Flash uses adaptive thinking with output_config.effort instead
// of Anthropic's token budget. For other models, the budget is counted against
// max_tokens, so raise max_tokens to leave room for the response. Temperature
// is intentionally left unset (Anthropic requires the default when thinking
// is on).
if effort := miniMaxM31FlashEffort(provider.model, request.ReasoningEffort); effort != "" {
mapped.Thinking = &thinkingConfig{Type: "adaptive"}
mapped.OutputConfig = &outputConfig{Effort: effort}
} else if budget, effectiveMax, enabled := resolveThinking(request.ReasoningEffort, provider.maxTokens); enabled {
mapped.MaxTokens = effectiveMax
mapped.Thinking = &thinkingConfig{Type: "enabled", BudgetTokens: budget}
}
@@ -266,6 +266,53 @@ func TestStreamCompletionEnablesThinkingWhenEffortRequested(t *testing.T) {
}
}
func TestStreamCompletionUsesMiniMaxAdaptiveEffort(t *testing.T) {
t.Parallel()
for _, effort := range []string{"low", "xhigh"} {
effort := effort
t.Run(effort, func(t *testing.T) {
t.Parallel()
var gotBody map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil {
t.Fatalf("decode request body: %v", err)
}
writeSSEEvent(w, "message_stop", `{"type":"message_stop"}`)
}))
defer server.Close()
provider, err := New(Options{BaseURL: server.URL + "/", Model: "MiniMax-M3.1-Flash", MaxTokens: 64_000})
if err != nil {
t.Fatalf("New returned error: %v", err)
}
stream, err := provider.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hi"}},
ReasoningEffort: effort,
})
if err != nil {
t.Fatalf("StreamCompletion returned error: %v", err)
}
drain(stream)
thinking, ok := gotBody["thinking"].(map[string]any)
if !ok || thinking["type"] != "adaptive" {
t.Fatalf("thinking = %#v, want adaptive", gotBody["thinking"])
}
if _, ok := thinking["budget_tokens"]; ok {
t.Fatalf("MiniMax thinking must not include budget_tokens: %#v", thinking)
}
outputConfig, ok := gotBody["output_config"].(map[string]any)
if !ok || outputConfig["effort"] != effort {
t.Fatalf("output_config = %#v, want effort %q", gotBody["output_config"], effort)
}
if mt, _ := gotBody["max_tokens"].(float64); int(mt) != 64_000 {
t.Fatalf("max_tokens = %#v, want unchanged 64000", gotBody["max_tokens"])
}
})
}
}
func TestStreamCompletionOmitsThinkingWithoutEffort(t *testing.T) {
t.Parallel()
var gotBody map[string]any
+11 -7
View File
@@ -10,17 +10,21 @@ type messagesRequest struct {
System any `json:"system,omitempty"`
// Thinking enables extended thinking when a reasoning effort was requested.
// Omitted (nil) for normal requests so behavior is unchanged.
Thinking *thinkingConfig `json:"thinking,omitempty"`
Tools []anthropicTool `json:"tools,omitempty"`
Stream bool `json:"stream"`
Thinking *thinkingConfig `json:"thinking,omitempty"`
OutputConfig *outputConfig `json:"output_config,omitempty"`
Tools []anthropicTool `json:"tools,omitempty"`
Stream bool `json:"stream"`
}
// thinkingConfig requests extended thinking with a token budget. Type is always
// "enabled"; BudgetTokens is how many tokens the model may spend thinking (it is
// counted against max_tokens).
// thinkingConfig requests extended thinking. Anthropic uses "enabled" with a
// token budget; MiniMax-M3.1-Flash uses "adaptive" with output_config.effort.
type thinkingConfig struct {
Type string `json:"type"`
BudgetTokens int `json:"budget_tokens"`
BudgetTokens int `json:"budget_tokens,omitempty"`
}
type outputConfig struct {
Effort string `json:"effort"`
}
type anthropicMessage struct {
@@ -288,10 +288,11 @@ func TestScanSSEDataWithContextContentResetsOnData(t *testing.T) {
pr, pw := io.Pipe()
go func() {
// One data line every 30ms for ~300ms (well past the 120ms content window,
// with comfortable slack under the 100ms idle timeout), so the content
// watchdog keeps resetting, then close cleanly.
for i := 0; i < 10; i++ {
// One data line every 30ms for ~750ms (well past the 600ms content window,
// with ~470ms of slack under the 500ms idle timeout, so a scheduling stall
// on a busy shared runner does not trip it), so the content watchdog keeps
// resetting, then close cleanly.
for i := 0; i < 25; i++ {
if _, err := io.WriteString(pw, "data: chunk\n\n"); err != nil {
return
}
@@ -303,7 +304,7 @@ func TestScanSSEDataWithContextContentResetsOnData(t *testing.T) {
n := 0
done := make(chan error, 1)
go func() {
done <- ScanSSEDataWithContext(context.Background(), func() {}, pr, 100*time.Millisecond, func(string) bool { n++; return true })
done <- ScanSSEDataWithContext(context.Background(), func() {}, pr, 500*time.Millisecond, func(string) bool { n++; return true })
}()
select {
@@ -314,8 +315,8 @@ func TestScanSSEDataWithContextContentResetsOnData(t *testing.T) {
case <-time.After(3 * time.Second):
t.Fatal("ScanSSEDataWithContext hung on a producing stream")
}
if n != 10 {
t.Fatalf("handled %d data lines, want 10", n)
if n != 25 {
t.Fatalf("handled %d data lines, want 25", n)
}
}
+127
View File
@@ -63,6 +63,14 @@ import (
// rest of a key from prose written after a whole one: that prose would be eaten,
// and a following JWT would lose its header.
//
// CALLER-SUPPLIED SECRETS (Options.ExtraSecretValues) get the same treatment
// from redactExtraSecretValues. Such a value is often opaque, matching no shape
// here, so without its own gap-tolerant match a separator inside it left both
// fragments in the output. Unlike a shape, the value is known exactly and has no
// word boundary to judge, so it is found in a copy with the separators removed
// and mapped back. RedactString runs it first, before any shape matcher, so a
// caller that clips or reformats the result afterwards never sees a fragment.
//
// WHAT THIS DELIBERATELY DOES NOT TREAT AS A SEPARATOR: tab, newline and
// carriage return. Those three are real text structure rather than invisible
// filler, so a credential broken across a line boundary is still two visible
@@ -435,6 +443,125 @@ func openAIShapeIsSecret(match string) bool {
return !strings.Contains(strings.TrimPrefix(match, "sk-"), "-")
}
// redactExtraSecretValues replaces every occurrence of every caller-supplied
// secret, whole or with runs of invisible separators between its characters.
//
// All occurrences, overlapping ones included, are found in the same text and
// overlapping spans are merged before anything is replaced. Replacing one at a
// time let each replacement consume bytes another occurrence needed: a shorter
// value inside a split longer one, a periodic value ("abcabc" in
// "abc<NUL>abcabc"), a value whose first and last characters match, or two
// values that overlap all left a fragment visible. Occurrences that only touch
// stay separate replacements, as they always were.
//
// Separators inside a match are covered by its replacement; a split match never
// starts or ends on one, so those beside it stay, unless the value as written
// starts or ends with them. Tab, newline and carriage return are not separators
// (see above).
func redactExtraSecretValues(value string, secrets []string, replacement string) string {
haystack, offsets := value, []int(nil)
if containsSplitSeparator(value) {
haystack, offsets = compactSplitSeparators(value)
}
var spans [][2]int
for _, secret := range secrets {
if strings.TrimSpace(secret) == "" {
continue
}
needle := secret
if offsets != nil {
// The value as written always matches, as it did before split
// matching existed. The compacted match cannot stand in for it: a
// value made only of separators compacts to nothing, and separators
// at either end of a value fall outside its compacted match.
if strings.Contains(value, secret) {
for _, at := range overlappingIndexes(value, secret) {
spans = append(spans, [2]int{at, at + len(secret)})
}
}
needle = stripSplitSeparators(secret)
}
if needle == "" || !strings.Contains(haystack, needle) {
continue
}
for _, at := range overlappingIndexes(haystack, needle) {
start, end := at, at+len(needle)
if offsets != nil {
start, end = offsets[start], offsets[end-1]+1
}
spans = append(spans, [2]int{start, end})
}
}
if len(spans) == 0 {
return value
}
sort.Slice(spans, func(i, j int) bool { return spans[i][0] < spans[j][0] })
var out strings.Builder
out.Grow(len(value))
written := 0
for i := 0; i < len(spans); {
start, end := spans[i][0], spans[i][1]
for i++; i < len(spans) && spans[i][0] < end; i++ {
end = max(end, spans[i][1])
}
out.WriteString(value[written:start])
out.WriteString(replacement)
written = end
}
out.WriteString(value[written:])
return out.String()
}
// compactSplitSeparators returns value with every splitSecretSeparator removed,
// and for each byte of the result the offset of that byte in value.
func compactSplitSeparators(value string) (string, []int) {
var builder strings.Builder
builder.Grow(len(value))
offsets := make([]int, 0, len(value))
for offset := 0; offset < len(value); {
r, width := utf8.DecodeRuneInString(value[offset:])
if !splitSecretSeparator(r) {
builder.WriteString(value[offset : offset+width])
for i := 0; i < width; i++ {
offsets = append(offsets, offset+i)
}
}
offset += width
}
return builder.String(), offsets
}
// overlappingIndexes returns the start of every occurrence of needle in
// haystack, overlapping ones included, in one linear Knuth-Morris-Pratt pass:
// restarting a search one byte after each match is quadratic for a periodic
// value over text full of it.
func overlappingIndexes(haystack, needle string) []int {
fallback := make([]int, len(needle))
for i, k := 1, 0; i < len(needle); i++ {
for k > 0 && needle[i] != needle[k] {
k = fallback[k-1]
}
if needle[i] == needle[k] {
k++
}
fallback[i] = k
}
var starts []int
for i, k := 0, 0; i < len(haystack); i++ {
for k > 0 && haystack[i] != needle[k] {
k = fallback[k-1]
}
if haystack[i] == needle[k] {
k++
}
if k == len(needle) {
starts = append(starts, i-len(needle)+1)
k = fallback[k-1]
}
}
return starts
}
// redactControlSplitSecrets replaces credentials whose body is interrupted by
// invisible characters. It is a no-op for text containing none of them, which is
// nearly all text, so the contiguous path keeps its behavior exactly.
+108
View File
@@ -105,6 +105,114 @@ func TestRedactStringCoversTheSeparatorsInsideASplitSecret(t *testing.T) {
}
}
// A CALLER-SUPPLIED SECRET IS A CREDENTIAL TOO.
//
// ExtraSecretValues names a value the caller knows is secret, whatever its
// shape: an opaque provider key matches none of the shape matchers, so the
// split pass above never saw it, and a separator inside it left both fragments
// in the output. It has to be matched across the same separators, cover them in
// its replacement, and leave the separators beside it where they were.
func TestRedactStringRedactsExtraSecretValuesSplitByInvisibleCharacters(t *testing.T) {
const secret = "zq7Mv2Lp9Xr4Tn8Wk3"
options := Options{ExtraSecretValues: []string{secret}}
if out := RedactString("key "+secret, Options{}); !strings.Contains(out, secret) {
t.Fatalf("SETUP INVALID: a shape matcher already redacts %q, so this would not reach the explicit-value pass: %q", secret, out)
}
for _, testCase := range []struct {
name string
value string
want string
}{
{"NUL", "key " + secret[:6] + "\x00" + secret[6:], "key [REDACTED]"},
{"zero width space", "key " + secret[:1] + "\u200b" + secret[1:], "key [REDACTED]"},
{"soft hyphen near the end", "key " + secret[:len(secret)-1] + "\u00ad" + secret[len(secret)-1:], "key [REDACTED]"},
{"every gap", "key " + strings.Join(strings.Split(secret, ""), "\x1b"), "key [REDACTED]"},
{"adjacent separators", "key " + secret[:9] + "\x00\x1b\u2060" + secret[9:], "key [REDACTED]"},
{"twice", secret[:3] + "\x00" + secret[3:] + " and " + secret[:12] + "\ufeff" + secret[12:], "[REDACTED] and [REDACTED]"},
{"separators beside it stay", "\x00" + secret[:5] + "\u200b" + secret[5:] + "\x00", "\x00[REDACTED]\x00"},
{"a line break is not a gap", secret[:5] + "\n" + secret[5:], secret[:5] + "\n" + secret[5:]},
} {
t.Run(testCase.name, func(t *testing.T) {
if got := RedactString(testCase.value, options); got != testCase.want {
t.Errorf("RedactString(%q) = %q, want %q", testCase.value, got, testCase.want)
}
})
}
// A SHORTER SECRET MUST NOT CLAIM PART OF A SPLIT LONGER ONE. Longest first
// has to hold across the contiguous and split matches together: replacing
// a contiguous shorter value first left the rest of the split longer one
// ("[REDACTED]\x00Lp9Xr4Tn8Wk3") visible.
for _, testCase := range []struct {
name string
secrets []string
value string
want string
}{
{"prefix", []string{secret[:6], secret}, "key " + secret[:6] + "\x00" + secret[6:], "key [REDACTED]"},
{"suffix", []string{secret[9:], secret}, "key " + secret[:9] + "\u200b" + secret[9:], "key [REDACTED]"},
{"inside", []string{secret[4:10], secret}, "key " + secret[:2] + "\x00" + secret[2:], "key [REDACTED]"},
{"shorter alone too", []string{secret, secret[:6]}, secret[:6] + " and " + secret[:6] + "\x00" + secret[6:], "[REDACTED] and [REDACTED]"},
} {
t.Run("overlap "+testCase.name, func(t *testing.T) {
if got := RedactString(testCase.value, Options{ExtraSecretValues: testCase.secrets}); got != testCase.want {
t.Errorf("RedactString(%q, %q) = %q, want %q", testCase.value, testCase.secrets, got, testCase.want)
}
})
}
// OVERLAPPING OCCURRENCES ARE ONE SPAN. Replacing one occurrence (either
// form) and then looking for the next consumed bytes the other needed:
// "abc\x00abcabc" holds a split "abc\x00abc" overlapping a whole "abcabc",
// and replacing the whole one first left "abc\x00" in front of it. Every
// occurrence of every value, in either form, has to be found in the same
// text and overlapping ones replaced together.
for _, testCase := range []struct {
name string
secrets []string
value string
want string
}{
{"split then whole, periodic secret", []string{"abcabc"}, "abc\x00abcabc", "[REDACTED]"},
{"whole then split, periodic secret", []string{"abcabc"}, "abcabc\x00abc", "[REDACTED]"},
{"shared first and last character", []string{"xY7qx"}, "key xY7q\u200bxY7qx", "key [REDACTED]"},
{"two secrets overlapping across a split", []string{"zq7Mv2Lp", "Lp9Xr4Tn"}, "key zq7Mv2L\x00p9Xr4Tn", "key [REDACTED]"},
{"whole occurrences overlapping", []string{"xY7qx"}, "xY7qxY7qx", "[REDACTED]"},
{"adjacent occurrences stay two", []string{"abc"}, "abc\x00abc", "[REDACTED]\x00[REDACTED]"},
} {
t.Run("overlapping occurrences "+testCase.name, func(t *testing.T) {
if got := RedactString(testCase.value, Options{ExtraSecretValues: testCase.secrets}); got != testCase.want {
t.Errorf("RedactString(%q, %q) = %q, want %q", testCase.value, testCase.secrets, got, testCase.want)
}
})
}
// A SECRET AS WRITTEN IS ALWAYS REDACTED. The split match compares the text
// and the value with separators removed; it adds to the whole match and
// must never replace it. Once the text held any separator, a value made only
// of separators compacted to nothing and was skipped, and separators at
// either end of a value fell outside its compacted match, so both stayed
// visible where a plain replacement had always caught them.
for _, testCase := range []struct {
name string
secrets []string
value string
want string
}{
{"separator-only secret", []string{"\x00"}, "a\x00b", "a[REDACTED]b"},
{"separator-only secret beside another separator", []string{"\u200b"}, "a\u200bb\x00", "a[REDACTED]b\x00"},
{"separators at both ends", []string{"\u200babc\x00"}, "x\u200babc\x00y", "x[REDACTED]y"},
} {
t.Run("as written "+testCase.name, func(t *testing.T) {
if got := RedactString(testCase.value, Options{ExtraSecretValues: testCase.secrets}); got != testCase.want {
t.Errorf("RedactString(%q, %q) = %q, want %q", testCase.value, testCase.secrets, got, testCase.want)
}
})
}
// A value the caller names that is NOT in the text leaves the text alone,
// separators and all.
if value := "unrelated\x00text \u200bhere"; RedactString(value, options) != value {
t.Errorf("RedactString rewrote text without the secret in it: %q", RedactString(value, options))
}
}
// NOTHING ELSE MOVES. The compaction exists to find a match; it must never
// reach text that has no secret in it, and text with no invisible character at
// all must take exactly the path it took before.
+1 -9
View File
@@ -463,15 +463,7 @@ func RedactString(value string, options Options) string {
replacement := replacement(options)
redacted := value
if len(options.ExtraSecretValues) > 0 {
secrets := append([]string{}, options.ExtraSecretValues...)
sort.SliceStable(secrets, func(i, j int) bool {
return len(secrets[i]) > len(secrets[j])
})
for _, secret := range secrets {
if strings.TrimSpace(secret) != "" {
redacted = strings.ReplaceAll(redacted, secret, replacement)
}
}
redacted = redactExtraSecretValues(redacted, options.ExtraSecretValues, replacement)
}
if strings.Contains(redacted, "-----BEGIN ") && strings.Contains(redacted, "PRIVATE KEY") {
+258
View File
@@ -0,0 +1,258 @@
package skills
import (
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"io"
"io/fs"
"os"
"path/filepath"
"sort"
"strings"
)
// A skill is a directory, not just its SKILL.md: the instructions often point at
// scripts, references or assets that sit beside them (the agentskills.io
// layout). Install copies that directory, and the skill tool says where it is.
// Install refuses a skill tree larger than this. A skill is a handful of
// scripts and documents; a source that is really a whole repository or a pile
// of binaries should fail loudly instead of filling the skills directory.
var (
maxSkillTreeFiles = 1000
maxSkillTreeBytes = int64(32 << 20)
)
// bundledFileListLimit caps how many bundled files the skill tool lists, so a
// large skill does not crowd the instructions out of the tool output.
const bundledFileListLimit = 50
// skipSkillEntry reports whether a skill tree entry is left out of an install,
// a tree hash and a listing alike. Hidden entries cover VCS metadata (.git,
// which a cloned source always has) and local clutter such as .DS_Store or
// .env; they are not part of what a skill ships.
func skipSkillEntry(name string) bool {
return strings.HasPrefix(name, ".")
}
// copySkillTree copies the skill directory src into dst, which must exist.
// Regular files are copied with normalized permissions (0o755 when the source
// is executable, 0o644 otherwise), so a script stays runnable without carrying
// over group- or world-writable bits. Symlinks are not followed or copied:
// their relative paths come back in skipped, so the caller can say so. Other
// non-regular entries (devices, sockets, pipes) are left out. bundled counts
// the files copied besides SKILL.md.
func copySkillTree(src, dst string) (bundled int, skipped []string, err error) {
files := 0
var total int64
err = filepath.WalkDir(src, func(path string, entry fs.DirEntry, walkErr error) error {
if walkErr != nil {
return walkErr
}
if path == src {
return nil
}
rel, err := filepath.Rel(src, path)
if err != nil {
return err
}
if skipSkillEntry(entry.Name()) {
if entry.IsDir() {
return filepath.SkipDir
}
return nil
}
if entry.Type()&fs.ModeSymlink != 0 {
skipped = append(skipped, filepath.ToSlash(rel))
return nil
}
target := filepath.Join(dst, rel)
if entry.IsDir() {
return os.MkdirAll(target, 0o755)
}
if !entry.Type().IsRegular() {
return nil
}
info, err := entry.Info()
if err != nil {
return err
}
files++
total += info.Size()
if files > maxSkillTreeFiles {
return fmt.Errorf("skill has more than %d files; a skill directory should hold only what the skill uses", maxSkillTreeFiles)
}
if total > maxSkillTreeBytes {
return fmt.Errorf("skill is larger than %d MiB; a skill directory should hold only what the skill uses", maxSkillTreeBytes>>20)
}
if filepath.ToSlash(rel) != skillFileName {
bundled++
}
mode := os.FileMode(0o644)
if info.Mode().Perm()&0o111 != 0 {
mode = 0o755
}
return copySkillFile(path, target, mode)
})
if err != nil {
return 0, nil, err
}
return bundled, skipped, nil
}
func copySkillFile(src, dst string, mode os.FileMode) error {
in, err := os.Open(src)
if err != nil {
return err
}
defer in.Close()
out, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_EXCL, mode)
if err != nil {
return err
}
if _, err := io.Copy(out, in); err != nil {
_ = out.Close()
return err
}
if err := out.Close(); err != nil {
return err
}
// The create mode is filtered by the umask; set it outright so an
// executable script is still executable wherever the install ran.
return os.Chmod(dst, mode)
}
// skillTreeHash hashes every file a skill ships, with the same entries
// copySkillTree copies: its relative path, whether it is executable, and its
// content. Sorted by path, so the hash depends only on the tree.
func skillTreeHash(dir string) (string, error) {
type entry struct {
rel string
exec bool
sum string
}
var entries []entry
err := filepath.WalkDir(dir, func(path string, item fs.DirEntry, walkErr error) error {
if walkErr != nil {
return walkErr
}
if path == dir {
return nil
}
if skipSkillEntry(item.Name()) {
if item.IsDir() {
return filepath.SkipDir
}
return nil
}
if !item.Type().IsRegular() {
return nil
}
info, err := item.Info()
if err != nil {
return err
}
data, err := os.ReadFile(path)
if err != nil {
return err
}
rel, err := filepath.Rel(dir, path)
if err != nil {
return err
}
sum := sha256.Sum256(data)
entries = append(entries, entry{rel: filepath.ToSlash(rel), exec: info.Mode().Perm()&0o111 != 0, sum: hex.EncodeToString(sum[:])})
return nil
})
if err != nil {
return "", err
}
sort.Slice(entries, func(i, j int) bool { return entries[i].rel < entries[j].rel })
hash := sha256.New()
for _, item := range entries {
flag := "-"
if item.exec {
flag = "x"
}
// Length-prefixed, so no path can be read as the start of another entry.
fmt.Fprintf(hash, "%d:%s|%s|%s;", len(item.rel), item.rel, flag, item.sum)
}
return "sha256:" + hex.EncodeToString(hash.Sum(nil)), nil
}
// BundledFiles lists the files a skill ships besides its SKILL.md, as
// slash-separated paths relative to the skill's directory, sorted. Hidden
// entries and symlinks are left out: the listing tells the model what the
// skill brings, and is not a way to reach anything outside it. total is the
// number of files found; at most limit come back.
func BundledFiles(skill Skill, limit int) (files []string, total int, err error) {
dir := filepath.Dir(skill.Path)
if strings.TrimSpace(skill.Path) == "" {
return nil, 0, errors.New("skill has no path")
}
err = filepath.WalkDir(dir, func(path string, entry fs.DirEntry, walkErr error) error {
if walkErr != nil {
return walkErr
}
if path == dir {
return nil
}
if skipSkillEntry(entry.Name()) {
if entry.IsDir() {
return filepath.SkipDir
}
return nil
}
if !entry.Type().IsRegular() {
return nil
}
rel, err := filepath.Rel(dir, path)
if err != nil {
return err
}
rel = filepath.ToSlash(rel)
if rel == skillFileName {
return nil
}
total++
if total > maxSkillTreeFiles {
// A directory this large is not something the listing should walk
// to the end of; the count stops being exact past the cap.
return filepath.SkipAll
}
files = append(files, rel)
return nil
})
if err != nil {
return nil, 0, err
}
sort.Strings(files)
if limit >= 0 && len(files) > limit {
files = files[:limit]
}
return files, total, nil
}
// ToolOutput is what the skill tool returns for a loaded skill: its SKILL.md
// body and, when the skill ships other files, the directory they are in and
// what they are, so instructions that say "run scripts/extract.py" can be
// followed. A skill that is only a SKILL.md gets its body back unchanged.
func ToolOutput(skill Skill) string {
files, total, err := BundledFiles(skill, bundledFileListLimit)
if err != nil || total == 0 {
return skill.Content
}
var builder strings.Builder
builder.WriteString(strings.TrimRight(skill.Content, "\n"))
builder.WriteString("\n\n---\n")
fmt.Fprintf(&builder, "This skill's files are in %s. Paths in the instructions above are relative to that directory.\n", filepath.Dir(skill.Path))
for _, file := range files {
builder.WriteString("- " + file + "\n")
}
if more := total - len(files); more > 0 {
fmt.Fprintf(&builder, "- (and %d more)\n", more)
}
return builder.String()
}
+351
View File
@@ -0,0 +1,351 @@
package skills
import (
"context"
"encoding/json"
"fmt"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
)
const bundledSkillManifest = "---\nname: pdf\ndescription: fill PDF forms\n---\nRun scripts/fill.py with the form path.\n"
// writeBundledSourceSkill lays down a skill in the agentskills.io shape: a
// SKILL.md plus the scripts, references and assets it points at, and the hidden
// clutter a cloned source carries.
func writeBundledSourceSkill(t *testing.T, dir string) string {
t.Helper()
writeSourceSkill(t, dir, bundledSkillManifest)
files := map[string]string{
"scripts/fill.py": "print('fill')\n",
"references/form.md": "# Form fields\n",
"assets/logo.svg": "<svg/>\n",
"LICENSE.txt": "MIT\n",
".git/config": "[core]\n",
".env": "TOKEN=local-only\n",
"scripts/.cache/tmp": "scratch\n",
}
for rel, content := range files {
path := filepath.Join(dir, filepath.FromSlash(rel))
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
t.Fatal(err)
}
mode := os.FileMode(0o644)
if strings.HasPrefix(rel, "scripts/") {
mode = 0o755
}
if err := os.WriteFile(path, []byte(content), mode); err != nil {
t.Fatal(err)
}
}
return dir
}
func TestInstallCopiesTheSkillDirectory(t *testing.T) {
destDir := t.TempDir()
src := writeBundledSourceSkill(t, filepath.Join(t.TempDir(), "pdf"))
result, err := Install(context.Background(), InstallOptions{Source: src, Dir: destDir})
if err != nil {
t.Fatalf("Install: %v", err)
}
installed := filepath.Join(destDir, "pdf")
for _, rel := range []string{"SKILL.md", "scripts/fill.py", "references/form.md", "assets/logo.svg", "LICENSE.txt"} {
if _, err := os.Stat(filepath.Join(installed, filepath.FromSlash(rel))); err != nil {
t.Errorf("%s was not installed: %v", rel, err)
}
}
for _, rel := range []string{".git", ".env", "scripts/.cache"} {
if _, err := os.Stat(filepath.Join(installed, filepath.FromSlash(rel))); err == nil {
t.Errorf("hidden %s was installed; it is not part of the skill", rel)
}
}
if !strings.HasPrefix(result.TreeHash, "sha256:") {
t.Fatalf("TreeHash = %q, want a sha256 tree hash", result.TreeHash)
}
lock, err := ReadLock(destDir)
if err != nil {
t.Fatal(err)
}
if lock["pdf"].TreeHash != result.TreeHash || lock["pdf"].Hash != result.Hash {
t.Fatalf("lock entry = %+v, want the install's hashes %q and %q", lock["pdf"], result.Hash, result.TreeHash)
}
if runtime.GOOS == "windows" {
return
}
script, err := os.Stat(filepath.Join(installed, "scripts", "fill.py"))
if err != nil {
t.Fatal(err)
}
if script.Mode().Perm() != 0o755 {
t.Errorf("installed script mode = %v, want 0755 so the skill can still run it", script.Mode().Perm())
}
reference, err := os.Stat(filepath.Join(installed, "references", "form.md"))
if err != nil {
t.Fatal(err)
}
if reference.Mode().Perm() != 0o644 {
t.Errorf("installed reference mode = %v, want 0644", reference.Mode().Perm())
}
}
func TestInstallSkipsSymlinksAndSaysSo(t *testing.T) {
destDir := t.TempDir()
src := writeBundledSourceSkill(t, filepath.Join(t.TempDir(), "pdf"))
outside := filepath.Join(t.TempDir(), "secret.txt")
if err := os.WriteFile(outside, []byte("not part of the skill\n"), 0o600); err != nil {
t.Fatal(err)
}
if err := os.Symlink(outside, filepath.Join(src, "references", "linked.md")); err != nil {
t.Skipf("symlinks unavailable: %v", err)
}
result, err := Install(context.Background(), InstallOptions{Source: src, Dir: destDir})
if err != nil {
t.Fatalf("Install: %v", err)
}
if len(result.SkippedLinks) != 1 || result.SkippedLinks[0] != "references/linked.md" {
t.Fatalf("SkippedLinks = %v, want [references/linked.md]", result.SkippedLinks)
}
if _, err := os.Lstat(filepath.Join(destDir, "pdf", "references", "linked.md")); err == nil {
t.Fatal("a symlink from the source was installed")
}
}
func TestInstallRefusesAnOversizedSkill(t *testing.T) {
for _, tc := range []struct {
name string
files int
bytes int64
want string
}{
{name: "too many files", files: 3, bytes: maxSkillTreeBytes, want: "more than 3 files"},
{name: "too many bytes", files: maxSkillTreeFiles, bytes: 1 << 20, want: "larger than 1 MiB"},
} {
t.Run(tc.name, func(t *testing.T) {
files, bytes := maxSkillTreeFiles, maxSkillTreeBytes
t.Cleanup(func() { maxSkillTreeFiles, maxSkillTreeBytes = files, bytes })
maxSkillTreeFiles, maxSkillTreeBytes = tc.files, tc.bytes
destDir := t.TempDir()
src := writeBundledSourceSkill(t, filepath.Join(t.TempDir(), "pdf"))
big := make([]byte, 2<<20)
if err := os.WriteFile(filepath.Join(src, "assets", "big.bin"), big, 0o644); err != nil {
t.Fatal(err)
}
_, err := Install(context.Background(), InstallOptions{Source: src, Dir: destDir})
if err == nil || !strings.Contains(err.Error(), tc.want) {
t.Fatalf("Install error = %v, want %q", err, tc.want)
}
if _, statErr := os.Stat(filepath.Join(destDir, "pdf")); statErr == nil {
t.Fatal("a refused install left a skill directory behind")
}
})
}
}
func TestReinstallWithOnlyABundledFileChangedIsAnUpdate(t *testing.T) {
destDir := t.TempDir()
src := writeBundledSourceSkill(t, filepath.Join(t.TempDir(), "pdf"))
first, err := Install(context.Background(), InstallOptions{Source: src, Dir: destDir})
if err != nil {
t.Fatalf("first Install: %v", err)
}
if err := os.WriteFile(filepath.Join(src, "scripts", "fill.py"), []byte("print('fill v2')\n"), 0o755); err != nil {
t.Fatal(err)
}
second, err := Install(context.Background(), InstallOptions{Source: src, Dir: destDir})
if err != nil {
t.Fatalf("second Install: %v", err)
}
if !second.Updated {
t.Fatal("a reinstall that changed a bundled script was not reported as an update")
}
if second.Hash != first.Hash || second.TreeHash == first.TreeHash {
t.Fatalf("hashes = %q/%q after %q/%q, want the same SKILL.md hash and a new tree hash", second.Hash, second.TreeHash, first.Hash, first.TreeHash)
}
if second.PreviousHash != first.Hash || second.PreviousTreeHash != first.TreeHash {
t.Fatalf("previous hashes = %q/%q, want the first install's %q/%q", second.PreviousHash, second.PreviousTreeHash, first.Hash, first.TreeHash)
}
got, err := os.ReadFile(filepath.Join(destDir, "pdf", "scripts", "fill.py"))
if err != nil || string(got) != "print('fill v2')\n" {
t.Fatalf("installed script = %q, %v; want the new version", got, err)
}
}
func TestInfoReportsDriftInABundledFile(t *testing.T) {
destDir := t.TempDir()
src := writeBundledSourceSkill(t, filepath.Join(t.TempDir(), "pdf"))
if _, err := Install(context.Background(), InstallOptions{Source: src, Dir: destDir}); err != nil {
t.Fatalf("Install: %v", err)
}
if info, ok := Info(destDir, "pdf"); !ok || info.HashDrift {
t.Fatalf("a fresh install reports drift: %+v, %v", info, ok)
}
script := filepath.Join(destDir, "pdf", "scripts", "fill.py")
if err := os.WriteFile(script, []byte("print('edited by hand')\n"), 0o755); err != nil {
t.Fatal(err)
}
if info, _ := Info(destDir, "pdf"); !info.HashDrift {
t.Fatal("an edited bundled script is not reported as drift")
}
}
// A lock entry written before installs carried their bundled files has no tree
// hash. Checking its tree would report every such skill as drifted, so only
// its SKILL.md is checked, as before.
func TestLockEntryWithoutATreeHashChecksOnlySKILLmd(t *testing.T) {
destDir := t.TempDir()
src := writeBundledSourceSkill(t, filepath.Join(t.TempDir(), "pdf"))
if _, err := Install(context.Background(), InstallOptions{Source: src, Dir: destDir}); err != nil {
t.Fatalf("Install: %v", err)
}
dropTreeHash(t, destDir, "pdf")
if err := os.WriteFile(filepath.Join(destDir, "pdf", "scripts", "fill.py"), []byte("print('edited')\n"), 0o755); err != nil {
t.Fatal(err)
}
if info, _ := Info(destDir, "pdf"); info.HashDrift {
t.Fatal("an entry without a tree hash reported drift for a bundled file")
}
if err := os.WriteFile(filepath.Join(destDir, "pdf", skillFileName), []byte(bundledSkillManifest+"edited\n"), 0o644); err != nil {
t.Fatal(err)
}
if info, _ := Info(destDir, "pdf"); !info.HashDrift {
t.Fatal("an edited SKILL.md under an old lock entry is not reported as drift")
}
}
// Reinstalling over an entry from before tree hashes compares with what that
// install left, SKILL.md alone: an unchanged skill that is only a SKILL.md is
// not an update, and one whose bundled files this install adds is.
func TestReinstallOverAnEntryWithoutATreeHash(t *testing.T) {
t.Run("SKILL.md only", func(t *testing.T) {
destDir := t.TempDir()
src := writeSourceSkill(t, filepath.Join(t.TempDir(), "pdf"), bundledSkillManifest)
if _, err := Install(context.Background(), InstallOptions{Source: src, Dir: destDir}); err != nil {
t.Fatalf("first Install: %v", err)
}
dropTreeHash(t, destDir, "pdf")
result, err := Install(context.Background(), InstallOptions{Source: src, Dir: destDir})
if err != nil {
t.Fatalf("second Install: %v", err)
}
if result.Updated {
t.Fatalf("reinstalling an unchanged SKILL.md-only skill over an old entry was reported as an update: %+v", result)
}
lock, err := ReadLock(destDir)
if err != nil {
t.Fatal(err)
}
if lock["pdf"].TreeHash != result.TreeHash {
t.Fatalf("lock tree hash = %q after the reinstall, want %q", lock["pdf"].TreeHash, result.TreeHash)
}
})
t.Run("with bundled files", func(t *testing.T) {
destDir := t.TempDir()
src := writeBundledSourceSkill(t, filepath.Join(t.TempDir(), "pdf"))
first, err := Install(context.Background(), InstallOptions{Source: src, Dir: destDir})
if err != nil {
t.Fatalf("first Install: %v", err)
}
dropTreeHash(t, destDir, "pdf")
result, err := Install(context.Background(), InstallOptions{Source: src, Dir: destDir})
if err != nil {
t.Fatalf("second Install: %v", err)
}
if !result.Updated || result.Hash != first.Hash || result.PreviousHash != first.Hash || result.PreviousTreeHash != "" {
t.Fatalf("result = %+v, want an update with SKILL.md hash %q unchanged and no previous tree hash", result, first.Hash)
}
})
}
// dropTreeHash rewrites the lock entry for name the way an install from before
// tree hashes left it: a source and a SKILL.md hash, nothing for the tree.
func dropTreeHash(t *testing.T, destDir, name string) {
t.Helper()
lock, err := ReadLock(destDir)
if err != nil {
t.Fatal(err)
}
entry := lock[name]
entry.TreeHash = ""
lock[name] = entry
data, err := json.Marshal(lock)
if err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(destDir, LockFileName), data, 0o644); err != nil {
t.Fatal(err)
}
}
func TestToolOutputSaysWhereTheBundledFilesAre(t *testing.T) {
root := t.TempDir()
writeBundledSourceSkill(t, filepath.Join(root, "pdf"))
skill, ok := Get(root, "pdf")
if !ok {
t.Fatal("skill not loaded")
}
output := ToolOutput(skill)
if !strings.HasPrefix(output, strings.TrimRight(skill.Content, "\n")) {
t.Fatalf("the instructions do not come first:\n%s", output)
}
if !strings.Contains(output, filepath.Dir(skill.Path)) {
t.Fatalf("the output does not say where the skill's files are:\n%s", output)
}
for _, want := range []string{"- LICENSE.txt", "- assets/logo.svg", "- references/form.md", "- scripts/fill.py"} {
if !strings.Contains(output, want+"\n") {
t.Errorf("output does not list %q:\n%s", want, output)
}
}
for _, hidden := range []string{".env", ".git", ".cache", "SKILL.md\n"} {
if strings.Contains(output, "- "+hidden) {
t.Errorf("output lists %q:\n%s", hidden, output)
}
}
}
func TestToolOutputCapsALongListing(t *testing.T) {
root := t.TempDir()
dir := writeSourceSkill(t, filepath.Join(root, "many"), "---\nname: many\n---\nbody\n")
for i := 0; i < bundledFileListLimit+7; i++ {
path := filepath.Join(dir, "references", fmt.Sprintf("part-%03d.md", i))
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("x\n"), 0o644); err != nil {
t.Fatal(err)
}
}
skill, ok := Get(root, "many")
if !ok {
t.Fatal("skill not loaded")
}
output := ToolOutput(skill)
if got := strings.Count(output, "\n- references/"); got != bundledFileListLimit {
t.Fatalf("listed %d files, want the first %d:\n%s", got, bundledFileListLimit, output)
}
if !strings.Contains(output, "- (and 7 more)\n") {
t.Fatalf("output does not say how many more there are:\n%s", output)
}
}
func TestToolOutputOfAManifestOnlySkillIsItsContent(t *testing.T) {
root := t.TempDir()
writeSourceSkill(t, filepath.Join(root, "plain"), "---\nname: plain\n---\nJust instructions.\n")
skill, ok := Get(root, "plain")
if !ok {
t.Fatal("skill not loaded")
}
if got := ToolOutput(skill); got != skill.Content {
t.Fatalf("ToolOutput = %q, want the unchanged content %q", got, skill.Content)
}
}
+2 -1
View File
@@ -15,9 +15,10 @@ func Info(dir string, name string) (SkillInfo, bool) {
if entry, found := lock[skill.Name]; found {
info.Source = entry.Source
info.Hash = entry.Hash
info.TreeHash = entry.TreeHash
}
}
info.HashDrift = skillHashDrift(skill, info.Hash)
info.HashDrift = skillHashDrift(skill, info.Hash, info.TreeHash)
return info, true
}
+93 -28
View File
@@ -1,9 +1,9 @@
package skills
// Distribution: install a skill from a git URL or a local path into the skills
// directory, with a content hash recorded in a lockfile (skills.lock) so every
// install/update is verifiable. Skills are markdown, so install NEVER executes
// fetched content — it copies and validates the SKILL.md and nothing else.
// directory, with content hashes recorded in a lockfile (skills.lock) so every
// install/update is verifiable. Install copies the skill's directory, its
// SKILL.md and the files beside it, and NEVER executes fetched content.
import (
"context"
@@ -61,16 +61,30 @@ type InstallResult struct {
Hash string `json:"hash"`
// Source echoes the source the skill was installed from.
Source string `json:"source"`
// TreeHash covers every file the skill ships, SKILL.md included; Hash is
// still the SKILL.md hash alone.
TreeHash string `json:"treeHash"`
// Updated is true when an existing install was replaced; PreviousHash then
// carries the prior recorded hash so a caller can show the change.
Updated bool `json:"updated"`
PreviousHash string `json:"previousHash,omitempty"`
// PreviousTreeHash is the tree hash the replaced install recorded. It is
// empty when that install predates tree hashes, and so copied SKILL.md
// alone.
PreviousTreeHash string `json:"previousTreeHash,omitempty"`
// SkippedLinks lists symlinks in the source, relative to the skill
// directory. Install neither follows nor copies them.
SkippedLinks []string `json:"skippedLinks,omitempty"`
}
// LockEntry records the source and content hash for one installed skill.
type LockEntry struct {
Source string `json:"source"`
Hash string `json:"hash"`
// TreeHash covers the whole installed skill directory. Entries written
// before skills installed their bundled files have none, and are checked on
// their SKILL.md alone.
TreeHash string `json:"treeHash,omitempty"`
}
// SkillInfo bundles a discovered skill with its recorded source and hash, for
@@ -79,12 +93,14 @@ type SkillInfo struct {
Skill Skill `json:"skill"`
Source string `json:"source,omitempty"`
Hash string `json:"hash,omitempty"`
TreeHash string `json:"treeHash,omitempty"`
HashDrift bool `json:"hashDrift"`
}
// Install fetches the skill at options.Source and copies its SKILL.md into
// options.Dir/<name>/, validating the frontmatter and recording a content hash
// in the lockfile. A git URL is fetched via the (injectable) GitRunner into a
// Install fetches the skill at options.Source and copies its directory into
// options.Dir/<name>/: SKILL.md and the scripts, references and assets beside
// it. It validates the frontmatter and records the SKILL.md hash and a hash of
// the whole tree in the lockfile. A git URL is fetched via the (injectable) GitRunner into a
// temp dir; a local path is read in place. Fetched content is never executed.
func Install(ctx context.Context, options InstallOptions) (InstallResult, error) {
source := strings.TrimSpace(options.Source)
@@ -133,10 +149,11 @@ func Install(ctx context.Context, options InstallOptions) (InstallResult, error)
if err := os.MkdirAll(staged, 0o755); err != nil {
return InstallResult{}, fmt.Errorf("create staged skill dir: %w", err)
}
stagedManifest := filepath.Join(staged, skillFileName)
if err := os.WriteFile(stagedManifest, data, 0o644); err != nil {
return InstallResult{}, fmt.Errorf("stage SKILL.md: %w", err)
bundled, skippedLinks, err := copySkillTree(skillDir, staged)
if err != nil {
return InstallResult{}, fmt.Errorf("stage skill directory: %w", err)
}
stagedManifest := filepath.Join(staged, skillFileName)
stagedData, err := os.ReadFile(stagedManifest)
if err != nil || hashContent(stagedData) != hash {
if err != nil {
@@ -169,8 +186,13 @@ func Install(ctx context.Context, options InstallOptions) (InstallResult, error)
return InstallResult{}, fmt.Errorf("%w: %q is installed from %s (use --force to overwrite)", ErrNameClash, name, previous.Source)
}
treeHash, err := skillTreeHash(staged)
if err != nil {
return InstallResult{}, fmt.Errorf("hash staged skill directory: %w", err)
}
target := filepath.Join(dir, name)
lock[name] = LockEntry{Source: source, Hash: hash}
lock[name] = LockEntry{Source: source, Hash: hash, TreeHash: treeHash}
if err := installtxn.CommitDir(target, staged, func() error {
return writeLock(dir, lock)
}); err != nil {
@@ -178,14 +200,26 @@ func Install(ctx context.Context, options InstallOptions) (InstallResult, error)
}
result := InstallResult{
Name: name,
Path: filepath.Join(target, skillFileName),
Hash: hash,
Source: source,
Name: name,
Path: filepath.Join(target, skillFileName),
Hash: hash,
TreeHash: treeHash,
Source: source,
SkippedLinks: skippedLinks,
}
if existed {
result.Updated = previous.Hash != hash
result.PreviousHash = previous.Hash
result.PreviousTreeHash = previous.TreeHash
if previous.TreeHash == "" {
// An entry from before tree hashes is from an install that copied
// SKILL.md alone. Against that, this install is an update when
// SKILL.md changed or it now brings bundled files; a skill that is
// still only its unchanged SKILL.md is not.
result.Updated = previous.Hash != hash || bundled > 0
} else {
// A changed script with the same SKILL.md is still an update.
result.Updated = previous.Hash != hash || previous.TreeHash != treeHash
}
}
return result, nil
}
@@ -193,8 +227,8 @@ func Install(ctx context.Context, options InstallOptions) (InstallResult, error)
// lockReconciler answers, for one interrupted workspace, which of the two trees
// beside each other the lockfile records. A kill between the tree swap and the
// lockfile publish leaves exactly what a kill after both leaves, so filesystem
// shape cannot tell them apart and only the recorded hash can: the tree it
// describes is the one the user was told they had.
// shape cannot tell them apart and only the recorded hashes can: the tree they
// describe is the one the user was told they had.
func lockReconciler(dir string) installtxn.Reconciler {
return func(name string, target string, backup string) (installtxn.Phase, error) {
lock, err := ReadLock(dir)
@@ -215,24 +249,45 @@ func lockReconciler(dir string) installtxn.Reconciler {
// recorded one and recovery must not pick between them.
return installtxn.PhaseUnknown, nil
}
targetHash, err := manifestHash(target)
committed, err := treeMatchesEntry(target, entry)
if err != nil {
return installtxn.PhaseUnknown, err
}
if targetHash == entry.Hash {
if committed {
return installtxn.PhaseCommitted, nil
}
backupHash, err := manifestHash(backup)
prePublish, err := treeMatchesEntry(backup, entry)
if err != nil {
return installtxn.PhaseUnknown, err
}
if backupHash == entry.Hash {
if prePublish {
return installtxn.PhasePrePublish, nil
}
return installtxn.PhaseUnknown, nil
}
}
// treeMatchesEntry reports whether the skill tree at path is the one entry
// records. A reinstall that changes only a bundled file leaves SKILL.md the same
// in both trees of an interrupted workspace, so an entry with a tree hash has to
// match on that as well. An entry from before tree hashes records SKILL.md
// alone, and that is all it can be matched on. As in manifestHash, a tree with
// no SKILL.md is a non-match and any other read error is returned.
func treeMatchesEntry(path string, entry LockEntry) (bool, error) {
sum, err := manifestHash(path)
if err != nil || sum != entry.Hash {
return false, err
}
if entry.TreeHash == "" {
return true, nil
}
tree, err := skillTreeHash(path)
if err != nil {
return false, err
}
return tree == entry.TreeHash, nil
}
// manifestHash hashes the SKILL.md under path the same way an install records
// it. A tree with no SKILL.md is a non-match rather than a failure: an install
// killed mid-copy can leave one, and it simply is not the tree the lockfile
@@ -346,9 +401,10 @@ func InfoFromRoots(primaryDir string, roots []string, name string) (SkillInfo, b
if entry, found := lock[skill.Name]; found {
info.Source = entry.Source
info.Hash = entry.Hash
info.TreeHash = entry.TreeHash
}
}
info.HashDrift = skillHashDrift(skill, info.Hash)
info.HashDrift = skillHashDrift(skill, info.Hash, info.TreeHash)
return info, true
}
@@ -556,11 +612,12 @@ func hashContent(data []byte) string {
return "sha256:" + hex.EncodeToString(sum[:])
}
// skillHashDrift reports whether the on-disk SKILL.md no longer matches the
// lockfile hash recorded at install time. Missing lock hashes never count as
// drift (agents-only / unlocked skills). A locked skill whose SKILL.md cannot
// be read is treated as drifted — not as a clean match.
func skillHashDrift(skill Skill, lockHash string) bool {
// skillHashDrift reports whether the installed skill no longer matches the
// lockfile hashes recorded at install time: the SKILL.md hash, and the tree
// hash when the entry has one, so an edited script counts as drift too. Missing
// lock hashes never count as drift (skills Zero did not install). A locked
// skill that cannot be read is treated as drifted, not as a clean match.
func skillHashDrift(skill Skill, lockHash string, lockTreeHash string) bool {
lockHash = strings.TrimSpace(lockHash)
if lockHash == "" {
return false
@@ -569,5 +626,13 @@ func skillHashDrift(skill Skill, lockHash string) bool {
if err != nil {
return true
}
return hashContent(data) != lockHash
if hashContent(data) != lockHash {
return true
}
lockTreeHash = strings.TrimSpace(lockTreeHash)
if lockTreeHash == "" {
return false
}
treeHash, err := skillTreeHash(filepath.Dir(skill.Path))
return err != nil || treeHash != lockTreeHash
}
+147 -6
View File
@@ -364,10 +364,11 @@ func TestInstallGitSourceUsesRunnerAndDoesNotExecuteContent(t *testing.T) {
if _, err := os.Stat(filepath.Join(cloneRoot, "PWNED")); err == nil {
t.Fatalf("install must never execute fetched content")
}
// The fetched executable must not be copied into the skills dir either (skills
// are markdown — only SKILL.md is installed, not arbitrary fetched files).
if _, err := os.Stat(filepath.Join(destDir, "remote", "install.sh")); err == nil {
t.Fatalf("install must not copy fetched executable into the skills dir")
// A script beside SKILL.md is part of the skill (#584), so it is installed,
// and still only copied: running it is up to whoever uses the skill later,
// through the normal tool permissions.
if _, err := os.Stat(filepath.Join(destDir, "remote", "install.sh")); err != nil {
t.Fatalf("the skill's bundled script was not installed: %v", err)
}
}
@@ -440,10 +441,10 @@ func TestInfoReturnsFrontmatterSourceAndHash(t *testing.T) {
func TestSkillHashDriftUnreadableLockedPath(t *testing.T) {
missing := filepath.Join(t.TempDir(), "missing", "SKILL.md")
if !skillHashDrift(Skill{Path: missing}, "sha256:deadbeef") {
if !skillHashDrift(Skill{Path: missing}, "sha256:deadbeef", "") {
t.Fatal("expected drift when locked SKILL.md cannot be read")
}
if skillHashDrift(Skill{Path: missing}, "") {
if skillHashDrift(Skill{Path: missing}, "", "") {
t.Fatal("missing lock hash must not count as drift")
}
}
@@ -861,6 +862,146 @@ func TestRecoveryResolvesEveryInterruptedUpdateStep(t *testing.T) {
}
}
const (
bundledScriptV1 = "print('v1')\n"
bundledScriptV2 = "print('v2')\n"
)
func writeBundledScript(t *testing.T, skillDir string, content string) string {
t.Helper()
path := filepath.Join(skillDir, "scripts", "run.py")
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte(content), 0o755); err != nil {
t.Fatal(err)
}
return path
}
// plantInterruptedBundledUpdate is plantInterruptedAlphaUpdate for a reinstall
// that changes only a bundled script, killed after the tree swap: alpha's
// SKILL.md is the same in the live tree and the backup, and only the tree hash
// in the lockfile tells them apart. published says whether the lockfile was
// rewritten for the new tree before the kill (S4) or not (S3).
func plantInterruptedBundledUpdate(t *testing.T, published bool) interruptedUpdate {
t.Helper()
dir := t.TempDir()
source := writeSourceSkill(t, filepath.Join(t.TempDir(), "alpha"), oldAlphaSkill)
writeBundledScript(t, source, bundledScriptV1)
if _, err := Install(context.Background(), InstallOptions{Source: source, Dir: dir}); err != nil {
t.Fatalf("seed alpha: %v", err)
}
gammaSource := writeSourceSkill(t, filepath.Join(t.TempDir(), "gamma"), gammaSkill)
if _, err := Install(context.Background(), InstallOptions{Source: gammaSource, Dir: dir}); err != nil {
t.Fatalf("seed gamma: %v", err)
}
staged, _, err := installtxn.StageDir(dir)
if err != nil {
t.Fatal(err)
}
workspace := filepath.Dir(staged)
writeWorkspaceMarker(t, workspace, "alpha")
writeSourceSkill(t, staged, oldAlphaSkill)
writeBundledScript(t, staged, bundledScriptV2)
target := filepath.Join(dir, "alpha")
if err := os.Rename(target, filepath.Join(workspace, "previous")); err != nil {
t.Fatal(err)
}
if err := os.Rename(staged, target); err != nil {
t.Fatal(err)
}
if published {
treeHash, err := skillTreeHash(target)
if err != nil {
t.Fatal(err)
}
publishLockEntry(t, dir, "alpha", LockEntry{Source: canonicalSource(source), Hash: hashContent([]byte(oldAlphaSkill)), TreeHash: treeHash})
}
return interruptedUpdate{dir: dir, workspace: workspace, oldSource: canonicalSource(source), newSource: canonicalSource(source)}
}
// SKILL.md cannot say which tree of an interrupted bundled-only update the
// lockfile records, since both carry the same one. The tree hash can: before the
// publish it still describes the backup, after it the live tree.
func TestRecoveryTellsABundledOnlyUpdateApartByTreeHash(t *testing.T) {
steps := []struct {
step string
published bool
want string
}{
{step: "S3", published: false, want: bundledScriptV1},
{step: "S4", published: true, want: bundledScriptV2},
}
for _, step := range steps {
for _, driver := range recoveryDrivers {
t.Run(step.step+"/"+driver.name, func(t *testing.T) {
update := plantInterruptedBundledUpdate(t, step.published)
if err := driver.run(t, update.dir); err != nil {
t.Fatalf("%s over a recoverable state: %v", driver.name, err)
}
target := filepath.Join(update.dir, "alpha")
got, err := os.ReadFile(filepath.Join(target, "scripts", "run.py"))
if err != nil || string(got) != step.want {
t.Fatalf("recovered script = %q, %v; want %q", got, err, step.want)
}
lock, err := ReadLock(update.dir)
if err != nil {
t.Fatal(err)
}
tree, err := skillTreeHash(target)
if err != nil {
t.Fatal(err)
}
if lock["alpha"].TreeHash != tree {
t.Errorf("lock tree hash %q does not describe the recovered tree %q", lock["alpha"].TreeHash, tree)
}
if _, err := os.Stat(update.workspace); !os.IsNotExist(err) {
t.Errorf("a resolved workspace must be cleared, got %v", err)
}
})
}
}
}
// A live tree whose files cannot be read cannot be shown to be the one the
// lockfile records, or not to be. Recovery reports it, as it does for a SKILL.md
// it cannot read, and decides nothing until a later pass can read it.
func TestRecoveryReportsALiveTreeItCannotHash(t *testing.T) {
for _, driver := range recoveryDrivers {
t.Run(driver.name, func(t *testing.T) {
update := plantInterruptedBundledUpdate(t, false)
lockBefore := readLockBytes(t, update.dir)
dirBefore := treeSnapshot(t, update.dir)
release := holdUnreadable(t, filepath.Join(update.dir, "alpha", "scripts", "run.py"))
err := driver.run(t, update.dir)
release()
if err == nil || !strings.Contains(err.Error(), "classify interrupted install alpha") {
t.Fatalf("%s over a live tree it cannot hash: %v, want the classification error", driver.name, err)
}
assertSnapshotUnchanged(t, dirBefore, treeSnapshot(t, update.dir))
if got := readLockBytes(t, update.dir); got != lockBefore {
t.Errorf("the lockfile was rewritten by a caller that must have aborted:\n%s", got)
}
// Once the tree can be read, the next pass resolves it: the lockfile
// still describes the backup, so the backup comes back.
if err := driver.run(t, update.dir); err != nil {
t.Fatalf("%s once the tree is readable: %v", driver.name, err)
}
got, err := os.ReadFile(filepath.Join(update.dir, "alpha", "scripts", "run.py"))
if err != nil || string(got) != bundledScriptV1 {
t.Fatalf("recovered script = %q, %v; want %q", got, err, bundledScriptV1)
}
})
}
}
// A recovery that could not be made is not the same as nothing to recover.
// Installing over or reporting the removal of a tree that is still owed a
// restore destroys the only copy of it, so the caller has to stop before it
+30
View File
@@ -0,0 +1,30 @@
//go:build !windows
package skills
import (
"os"
"sync"
"testing"
)
// holdUnreadable takes every permission off path, so reading it fails until
// release puts them back. Cleanup releases it too. Root reads a file whatever
// its mode, so a test that needs the read to fail is skipped there.
func holdUnreadable(t *testing.T, path string) (release func()) {
t.Helper()
if os.Geteuid() == 0 {
t.Skip("root reads a file whatever its mode")
}
info, err := os.Stat(path)
if err != nil {
t.Fatal(err)
}
if err := os.Chmod(path, 0); err != nil {
t.Fatal(err)
}
var once sync.Once
release = func() { once.Do(func() { _ = os.Chmod(path, info.Mode().Perm()) }) }
t.Cleanup(release)
return release
}
@@ -0,0 +1,28 @@
//go:build windows
package skills
import (
"sync"
"testing"
"golang.org/x/sys/windows"
)
// holdUnreadable keeps path open with no sharing, so every other open of it
// fails until release runs. Cleanup releases it too.
func holdUnreadable(t *testing.T, path string) (release func()) {
t.Helper()
name, err := windows.UTF16PtrFromString(path)
if err != nil {
t.Fatal(err)
}
handle, err := windows.CreateFile(name, windows.GENERIC_READ, 0, nil, windows.OPEN_EXISTING, windows.FILE_ATTRIBUTE_NORMAL, 0)
if err != nil {
t.Fatalf("hold %s open: %v", path, err)
}
var once sync.Once
release = func() { once.Do(func() { _ = windows.CloseHandle(handle) }) }
t.Cleanup(release)
return release
}
+179 -17
View File
@@ -1,7 +1,9 @@
package specialist
import (
"context"
"encoding/json"
"os"
"strings"
"sync"
@@ -31,18 +33,41 @@ type specialistAccountingInput struct {
Mode string
Background bool
PID int
// LaunchFailed marks a child that never started (e.g. the command failed to
// exec). The internal start/stop accounting still records the attempt, but no
// public specialistStart/specialistStop hook fires for it: there was no
// sub-agent lifecycle to observe. The zero value is "started", so a caller
// that forgets to set it never silently suppresses hooks.
LaunchFailed bool
}
// recordSpecialistStart records the internal attempt accounting. It deliberately
// does NOT dispatch the public specialistStart hook: the hook fires only once
// the child has actually started (dispatchSpecialistStartHook), so a failed
// launch cannot announce a sub-agent that never existed.
func (executor Executor) recordSpecialistStart(input specialistAccountingInput) {
payload := baseSpecialistPayload(input)
_, _ = appendSpecialistSessionEvent(executor.SessionStore, input.ParentSessionID, sessions.EventSpecialistStart, payload)
}
// dispatchSpecialistStartHook publishes the specialistStart hook for a child
// whose startup was confirmed. Callers pass the same accounting input used for
// the attempt record so payload and matcher subject stay identical.
func (executor Executor) dispatchSpecialistStartHook(input specialistAccountingInput) {
if input.LaunchFailed || executor.LifecycleHooks == nil || executor.LifecycleHooks.Dispatch == nil {
return
}
executor.dispatchLifecycleHook("specialistStart", input, baseSpecialistPayload(input))
}
func (executor Executor) recordSpecialistStop(input specialistAccountingInput, summary StreamResult, status string, exitCode int, runErr error, usageRolledUp bool) {
store := accountingStore(executor.SessionStore)
accountingMu.Lock()
defer accountingMu.Unlock()
payload := baseSpecialistPayload(input)
// childStarted is the durable eligibility flag for the public stop hook: a
// later reaper (TaskOutput poll, recovered TaskStop, onExit in another
// process) derives whether the terminal transition deserves a hook from the
// recorded stop event, since it cannot know whether the child ever started.
payload["childStarted"] = !input.LaunchFailed
if summary.RunID != "" {
payload["runId"] = summary.RunID
}
@@ -55,9 +80,127 @@ func (executor Executor) recordSpecialistStop(input specialistAccountingInput, s
if len(summary.Errors) > 0 {
payload["errors"] = append([]string(nil), summary.Errors...)
}
// Atomic check+append under the session lock so a concurrent stop path cannot
// also pass the existence check and write a duplicate stop event.
// Hold accountingMu only for the deduped append so a slow lifecycle hook cannot
// stall concurrent specialist stop/usage accounting.
accountingMu.Lock()
_, _ = appendSpecialistEventOnce(store, input.ParentSessionID, sessions.EventSpecialistStop, payload, input.ChildSessionID, summary.RunID)
accountingMu.Unlock()
executor.dispatchSpecialistStopHook(store, input, payload, summary.RunID)
}
// specialistStopHookEvent durably records that the specialistStop hook dispatch
// decision for one child/run was settled. It lives in the parent session's event
// log - the same cross-process, lock-guarded ledger the accounting dedup uses -
// so a later Zero process reaping the same task (fresh LifecycleHooks bridge)
// does not re-run the hook. This is internal bookkeeping, not a lifecycle event.
const specialistStopHookEvent sessions.EventType = "specialist_stop_hook"
// dispatchSpecialistStopHook runs the specialistStop hook at most once per
// child/run. The claim follows the accounting equivalence rules (an empty runId
// is a catch-all for the same child), so concurrent or sequential finishers -
// Task.onExit, a TaskOutput poll in another process, a recovered TaskStop - share
// one durable claim. Swarm members have no parent session (nothing durable to
// key on), so they fall back to this bridge's in-process claim, which mirrors
// the same run-ID equivalence.
//
// Hooks are external side effects, so on a store error the durable claim cannot
// be established; the hook then falls back to the in-process claim rather than
// running unchecked, keeping at-most-once within this process. A won durable
// claim also takes the in-process claim, so a transient store error before or
// after it on the same bridge cannot dispatch the same stop a second time.
func (executor Executor) dispatchSpecialistStopHook(store *sessions.Store, input specialistAccountingInput, payload map[string]any, runID string) {
hooks := executor.LifecycleHooks
if hooks == nil || hooks.Dispatch == nil {
return
}
started := !input.LaunchFailed
parentSessionID := strings.TrimSpace(input.ParentSessionID)
if parentSessionID == "" || !sessions.ValidSessionID(parentSessionID) {
if started && hooks.claimInProcessStopDispatch(input.ChildSessionID, runID) {
executor.dispatchLifecycleHook("specialistStop", input, payload)
}
return
}
if recorded, ok := recordedStopChildStarted(store, parentSessionID, input.ChildSessionID, runID); ok {
started = recorded
}
claimed, err := appendSpecialistEventOnce(store, parentSessionID, specialistStopHookEvent, stopHookClaimPayload(input, runID, started), input.ChildSessionID, runID)
if err != nil {
// The durable ledger is unavailable. Keep the hook advisory but bounded to
// a single dispatch in this process rather than dropping it or risking an
// unbounded duplicate.
if started && hooks.claimInProcessStopDispatch(input.ChildSessionID, runID) {
executor.dispatchLifecycleHook("specialistStop", input, payload)
}
return
}
if !claimed {
return
}
// Take the in-process claim even for an ineligible stop: this bridge has
// settled the decision, and a later fallback must not re-decide it.
if hooks.claimInProcessStopDispatch(input.ChildSessionID, runID) && started {
executor.dispatchLifecycleHook("specialistStop", input, payload)
}
}
// recordedStopChildStarted reads the first recorded stop event matching the
// child/run equivalence rules and reports its childStarted eligibility flag.
// ok is false when no stop event exists yet or it predates the flag, in which
// case the caller keeps its locally known eligibility.
func recordedStopChildStarted(store *sessions.Store, parentSessionID, childSessionID, runID string) (started bool, ok bool) {
store = accountingStore(store)
childSessionID = strings.TrimSpace(childSessionID)
if childSessionID == "" {
return false, false
}
runID = strings.TrimSpace(runID)
events, err := store.ReadEvents(parentSessionID)
if err != nil {
return false, false
}
for _, event := range events {
if event.Type != sessions.EventSpecialistStop {
continue
}
var payload map[string]any
if err := json.Unmarshal(event.Payload, &payload); err != nil {
continue
}
if !specialistPayloadMatches(payload, childSessionID, runID) {
continue
}
if started, ok := payload["childStarted"].(bool); ok {
return started, true
}
return false, false
}
return false, false
}
func stopHookClaimPayload(input specialistAccountingInput, runID string, started bool) map[string]any {
payload := baseSpecialistPayload(input)
payload["hook"] = "specialistStop"
// childStarted is the resolved eligibility (locally known or read back from
// the durable stop event); dispatched records that a hook was run for it.
payload["childStarted"] = started
payload["dispatched"] = started
if runID = strings.TrimSpace(runID); runID != "" {
payload["runId"] = runID
}
return payload
}
func (executor Executor) dispatchLifecycleHook(event string, input specialistAccountingInput, payload map[string]any) {
if executor.LifecycleHooks == nil || executor.LifecycleHooks.Dispatch == nil {
return
}
hookPayload := make(map[string]any, len(payload)+1)
for key, value := range payload {
hookPayload[key] = value
}
hookPayload["event"] = event
executor.LifecycleHooks.Dispatch(context.Background(), event, strings.TrimSpace(input.SpecialistName), hookPayload)
}
func (executor Executor) rollUpSpecialistUsage(input specialistAccountingInput, summary StreamResult) bool {
@@ -82,6 +225,18 @@ func (executor Executor) recordBackgroundTaskAccounting(task background.Task, su
executor.recordSpecialistStop(input, summary, string(task.Status), task.ExitCode, nil, rolledUp)
}
// taskStreamSummary builds the stream summary for a finished background task
// from its output file. Every reaper (Task.onExit, TaskOutput, TaskStop) derives
// its runId this way, so the durable stop event and stop-hook claim agree on the
// run identity regardless of which path observed the task first.
func taskStreamSummary(task background.Task) StreamResult {
summary := StreamResult{ExitCode: task.ExitCode}
if data, err := os.ReadFile(task.OutputFile); err == nil {
summary, _ = summarizeTaskData(string(data), task.ExitCode)
}
return summary
}
func appendSpecialistUsageRollup(store *sessions.Store, input specialistAccountingInput, summary StreamResult) (bool, error) {
if !summary.Usage.HasUsage() {
return false, nil
@@ -132,19 +287,7 @@ func specialistEventMatcher(eventType sessions.EventType, childSessionID, runID
if err := json.Unmarshal(event.Payload, &payload); err != nil {
continue
}
if specialistPayloadString(payload, "source") != specialistAccountingSource {
continue
}
if specialistPayloadString(payload, "childSessionId") != childSessionID && specialistPayloadString(payload, "taskId") != childSessionID {
continue
}
// An already-recorded event with NO runId is a catch-all for this child
// (e.g. the immediate stop written when a PID couldn't be registered) and
// must match a later event that does carry a runId — otherwise the same
// child gets two stop/usage events. We match when: querying without a runId,
// the existing event has none, or the runIds are equal.
existingRunID := specialistPayloadString(payload, "runId")
if runID == "" || existingRunID == "" || existingRunID == runID {
if specialistPayloadMatches(payload, childSessionID, runID) {
return true
}
}
@@ -152,6 +295,25 @@ func specialistEventMatcher(eventType sessions.EventType, childSessionID, runID
}
}
// specialistPayloadMatches applies the shared "same specialist stop" identity to
// one event payload. The accounting dedup (specialistEventMatcher) and the
// durable stop-hook claim both use it, so the audit event and the hook side
// effect can never disagree about what counts as a repeat. An already-recorded
// event with NO runId is a catch-all for this child (e.g. the immediate stop
// written when a PID couldn't be registered) and matches a later event that
// carries a runId - otherwise the same child gets two stop/usage events. A query
// without a runId likewise matches any recorded run.
func specialistPayloadMatches(payload map[string]any, childSessionID, runID string) bool {
if specialistPayloadString(payload, "source") != specialistAccountingSource {
return false
}
if specialistPayloadString(payload, "childSessionId") != childSessionID && specialistPayloadString(payload, "taskId") != childSessionID {
return false
}
existingRunID := specialistPayloadString(payload, "runId")
return runID == "" || existingRunID == "" || existingRunID == runID
}
// appendSpecialistEventOnce atomically records a "record once" specialist event:
// it appends payload only if no matching event already exists, under the session
// store's lock so concurrent stop/usage callers (even cross-process) cannot both
+836 -2
View File
@@ -3,10 +3,15 @@ package specialist
import (
"context"
"encoding/json"
"errors"
"os"
"path/filepath"
"reflect"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/Twigpine/zero/internal/background"
"github.com/Twigpine/zero/internal/sessions"
@@ -21,10 +26,21 @@ func TestExecutorRecordsForegroundLifecycleAndUsageRollup(t *testing.T) {
t.Fatalf("Create parent returned error: %v", err)
}
zero := 0
var hookEvents []string
var hookNames []string
executor := Executor{
BinaryPath: "/usr/local/bin/zero",
SessionStore: store,
NewSessionID: func() (string, error) { return "child_task", nil },
LifecycleHooks: &LifecycleHooks{
Dispatch: func(_ context.Context, event string, specialistName string, payload map[string]any) {
hookEvents = append(hookEvents, event)
hookNames = append(hookNames, specialistName)
if payload["event"] != event {
t.Errorf("payload event = %#v, want %q", payload["event"], event)
}
},
},
Load: func(LoadOptions) (LoadResult, error) {
return LoadResult{Specialists: []Manifest{{
Metadata: Metadata{Name: "worker", Description: "Does focused work"},
@@ -59,12 +75,18 @@ func TestExecutorRecordsForegroundLifecycleAndUsageRollup(t *testing.T) {
if result.Result.Status != tools.StatusOK {
t.Fatalf("Run status = %s, output=%s", result.Result.Status, result.Result.Output)
}
if got, want := strings.Join(hookEvents, ","), "specialistStart,specialistStop"; got != want {
t.Fatalf("lifecycle hooks = %q, want %q", got, want)
}
if got, want := strings.Join(hookNames, ","), "worker,worker"; got != want {
t.Fatalf("lifecycle specialist names = %q, want %q", got, want)
}
events, err := store.ReadEvents(parent.SessionID)
if err != nil {
t.Fatalf("ReadEvents returned error: %v", err)
}
if got, want := eventTypes(events), []sessions.EventType{sessions.EventSpecialistStart, sessions.EventUsage, sessions.EventSpecialistStop}; strings.Join(eventTypesString(got), ",") != strings.Join(eventTypesString(want), ",") {
if got, want := eventTypes(events), []sessions.EventType{sessions.EventSpecialistStart, sessions.EventUsage, sessions.EventSpecialistStop, specialistStopHookEvent}; strings.Join(eventTypesString(got), ",") != strings.Join(eventTypesString(want), ",") {
t.Fatalf("event types = %#v, want %#v", got, want)
}
startPayload := eventPayload(t, events[0])
@@ -88,6 +110,13 @@ func TestExecutorRecordsForegroundLifecycleAndUsageRollup(t *testing.T) {
requirePayloadString(t, stopPayload, "status", "success")
requirePayloadInt(t, stopPayload, "exitCode", 0)
requirePayloadBool(t, stopPayload, "usageRolledUp", true)
requirePayloadBool(t, stopPayload, "childStarted", true)
claimPayload := eventPayload(t, events[3])
requirePayloadString(t, claimPayload, "hook", "specialistStop")
requirePayloadString(t, claimPayload, "runId", "run_1")
requirePayloadBool(t, claimPayload, "dispatched", true)
requirePayloadBool(t, claimPayload, "childStarted", true)
}
func TestExecutorRecordsStartedChildErrorExitCode(t *testing.T) {
@@ -286,7 +315,17 @@ func TestRecordSpecialistStopDedupesUnderConcurrency(t *testing.T) {
if err != nil {
t.Fatalf("Create parent returned error: %v", err)
}
executor := Executor{SessionStore: store}
var hookStops atomic.Int32
executor := Executor{
SessionStore: store,
LifecycleHooks: &LifecycleHooks{
Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) {
if event == "specialistStop" {
hookStops.Add(1)
}
},
},
}
input := specialistAccountingInput{
ParentSessionID: parent.SessionID,
ChildSessionID: "child_task",
@@ -321,4 +360,799 @@ func TestRecordSpecialistStopDedupesUnderConcurrency(t *testing.T) {
if stops != 1 {
t.Fatalf("expected exactly 1 stop event under concurrency, got %d", stops)
}
if got := hookStops.Load(); got != 1 {
t.Fatalf("specialistStop hooks = %d, want 1 under concurrency", got)
}
}
func TestRecordSpecialistStopFiresWithoutParentSession(t *testing.T) {
var hookStops atomic.Int32
executor := Executor{
LifecycleHooks: &LifecycleHooks{
Dispatch: func(_ context.Context, event string, name string, _ map[string]any) {
if event == "specialistStop" {
hookStops.Add(1)
if name != "subagent" {
t.Errorf("specialist name = %q, want subagent", name)
}
}
},
},
}
executor.recordSpecialistStop(specialistAccountingInput{
// Swarm members omit ParentSessionID; append never succeeds.
ChildSessionID: "member_child",
SpecialistName: "subagent",
Mode: "foreground",
}, StreamResult{RunID: "run_swarm"}, "success", 0, nil, false)
if got := hookStops.Load(); got != 1 {
t.Fatalf("specialistStop hooks = %d, want 1 with empty parent", got)
}
// Second observation of the same child/run must not re-fire.
executor.recordSpecialistStop(specialistAccountingInput{
ChildSessionID: "member_child",
SpecialistName: "subagent",
Mode: "foreground",
}, StreamResult{RunID: "run_swarm"}, "success", 0, nil, false)
if got := hookStops.Load(); got != 1 {
t.Fatalf("specialistStop hooks = %d after duplicate, want 1", got)
}
}
func TestRecordSpecialistStopHookLessWinnerDoesNotSuppressLaterHooks(t *testing.T) {
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"})
if err != nil {
t.Fatalf("Create parent returned error: %v", err)
}
input := specialistAccountingInput{
ParentSessionID: parent.SessionID,
ChildSessionID: "child_task",
SpecialistName: "worker",
Mode: "background",
Background: true,
}
summary := StreamResult{RunID: "run_1"}
// Hook-less TaskOutput path wins the session append race.
Executor{SessionStore: store}.recordSpecialistStop(input, summary, "success", 0, nil, true)
var hookStops atomic.Int32
hooks := &LifecycleHooks{
Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) {
if event == "specialistStop" {
hookStops.Add(1)
}
},
}
Executor{SessionStore: store, LifecycleHooks: hooks}.recordSpecialistStop(input, summary, "success", 0, nil, true)
if got := hookStops.Load(); got != 1 {
t.Fatalf("specialistStop hooks = %d, want 1 after hook-less winner", got)
}
}
// A durable stop-hook claim must survive a process boundary: a later Zero
// process (fresh LifecycleHooks bridge over the same session store) reaping the
// same task must not re-run the hook even though its session-event append is
// rejected as a duplicate.
func TestRecordSpecialistStopHookClaimDurableAcrossBridges(t *testing.T) {
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"})
if err != nil {
t.Fatalf("Create parent returned error: %v", err)
}
input := specialistAccountingInput{
ParentSessionID: parent.SessionID,
ChildSessionID: "child_task",
SpecialistName: "worker",
Mode: "background",
Background: true,
}
summary := StreamResult{RunID: "run_1"}
var firstBridge atomic.Int32
Executor{SessionStore: store, LifecycleHooks: stopCounterHooks(&firstBridge)}.recordSpecialistStop(input, summary, "success", 0, nil, true)
if got := firstBridge.Load(); got != 1 {
t.Fatalf("first bridge stop hooks = %d, want 1", got)
}
var secondBridge atomic.Int32
Executor{SessionStore: store, LifecycleHooks: stopCounterHooks(&secondBridge)}.recordSpecialistStop(input, summary, "success", 0, nil, true)
if got := secondBridge.Load(); got != 0 {
t.Fatalf("second bridge stop hooks = %d, want 0: the durable claim must suppress a re-dispatch", got)
}
events, err := store.ReadEvents(parent.SessionID)
if err != nil {
t.Fatalf("ReadEvents returned error: %v", err)
}
if claims := eventsOfType(events, specialistStopHookEvent); len(claims) != 1 {
t.Fatalf("stop-hook claim events = %d, want 1", len(claims))
}
if stops := eventsOfType(events, sessions.EventSpecialistStop); len(stops) != 1 {
t.Fatalf("stop events = %d, want 1", len(stops))
}
}
// The durable claim must follow the accounting run-ID equivalence: an empty
// runId is a catch-all in both directions, while genuinely distinct runs stay
// distinct stop transitions.
func TestRecordSpecialistStopHookClaimRunIDEquivalence(t *testing.T) {
for _, tc := range []struct {
name string
first string
second string
want int32
}{
{name: "empty run id is a catch-all for a later run", first: "", second: "run_1", want: 1},
{name: "a recorded run covers a later empty query", first: "run_1", second: "", want: 1},
{name: "distinct runs are distinct stops", first: "run_1", second: "run_2", want: 2},
} {
t.Run(tc.name, func(t *testing.T) {
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"})
if err != nil {
t.Fatalf("Create parent returned error: %v", err)
}
var stops atomic.Int32
executor := Executor{SessionStore: store, LifecycleHooks: stopCounterHooks(&stops)}
input := specialistAccountingInput{
ParentSessionID: parent.SessionID,
ChildSessionID: "child_task",
SpecialistName: "worker",
Mode: "background",
Background: true,
}
executor.recordSpecialistStop(input, StreamResult{RunID: tc.first}, "success", 0, nil, false)
executor.recordSpecialistStop(input, StreamResult{RunID: tc.second}, "success", 0, nil, false)
if got := stops.Load(); got != tc.want {
t.Fatalf("stop hooks = %d, want %d", got, tc.want)
}
})
}
}
// The in-process fallback (swarm members have no parent session) applies the
// same run-ID equivalence, so a later observation carrying a runId cannot
// re-fire a stop that was already claimed without one.
func TestRecordSpecialistStopInProcessClaimRunIDEquivalence(t *testing.T) {
var stops atomic.Int32
executor := Executor{LifecycleHooks: stopCounterHooks(&stops)}
input := specialistAccountingInput{ChildSessionID: "member_child", SpecialistName: "subagent", Mode: "foreground"}
executor.recordSpecialistStop(input, StreamResult{}, "success", 0, nil, false)
executor.recordSpecialistStop(input, StreamResult{RunID: "run_later"}, "success", 0, nil, false)
if got := stops.Load(); got != 1 {
t.Fatalf("stop hooks = %d, want 1", got)
}
}
// When the durable ledger is briefly unavailable, the in-process fallback and a
// durable claim on the same bridge must make one dispatch decision between
// them, whichever order they happen in.
func TestRecordSpecialistStopDurableAndFallbackClaimsShareOneDispatch(t *testing.T) {
for _, tc := range []struct {
name string
fallbackFirst bool
}{
{name: "fallback then durable", fallbackFirst: true},
{name: "durable then fallback", fallbackFirst: false},
} {
t.Run(tc.name, func(t *testing.T) {
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"})
if err != nil {
t.Fatalf("Create parent returned error: %v", err)
}
eventsFile := filepath.Join(store.RootDir, parent.SessionID, sessions.EventsFile)
asideFile := eventsFile + ".aside"
// A directory where the event log lives makes every read and append
// fail on all platforms, the way a transient store error would.
breakLedger := func() {
t.Helper()
if err := os.Rename(eventsFile, asideFile); err != nil && !errors.Is(err, os.ErrNotExist) {
t.Fatalf("move event log aside: %v", err)
}
if err := os.Mkdir(eventsFile, 0o700); err != nil {
t.Fatalf("block event log: %v", err)
}
}
restoreLedger := func() {
t.Helper()
if err := os.Remove(eventsFile); err != nil {
t.Fatalf("unblock event log: %v", err)
}
if err := os.Rename(asideFile, eventsFile); err != nil && !errors.Is(err, os.ErrNotExist) {
t.Fatalf("restore event log: %v", err)
}
}
var stops atomic.Int32
executor := Executor{SessionStore: store, LifecycleHooks: stopCounterHooks(&stops)}
input := specialistAccountingInput{
ParentSessionID: parent.SessionID,
ChildSessionID: "child_task",
SpecialistName: "worker",
Mode: "background",
Background: true,
}
summary := StreamResult{RunID: "run_1"}
if tc.fallbackFirst {
breakLedger()
executor.recordSpecialistStop(input, summary, "success", 0, nil, false)
restoreLedger()
executor.recordSpecialistStop(input, summary, "success", 0, nil, false)
} else {
executor.recordSpecialistStop(input, summary, "success", 0, nil, false)
breakLedger()
executor.recordSpecialistStop(input, summary, "success", 0, nil, false)
restoreLedger()
}
if got := stops.Load(); got != 1 {
t.Fatalf("stop hooks = %d, want 1 across durable and fallback claims", got)
}
})
}
}
// A child that never started must not publish public lifecycle hooks, while the
// internal attempt/stop accounting is retained and durably marks the stop as
// ineligible so a later reaper cannot fire the hook either.
func TestExecutorSkipsHooksWhenChildNeverStarts(t *testing.T) {
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"})
if err != nil {
t.Fatalf("Create parent returned error: %v", err)
}
var hooks []string
executor := Executor{
BinaryPath: "/usr/local/bin/zero",
SessionStore: store,
NewSessionID: func() (string, error) { return "child_task", nil },
LifecycleHooks: &LifecycleHooks{
Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) {
hooks = append(hooks, event)
},
},
Load: func(LoadOptions) (LoadResult, error) {
return LoadResult{Specialists: []Manifest{{
Metadata: Metadata{Name: "worker", Description: "Does focused work"},
SystemPrompt: "Work carefully.",
ResolvedTools: []string{"read_file"},
}}}, nil
},
RunChild: func(context.Context, string, []string, func(streamjson.Event)) (ChildRunResult, error) {
return ChildRunResult{}, errors.New("exec: no such file or directory")
},
}
if _, err := executor.Run(context.Background(), TaskParameters{Name: "worker", Prompt: "inspect auth"}, TaskRunOptions{ParentSessionID: parent.SessionID}); err == nil {
t.Fatal("Run returned nil error for a failed child start")
}
if len(hooks) != 0 {
t.Fatalf("public lifecycle hooks fired for a child that never started: %#v", hooks)
}
events, err := store.ReadEvents(parent.SessionID)
if err != nil {
t.Fatalf("ReadEvents returned error: %v", err)
}
if starts := eventsOfType(events, sessions.EventSpecialistStart); len(starts) != 1 {
t.Fatalf("attempt start events = %d, want 1", len(starts))
}
stopEvents := eventsOfType(events, sessions.EventSpecialistStop)
if len(stopEvents) != 1 {
t.Fatalf("stop events = %#v", events)
}
requirePayloadBool(t, eventPayload(t, stopEvents[0]), "childStarted", false)
claims := eventsOfType(events, specialistStopHookEvent)
if len(claims) != 1 {
t.Fatalf("stop-hook claim events = %d, want 1", len(claims))
}
requirePayloadBool(t, eventPayload(t, claims[0]), "dispatched", false)
}
// A never-started background child leaves an errored task behind; a later
// TaskOutput reaper must not fire the stop hook from the recorded eligibility.
func TestBackgroundLaunchFailureDoesNotFireStopHookOnReap(t *testing.T) {
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"})
if err != nil {
t.Fatalf("Create parent returned error: %v", err)
}
manager, err := background.NewManager(t.TempDir())
if err != nil {
t.Fatalf("NewManager returned error: %v", err)
}
var hooks []string
executor := Executor{
BinaryPath: "/usr/local/bin/zero",
SessionStore: store,
BackgroundManager: manager,
NewSessionID: func() (string, error) { return "child_task", nil },
LifecycleHooks: &LifecycleHooks{
Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) {
hooks = append(hooks, event)
},
},
Load: func(LoadOptions) (LoadResult, error) {
return LoadResult{Specialists: []Manifest{{
Metadata: Metadata{Name: "worker", Description: "Does focused work"},
SystemPrompt: "Work carefully.",
ResolvedTools: []string{"read_file"},
}}}, nil
},
LaunchBackground: func(string, []string, string, func(int)) (int, error) {
return 0, errors.New("launch failed")
},
}
_, err = executor.Run(context.Background(), TaskParameters{Name: "worker", Prompt: "inspect auth", RunInBackground: true}, TaskRunOptions{ParentSessionID: parent.SessionID})
if err == nil {
t.Fatal("Run returned nil error for a failed background launch")
}
if len(hooks) != 0 {
t.Fatalf("public lifecycle hooks fired for a child that never started: %#v", hooks)
}
// A separate process reaping the errored task must stay silent too.
var reapHooks atomic.Int32
tool := NewOutputTool(manager)
tool.SessionStore = store
tool.LifecycleHooks = stopCounterHooks(&reapHooks)
if result := tool.Run(context.Background(), map[string]any{"task_id": "child_task"}); result.Status != tools.StatusOK {
t.Fatalf("TaskOutput status = %s output=%s", result.Status, result.Output)
}
if got := reapHooks.Load(); got != 0 {
t.Fatalf("stop hooks from reaping a never-started task = %d, want 0", got)
}
events, err := store.ReadEvents(parent.SessionID)
if err != nil {
t.Fatalf("ReadEvents returned error: %v", err)
}
if stopEvents := eventsOfType(events, sessions.EventSpecialistStop); len(stopEvents) != 1 {
t.Fatalf("stop events = %d, want 1", len(stopEvents))
}
if claims := eventsOfType(events, specialistStopHookEvent); len(claims) != 1 {
t.Fatalf("stop-hook claim events = %d, want 1", len(claims))
}
}
// A background child that exits before the launcher returns must not run its
// terminal accounting - and therefore must not fire the specialistStop hook -
// until the start decision has completed. Otherwise a fast-exiting child
// publishes specialistStop before specialistStart, and the recorded stop can
// miss the pid that SetPID has not stored yet. The synchronous launcher below
// also pins the deadlock requirement: the exit callback may not wait inline,
// because onExit runs before the launcher (and the start decision) can return.
func TestRunBackgroundDefersExitUntilStartDecisionCompletes(t *testing.T) {
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"})
if err != nil {
t.Fatalf("Create parent returned error: %v", err)
}
manager, err := background.NewManager(t.TempDir())
if err != nil {
t.Fatalf("NewManager returned error: %v", err)
}
var hookMu sync.Mutex
var hookEvents []string
stopDispatched := make(chan struct{}, 1)
executor := Executor{
BinaryPath: "/usr/local/bin/zero",
SessionStore: store,
BackgroundManager: manager,
NewSessionID: func() (string, error) { return "child_task", nil },
LifecycleHooks: &LifecycleHooks{
Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) {
hookMu.Lock()
hookEvents = append(hookEvents, event)
hookMu.Unlock()
if event == "specialistStop" {
select {
case stopDispatched <- struct{}{}:
default:
}
}
},
},
Load: func(LoadOptions) (LoadResult, error) {
return LoadResult{Specialists: []Manifest{{
Metadata: Metadata{Name: "worker", Description: "Does focused work"},
SystemPrompt: "Work carefully.",
ResolvedTools: []string{"read_file"},
}}}, nil
},
// The child exits before the launcher returns, so the exit callback runs
// synchronously - exactly the ordering hazard this test guards.
LaunchBackground: func(_ string, _ []string, _ string, onExit func(exitCode int)) (int, error) {
onExit(0)
return 1234, nil
},
}
runErr := make(chan error, 1)
go func() {
_, err := executor.Run(context.Background(), TaskParameters{
Name: "worker",
Prompt: "inspect auth",
RunInBackground: true,
}, TaskRunOptions{ParentSessionID: parent.SessionID})
runErr <- err
}()
select {
case err := <-runErr:
if err != nil {
t.Fatalf("Run returned error: %v", err)
}
case <-time.After(10 * time.Second):
t.Fatal("synchronous launcher deadlocked the background start decision")
}
// Run publishes specialistStart before it returns; the deferred exit handling
// publishes specialistStop afterwards.
select {
case <-stopDispatched:
case <-time.After(10 * time.Second):
t.Fatal("deferred exit accounting never dispatched specialistStop")
}
hookMu.Lock()
gotHooks := strings.Join(hookEvents, ",")
hookMu.Unlock()
if gotHooks != "specialistStart,specialistStop" {
t.Fatalf("lifecycle hooks = %q, want start before stop", gotHooks)
}
events, err := store.ReadEvents(parent.SessionID)
if err != nil {
t.Fatalf("ReadEvents returned error: %v", err)
}
stopEvents := eventsOfType(events, sessions.EventSpecialistStop)
if len(stopEvents) != 1 {
t.Fatalf("stop events = %#v", events)
}
// The deferred stop must observe the task after SetPID; a stop recorded
// before the start decision completes would omit the pid.
requirePayloadInt(t, eventPayload(t, stopEvents[0]), "pid", 1234)
}
// When SetPID fails after the child started, Run kills the unrecorded child. Only
// a successful kill is a terminal stop. A failed kill may leave the child
// running, so Run must keep the task running (no reaper may claim it) and must
// not record the stop or take the one-shot specialistStop claim; the launcher's
// exit callback does both when the child really ends.
func TestRunBackgroundSetPIDFailureRecordsStopOnlyAfterTermination(t *testing.T) {
for _, tc := range []struct {
name string
killErr error
}{
{name: "kill succeeds"},
{name: "kill fails", killErr: errors.New("operation not permitted")},
} {
t.Run(tc.name, func(t *testing.T) {
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"})
if err != nil {
t.Fatalf("Create parent returned error: %v", err)
}
manager, err := background.NewManager(t.TempDir())
if err != nil {
t.Fatalf("NewManager returned error: %v", err)
}
var killed []int
originalTerminate := terminateOrphanedProcess
terminateOrphanedProcess = func(pid int) error {
killed = append(killed, pid)
return tc.killErr
}
t.Cleanup(func() { terminateOrphanedProcess = originalTerminate })
var hookMu sync.Mutex
var hookEvents []string
hooksSoFar := func() string {
hookMu.Lock()
defer hookMu.Unlock()
return strings.Join(hookEvents, ",")
}
stopDispatched := make(chan struct{}, 1)
metadataFile := filepath.Join(manager.RootDir(), "child_task.json")
// The child reads its prompt from this file, so it must outlive any
// child that may still be running.
promptFile := filepath.Join(t.TempDir(), "prompt.md")
promptExists := func() bool {
_, err := os.Stat(promptFile)
return err == nil
}
var onExitChild func(exitCode int)
executor := Executor{
BinaryPath: "/usr/local/bin/zero",
SessionStore: store,
BackgroundManager: manager,
NewSessionID: func() (string, error) { return "child_task", nil },
PromptFileMaxSize: 1,
WritePromptFile: func(prompt string) (string, error) {
return promptFile, os.WriteFile(promptFile, []byte(prompt), 0o600)
},
LifecycleHooks: &LifecycleHooks{
Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) {
hookMu.Lock()
hookEvents = append(hookEvents, event)
hookMu.Unlock()
if event == "specialistStop" {
select {
case stopDispatched <- struct{}{}:
default:
}
}
},
},
Load: func(LoadOptions) (LoadResult, error) {
return LoadResult{Specialists: []Manifest{{
Metadata: Metadata{Name: "worker", Description: "Does focused work"},
SystemPrompt: "Work carefully.",
ResolvedTools: []string{"read_file"},
}}}, nil
},
LaunchBackground: func(_ string, _ []string, _ string, onExit func(exitCode int)) (int, error) {
onExitChild = onExit
// The child has started. Put a non-empty directory where the
// task metadata lives so SetPID's atomic replace fails.
if err := os.Remove(metadataFile); err != nil {
return 0, err
}
if err := os.MkdirAll(filepath.Join(metadataFile, "blocker"), 0o700); err != nil {
return 0, err
}
return 1234, nil
},
}
_, runErr := executor.Run(context.Background(), TaskParameters{
Name: "worker",
Prompt: "inspect auth",
RunInBackground: true,
}, TaskRunOptions{ParentSessionID: parent.SessionID})
if runErr == nil {
t.Fatal("Run returned nil error after SetPID failed")
}
if !reflect.DeepEqual(killed, []int{1234}) {
t.Fatalf("terminated pids = %v, want [1234]", killed)
}
readStops := func() ([]sessions.Event, []sessions.Event) {
t.Helper()
events, err := store.ReadEvents(parent.SessionID)
if err != nil {
t.Fatalf("ReadEvents returned error: %v", err)
}
return eventsOfType(events, sessions.EventSpecialistStop), eventsOfType(events, specialistStopHookEvent)
}
if tc.killErr == nil {
// The child is gone, so the stop is terminal and recorded now.
stops, claims := readStops()
if len(stops) != 1 || len(claims) != 1 {
t.Fatalf("after successful kill: stop events = %d, claims = %d, want 1 and 1", len(stops), len(claims))
}
if got := hooksSoFar(); got != "specialistStart,specialistStop" {
t.Fatalf("lifecycle hooks = %q, want specialistStart,specialistStop", got)
}
if promptExists() {
t.Fatal("prompt file survived after the orphan was killed")
}
return
}
if !strings.Contains(runErr.Error(), "terminate orphaned pid 1234") {
t.Fatalf("Run error = %v, want the failed kill surfaced", runErr)
}
if !promptExists() {
t.Fatal("prompt file removed while the child may still be running")
}
if got := hooksSoFar(); got != "specialistStart" {
t.Fatalf("lifecycle hooks after failed kill = %q, want only specialistStart", got)
}
if stops, claims := readStops(); len(stops) != 0 || len(claims) != 0 {
t.Fatalf("after failed kill: stop events = %d, claims = %d, want none while the child may run", len(stops), len(claims))
}
task, ok := manager.Get("child_task")
if !ok || task.Status != background.StatusRunning {
t.Fatalf("task after failed kill = %#v (found %v), want still running", task, ok)
}
// A TaskOutput or TaskStop reap in the meantime must not claim it.
executor.recordBackgroundTaskAccounting(task, taskStreamSummary(task))
if stops, claims := readStops(); len(stops) != 0 || len(claims) != 0 {
t.Fatalf("reap of running task recorded stop events = %d, claims = %d", len(stops), len(claims))
}
// The surviving child finally exits: its exit callback records the stop
// and fires the hook exactly once.
if err := os.RemoveAll(metadataFile); err != nil {
t.Fatalf("RemoveAll metadata blocker returned error: %v", err)
}
onExitChild(1)
select {
case <-stopDispatched:
case <-time.After(10 * time.Second):
t.Fatal("exit callback never dispatched specialistStop")
}
if got := hooksSoFar(); got != "specialistStart,specialistStop" {
t.Fatalf("lifecycle hooks = %q, want specialistStart,specialistStop", got)
}
stops, claims := readStops()
if len(stops) != 1 || len(claims) != 1 {
t.Fatalf("after child exit: stop events = %d, claims = %d, want 1 and 1", len(stops), len(claims))
}
requirePayloadInt(t, eventPayload(t, stops[0]), "exitCode", 1)
// The exit callback owns the prompt cleanup once the child is gone.
deadline := time.Now().Add(10 * time.Second)
for promptExists() {
if time.Now().After(deadline) {
t.Fatal("exit callback never removed the prompt file")
}
time.Sleep(10 * time.Millisecond)
}
})
}
}
// A recovered task stopped explicitly (launcher process gone) must record the
// terminal transition and fire the stop hook exactly once through the shared
// durable claim, even when a later TaskOutput poll reaps it again.
func TestStopToolRecordsRecoveredTaskStopOnce(t *testing.T) {
killed := []int{}
manager, err := background.NewManagerWithOptions(background.ManagerOptions{
RootDir: t.TempDir(),
KillProcess: func(pid int) error {
killed = append(killed, pid)
return nil
},
})
if err != nil {
t.Fatalf("NewManagerWithOptions returned error: %v", err)
}
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"})
if err != nil {
t.Fatalf("Create parent returned error: %v", err)
}
outputFile, err := manager.Register(background.RegisterInput{
TaskID: "child_task",
Type: specialistTaskType,
SpecialistName: "worker",
Description: "Auth check",
ParentID: parent.SessionID,
PID: 4321,
})
if err != nil {
t.Fatalf("Register returned error: %v", err)
}
if err := os.WriteFile(outputFile, []byte(strings.Join([]string{
`{"schemaVersion":2,"type":"run_start","runId":"run_1","sessionId":"child_task"}`,
`{"schemaVersion":2,"type":"final","runId":"run_1","text":"done"}`,
`{"schemaVersion":2,"type":"run_end","runId":"run_1","status":"success","exitCode":0}`,
"",
}, "\n")), 0o600); err != nil {
t.Fatalf("WriteFile returned error: %v", err)
}
var counter atomic.Int32
tool := &StopTool{manager: manager, SessionStore: store, LifecycleHooks: stopCounterHooks(&counter)}
result := tool.Run(context.Background(), map[string]any{"task_id": "child_task"})
if result.Status != tools.StatusOK || result.Meta["status"] != string(background.StatusKilled) {
t.Fatalf("TaskStop result = %#v", result)
}
if !reflect.DeepEqual(killed, []int{4321}) {
t.Fatalf("killed pids = %#v", killed)
}
if got := counter.Load(); got != 1 {
t.Fatalf("stop hooks = %d, want 1 from TaskStop", got)
}
// A later TaskOutput reaper in another process (fresh bridge) must not
// re-dispatch; the durable claim keys on the runId read from the output.
var reapCounter atomic.Int32
outputTool := NewOutputTool(manager)
outputTool.SessionStore = store
outputTool.LifecycleHooks = stopCounterHooks(&reapCounter)
if result := outputTool.Run(context.Background(), map[string]any{"task_id": "child_task"}); result.Status != tools.StatusOK {
t.Fatalf("TaskOutput status = %s output=%s", result.Status, result.Output)
}
if got := reapCounter.Load(); got != 0 {
t.Fatalf("stop hooks after later reap = %d, want 0", got)
}
events, err := store.ReadEvents(parent.SessionID)
if err != nil {
t.Fatalf("ReadEvents returned error: %v", err)
}
if stopEvents := eventsOfType(events, sessions.EventSpecialistStop); len(stopEvents) != 1 {
t.Fatalf("stop events = %d, want 1", len(stopEvents))
}
claims := eventsOfType(events, specialistStopHookEvent)
if len(claims) != 1 {
t.Fatalf("stop-hook claim events = %d, want 1", len(claims))
}
requirePayloadString(t, eventPayload(t, claims[0]), "runId", "run_1")
}
func stopCounterHooks(counter *atomic.Int32) *LifecycleHooks {
return &LifecycleHooks{
Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) {
if event == "specialistStop" {
counter.Add(1)
}
},
}
}
func TestOutputToolSharesLifecycleHooksWithStopAccounting(t *testing.T) {
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
parent, err := store.Create(sessions.CreateInput{SessionID: "parent_session"})
if err != nil {
t.Fatalf("Create parent returned error: %v", err)
}
manager, err := background.NewManager(t.TempDir())
if err != nil {
t.Fatalf("NewManager returned error: %v", err)
}
outputFile, err := manager.Register(background.RegisterInput{
TaskID: "child_task",
Type: "specialist",
SpecialistName: "worker",
Description: "Auth check",
ParentID: parent.SessionID,
})
if err != nil {
t.Fatalf("Register returned error: %v", err)
}
if err := os.WriteFile(outputFile, []byte(strings.Join([]string{
`{"schemaVersion":2,"type":"run_start","runId":"run_1","sessionId":"child_task"}`,
`{"schemaVersion":2,"type":"final","runId":"run_1","text":"done"}`,
`{"schemaVersion":2,"type":"run_end","runId":"run_1","status":"success","exitCode":0}`,
"",
}, "\n")), 0o600); err != nil {
t.Fatalf("WriteFile returned error: %v", err)
}
if err := manager.UpdateStatus("child_task", background.StatusCompleted, 0); err != nil {
t.Fatalf("UpdateStatus returned error: %v", err)
}
var hookStops atomic.Int32
hooks := &LifecycleHooks{
Dispatch: func(_ context.Context, event string, _ string, _ map[string]any) {
if event == "specialistStop" {
hookStops.Add(1)
}
},
}
tool := NewOutputTool(manager)
tool.SessionStore = store
tool.LifecycleHooks = hooks
if result := tool.Run(context.Background(), map[string]any{"task_id": "child_task"}); result.Status != tools.StatusOK {
t.Fatalf("TaskOutput status = %s output=%s", result.Status, result.Output)
}
if got := hookStops.Load(); got != 1 {
t.Fatalf("specialistStop hooks = %d, want 1 from TaskOutput", got)
}
// Shared bridge: a later Task.onExit-style caller must not re-fire.
Executor{SessionStore: store, LifecycleHooks: hooks}.recordBackgroundTaskAccounting(
mustGetBackgroundTask(t, manager, "child_task"),
StreamResult{RunID: "run_1"},
)
if got := hookStops.Load(); got != 1 {
t.Fatalf("specialistStop hooks = %d after second path, want 1", got)
}
}
func mustGetBackgroundTask(t *testing.T, manager *background.Manager, taskID string) background.Task {
t.Helper()
task, ok := manager.Get(taskID)
if !ok {
t.Fatalf("background task %q not found", taskID)
}
return task
}
+133 -22
View File
@@ -15,6 +15,7 @@ import (
"path/filepath"
"strconv"
"strings"
"sync"
"github.com/Twigpine/zero/internal/background"
"github.com/Twigpine/zero/internal/sessions"
@@ -24,12 +25,17 @@ import (
const (
sessionTagSpecialist = "specialist"
specialistTaskType = "specialist"
promptFileThresholdBytes = 4 * 1024
maxSpecialistDepth = 8 // hard cap to prevent infinite recursion/resource exhaustion via Task
)
const SessionTagSpecialist = sessionTagSpecialist
// terminateOrphanedProcess stops a launched background child whose PID could not
// be recorded. Tests replace it to drive the failed-kill path.
var terminateOrphanedProcess = background.TerminateProcess
type NewSessionIDFunc func() (string, error)
type WritePromptFileFunc func(prompt string) (string, error)
type LoadFunc func(LoadOptions) (LoadResult, error)
@@ -50,6 +56,64 @@ type Executor struct {
BackgroundManager *background.Manager
BackgroundManagerFunc BackgroundManagerFunc
BackgroundRuntime *Runtime
// LifecycleHooks optionally dispatches specialistStart/specialistStop hooks.
// The shared pointer lets the CLI attach the hooks dispatcher after specialist
// tools are registered (hooks are constructed later in startup).
LifecycleHooks *LifecycleHooks
}
// LifecycleHooks bridges specialist accounting into the hooks dispatcher without
// importing the hooks package into every specialist call site.
type LifecycleHooks struct {
// Dispatch runs configured hooks for event ("specialistStart"/"specialistStop").
// specialistName is used as the matcher subject. nil Dispatch is a no-op.
Dispatch func(ctx context.Context, event string, specialistName string, payload map[string]any)
// stopMu guards stopClaims, the in-process fallback used when no durable
// claim can be written (swarm members have no parent session) or the session
// store is unavailable.
stopMu sync.Mutex
stopClaims map[string]map[string]struct{}
}
// claimInProcessStopDispatch reports whether this bridge should run
// specialistStop for the given child/run. It is the process-local fallback for
// callers without a durable ledger; the run-ID equivalence mirrors
// specialistPayloadMatches: an empty runId is a catch-all for both the query and
// any recorded claim, so a later observation with a runId cannot re-fire a stop
// that was already claimed without one. An empty child session never dedups
// (there is nothing to key on), matching the accounting matcher.
func (h *LifecycleHooks) claimInProcessStopDispatch(childSessionID, runID string) bool {
if h == nil || h.Dispatch == nil {
return false
}
childSessionID = strings.TrimSpace(childSessionID)
if childSessionID == "" {
return true
}
runID = strings.TrimSpace(runID)
h.stopMu.Lock()
defer h.stopMu.Unlock()
if h.stopClaims == nil {
h.stopClaims = map[string]map[string]struct{}{}
}
claims := h.stopClaims[childSessionID]
if claims == nil {
claims = map[string]struct{}{}
h.stopClaims[childSessionID] = claims
}
claimed := false
if runID == "" {
claimed = len(claims) > 0
} else {
_, emptyClaim := claims[""]
_, exactClaim := claims[runID]
claimed = emptyClaim || exactClaim
}
if !claimed {
claims[runID] = struct{}{}
}
return !claimed
}
type BuildArgsInput struct {
@@ -478,7 +542,7 @@ func (executor Executor) runBackground(ctx context.Context, built BuildArgsResul
}
outputFile, err := manager.Register(background.RegisterInput{
TaskID: built.SessionID,
Type: "specialist",
Type: specialistTaskType,
SpecialistName: manifest.Metadata.Name,
Description: params.Description,
ParentID: options.ParentSessionID,
@@ -516,38 +580,57 @@ func (executor Executor) runBackground(ctx context.Context, built BuildArgsResul
Background: true,
}
executor.recordSpecialistStart(accounting)
// The launcher's exit callback runs the terminal accounting, which can claim
// and dispatch the specialistStop hook. A child that exits immediately would
// otherwise publish specialistStop before the start decision below completes
// (and before SetPID records the pid). The callback waits for startSettled
// from its own goroutine: waiting inline would deadlock a synchronous
// injected launcher whose onExit runs before it returns.
startSettled := make(chan struct{})
defer close(startSettled)
pid, err := executor.launchBackground(binaryPath, built.Args, outputFile, func(exitCode int) {
status := background.StatusCompleted
if exitCode != 0 {
status = background.StatusError
}
_ = manager.MarkExited(built.SessionID, status, exitCode)
if task, ok := manager.Get(built.SessionID); ok {
summary := StreamResult{ExitCode: task.ExitCode}
if data, err := os.ReadFile(task.OutputFile); err == nil {
summary, _ = summarizeTaskData(string(data), task.ExitCode)
go func() {
<-startSettled
status := background.StatusCompleted
if exitCode != 0 {
status = background.StatusError
}
executor.recordBackgroundTaskAccounting(task, summary)
}
executor.cleanupBackgroundPromptFile(built.SessionID, built.PromptFile)
_ = manager.MarkExited(built.SessionID, status, exitCode)
if task, ok := manager.Get(built.SessionID); ok {
executor.recordBackgroundTaskAccounting(task, taskStreamSummary(task))
}
executor.cleanupBackgroundPromptFile(built.SessionID, built.PromptFile)
}()
})
if err != nil {
// The child never started: keep the attempt accounting, but publish no
// public lifecycle hooks for a sub-agent that did not exist.
accounting.LaunchFailed = true
_ = manager.UpdateStatus(built.SessionID, background.StatusError, -1)
executor.cleanupBackgroundPromptFile(built.SessionID, built.PromptFile)
executor.recordSpecialistStop(accounting, StreamResult{ExitCode: -1}, "error", -1, err, false)
return ExecResult{}, err
}
// Startup is confirmed (launchBackground only returns without error after the
// child process started), so the specialistStart hook can fire now.
executor.dispatchSpecialistStartHook(accounting)
if pid > 0 {
if err := manager.SetPID(built.SessionID, pid); err != nil {
// The child is running but its PID was never recorded, so nothing can
// track or stop it later. Kill it (and its process group) now, mark the
// task errored, and record the stop so it cannot become an untracked
// orphan. The dedup in recordSpecialistStop makes the later onExit
// accounting a no-op.
if killErr := background.TerminateProcess(pid); killErr != nil {
// orphan. Once the kill succeeds the stop is a real terminal transition:
// both the accounting dedup and the shared stop-hook claim make the
// later onExit accounting a no-op.
if killErr := terminateOrphanedProcess(pid); killErr != nil {
// Don't mask a failed kill: a surviving orphan is worse than the
// SetPID error alone, so surface both to the caller.
// SetPID error alone, so surface both to the caller. The child may
// still be running, so this is not a terminal stop: leave the task
// running and let the launcher's exit callback mark it exited,
// record the stop (and claim specialistStop), and remove the prompt
// file it may still need to read when it really ends.
err = errors.Join(err, fmt.Errorf("terminate orphaned pid %d: %w", pid, killErr))
return ExecResult{}, err
}
_ = manager.UpdateStatus(built.SessionID, background.StatusError, -1)
executor.cleanupBackgroundPromptFile(built.SessionID, built.PromptFile)
@@ -627,7 +710,23 @@ func (executor Executor) runBuiltArgs(ctx context.Context, built BuildArgsResult
Background: false,
}
executor.recordSpecialistStart(accounting)
run, err := executor.runChild(ctx, binaryPath, built.Args, progress)
startDispatched := false
onStarted := func() {
if startDispatched {
return
}
startDispatched = true
executor.dispatchSpecialistStartHook(accounting)
}
run, err := executor.runChild(ctx, binaryPath, built.Args, progress, onStarted)
started := run.Started || err == nil
if started {
// A seam that returns a completed run without the explicit Started flag
// still proves the child existed; make sure the hook fired exactly once.
onStarted()
} else {
accounting.LaunchFailed = true
}
if err != nil {
exitCode := run.exitCodeOr(-1)
summary := SummarizeStream(run.Events, exitCode)
@@ -697,11 +796,20 @@ func (executor Executor) binaryPath() (string, error) {
return path, nil
}
func (executor Executor) runChild(ctx context.Context, binaryPath string, args []string, progress func(streamjson.Event)) (ChildRunResult, error) {
// runChild runs one child process. onStarted, when non-nil, is invoked exactly
// once as soon as the child process has actually started (after os/exec Start
// succeeds), so public lifecycle hooks cannot fire for a launch that failed.
// The injected RunChild seam cannot stream that signal, so a seam result marked
// Started reports it when the seam returns.
func (executor Executor) runChild(ctx context.Context, binaryPath string, args []string, progress func(streamjson.Event), onStarted func()) (ChildRunResult, error) {
if executor.RunChild != nil {
return executor.RunChild(ctx, binaryPath, append([]string(nil), args...), progress)
result, err := executor.RunChild(ctx, binaryPath, append([]string(nil), args...), progress)
if result.Started && onStarted != nil {
onStarted()
}
return result, err
}
return runChildProcess(ctx, binaryPath, args, progress)
return runChildProcess(ctx, binaryPath, args, progress, onStarted)
}
func (executor Executor) launchBackground(binaryPath string, args []string, outputFile string, onExit func(exitCode int)) (int, error) {
@@ -830,7 +938,7 @@ func cleanupPromptFile(promptFile string) {
_ = os.Remove(promptFile)
}
func runChildProcess(ctx context.Context, binaryPath string, args []string, progress func(streamjson.Event)) (ChildRunResult, error) {
func runChildProcess(ctx context.Context, binaryPath string, args []string, progress func(streamjson.Event), onStarted func()) (ChildRunResult, error) {
var stderr bytes.Buffer
command := osexec.CommandContext(ctx, binaryPath, args...)
// Put the child in its own process group and group-kill on cancel/timeout, so a
@@ -844,6 +952,9 @@ func runChildProcess(ctx context.Context, binaryPath string, args []string, prog
if err := command.Start(); err != nil {
return ChildRunResult{Stderr: stderr.String()}, fmt.Errorf("start specialist child: %w", err)
}
if onStarted != nil {
onStarted()
}
events := []streamjson.Event{}
reader := bufio.NewReader(stdout)
for {
+12 -6
View File
@@ -21,9 +21,12 @@ const (
)
type OutputTool struct {
manager *background.Manager
managerFunc BackgroundManagerFunc
SessionStore *sessions.Store
manager *background.Manager
managerFunc BackgroundManagerFunc
SessionStore *sessions.Store
// LifecycleHooks must be the same bridge as Task so onExit and TaskOutput
// poll share stop-hook dedup. nil is safe (accounting still runs).
LifecycleHooks *LifecycleHooks
PollInterval time.Duration
DefaultTimeout time.Duration
MaxTimeout time.Duration
@@ -35,8 +38,8 @@ type outputParameters struct {
TimeoutMS int
}
func newOutputToolWithManagerFunc(managerFunc BackgroundManagerFunc, sessionStore *sessions.Store) *OutputTool {
return &OutputTool{managerFunc: managerFunc, SessionStore: sessionStore}
func newOutputToolWithManagerFunc(managerFunc BackgroundManagerFunc, sessionStore *sessions.Store, lifecycleHooks *LifecycleHooks) *OutputTool {
return &OutputTool{managerFunc: managerFunc, SessionStore: sessionStore, LifecycleHooks: lifecycleHooks}
}
func (tool *OutputTool) Name() string {
@@ -143,7 +146,10 @@ func (tool *OutputTool) readOutput(task background.Task) tools.Result {
dataString := string(data)
summary, rawLines := summarizeTaskData(dataString, task.ExitCode)
if task.Status != background.StatusRunning {
Executor{SessionStore: tool.SessionStore}.recordBackgroundTaskAccounting(task, summary)
Executor{
SessionStore: tool.SessionStore,
LifecycleHooks: tool.LifecycleHooks,
}.recordBackgroundTaskAccounting(task, summary)
}
return tools.Result{
Status: tools.StatusOK,
+2 -2
View File
@@ -14,8 +14,8 @@ func RegisterTools(registry *tools.Registry, executor Executor) (*Runtime, error
toolExecutor.BackgroundRuntime = runtime
toolExecutor.BackgroundManagerFunc = runtime.Manager
registry.Register(NewTaskTool(toolExecutor))
registry.Register(newOutputToolWithManagerFunc(runtime.Manager, executor.SessionStore))
registry.Register(newStopToolWithManagerFunc(runtime.Manager))
registry.Register(newOutputToolWithManagerFunc(runtime.Manager, executor.SessionStore, executor.LifecycleHooks))
registry.Register(newStopToolWithManagerFunc(runtime.Manager, executor.SessionStore, executor.LifecycleHooks))
registry.Register(NewGenerateTool(NewStorage(executor.Paths)))
return runtime, nil
}
+20 -2
View File
@@ -6,20 +6,26 @@ import (
"strings"
"github.com/Twigpine/zero/internal/background"
"github.com/Twigpine/zero/internal/sessions"
"github.com/Twigpine/zero/internal/tools"
)
type StopTool struct {
manager *background.Manager
managerFunc BackgroundManagerFunc
// SessionStore and LifecycleHooks participate in the shared, deduped stop
// transition so an explicitly stopped recovered task produces exactly one
// specialistStop hook, the same as a launcher-owned exit.
SessionStore *sessions.Store
LifecycleHooks *LifecycleHooks
}
type stopParameters struct {
TaskID string
}
func newStopToolWithManagerFunc(managerFunc BackgroundManagerFunc) *StopTool {
return &StopTool{managerFunc: managerFunc}
func newStopToolWithManagerFunc(managerFunc BackgroundManagerFunc, sessionStore *sessions.Store, lifecycleHooks *LifecycleHooks) *StopTool {
return &StopTool{managerFunc: managerFunc, SessionStore: sessionStore, LifecycleHooks: lifecycleHooks}
}
func (tool *StopTool) Name() string {
@@ -65,6 +71,18 @@ func (tool *StopTool) Run(ctx context.Context, args map[string]any) tools.Result
if err := manager.Kill(params.TaskID); err != nil {
return taskError(err)
}
// A recovered task (launcher process gone) may only ever be reaped here: the
// in-memory onExit callback does not exist in this process, and a later
// TaskOutput poll is optional. Record the terminal transition now so the
// specialistStop hook fires exactly once through the shared durable claim.
// recordBackgroundTaskAccounting is a no-op for non-specialist tasks and for
// tasks that are still running (Kill raced a status change).
if task, ok := manager.Get(params.TaskID); ok && task.Type == specialistTaskType {
Executor{
SessionStore: tool.SessionStore,
LifecycleHooks: tool.LifecycleHooks,
}.recordBackgroundTaskAccounting(task, taskStreamSummary(task))
}
return tools.Result{
Status: tools.StatusOK,
Output: fmt.Sprintf("task_id: %s\nstatus: killed", params.TaskID),
+7 -4
View File
@@ -148,10 +148,13 @@ func TestConcurrencyCapAndQueueDrains(t *testing.T) {
}
// Release everyone; the queue should drain one-per-slot until all are done.
close(gate)
testutil.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())
}
// A member's task reads Done before its slot is released: the supervisor
// completes the task, then removes the member and only then frees the slot
// in afterExit. So the counts are settled once the drain has finished, not
// the moment the fifth task reads Done, and the wait covers both.
testutil.WaitFor(t, "all tasks done and every slot released", func() bool {
return sw.Coordinator().Summarize().Done == 5 && team.Running() == 0 && team.QueueDepth() == 0
})
if got := len(l.recorded()); got != 5 {
t.Fatalf("launched %d members, want 5", got)
}
+4
View File
@@ -35,6 +35,10 @@ func TestMain(m *testing.M) {
return
}
// A developer's ZERO_WEBFETCH_PROVIDER would route every web_fetch test
// through the billed Toolbox backend and its real key lookup; tests that
// need it set it with t.Setenv.
_ = os.Unsetenv("ZERO_WEBFETCH_PROVIDER")
os.Exit(m.Run())
}
+9 -2
View File
@@ -936,8 +936,15 @@ func TestExecCommandInteractiveBlockSuggestsTTY(t *testing.T) {
}
func TestExecCommandTTYSkipsInteractiveBlock(t *testing.T) {
tool := NewScopedExecCommandTool(t.TempDir(), nil, newExecSessionManager())
// ssh to a bogus host fails fast on its own; the point is that the
// The ssh can still be resolving the bogus host when the 250ms yield
// returns, which leaves its session running in the test root. Stop it
// before the root goes, and use execTestRoot for the moment Windows still
// holds the directory after the stop.
root := execTestRoot(t)
manager := newExecSessionManager()
t.Cleanup(func() { manager.StopAll() })
tool := NewScopedExecCommandTool(root, nil, manager)
// ssh to a bogus host fails on its own; the point is that the
// interactive guard must not fire ahead of the tty path.
result := tool.Run(context.Background(), map[string]any{
"cmd": "ssh host.example.com", "tty": true, "yield_time_ms": 250,
+2
View File
@@ -175,6 +175,7 @@ func toolNames(toolset []Tool) []string {
func TestCoreNetworkToolsExposeSafetyMetadata(t *testing.T) {
// web_search is only registered when a backend is configured.
t.Setenv("ZERO_WEBSEARCH_BASE_URL", "https://search.example/api")
t.Setenv("ZERO_WEBSEARCH_PROVIDER", "")
byName := map[string]Tool{}
for _, tool := range CoreNetworkTools() {
byName[tool.Name()] = tool
@@ -213,6 +214,7 @@ func TestCoreNetworkToolsOmitWebSearchWhenUnconfigured(t *testing.T) {
// No backend configured → don't offer web_search (it could only error, which
// makes the model waste calls + prompts before falling back to an MCP search).
t.Setenv("ZERO_WEBSEARCH_BASE_URL", "")
t.Setenv("ZERO_WEBSEARCH_PROVIDER", "")
names := map[string]bool{}
for _, tool := range CoreNetworkTools() {
names[tool.Name()] = true
+1 -1
View File
@@ -73,7 +73,7 @@ func (tool *skillTool) Run(_ context.Context, args map[string]any) Result {
names := make([]string, 0, len(loaded))
for _, skill := range loaded {
if skill.Name == name {
return okResult(skill.Content)
return okResult(skills.ToolOutput(skill))
}
names = append(names, skill.Name)
}
+24
View File
@@ -96,3 +96,27 @@ func TestSkillToolNoSkillsAvailable(t *testing.T) {
t.Fatalf("expected a no-skills message, got: %q", result.Output)
}
}
// A skill that ships files beside its SKILL.md (#584) is no use if the model
// cannot find them, so the output names their directory and lists them.
func TestSkillToolPointsAtTheSkillsBundledFiles(t *testing.T) {
dir := t.TempDir()
writeSkillFile(t, dir, "pdf", "---\nname: pdf\n---\nRun scripts/fill.py.")
script := filepath.Join(dir, "pdf", "scripts", "fill.py")
if err := os.MkdirAll(filepath.Dir(script), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(script, []byte("print('fill')\n"), 0o755); err != nil {
t.Fatal(err)
}
result := NewSkillTool(dir).Run(context.Background(), map[string]any{"name": "pdf"})
if result.Status != StatusOK {
t.Fatalf("Status = %s (output: %s)", result.Status, result.Output)
}
for _, want := range []string{"Run scripts/fill.py.", filepath.Join(dir, "pdf"), "- scripts/fill.py"} {
if !strings.Contains(result.Output, want) {
t.Fatalf("output does not contain %q:\n%s", want, result.Output)
}
}
}
+17 -2
View File
@@ -36,6 +36,7 @@ type webFetchTool struct {
baseTool
client *http.Client
resolver webFetchResolver
aimlapi webFetchAIMLAPIBackend
}
type webFetchResolver interface {
@@ -118,14 +119,14 @@ func newWebFetchToolWithClientAndResolver(client *http.Client, resolver webFetch
},
"max_bytes": {
Type: "integer",
Description: "Maximum raw response body bytes to download before conversion.",
Description: "Maximum raw response body bytes to download before conversion. With the AI/ML API Toolbox backend (ZERO_WEBFETCH_PROVIDER=aimlapi), the maximum bytes of crawled Markdown returned.",
Default: defaultWebFetchMaxBytes,
Minimum: intPtr(1),
Maximum: intPtr(maxWebFetchMaxBytes),
},
"format": {
Type: "string",
Description: "auto (default): HTML responses are converted to compact markdown, everything else is returned as-is. raw: never convert. markdown: force conversion.",
Description: "auto (default): HTML responses are converted to compact markdown, everything else is returned as-is. raw: never convert. markdown: force conversion. With the AI/ML API Toolbox backend (ZERO_WEBFETCH_PROVIDER=aimlapi), the page is always returned as crawled Markdown and raw is rejected.",
Enum: []string{"auto", "raw", "markdown"},
Default: "auto",
},
@@ -147,6 +148,7 @@ func newWebFetchToolWithClientAndResolver(client *http.Client, resolver webFetch
},
client: client,
resolver: resolver,
aimlapi: defaultWebFetchAIMLAPIBackend(),
}
}
@@ -165,6 +167,15 @@ func (tool webFetchTool) RejectBeforePermission(args map[string]any) (Result, bo
if err := validateWebFetchURLBeforePermission(rawURL); err != nil {
return errorResult("Error: Unsafe URL for web_fetch: " + err.Error()), true
}
if webFetchAIMLAPISelected() {
format, err := stringArg(args, "format", "auto", false)
if err != nil {
return errorResult("Error: Invalid arguments for web_fetch: " + err.Error()), true
}
if strings.EqualFold(strings.TrimSpace(format), "raw") {
return webFetchAIMLAPIRawFormatError(), true
}
}
return Result{}, false
}
@@ -199,6 +210,10 @@ func (tool webFetchTool) run(ctx context.Context, args map[string]any) Result {
if err := validateWebFetchURLBeforePermission(rawURL); err != nil {
return errorResult("Error: Unsafe URL for web_fetch: " + err.Error())
}
// Opt-in only: an AI/ML API key alone never switches the backend.
if webFetchAIMLAPISelected() {
return tool.runAIMLAPI(ctx, rawURL, maxBytes, format)
}
parsedURL, err := validateWebFetchURL(ctx, rawURL, tool.resolver)
if err != nil {
+523
View File
@@ -0,0 +1,523 @@
package tools
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"os"
"strconv"
"strings"
"time"
"unicode/utf8"
"github.com/Twigpine/zero/internal/aimlapi"
"github.com/Twigpine/zero/internal/config"
)
// The AI/ML API Toolbox backend fetches the page as Markdown through a
// Toolbox tool run instead of a direct GET. Toolbox runs are billed to the
// user's AI/ML API balance, so it is opt-in only: it runs when
// ZERO_WEBFETCH_PROVIDER=aimlapi is set, never just because a key exists.
const (
webFetchProviderEnv = "ZERO_WEBFETCH_PROVIDER"
webFetchProviderAIMLAPI = "aimlapi"
// webFetchAIMLAPITool is the Toolbox tool web_fetch runs: the lowest fixed
// price ($0.0013 per run), main-content Markdown that keeps links, and a 404
// or unresolvable host comes back as an is_error run, reported as an error.
webFetchAIMLAPITool = "search1api/crawl"
// webFetchAIMLAPIMaxCostUSD is sent with every run. It leaves room for a
// small price change; a run quoted above it is refused before it starts
// (400 max_cost_exceeded) instead of being charged.
webFetchAIMLAPIMaxCostUSD = 0.005
// webFetchAIMLAPIWaitMS is how long POST /run may hold the request open
// before answering 202 with a run to poll.
webFetchAIMLAPIWaitMS = 30000
// webFetchAIMLAPITimeout bounds the whole scrape, retries and polling
// included. Each request is also bounded, with room for wait_ms.
webFetchAIMLAPITimeout = 120 * time.Second
webFetchAIMLAPIRequestTimeout = 60 * time.Second
webFetchAIMLAPIMaxAttempts = 3
webFetchAIMLAPIBackoff = 500 * time.Millisecond
webFetchAIMLAPIMaxRetryDelay = 10 * time.Second
webFetchAIMLAPIDefaultPoll = time.Second
webFetchAIMLAPIMinPoll = 250 * time.Millisecond
webFetchAIMLAPIMaxPoll = 5 * time.Second
// webFetchAIMLAPIResponseLimit bounds a run object; it holds at most a page
// of Markdown, which max_bytes then cuts to the caller's budget.
webFetchAIMLAPIResponseLimit = 8 * 1024 * 1024
// webFetchAIMLAPIErrorReadLimit is read from an error body before the key
// is scrubbed; the shown text is cut to webFetchErrorBodyLimit afterwards,
// so the cut can never leave a fragment of the key behind.
webFetchAIMLAPIErrorReadLimit = 64 * 1024
)
var errWebFetchAIMLAPIRedirect = errors.New("the AI/ML API Toolbox answered with a redirect; redirects are refused so the key never leaves tools.aimlapi.com")
// webFetchAIMLAPIBackend holds the seams of the Toolbox backend. baseURL is
// always aimlapi.ToolboxBaseURL outside tests; it is not configurable.
type webFetchAIMLAPIBackend struct {
client *http.Client
baseURL string
serviceKey func() (string, error)
idempotencyKey func() (string, error)
sleep func(context.Context, time.Duration) error
}
func defaultWebFetchAIMLAPIBackend() webFetchAIMLAPIBackend {
return webFetchAIMLAPIBackend{
baseURL: aimlapi.ToolboxBaseURL,
serviceKey: config.LoadAIMLAPIServiceKey,
// A UUIDv4 generator; it is used here for the run's idempotency key.
idempotencyKey: aimlapi.NewPaymentSessionID,
sleep: sleepWebFetchContext,
}
}
// webFetchAIMLAPISelected reads ZERO_WEBFETCH_PROVIDER on every call. Unset,
// empty and unknown values keep the direct fetch.
func webFetchAIMLAPISelected() bool {
return strings.EqualFold(strings.TrimSpace(os.Getenv(webFetchProviderEnv)), webFetchProviderAIMLAPI)
}
type webFetchToolboxRunRequest struct {
Tool string `json:"tool"`
Input webFetchToolboxRunInput `json:"input"`
MaxCostUSD float64 `json:"max_cost_usd"`
WaitMS int `json:"wait_ms"`
IdempotencyKey string `json:"idempotency_key"`
}
type webFetchToolboxRunInput struct {
URL string `json:"url"`
}
type webFetchToolboxRun struct {
ID string `json:"id"`
Status string `json:"status"`
IsError bool `json:"is_error"`
Output json.RawMessage `json:"output"`
Error json.RawMessage `json:"error"`
NextPollHintMS *float64 `json:"next_poll_hint_ms"`
Cost struct {
ActualUSD *float64 `json:"actual_usd"`
OverMax bool `json:"over_max"`
} `json:"cost"`
}
type webFetchToolboxRunError struct {
Message string `json:"message"`
Code string `json:"code"`
Kind string `json:"kind"`
Retryable bool `json:"retryable"`
}
// runAIMLAPI fetches target through the Toolbox. Any failure is an error
// result: there is no fallback to a direct fetch.
func (tool webFetchTool) runAIMLAPI(ctx context.Context, rawURL string, maxBytes int, format string) Result {
// The pre-permission and resolved-host checks run before the URL leaves
// the machine, so localhost, private and metadata URLs never reach a
// third-party crawler. A resolver is required here: there is no dial-time
// pin to fall back on, since the Toolbox dials the target.
parsedURL, err := validateWebFetchURL(ctx, rawURL, firstWebFetchResolver(tool.resolver))
if err != nil {
return errorResult("Error: Unsafe URL for web_fetch: " + err.Error())
}
// The Toolbox returns the crawler's Markdown, never the original response,
// so format "raw" cannot be honoured. Refusing it (before any billed run)
// is clearer than returning Markdown the caller asked not to get.
if format == "raw" {
return webFetchAIMLAPIRawFormatError()
}
backend := tool.aimlapi
key, err := backend.serviceKey()
if err == nil && strings.TrimSpace(key) == "" {
err = config.ErrAIMLAPIKeyMissing
}
if err != nil {
return errorResult("Error: ZERO_WEBFETCH_PROVIDER=aimlapi needs a usable AI/ML API key: " + redactWebFetchText(scrubWebFetchKey(err.Error(), key)) + ". web_fetch does not fall back to a direct fetch; unset ZERO_WEBFETCH_PROVIDER to fetch pages directly.")
}
scrub := func(value string) string {
return redactWebFetchText(scrubWebFetchKey(value, key))
}
run, statusCode, err := backend.scrape(ctx, key, parsedURL.String())
if err != nil {
return errorResult("Error fetching URL via AI/ML API Toolbox (" + webFetchAIMLAPITool + "): " + scrub(err.Error()))
}
capNote := ""
if run.Cost.OverMax {
capNote = fmt.Sprintf(" The run hit the $%g max_cost_usd cap.", webFetchAIMLAPIMaxCostUSD)
}
if run.IsError || run.Status != "completed" {
reason := "run ended with status " + firstNonEmptyString(run.Status, "unknown")
if runError := decodeWebFetchToolboxRunError(run.Error); runError.Message != "" {
reason = runError.Message
if runError.Retryable {
reason += " (retryable)"
}
}
return errorResult("Error fetching URL via AI/ML API Toolbox (" + webFetchAIMLAPITool + "): " + scrub(reason) + "." + capNote)
}
content := webFetchToolboxContent(run.Output)
if strings.TrimSpace(content) == "" {
return errorResult("Error fetching URL via AI/ML API Toolbox (" + webFetchAIMLAPITool + "): the run returned no page content." + capNote)
}
// Redact before cutting to max_bytes, so the cut cannot split a secret
// into a fragment the patterns no longer match.
body := scrub(content)
truncated := false
if len(body) > maxBytes {
cut := maxBytes
for cut > 0 && !utf8.RuneStart(body[cut]) {
cut--
}
body = body[:cut]
truncated = true
}
finalURL := redactWebFetchURL(parsedURL)
headers := []string{
"URL: " + finalURL,
"Status: completed (AI/ML API Toolbox run)",
"Content-Type: text/markdown",
"Bytes: " + strconv.Itoa(len(body)),
"Backend: " + webFetchProviderAIMLAPI + " (" + webFetchAIMLAPITool + ")",
}
meta := map[string]string{
"url": finalURL,
"status_code": strconv.Itoa(statusCode),
"content_type": "text/markdown",
"bytes": strconv.Itoa(len(body)),
"truncated": strconv.FormatBool(truncated),
"converted": "false",
"backend": webFetchProviderAIMLAPI,
"tool": webFetchAIMLAPITool,
}
if run.ID != "" {
meta["run_id"] = scrub(run.ID)
}
if run.Cost.ActualUSD != nil {
meta["cost_usd"] = strconv.FormatFloat(*run.Cost.ActualUSD, 'f', -1, 64)
}
return Result{
Status: StatusOK,
Output: strings.Join(append(headers, "", body), "\n"),
Truncated: truncated,
Meta: meta,
}
}
// scrape starts one run and waits for it to finish. It returns the terminal
// run and the HTTP status of the response that carried it.
func (backend webFetchAIMLAPIBackend) scrape(ctx context.Context, key string, target string) (webFetchToolboxRun, int, error) {
ctx, cancel := context.WithTimeout(ctx, webFetchAIMLAPITimeout)
defer cancel()
idempotencyKey, err := backend.idempotencyKey()
if err != nil {
return webFetchToolboxRun{}, 0, fmt.Errorf("create idempotency key: %w", err)
}
// One body for every attempt: a retry carries the same idempotency key, so
// a run that did start is returned again instead of being billed twice.
body, err := json.Marshal(webFetchToolboxRunRequest{
Tool: webFetchAIMLAPITool,
Input: webFetchToolboxRunInput{URL: target},
MaxCostUSD: webFetchAIMLAPIMaxCostUSD,
WaitMS: webFetchAIMLAPIWaitMS,
IdempotencyKey: idempotencyKey,
})
if err != nil {
return webFetchToolboxRun{}, 0, err
}
client := backend.httpClient()
var run webFetchToolboxRun
statusCode := 0
for attempt := 1; ; attempt++ {
response, err := backend.send(ctx, &client, http.MethodPost, backend.baseURL+"/run", body, key)
if err != nil {
return webFetchToolboxRun{}, 0, err
}
if response.StatusCode >= 200 && response.StatusCode < 300 {
run, err = decodeWebFetchToolboxRun(response)
if err != nil {
return webFetchToolboxRun{}, 0, err
}
statusCode = response.StatusCode
break
}
// 400 max_cost_exceeded and other refusals are final; only rate limits
// and gateway errors are retried.
if !webFetchAIMLAPIRetryable(response.StatusCode) || attempt >= webFetchAIMLAPIMaxAttempts {
return webFetchToolboxRun{}, 0, webFetchToolboxRequestError(response, "run", key)
}
delay := webFetchAIMLAPIRetryDelay(response.Header.Get("Retry-After"), attempt)
discardWebFetchBody(response)
if err := backend.sleep(ctx, delay); err != nil {
return webFetchToolboxRun{}, 0, err
}
}
pollRetryDelay := time.Duration(0)
for !webFetchToolboxTerminal(run.Status) {
if run.ID == "" {
return webFetchToolboxRun{}, 0, errors.New("the run is still pending but has no run id")
}
if err := backend.sleep(ctx, max(webFetchAIMLAPIPollDelay(run.NextPollHintMS), pollRetryDelay)); err != nil {
return webFetchToolboxRun{}, 0, err
}
response, err := backend.send(ctx, &client, http.MethodGet, backend.baseURL+"/runs/"+url.PathEscape(run.ID), nil, key)
if err != nil {
return webFetchToolboxRun{}, 0, err
}
pollRetryDelay = 0
if webFetchAIMLAPIRetryable(response.StatusCode) {
if retryAfter := strings.TrimSpace(response.Header.Get("Retry-After")); retryAfter != "" {
pollRetryDelay = webFetchAIMLAPIRetryDelay(retryAfter, 1)
}
discardWebFetchBody(response)
continue
}
if response.StatusCode < 200 || response.StatusCode >= 300 {
return webFetchToolboxRun{}, 0, webFetchToolboxRequestError(response, "run status", key)
}
next, err := decodeWebFetchToolboxRun(response)
if err != nil {
return webFetchToolboxRun{}, 0, err
}
if next.ID == "" {
next.ID = run.ID
}
run, statusCode = next, response.StatusCode
}
return run, statusCode, nil
}
// httpClient is a plain client for the fixed Toolbox endpoint. It does not use
// the SSRF-pinning transport meant for target sites, and like Zero's other
// aimlapi.com clients it uses the default transport, which honours
// HTTP_PROXY/HTTPS_PROXY/NO_PROXY.
func (backend webFetchAIMLAPIBackend) httpClient() http.Client {
client := http.Client{Timeout: webFetchAIMLAPIRequestTimeout}
if backend.client != nil {
client = *backend.client
}
// Never follow a redirect: the key and attribution headers must not reach
// any host but the Toolbox.
client.CheckRedirect = func(*http.Request, []*http.Request) error {
return errWebFetchAIMLAPIRedirect
}
return client
}
func (backend webFetchAIMLAPIBackend) send(ctx context.Context, client *http.Client, method string, endpoint string, body []byte, key string) (*http.Response, error) {
var reader io.Reader
if body != nil {
reader = bytes.NewReader(body)
}
request, err := http.NewRequestWithContext(ctx, method, endpoint, reader)
if err != nil {
return nil, err
}
// Fail closed: the key and attribution go only to an aimlapi.com service.
if !aimlapi.IsAttributedURL(request.URL) {
return nil, errors.New("refusing to send the AI/ML API key to a host other than tools.aimlapi.com")
}
request.Header.Set("Authorization", "Bearer "+key)
aimlapi.ApplyAttributionHeaders(request.Header, request.URL)
request.Header.Set("User-Agent", "zero-web-fetch/0.1")
request.Header.Set("Accept", "application/json")
if body != nil {
request.Header.Set("Content-Type", "application/json")
}
response, err := client.Do(request)
if err != nil {
// Report the cause without the *url.Error around it, as web_search
// does: that quotes the request URL, and a polling URL carries the run
// id the Toolbox returned percent-escaped, a form of any key echoed in
// it that scrubbing cannot match.
var urlErr *url.Error
if errors.As(err, &urlErr) {
return nil, urlErr.Err
}
return nil, err
}
if response.Body == nil {
response.Body = io.NopCloser(strings.NewReader(""))
}
return response, nil
}
func decodeWebFetchToolboxRun(response *http.Response) (webFetchToolboxRun, error) {
defer response.Body.Close()
raw, truncated, err := readWebFetchBody(response.Body, webFetchAIMLAPIResponseLimit)
if err != nil {
return webFetchToolboxRun{}, fmt.Errorf("read the run: %w", err)
}
if truncated {
return webFetchToolboxRun{}, fmt.Errorf("the run response is larger than %d bytes", webFetchAIMLAPIResponseLimit)
}
var run webFetchToolboxRun
if err := json.Unmarshal([]byte(raw), &run); err != nil {
return webFetchToolboxRun{}, errors.New("the run response is not valid JSON")
}
return run, nil
}
// webFetchToolboxRequestError reports a non-2xx answer with the error envelope
// the Toolbox sends ({"error":{"code","message","request_id"}}). The key is
// scrubbed from the decoded text before it is cut to size.
func webFetchToolboxRequestError(response *http.Response, action string, key string) error {
defer response.Body.Close()
raw, _, _ := readWebFetchBody(response.Body, webFetchAIMLAPIErrorReadLimit)
var envelope struct {
Error struct {
Code string `json:"code"`
Message string `json:"message"`
RequestID string `json:"request_id"`
} `json:"error"`
}
detail := ""
if json.Unmarshal([]byte(raw), &envelope) == nil && (envelope.Error.Code != "" || envelope.Error.Message != "") {
if envelope.Error.Code != "" {
detail += " (" + envelope.Error.Code + ")"
}
if envelope.Error.Message != "" {
detail += ": " + envelope.Error.Message
}
if envelope.Error.RequestID != "" {
detail += " [request_id " + envelope.Error.RequestID + "]"
}
} else if text := strings.TrimSpace(raw); text != "" {
detail = ": " + text
}
detail = scrubWebFetchKey(detail, key)
if len(detail) > webFetchErrorBodyLimit {
detail = detail[:webFetchErrorBodyLimit]
}
message := "AI/ML API Toolbox " + action + " returned HTTP " + webFetchStatusLine(response) + detail
if envelope.Error.Code == "max_cost_exceeded" {
message += fmt.Sprintf(". The run was refused before it started (cap $%g) and was not billed", webFetchAIMLAPIMaxCostUSD)
}
return errors.New(message)
}
func decodeWebFetchToolboxRunError(raw json.RawMessage) webFetchToolboxRunError {
var runError webFetchToolboxRunError
if len(raw) == 0 {
return runError
}
if json.Unmarshal(raw, &runError) == nil {
if runError.Message == "" {
runError.Message = firstNonEmptyString(runError.Code, runError.Kind)
}
return runError
}
var text string
if json.Unmarshal(raw, &text) == nil {
runError.Message = text
}
return runError
}
// webFetchToolboxContent returns output.results.content, the crawler's
// Markdown, or "" when the output has another shape.
func webFetchToolboxContent(raw json.RawMessage) string {
var output struct {
Results struct {
Content json.RawMessage `json:"content"`
} `json:"results"`
}
if len(raw) == 0 || json.Unmarshal(raw, &output) != nil {
return ""
}
var content string
if json.Unmarshal(output.Results.Content, &content) != nil {
return ""
}
return content
}
func webFetchToolboxTerminal(status string) bool {
switch status {
case "completed", "failed", "stopped", "timed_out":
return true
}
return false
}
func webFetchAIMLAPIRetryable(statusCode int) bool {
switch statusCode {
case http.StatusTooManyRequests, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout:
return true
}
return false
}
// webFetchAIMLAPIRetryDelay honours Retry-After (seconds or an HTTP date) and
// otherwise backs off exponentially; either way it is capped.
func webFetchAIMLAPIRetryDelay(retryAfter string, attempt int) time.Duration {
delay := webFetchAIMLAPIBackoff << (attempt - 1)
if value := strings.TrimSpace(retryAfter); value != "" {
if seconds, err := strconv.Atoi(value); err == nil && seconds >= 0 {
// Capped before the multiplication, so a huge value cannot overflow.
delay = time.Duration(min(seconds, int(webFetchAIMLAPIMaxRetryDelay/time.Second))) * time.Second
} else if when, err := http.ParseTime(value); err == nil {
delay = time.Until(when)
}
}
return max(0, min(delay, webFetchAIMLAPIMaxRetryDelay))
}
// webFetchAIMLAPIPollDelay honours next_poll_hint_ms, clamped to 250 ms..5 s.
func webFetchAIMLAPIPollDelay(hintMS *float64) time.Duration {
if hintMS == nil {
return webFetchAIMLAPIDefaultPoll
}
// Clamp in milliseconds first, so a huge hint cannot overflow a Duration.
ms := max(float64(webFetchAIMLAPIMinPoll/time.Millisecond), min(*hintMS, float64(webFetchAIMLAPIMaxPoll/time.Millisecond)))
return time.Duration(ms) * time.Millisecond
}
// scrubWebFetchKey removes an echoed AI/ML API key from text shown to the
// model, whole or split by invisible characters, as web_search does. An empty
// key is no key: replacing "" would rewrite every byte gap.
func scrubWebFetchKey(value string, key string) string {
if key == "" {
return value
}
return scrubAIMLAPIResult(value, key)
}
func discardWebFetchBody(response *http.Response) {
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, webFetchErrorBodyLimit))
_ = response.Body.Close()
}
func sleepWebFetchContext(ctx context.Context, delay time.Duration) error {
if delay <= 0 {
return ctx.Err()
}
timer := time.NewTimer(delay)
defer timer.Stop()
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
return nil
}
}
// webFetchAIMLAPIRawFormatError is shared by preflight and execution.
func webFetchAIMLAPIRawFormatError() Result {
return errorResult(`Error: format "raw" is not available with ZERO_WEBFETCH_PROVIDER=aimlapi: the AI/ML API Toolbox returns the page as Markdown, not the original response. Use format "auto" or "markdown", or unset ZERO_WEBFETCH_PROVIDER to fetch the page directly.`)
}
+842
View File
@@ -0,0 +1,842 @@
package tools
import (
"context"
"encoding/json"
"errors"
"io"
"net/http"
"net/netip"
"net/url"
"path/filepath"
"reflect"
"regexp"
"strconv"
"strings"
"sync"
"testing"
"time"
"github.com/Twigpine/zero/internal/aimlapi"
"github.com/Twigpine/zero/internal/config"
)
const webFetchAIMLAPITestKey = "aiml-test-key-0042"
// Trimmed from a real search1api/crawl run.
const webFetchAIMLAPICompletedRun = `{"id":"5SJmhg3ePaXxHh_BLdlS9","tool":"search1api/crawl","status":"completed","is_error":false,"untrusted":true,` +
`"output":{"crawlParameters":{"url":"https://example.com/"},"results":{"title":"Example Domain","link":"https://example.com/",` +
`"content":"This domain is for use in documentation examples without needing permission.","metadata":{"sourceUrl":"https://example.com/"}}},` +
`"cost":{"quoted_usd":0.0013,"actual_usd":0.0013,"credits":2600,"over_max":false}}`
// Trimmed from a real 400 max_cost_exceeded refusal (same envelope for every tool).
const webFetchAIMLAPIMaxCostBody = `{"error":{"message":"The cheapest source costs 0.0091 USD, above the limit of 0.001 USD.","type":"invalid_request_error",` +
`"code":"max_cost_exceeded","request_id":"hnbQUvsseM1EX9RrgPjBT","refused_by":"request","estimated_cost_usd":0.0091,"max_cost_usd":0.001}}`
var webFetchUUIDPattern = regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$`)
type recordedToolboxRequest struct {
method string
url string
header http.Header
body []byte
}
// webFetchFakeToolbox records every request that reaches the Toolbox client and
// answers from respond, indexed by request order.
type webFetchFakeToolbox struct {
t *testing.T
mu sync.Mutex
requests []recordedToolboxRequest
respond func(index int, request *http.Request) *http.Response
}
func (fake *webFetchFakeToolbox) client() *http.Client {
return &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
var body []byte
if request.Body != nil {
body, _ = io.ReadAll(request.Body)
}
fake.mu.Lock()
index := len(fake.requests)
fake.requests = append(fake.requests, recordedToolboxRequest{
method: request.Method,
url: request.URL.String(),
header: request.Header.Clone(),
body: body,
})
fake.mu.Unlock()
if fake.respond == nil {
fake.t.Errorf("unexpected Toolbox request %s %s", request.Method, request.URL)
return webFetchTestResponse(request, http.StatusInternalServerError, "", ""), nil
}
return fake.respond(index, request), nil
})}
}
func (fake *webFetchFakeToolbox) recorded() []recordedToolboxRequest {
fake.mu.Lock()
defer fake.mu.Unlock()
return append([]recordedToolboxRequest(nil), fake.requests...)
}
type webFetchAIMLAPIHarness struct {
tool webFetchTool
toolbox *webFetchFakeToolbox
direct *[]string
sleeps *[]time.Duration
}
// hermeticWebFetchAIMLAPIEnv redirects every root the real key lookup reads.
func hermeticWebFetchAIMLAPIEnv(t *testing.T) {
t.Helper()
home := t.TempDir()
t.Setenv("HOME", home)
t.Setenv("USERPROFILE", home)
t.Setenv("APPDATA", filepath.Join(home, "AppData", "Roaming"))
t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, ".config"))
t.Setenv("ZERO_CRED_STORAGE", "file")
t.Setenv("ZERO_PROVIDER", "")
t.Setenv("AIMLAPI_API_KEY", "")
t.Setenv("AIMLAPI_INFERENCE_URL", "")
t.Setenv("AIMLAPI_PARTNER_ID", "")
}
// newWebFetchAIMLAPIHarness builds web_fetch with provider set to provider, a
// recording direct client, a fake Toolbox, a fake resolver, a fixed key and a
// sleep that only records.
func newWebFetchAIMLAPIHarness(t *testing.T, provider string, respond func(int, *http.Request) *http.Response) webFetchAIMLAPIHarness {
t.Helper()
hermeticWebFetchAIMLAPIEnv(t)
t.Setenv(webFetchProviderEnv, provider)
var mu sync.Mutex
direct := []string{}
sleeps := []time.Duration{}
directClient := webFetchTestClient(func(request *http.Request) (*http.Response, error) {
mu.Lock()
direct = append(direct, request.URL.String())
mu.Unlock()
return webFetchTestResponse(request, http.StatusOK, "text/plain", "direct body"), nil
})
resolver := webFetchResolverFunc(func(_ context.Context, _ string, host string) ([]netip.Addr, error) {
if host == "private.example" {
return []netip.Addr{netip.MustParseAddr("10.0.0.5")}, nil
}
return []netip.Addr{netip.MustParseAddr("93.184.216.34")}, nil
})
tool := newWebFetchToolWithClientAndResolver(directClient, resolver).(webFetchTool)
toolbox := &webFetchFakeToolbox{t: t, respond: respond}
tool.aimlapi.client = toolbox.client()
tool.aimlapi.serviceKey = func() (string, error) { return webFetchAIMLAPITestKey, nil }
tool.aimlapi.sleep = func(ctx context.Context, delay time.Duration) error {
mu.Lock()
sleeps = append(sleeps, delay)
mu.Unlock()
return ctx.Err()
}
return webFetchAIMLAPIHarness{tool: tool, toolbox: toolbox, direct: &direct, sleeps: &sleeps}
}
func toolboxJSON(request *http.Request, statusCode int, body string) *http.Response {
return webFetchTestResponse(request, statusCode, "application/json", body)
}
func completedToolboxRun(content string) string {
encoded, _ := json.Marshal(content)
return `{"id":"run-1","status":"completed","is_error":false,"output":{"results":{"title":"T","link":"https://example.com/","content":` +
string(encoded) + `}},"cost":{"quoted_usd":0.0013,"actual_usd":0.0013,"over_max":false}}`
}
func decodeToolboxBody(t *testing.T, body []byte) map[string]any {
t.Helper()
var decoded map[string]any
if err := json.Unmarshal(body, &decoded); err != nil {
t.Fatalf("Toolbox body is not JSON: %v: %s", err, body)
}
return decoded
}
func TestWebFetchAIMLAPIKeyAloneKeepsDirectFetch(t *testing.T) {
for _, provider := range []string{"", " ", "firecrawl", "aimlapi-x"} {
t.Run("provider="+provider, func(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, provider, nil)
// A usable key from every source must not switch the backend.
t.Setenv("AIMLAPI_API_KEY", webFetchAIMLAPITestKey)
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/page"})
if result.Status != StatusOK || !strings.Contains(result.Output, "direct body") {
t.Fatalf("expected the direct fetch, got %s: %s", result.Status, result.Output)
}
if strings.Contains(result.Output, "Backend:") || result.Meta["backend"] != "" {
t.Fatalf("direct fetch must not report a backend: %q %#v", result.Output, result.Meta)
}
if len(*h.direct) != 1 || len(h.toolbox.recorded()) != 0 {
t.Fatalf("direct=%v toolbox=%d, want one direct request and no Toolbox request", *h.direct, len(h.toolbox.recorded()))
}
})
}
}
func TestWebFetchAIMLAPISelectedRunsToolboxCrawl(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, " AIMLAPI ", func(_ int, request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusOK, webFetchAIMLAPICompletedRun)
})
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/guide?page=2&q=go"})
if result.Status != StatusOK {
t.Fatalf("expected ok, got %s: %s", result.Status, result.Output)
}
if len(*h.direct) != 0 {
t.Fatalf("the direct fetch must not run, got %v", *h.direct)
}
requests := h.toolbox.recorded()
if len(requests) != 1 {
t.Fatalf("Toolbox requests = %d, want 1", len(requests))
}
request := requests[0]
if request.method != http.MethodPost || request.url != "https://tools.aimlapi.com/v1/tools/run" {
t.Fatalf("request = %s %s, want POST https://tools.aimlapi.com/v1/tools/run", request.method, request.url)
}
body := decodeToolboxBody(t, request.body)
idempotencyKey, _ := body["idempotency_key"].(string)
if !webFetchUUIDPattern.MatchString(idempotencyKey) {
t.Fatalf("idempotency_key = %q, want a UUIDv4", idempotencyKey)
}
want := map[string]any{
"tool": "search1api/crawl",
"input": map[string]any{"url": "https://example.com/guide?page=2&q=go"},
"max_cost_usd": 0.005,
"wait_ms": float64(30000),
"idempotency_key": idempotencyKey,
}
if !reflect.DeepEqual(body, want) {
t.Fatalf("body = %#v\nwant %#v", body, want)
}
for name, value := range map[string]string{
"Authorization": "Bearer " + webFetchAIMLAPITestKey,
"X-AIMLAPI-Source": "agent/zero",
"X-AIMLAPI-Partner-ID": aimlapi.DefaultPartnerID,
"User-Agent": "zero-web-fetch/0.1",
"Content-Type": "application/json",
} {
if got := request.header.Get(name); got != value {
t.Fatalf("header %s = %q, want %q", name, got, value)
}
}
content := "This domain is for use in documentation examples without needing permission."
for _, line := range []string{
"URL: https://example.com/guide?page=2&q=go",
"Status: completed (AI/ML API Toolbox run)",
"Content-Type: text/markdown",
"Bytes: 76",
"Backend: aimlapi (search1api/crawl)",
"\n\n" + content,
} {
if !strings.Contains(result.Output, line) {
t.Fatalf("output missing %q:\n%s", line, result.Output)
}
}
if result.Truncated {
t.Fatal("result must not be truncated")
}
for key, value := range map[string]string{
"status_code": "200", "content_type": "text/markdown", "bytes": "76", "truncated": "false",
"backend": "aimlapi", "tool": "search1api/crawl", "run_id": "5SJmhg3ePaXxHh_BLdlS9", "cost_usd": "0.0013",
} {
if result.Meta[key] != value {
t.Fatalf("meta[%s] = %q, want %q (meta %#v)", key, result.Meta[key], value, result.Meta)
}
}
}
func TestWebFetchAIMLAPIResolvesTheKeyOnEveryCall(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(_ int, request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusOK, webFetchAIMLAPICompletedRun)
})
keys := []string{"first-key", "second-key"}
calls := 0
h.tool.aimlapi.serviceKey = func() (string, error) {
key := keys[calls]
calls++
return key, nil
}
for range keys {
if result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"}); result.Status != StatusOK {
t.Fatalf("expected ok, got %s: %s", result.Status, result.Output)
}
}
requests := h.toolbox.recorded()
if len(requests) != 2 || requests[0].header.Get("Authorization") != "Bearer first-key" || requests[1].header.Get("Authorization") != "Bearer second-key" {
t.Fatalf("each call must resolve the key again, got %d requests", len(requests))
}
}
func TestWebFetchAIMLAPIWithoutUsableKeySendsNothing(t *testing.T) {
for name, keyErr := range map[string]error{
"missing": config.ErrAIMLAPIKeyMissing,
"proxy": config.ErrAIMLAPIProxyProfile,
"invalid": config.ErrAIMLAPIKeyInvalid,
"empty key": nil,
} {
t.Run(name, func(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", nil)
h.tool.aimlapi.serviceKey = func() (string, error) { return "", keyErr }
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"})
assertWebFetchAIMLAPIKeyRefusal(t, h, result)
if keyErr != nil && !strings.Contains(result.Output, keyErr.Error()) {
t.Fatalf("expected the key error in the output, got %q", result.Output)
}
})
}
// The real lookup, against an empty hermetic home.
t.Run("real lookup without a key", func(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", nil)
h.tool.aimlapi.serviceKey = defaultWebFetchAIMLAPIBackend().serviceKey
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"})
assertWebFetchAIMLAPIKeyRefusal(t, h, result)
if !strings.Contains(result.Output, config.ErrAIMLAPIKeyMissing.Error()) {
t.Fatalf("expected the missing-key error, got %q", result.Output)
}
})
// The real lookup withholds an environment key once the inference URL
// points at a proxy, and web_fetch then sends nothing.
t.Run("real lookup with a proxied key", func(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", nil)
h.tool.aimlapi.serviceKey = defaultWebFetchAIMLAPIBackend().serviceKey
t.Setenv("AIMLAPI_API_KEY", webFetchAIMLAPITestKey)
t.Setenv("AIMLAPI_INFERENCE_URL", "https://proxy.example/v1")
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"})
assertWebFetchAIMLAPIKeyRefusal(t, h, result)
if !strings.Contains(result.Output, config.ErrAIMLAPIProxyProfile.Error()) {
t.Fatalf("expected the proxy-profile error, got %q", result.Output)
}
})
}
func assertWebFetchAIMLAPIKeyRefusal(t *testing.T, h webFetchAIMLAPIHarness, result Result) {
t.Helper()
if result.Status != StatusError {
t.Fatalf("expected an error result, got %s: %s", result.Status, result.Output)
}
if !strings.Contains(result.Output, "needs a usable AI/ML API key") || !strings.Contains(result.Output, "does not fall back") {
t.Fatalf("unexpected refusal: %q", result.Output)
}
if len(h.toolbox.recorded()) != 0 || len(*h.direct) != 0 {
t.Fatalf("no request may be sent: toolbox=%d direct=%v", len(h.toolbox.recorded()), *h.direct)
}
}
func TestWebFetchAIMLAPIRefusesUnsafeURLsBeforeTheToolbox(t *testing.T) {
for _, rawURL := range []string{
"http://127.0.0.1/admin",
"http://localhost/status",
"http://169.254.169.254/latest/meta-data",
"http://metadata.google.internal/computeMetadata/v1/",
"https://private.example/internal",
"ftp://example.com/file",
"https://example.com:8443/",
} {
t.Run(rawURL, func(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", nil)
result := h.tool.Run(context.Background(), map[string]any{"url": rawURL})
if result.Status != StatusError || !strings.Contains(result.Output, "Unsafe URL") {
t.Fatalf("expected an unsafe URL error, got %s: %s", result.Status, result.Output)
}
if len(h.toolbox.recorded()) != 0 || len(*h.direct) != 0 {
t.Fatalf("an unsafe URL must not reach anyone: toolbox=%d direct=%v", len(h.toolbox.recorded()), *h.direct)
}
})
}
h := newWebFetchAIMLAPIHarness(t, "aimlapi", nil)
if _, rejected := h.tool.RejectBeforePermission(map[string]any{"url": "http://localhost/"}); !rejected {
t.Fatal("RejectBeforePermission must still refuse localhost with the Toolbox backend")
}
}
func TestWebFetchAIMLAPIRefusesNonToolboxHosts(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", nil)
h.tool.aimlapi.baseURL = "https://tools.aimlapi.com.evil.example/v1/tools"
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"})
if result.Status != StatusError || !strings.Contains(result.Output, "refusing to send the AI/ML API key") {
t.Fatalf("expected a fail-closed refusal, got %s: %s", result.Status, result.Output)
}
if len(h.toolbox.recorded()) != 0 {
t.Fatalf("the key must not leave for a non-Toolbox host, got %d requests", len(h.toolbox.recorded()))
}
}
func TestWebFetchAIMLAPIDefaultClientIsPlainAndRefusesRedirects(t *testing.T) {
client := defaultWebFetchAIMLAPIBackend().httpClient()
// A nil transport is http.DefaultTransport: the proxy environment applies
// and the target-site SSRF pin does not.
if client.Transport != nil {
t.Fatalf("Toolbox client transport = %T, want the default transport", client.Transport)
}
if client.CheckRedirect == nil || client.CheckRedirect(nil, nil) == nil {
t.Fatal("Toolbox client must refuse redirects")
}
if defaultWebFetchAIMLAPIBackend().baseURL != aimlapi.ToolboxBaseURL {
t.Fatalf("baseURL = %q, want %q", defaultWebFetchAIMLAPIBackend().baseURL, aimlapi.ToolboxBaseURL)
}
}
func TestWebFetchAIMLAPIRefusesRedirects(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(index int, request *http.Request) *http.Response {
if index == 0 {
response := toolboxJSON(request, http.StatusFound, "")
response.Header.Set("Location", "https://collector.example/steal")
return response
}
return toolboxJSON(request, http.StatusOK, webFetchAIMLAPICompletedRun)
})
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"})
if result.Status != StatusError || !strings.Contains(result.Output, "redirects are refused") {
t.Fatalf("expected a redirect refusal, got %s: %s", result.Status, result.Output)
}
if requests := h.toolbox.recorded(); len(requests) != 1 {
t.Fatalf("the redirect must not be followed, got %d requests (last %s)", len(requests), requests[len(requests)-1].url)
}
}
func TestWebFetchAIMLAPIPollsAQueuedRun(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(index int, request *http.Request) *http.Response {
switch index {
case 0:
return toolboxJSON(request, http.StatusAccepted, `{"id":"run/7","status":"queued","next_poll_hint_ms":10}`)
case 1:
return toolboxJSON(request, http.StatusServiceUnavailable, `{"error":{"code":"unavailable"}}`)
case 2:
return toolboxJSON(request, http.StatusOK, `{"id":"run/7","status":"running","next_poll_hint_ms":60000}`)
default:
return toolboxJSON(request, http.StatusOK, completedToolboxRun("polled page"))
}
})
deadlineChecked := false
h.tool.aimlapi.sleep = func(ctx context.Context, delay time.Duration) error {
deadline, ok := ctx.Deadline()
if !ok || time.Until(deadline) > webFetchAIMLAPITimeout {
t.Errorf("polling must run under the %s overall deadline (deadline %v, set %v)", webFetchAIMLAPITimeout, deadline, ok)
}
deadlineChecked = true
*h.sleeps = append(*h.sleeps, delay)
return nil
}
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"})
if result.Status != StatusOK || !strings.Contains(result.Output, "polled page") {
t.Fatalf("expected the polled content, got %s: %s", result.Status, result.Output)
}
requests := h.toolbox.recorded()
if len(requests) != 4 {
t.Fatalf("requests = %d, want POST and three polls", len(requests))
}
for _, request := range requests[1:] {
if request.method != http.MethodGet || request.url != "https://tools.aimlapi.com/v1/tools/runs/run%2F7" {
t.Fatalf("poll = %s %s, want GET of the escaped run id", request.method, request.url)
}
if request.header.Get("Authorization") != "Bearer "+webFetchAIMLAPITestKey || request.header.Get("X-AIMLAPI-Source") != "agent/zero" {
t.Fatalf("poll must carry auth and attribution, got %v", request.header)
}
}
// Hints are clamped to 250 ms..5 s; the 503 poll reuses the queued run's hint.
wantSleeps := []time.Duration{250 * time.Millisecond, 250 * time.Millisecond, 5 * time.Second}
if !deadlineChecked || !reflect.DeepEqual(*h.sleeps, wantSleeps) {
t.Fatalf("sleeps = %v, want %v", *h.sleeps, wantSleeps)
}
}
func TestWebFetchAIMLAPIStopsPollingWhenTheContextEnds(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(_ int, request *http.Request) *http.Response {
if request.Method == http.MethodGet {
t.Error("poll GET sent after canceled sleep")
return toolboxJSON(request, http.StatusOK, `{"id":"run-1","status":"failed"}`)
}
return toolboxJSON(request, http.StatusAccepted, `{"id":"run-1","status":"queued"}`)
})
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
h.tool.aimlapi.sleep = func(sleepCtx context.Context, delay time.Duration) error {
cancel()
return sleepWebFetchContext(sleepCtx, delay)
}
result := h.tool.Run(ctx, map[string]any{"url": "https://example.com/"})
if result.Status != StatusError || !strings.Contains(result.Output, "context canceled") {
t.Fatalf("expected a cancellation error, got %s: %s", result.Status, result.Output)
}
}
func TestWebFetchAIMLAPIReportsFailedRuns(t *testing.T) {
for name, body := range map[string]string{
"is_error with completed": `{"id":"run-1","status":"completed","is_error":true,"error":{"message":"Target returned HTTP 404","kind":"provider_error"},` +
`"output":{"results":{"content":"stale page"}}}`,
"failed status": `{"id":"run-1","status":"failed","is_error":false,"output":{"results":{"content":"stale page"}}}`,
"no content": `{"id":"run-1","status":"completed","is_error":false,"output":{"results":{"content":""}}}`,
} {
t.Run(name, func(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(_ int, request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusOK, body)
})
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/missing"})
if result.Status != StatusError || strings.Contains(result.Output, "stale page") {
t.Fatalf("expected an error without content, got %s: %s", result.Status, result.Output)
}
if len(*h.direct) != 0 {
t.Fatalf("a failed scrape must not fall back to a direct fetch, got %v", *h.direct)
}
if name == "is_error with completed" && !strings.Contains(result.Output, "Target returned HTTP 404") {
t.Fatalf("expected the run error message, got %q", result.Output)
}
})
}
}
func TestWebFetchAIMLAPIDoesNotRetryMaxCostExceeded(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(_ int, request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusBadRequest, webFetchAIMLAPIMaxCostBody)
})
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"})
if result.Status != StatusError {
t.Fatalf("expected an error, got %s: %s", result.Status, result.Output)
}
for _, want := range []string{"HTTP 400 Bad Request", "(max_cost_exceeded)", "request_id hnbQUvsseM1EX9RrgPjBT", "not billed"} {
if !strings.Contains(result.Output, want) {
t.Fatalf("output missing %q: %q", want, result.Output)
}
}
if len(h.toolbox.recorded()) != 1 || len(*h.sleeps) != 0 {
t.Fatalf("max_cost_exceeded must not be retried: requests=%d sleeps=%v", len(h.toolbox.recorded()), *h.sleeps)
}
}
func TestWebFetchAIMLAPIRetriesWithTheSameRunRequest(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(index int, request *http.Request) *http.Response {
switch index {
case 0:
response := toolboxJSON(request, http.StatusTooManyRequests, `{"error":{"code":"rate_limited","message":"slow down"}}`)
response.Header.Set("Retry-After", "2")
return response
case 1:
return toolboxJSON(request, http.StatusBadGateway, `{"error":{"code":"bad_gateway"}}`)
default:
return toolboxJSON(request, http.StatusOK, webFetchAIMLAPICompletedRun)
}
})
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"})
if result.Status != StatusOK {
t.Fatalf("expected success after retries, got %s: %s", result.Status, result.Output)
}
requests := h.toolbox.recorded()
if len(requests) != 3 {
t.Fatalf("requests = %d, want 3", len(requests))
}
for _, request := range requests[1:] {
if request.method != http.MethodPost || string(request.body) != string(requests[0].body) {
t.Fatalf("a retry must resend the identical body:\n%s\n%s", requests[0].body, request.body)
}
}
// Retry-After: 2 is honoured; the 502 without one backs off 500 ms << 1.
if want := []time.Duration{2 * time.Second, time.Second}; !reflect.DeepEqual(*h.sleeps, want) {
t.Fatalf("sleeps = %v, want %v", *h.sleeps, want)
}
}
func TestWebFetchAIMLAPIGivesUpAfterTheLastAttempt(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(_ int, request *http.Request) *http.Response {
response := toolboxJSON(request, http.StatusTooManyRequests, `{"error":{"code":"rate_limited","message":"too many concurrent runs"}}`)
response.Header.Set("Retry-After", "3600")
return response
})
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"})
if result.Status != StatusError || !strings.Contains(result.Output, "(rate_limited): too many concurrent runs") {
t.Fatalf("expected the rate-limit error, got %s: %s", result.Status, result.Output)
}
if len(h.toolbox.recorded()) != 3 {
t.Fatalf("requests = %d, want 3 attempts", len(h.toolbox.recorded()))
}
// A huge Retry-After is capped.
if want := []time.Duration{10 * time.Second, 10 * time.Second}; !reflect.DeepEqual(*h.sleeps, want) {
t.Fatalf("sleeps = %v, want %v", *h.sleeps, want)
}
}
func TestWebFetchAIMLAPIUsesAFreshIdempotencyKeyPerCall(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(_ int, request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusOK, webFetchAIMLAPICompletedRun)
})
for range 2 {
if result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"}); result.Status != StatusOK {
t.Fatalf("expected ok, got %s: %s", result.Status, result.Output)
}
}
requests := h.toolbox.recorded()
first, _ := decodeToolboxBody(t, requests[0].body)["idempotency_key"].(string)
second, _ := decodeToolboxBody(t, requests[1].body)["idempotency_key"].(string)
if !webFetchUUIDPattern.MatchString(first) || !webFetchUUIDPattern.MatchString(second) || first == second {
t.Fatalf("each web_fetch call needs its own UUID idempotency key, got %q and %q", first, second)
}
}
func TestWebFetchAIMLAPITruncatesAtMaxBytes(t *testing.T) {
for name, tc := range map[string]struct {
content string
maxBytes int
want string
}{
"ascii": {content: "abcdefghij", maxBytes: 4, want: "abcd"},
"rune boundary": {content: "aé-tail", maxBytes: 2, want: "a"},
} {
t.Run(name, func(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(_ int, request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusOK, completedToolboxRun(tc.content))
})
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/", "max_bytes": tc.maxBytes})
if result.Status != StatusOK || !result.Truncated || result.Meta["truncated"] != "true" {
t.Fatalf("expected a truncated result, got %s truncated=%v meta=%#v", result.Status, result.Truncated, result.Meta)
}
wantBytes := len(tc.want)
if result.Meta["bytes"] != strconv.Itoa(wantBytes) || !strings.Contains(result.Output, "Bytes: "+strconv.Itoa(wantBytes)+"\n") {
t.Fatalf("byte count must describe the cut body, meta=%#v output=%q", result.Meta, result.Output)
}
if !strings.HasSuffix(result.Output, "\n\n"+tc.want) {
t.Fatalf("output must end with %q, got %q", tc.want, result.Output)
}
})
}
}
func TestWebFetchAIMLAPIFormats(t *testing.T) {
content := "# Title\n\n<b>kept</b> as the crawler sent it"
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(_ int, request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusOK, completedToolboxRun(content))
})
raw := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/", "format": "raw"})
if raw.Status != StatusError || !strings.Contains(raw.Output, `format "raw" is not available`) {
t.Fatalf("format raw must be refused, got %s: %s", raw.Status, raw.Output)
}
if len(h.toolbox.recorded()) != 0 {
t.Fatal("format raw must be refused before a billed run")
}
for _, format := range []string{"auto", "markdown"} {
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/", "format": format})
if result.Status != StatusOK || !strings.HasSuffix(result.Output, "\n\n"+content) {
t.Fatalf("format %s must return the crawler's Markdown unchanged, got %s: %q", format, result.Status, result.Output)
}
}
}
func TestWebFetchAIMLAPIRedactsTheKey(t *testing.T) {
// The key contains "/", which a JSON encoder may write as "\/": the key is
// scrubbed from the decoded text, so the escaped echo is caught as well.
const key = "aiml/Zx9-echo-key"
for name, respond := range map[string]func(*http.Request) *http.Response{
"error envelope": func(request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusUnauthorized, `{"error":{"code":"invalid_api_key","message":"key aiml\/Zx9-echo-key is not valid","request_id":"r-1"}}`)
},
"plain error body": func(request *http.Request) *http.Response {
return webFetchTestResponse(request, http.StatusForbidden, "text/plain", "denied for Bearer "+key)
},
"run error": func(request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusOK, `{"id":"r","status":"completed","is_error":true,"error":{"message":"bad key `+key+`"}}`)
},
"content": func(request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusOK, completedToolboxRun("page echoing "+key))
},
// The shown body is cut to 4 KiB; the key straddles that cut, so it
// must be scrubbed before the cut or a fragment of it survives.
"key at the cut": func(request *http.Request) *http.Response {
return webFetchTestResponse(request, http.StatusForbidden, "text/plain", strings.Repeat("x", webFetchErrorBodyLimit-8)+key)
},
// A key split by an invisible character reads whole once a terminal or
// log viewer drops that character, so it must be scrubbed as well.
"split key in error body": func(request *http.Request) *http.Response {
return webFetchTestResponse(request, http.StatusForbidden, "text/plain", "denied for Bearer aiml/Zx9\u200b-echo-key")
},
"split key in run error": func(request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusOK, `{"id":"r","status":"completed","is_error":true,"error":{"message":"bad key aiml/Zx9\u0000-echo-key"}}`)
},
"split key in content": func(request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusOK, completedToolboxRun("page echoing aiml/Zx9\u200b-echo-key"))
},
"split key in run id": func(request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusOK, strings.Replace(completedToolboxRun("page"), `"id":"run-1"`, `"id":"run aiml/Zx9\u0000-echo-key"`, 1))
},
} {
t.Run(name, func(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(_ int, request *http.Request) *http.Response { return respond(request) })
h.tool.aimlapi.serviceKey = func() (string, error) { return key, nil }
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"})
if output := withoutInvisibleRunes(result.Output); strings.Contains(output, key) || strings.Contains(output, "Zx9-echo-key") || strings.Contains(output, "aiml/Z") {
t.Fatalf("the key leaked: %q", result.Output)
}
for metaKey, value := range result.Meta {
if strings.Contains(withoutInvisibleRunes(value), "Zx9-echo-key") {
t.Fatalf("the key leaked in meta[%s]: %q", metaKey, value)
}
}
// The cut case keeps only the start of the marker.
if !strings.Contains(result.Output, "[REDAC") && !strings.Contains(result.Meta["run_id"], "[REDAC") {
t.Fatalf("expected a redaction marker, got %q", result.Output)
}
})
}
}
// A transport failure is reported by its cause, not by the *url.Error around
// it: that quotes the polling URL, whose run id the Toolbox chose and which
// carries any key echoed in it percent-escaped, past the key scrubbing.
func TestWebFetchAIMLAPITransportErrorOmitsPollingURL(t *testing.T) {
for name, separator := range map[string]string{"NUL": `\u0000`, "ZWSP": `\u200b`} {
t.Run(name, func(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", nil)
queued := `{"id":"run_` + splitAIMLAPIKey(webFetchAIMLAPITestKey, separator) + `","status":"queued","next_poll_hint_ms":1}`
polls := 0
h.tool.aimlapi.client = &http.Client{Transport: roundTripFunc(func(request *http.Request) (*http.Response, error) {
if request.Method == http.MethodPost {
return toolboxJSON(request, http.StatusAccepted, queued), nil
}
polls++
return nil, errors.New("connection reset by peer")
})}
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"})
if polls != 1 {
t.Fatalf("polls = %d, want the failing poll to be sent once", polls)
}
decoded, err := url.PathUnescape(result.Output)
if err != nil {
decoded = result.Output
}
if strings.Contains(withoutInvisibleRunes(decoded), webFetchAIMLAPITestKey) {
t.Fatalf("the AI/ML API key leaked into the output: %q", result.Output)
}
if strings.Contains(result.Output, "/runs/") {
t.Fatalf("the polling URL leaked into the output: %q", result.Output)
}
if result.Status != StatusError || !strings.Contains(result.Output, "connection reset by peer") {
t.Fatalf("status = %v, output = %q, want the transport failure", result.Status, result.Output)
}
})
}
}
func TestWebFetchAIMLAPIRetryDelay(t *testing.T) {
for _, tc := range []struct {
header string
attempt int
want time.Duration
}{
{"", 1, 500 * time.Millisecond},
{"", 2, time.Second},
{"3", 1, 3 * time.Second},
{"0", 1, 0},
{"99999999999999999999", 1, 500 * time.Millisecond},
{"99999999999999999", 1, 10 * time.Second},
{"86400", 1, 10 * time.Second},
{"soon", 2, time.Second},
} {
if got := webFetchAIMLAPIRetryDelay(tc.header, tc.attempt); got != tc.want {
t.Errorf("retry delay(%q, %d) = %v, want %v", tc.header, tc.attempt, got, tc.want)
}
}
if got := webFetchAIMLAPIRetryDelay(time.Now().Add(time.Hour).UTC().Format(http.TimeFormat), 1); got != 10*time.Second {
t.Errorf("an HTTP-date Retry-After an hour out = %v, want the 10s cap", got)
}
if err := sleepWebFetchContext(context.Background(), 0); err != nil {
t.Fatalf("zero sleep = %v", err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
if err := sleepWebFetchContext(ctx, time.Hour); !errors.Is(err, context.Canceled) {
t.Fatalf("cancelled sleep = %v, want context.Canceled", err)
}
}
func TestWebFetchAIMLAPIRawRejectedBeforePermission(t *testing.T) {
tool := newWebFetchToolWithClientAndResolver(nil, nil)
for _, provider := range []string{"aimlapi", ""} {
t.Setenv("ZERO_WEBFETCH_PROVIDER", provider)
for _, format := range []string{"raw", " RAW ", "markdown", "auto"} {
result, rejected := tool.(interface {
RejectBeforePermission(map[string]any) (Result, bool)
}).RejectBeforePermission(map[string]any{"url": "https://example.com/", "format": format})
want := provider == "aimlapi" && strings.EqualFold(strings.TrimSpace(format), "raw")
if rejected != want {
t.Fatalf("provider=%q format=%q: rejected=%v result=%s", provider, format, rejected, result.Output)
}
}
}
}
func TestWebFetchAIMLAPIPollHonorsRetryAfter(t *testing.T) {
for _, header := range []string{"2", time.Now().Add(4 * time.Second).UTC().Format(http.TimeFormat)} {
t.Run(header, func(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(index int, request *http.Request) *http.Response {
if index == 0 {
return toolboxJSON(request, http.StatusAccepted, `{"id":"run-1","status":"queued","next_poll_hint_ms":250}`)
}
if index == 1 {
response := toolboxJSON(request, http.StatusTooManyRequests, `{}`)
response.Header.Set("Retry-After", header)
return response
}
return toolboxJSON(request, http.StatusOK, webFetchAIMLAPICompletedRun)
})
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/"})
if result.Status != StatusOK {
t.Fatalf("fetch failed: %s", result.Output)
}
if sleeps := *h.sleeps; len(sleeps) != 2 || sleeps[1] < 2*time.Second {
t.Fatalf("Retry-After not honored: %v", sleeps)
}
})
}
}
func TestWebFetchAIMLAPIOversizedRunIsRejected(t *testing.T) {
h := newWebFetchAIMLAPIHarness(t, "aimlapi", func(_ int, request *http.Request) *http.Response {
return toolboxJSON(request, http.StatusOK, completedToolboxRun(strings.Repeat("x", webFetchAIMLAPIResponseLimit)))
})
result := h.tool.Run(context.Background(), map[string]any{"url": "https://example.com/", "max_bytes": 32})
if result.Status != StatusError || !strings.Contains(result.Output, "larger than") {
t.Fatalf("oversized envelope accepted: %s", result.Output)
}
}
+43 -8
View File
@@ -14,6 +14,7 @@ import (
"strings"
"time"
"github.com/Twigpine/zero/internal/config"
"github.com/Twigpine/zero/internal/redaction"
zeroSandbox "github.com/Twigpine/zero/internal/sandbox"
)
@@ -43,6 +44,18 @@ type searchBackend interface {
Search(ctx context.Context, query string, limit int) ([]searchResult, error)
}
// timedSearchBackend is implemented by a backend that needs a time budget other
// than webSearchTimeout.
type timedSearchBackend interface {
searchTimeout() time.Duration
}
// domainScopedSearchBackend is implemented by a backend that can restrict a
// search to an allowlist of hostnames itself. Run still filters the results.
type domainScopedSearchBackend interface {
searchDomains(ctx context.Context, query string, limit int, domains []string) ([]searchResult, error)
}
type webSearchTool struct {
baseTool
backend searchBackend
@@ -132,20 +145,33 @@ func (tool webSearchTool) Run(ctx context.Context, args map[string]any) Result {
}
if tool.backend == nil {
return errorResult("Error: no search backend configured. Set ZERO_WEBSEARCH_BASE_URL (and ZERO_WEBSEARCH_API_KEY) to enable web_search.")
return errorResult("Error: no search backend configured. Set ZERO_WEBSEARCH_BASE_URL (and ZERO_WEBSEARCH_API_KEY), or ZERO_WEBSEARCH_PROVIDER=aimlapi to use the AI/ML API Toolbox (billed per search), to enable web_search.")
}
runCtx, cancel := context.WithTimeout(ctx, webSearchTimeout)
timeout := webSearchTimeout
if timed, ok := tool.backend.(timedSearchBackend); ok {
timeout = timed.searchTimeout()
}
runCtx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
results, err := tool.backend.Search(runCtx, query, limit)
var results []searchResult
scopedSearch := false
if scoped, ok := tool.backend.(domainScopedSearchBackend); ok && len(domains) > 0 {
scopedSearch = true
results, err = scoped.searchDomains(runCtx, query, limit, domains)
} else {
results, err = tool.backend.Search(runCtx, query, limit)
}
if err != nil {
return errorResult("Error performing web search: " + redactWebSearchText(err.Error()))
}
// Filter by domains BEFORE the empty-result short-circuit so a non-empty
// allowlist that ate every hit surfaces as a clear "no results matched
// domains" error rather than a misleading "no results" message.
if len(results) > 0 && len(domains) > 0 {
// domains" error rather than a misleading "no results" message. A scoped
// backend applied the allowlist itself, so its empty result gets the same
// error; an unscoped backend's empty result is an ordinary "no results".
if len(domains) > 0 && (len(results) > 0 || scopedSearch) {
filtered, err := filterWebSearchByDomains(results, domains)
if err != nil {
return errorResult("Error: " + redactWebSearchText(err.Error()))
@@ -188,9 +214,18 @@ func formatSearchResults(results []searchResult) string {
return strings.Join(lines, "\n")
}
// defaultSearchBackend returns the env-configured generic backend, or nil when
// ZERO_WEBSEARCH_BASE_URL is unset (the tool then reports it as unconfigured).
// defaultSearchBackend returns the env-configured backend, or nil when none is
// configured (the tool then reports it as unconfigured).
//
// ZERO_WEBSEARCH_PROVIDER=aimlapi selects the AI/ML API Toolbox. It ignores
// ZERO_WEBSEARCH_BASE_URL and ZERO_WEBSEARCH_API_KEY: the Toolbox host is fixed
// and the key is the AI/ML API service key. Otherwise the generic backend is
// used when ZERO_WEBSEARCH_BASE_URL is set.
func defaultSearchBackend() searchBackend {
provider := strings.TrimSpace(os.Getenv("ZERO_WEBSEARCH_PROVIDER"))
if strings.EqualFold(provider, aimlapiSearchProvider) {
return newAIMLAPISearchBackend(nil, config.LoadAIMLAPIServiceKey)
}
baseURL := strings.TrimSpace(os.Getenv("ZERO_WEBSEARCH_BASE_URL"))
if baseURL == "" {
return nil
@@ -202,7 +237,7 @@ func defaultSearchBackend() searchBackend {
},
baseURL: baseURL,
apiKey: strings.TrimSpace(os.Getenv("ZERO_WEBSEARCH_API_KEY")),
provider: strings.TrimSpace(os.Getenv("ZERO_WEBSEARCH_PROVIDER")),
provider: provider,
}
}
+325
View File
@@ -0,0 +1,325 @@
package tools
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"strconv"
"strings"
"time"
"github.com/Twigpine/zero/internal/aimlapi"
"github.com/Twigpine/zero/internal/redaction"
)
const (
// aimlapiSearchProvider is the ZERO_WEBSEARCH_PROVIDER value that selects the
// AI/ML API Toolbox. It is the only way to select it: every run is billed to
// the user, so a configured AI/ML API key never enables it by itself.
aimlapiSearchProvider = "aimlapi"
aimlapiSearchTool = "web/search"
// aimlapiSearchMaxCostUSD caps the spend of one search. web/search costs
// $0.0013 per result, at least $0.0091 and at most $0.0221 (20 results), so
// the cap always covers a run; the Toolbox refuses a run it would exceed
// before billing it.
aimlapiSearchMaxCostUSD = 0.025
// Toolbox searches typically finish in about 2-3 s, but a run can queue.
// wait_ms asks the Toolbox to hold the POST open until the run ends or that
// time passes, so most searches need no polling; the rest of the budget is
// left for polling a run that is still queued. The generic backends keep
// webSearchTimeout.
aimlapiSearchTimeout = 30 * time.Second
aimlapiSearchWaitMS = 20_000
aimlapiDefaultPollDelay = time.Second
aimlapiMinPollDelay = 250 * time.Millisecond
aimlapiMaxPollDelay = 5 * time.Second
aimlapiMaxErrorTextRunes = 300
)
// aimlapiSearchBackend runs the Toolbox tool web/search. The destination is
// the fixed aimlapi.ToolboxBaseURL (ZERO_WEBSEARCH_BASE_URL is ignored), and
// the key is the AI/ML API service key, resolved again on every search so a
// provider change is seen by the next call.
type aimlapiSearchBackend struct {
client *http.Client
loadKey func() (string, error)
}
// newAIMLAPISearchBackend builds the backend; a nil transport means
// http.DefaultTransport.
func newAIMLAPISearchBackend(transport http.RoundTripper, loadKey func() (string, error)) *aimlapiSearchBackend {
return &aimlapiSearchBackend{
client: &http.Client{
Transport: transport,
Timeout: aimlapiSearchTimeout,
// Never follow a redirect: the key and the attribution headers must
// reach no host but the Toolbox.
CheckRedirect: func(*http.Request, []*http.Request) error {
return errors.New("refusing redirect from the AI/ML API Toolbox")
},
},
loadKey: loadKey,
}
}
func (backend *aimlapiSearchBackend) searchTimeout() time.Duration {
return aimlapiSearchTimeout
}
func (backend *aimlapiSearchBackend) Search(ctx context.Context, query string, limit int) ([]searchResult, error) {
return backend.searchDomains(ctx, query, limit, nil)
}
// searchDomains passes the normalized allowlist as include_domains so the
// Toolbox filters server-side; Run still applies filterWebSearchByDomains.
func (backend *aimlapiSearchBackend) searchDomains(ctx context.Context, query string, limit int, domains []string) ([]searchResult, error) {
key, err := backend.loadKey()
if err != nil {
return nil, fmt.Errorf("AI/ML API search is unavailable: %w", err)
}
idempotencyKey, err := aimlapi.NewPaymentSessionID()
if err != nil {
return nil, fmt.Errorf("generate idempotency key: %w", err)
}
input := map[string]any{"query": query, "max_results": limit}
if len(domains) > 0 {
input["include_domains"] = domains
}
payload, err := json.Marshal(map[string]any{
"tool": aimlapiSearchTool,
"input": input,
"max_cost_usd": aimlapiSearchMaxCostUSD,
"wait_ms": aimlapiSearchWaitMS,
// A fresh key per search: the Toolbox coalesces a retried POST onto the
// same run instead of billing a second one.
"idempotency_key": idempotencyKey,
})
if err != nil {
return nil, fmt.Errorf("encode AI/ML API search request: %w", err)
}
run, err := backend.do(ctx, http.MethodPost, aimlapi.ToolboxBaseURL+"/run", payload, key)
if err != nil {
return nil, err
}
for !aimlapiTerminalRunStatus(run.Status) && !run.IsError {
if strings.TrimSpace(run.ID) == "" {
return nil, errors.New("AI/ML API search returned a queued run without an id")
}
timer := time.NewTimer(aimlapiPollDelay(run.NextPollHintMS))
select {
case <-ctx.Done():
timer.Stop()
return nil, fmt.Errorf("AI/ML API search run %s did not finish: %w", aimlapiErrorText(run.ID, key), ctx.Err())
case <-timer.C:
}
run, err = backend.do(ctx, http.MethodGet, aimlapi.ToolboxBaseURL+"/runs/"+url.PathEscape(run.ID), nil, key)
if err != nil {
return nil, err
}
}
return aimlapiRunResults(run, key)
}
type aimlapiErrorBody struct {
Code string `json:"code"`
Message string `json:"message"`
RequestID string `json:"request_id"`
}
type aimlapiRun struct {
ID string `json:"id"`
Status string `json:"status"`
IsError bool `json:"is_error"`
NextPollHintMS *float64 `json:"next_poll_hint_ms"`
Error *aimlapiErrorBody `json:"error"`
Output *struct {
Results []struct {
URL string `json:"url"`
Title string `json:"title"`
Snippet string `json:"snippet"`
} `json:"results"`
} `json:"output"`
}
// do sends one Toolbox request and decodes the run object it returns.
func (backend *aimlapiSearchBackend) do(ctx context.Context, method, target string, payload []byte, key string) (aimlapiRun, error) {
var body io.Reader
if payload != nil {
body = bytes.NewReader(payload)
}
request, err := newAIMLAPIToolboxRequest(ctx, method, target, body, key)
if err != nil {
return aimlapiRun{}, err
}
response, err := backend.client.Do(request)
if err != nil {
// Report the cause without the *url.Error around it: that quotes the
// request URL, and a polling URL carries the run id the Toolbox returned
// percent-escaped, a form of any key echoed in it that scrubbing cannot
// match.
cause := err
var urlErr *url.Error
if errors.As(err, &urlErr) {
cause = urlErr.Err
}
return aimlapiRun{}, fmt.Errorf("AI/ML API search request failed: %s", aimlapiErrorText(cause.Error(), key))
}
defer response.Body.Close()
raw, err := io.ReadAll(io.LimitReader(response.Body, webSearchBodyLimit))
if err != nil {
return aimlapiRun{}, fmt.Errorf("read AI/ML API search response: %s", aimlapiErrorText(err.Error(), key))
}
if response.StatusCode < 200 || response.StatusCode >= 300 {
return aimlapiRun{}, aimlapiHTTPError(response, raw, key)
}
var run aimlapiRun
if err := json.Unmarshal(raw, &run); err != nil {
return aimlapiRun{}, fmt.Errorf("AI/ML API search returned an invalid response (HTTP %d)", response.StatusCode)
}
return run, nil
}
// newAIMLAPIToolboxRequest builds a Toolbox request and attaches the key and
// the attribution headers only after checking that target is an aimlapi.com
// service URL; any other target is refused before anything is sent.
func newAIMLAPIToolboxRequest(ctx context.Context, method, target string, body io.Reader, key string) (*http.Request, error) {
request, err := http.NewRequestWithContext(ctx, method, target, body)
if err != nil {
return nil, fmt.Errorf("build AI/ML API search request: %w", err)
}
if !aimlapi.IsAttributedURL(request.URL) {
return nil, fmt.Errorf("AI/ML API search refused to send credentials to %q", request.URL.Redacted())
}
request.Header.Set("Accept", "application/json")
if body != nil {
request.Header.Set("Content-Type", "application/json")
}
request.Header.Set("User-Agent", "zero-web-search/0.1")
request.Header.Set("Authorization", "Bearer "+key)
aimlapi.ApplyAttributionHeaders(request.Header, request.URL)
return request, nil
}
// aimlapiHTTPError reports a refused request by status, error code, message
// and request id. The raw body is never surfaced, and the key is scrubbed from
// every field taken from it.
func aimlapiHTTPError(response *http.Response, raw []byte, key string) error {
var envelope struct {
Error *aimlapiErrorBody `json:"error"`
}
_ = json.Unmarshal(raw, &envelope)
var parts aimlapiErrorBody
if envelope.Error != nil {
parts = *envelope.Error
}
text := fmt.Sprintf("AI/ML API search error HTTP %d", response.StatusCode)
if code := aimlapiErrorText(parts.Code, key); code != "" {
text += " " + code
}
switch {
case response.StatusCode == http.StatusUnauthorized:
text += " (check the AI/ML API key)"
case parts.Code == "insufficient_funds":
text += " (top up your AI/ML API balance)"
case response.StatusCode == http.StatusTooManyRequests:
if seconds, err := strconv.Atoi(strings.TrimSpace(response.Header.Get("Retry-After"))); err == nil && seconds > 0 {
text += fmt.Sprintf(" (rate limited; retry after %ds)", seconds)
} else {
text += " (rate limited)"
}
}
if message := aimlapiErrorText(parts.Message, key); message != "" {
text += ": " + message
}
if requestID := aimlapiErrorText(parts.RequestID, key); requestID != "" {
text += " (request_id: " + requestID + ")"
}
return errors.New(text)
}
// aimlapiRunResults maps a finished run onto search results. A run that ended
// in any status but completed, or reports is_error (which a completed run can),
// is a failure.
func aimlapiRunResults(run aimlapiRun, key string) ([]searchResult, error) {
status := strings.TrimSpace(run.Status)
if status == "" {
status = "unknown"
}
if run.IsError || status != "completed" {
text := fmt.Sprintf("AI/ML API %s run %s", aimlapiSearchTool, aimlapiErrorText(status, key))
if run.IsError && status == "completed" {
text += " with an error"
}
if run.Error != nil {
if code := aimlapiErrorText(run.Error.Code, key); code != "" {
text += " " + code
}
if message := aimlapiErrorText(run.Error.Message, key); message != "" {
text += ": " + message
}
}
if id := aimlapiErrorText(run.ID, key); id != "" {
text += " (run_id: " + id + ")"
}
return nil, errors.New(text)
}
if run.Output == nil || run.Output.Results == nil {
return nil, fmt.Errorf("AI/ML API %s returned no results list", aimlapiSearchTool)
}
out := make([]searchResult, 0, len(run.Output.Results))
for _, item := range run.Output.Results {
if strings.TrimSpace(item.URL) == "" {
continue
}
out = append(out, searchResult{Title: scrubAIMLAPIResult(item.Title, key), URL: scrubAIMLAPIResult(item.URL, key), Snippet: scrubAIMLAPIResult(item.Snippet, key)})
}
return out, nil
}
func aimlapiTerminalRunStatus(status string) bool {
switch strings.TrimSpace(status) {
case "completed", "failed", "stopped", "timed_out":
return true
}
return false
}
// aimlapiPollDelay honours next_poll_hint_ms, clamped to 250 ms..5 s.
func aimlapiPollDelay(hintMS *float64) time.Duration {
if hintMS == nil {
return aimlapiDefaultPollDelay
}
delay := time.Duration(*hintMS * float64(time.Millisecond))
if delay < aimlapiMinPollDelay {
return aimlapiMinPollDelay
}
if delay > aimlapiMaxPollDelay {
return aimlapiMaxPollDelay
}
return delay
}
// aimlapiErrorText scrubs the key from a field taken from a response, then
// collapses whitespace and clips it. Scrubbing runs first so clipping can never
// leave a fragment of the key behind.
func aimlapiErrorText(value, key string) string {
text := redaction.RedactString(value, redaction.Options{ExtraSecretValues: []string{key}})
text = strings.Join(strings.Fields(text), " ")
if runes := []rune(text); len(runes) > aimlapiMaxErrorTextRunes {
text = string(runes[:aimlapiMaxErrorTextRunes]) + "…"
}
return text
}
// scrubAIMLAPIResult redacts the active credential without clipping page content.
func scrubAIMLAPIResult(value, key string) string {
return redaction.RedactString(value, redaction.Options{ExtraSecretValues: []string{key}})
}
+702
View File
@@ -0,0 +1,702 @@
package tools
import (
"context"
"encoding/json"
"errors"
"io"
"net/http"
"net/url"
"path/filepath"
"regexp"
"strings"
"sync"
"testing"
"time"
"unicode"
"github.com/Twigpine/zero/internal/aimlapi"
"github.com/Twigpine/zero/internal/config"
)
// Trimmed from a real 200 web/search run.
const aimlapiCompletedRun = `{"id":"h3X2expTh3D7vd4pwyxiY","tool":"web/search","status":"completed","is_error":false,"untrusted":true,
"output":{"results":[
{"url":"https://pkg.go.dev/net/http","title":"http package - net/http - Go Packages","snippet":"For control over proxies, TLS configuration, keep-alives..."},
{"url":"https://go.dev/src/net/http/client.go?m=text","title":"go.dev","snippet":"// Timeout specifies a time limit for requests made by this Client."},
{"url":"https://stackoverflow.com/questions/16895294/how-to-set-timeout-for-http-get-requests-in-golang","title":"How to set timeout for http.Get() requests in Golang?","snippet":"Apparently in Go 1.3 http.Client has Timeout field","published_at":"2013-11-02T00:00:00.000Z"}
]},
"cost":{"quoted_usd":0.0091,"actual_usd":0.0091,"over_max":false}}`
// A real 400 upfront refusal (not billed).
const aimlapiMaxCostExceeded = `{"error":{"message":"The cheapest source costs 0.0091 USD, above the limit of 0.001 USD.","type":"invalid_request_error","code":"max_cost_exceeded","request_id":"hnbQUvsseM1EX9RrgPjBT","refused_by":"request","estimated_cost_usd":0.0091,"max_cost_usd":0.001}}`
type toolboxCall struct {
method string
url string
header http.Header
body map[string]any
remaining time.Duration
}
// fakeToolbox is an http.RoundTripper that records every request and answers
// from respond, so no request leaves the process.
type fakeToolbox struct {
mu sync.Mutex
calls []toolboxCall
respond func(index int, request *http.Request) *http.Response
}
func (fake *fakeToolbox) RoundTrip(request *http.Request) (*http.Response, error) {
call := toolboxCall{method: request.Method, url: request.URL.String(), header: request.Header.Clone()}
if deadline, ok := request.Context().Deadline(); ok {
call.remaining = time.Until(deadline)
}
if request.Body != nil {
raw, _ := io.ReadAll(request.Body)
_ = request.Body.Close()
if len(raw) > 0 {
_ = json.Unmarshal(raw, &call.body)
}
}
fake.mu.Lock()
index := len(fake.calls)
fake.calls = append(fake.calls, call)
fake.mu.Unlock()
// Like a real transport, refuse a request whose context has already ended.
if err := request.Context().Err(); err != nil {
return nil, err
}
response := fake.respond(index, request)
response.Request = request
return response, nil
}
func (fake *fakeToolbox) recorded() []toolboxCall {
fake.mu.Lock()
defer fake.mu.Unlock()
return append([]toolboxCall(nil), fake.calls...)
}
func toolboxResponse(status int, body string) *http.Response {
return &http.Response{
StatusCode: status,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(body)),
}
}
func alwaysRespond(status int, body string) func(int, *http.Request) *http.Response {
return func(int, *http.Request) *http.Response { return toolboxResponse(status, body) }
}
func staticKey(key string) func() (string, error) {
return func() (string, error) { return key, nil }
}
// isolateAIMLAPISearchEnv points every config, cache and credential root at a
// temporary home and clears the variables the backend selection and the key
// resolution read.
func isolateAIMLAPISearchEnv(t *testing.T) {
t.Helper()
home := t.TempDir()
t.Setenv("HOME", home)
t.Setenv("USERPROFILE", home)
t.Setenv("XDG_CONFIG_HOME", filepath.Join(home, ".config"))
t.Setenv("XDG_CACHE_HOME", filepath.Join(home, ".cache"))
t.Setenv("APPDATA", filepath.Join(home, "AppData", "Roaming"))
t.Setenv("LOCALAPPDATA", filepath.Join(home, "AppData", "Local"))
t.Setenv("ZERO_CRED_STORAGE", "file")
t.Setenv("ZERO_PROVIDER", "")
t.Setenv("AIMLAPI_API_KEY", "")
t.Setenv("AIMLAPI_INFERENCE_URL", "")
t.Setenv("AIMLAPI_PARTNER_ID", "")
t.Setenv("ZERO_WEBSEARCH_PROVIDER", "")
t.Setenv("ZERO_WEBSEARCH_BASE_URL", "")
t.Setenv("ZERO_WEBSEARCH_API_KEY", "")
}
func hasWebSearchTool(tools []Tool) bool {
for _, tool := range tools {
if tool.Name() == "web_search" {
return true
}
}
return false
}
func TestWebSearchAIMLAPISelectedOnlyByExplicitProvider(t *testing.T) {
cases := []struct {
name string
provider string
baseURL string
want bool
}{
{name: "key alone", provider: "", want: false},
{name: "other provider", provider: "searxng", want: false},
{name: "explicit", provider: "aimlapi", want: true},
{name: "case-insensitive", provider: " AIMLAPI ", want: true},
{name: "base URL ignored", provider: "aimlapi", baseURL: "https://search.example/api", want: true},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
isolateAIMLAPISearchEnv(t)
t.Setenv("AIMLAPI_API_KEY", "aiml-env-key")
t.Setenv("ZERO_WEBSEARCH_PROVIDER", tc.provider)
t.Setenv("ZERO_WEBSEARCH_BASE_URL", tc.baseURL)
_, isAIMLAPI := defaultSearchBackend().(*aimlapiSearchBackend)
if isAIMLAPI != tc.want {
t.Fatalf("defaultSearchBackend() is aimlapi = %v, want %v", isAIMLAPI, tc.want)
}
if registered := hasWebSearchTool(CoreNetworkTools()); registered != tc.want && tc.baseURL == "" {
t.Fatalf("web_search registered = %v, want %v", registered, tc.want)
}
})
}
}
func TestWebSearchAIMLAPIDefaultBackendUsesServiceKeyAndFixedHost(t *testing.T) {
isolateAIMLAPISearchEnv(t)
t.Setenv("ZERO_WEBSEARCH_PROVIDER", "aimlapi")
t.Setenv("ZERO_WEBSEARCH_BASE_URL", "https://search.example/api")
t.Setenv("ZERO_WEBSEARCH_API_KEY", "zws-generic-key")
t.Setenv("AIMLAPI_API_KEY", "aiml-env-key")
backend, ok := defaultSearchBackend().(*aimlapiSearchBackend)
if !ok {
t.Fatalf("defaultSearchBackend() = %T, want *aimlapiSearchBackend", defaultSearchBackend())
}
fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, aimlapiCompletedRun)}
backend.client.Transport = fake // keeps the production redirect policy and key loader
res := newWebSearchToolWithBackend(backend).Run(context.Background(), map[string]any{"query": "go http timeout"})
if res.Status != StatusOK {
t.Fatalf("status = %v, output = %q", res.Status, res.Output)
}
calls := fake.recorded()
if len(calls) != 1 {
t.Fatalf("requests = %d, want 1", len(calls))
}
if calls[0].url != aimlapi.ToolboxBaseURL+"/run" {
t.Fatalf("request URL = %q, want the fixed Toolbox URL", calls[0].url)
}
if got := calls[0].header.Get("Authorization"); got != "Bearer aiml-env-key" {
t.Fatalf("Authorization = %q, want the AI/ML API service key", got)
}
}
func TestAIMLAPISearchSendsExactRequest(t *testing.T) {
isolateAIMLAPISearchEnv(t)
fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, aimlapiCompletedRun)}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key")))
for range 2 {
res := tool.Run(context.Background(), map[string]any{"query": "go http timeout", "limit": 3})
if res.Status != StatusOK {
t.Fatalf("status = %v, output = %q", res.Status, res.Output)
}
for _, want := range []string{
"1. http package - net/http - Go Packages — https://pkg.go.dev/net/http",
" For control over proxies",
"3. How to set timeout for http.Get() requests in Golang?",
} {
if !strings.Contains(res.Output, want) {
t.Fatalf("output missing %q:\n%s", want, res.Output)
}
}
}
calls := fake.recorded()
if len(calls) != 2 {
t.Fatalf("requests = %d, want 2", len(calls))
}
uuid := regexp.MustCompile(`^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$`)
for _, call := range calls {
if call.method != http.MethodPost || call.url != "https://tools.aimlapi.com/v1/tools/run" {
t.Fatalf("request = %s %s, want POST https://tools.aimlapi.com/v1/tools/run", call.method, call.url)
}
for header, want := range map[string]string{
"Authorization": "Bearer test-aiml-key",
aimlapi.SourceHeaderName: "agent/zero",
aimlapi.PartnerHeaderName: aimlapi.DefaultPartnerID,
"User-Agent": "zero-web-search/0.1",
"Content-Type": "application/json",
} {
if got := call.header.Get(header); got != want {
t.Fatalf("%s = %q, want %q", header, got, want)
}
}
body := call.body
if body["tool"] != "web/search" {
t.Fatalf("tool = %v, want web/search", body["tool"])
}
if body["max_cost_usd"] != 0.025 {
t.Fatalf("max_cost_usd = %v, want 0.025", body["max_cost_usd"])
}
if body["wait_ms"] != float64(20000) {
t.Fatalf("wait_ms = %v, want 20000", body["wait_ms"])
}
input, _ := body["input"].(map[string]any)
if input["query"] != "go http timeout" || input["max_results"] != float64(3) {
t.Fatalf("input = %v, want query and max_results 3", input)
}
if _, ok := input["include_domains"]; ok {
t.Fatalf("include_domains sent without a domains argument: %v", input)
}
if key, _ := body["idempotency_key"].(string); !uuid.MatchString(key) {
t.Fatalf("idempotency_key = %q, want a UUIDv4", key)
}
// The aimlapi budget replaces the generic 10 s one, and the server-side
// wait stays inside it.
if call.remaining <= webSearchTimeout || call.remaining > aimlapiSearchTimeout {
t.Fatalf("request deadline in %v, want within (%v, %v]", call.remaining, webSearchTimeout, aimlapiSearchTimeout)
}
}
if calls[0].body["idempotency_key"] == calls[1].body["idempotency_key"] {
t.Fatalf("idempotency_key reused across searches: %v", calls[0].body["idempotency_key"])
}
if time.Duration(aimlapiSearchWaitMS)*time.Millisecond >= aimlapiSearchTimeout {
t.Fatalf("wait_ms %d must stay below the %v budget", aimlapiSearchWaitMS, aimlapiSearchTimeout)
}
}
func TestAIMLAPISearchResolvesKeyOnEveryCall(t *testing.T) {
isolateAIMLAPISearchEnv(t)
fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, aimlapiCompletedRun)}
loads := 0
backend := newAIMLAPISearchBackend(fake, func() (string, error) {
loads++
if loads == 1 {
return "first-key", nil
}
return "", config.ErrAIMLAPIKeyMissing
})
tool := newWebSearchToolWithBackend(backend)
if res := tool.Run(context.Background(), map[string]any{"query": "q"}); res.Status != StatusOK {
t.Fatalf("first search: status = %v, output = %q", res.Status, res.Output)
}
// The key was removed after the first search: the next one must see that.
res := tool.Run(context.Background(), map[string]any{"query": "q"})
if res.Status != StatusError || !strings.Contains(res.Output, "no AI/ML API key is configured") {
t.Fatalf("second search: status = %v, output = %q", res.Status, res.Output)
}
if calls := fake.recorded(); len(calls) != 1 || calls[0].header.Get("Authorization") != "Bearer first-key" {
t.Fatalf("requests = %+v, want exactly the first search", calls)
}
}
func TestAIMLAPISearchSendsNothingWithoutUsableKey(t *testing.T) {
cases := []struct {
name string
env map[string]string
want string
}{
{name: "missing", want: config.ErrAIMLAPIKeyMissing.Error()},
{name: "proxy", env: map[string]string{"AIMLAPI_API_KEY": "proxy-key", "AIMLAPI_INFERENCE_URL": "https://proxy.example/v1"}, want: config.ErrAIMLAPIProxyProfile.Error()},
{name: "invalid", env: map[string]string{"AIMLAPI_API_KEY": "two words"}, want: config.ErrAIMLAPIKeyInvalid.Error()},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
isolateAIMLAPISearchEnv(t)
for name, value := range tc.env {
t.Setenv(name, value)
}
fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, aimlapiCompletedRun)}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, config.LoadAIMLAPIServiceKey))
res := tool.Run(context.Background(), map[string]any{"query": "q"})
if res.Status != StatusError || !strings.Contains(res.Output, tc.want) {
t.Fatalf("status = %v, output = %q, want an error mentioning %q", res.Status, res.Output, tc.want)
}
if calls := fake.recorded(); len(calls) != 0 {
t.Fatalf("sent %d requests without a usable key", len(calls))
}
})
}
}
func TestAIMLAPISearchPollsQueuedRun(t *testing.T) {
isolateAIMLAPISearchEnv(t)
fake := &fakeToolbox{respond: func(index int, _ *http.Request) *http.Response {
switch index {
case 0:
return toolboxResponse(http.StatusAccepted, `{"id":"run/1","status":"queued","next_poll_hint_ms":1}`)
case 1:
return toolboxResponse(http.StatusOK, `{"id":"run/1","status":"running","is_error":false,"next_poll_hint_ms":1}`)
default:
return toolboxResponse(http.StatusOK, aimlapiCompletedRun)
}
}}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key")))
res := tool.Run(context.Background(), map[string]any{"query": "q"})
if res.Status != StatusOK || !strings.Contains(res.Output, "https://pkg.go.dev/net/http") {
t.Fatalf("status = %v, output = %q", res.Status, res.Output)
}
calls := fake.recorded()
if len(calls) != 3 {
t.Fatalf("requests = %d, want POST then two polls", len(calls))
}
for _, call := range calls[1:] {
if call.method != http.MethodGet || call.url != "https://tools.aimlapi.com/v1/tools/runs/run%2F1" {
t.Fatalf("poll = %s %s, want GET of the run status", call.method, call.url)
}
if call.header.Get("Authorization") != "Bearer test-aiml-key" || call.header.Get(aimlapi.SourceHeaderName) != "agent/zero" {
t.Fatalf("poll headers = %v, want the key and attribution", call.header)
}
}
}
func TestAIMLAPIPollDelayIsClamped(t *testing.T) {
ms := func(value float64) *float64 { return &value }
cases := []struct {
hint *float64
want time.Duration
}{
{hint: nil, want: time.Second},
{hint: ms(1), want: 250 * time.Millisecond},
{hint: ms(1200), want: 1200 * time.Millisecond},
{hint: ms(60_000), want: 5 * time.Second},
}
for _, tc := range cases {
if got := aimlapiPollDelay(tc.hint); got != tc.want {
t.Fatalf("aimlapiPollDelay(%v) = %v, want %v", tc.hint, got, tc.want)
}
}
}
func TestAIMLAPISearchStopsPollingWhenContextEnds(t *testing.T) {
isolateAIMLAPISearchEnv(t)
fake := &fakeToolbox{respond: alwaysRespond(http.StatusAccepted, `{"id":"run_1","status":"queued","next_poll_hint_ms":5000}`)}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key")))
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
started := time.Now()
res := tool.Run(ctx, map[string]any{"query": "q"})
if elapsed := time.Since(started); elapsed > 2*time.Second {
t.Fatalf("search took %v after the context ended", elapsed)
}
if res.Status != StatusError || !strings.Contains(res.Output, "did not finish") {
t.Fatalf("status = %v, output = %q", res.Status, res.Output)
}
}
func TestAIMLAPISearchCompletedRunWithIsErrorFails(t *testing.T) {
isolateAIMLAPISearchEnv(t)
fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, `{"id":"run_9","status":"completed","is_error":true,
"error":{"code":"source_failed","message":"upstream search failed"},
"output":{"results":[{"url":"https://example.com/partial","title":"partial"}]}}`)}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key")))
res := tool.Run(context.Background(), map[string]any{"query": "q"})
if res.Status != StatusError {
t.Fatalf("status = %v, output = %q, want an error", res.Status, res.Output)
}
for _, want := range []string{"source_failed", "upstream search failed", "run_9"} {
if !strings.Contains(res.Output, want) {
t.Fatalf("output missing %q: %q", want, res.Output)
}
}
if strings.Contains(res.Output, "example.com/partial") {
t.Fatalf("results of a failed run were shown: %q", res.Output)
}
}
func TestAIMLAPISearchTerminalFailureStatuses(t *testing.T) {
isolateAIMLAPISearchEnv(t)
for _, status := range []string{"failed", "stopped", "timed_out"} {
t.Run(status, func(t *testing.T) {
fake := &fakeToolbox{respond: func(index int, _ *http.Request) *http.Response {
if index > 0 { // a poll would mean the status was not treated as terminal
return toolboxResponse(http.StatusInternalServerError, `{}`)
}
return toolboxResponse(http.StatusOK, `{"id":"run_2","status":"`+status+`","is_error":false}`)
}}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key")))
res := tool.Run(context.Background(), map[string]any{"query": "q"})
if res.Status != StatusError || !strings.Contains(res.Output, "run "+status) {
t.Fatalf("status = %v, output = %q", res.Status, res.Output)
}
if calls := fake.recorded(); len(calls) != 1 {
t.Fatalf("requests = %d, want no polling after a terminal status", len(calls))
}
})
}
}
func TestAIMLAPISearchMalformedResponses(t *testing.T) {
isolateAIMLAPISearchEnv(t)
cases := map[string]struct{ body, want string }{
"not json": {body: `<html>oops</html>`, want: "invalid response"},
"no results": {body: `{"id":"run_3","status":"completed","is_error":false,"output":{}}`, want: "no results list"},
}
for name, tc := range cases {
t.Run(name, func(t *testing.T) {
fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, tc.body)}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key")))
res := tool.Run(context.Background(), map[string]any{"query": "q"})
if res.Status != StatusError || !strings.Contains(res.Output, tc.want) {
t.Fatalf("status = %v, output = %q, want an error mentioning %q", res.Status, res.Output, tc.want)
}
})
}
}
func TestAIMLAPISearchMaxCostExceededReportsCodeAndRequestID(t *testing.T) {
isolateAIMLAPISearchEnv(t)
fake := &fakeToolbox{respond: alwaysRespond(http.StatusBadRequest, aimlapiMaxCostExceeded)}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key")))
res := tool.Run(context.Background(), map[string]any{"query": "q"})
if res.Status != StatusError {
t.Fatalf("status = %v, output = %q", res.Status, res.Output)
}
for _, want := range []string{"HTTP 400", "max_cost_exceeded", "cheapest source costs", "request_id: hnbQUvsseM1EX9RrgPjBT"} {
if !strings.Contains(res.Output, want) {
t.Fatalf("output missing %q: %q", want, res.Output)
}
}
}
func TestAIMLAPISearchRateLimitedReportsRetryAfter(t *testing.T) {
isolateAIMLAPISearchEnv(t)
fake := &fakeToolbox{respond: func(int, *http.Request) *http.Response {
response := toolboxResponse(http.StatusTooManyRequests, `{"error":{"code":"rate_limited","message":"slow down","request_id":"req_429"}}`)
response.Header.Set("Retry-After", "7")
return response
}}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key")))
res := tool.Run(context.Background(), map[string]any{"query": "q"})
for _, want := range []string{"HTTP 429 rate_limited", "retry after 7s", "request_id: req_429"} {
if !strings.Contains(res.Output, want) {
t.Fatalf("output missing %q: %q", want, res.Output)
}
}
}
func TestAIMLAPISearchRefusesRedirect(t *testing.T) {
isolateAIMLAPISearchEnv(t)
for _, location := range []string{"https://collector.example/run", "https://tools.aimlapi.com/v1/tools/run"} {
t.Run(location, func(t *testing.T) {
fake := &fakeToolbox{respond: func(index int, _ *http.Request) *http.Response {
if index == 0 {
response := toolboxResponse(http.StatusTemporaryRedirect, ``)
response.Header.Set("Location", location)
return response
}
return toolboxResponse(http.StatusOK, aimlapiCompletedRun)
}}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key")))
res := tool.Run(context.Background(), map[string]any{"query": "q"})
if res.Status != StatusError || !strings.Contains(res.Output, "refusing redirect") {
t.Fatalf("status = %v, output = %q, want a refused redirect", res.Status, res.Output)
}
if calls := fake.recorded(); len(calls) != 1 {
t.Fatalf("requests = %d, want the redirect not followed", len(calls))
}
})
}
}
func TestAIMLAPISearchRedactsKeyEchoedInErrors(t *testing.T) {
isolateAIMLAPISearchEnv(t)
const key = "zq7Mv2Lp9Xr4Tn8Wk3"
if redacted := redactWebSearchText("key " + key); !strings.Contains(redacted, key) {
t.Fatalf("precondition: the generic redactor already hides this key shape (%q); pick another", redacted)
}
bodies := map[string]struct {
status int
body string
}{
"http error": {status: http.StatusUnauthorized, body: `{"error":{"code":"invalid_api_key","message":"API key ` + key + ` is not valid","request_id":"req_` + key + `"}}`},
"run error": {status: http.StatusOK, body: `{"id":"run_` + key + `","status":"failed","is_error":true,"error":{"code":"bad_auth","message":"rejected ` + key + `"}}`},
// The same key split by characters nobody sees: a reader that drops
// them gets the key back, so the split form has to be redacted too.
"http error split by NUL": {status: http.StatusUnauthorized, body: `{"error":{"code":"invalid_api_key","message":"API key ` + splitAIMLAPIKey(key, `\u0000`) + ` is not valid","request_id":"req_` + splitAIMLAPIKey(key, `\u200b`) + `"}}`},
"run error split by ZWSP": {status: http.StatusOK, body: `{"id":"run_` + splitAIMLAPIKey(key, `\u200b`) + `","status":"failed","is_error":true,"error":{"code":"bad_auth","message":"rejected ` + splitAIMLAPIKey(key, `\u0000`) + `"}}`},
}
for name, tc := range bodies {
t.Run(name, func(t *testing.T) {
fake := &fakeToolbox{respond: alwaysRespond(tc.status, tc.body)}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey(key)))
res := tool.Run(context.Background(), map[string]any{"query": "q"})
if res.Status != StatusError {
t.Fatalf("status = %v, output = %q", res.Status, res.Output)
}
if strings.Contains(withoutInvisibleRunes(res.Output), key) {
t.Fatalf("the AI/ML API key leaked into the output: %q", res.Output)
}
if !strings.Contains(res.Output, "req_") && !strings.Contains(res.Output, "run_") {
t.Fatalf("output lost the request or run id: %q", res.Output)
}
})
}
}
// A failed request must not echo its URL. The polling URL carries the run id
// the Toolbox returned, percent-escaped, and an escaped key ("%00",
// "%E2%80%8B") is not the text the key scrubber matches.
func TestAIMLAPISearchTransportErrorOmitsPollingURL(t *testing.T) {
isolateAIMLAPISearchEnv(t)
const key = "zq7Mv2Lp9Xr4Tn8Wk3"
for name, separator := range map[string]string{"NUL": `\u0000`, "ZWSP": `\u200b`} {
t.Run(name, func(t *testing.T) {
queued := `{"id":"run_` + splitAIMLAPIKey(key, separator) + `","status":"queued","next_poll_hint_ms":1}`
var polls int
transport := roundTripFunc(func(request *http.Request) (*http.Response, error) {
if request.Method == http.MethodPost {
response := toolboxResponse(http.StatusAccepted, queued)
response.Request = request
return response, nil
}
polls++
return nil, errors.New("connection reset by peer")
})
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(transport, staticKey(key)))
res := tool.Run(context.Background(), map[string]any{"query": "q"})
if polls != 1 {
t.Fatalf("polls = %d, want the failing poll to be sent once", polls)
}
decoded, err := url.PathUnescape(res.Output)
if err != nil {
decoded = res.Output
}
if strings.Contains(withoutInvisibleRunes(decoded), key) {
t.Fatalf("the AI/ML API key leaked into the output: %q", res.Output)
}
if strings.Contains(res.Output, "/runs/") {
t.Fatalf("the polling URL leaked into the output: %q", res.Output)
}
if res.Status != StatusError || !strings.Contains(res.Output, "AI/ML API search request failed: connection reset by peer") {
t.Fatalf("status = %v, output = %q, want the transport failure", res.Status, res.Output)
}
})
}
}
func TestAIMLAPIToolboxRequestRefusesOtherDestinations(t *testing.T) {
isolateAIMLAPISearchEnv(t)
for _, target := range []string{
"https://tools.aimlapi.com.evil.example/v1/tools/run",
"http://tools.aimlapi.com/v1/tools/run",
"https://tools.aimlapi.com:8443/v1/tools/run",
"https://[email protected]/v1/tools/run",
"https://search.example/api",
} {
request, err := newAIMLAPIToolboxRequest(context.Background(), http.MethodPost, target, strings.NewReader(`{}`), "test-aiml-key")
if err == nil || request != nil {
t.Fatalf("newAIMLAPIToolboxRequest(%q) = %v, %v; want a refusal", target, request, err)
}
if strings.Contains(err.Error(), "test-aiml-key") {
t.Fatalf("refusal leaked the key: %v", err)
}
}
request, err := newAIMLAPIToolboxRequest(context.Background(), http.MethodGet, aimlapi.ToolboxBaseURL+"/runs/x", nil, "test-aiml-key")
if err != nil {
t.Fatalf("Toolbox URL refused: %v", err)
}
if request.Header.Get("Authorization") != "Bearer test-aiml-key" || request.Header.Get(aimlapi.SourceHeaderName) != "agent/zero" {
t.Fatalf("headers = %v, want key and attribution on the Toolbox URL", request.Header)
}
}
func TestAIMLAPISearchDomainsAreSentAndStillFiltered(t *testing.T) {
isolateAIMLAPISearchEnv(t)
fake := &fakeToolbox{respond: alwaysRespond(http.StatusOK, aimlapiCompletedRun)}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-aiml-key")))
res := tool.Run(context.Background(), map[string]any{
"query": "go http timeout",
"domains": []any{"https://WWW.Go.dev/doc", "pkg.go.dev", "go.dev"},
})
if res.Status != StatusOK {
t.Fatalf("status = %v, output = %q", res.Status, res.Output)
}
calls := fake.recorded()
if len(calls) != 1 {
t.Fatalf("requests = %d, want 1", len(calls))
}
input, _ := calls[0].body["input"].(map[string]any)
if got, _ := json.Marshal(input["include_domains"]); string(got) != `["go.dev","pkg.go.dev"]` {
t.Fatalf("include_domains = %s, want the normalized allowlist", got)
}
// The Toolbox filter is not trusted alone: a result outside the allowlist
// is still dropped.
if strings.Contains(res.Output, "stackoverflow.com") {
t.Fatalf("result outside the allowlist was shown: %q", res.Output)
}
if !strings.Contains(res.Output, "https://pkg.go.dev/net/http") || !strings.Contains(res.Output, "https://go.dev/src/net/http/client.go") {
t.Fatalf("allowed results missing: %q", res.Output)
}
}
func TestWebSearchUnconfiguredHintNamesAIMLAPI(t *testing.T) {
res := newWebSearchToolWithBackend(nil).Run(context.Background(), map[string]any{"query": "q"})
if res.Status != StatusError || !strings.Contains(res.Output, "ZERO_WEBSEARCH_PROVIDER=aimlapi") {
t.Fatalf("status = %v, output = %q", res.Status, res.Output)
}
}
func TestAIMLAPISearchRedactsKeyInSuccessfulResults(t *testing.T) {
isolateAIMLAPISearchEnv(t)
const key = "zq7Mv2Lp9Xr4Tn8Wk3"
fake := &fakeToolbox{respond: func(_ int, _ *http.Request) *http.Response {
return toolboxResponse(http.StatusOK, `{"status":"completed","output":{"results":[{"title":"`+key+`","url":"https://example.com/`+key+`","snippet":"echo `+key+`"}]}}`)
}}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey(key)))
result := tool.Run(context.Background(), map[string]any{"query": "q"})
if result.Status != StatusOK || strings.Contains(result.Output, key) {
t.Fatalf("successful result leaked key: %s", result.Output)
}
}
func TestAIMLAPISearchRedactsSplitKeyInSuccessfulResults(t *testing.T) {
isolateAIMLAPISearchEnv(t)
const key = "zq7Mv2Lp9Xr4Tn8Wk3"
fake := &fakeToolbox{respond: func(_ int, _ *http.Request) *http.Response {
return toolboxResponse(http.StatusOK, `{"status":"completed","output":{"results":[{"title":"`+splitAIMLAPIKey(key, `\u0000`)+`","url":"https://example.com/`+splitAIMLAPIKey(key, `\u200b`)+`","snippet":"echo `+splitAIMLAPIKey(key, `\u2060`)+`"}]}}`)
}}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey(key)))
result := tool.Run(context.Background(), map[string]any{"query": "q"})
if result.Status != StatusOK {
t.Fatalf("status = %v, output = %q", result.Status, result.Output)
}
if strings.Contains(withoutInvisibleRunes(result.Output), key) {
t.Fatalf("successful result leaked the split key: %q", result.Output)
}
}
// splitAIMLAPIKey puts a JSON-escaped invisible character into the middle of
// key, as a response that echoes it that way would.
func splitAIMLAPIKey(key, escapedSeparator string) string {
return key[:len(key)/2] + escapedSeparator + key[len(key)/2:]
}
// withoutInvisibleRunes drops the control and format characters a terminal,
// log viewer or copy-paste would never show, so a key split by one reads whole.
func withoutInvisibleRunes(value string) string {
return strings.Map(func(r rune) rune {
if r != '\t' && r != '\n' && r != '\r' && (unicode.Is(unicode.Cc, r) || unicode.Is(unicode.Cf, r)) {
return -1
}
return r
}, value)
}
func TestAIMLAPISearchEmptyScopedResultsReportAllowlist(t *testing.T) {
isolateAIMLAPISearchEnv(t)
fake := &fakeToolbox{respond: func(_ int, _ *http.Request) *http.Response {
return toolboxResponse(http.StatusOK, `{"status":"completed","output":{"results":[]}}`)
}}
tool := newWebSearchToolWithBackend(newAIMLAPISearchBackend(fake, staticKey("test-key")))
result := tool.Run(context.Background(), map[string]any{"query": "q", "domains": []any{"example.com"}})
if result.Status != StatusError || !strings.Contains(result.Output, "loosen the allowlist") {
t.Fatalf("missing allowlist guidance: %s", result.Output)
}
}
+15
View File
@@ -105,6 +105,7 @@ func TestWebSearchRedactsBackendError(t *testing.T) {
func TestWebSearchRegisteredInCoreNetworkTools(t *testing.T) {
// web_search is registered only when a backend is configured.
t.Setenv("ZERO_WEBSEARCH_BASE_URL", "https://search.example/api")
t.Setenv("ZERO_WEBSEARCH_PROVIDER", "")
found := false
for _, tool := range CoreNetworkTools() {
if tool.Name() == "web_search" {
@@ -300,6 +301,20 @@ func TestWebSearchDomainsFilterEmptyResultIsError(t *testing.T) {
}
}
func TestWebSearchDomainsUnscopedEmptyResultIsNoResults(t *testing.T) {
// A backend that cannot scope by domain never saw the allowlist, so an empty
// result is an ordinary "no results", not an allowlist that ate every hit.
tool := newWebSearchToolWithBackend(&fakeSearchBackend{})
res := tool.Run(context.Background(), map[string]any{
"query": "x",
"domains": []any{"react.dev"},
})
if res.Status != StatusOK || !strings.Contains(res.Output, "No results for query") {
t.Fatalf("want ordinary no-results for an unscoped backend, got %v: %s", res.Status, res.Output)
}
}
func TestWebSearchDomainsFilterToleratesStringSlice(t *testing.T) {
// Some providers send []string instead of []any. The tool must accept both.
backend := &fakeSearchBackend{results: []searchResult{
+76 -8
View File
@@ -91,6 +91,20 @@ func (m model) handleBTWCommand(question string) (model, tea.Cmd) {
parent.flushQueue = nil
side := parent
// The side surface must own its file-view lifecycle. Because `side` is a
// value copy of `parent`, its liveSeq pointer would otherwise alias the
// hidden parent's: closing or switching the view in BTW would revoke the
// parent's in-flight load. Detach it and give the side a fresh lifetime so
// inherited markers cannot masquerade as the parent's snapshot.
side.fileView.liveSeq = nil
if side.fileView.active {
side.fileView.lifetimeToken = nextFileViewLifetimeToken()
side.fileView.loadedToken = [16]byte{}
side.fileView.loadedSeq = 0
side.fileView.loadedRev = 0
side.fileView.snapshotReady = false
side.fileView.renderedContent = ""
}
side.activeSession = fork
side.sessionEvents = events
side.transcript = initialTranscript()
@@ -156,10 +170,16 @@ func (m model) handleBTWCommand(question string) (model, tea.Cmd) {
sideRunIDBase: side.runID,
}
if question == "" {
return side, nil
var btwFileViewCmd tea.Cmd
if side.fileView.active && side.fileView.mode == fileViewFull {
side, btwFileViewCmd = side.startFileViewLoadCmd(side.chatColumnWidth())
}
return side.launchPrompt(question)
if question == "" {
return side, btwFileViewCmd
}
next, launchCmd := side.launchPrompt(question)
return next, batchCommands(btwFileViewCmd, launchCmd)
}
func (m model) leaveBTW() (model, tea.Cmd) {
@@ -175,6 +195,10 @@ func (m model) leaveBTW() (model, tea.Cmd) {
if m.compactInFlight {
return m.appendSystemNotice("BTW compaction is still running. Wait for it to finish before returning."), nil
}
// The side surface is being discarded. Revoke its file-view lifetime token so
// a load still queued behind BTW cannot run or land after the view closes.
// The restored parent keeps its own request and snapshot untouched.
m.revokeFileViewRequest()
m, _ = m.clearLoopsForSessionSwitch()
parent := *m.btw.parent
parent.goalContinuationsSuspended = false
@@ -201,7 +225,9 @@ func (m model) leaveBTW() (model, tea.Cmd) {
parent.resetFlushFrontier("· returned from btw ·")
var goalCmd tea.Cmd
parent, goalCmd = parent.launchGoalContinuationIfReady()
return parent, batchCommands(sweepCmd, spinnerCmd, goalCmd)
var fileCmd tea.Cmd
parent, fileCmd = parent.recoverInvalidatedFileView()
return parent, batchCommands(sweepCmd, spinnerCmd, goalCmd, fileCmd)
}
func btwCommandUnavailable(command parsedCommand) bool {
@@ -236,9 +262,9 @@ func btwCommandUnavailable(command parsedCommand) bool {
// messages have no run ID, so normal BTW routing intentionally leaves them on
// the visible side surface; copying the layout fields prevents stale wrapping
// after the parent is restored.
func (m model) resizeBTWParent(msg tea.WindowSizeMsg) model {
func (m model) resizeBTWParent(msg tea.WindowSizeMsg) (model, tea.Cmd) {
if !m.btw.active || m.btw.parent == nil {
return m
return m, nil
}
parent := *m.btw.parent
parent.width = msg.Width
@@ -247,8 +273,15 @@ func (m model) resizeBTWParent(msg tea.WindowSizeMsg) model {
parent.lineAges = nil
parent.lastStreamActivity = parent.now()
parent.input.SetWidth(maxInt(20, chatWidth(msg.Width)-14))
// The parent's full-file view wrapped at the old width; reload it at the new
// one. Its completion has no run ID, so routeBTWParentMessage matches it by
// lifetime token and delivers it to the hidden parent.
var cmd tea.Cmd
if parent.fileView.active && parent.fileView.mode == fileViewFull {
parent, cmd = parent.startFileViewLoadCmd(parent.chatColumnWidth())
}
m.btw.parent = &parent
return m
return m, cmd
}
func btwTitle(parent string) string {
@@ -272,6 +305,22 @@ func (m model) routeBTWParentMessage(msg tea.Msg) (model, tea.Cmd, bool) {
}
return m.routeBTWMessageToParent(msg)
}
if loaded, ok := msg.(fileViewLoadedMsg); ok {
// File-view completions carry no run ID. Deliver the hidden parent's
// load (matched by lifetime token) so it does not stay stuck on
// "Loading…" while BTW is active; side-owned loads fall through.
if m.btw.parent.fileView.active && m.btw.parent.fileView.lifetimeToken == loaded.lifetimeToken {
return m.routeBTWMessageToParent(msg)
}
return m, nil, false
}
if revalidated, ok := msg.(fileViewRevalidatedMsg); ok {
// A revalidation the parent started before /btw is routed the same way.
if m.btw.parent.fileView.active && m.btw.parent.fileView.lifetimeToken == revalidated.lifetimeToken {
return m.routeBTWMessageToParent(msg)
}
return m, nil, false
}
runID, ok := btwMessageRunID(msg)
if !ok || runID <= 0 || runID >= m.btw.sideRunIDBase {
return m, nil, false
@@ -279,6 +328,21 @@ func (m model) routeBTWParentMessage(msg tea.Msg) (model, tea.Cmd, bool) {
return m.routeBTWMessageToParent(msg)
}
// refreshHiddenParentFileView reloads the hidden parent's full-file view when a
// side-surface mutation advanced the shared path revision. The load completion
// carries the parent's lifetime token, so routeBTWParentMessage delivers it
// back to the hidden model without waiting for a git sweep or for /btw return.
func (m model) refreshHiddenParentFileView() (model, tea.Cmd) {
if !m.btw.active || m.btw.parent == nil {
return m, nil
}
parent := *m.btw.parent
var cmd tea.Cmd
parent, cmd = parent.recoverInvalidatedFileView()
m.btw.parent = &parent
return m, cmd
}
func (m model) routeBTWMessageToParent(msg tea.Msg) (model, tea.Cmd, bool) {
parentNext, cmd := m.btw.parent.updateModel(msg)
parent, ok := parentNext.(model)
@@ -296,7 +360,9 @@ func (m model) routeBTWMessageToParent(msg tea.Msg) (model, tea.Cmd, bool) {
case agentResponseMsg:
m.btw.parentNeedsInput = parent.pendingPermission != nil || parent.pendingAskUser != nil
}
return m, cmd, true
var fileCmd tea.Cmd
m, fileCmd = m.recoverInvalidatedFileView()
return m, batchCommands(cmd, fileCmd), true
}
func btwMessageRunID(msg tea.Msg) (int, bool) {
@@ -327,6 +393,8 @@ func btwMessageRunID(msg tea.Msg) (int, bool) {
return typed.runID, true
case specialistProgressMsg:
return typed.runID, true
case unknownScopeMutationMsg:
return typed.runID, true
case swarmSessionsMsg:
return typed.runID, true
case permissionRequestMsg:
+110
View File
@@ -3,6 +3,7 @@ package tui
import (
"bytes"
"context"
"errors"
"os"
"path/filepath"
"strings"
@@ -535,3 +536,112 @@ func TestBTWCtrlCDuringRunDoesNotClearDraft(t *testing.T) {
t.Fatalf("missing in-flight return guidance: %#v", got.transcript)
}
}
// TestBTWLeaveRevokesSideFileLoad verifies that leaving BTW revokes the side
// surface's file-view lifetime so a still-queued side load cannot run or land
// after the view closes, without cancelling the restored parent's own load.
func TestBTWLeaveRevokesSideFileLoad(t *testing.T) {
resetFileViewCacheForTest()
m := newBTWTestModel(t)
dir := t.TempDir()
m.cwd = dir
name := "side_view.go"
if err := os.WriteFile(filepath.Join(dir, name), []byte("package side\n"), 0o644); err != nil {
t.Fatal(err)
}
// Open a full view so the side surface inherits an active file view.
m, parentLoadCmd := m.openFileView(name)
if parentLoadCmd == nil {
t.Fatal("expected a pending parent load command")
}
parentLiveSeq := m.fileView.liveSeq
if parentLiveSeq == nil {
t.Fatal("parent load token missing")
}
parentSeq := parentLiveSeq.Load()
// Enter BTW: the side surface detaches and schedules its own load.
side, sideLoadCmd := m.handleBTWCommand("")
if !side.btw.active || side.btw.parent == nil {
t.Fatal("expected an active BTW conversation")
}
if sideLoadCmd == nil {
t.Fatal("side surface should schedule its own file load")
}
sideLiveSeq := side.fileView.liveSeq
if sideLiveSeq == nil || sideLiveSeq == parentLiveSeq {
t.Fatal("side load must own a detached lifetime token")
}
sideSeq := sideLiveSeq.Load()
// Return while the side load is still queued.
parent, _ := side.leaveBTW()
if sideLiveSeq.Load() == sideSeq {
t.Fatal("leaveBTW must revoke the side file-view request")
}
if parentLiveSeq.Load() != parentSeq || parentLiveSeq == sideLiveSeq {
t.Fatal("leaveBTW must not revoke the restored parent's request")
}
// The held side command must now be a superseded no-op.
held := sideLoadCmd()
loaded, ok := held.(fileViewLoadedMsg)
if !ok {
t.Fatalf("held side command produced %T, want fileViewLoadedMsg", held)
}
if !errors.Is(loaded.err, errFileViewSuperseded) {
t.Fatalf("queued side load must be superseded after leaving BTW, got %v", loaded.err)
}
// The parent's own load still completes and settles its snapshot.
parent = deliverCommandMessages(t, parent, parentLoadCmd)
if parent.fileView.loading {
t.Fatal("parent load must remain valid after the side revocation")
}
if !strings.Contains(plainRender(t, parent.renderFileViewFull(80)), "package side") {
t.Fatal("restored parent must keep its snapshot")
}
}
// TestBTWForkOfErroredFileViewShowsLoadingPlaceholder verifies that a side
// surface forked from a full view whose last snapshot was a read error does
// not render an empty body while its own load is pending. The inherited error
// snapshot is invalidated by the fork, so the side must show the loading
// placeholder until its load lands, then the side's own error snapshot.
func TestBTWForkOfErroredFileViewShowsLoadingPlaceholder(t *testing.T) {
resetFileViewCacheForTest()
m := newBTWTestModel(t)
m.cwd = t.TempDir()
// A missing file opens straight in full mode and settles on a read error.
m, loadCmd := m.openFileView("missing.go")
if loadCmd == nil || m.fileView.mode != fileViewFull {
t.Fatal("expected a full-view load for a file without recorded edits")
}
m = deliverCommandMessages(t, m, loadCmd)
if !m.fileView.hasError {
t.Fatal("parent full view should settle on a read error")
}
if got := plainRender(t, m.renderFileViewFull(80)); !strings.Contains(got, "Could not read file") {
t.Fatalf("parent should show the read error, got %q", got)
}
side, sideLoadCmd := m.handleBTWCommand("")
if !side.btw.active || sideLoadCmd == nil {
t.Fatal("expected an active BTW surface with its own pending file load")
}
if got := plainRender(t, side.renderFileViewFull(80)); !strings.Contains(got, fileViewLoadingPlaceholder) {
t.Fatalf("pending side load must show the loading placeholder, got %q", got)
}
if got := plainRender(t, side.btw.parent.renderFileViewFull(80)); !strings.Contains(got, "Could not read file") {
t.Fatalf("hidden parent must keep its own error snapshot, got %q", got)
}
side = deliverCommandMessages(t, side, sideLoadCmd)
if side.fileView.loading {
t.Fatal("side load should settle")
}
if got := plainRender(t, side.renderFileViewFull(80)); !strings.Contains(got, "Could not read file") {
t.Fatalf("settled side load should show its read error, got %q", got)
}
}
+9
View File
@@ -128,6 +128,15 @@ func (c *staticRenderCache) stats() renderCacheStats {
return c.statsData
}
func fileViewCacheStatsForTest() fileViewCacheStats {
return defaultFileViewCache.stats()
}
func resetFileViewCacheForTest() {
defaultFileViewCache.clear()
defaultFileViewCache.resetStats()
}
func renderSelectableList(options selectableListOptions) string {
if len(options.Items) == 0 {
return ""
+1356 -70
View File
File diff suppressed because it is too large. Load diff
File diff suppressed because it is too large. Load diff
+55
View File
@@ -0,0 +1,55 @@
//go:build !windows
package tui
import (
"path/filepath"
"strings"
"syscall"
"testing"
"time"
)
func TestFileView_SpecialTargetNonRegularFileRejected(t *testing.T) {
resetFileViewCacheForTest()
dir := t.TempDir()
pipePath := filepath.Join(dir, "test.pipe")
// Create a FIFO if supported on the platform
if err := syscall.Mkfifo(pipePath, 0o600); err != nil {
t.Skipf("mkfifo not supported on this environment: %v", err)
}
res := readFileViewBounded(pipePath, 100, 1000, 10000)
if res.err == nil {
t.Fatal("expected error reading FIFO non-regular file, got nil")
}
if !strings.Contains(res.err.Error(), "non-regular") {
t.Fatalf("expected non-regular error message, got: %v", res.err)
}
}
func TestFileView_FIFONeverBlocksOpen(t *testing.T) {
resetFileViewCacheForTest()
dir := t.TempDir()
pipePath := filepath.Join(dir, "hang_test.pipe")
if err := syscall.Mkfifo(pipePath, 0o600); err != nil {
t.Skipf("mkfifo not supported on this environment: %v", err)
}
done := make(chan fileViewReadResult, 1)
go func() {
res := readFileViewBounded(pipePath, 100, 1000, 10000)
done <- res
}()
select {
case res := <-done:
if res.err == nil || !strings.Contains(res.err.Error(), "non-regular") {
t.Fatalf("expected non-regular error, got: %v", res.err)
}
case <-time.After(500 * time.Millisecond):
t.Fatal("readFileViewBounded blocked on open(2) for FIFO without writer (TOCTOU / missing O_NONBLOCK)")
}
}
+2 -2
View File
@@ -152,12 +152,12 @@ func TestTouchedFilesMergesGitSweep(t *testing.T) {
func TestOpenFileViewGitOnlyFallsBackToFull(t *testing.T) {
m := filesPanelTestModel()
m.gitTouched = []gitSweepFile{{path: "kanban/board.tsx", created: true}}
m = m.openFileView("kanban/board.tsx")
m, _ = m.openFileView("kanban/board.tsx")
if m.fileView.mode != fileViewFull {
t.Fatal("git-only file should open in full mode")
}
m = m.exitFileView()
m = m.openFileView("web/app.js")
m, _ = m.openFileView("web/app.js")
if m.fileView.mode != fileViewDiff {
t.Fatal("a file with edit cards still opens in diff mode")
}
+49 -3
View File
@@ -15,6 +15,8 @@ import (
"sort"
"strings"
tea "charm.land/bubbletea/v2"
"github.com/Twigpine/zero/internal/tools"
)
@@ -357,8 +359,13 @@ func (m *model) setSelectedFile(path string) {
// selectFile marks path as the selected file and scrolls the transcript so its
// most recent edit card is in view; the card tint comes from the renderers
// reading selectedFile (rowTouchesSelectedFile).
func (m model) selectFile(path string) model {
// reading selectedFile (rowTouchesSelectedFile). A second activation of the
// same path drills into the file view and returns any load command.
func (m model) selectFile(path string) (model, tea.Cmd) {
if m.selectedFile == path {
m.runDetailsOpen = false
return m.openFileView(path)
}
rowIndex := m.lastRowIndexForFile(path)
m.setSelectedFile(path)
if offset, ok := m.scrollOffsetForTranscriptRow(rowIndex); ok {
@@ -367,7 +374,46 @@ func (m model) selectFile(path string) model {
m.chatBodyLines = 0
}
}
return m
return m, nil
}
func (m model) runDetailsInnerWidth() int {
overlayWidth := minInt(72, maxInt(40, m.width-8))
overlayWidth = minInt(overlayWidth, m.width)
return maxInt(12, overlayWidth-4)
}
func (m model) runDetailsFileAtMouse(msg tea.MouseMsg) (string, bool) {
if !m.runDetailsVisible() {
return "", false
}
width := chatWidth(m.width)
overlay := m.runDetailsOverlay(width)
hit, ok := m.overlayMouseHit(msg, overlay, width)
if !ok {
return "", false
}
inner := m.runDetailsInnerWidth()
layout := m.runDetailsLayout(inner)
if layout.fileStart < 0 {
return "", false
}
// styledBlockFillTitle contributes a 1-line top border with title before the content rows.
const topBorderHeight = 1
y := hit.y - topBorderHeight
var match fileHit
found := false
for _, h := range layout.fileHits {
if y == layout.fileStart+h.lineOffset {
match = h
found = true
break
}
}
if !found {
return "", false
}
return match.path, true
}
// scrollOffsetForTranscriptRow computes the chatScrollOffset that places the
+252
View File
@@ -6,6 +6,8 @@ import (
tea "charm.land/bubbletea/v2"
"github.com/Twigpine/zero/internal/config"
"github.com/Twigpine/zero/internal/dictation"
"github.com/Twigpine/zero/internal/execution"
"github.com/Twigpine/zero/internal/tools"
)
@@ -296,6 +298,24 @@ func TestSidebarFileLinesOverflowExcludesLiveRow(t *testing.T) {
// TestSidebarHasContentForLiveWrite: the sidebar counts an in-flight first
// write as content, so the FILES pulse is visible before any result row lands.
func TestSidebarFileHitsDistinguishSuffixPaths(t *testing.T) {
m := sidebarTestModel()
m.transcript = append(m.transcript,
transcriptRow{kind: rowToolResult, tool: "edit_file", id: "a", status: tools.StatusOK, changedFiles: []string{"a.go"}},
transcriptRow{kind: rowToolResult, tool: "edit_file", id: "b", status: tools.StatusOK, changedFiles: []string{"dir/a.go"}},
)
_, hits := m.sidebarFileLines(80)
if len(hits) < 2 {
t.Fatalf("want two fileHit rows, got %#v", hits)
}
if hits[0].path == hits[1].path {
t.Fatal("fileHit paths must stay distinct")
}
if hits[0].lineOffset == hits[1].lineOffset {
t.Fatal("fileHit rows must not share a line offset")
}
}
func TestSidebarHasContentForLiveWrite(t *testing.T) {
m := sidebarTestModel()
m.plan.steps = nil // drop the helper's seeded plan: no agents/plan/files now
@@ -309,3 +329,235 @@ func TestSidebarHasContentForLiveWrite(t *testing.T) {
t.Fatal("a live in-flight write must count as sidebar content")
}
}
func TestRunDetailsFileMouseSelectionEndToEnd(t *testing.T) {
m := filesPanelTestModel()
m.width = 80
m.height = 24
m.altScreen = true
m.runDetailsOpen = true
overlay := m.runDetailsOverlay(m.width)
if overlay == "" {
t.Fatal("expected non-empty runDetailsOverlay")
}
inner := m.runDetailsInnerWidth()
layout := m.runDetailsLayout(inner)
if len(layout.fileHits) == 0 || layout.fileStart < 0 {
t.Fatalf("expected file hits in layout, got: %+v", layout)
}
targetHit := layout.fileHits[0] // "internal/tui/sidebar.go"
overlayLines := viewLines(overlay)
left, trimmedLines, overlayWidth := normalizeOverlayBlock(overlayLines, m.width)
if overlayWidth <= 0 || len(trimmedLines) == 0 {
t.Fatalf("invalid overlay block normalization: left=%d, width=%d", left, overlayWidth)
}
rect := m.overlayMouseRect(len(trimmedLines), m.width)
// Click on the target file row: top border is row 0 of the overlay, content starts at row 1
clickY := rect.y + 1 + layout.fileStart + targetHit.lineOffset
clickX := left + 4 // inside the border and padding
clickMsg := testMouseClick(tea.MouseLeft, clickX, clickY)
path, ok := m.runDetailsFileAtMouse(clickMsg)
if !ok || path != targetHit.path {
t.Fatalf("runDetailsFileAtMouse failed: ok=%v, got=%q, want=%q", ok, path, targetHit.path)
}
// Route the click through m.Update to verify end-to-end selection
updated, _ := m.Update(clickMsg)
m = updated.(model)
if m.selectedFile != targetHit.path {
t.Fatalf("expected selectedFile %q after mouse click, got %q", targetHit.path, m.selectedFile)
}
}
func TestRunDetailsFileMouseBlockedBehindHigherPriorityOverlays(t *testing.T) {
baseModel := func(t *testing.T) (model, tea.MouseClickMsg, string) {
t.Helper()
m := filesPanelTestModel()
m.width = 80
m.height = 24
m.altScreen = true
m.runDetailsOpen = true
overlay := m.runDetailsOverlay(m.width)
if overlay == "" {
t.Fatal("expected non-empty runDetailsOverlay")
}
inner := m.runDetailsInnerWidth()
layout := m.runDetailsLayout(inner)
if len(layout.fileHits) == 0 || layout.fileStart < 0 {
t.Fatalf("expected file hits in layout, got: %+v", layout)
}
targetHit := layout.fileHits[0]
overlayLines := viewLines(overlay)
left, trimmedLines, overlayWidth := normalizeOverlayBlock(overlayLines, m.width)
if overlayWidth <= 0 || len(trimmedLines) == 0 {
t.Fatalf("invalid overlay block normalization: left=%d, width=%d", left, overlayWidth)
}
rect := m.overlayMouseRect(len(trimmedLines), m.width)
clickY := rect.y + 1 + layout.fileStart + targetHit.lineOffset
clickX := left + 4
clickMsg := testMouseClick(tea.MouseLeft, clickX, clickY)
if path, ok := m.runDetailsFileAtMouse(clickMsg); !ok || path != targetHit.path {
t.Fatalf("precondition failed: visible Run details click should hit %q, got ok=%v path=%q", targetHit.path, ok, path)
}
return m, clickMsg, targetHit.path
}
t.Run("speech_api_key_prompt_blocks_hidden_run_details_click", func(t *testing.T) {
m, clickMsg, wantPath := baseModel(t)
m.dictation.cfg = config.STTConfig{Provider: config.STTProviderGroq, Model: "whisper-large-v3-turbo"}
m.dictation.saveKey = func(string, string) error { return nil }
m.dictation.phase = dictTranscribing
// Asynchronous dictation auth error opens the speech API-key prompt while Run details stays open.
updated, _ := m.Update(dictationTranscribedMsg{
sessionID: m.dictation.sessionID,
err: &dictation.AuthError{Provider: "groq", Message: "auth error: your API key is missing or invalid"},
})
m = updated.(model)
if !m.runDetailsOpen || m.sttKeyPrompt == nil {
t.Fatalf("expected runDetailsOpen=true and sttKeyPrompt!=nil, got runDetailsOpen=%v sttKeyPrompt=%v", m.runDetailsOpen, m.sttKeyPrompt)
}
if m.runDetailsVisible() {
t.Fatal("Run details must not report visible while the speech API-key prompt owns the viewport")
}
if path, ok := m.runDetailsFileAtMouse(clickMsg); ok || path != "" {
t.Fatalf("runDetailsFileAtMouse must reject clicks behind speech prompt, got ok=%v path=%q", ok, path)
}
// First and second clicks at the shared coordinate must neither select the file nor open fileView.
for i := 0; i < 2; i++ {
var cmd tea.Cmd
updated, cmd = m.Update(clickMsg)
m = updated.(model)
if cmd != nil {
t.Fatalf("click %d behind speech prompt should not return a command", i+1)
}
if m.selectedFile != "" || m.fileView.active {
t.Fatalf("click %d behind speech prompt leaked to Run details: selectedFile=%q fileView.active=%v", i+1, m.selectedFile, m.fileView.active)
}
}
// Dismissing the speech prompt restores Run details to the foreground so the click works again.
updated, _ = m.Update(testKey(tea.KeyEsc))
m = updated.(model)
if m.sttKeyPrompt != nil || !m.runDetailsVisible() {
t.Fatal("Run details should be visible again after dismissing the speech prompt")
}
updated, _ = m.Update(clickMsg)
m = updated.(model)
if m.selectedFile != wantPath {
t.Fatalf("expected selectedFile %q after speech prompt closed, got %q", wantPath, m.selectedFile)
}
})
t.Run("leader_help_overlay_blocks_hidden_run_details_click", func(t *testing.T) {
m, clickMsg, wantPath := baseModel(t)
// Ctrl+X then ? opens the leader help overlay in front of Run details.
updated, _ := m.Update(testKeyCtrl('x'))
m = updated.(model)
updated, _ = m.Update(tea.KeyPressMsg{Code: '?', Text: "?"})
m = updated.(model)
if !m.runDetailsOpen || !m.leaderHelpOverlay {
t.Fatalf("expected runDetailsOpen=true and leaderHelpOverlay=true, got runDetailsOpen=%v leaderHelpOverlay=%v", m.runDetailsOpen, m.leaderHelpOverlay)
}
if m.runDetailsVisible() {
t.Fatal("Run details must not report visible while leader help owns the viewport")
}
if path, ok := m.runDetailsFileAtMouse(clickMsg); ok || path != "" {
t.Fatalf("runDetailsFileAtMouse must reject clicks behind leader help, got ok=%v path=%q", ok, path)
}
for i := 0; i < 2; i++ {
var cmd tea.Cmd
updated, cmd = m.Update(clickMsg)
m = updated.(model)
if cmd != nil {
t.Fatalf("click %d behind leader help should not return a command", i+1)
}
if m.selectedFile != "" || m.fileView.active {
t.Fatalf("click %d behind leader help leaked to Run details: selectedFile=%q fileView.active=%v", i+1, m.selectedFile, m.fileView.active)
}
}
// Dismissing leader help with Esc restores Run details to the foreground.
updated, _ = m.Update(testKey(tea.KeyEsc))
m = updated.(model)
if m.leaderHelpOverlay || !m.runDetailsVisible() {
t.Fatal("Run details should be visible again after dismissing leader help")
}
updated, _ = m.Update(clickMsg)
m = updated.(model)
if m.selectedFile != wantPath {
t.Fatalf("expected selectedFile %q after leader help closed, got %q", wantPath, m.selectedFile)
}
})
}
// TestRunDetailsEnterBlockedWhenRunDetailsHidden pins the keyboard twin of the
// mouse gate above: runDetailsOpen can outlive the visible overlay, and Enter
// must not drill into a FILES row the user can no longer see.
func TestRunDetailsEnterBlockedWhenRunDetailsHidden(t *testing.T) {
// setup returns a model with Run details visible and a FILES row selected
// that Enter would open (verified on a copy), so a later denial can only
// come from the visibility gate.
setup := func(t *testing.T, m model) (model, string) {
t.Helper()
m.transcript = append(m.transcript, transcriptRow{kind: rowToolResult, tool: "write_file", id: "f1",
status: tools.StatusOK, changedFiles: []string{"web/app.js"}})
m.runDetailsOpen = true
layout := m.runDetailsLayout(m.runDetailsInnerWidth())
if !m.runDetailsVisible() || len(layout.fileHits) == 0 {
t.Fatalf("precondition: Run details must be visible with FILES hits, visible=%v hits=%d", m.runDetailsVisible(), len(layout.fileHits))
}
path := layout.fileHits[0].path
m, _ = m.selectFile(path)
if m.selectedFile != path || m.fileView.active {
t.Fatalf("precondition: first activation must only select %q, got selected=%q fileView=%v", path, m.selectedFile, m.fileView.active)
}
control, _ := m.Update(testKey(tea.KeyEnter))
if fv := control.(model).fileView; !fv.active || fv.path != path {
t.Fatalf("precondition: Enter on visible Run details must open %q, got active=%v path=%q", path, fv.active, fv.path)
}
return m, path
}
assertEnterIgnored := func(t *testing.T, m model, path string) {
t.Helper()
if !m.runDetailsOpen || m.runDetailsVisible() {
t.Fatalf("expected runDetailsOpen=true with Run details hidden, got open=%v visible=%v", m.runDetailsOpen, m.runDetailsVisible())
}
updated, cmd := m.Update(testKey(tea.KeyEnter))
m = updated.(model)
if m.fileView.active || cmd != nil {
t.Fatalf("Enter behind hidden Run details opened %q: fileView.active=%v path=%q cmd=%v", path, m.fileView.active, m.fileView.path, cmd != nil)
}
}
t.Run("specialist_card_click_enters_subchat", func(t *testing.T) {
m, path := setup(t, specialistClickTestModel(t, "child-sess-1"))
x, y, ok := specialistCardClickPoint(m, 2)
if !ok {
t.Fatal("expected a specialist card selectable line in the layout")
}
updated, _ := m.Update(testMouseClick(tea.MouseLeft, x, y))
m = updated.(model)
if !m.subchat.active {
t.Fatal("clicking the specialist card beside Run details must enter the subchat")
}
assertEnterIgnored(t, m, path)
})
t.Run("detailed_transcript", func(t *testing.T) {
m, path := setup(t, filesPanelTestModel())
m.transcriptDetailed = true
assertEnterIgnored(t, m, path)
})
}
+1 -1
View File
@@ -111,7 +111,7 @@ func TestAltScreenSettledCacheInvalidatesWithFileSelection(t *testing.T) {
t.Fatal("precondition: settled cache should be populated")
}
m = m.selectFile("internal/tui/sidebar.go")
m, _ = m.selectFile("internal/tui/sidebar.go")
if m.altScreenSettledWidth != 0 {
t.Fatal("selecting a file touched by a settled row must invalidate the cache")
}
+228 -29
View File
@@ -2,7 +2,6 @@ package tui
import (
"context"
"encoding/json"
"fmt"
"os"
"path/filepath"
@@ -780,6 +779,13 @@ type specialistCompleteMsg struct {
errorMsg string
}
// unknownScopeMutationMsg carries the completion of a suppressed tool result
// (e.g. TaskOutput) that may have mutated workspace files without reporting the
// affected paths. It triggers the conservative file-view refresh.
type unknownScopeMutationMsg struct {
runID int
}
// swarmSessionsMsg carries swarm task_id -> member session_id pairs (from
// swarm_collect's Meta) so the AGENTS sidebar rows can drill into a member's
// session like a specialist card.
@@ -1440,13 +1446,18 @@ func batchCommands(cmds ...tea.Cmd) tea.Cmd {
}
func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
var resizeCmd tea.Cmd
if size, ok := msg.(tea.WindowSizeMsg); ok {
m = m.resizeBTWParent(size)
m, resizeCmd = m.resizeBTWParent(size)
}
if next, cmd, routed := m.routeBTWParentMessage(msg); routed {
return next, cmd
return next, batchCommands(cmd, resizeCmd)
}
switch msg := msg.(type) {
case fileViewLoadedMsg:
return m.handleFileViewLoaded(msg)
case fileViewRevalidatedMsg:
return m.handleFileViewRevalidated(msg)
case uv.CellSizeEvent:
if msg.Width > 0 && msg.Height > 0 {
m.petCellPixelWidth = msg.Width
@@ -1509,6 +1520,9 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
m.hasDarkBg = msg.IsDark()
if m.themeMode != themeSystem {
applyTheme(m.themeMode, m.hasDarkBg)
if m.fileView.active && m.fileView.mode == fileViewFull {
return m.startFileViewLoadCmd(m.chatColumnWidth())
}
}
return m, nil
case tea.MouseMsg:
@@ -1764,6 +1778,19 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
case m.runDetailsOpen:
if keyIs(msg, tea.KeyEsc) || m.keyMatch(m.keyBindings.toggleSidebar, msg, func(tea.KeyMsg) bool { return keyCtrl(msg, 'b') }) {
m.runDetailsOpen = false
return m, nil
}
// Enter drills into the selected FILES row only while Run details is
// the overlay on screen, the same gate mouse picks use. The flag can
// outlive the visible overlay (e.g. a click entered a subchat).
if keyIs(msg, tea.KeyEnter) && m.selectedFile != "" && m.runDetailsVisible() {
layout := m.runDetailsLayout(m.runDetailsInnerWidth())
for _, h := range layout.fileHits {
if h.path == m.selectedFile {
return m.selectFile(m.selectedFile)
}
}
return m, nil
}
return m, nil
case m.keyMatch(m.keyBindings.toggleDetailed, msg, func(tea.KeyMsg) bool { return keyCtrl(msg, 'o') }):
@@ -1773,9 +1800,9 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
// (so mid-sentence typing is never hijacked) and no modal is up (so a
// permission prompt / ask-user / wizard keeps its own key handling).
if keyText(msg) == "f" {
return m.setFileViewMode(fileViewFull), nil
return m.setFileViewMode(fileViewFull)
}
return m.setFileViewMode(fileViewDiff), nil
return m.setFileViewMode(fileViewDiff)
case m.keyMatch(m.keyBindings.toggleMouse, msg, func(tea.KeyMsg) bool { return keyCtrl(msg, 'e') }) && canFireComposerGatedToggle(m.keyBindings.toggleMouse, defaultToggleMouseChord, m.composerValue() == ""):
// Release/recapture the mouse so the user can drag-select and copy text
// natively (mouse capture otherwise intercepts terminal selection). The
@@ -2562,7 +2589,12 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
}
// A resumed/idle session may already hold agents; keep their short lifecycle
// fade alive. No-op when the loop is already running or nothing animates.
return m, m.ensureSpinnerTick()
if m.fileView.active && m.fileView.mode == fileViewFull {
var cmd tea.Cmd
m, cmd = m.startFileViewLoadCmd(m.chatColumnWidth())
return m, batchCommands(m.ensureSpinnerTick(), cmd, resizeCmd)
}
return m, batchCommands(m.ensureSpinnerTick(), resizeCmd)
case permissionRequestMsg:
// The agent goroutine that raised this request is BLOCKED waiting on the
// decision callback, so every branch below must resolve it exactly once —
@@ -2692,6 +2724,11 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
if m.exiting && len(m.flushRunIDs) == 0 {
return m.quit()
}
// The cancelled run kept going until it finished, and agentRowMsg
// ignored its late tool results because the run is no longer active.
// A write that completed in that window reported nothing, so give an
// open file view the conservative unknown-scope refresh.
return m.invalidateFileViewForUnknownMutation()
}
return m, nil
}
@@ -2840,6 +2877,11 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
// complete once the turn settles.
var sweepCmd tea.Cmd
m, sweepCmd = m.maybeGitSweep()
// The sweep is a no-op outside a git workspace, so an open full-file view
// still needs the conservative invalidation to pick up mutations whose
// paths were never reported (e.g. a hidden TaskOutput success).
var fileRefreshCmd tea.Cmd
m, fileRefreshCmd = m.invalidateFileViewForUnknownMutation()
// If this run was a loop iteration, advance that loop (schedule its next wake
// or stop it). Done before launchQueuedMessageIfReady so a user's queued prompt
// still wins the immediate re-launch; the loop fires on the next idle tick.
@@ -2871,7 +2913,7 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
if msg.goalAware && !hadQueuedMessage && peerCmd == nil && next.pendingPermission == nil && msg.specReview == nil {
next, goalCmd = next.launchGoalContinuationIfReady()
}
return next, tea.Batch(pendingClearCmd, titleCmd, recapCmd, sweepCmd, queuedCmd, peerCmd, peerApprovalCmd, loopTickCmd, goalCmd)
return next, tea.Batch(pendingClearCmd, titleCmd, recapCmd, sweepCmd, fileRefreshCmd, queuedCmd, peerCmd, peerApprovalCmd, loopTickCmd, goalCmd)
case sessionTitleGeneratedMsg:
return m.handleSessionTitleGenerated(msg)
case recapIdleMsg:
@@ -2899,6 +2941,9 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
}
m.lastCompactResult = &msg.result
m = m.setCompactStatusRow(m.compactText(true))
if msg.hasSessionSnapshot {
return m.refreshFileViewMarkers()
}
return m, nil
case planUpdateMsg:
if msg.runID != m.activeRunID {
@@ -2977,7 +3022,13 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
}
m.transcript = appendTranscriptRow(m.transcript, cardRow)
}
return m, nil
// A specialist (Task) works in a separate workspace-visible process and
// its successful result is suppressed from the transcript, so this is the
// only boundary that sees the mutation. Conservatively invalidate the open
// snapshot: the child may have changed files without a path report.
var refresh tea.Cmd
m, refresh = m.invalidateFileViewForUnknownMutation()
return m, refresh
case specialistProgressMsg:
if msg.runID != m.activeRunID {
return m, nil
@@ -3032,10 +3083,65 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
// A finished command tool may have mutated files git can see but no
// changedFiles reports (npm create, heredoc writes, subagent edits) —
// re-sweep so the FILES sidebar picks them up mid-turn.
if msg.row.kind == rowToolResult && isPlanCommandTool(msg.row.tool) {
var sweep tea.Cmd
m, sweep = m.maybeGitSweep()
return m, sweep
if msg.row.kind == rowToolResult {
var loadCmds []tea.Cmd
invalidated := false
selfNeedsRefresh := false
for _, p := range msg.row.changedFiles {
invalidated = true
target := p
if !filepath.IsAbs(target) {
target = filepath.Join(m.cwd, target)
}
rev := defaultFileViewCache.invalidatePath(target)
if m.fileView.active && (p == m.fileView.path || target == m.fileView.path) {
m.fileView.requiredSourceRev = rev
if m.fileView.mode == fileViewFull {
selfNeedsRefresh = true
}
}
}
// A tool that can mutate files outside exact ChangedFiles (command
// tools, terminal sessions, swarm status/collect, or results with
// aggregated directory summaries) must also run the conservative
// unknown-scope refresh even when exact paths were reported.
unknownScope := isPlanCommandTool(msg.row.tool) ||
toolResultMayMutateUnknownScope(msg.row.tool) ||
len(msg.row.changeSummaries) > 0
if unknownScope {
var sweep tea.Cmd
m, sweep = m.maybeGitSweep()
if sweep != nil {
loadCmds = append(loadCmds, sweep)
}
var refresh tea.Cmd
m, refresh = m.invalidateFileViewForUnknownMutation()
if refresh != nil {
loadCmds = append(loadCmds, refresh)
}
selfNeedsRefresh = false
}
if selfNeedsRefresh {
var cmd tea.Cmd
m, cmd = m.startFileViewRefreshCmd(m.chatColumnWidth())
if cmd != nil {
loadCmds = append(loadCmds, cmd)
}
}
// The handler above only reloads the surface that received the result.
// A side mutation must also wake the hidden parent when that parent is
// showing the same path; the reverse direction is recovered after the
// parent message is routed. Neither wait depends on a git sweep.
if invalidated && !unknownScope {
var peerCmd tea.Cmd
m, peerCmd = m.refreshHiddenParentFileView()
if peerCmd != nil {
loadCmds = append(loadCmds, peerCmd)
}
}
if len(loadCmds) > 0 {
return m, batchCommands(loadCmds...)
}
}
return m, nil
case swarmSessionsMsg:
@@ -3051,6 +3157,13 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
}
}
return m, nil
case unknownScopeMutationMsg:
// A suppressed mutating result (TaskOutput) completed. Its success row is
// hidden, so trigger the conservative refresh on the owning surface.
if msg.runID != m.activeRunID {
return m, nil
}
return m.invalidateFileViewForUnknownMutation()
case doctorCommandResultMsg:
if msg.id == 0 || msg.id == m.doctorCommandSeq {
m.doctorInFlight = false
@@ -3068,7 +3181,13 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
m.prState = msg.state
return m, nil
case gitSweepMsg:
return m.handleGitSweepMsg(msg), nil
m = m.handleGitSweepMsg(msg)
// The swept list is cumulative since the baseline, so it says nothing
// about what changed since the open snapshot, and each live sweep starts
// beside tool-result, turn-end, or BTW-return handling that already
// reloads or rechecks the open view. Revalidate the open view's content
// rather than force a second reload.
return m.revalidateFileView()
case prWatcherStartedMsg:
if msg.stop == nil {
return m, nil
@@ -3080,7 +3199,8 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
return m, nil
case bashResultMsg:
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: msg.output})
return m, nil
// invalidateFileViewForUnknownMutation schedules startFileViewRefreshCmd when fileView is active
return m.invalidateFileViewForUnknownMutation()
case providerModelsDiscoveredMsg:
return m.applyProviderModelsDiscovered(msg), nil
case setupModelsDiscoveredMsg:
@@ -3716,6 +3836,16 @@ func (m model) scrollChat(delta int) model {
if !m.altScreen || delta == 0 {
return m
}
// A pending full-file reload renders a one-line placeholder whose geometry
// cannot represent the reader's position. Apply the delta to the pending
// intent against the last completed body instead, so a Page Down, arrow, or
// wheel during the reload is not discarded by the placeholder's zero height.
if m.fileView.active && m.fileView.mode == fileViewFull && m.fileView.loading {
maxOffset := maxInt(0, m.fileView.completedBodyLines-m.chatViewportHeight())
m.fileView.preservedScrollOffset = clampInt(m.fileView.preservedScrollOffset+delta, 0, maxOffset)
m.chatScrollOffset = m.fileView.preservedScrollOffset
return m
}
viewport, ok := m.chatTranscriptViewport()
if !ok {
return m
@@ -3727,6 +3857,28 @@ func (m model) scrollChat(delta int) model {
return m
}
// chatViewportHeight returns the height of the transcript body viewport, or 0
// when it cannot be measured. It lets pending-scroll math run against the last
// completed body even while the visible body is the loading placeholder.
func (m model) chatViewportHeight() int {
viewport, ok := m.chatTranscriptViewport()
if !ok {
return 0
}
return viewport.height
}
// resetChatScrollToBottom snaps the viewport to the tail for a real submission
// and, if a full-file reload is pending, records the same intent so completion
// cannot restore the reader's pre-submission position.
func (m model) resetChatScrollToBottom() model {
m.chatScrollOffset = 0
if m.fileView.loading {
m.fileView.preservedScrollOffset = 0
}
return m
}
func (m model) chatScrollMetrics() (int, int) {
viewport, ok := m.chatTranscriptViewport()
if !ok {
@@ -3761,6 +3913,16 @@ func (m model) chatTranscriptViewport() (transcriptViewport, bool) {
// so the absolute view holds; at the bottom (offset 0) it follows normally. Only the
// scrolled-up path renders the body, so the common case stays cheap.
func (m model) syncChatScroll() model {
// A pending full-file reload temporarily renders a one-line loading
// placeholder. Measuring it would clamp the reader's offset to 0 and lose
// their place, so hold the pending intent — which scroll actions and prompt
// submission keep up to date — until handleFileViewLoaded reconciles it
// against the real body.
if m.altScreen && m.fileView.active && m.fileView.mode == fileViewFull &&
m.fileView.loading && m.fileView.preservedScrollOffset > 0 {
m.chatScrollOffset = m.fileView.preservedScrollOffset
return m
}
if !m.altScreen || m.chatScrollOffset <= 0 {
// At the bottom (or inline mode): follow the tail; reset the pin baseline.
m.chatBodyLines = 0
@@ -4627,6 +4789,9 @@ func (m model) choosePicker() (tea.Model, tea.Cmd) {
if text != "" {
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: text})
}
if text == "" {
m, cmd = m.refreshFileViewMarkers()
}
case pickerSkill:
// Fill the composer with "/name " so the user adds their request before
// submitting (a bare second Enter runs it without one); names the slash
@@ -4652,10 +4817,21 @@ func (m model) choosePicker() (tea.Model, tea.Cmd) {
// local preview and never changes the active palette.
text := ""
m, text = m.handleThemeCommand(item.Value)
var loadCmd tea.Cmd
if m.fileView.active && m.fileView.mode == fileViewFull {
m, loadCmd = m.startFileViewLoadCmd(m.chatColumnWidth())
}
if validThemeMode(item.Value) && !strings.Contains(text, "could not save theme preference") {
return m.showTransientNotice(m.themeAppliedNotice(), transientNoticeSuccess)
next, noticeCmd := m.showTransientNotice(m.themeAppliedNotice(), transientNoticeSuccess)
if loadCmd != nil {
return next, tea.Batch(noticeCmd, loadCmd)
}
return next, noticeCmd
}
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: text})
if loadCmd != nil {
return m, loadCmd
}
case pickerNotify:
// The picker item's Value is "<mode> <focus>"; reusing the text handler
// keeps validation, persistence, and the user-facing message in one
@@ -4731,9 +4907,10 @@ func (m model) handleSubmit() (tea.Model, tea.Cmd) {
m.clearSuggestions()
// Snap the viewport back to the bottom for a real submission, but not for an
// empty Enter (a no-op) — that would yank the user away from wherever they
// had scrolled without anything actually being submitted.
// had scrolled without anything actually being submitted. A pending file
// reload records the same intent so its completion cannot undo the reset.
if command.kind != commandEmpty {
m.chatScrollOffset = 0
m = m.resetChatScrollToBottom()
}
return m.dispatchCommand(command)
@@ -4792,7 +4969,9 @@ func (m model) dispatchCommand(command parsedCommand) (tea.Model, tea.Cmd) {
if m.loopActive() {
m = m.appendLoopSystem(m.loopFooterSummary() + " still running — /loop stop all to end them.")
}
return m, nil
var clearCmd tea.Cmd
m, clearCmd = m.refreshFileViewMarkers()
return m, clearCmd
case commandNew:
// A fresh session mid-run would strand the in-flight turn's events; make the
// user cancel first. Idle, /new saves the current session (already on disk)
@@ -4801,7 +4980,10 @@ func (m model) dispatchCommand(command parsedCommand) (tea.Model, tea.Cmd) {
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: "A run is in progress. Press Esc to cancel it first, then /new."})
return m, nil
}
return m.startNewSession(), nil
next := m.startNewSession()
var newCmd tea.Cmd
next, newCmd = next.refreshFileViewMarkers()
return next, newCmd
case commandBTW:
return m.handleBTWCommand(command.text)
case commandLoop:
@@ -4963,7 +5145,11 @@ func (m model) dispatchCommand(command parsedCommand) (tea.Model, tea.Cmd) {
} else if text != "" {
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: text})
}
return m, nil
var resumeCmd tea.Cmd
if text == "" {
m, resumeCmd = m.refreshFileViewMarkers()
}
return m, resumeCmd
case commandRename:
if title := strings.TrimSpace(command.text); title != "" {
return m.renameActiveSession(title), nil
@@ -4983,9 +5169,10 @@ func (m model) dispatchCommand(command parsedCommand) (tea.Model, tea.Cmd) {
return m.toggleDetailedTranscript(), nil
case commandRewind:
text := ""
m, text = m.handleRewindCommand(command.text)
var rewindCmd tea.Cmd
m, text, rewindCmd = m.handleRewindCommand(command.text)
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: text})
return m, nil
return m, rewindCmd
case commandEffort:
if strings.TrimSpace(command.text) == "" {
if m.pending {
@@ -5073,10 +5260,21 @@ func (m model) dispatchCommand(command parsedCommand) (tea.Model, tea.Cmd) {
}
text := ""
m, text = m.handleThemeCommand(command.text)
var loadCmd tea.Cmd
if m.fileView.active && m.fileView.mode == fileViewFull {
m, loadCmd = m.startFileViewLoadCmd(m.chatColumnWidth())
}
if validThemeMode(command.text) && !strings.Contains(text, "could not save theme preference") {
return m.showTransientNotice(m.themeAppliedNotice(), transientNoticeSuccess)
next, noticeCmd := m.showTransientNotice(m.themeAppliedNotice(), transientNoticeSuccess)
if loadCmd != nil {
return next, tea.Batch(noticeCmd, loadCmd)
}
return next, noticeCmd
}
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: text})
if loadCmd != nil {
return m, loadCmd
}
return m, nil
case commandNotify:
// Bare `/notify` opens the popup picker so the user can pick mode + focus
@@ -5213,7 +5411,7 @@ func (m model) executeSlash(input string) (tea.Model, tea.Cmd) {
}
m.rememberInput(input)
m.clearSuggestions()
m.chatScrollOffset = 0
m = m.resetChatScrollToBottom()
return m.dispatchCommand(command)
}
@@ -5965,11 +6163,7 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
// them is flushed at end-of-run AND on cancel (flushRunIDs), so the blobs
// never stay orphaned — see its contract in internal/sessions.
if m.sessionStore != nil && m.activeSession.SessionID != "" {
var args map[string]any
if call.Arguments != "" {
_ = json.Unmarshal([]byte(call.Arguments), &args)
}
if targets := tools.MutationTargets(m.cwd, call.Name, args); len(targets) > 0 {
if targets := agent.CheckpointTargets(m.cwd, options.Registry, call); len(targets) > 0 {
if payload, ok := m.sessionStore.SnapshotForCheckpoint(m.activeSession.SessionID, m.cwd, call.Name, targets); ok {
sessionEvents = append(sessionEvents, pendingSessionEvent{
Type: sessions.EventSessionCheckpoint,
@@ -6019,6 +6213,11 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
if !toolResultCardSuppressedInTranscript(result.Name, result.Status) {
rows = append(rows, row)
m.sendAgentRow(runID, row)
} else if result.Name != "Task" && toolResultMayMutateUnknownScope(result.Name) && m.runtimeMessageSink != nil {
// A hidden successful mutating result (TaskOutput) still needs the
// conservative refresh; Task completion carries it via
// specialistCompleteMsg.
m.runtimeMessageSink(unknownScopeMutationMsg{runID: runID})
}
// Keep the latest plan state in sync for run details and step drill-in.
if result.Name == "update_plan" && m.registry != nil {
+5
View File
@@ -154,6 +154,11 @@ func (m model) handleMouse(msg tea.MouseMsg) (tea.Model, tea.Cmd) {
if next, cmd, handled := m.handlePetMouse(msg); handled {
return next, cmd
}
if mouseLeftPress(msg) && m.runDetailsVisible() {
if path, ok := m.runDetailsFileAtMouse(msg); ok {
return m.selectFile(path)
}
}
if mouseLeftPress(msg) {
switch {
case m.providerWizard != nil:
+3 -3
View File
@@ -258,7 +258,7 @@ func TestProviderWizardChatGPTOAuthModelsUseCodexSet(t *testing.T) {
for _, model := range models {
got[model.ID] = true
}
for _, want := range []string{"gpt-5.5", "gpt-5.4", "gpt-5.4-mini", "gpt-5.3-codex-spark"} {
for _, want := range []string{"gpt-6.1-sol", "gpt-6.1-terra", "gpt-6-astra", "gpt-6-luna"} {
if !got[want] {
t.Fatalf("ChatGPT OAuth models missing %q; got %#v", want, providerWizardModelIDs(models))
}
@@ -266,8 +266,8 @@ func TestProviderWizardChatGPTOAuthModelsUseCodexSet(t *testing.T) {
if got["gpt-5"] {
t.Fatalf("ChatGPT OAuth models should not include stale gpt-5 fallback; got %#v", providerWizardModelIDs(models))
}
if models[0].ID != "gpt-5.5" {
t.Fatalf("default ChatGPT OAuth model = %q, want gpt-5.5", models[0].ID)
if models[0].ID != "gpt-6.1-sol" {
t.Fatalf("default ChatGPT OAuth model = %q, want gpt-6.1-sol", models[0].ID)
}
}
Loaded 100 of 110 files, more files were not shown because too many files have changed in this diff. Show more