mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-09 13:31:26 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1dfe275e50 | ||
|
|
3ad039e937 | ||
|
|
3123b8d36e | ||
|
|
a1d4ac3b9a | ||
|
|
0899b1132e | ||
|
|
433653b873 | ||
|
|
e09c43e20a | ||
|
|
066a3a2ff0 | ||
|
|
b868a92290 | ||
|
|
daf9b8c65d | ||
|
|
c022e6afe7 | ||
|
|
c1b4e27749 | ||
|
|
937631a4ea | ||
|
|
9d8059cd49 | ||
|
|
7dfdbcb686 | ||
|
|
f1ff15eb94 | ||
|
|
eb3bbb2f12 | ||
|
|
9e09d0a7e6 | ||
|
|
f984082c82 | ||
|
|
33f62e7f94 | ||
|
|
8535e4a6b1 | ||
|
|
6bce7d3883 | ||
|
|
3b7c741480 | ||
|
|
b11a8a6231 | ||
|
|
105e955fa6 | ||
|
|
67de11f0ad | ||
|
|
2f8dc76666 | ||
|
|
886567fe19 | ||
|
|
b16971f5c3 | ||
|
|
03ee13e480 | ||
|
|
b69ff6733d | ||
|
|
4ccdab8f08 | ||
|
|
600bbc4c44 | ||
|
|
b8d344d48c | ||
|
|
2176c98271 | ||
|
|
0ae48e400d | ||
|
|
47d6ab7015 | ||
|
|
591874f1f8 | ||
|
|
b619016f08 | ||
|
|
60a1452eeb | ||
|
|
98fdb3e8e5 | ||
|
|
fc06cc6aaa | ||
|
|
b44d56f2b2 | ||
|
|
6bf5a86b7f | ||
|
|
e2ac98b906 | ||
|
|
b60359e2d6 | ||
|
|
f889b7e3b3 | ||
|
|
e0e8db306e | ||
|
|
f62ae776c9 | ||
|
|
427df5fdb9 | ||
|
|
22c9194bce | ||
|
|
85526b4978 | ||
|
|
1ee7aa762a | ||
|
|
4f8aa57314 | ||
|
|
76e5e35a1c | ||
|
|
cd76c67050 | ||
|
|
081aabdbc4 | ||
|
|
0b646ad7f9 | ||
|
|
4e2231e5de | ||
|
|
f7814819cb | ||
|
|
5c64cef7f6 | ||
|
|
55621d2176 | ||
|
|
915e1df103 | ||
|
|
3bfd898bfc | ||
|
|
c4905d296b | ||
|
|
93f9505b72 | ||
|
|
6a4eff21a2 |
No files matched your search
@@ -1 +1 @@
|
||||
{".":"0.13.0"}
|
||||
{".":"0.11.0"}
|
||||
@@ -6,52 +6,6 @@ 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.13.0](https://github.com/Twigpine/zero/compare/v0.12.0...v0.13.0) (2026-11-09)
|
||||
|
||||
|
||||
### Features
|
||||
|
||||
* **mcp:** add AI/ML API as an opt-in built-in server ([#1193](https://github.com/Twigpine/zero/issues/1193)) ([16ccd78](https://github.com/Twigpine/zero/commit/16ccd78765b6a2a56b7370e3031710d718aa4757))
|
||||
* **plugins:** add CLI enable/disable for plugin.json ([#1202](https://github.com/Twigpine/zero/issues/1202)) ([05c1b5e](https://github.com/Twigpine/zero/commit/05c1b5e16098e1c60212316d20ae8e055a27b5ae))
|
||||
* **providers:** point OpenGateway at api.twigpine.com and trust its live model list ([#1208](https://github.com/Twigpine/zero/issues/1208)) ([47f78c6](https://github.com/Twigpine/zero/commit/47f78c66247aa96c8753d6bb95e935126ff1761e))
|
||||
* **web-fetch:** add an opt-in AI/ML API Toolbox scraping backend ([#1192](https://github.com/Twigpine/zero/issues/1192)) ([c9f1794](https://github.com/Twigpine/zero/commit/c9f179482a7f6d7111ee64fa9cb8ffa2517c2397))
|
||||
* **web-search:** add AI/ML API Toolbox as an opt-in backend ([#1189](https://github.com/Twigpine/zero/issues/1189)) ([9873645](https://github.com/Twigpine/zero/commit/98736457bdac6b0272e032728a5485cf557780fe))
|
||||
|
||||
|
||||
### Bug Fixes
|
||||
|
||||
* **hooks:** dispatch specialistStart/Stop lifecycle hooks ([#1204](https://github.com/Twigpine/zero/issues/1204)) ([2ca84cc](https://github.com/Twigpine/zero/commit/2ca84cc127d9369765004b72cf95049fecc6c28b))
|
||||
* **tui:** make reasoning stream rendering linear (coalesced deltas, O(1) append, tail-bounded preview) ([#1198](https://github.com/Twigpine/zero/issues/1198)) ([bbf779e](https://github.com/Twigpine/zero/commit/bbf779e69168d62bb1103aefdd1b10e53433209e))
|
||||
|
||||
|
||||
### Performance Improvements
|
||||
|
||||
* **tui:** bounded file read, LRU rendering cache, and size caps for full-file view ([#1203](https://github.com/Twigpine/zero/issues/1203)) ([e034a93](https://github.com/Twigpine/zero/commit/e034a93815629c68976dcdab7dcc08bbe9b03693))
|
||||
|
||||
## [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)
|
||||
|
||||
|
||||
|
||||
@@ -15,7 +15,6 @@
|
||||
<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>
|
||||
@@ -88,18 +87,9 @@ 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>
|
||||
|
||||
@@ -397,76 +387,6 @@ 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
|
||||
|
||||
@@ -15,7 +15,6 @@
|
||||
<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>
|
||||
@@ -86,18 +85,9 @@ 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>
|
||||
|
||||
|
||||
+1
-13
@@ -33,16 +33,4 @@ 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`). 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).
|
||||
keyless web search, to Exa's hosted MCP endpoint at `https://mcp.exa.ai/mcp`).
|
||||
+2
-69
@@ -110,8 +110,6 @@ 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/
|
||||
@@ -142,15 +140,6 @@ 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:
|
||||
@@ -203,11 +192,7 @@ Events the agent emits (in dispatch order):
|
||||
| `specialistStart` | A sub-agent is spawned | yes (specialist name) |
|
||||
| `specialistStop` | A sub-agent ends | yes (specialist name) |
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
> **Roadmap.** An in-UI hooks manager is on the backlog. Today you edit the JSON directly.
|
||||
|
||||
@@ -261,56 +246,6 @@ 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.
|
||||
|
||||
If the server echoes back the key Zero sent it, Zero replaces the key with
|
||||
`[REDACTED]` in that server's error messages and tool results before anything
|
||||
reads them, and reports an HTTP error status without the response body.
|
||||
|
||||
### As a server — expose Zero's tools to another agent
|
||||
|
||||
```bash
|
||||
@@ -355,12 +290,10 @@ 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 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.
|
||||
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.
|
||||
|
||||
Plugin commands run with the plugin directory as their working directory. Use relative paths; the loader resolves them at activation time.
|
||||
|
||||
|
||||
@@ -634,10 +634,6 @@ 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 +0,0 @@
|
||||
<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>
|
||||
|
Before Width: | Height: | Size: 1.6 KiB |
@@ -1 +0,0 @@
|
||||
<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>
|
||||
|
Before Width: | Height: | Size: 1.6 KiB |
+45
-20
@@ -962,16 +962,25 @@ func toolCallEvent(call agent.ToolCall) sessions.AppendEventInput {
|
||||
}
|
||||
}
|
||||
|
||||
// toolResultEvent persists a tool result through the ONE session contract every
|
||||
// writer shares, agent.ToolResultSessionPayload.
|
||||
//
|
||||
// This used to spell its own payload and wrote result.Output. After the
|
||||
// output/notice split that field is the UNDECORATED text, and nothing else in
|
||||
// the payload carried the enforcement notices, so the disclosure was dropped at
|
||||
// the moment of persisting. The live tool_call_update for the same result had
|
||||
// shown it (toolResultContent reads ModelOutput), and then session/load
|
||||
// replayed the call with the sandbox's narrowing gone. Fixing only the live
|
||||
// translation could not reach this: the text was already lost on disk.
|
||||
//
|
||||
// The session store is shared with the TUI and headless exec, so the same
|
||||
// payload is also what lets a TUI resume of an ACP-written session draw the
|
||||
// notice as card furniture, from the typed field and the undecorated body.
|
||||
// Reported by @jatmn.
|
||||
func toolResultEvent(result agent.ToolResult) sessions.AppendEventInput {
|
||||
return sessions.AppendEventInput{
|
||||
Type: sessions.EventToolResult,
|
||||
Payload: map[string]any{
|
||||
"toolCallId": result.ToolCallID,
|
||||
"name": result.Name,
|
||||
"status": result.Status,
|
||||
"output": result.Output,
|
||||
"changedFiles": append([]string(nil), result.ChangedFiles...),
|
||||
},
|
||||
Type: sessions.EventToolResult,
|
||||
Payload: agent.ToolResultSessionPayload(result),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1144,13 +1153,14 @@ func replayToolUpdate(event sessions.Event) *ToolCallUpdate {
|
||||
return nil
|
||||
}
|
||||
var payload struct {
|
||||
Name string `json:"name"`
|
||||
ToolCallID string `json:"toolCallId"`
|
||||
ID string `json:"id"`
|
||||
Arguments string `json:"arguments"`
|
||||
Status string `json:"status"`
|
||||
Output string `json:"output"`
|
||||
ChangedFiles []string `json:"changedFiles"`
|
||||
Name string `json:"name"`
|
||||
ToolCallID string `json:"toolCallId"`
|
||||
ID string `json:"id"`
|
||||
Arguments string `json:"arguments"`
|
||||
Status string `json:"status"`
|
||||
Output string `json:"output"`
|
||||
ChangedFiles []string `json:"changedFiles"`
|
||||
EnforcementNotices []string `json:"enforcementNotices"`
|
||||
}
|
||||
if json.Unmarshal(raw, &payload) != nil {
|
||||
return nil
|
||||
@@ -1170,12 +1180,27 @@ func replayToolUpdate(event sessions.Event) *ToolCallUpdate {
|
||||
if status == "" {
|
||||
status = tools.StatusOK
|
||||
}
|
||||
// THE DISCLOSURE IS RESTORED AS WHAT IT WAS: typed notices beside undecorated
|
||||
// text, so toolCallResult composes it the way it did live and the client
|
||||
// sees it exactly once. The stored output already has the notices composed
|
||||
// in (that is the shared contract, because output is the provider-facing
|
||||
// text), so handing it to ModelOutput together with the typed field would
|
||||
// draw the disclosure twice, and dropping the typed field would make this
|
||||
// reader depend on every writer having decorated the text. Taking the
|
||||
// notices back off the front covers both: a record whose output was never
|
||||
// decorated comes through unchanged and still gets its notice.
|
||||
//
|
||||
// displayPreview is deliberately not read. It is the card body for a
|
||||
// terminal, and the ACP wire has no card: the live update sends the model
|
||||
// text, so a replay that sent the preview instead would not match it.
|
||||
notices := append([]string(nil), payload.EnforcementNotices...)
|
||||
upd := toolCallResult(agent.ToolResult{
|
||||
ToolCallID: id,
|
||||
Name: payload.Name,
|
||||
Status: status,
|
||||
Output: payload.Output,
|
||||
ChangedFiles: append([]string(nil), payload.ChangedFiles...),
|
||||
ToolCallID: id,
|
||||
Name: payload.Name,
|
||||
Status: status,
|
||||
Output: tools.WithoutEnforcementNotices(payload.Output, notices),
|
||||
EnforcementNotices: notices,
|
||||
ChangedFiles: append([]string(nil), payload.ChangedFiles...),
|
||||
})
|
||||
return &upd
|
||||
}
|
||||
|
||||
@@ -0,0 +1,195 @@
|
||||
package acp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Twigpine/zero/internal/agent"
|
||||
"github.com/Twigpine/zero/internal/sessions"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
"github.com/Twigpine/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
const persistedNotice = "least-privilege notice: read access was narrowed"
|
||||
|
||||
func mustRawPayload(t *testing.T, payload map[string]any) json.RawMessage {
|
||||
t.Helper()
|
||||
encoded, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return encoded
|
||||
}
|
||||
|
||||
func toolUpdateText(update ToolCallUpdate) string {
|
||||
var text strings.Builder
|
||||
for _, part := range update.Content {
|
||||
if part.Content != nil {
|
||||
text.WriteString(part.Content.Text)
|
||||
}
|
||||
}
|
||||
return text.String()
|
||||
}
|
||||
|
||||
// nextToolResultUpdate returns the next tool_call_update, skipping the
|
||||
// tool_call that announces it.
|
||||
func nextToolResultUpdate(t *testing.T, ctx context.Context, updates <-chan ToolCallUpdate, stage string) ToolCallUpdate {
|
||||
t.Helper()
|
||||
for {
|
||||
select {
|
||||
case update := <-updates:
|
||||
if update.SessionUpdate == UpdateToolCallUpdate {
|
||||
return update
|
||||
}
|
||||
case <-ctx.Done():
|
||||
t.Fatalf("%s: the tool result update never arrived", stage)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A DISCLOSURE SHOWN LIVE HAS TO SURVIVE BEING WRITTEN DOWN.
|
||||
//
|
||||
// enforcement_notice_test.go pins the live translation: toolResultContent reads
|
||||
// ModelOutput, so a connected ACP client sees the sandbox's narrowing. That was
|
||||
// the only half fixed. The turn also persists each result, and the persisted
|
||||
// payload was spelled by hand in this package with result.Output, which after
|
||||
// the output/notice split is the undecorated text, and with no notices field at
|
||||
// all. The disclosure was gone from disk, so session/load replayed a sandboxed
|
||||
// command as if nothing had constrained it.
|
||||
//
|
||||
// This runs a real turn, so the result reaches the store through the same
|
||||
// OnToolResult callback production uses, then loads the session in a fresh
|
||||
// agent and reads the replayed update. It fails if only translate.go is fixed,
|
||||
// which is the state it was written against. Reported by @jatmn.
|
||||
func TestACPReloadedToolResultCarriesTheEnforcementNoticeOnce(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
output string
|
||||
}{
|
||||
{name: "a command with output", output: "the command output"},
|
||||
// An enforced command that printed nothing still has a disclosure, and
|
||||
// an empty body is where a reader that falls back to the decorated text
|
||||
// draws it twice.
|
||||
{name: "a command that printed nothing", output: ""},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
deps := testDeps(t)
|
||||
deps.RunAgent = func(_ context.Context, _ string, _ zeroruntime.Provider, opts agent.Options) (agent.Result, error) {
|
||||
call := agent.ToolCall{ID: "call-sandboxed", Name: "bash", Arguments: `{"command":"ls"}`}
|
||||
opts.OnToolCall(call)
|
||||
opts.OnToolResult(agent.ToolResult{
|
||||
ToolCallID: call.ID,
|
||||
Name: call.Name,
|
||||
Status: tools.StatusOK,
|
||||
Output: tc.output,
|
||||
EnforcementNotices: []string{persistedNotice},
|
||||
})
|
||||
return agent.Result{FinalAnswer: "done"}, nil
|
||||
}
|
||||
workspace := t.TempDir()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
live := newHarness(t, deps)
|
||||
var created NewSessionResult
|
||||
if err := live.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: workspace, McpServers: []McpServer{}}, &created); err != nil {
|
||||
t.Fatalf("session/new: %v", err)
|
||||
}
|
||||
if err := live.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: created.SessionID, Prompt: []ContentBlock{TextBlock("list the files")}}, &PromptResult{}); err != nil {
|
||||
t.Fatalf("session/prompt: %v", err)
|
||||
}
|
||||
liveText := toolUpdateText(nextToolResultUpdate(t, ctx, live.tools, "live"))
|
||||
if count := strings.Count(liveText, persistedNotice); count != 1 {
|
||||
t.Fatalf("premise: the live update carries the notice %d times, want 1:\n%s", count, liveText)
|
||||
}
|
||||
live.stop()
|
||||
|
||||
loader := newHarness(t, deps)
|
||||
defer loader.stop()
|
||||
if err := loader.client.Call(ctx, MethodSessionLoad, LoadSessionParams{SessionID: created.SessionID, Cwd: workspace, McpServers: []McpServer{}}, &LoadSessionResult{}); err != nil {
|
||||
t.Fatalf("session/load: %v", err)
|
||||
}
|
||||
replayed := toolUpdateText(nextToolResultUpdate(t, ctx, loader.tools, "replay"))
|
||||
if count := strings.Count(replayed, persistedNotice); count != 1 {
|
||||
t.Errorf("the replayed result carries the notice %d times, want exactly 1:\n%s", count, replayed)
|
||||
}
|
||||
if replayed != liveText {
|
||||
t.Errorf("the replayed result does not match what the client saw live\nlive: %q\nreplayed: %q", liveText, replayed)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// THE ACP WRITER IS THE SHARED CONTRACT, NOT A COPY OF IT. The session store is
|
||||
// the one the TUI and headless exec also write to and resume from, so a payload
|
||||
// spelled here by hand is a third representation waiting to drift. It already
|
||||
// had: it was the only writer with no notices and no undecorated body.
|
||||
func TestToolResultEventIsTheSharedSessionPayload(t *testing.T) {
|
||||
result := agent.ToolResult{
|
||||
ToolCallID: "call-1",
|
||||
Name: "bash",
|
||||
Status: tools.StatusOK,
|
||||
Output: "the command output",
|
||||
EnforcementNotices: []string{persistedNotice},
|
||||
ChangedFiles: []string{"a.go"},
|
||||
}
|
||||
event := toolResultEvent(result)
|
||||
if event.Type != sessions.EventToolResult {
|
||||
t.Fatalf("event type = %q", event.Type)
|
||||
}
|
||||
if want := agent.ToolResultSessionPayload(result); !reflect.DeepEqual(event.Payload, want) {
|
||||
t.Fatalf("the ACP payload is not the shared one\n got: %#v\nwant: %#v", event.Payload, want)
|
||||
}
|
||||
payload, ok := event.Payload.(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("payload is %T", event.Payload)
|
||||
}
|
||||
if notices, _ := payload["enforcementNotices"].([]string); len(notices) != 1 || notices[0] != persistedNotice {
|
||||
t.Errorf("typed notices were not persisted: %#v", payload["enforcementNotices"])
|
||||
}
|
||||
if body, present := payload["displayPreview"]; !present || body != "the command output" {
|
||||
t.Errorf("the undecorated body was not persisted: %#v (present=%v)", body, present)
|
||||
}
|
||||
}
|
||||
|
||||
// A record can reach the reader in either shape, and both have to come out as
|
||||
// one disclosure: output with the notices already composed in, which is what
|
||||
// every current writer stores, and output left undecorated beside the typed
|
||||
// field, which a reader must not answer by dropping the notice.
|
||||
func TestReplayedToolResultComposesTheNoticeOnceForEitherStoredShape(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
output string
|
||||
want string
|
||||
}{
|
||||
{name: "output stored decorated", output: persistedNotice + "\n\nthe command output", want: persistedNotice + "\n\nthe command output"},
|
||||
{name: "output stored undecorated", output: "the command output", want: persistedNotice + "\n\nthe command output"},
|
||||
{name: "decorated and otherwise empty", output: persistedNotice, want: persistedNotice},
|
||||
{name: "undecorated and empty", output: "", want: persistedNotice},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
update := replayToolUpdate(sessions.Event{Type: sessions.EventToolResult, Payload: mustRawPayload(t, map[string]any{
|
||||
"toolCallId": "call-1", "name": "bash", "status": "ok",
|
||||
"output": tc.output, "enforcementNotices": []string{persistedNotice},
|
||||
})})
|
||||
if update == nil {
|
||||
t.Fatal("the stored result was not replayed")
|
||||
}
|
||||
if got := toolUpdateText(*update); got != tc.want {
|
||||
t.Errorf("replayed text = %q, want %q", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// And a stored result with no notices is replayed as it was written.
|
||||
plain := replayToolUpdate(sessions.Event{Type: sessions.EventToolResult, Payload: mustRawPayload(t, map[string]any{
|
||||
"toolCallId": "call-2", "name": "bash", "status": "ok", "output": "plain output",
|
||||
})})
|
||||
if plain == nil || toolUpdateText(*plain) != "plain output" {
|
||||
t.Errorf("an ordinary stored result changed on replay: %+v", plain)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package acp
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/agent"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// AN ACP CLIENT MUST SEE THE DISCLOSURE THE TUI SEES.
|
||||
//
|
||||
// agent.ToolResult stores the UNDECORATED model text alongside the typed
|
||||
// enforcement notices; ModelOutput is what composes them. Reading .Output
|
||||
// directly compiles and looks right, and silently drops the notice for every
|
||||
// ACP client, which is the one surface with no other way to learn the sandbox
|
||||
// narrowed what the command could do.
|
||||
func TestToolResultContentCarriesTheEnforcementNotice(t *testing.T) {
|
||||
const notice = "least-privilege notice: read access was narrowed"
|
||||
result := agent.ToolResult{
|
||||
Name: "bash",
|
||||
Status: tools.StatusOK,
|
||||
Output: "the command output",
|
||||
EnforcementNotices: []string{notice},
|
||||
}
|
||||
|
||||
content := toolResultContent(result)
|
||||
if len(content) == 0 {
|
||||
t.Fatal("no content produced for a successful tool result")
|
||||
}
|
||||
var text strings.Builder
|
||||
for _, part := range content {
|
||||
if part.Content != nil {
|
||||
text.WriteString(part.Content.Text)
|
||||
}
|
||||
}
|
||||
got := text.String()
|
||||
|
||||
if count := strings.Count(got, notice); count != 1 {
|
||||
t.Errorf("the notice appears %d times, want exactly 1:\n%s", count, got)
|
||||
}
|
||||
if !strings.Contains(got, "the command output") {
|
||||
t.Errorf("the underlying output was lost:\n%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
// And a result with no notice is unchanged, so the accessor is not adding
|
||||
// anything to ordinary output.
|
||||
func TestToolResultContentLeavesAnOrdinaryResultAlone(t *testing.T) {
|
||||
result := agent.ToolResult{
|
||||
Name: "bash",
|
||||
Status: tools.StatusOK,
|
||||
Output: "plain output",
|
||||
}
|
||||
content := toolResultContent(result)
|
||||
if len(content) == 0 {
|
||||
t.Fatal("no content produced")
|
||||
}
|
||||
if content[0].Content == nil {
|
||||
t.Fatal("content block missing")
|
||||
}
|
||||
if got := content[0].Content.Text; got != "plain output" {
|
||||
t.Errorf("ordinary output = %q, want it untouched", got)
|
||||
}
|
||||
}
|
||||
@@ -248,7 +248,11 @@ func toolCallResult(result agent.ToolResult) ToolCallUpdate {
|
||||
|
||||
func toolResultContent(result agent.ToolResult) []ToolCallContent {
|
||||
content := make([]ToolCallContent, 0, 1+len(result.FileDiffs))
|
||||
text := strings.TrimRight(result.Output, "\n")
|
||||
// ModelOutput, not the raw field. agent.ToolResult stores the undecorated
|
||||
// model text alongside the typed enforcement notices, and the accessor is
|
||||
// what composes the two; reading Output directly sends an ACP client the
|
||||
// output with the disclosure missing.
|
||||
text := strings.TrimRight(result.ModelOutput(), "\n")
|
||||
if text == "" {
|
||||
text = result.Display.Summary
|
||||
}
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/execution"
|
||||
"github.com/Twigpine/zero/internal/hooks"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// The tool's own disclosure and the afterTool hook's are the SAME STRING, which
|
||||
// is the real case: both come from the fixed Windows deny_read warning, so a hook
|
||||
// running under the same token shape as the tool it follows reports exactly what
|
||||
// the tool reported.
|
||||
const sharedEnforcementNotice = "least-privilege notice"
|
||||
|
||||
const afterToolChatter = "vet-found-nothing"
|
||||
|
||||
// sharedNoticeHookPreparer runs a hook that prints ordinary output and carries
|
||||
// the same enforcement notice the tool carries.
|
||||
type sharedNoticeHookPreparer struct{}
|
||||
|
||||
func (sharedNoticeHookPreparer) PrepareExecution(_ context.Context, _ execution.Request) (execution.PreparedCommand, error) {
|
||||
var command *exec.Cmd
|
||||
if runtime.GOOS == "windows" {
|
||||
command = exec.Command("cmd.exe", "/c", "echo "+afterToolChatter)
|
||||
} else {
|
||||
command = exec.Command("/bin/sh", "-c", "echo "+afterToolChatter)
|
||||
}
|
||||
return execution.PreparedCommand{
|
||||
Command: command,
|
||||
Enforcement: execution.Enforcement{Notices: []string{sharedEnforcementNotice}},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func afterToolNoticeDispatcher(t *testing.T) *hooks.Dispatcher {
|
||||
t.Helper()
|
||||
audit, err := hooks.NewAuditStore(hooks.AuditStoreOptions{AuditPath: filepath.Join(t.TempDir(), "audit.jsonl")})
|
||||
if err != nil {
|
||||
t.Fatalf("NewAuditStore: %v", err)
|
||||
}
|
||||
return hooks.NewDispatcher(hooks.DispatcherOptions{
|
||||
Config: hooks.Config{
|
||||
Enabled: true,
|
||||
Hooks: []hooks.Definition{
|
||||
{ID: "zero.after-tool", Event: hooks.EventAfterTool, Matcher: "notice_projection", Command: "vet", Enabled: true},
|
||||
},
|
||||
},
|
||||
Audit: audit,
|
||||
Cwd: t.TempDir(),
|
||||
Execution: execution.NewRunner(sharedNoticeHookPreparer{}),
|
||||
})
|
||||
}
|
||||
|
||||
// ONE DISCLOSURE, ONE DELIVERY, FROM EITHER HOOK PHASE.
|
||||
//
|
||||
// beforeTool was moved onto the typed EnforcementNotices slice and afterTool was
|
||||
// left folding its notices into the prose feedback. Both halves then wrote the
|
||||
// same fact: the typed slice, which every surface composes through ModelOutput
|
||||
// and HumanDisplay, and the "Hook output:" block appended to the body. Since the
|
||||
// two carry the identical fixed string, the model saw the disclosure twice and a
|
||||
// bash or exec card showed it in the amber furniture and again in the body.
|
||||
//
|
||||
// Half a symmetry is its own defect, and this is the composition that catches it:
|
||||
// a tool that reports a notice, followed by an afterTool hook that reports the
|
||||
// same one.
|
||||
func TestAnAfterToolNoticeIsDeliveredOnce(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(noticeProjectionTool{})
|
||||
|
||||
result, err := executeToolCall(context.Background(), registry, ToolCall{
|
||||
ID: "call-1", Name: "notice_projection", Arguments: `{}`,
|
||||
}, PermissionModeAuto, Options{Cwd: t.TempDir(), Hooks: afterToolNoticeDispatcher(t)})
|
||||
if err != nil {
|
||||
t.Fatalf("executeToolCall: %v", err)
|
||||
}
|
||||
|
||||
// SETUP: the hook really ran, or the silence below is the silence of a hook
|
||||
// that never executed.
|
||||
if !strings.Contains(result.ModelOutput(), afterToolChatter) {
|
||||
t.Fatalf("SETUP INVALID: the afterTool hook's own output never arrived, so nothing here is under test:\n%s", result.ModelOutput())
|
||||
}
|
||||
|
||||
if got := strings.Count(result.ModelOutput(), sharedEnforcementNotice); got != 1 {
|
||||
t.Errorf("the disclosure reached the model %d times, want once:\n%s", got, result.ModelOutput())
|
||||
}
|
||||
if got := strings.Count(strings.Join(result.EnforcementNotices, "\n"), sharedEnforcementNotice); got != 1 {
|
||||
t.Errorf("the typed slice carries the disclosure %d times, want once: %v", got, result.EnforcementNotices)
|
||||
}
|
||||
// The body must not carry it at all: the surfaces draw it from the slice, so a
|
||||
// copy in the body is what renders it twice.
|
||||
if strings.Contains(result.BaseModelOutput(), sharedEnforcementNotice) {
|
||||
t.Errorf("the disclosure is in the result body as well as the typed slice:\n%s", result.BaseModelOutput())
|
||||
}
|
||||
if strings.Count(result.HumanDisplay().Summary, sharedEnforcementNotice) != 1 {
|
||||
t.Errorf("the card summary shows the disclosure %d times, want once:\n%s",
|
||||
strings.Count(result.HumanDisplay().Summary, sharedEnforcementNotice), result.HumanDisplay().Summary)
|
||||
}
|
||||
}
|
||||
|
||||
// And an afterTool hook that discloses something the tool did not still gets
|
||||
// through, on the typed channel rather than as prose.
|
||||
func TestAnAfterToolNoticeReachesTheTypedSlice(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(quietTool{})
|
||||
|
||||
result, err := executeToolCall(context.Background(), registry, ToolCall{
|
||||
ID: "call-1", Name: "quiet_tool", Arguments: `{}`,
|
||||
}, PermissionModeAuto, Options{Cwd: t.TempDir(), Hooks: quietToolAfterHookDispatcher(t)})
|
||||
if err != nil {
|
||||
t.Fatalf("executeToolCall: %v", err)
|
||||
}
|
||||
if !strings.Contains(result.ModelOutput(), afterToolChatter) {
|
||||
t.Fatalf("SETUP INVALID: the afterTool hook never ran:\n%s", result.ModelOutput())
|
||||
}
|
||||
if len(result.EnforcementNotices) != 1 || result.EnforcementNotices[0] != sharedEnforcementNotice {
|
||||
t.Errorf("an afterTool hook's disclosure did not reach the typed slice, so no card renders it: %v", result.EnforcementNotices)
|
||||
}
|
||||
if strings.Contains(result.BaseModelOutput(), sharedEnforcementNotice) {
|
||||
t.Errorf("the disclosure travelled as prose in the body instead:\n%s", result.BaseModelOutput())
|
||||
}
|
||||
}
|
||||
|
||||
// quietTool carries no enforcement notice of its own.
|
||||
type quietTool struct{}
|
||||
|
||||
func (quietTool) Name() string { return "quiet_tool" }
|
||||
func (quietTool) Description() string { return "test tool with no enforcement notice" }
|
||||
func (quietTool) Parameters() tools.Schema { return tools.Schema{Type: "object"} }
|
||||
func (quietTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow}
|
||||
}
|
||||
|
||||
func (quietTool) Run(ctx context.Context, args map[string]any) tools.Result {
|
||||
return tools.Result{Status: tools.StatusOK, Output: "the command output"}
|
||||
}
|
||||
|
||||
func quietToolAfterHookDispatcher(t *testing.T) *hooks.Dispatcher {
|
||||
t.Helper()
|
||||
audit, err := hooks.NewAuditStore(hooks.AuditStoreOptions{AuditPath: filepath.Join(t.TempDir(), "audit.jsonl")})
|
||||
if err != nil {
|
||||
t.Fatalf("NewAuditStore: %v", err)
|
||||
}
|
||||
return hooks.NewDispatcher(hooks.DispatcherOptions{
|
||||
Config: hooks.Config{
|
||||
Enabled: true,
|
||||
Hooks: []hooks.Definition{
|
||||
{ID: "zero.after-tool", Event: hooks.EventAfterTool, Matcher: "quiet_tool", Command: "vet", Enabled: true},
|
||||
},
|
||||
},
|
||||
Audit: audit,
|
||||
Cwd: t.TempDir(),
|
||||
Execution: execution.NewRunner(sharedNoticeHookPreparer{}),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,210 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/execution"
|
||||
"github.com/Twigpine/zero/internal/hooks"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
zeroruntime "github.com/Twigpine/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
const beforeToolNotice = "denyRead is configured, so the write jail is not confining writes"
|
||||
|
||||
// beforeToolChatter is what a hook prints for its own reasons. It must never
|
||||
// reach the model: main is silent for a successful hook, and a hook that logs is
|
||||
// not asking to be heard by anything but the operator's terminal.
|
||||
const beforeToolChatter = "hook-ran-and-logged-this"
|
||||
|
||||
// noticeHookPreparer plans the hook command with an enforcement notice attached,
|
||||
// the way the sandbox does for a command it weakened. The prepared child prints
|
||||
// ordinary output as well, so one run carries both kinds of text and the
|
||||
// delivery decision has to tell them apart.
|
||||
type noticeHookPreparer struct{}
|
||||
|
||||
func (noticeHookPreparer) PrepareExecution(_ context.Context, _ execution.Request) (execution.PreparedCommand, error) {
|
||||
var command *exec.Cmd
|
||||
if runtime.GOOS == "windows" {
|
||||
command = exec.Command("cmd.exe", "/c", "echo "+beforeToolChatter)
|
||||
} else {
|
||||
command = exec.Command("/bin/sh", "-c", "echo "+beforeToolChatter)
|
||||
}
|
||||
return execution.PreparedCommand{
|
||||
Command: command,
|
||||
Enforcement: execution.Enforcement{Notices: []string{beforeToolNotice}},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func beforeToolDispatcher(t *testing.T, event hooks.Event, exitCode int) *hooks.Dispatcher {
|
||||
t.Helper()
|
||||
audit, err := hooks.NewAuditStore(hooks.AuditStoreOptions{AuditPath: filepath.Join(t.TempDir(), "audit.jsonl")})
|
||||
if err != nil {
|
||||
t.Fatalf("NewAuditStore: %v", err)
|
||||
}
|
||||
return hooks.NewDispatcher(hooks.DispatcherOptions{
|
||||
Config: hooks.Config{
|
||||
Enabled: true,
|
||||
Hooks: []hooks.Definition{
|
||||
{ID: "zero.before-tool", Event: event, Matcher: "read_file", Command: "hook", Enabled: true},
|
||||
},
|
||||
},
|
||||
Audit: audit,
|
||||
Cwd: t.TempDir(),
|
||||
Execution: execution.NewRunner(noticeHookPreparer{}),
|
||||
})
|
||||
}
|
||||
|
||||
func readFileRunOptions(t *testing.T, dispatcher *hooks.Dispatcher) (Options, *mockProvider, string) {
|
||||
t.Helper()
|
||||
root := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(root, "notes.txt"), []byte("hello"), 0o644); err != nil {
|
||||
t.Fatalf("write notes.txt: %v", err)
|
||||
}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.NewReadFileTool(root))
|
||||
provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{
|
||||
{
|
||||
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call-1", ToolName: "read_file"},
|
||||
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call-1", ArgumentsFragment: `{"path":"notes.txt"}`},
|
||||
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call-1"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
},
|
||||
{
|
||||
{Type: zeroruntime.StreamEventText, Content: "read it"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
},
|
||||
}}
|
||||
return Options{
|
||||
SessionID: "session-hook",
|
||||
Cwd: root,
|
||||
Registry: registry,
|
||||
ProviderName: "test-provider",
|
||||
Model: "test-model",
|
||||
Hooks: dispatcher,
|
||||
MaxTurns: 2,
|
||||
}, provider, root
|
||||
}
|
||||
|
||||
// countRequestsContaining reports how many provider requests carry needle, so a
|
||||
// notice delivered twice is distinguishable from one delivered once.
|
||||
func countRequestsContaining(requests []zeroruntime.CompletionRequest, needle string) int {
|
||||
total := 0
|
||||
for _, request := range requests {
|
||||
for _, message := range request.Messages {
|
||||
if strings.Contains(message.Content, needle) {
|
||||
total++
|
||||
}
|
||||
}
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
// THE NOTICE CROSSES TO THE MODEL. THE HOOK'S OWN OUTPUT DOES NOT.
|
||||
//
|
||||
// executeToolCall used to read the beforeTool outcome only when Blocked was
|
||||
// true, so a hook that ran under the weakened DenyRead token said so to nobody.
|
||||
// Delivering DispatchOutcome.Messages fixed that and overshot: hookMessage folds
|
||||
// the notice together with the hook's ordinary stdout, so every successful
|
||||
// hook's routine logging became a standing input channel into the next model
|
||||
// request, which is not what main does.
|
||||
//
|
||||
// One hook run produces both kinds of text here, because the bug is exactly a
|
||||
// failure to tell them apart. Asserted on what the PROVIDER received, since that
|
||||
// is the boundary that matters; a unit test on the joining helper cannot see
|
||||
// which slice the loop passes it.
|
||||
func TestSuccessfulBeforeToolHookDeliversItsNoticeAndNotItsOutput(t *testing.T) {
|
||||
options, provider, _ := readFileRunOptions(t, beforeToolDispatcher(t, hooks.EventBeforeTool, 0))
|
||||
if _, err := Run(context.Background(), "read the notes", provider, options); err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
|
||||
// The tool ran, so this is the successful-hook path rather than a blocked
|
||||
// call that never reached the tool.
|
||||
if !someRequestContains(provider.requests, "hello") {
|
||||
t.Fatal("SETUP INVALID: the tool result never reached the model, so nothing was delivered to check")
|
||||
}
|
||||
// And the hook really did run and really did print, or the silence asserted
|
||||
// below would be the silence of a hook that never executed.
|
||||
if !someRequestContains(provider.requests, beforeToolNotice) {
|
||||
t.Fatal("the enforcement notice never reached the model, so a hook could run under the weakened token and say so to nobody")
|
||||
}
|
||||
if got := countRequestsContaining(provider.requests, beforeToolNotice); got != 1 {
|
||||
t.Errorf("the notice reached the model %d times, want exactly once", got)
|
||||
}
|
||||
if someRequestContains(provider.requests, beforeToolChatter) {
|
||||
t.Error("the hook's ordinary output reached the model; main is silent for a successful hook and routine logging must not become model input")
|
||||
}
|
||||
}
|
||||
|
||||
// A VETO MUST NOT SWALLOW A NOTICE FROM A HOOK THAT ALREADY RAN.
|
||||
//
|
||||
// Dispatch runs hooks in order and returns at the first veto. The successful
|
||||
// hook ahead of it may already have run under the weakened token, and that is a
|
||||
// fact about something that happened. The veto result used to be built from the
|
||||
// blocking hook's Reason alone, so the earlier disclosure existed only in the
|
||||
// audit record.
|
||||
func TestABlockedCallStillCarriesTheEarlierHooksNotice(t *testing.T) {
|
||||
audit, err := hooks.NewAuditStore(hooks.AuditStoreOptions{AuditPath: filepath.Join(t.TempDir(), "audit.jsonl")})
|
||||
if err != nil {
|
||||
t.Fatalf("NewAuditStore: %v", err)
|
||||
}
|
||||
dispatcher := hooks.NewDispatcher(hooks.DispatcherOptions{
|
||||
Config: hooks.Config{
|
||||
Enabled: true,
|
||||
Hooks: []hooks.Definition{
|
||||
{ID: "zero.first", Event: hooks.EventBeforeTool, Matcher: "read_file", Command: "hook", Enabled: true},
|
||||
{ID: "zero.veto", Event: hooks.EventBeforeTool, Matcher: "read_file", Command: "veto", Enabled: true},
|
||||
},
|
||||
},
|
||||
Audit: audit,
|
||||
Cwd: t.TempDir(),
|
||||
Execution: execution.NewRunner(vetoSecondPreparer{}),
|
||||
})
|
||||
options, provider, _ := readFileRunOptions(t, dispatcher)
|
||||
if _, err := Run(context.Background(), "read the notes", provider, options); err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
|
||||
// SETUP: the second hook really did veto, or this is the ordinary path.
|
||||
if !someRequestContains(provider.requests, "was blocked by hook") {
|
||||
t.Fatal("SETUP INVALID: the call was not blocked, so the veto path is not under test")
|
||||
}
|
||||
if !someRequestContains(provider.requests, beforeToolNotice) {
|
||||
t.Error("the veto result dropped the notice from the hook that had already run under the weakened token")
|
||||
}
|
||||
if got := countRequestsContaining(provider.requests, beforeToolNotice); got != 1 {
|
||||
t.Errorf("the notice reached the model %d times, want exactly once", got)
|
||||
}
|
||||
if someRequestContains(provider.requests, beforeToolChatter) {
|
||||
t.Error("the vetoed result carried the earlier hook's ordinary output")
|
||||
}
|
||||
}
|
||||
|
||||
// vetoSecondPreparer runs the first hook successfully with a notice and makes
|
||||
// the second one exit non-zero, which is a veto for a blocking event.
|
||||
type vetoSecondPreparer struct{}
|
||||
|
||||
func (vetoSecondPreparer) PrepareExecution(_ context.Context, request execution.Request) (execution.PreparedCommand, error) {
|
||||
script := "echo " + beforeToolChatter
|
||||
notices := []string{beforeToolNotice}
|
||||
if request.Command.Name == "veto" {
|
||||
script = "exit 2"
|
||||
notices = nil
|
||||
}
|
||||
var command *exec.Cmd
|
||||
if runtime.GOOS == "windows" {
|
||||
command = exec.Command("cmd.exe", "/c", script)
|
||||
} else {
|
||||
command = exec.Command("/bin/sh", "-c", script)
|
||||
}
|
||||
return execution.PreparedCommand{
|
||||
Command: command,
|
||||
Enforcement: execution.Enforcement{Notices: notices},
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/execution"
|
||||
"github.com/Twigpine/zero/internal/hooks"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
zeroruntime "github.com/Twigpine/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// editMatchingDispatcher is beforeToolDispatcher for the edit tool, so the same
|
||||
// notice-and-chatter hook runs ahead of a call whose result carries a diff.
|
||||
func editMatchingDispatcher(t *testing.T) *hooks.Dispatcher {
|
||||
t.Helper()
|
||||
audit, err := hooks.NewAuditStore(hooks.AuditStoreOptions{AuditPath: filepath.Join(t.TempDir(), "audit.jsonl")})
|
||||
if err != nil {
|
||||
t.Fatalf("NewAuditStore: %v", err)
|
||||
}
|
||||
return hooks.NewDispatcher(hooks.DispatcherOptions{
|
||||
Config: hooks.Config{
|
||||
Enabled: true,
|
||||
Hooks: []hooks.Definition{
|
||||
{ID: "zero.before-tool", Event: hooks.EventBeforeTool, Matcher: "edit_file", Command: "hook", Enabled: true},
|
||||
},
|
||||
},
|
||||
Audit: audit,
|
||||
Cwd: t.TempDir(),
|
||||
Execution: execution.NewRunner(noticeHookPreparer{}),
|
||||
})
|
||||
}
|
||||
|
||||
// editFileRunOptions is readFileRunOptions for a tool whose result carries a
|
||||
// rich preview, which is the case the disclosure went missing on.
|
||||
func editFileRunOptions(t *testing.T, dispatcher *hooks.Dispatcher, onResult func(ToolResult)) (Options, *mockProvider) {
|
||||
t.Helper()
|
||||
root := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(root, "x.go"), []byte("package main\n\nconst answer = 41\n"), 0o644); err != nil {
|
||||
t.Fatalf("write x.go: %v", err)
|
||||
}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.NewScopedEditFileTool(root, nil))
|
||||
provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{
|
||||
{
|
||||
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call-1", ToolName: "edit_file"},
|
||||
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call-1", ArgumentsFragment: `{"path":"x.go","old_string":"41","new_string":"42"}`},
|
||||
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call-1"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
},
|
||||
{
|
||||
{Type: zeroruntime.StreamEventText, Content: "edited it"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
},
|
||||
}}
|
||||
return Options{
|
||||
SessionID: "session-hook-edit",
|
||||
Cwd: root,
|
||||
Registry: registry,
|
||||
ProviderName: "test-provider",
|
||||
Model: "test-model",
|
||||
Hooks: dispatcher,
|
||||
OnToolResult: onResult,
|
||||
PermissionMode: PermissionModeUnsafe,
|
||||
MaxTurns: 2,
|
||||
}, provider
|
||||
}
|
||||
|
||||
// THE NOTICE HAS TO SURVIVE A RESULT WHOSE BODY IS A DIFF.
|
||||
//
|
||||
// This is the boundary the component suites left between them. The delivery test
|
||||
// above proves a beforeTool notice reaches the provider with read_file, whose
|
||||
// body is its output. The card tests prove a card renders a notice when the
|
||||
// result already carries one typed. Neither composed the two, and in between them
|
||||
// the notice was being folded into result.Output as hook prose.
|
||||
//
|
||||
// For an edit or a write that is exactly where it disappears: the card builds its
|
||||
// body from Display.Preview rather than Output, and draws enforcement furniture
|
||||
// only from the typed slice, so the operator saw the diff and no disclosure at
|
||||
// all, live and restored, while the model saw the disclosure. A notice about a
|
||||
// weakened token was hidden on precisely the results where something was written.
|
||||
func TestABeforeToolNoticeSurvivesARichPreviewResult(t *testing.T) {
|
||||
dispatcher := editMatchingDispatcher(t)
|
||||
|
||||
var results []ToolResult
|
||||
options, provider := editFileRunOptions(t, dispatcher, func(result ToolResult) {
|
||||
results = append(results, result)
|
||||
})
|
||||
if _, err := Run(context.Background(), "bump the answer", provider, options); err != nil {
|
||||
t.Fatalf("Run: %v", err)
|
||||
}
|
||||
|
||||
if len(results) != 1 {
|
||||
t.Fatalf("SETUP INVALID: %d tool results, want the one edit", len(results))
|
||||
}
|
||||
result := results[0]
|
||||
if result.Status != tools.StatusOK {
|
||||
t.Fatalf("SETUP INVALID: the edit failed, so this is not the successful rich-preview path: %s", result.Output)
|
||||
}
|
||||
preview := result.BaseDisplay().Preview
|
||||
if !strings.Contains(preview, "42") {
|
||||
t.Fatalf("SETUP INVALID: the result carries no diff preview, which is the whole case under test: %q", preview)
|
||||
}
|
||||
|
||||
// TYPED, which is what every interactive surface reads.
|
||||
if len(result.EnforcementNotices) != 1 || result.EnforcementNotices[0] != beforeToolNotice {
|
||||
t.Errorf("the hook's disclosure did not reach the typed field, so the card and the restored card show the diff and no disclosure: %v", result.EnforcementNotices)
|
||||
}
|
||||
// AND NOT IN THE BODY AS WELL. Decoration has one owner per surface; carrying
|
||||
// it in both places renders it twice wherever the surface draws the slice.
|
||||
if strings.Contains(result.BaseModelOutput(), beforeToolNotice) {
|
||||
t.Errorf("the disclosure is in the result body as well as the typed field:\n%s", result.BaseModelOutput())
|
||||
}
|
||||
// The diff itself stays a parseable diff.
|
||||
if strings.Contains(preview, beforeToolNotice) {
|
||||
t.Errorf("the disclosure was glued onto the diff:\n%s", preview)
|
||||
}
|
||||
|
||||
// And the model still sees it, exactly once, with none of the hook's chatter.
|
||||
if got := countRequestsContaining(provider.requests, beforeToolNotice); got != 1 {
|
||||
t.Errorf("the notice reached the model %d times, want exactly once", got)
|
||||
}
|
||||
if someRequestContains(provider.requests, beforeToolChatter) {
|
||||
t.Error("the hook's ordinary output reached the model")
|
||||
}
|
||||
}
|
||||
@@ -1,143 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
"github.com/Twigpine/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// A CONFIGURATION REFUSAL IS NOT A RETRIABLE FAILURE, EVEN WHEN IT ARRIVES
|
||||
// EARLY.
|
||||
//
|
||||
// capture_artifact rejects in RejectBeforePermission, which the registry returns
|
||||
// straight back before any of the gates that attach provenance. Its
|
||||
// valid-but-unavailable calls therefore reached the classifier with no denial
|
||||
// category, no permission metadata and no refusal marker, and were read as
|
||||
// ordinary retriable failures: the model got the schema hint and the call could
|
||||
// consume the profile failure-streak escalation, for a tool that never executed
|
||||
// and that no argument change can enable.
|
||||
func TestDisabledCaptureArtifactIsAPolicyRefusalNotARetriableFailure(t *testing.T) {
|
||||
// No artifacts directory and no enabled driver: valid arguments, unavailable
|
||||
// tool. This is the shape an operator produces by configuration alone.
|
||||
registry := tools.NewRegistry()
|
||||
for _, tool := range tools.NewLocalControlArtifactTools(tools.LocalControlArtifactOptions{}) {
|
||||
registry.Register(tool)
|
||||
}
|
||||
|
||||
result := registry.RunWithOptions(context.Background(), "capture_artifact", map[string]any{
|
||||
"action": "browser_screenshot",
|
||||
"name": "shot",
|
||||
}, tools.RunOptions{PermissionGranted: true})
|
||||
|
||||
if result.Status != tools.StatusError {
|
||||
t.Fatalf("SETUP INVALID: expected the disabled tool to refuse, got %s: %s", result.Status, result.Output)
|
||||
}
|
||||
if !tools.IsPolicyRefusalResult(result) {
|
||||
t.Fatalf("the refusal carries no provenance, so nothing downstream can tell it from a failed command: %#v", result.Meta)
|
||||
}
|
||||
|
||||
// Through the conversion the loop performs, which is where the marker has to
|
||||
// survive to be worth anything.
|
||||
converted := ToolResult{
|
||||
Status: result.Status,
|
||||
Output: result.Output,
|
||||
Meta: result.Meta,
|
||||
}
|
||||
if !isPolicyRefusal(converted) {
|
||||
t.Error("the marker did not survive into the agent-facing result")
|
||||
}
|
||||
if isRetriableToolError(converted) {
|
||||
t.Error("a tool that never executed, and that no argument change can enable, was marked retriable: the model gets a schema hint telling it to fix arguments that were already valid")
|
||||
}
|
||||
}
|
||||
|
||||
// The malformed-argument branch must stay retriable. That one IS fixable by
|
||||
// trying again differently, which is what the hint exists for, so marking every
|
||||
// early rejection would trade one wrong answer for another.
|
||||
func TestMalformedCaptureArtifactArgumentsStayRetriable(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
for _, tool := range tools.NewLocalControlArtifactTools(tools.LocalControlArtifactOptions{}) {
|
||||
registry.Register(tool)
|
||||
}
|
||||
|
||||
result := registry.RunWithOptions(context.Background(), "capture_artifact", map[string]any{
|
||||
"action": "not_a_real_action",
|
||||
}, tools.RunOptions{PermissionGranted: true})
|
||||
|
||||
if result.Status != tools.StatusError {
|
||||
t.Fatalf("SETUP INVALID: expected invalid arguments to fail, got %s", result.Status)
|
||||
}
|
||||
if tools.IsPolicyRefusalResult(result) {
|
||||
t.Fatalf("a malformed-argument error was marked a policy refusal, so the model is denied the hint that would let it fix the call: %q", result.Output)
|
||||
}
|
||||
converted := ToolResult{Status: result.Status, Output: result.Output, Meta: result.Meta}
|
||||
if !isRetriableToolError(converted) {
|
||||
t.Error("invalid arguments should stay retriable")
|
||||
}
|
||||
}
|
||||
|
||||
// alwaysRefusingCaptureTool stands in for the disabled tool at Run level, so the
|
||||
// loop consequence can be observed rather than inferred.
|
||||
type alwaysRefusingCaptureTool struct{ ran int }
|
||||
|
||||
func (tool *alwaysRefusingCaptureTool) Name() string { return "capture_artifact" }
|
||||
func (tool *alwaysRefusingCaptureTool) Description() string { return "test capture tool" }
|
||||
func (tool *alwaysRefusingCaptureTool) Parameters() tools.Schema {
|
||||
return tools.Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]tools.PropertySchema{"action": {Type: "string"}},
|
||||
Required: []string{"action"},
|
||||
AdditionalProperties: false,
|
||||
}
|
||||
}
|
||||
func (tool *alwaysRefusingCaptureTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow, Reason: "captures artifacts"}
|
||||
}
|
||||
func (tool *alwaysRefusingCaptureTool) Run(context.Context, map[string]any) tools.Result {
|
||||
tool.ran++
|
||||
return tools.Result{}
|
||||
}
|
||||
func (tool *alwaysRefusingCaptureTool) RejectBeforePermission(map[string]any) (tools.Result, bool) {
|
||||
return tools.Result{
|
||||
Status: tools.StatusError,
|
||||
Output: "Error: capture_artifact is disabled because no artifact directory is configured.",
|
||||
Meta: map[string]string{tools.PolicyRefusalMeta: tools.PolicyRefusalToolNotEnabled},
|
||||
}, true
|
||||
}
|
||||
|
||||
// THE LOOP CONSEQUENCE. A refusal must not draw the retry hint, because the hint
|
||||
// tells the model to fix arguments that were already valid and the tool will
|
||||
// refuse identically next time.
|
||||
func TestRunDoesNotHintARefusedCaptureArtifact(t *testing.T) {
|
||||
tool := &alwaysRefusingCaptureTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
|
||||
calls := toolFailureHintAt + 1
|
||||
turns := make([][]zeroruntime.StreamEvent, 0, calls+1)
|
||||
for i := range calls {
|
||||
turns = append(turns, toolTurn("call-"+strconv.Itoa(i), "capture_artifact", `{"action":"browser_screenshot"}`))
|
||||
}
|
||||
turns = append(turns, textTurn("gave up on the screenshot"))
|
||||
|
||||
result, err := Run(context.Background(), "take a screenshot", &mockProvider{turns: turns}, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAsk,
|
||||
MaxTurns: len(turns) + 5,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tool.ran != 0 {
|
||||
t.Fatalf("SETUP INVALID: the tool executed %d times; it must be refused before Run", tool.ran)
|
||||
}
|
||||
for _, message := range result.Messages {
|
||||
if strings.Contains(message.Content, toolFailureHintMarker) {
|
||||
t.Fatal("a refused, never-executed tool drew the retry hint, which tells the model to fix arguments that were already valid")
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,112 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
"github.com/Twigpine/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// alternatingRefusalCaptureTool refuses every call with the same CATEGORY and a
|
||||
// different MESSAGE, which is what a real disabled driver does: the refusal
|
||||
// names the action, and the model is free to alternate valid actions.
|
||||
type alternatingRefusalCaptureTool struct{ ran int }
|
||||
|
||||
func (tool *alternatingRefusalCaptureTool) Name() string { return "capture_artifact" }
|
||||
func (tool *alternatingRefusalCaptureTool) Description() string { return "test capture tool" }
|
||||
func (tool *alternatingRefusalCaptureTool) Parameters() tools.Schema {
|
||||
return tools.Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]tools.PropertySchema{"action": {Type: "string"}},
|
||||
Required: []string{"action"},
|
||||
AdditionalProperties: false,
|
||||
}
|
||||
}
|
||||
func (tool *alternatingRefusalCaptureTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow, Reason: "captures artifacts"}
|
||||
}
|
||||
func (tool *alternatingRefusalCaptureTool) Run(context.Context, map[string]any) tools.Result {
|
||||
tool.ran++
|
||||
return tools.Result{}
|
||||
}
|
||||
func (tool *alternatingRefusalCaptureTool) RejectBeforePermission(args map[string]any) (tools.Result, bool) {
|
||||
action, _ := args["action"].(string)
|
||||
return tools.Result{
|
||||
Status: tools.StatusError,
|
||||
// Action-specific wording, same category. This is the whole point: the
|
||||
// prose differs on every call while the refusal never changes.
|
||||
Output: "Error: Local control driver for " + action + " is disabled.",
|
||||
Meta: map[string]string{tools.PolicyRefusalMeta: tools.PolicyRefusalToolNotEnabled},
|
||||
}, true
|
||||
}
|
||||
|
||||
// A REFUSAL KEYS ON ITS CATEGORY, INCLUDING WHEN THE CATEGORY IS ONLY A MARKER.
|
||||
//
|
||||
// The registry marks pre-execution refusals in metadata, and isPolicyRefusal
|
||||
// read that marker while observeToolResult keyed on DenialReason, which those
|
||||
// paths leave empty. The guard fell back to errorSignature(output), so two
|
||||
// refusals of the same category with different wording looked like two
|
||||
// different failures and the streak restarted at 1 every call.
|
||||
//
|
||||
// A model alternating capture_artifact's browser_screenshot and browser_pdf
|
||||
// against a disabled driver is refused identically each time and never tripped
|
||||
// the six-call halt. Only the generic twelve-error fallback stopped the run,
|
||||
// reporting varied errors rather than a repeated refusal.
|
||||
func TestAlternatingRefusedActionsStillTripTheRefusalHalt(t *testing.T) {
|
||||
tool := &alternatingRefusalCaptureTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
|
||||
// More calls than the refusal halt but FEWER than the generic any-error
|
||||
// fallback, so only the category-keyed streak can stop this.
|
||||
calls := toolFailureAnyErrorStopAt - 1
|
||||
if calls <= toolFailureStopAt {
|
||||
t.Fatalf("SETUP INVALID: %d calls cannot distinguish the refusal halt from the generic fallback", calls)
|
||||
}
|
||||
actions := []string{"browser_screenshot", "browser_pdf"}
|
||||
turns := make([][]zeroruntime.StreamEvent, 0, calls+1)
|
||||
for i := range calls {
|
||||
action := actions[i%len(actions)]
|
||||
turns = append(turns, toolTurn("call-"+strconv.Itoa(i), "capture_artifact", `{"action":"`+action+`"}`))
|
||||
}
|
||||
turns = append(turns, textTurn("gave up"))
|
||||
|
||||
result, err := Run(context.Background(), "capture something", &mockProvider{turns: turns}, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAsk,
|
||||
MaxTurns: len(turns) + 5,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tool.ran != 0 {
|
||||
t.Fatalf("SETUP INVALID: the tool executed %d times; it must be refused before Run", tool.ran)
|
||||
}
|
||||
|
||||
refusals := 0
|
||||
for _, message := range result.Messages {
|
||||
if strings.Contains(message.Content, "is disabled") {
|
||||
refusals++
|
||||
}
|
||||
}
|
||||
if refusals > toolFailureStopAt {
|
||||
t.Errorf("the run made %d refused calls; the six-call refusal halt never tripped because the streak re-keyed on each action's wording", refusals)
|
||||
}
|
||||
|
||||
// And the halt has to read as a repeated refusal, not as varied errors.
|
||||
stop := strings.ToLower(strings.Join(messageContents(result.Messages), "\n"))
|
||||
if strings.Contains(stop, toolFailureHintMarker) {
|
||||
t.Error("a refused, never-executed tool drew the retry hint")
|
||||
}
|
||||
}
|
||||
|
||||
func messageContents(messages []zeroruntime.Message) []string {
|
||||
out := make([]string, 0, len(messages))
|
||||
for _, message := range messages {
|
||||
out = append(out, message.Content)
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -1,93 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/localcontrol"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// THE DISABLED-DRIVER BRANCH, REACHED FOR REAL.
|
||||
//
|
||||
// RejectBeforePermission refuses on two separate conditions, and its neighbour
|
||||
// test constructs empty options, so it returns at the FIRST one (no artifact
|
||||
// directory configured) and the second is never evaluated. Reverting the
|
||||
// disabled-driver branch in local_capture.go to a plain errorResult therefore
|
||||
// left the whole agent and tools suites green.
|
||||
//
|
||||
// This is the shape an operator actually produces: an artifact root IS
|
||||
// configured and one driver IS enabled, and the model asks for an action owned
|
||||
// by a different, disabled driver. Without provenance that call is read as a
|
||||
// retriable failure, so it collects a schema hint and can spend the profile's
|
||||
// failure-streak escalation, for a tool no argument change can enable.
|
||||
func TestCaptureArtifactDisabledDriverIsAPolicyRefusal(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
options := tools.LocalControlArtifactOptions{
|
||||
ArtifactsDir: t.TempDir(),
|
||||
// Configured and enabled, so the tool as a whole is available and the
|
||||
// missing-directory branch cannot fire.
|
||||
Browser: localcontrol.BrowserOptions{Enabled: true, Driver: "test-browser", HelperPath: "browser-helper"},
|
||||
// Terminal deliberately left disabled.
|
||||
}
|
||||
for _, tool := range tools.NewLocalControlArtifactTools(options) {
|
||||
registry.Register(tool)
|
||||
}
|
||||
|
||||
// An enabled driver's action must still be accepted, or the test below would
|
||||
// pass for a tool that refuses everything.
|
||||
if result := registry.RunWithOptions(context.Background(), "capture_artifact", map[string]any{
|
||||
"action": "browser_screenshot", "name": "shot",
|
||||
}, tools.RunOptions{PermissionGranted: true}); tools.IsPolicyRefusalResult(result) {
|
||||
t.Fatalf("SETUP INVALID: the enabled browser driver was refused as policy: %s", result.Output)
|
||||
}
|
||||
|
||||
result := registry.RunWithOptions(context.Background(), "capture_artifact", map[string]any{
|
||||
"action": "terminal_snapshot", "name": "snap", "session": "test-session",
|
||||
}, tools.RunOptions{PermissionGranted: true})
|
||||
|
||||
if result.Status != tools.StatusError {
|
||||
t.Fatalf("SETUP INVALID: the disabled terminal driver did not refuse: %s / %s", result.Status, result.Output)
|
||||
}
|
||||
if !tools.IsPolicyRefusalResult(result) {
|
||||
t.Fatalf("a driver disabled by configuration refuses without provenance, so it reads as a retriable failure: %#v", result.Meta)
|
||||
}
|
||||
if got := result.Meta["policy_refusal"]; got != tools.PolicyRefusalToolNotEnabled {
|
||||
t.Errorf("policy_refusal = %q, want %q", got, tools.PolicyRefusalToolNotEnabled)
|
||||
}
|
||||
if isRetriableToolError(ToolResult{Status: result.Status, Output: result.ModelOutput(), Meta: result.Meta}) {
|
||||
t.Error("the refusal is still classified as retriable, so the model gets a schema hint for a configuration decision")
|
||||
}
|
||||
}
|
||||
|
||||
// AND THE EARLY REJECTIONS MUST NOT COLLAPSE INTO ONE ANSWER.
|
||||
//
|
||||
// RejectBeforePermission refuses on configuration; argument validation refuses
|
||||
// on a fixable mistake, and it runs FIRST, which is how the disabled-driver
|
||||
// branch stayed uncovered. A malformed call must stay an ordinary retriable
|
||||
// error, because trying again differently is exactly the right response to it.
|
||||
func TestCaptureArtifactMalformedArgumentsStayRetriable(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
for _, tool := range tools.NewLocalControlArtifactTools(tools.LocalControlArtifactOptions{
|
||||
ArtifactsDir: t.TempDir(),
|
||||
Terminal: localcontrol.TerminalOptions{Enabled: true, Driver: "test-terminal", HelperPath: "terminal-helper"},
|
||||
}) {
|
||||
registry.Register(tool)
|
||||
}
|
||||
|
||||
// The driver IS enabled, so nothing here is a configuration decision. The
|
||||
// call is simply missing the session the action requires.
|
||||
result := registry.RunWithOptions(context.Background(), "capture_artifact", map[string]any{
|
||||
"action": "terminal_snapshot", "name": "snap",
|
||||
}, tools.RunOptions{PermissionGranted: true})
|
||||
|
||||
if result.Status != tools.StatusError {
|
||||
t.Fatalf("SETUP INVALID: the malformed call did not fail: %s", result.Output)
|
||||
}
|
||||
if tools.IsPolicyRefusalResult(result) {
|
||||
t.Errorf("a fixable argument mistake was marked a policy refusal, so the model loses the schema hint that would fix it: %s", result.Output)
|
||||
}
|
||||
if !isRetriableToolError(ToolResult{Status: result.Status, Output: result.ModelOutput(), Meta: result.Meta}) {
|
||||
t.Errorf("a fixable argument mistake is not retriable: %s", result.Output)
|
||||
}
|
||||
}
|
||||
@@ -1,34 +0,0 @@
|
||||
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)
|
||||
}
|
||||
@@ -1,99 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// noticeProjectionTool stands in for any sandboxed command tool: it reports the
|
||||
// enforcement disclosure the way the real ones do, through the sandbox metadata
|
||||
// key that finalizeToolOutcome promotes into the typed notice slice.
|
||||
type noticeProjectionTool struct{}
|
||||
|
||||
func (noticeProjectionTool) Name() string { return "notice_projection" }
|
||||
func (noticeProjectionTool) Description() string { return "test tool carrying an enforcement notice" }
|
||||
func (noticeProjectionTool) Parameters() tools.Schema { return tools.Schema{Type: "object"} }
|
||||
func (noticeProjectionTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow}
|
||||
}
|
||||
|
||||
func (noticeProjectionTool) Run(ctx context.Context, args map[string]any) tools.Result {
|
||||
return tools.Result{
|
||||
Status: tools.StatusOK,
|
||||
Output: "the command output",
|
||||
Display: tools.Display{Summary: "the human summary"},
|
||||
Meta: map[string]string{"sandbox_notices": "least-privilege notice"},
|
||||
}
|
||||
}
|
||||
|
||||
// THE PROJECTION MUST CARRY ONE REPRESENTATION, NOT TWO.
|
||||
//
|
||||
// executeToolCall copies the typed notice slice into agent.ToolResult, so the
|
||||
// text fields it copies alongside must be the UNDECORATED base. Storing the
|
||||
// already-rendered text there instead leaves the same fact in two places with no
|
||||
// contract between them: it happens to render once today only because the
|
||||
// outcome arrives finalized and the agent accessor then reads Outcome.ModelView
|
||||
// rather than the stored field. Any result that reaches the accessor without a
|
||||
// finalized outcome renders the disclosure twice, and every raw reader of
|
||||
// .Output sees text that disagrees with Outcome.ModelView.
|
||||
func TestEnforcementNoticeIsStoredOnceAndRenderedOnce(t *testing.T) {
|
||||
const notice = "least-privilege notice"
|
||||
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(noticeProjectionTool{})
|
||||
|
||||
result, err := executeToolCall(context.Background(), registry, ToolCall{
|
||||
ID: "call-1", Name: "notice_projection", Arguments: `{}`,
|
||||
}, PermissionModeAuto, Options{Cwd: t.TempDir()})
|
||||
if err != nil {
|
||||
t.Fatalf("executeToolCall: %v", err)
|
||||
}
|
||||
if len(result.EnforcementNotices) == 0 {
|
||||
t.Fatalf("the notice never reached the agent result: %#v", result)
|
||||
}
|
||||
|
||||
// The stored fields are the canonical undecorated base, and they agree with
|
||||
// the finalized outcome they were projected from.
|
||||
if strings.Contains(result.Output, notice) {
|
||||
t.Errorf("ToolResult.Output stores the rendered notice as well as the slice: %q", result.Output)
|
||||
}
|
||||
if strings.Contains(result.Display.Summary, notice) {
|
||||
t.Errorf("ToolResult.Display.Summary stores the rendered notice as well as the slice: %q", result.Display.Summary)
|
||||
}
|
||||
if result.Output != result.Outcome.ModelView {
|
||||
t.Errorf("stored output %q disagrees with the finalized model view %q", result.Output, result.Outcome.ModelView)
|
||||
}
|
||||
if result.Display.Summary != result.Outcome.HumanView.Summary {
|
||||
t.Errorf("stored summary %q disagrees with the finalized human view %q", result.Display.Summary, result.Outcome.HumanView.Summary)
|
||||
}
|
||||
|
||||
// And every consumer that renders goes through the accessors, which show the
|
||||
// disclosure exactly once without hiding the output it is attached to.
|
||||
// loop.go builds the provider transcript from ModelOutput; the CLI writer and
|
||||
// the TUI cards use both accessors.
|
||||
transcript := result.ModelOutput()
|
||||
if got := strings.Count(transcript, notice); got != 1 {
|
||||
t.Errorf("the transcript shows the notice %d times, want 1: %q", got, transcript)
|
||||
}
|
||||
if !strings.Contains(transcript, "the command output") {
|
||||
t.Errorf("the transcript lost the command output: %q", transcript)
|
||||
}
|
||||
summary := result.HumanDisplay().Summary
|
||||
if got := strings.Count(summary, notice); got != 1 {
|
||||
t.Errorf("the human summary shows the notice %d times, want 1: %q", got, summary)
|
||||
}
|
||||
if !strings.Contains(summary, "the human summary") {
|
||||
t.Errorf("the human summary lost the tool summary: %q", summary)
|
||||
}
|
||||
}
|
||||
|
||||
// A result that never crossed the registry has no finalized outcome, so the
|
||||
// accessor falls back to the stored field. That is the path on which a stored
|
||||
// rendering would double, and it is the reason the contract above is stated on
|
||||
// the stored fields rather than only on the accessors.
|
||||
func TestUnfinalizedResultStillRendersTheNoticeOnce(t *testing.T) {
|
||||
const notice = "least-privilege notice"
|
||||
result := ToolResult{
|
||||
Status: tools.StatusOK,
|
||||
Output: "the command output",
|
||||
Display: tools.Display{Summary: "the human summary"},
|
||||
EnforcementNotices: []string{notice},
|
||||
}
|
||||
if result.Outcome.Finalized() {
|
||||
t.Fatal("fixture is finalized; it no longer covers the fallback path")
|
||||
}
|
||||
if got := strings.Count(result.ModelOutput(), notice); got != 1 {
|
||||
t.Errorf("ModelOutput shows the notice %d times, want 1: %q", got, result.ModelOutput())
|
||||
}
|
||||
if got := strings.Count(result.HumanDisplay().Summary, notice); got != 1 {
|
||||
t.Errorf("HumanDisplay shows the notice %d times, want 1: %q", got, result.HumanDisplay().Summary)
|
||||
}
|
||||
}
|
||||
+13
-192
@@ -55,22 +55,6 @@ const (
|
||||
// error, so this only affects true same-error loops.
|
||||
toolFailureStopAt = 6
|
||||
|
||||
// toolFailureAnyErrorStopAt halts a tool that keeps failing with DIFFERENT
|
||||
// errors. The streak above cannot see that case by construction: a changed
|
||||
// signature rebuilds the record at 1, so a tool whose error text varies every
|
||||
// call never reaches toolFailureStopAt however long it loops.
|
||||
//
|
||||
// That is not hypothetical. A headless run made 384 denied calls over 26
|
||||
// minutes without tripping a halt set to 6, because each denial named a
|
||||
// different path. Keying denials on their category (see observeToolResult)
|
||||
// fixes that case; this bound covers every error with no category to key on.
|
||||
//
|
||||
// Deliberately well above toolFailureStopAt. A model iterating on a genuinely
|
||||
// tricky edit legitimately fails several times with different errors while it
|
||||
// converges, and this must not cut those runs short — the same reasoning that
|
||||
// moved toolFailureStopAt from 4 to 6. Both counters reset on success.
|
||||
toolFailureAnyErrorStopAt = 12
|
||||
|
||||
// maxContinueNudges bounds how many times the headless completion gate
|
||||
// (Options.RequireCompletionSignal) re-prompts a model that stopped without a
|
||||
// tool call while work clearly remained. Once spent, the run finalizes as
|
||||
@@ -340,91 +324,15 @@ func acceptanceVerificationNudge(objective string) string {
|
||||
const toolFailureHintMarker = "kept failing with the same error"
|
||||
|
||||
type toolFailureRecord struct {
|
||||
count int
|
||||
identity failureIdentity
|
||||
// sawExecuted and sawRefused are the AGGREGATE over the whole streak, not
|
||||
// the current identity. The content-blind bound needs them because it fires
|
||||
// exactly when no identity repeated, so the last one is unrepresentative.
|
||||
sawExecuted bool
|
||||
sawRefused bool
|
||||
hintShown bool
|
||||
// anyErrorCount counts consecutive failures of this tool REGARDLESS of the
|
||||
// error, and is cleared only by a success. count above restarts whenever the
|
||||
// signature changes, which is exactly what a varying error message defeats;
|
||||
// this one cannot be reset by changing the text.
|
||||
anyErrorCount int
|
||||
count int
|
||||
errSig string
|
||||
hintShown bool
|
||||
}
|
||||
|
||||
// failureKind separates a failure the TOOL produced from one POLICY produced.
|
||||
//
|
||||
// The two used to share one string namespace, with provenance recovered
|
||||
// afterwards by testing that string for a "denial:" prefix. Trusted provenance
|
||||
// and untrusted content in one namespace is a namespace the untrusted side can
|
||||
// write into: a command that printed exactly "denial:permission_denied" and
|
||||
// exited non-zero acquired the identity of a real permission refusal, merged
|
||||
// its streak with one, and had the run report it as refused although it ran
|
||||
// every time. The kind is carried as its own field so no output can spell it.
|
||||
type failureKind uint8
|
||||
|
||||
const (
|
||||
// The zero value: no failure recorded.
|
||||
_ failureKind = iota
|
||||
// failureKindExecuted: the tool ran and returned an error. The key is the
|
||||
// normalized error signature, which is tool-controlled text.
|
||||
failureKindExecuted
|
||||
// failureKindRefused: policy refused the call before the tool ran. The key is
|
||||
// the denial category, a small closed enum the loop sets.
|
||||
failureKindRefused
|
||||
)
|
||||
|
||||
// failureIdentity is what the same-identity streak counts. Comparing the pair
|
||||
// means an executed error can never equal a refusal however it is spelled.
|
||||
type failureIdentity struct {
|
||||
kind failureKind
|
||||
key string
|
||||
}
|
||||
|
||||
// toolFailureCause is why the guard stopped, decided where the facts are rather
|
||||
// than reconstructed by the caller from overlapping booleans.
|
||||
type toolFailureCause uint8
|
||||
|
||||
const (
|
||||
// The zero value: the guard did not stop.
|
||||
_ toolFailureCause = iota
|
||||
// toolFailureCauseSameError: the same executed error signature, six times.
|
||||
toolFailureCauseSameError
|
||||
// toolFailureCauseSameRefusal: the same denial category, six times. The prose
|
||||
// differs on every call because it names what was refused, so this is NOT the
|
||||
// same-error case.
|
||||
toolFailureCauseSameRefusal
|
||||
// toolFailureCauseVariedExecuted: the content-blind bound, and every failure
|
||||
// was the tool executing and failing.
|
||||
toolFailureCauseVariedExecuted
|
||||
// toolFailureCauseVariedRefused: the content-blind bound, and every failure
|
||||
// was a policy refusal. Alternating two categories reaches twelve without
|
||||
// either reaching six, and the tool never ran once.
|
||||
toolFailureCauseVariedRefused
|
||||
// toolFailureCauseVariedMixed: the content-blind bound over both kinds.
|
||||
// Named deliberately rather than left to whichever field a switch tested
|
||||
// first.
|
||||
toolFailureCauseVariedMixed
|
||||
)
|
||||
|
||||
type toolFailureOutcome struct {
|
||||
InjectHint bool
|
||||
Stop bool
|
||||
Count int
|
||||
// Cause is set only when Stop is. It carries the aggregate fact about the
|
||||
// sequence that tripped a bound, which the last record alone cannot supply:
|
||||
// the content-blind bound is reached precisely when no single identity
|
||||
// repeated enough, so the identity present at the end says nothing about the
|
||||
// eleven before it.
|
||||
Cause toolFailureCause
|
||||
}
|
||||
|
||||
// Refused reports whether policy, rather than the tool, produced the failures.
|
||||
func (cause toolFailureCause) Refused() bool {
|
||||
return cause == toolFailureCauseSameRefusal || cause == toolFailureCauseVariedRefused
|
||||
}
|
||||
|
||||
// errorSignature normalizes a tool error to a short, comparable signature so
|
||||
@@ -450,47 +358,12 @@ func toolFailureHint(toolName, schemaJSON, errOutput string) string {
|
||||
|
||||
// toolFailureStopAnswer is the final answer when the repeated-failure guard halts
|
||||
// a run.
|
||||
func toolFailureStopAnswer(toolName string, count int, cause toolFailureCause) string {
|
||||
// Each branch claims only what its counter established, and the mixed case is
|
||||
// spelled out rather than decided by whichever field a switch happened to
|
||||
// test first. The wording used to overclaim in both directions: the
|
||||
// content-blind bound said every failure differed, which it does not track,
|
||||
// and it described a run of policy refusals as the tool failing, when the
|
||||
// tool had not run at all.
|
||||
verb, tail := " tool failed ", " times in a row with the same error, "
|
||||
switch cause {
|
||||
case toolFailureCauseSameRefusal:
|
||||
verb, tail = " tool was refused ", " times in a row, "
|
||||
case toolFailureCauseVariedExecuted:
|
||||
tail = " times in a row with varying errors, "
|
||||
case toolFailureCauseVariedRefused:
|
||||
verb, tail = " tool was refused ", " times in a row for different reasons, "
|
||||
case toolFailureCauseVariedMixed:
|
||||
verb, tail = " tool failed or was refused ", " times in a row, "
|
||||
}
|
||||
return "Agent stopped: the `" + toolName + "`" + verb + strconv.Itoa(count) +
|
||||
tail + "so I halted instead of looping further. " +
|
||||
func toolFailureStopAnswer(toolName string, count int) string {
|
||||
return "Agent stopped: the `" + toolName + "` tool failed " + strconv.Itoa(count) +
|
||||
" times in a row with the same error, so I halted instead of looping further. " +
|
||||
"Please check the request or adjust the tool arguments."
|
||||
}
|
||||
|
||||
// toolFailureIncompleteReason describes a guard halt for the headless
|
||||
// completion gate. A guard halt is BY DEFINITION unfinished work: the run
|
||||
// stopped because a tool would not stop failing, not because the task was done.
|
||||
//
|
||||
// It makes the same claim as toolFailureStopAnswer for the same halt, including
|
||||
// the mixed case: a headless caller reads this reason and an interactive one
|
||||
// reads the stop answer, and the two must not describe one halt differently.
|
||||
// Reported by @jatmn.
|
||||
func toolFailureIncompleteReason(toolName string, cause toolFailureCause) string {
|
||||
switch {
|
||||
case cause.Refused():
|
||||
return "halted after `" + toolName + "` was refused repeatedly, without completing the task"
|
||||
case cause == toolFailureCauseVariedMixed:
|
||||
return "halted after `" + toolName + "` failed or was refused repeatedly, without completing the task"
|
||||
}
|
||||
return "halted after `" + toolName + "` failed repeatedly, without completing the task"
|
||||
}
|
||||
|
||||
// The no-output stop answer is assembled from these fixed parts (only the turn
|
||||
// count varies). IsNoProgressStop matches all three so a legitimate message that
|
||||
// merely quotes the marker substring is not misclassified as a failed empty run.
|
||||
@@ -599,68 +472,28 @@ func newGuardState() *guardState {
|
||||
// observeToolResult tracks repeated identical failures of a tool. A successful
|
||||
// result clears that tool's failure streak. Returns whether to inject a one-shot
|
||||
// corrective hint and/or stop the run.
|
||||
// hintable is separate from failed on purpose. A categorized denial MUST count
|
||||
// toward the streaks — that is the whole point of keying on the category — but a
|
||||
// schema hint is the wrong response to a policy refusal: the call shape is fine,
|
||||
// the answer was no. Collapsing the two is what made the earlier version of this
|
||||
// fix a no-op, since the caller could only pass a flag that excluded denials.
|
||||
func (state *guardState) observeToolResult(name string, failed bool, hintable bool, output string, denial DenialCategory) toolFailureOutcome {
|
||||
func (state *guardState) observeToolResult(name string, failed bool, output string) toolFailureOutcome {
|
||||
if state.toolFailures == nil {
|
||||
state.toolFailures = map[string]*toolFailureRecord{}
|
||||
}
|
||||
if !failed {
|
||||
delete(state.toolFailures, name) // success resets both counters
|
||||
delete(state.toolFailures, name) // success resets the streak
|
||||
return toolFailureOutcome{}
|
||||
}
|
||||
// A denial keys on its CATEGORY, not its prose. The message embeds the path
|
||||
// or command that was refused, so it differs on every call while describing
|
||||
// the same unchanging refusal — which rebuilt the record at 1 each time and
|
||||
// let a denied tool loop indefinitely under a halt set to 6. The category is
|
||||
// a small closed enum the loop already sets on the result.
|
||||
identity := failureIdentity{kind: failureKindExecuted, key: errorSignature(output)}
|
||||
if denial != DenialNone {
|
||||
identity = failureIdentity{kind: failureKindRefused, key: string(denial)}
|
||||
}
|
||||
sig := errorSignature(output)
|
||||
record := state.toolFailures[name]
|
||||
if record == nil {
|
||||
record = &toolFailureRecord{identity: identity}
|
||||
if record == nil || record.errSig != sig {
|
||||
record = &toolFailureRecord{count: 1, errSig: sig}
|
||||
state.toolFailures[name] = record
|
||||
}
|
||||
if record.identity != identity {
|
||||
// A different failure restarts the same-identity streak but NOT the
|
||||
// content-blind one: changing how a tool fails is not progress.
|
||||
record.count = 1
|
||||
record.identity = identity
|
||||
record.hintShown = false
|
||||
} else {
|
||||
record.count++
|
||||
}
|
||||
record.anyErrorCount++
|
||||
if identity.kind == failureKindRefused {
|
||||
record.sawRefused = true
|
||||
} else {
|
||||
record.sawExecuted = true
|
||||
}
|
||||
|
||||
outcome := toolFailureOutcome{Count: record.count}
|
||||
switch {
|
||||
case record.count >= toolFailureStopAt:
|
||||
if record.count >= toolFailureStopAt {
|
||||
outcome.Stop = true
|
||||
outcome.Cause = toolFailureCauseSameError
|
||||
if identity.kind == failureKindRefused {
|
||||
outcome.Cause = toolFailureCauseSameRefusal
|
||||
}
|
||||
return outcome
|
||||
case record.anyErrorCount >= toolFailureAnyErrorStopAt:
|
||||
// Report the counter that actually tripped. record.count is the
|
||||
// same-identity streak and is often 1 here, which would describe a tool
|
||||
// that failed a dozen different ways as having failed once.
|
||||
outcome.Stop = true
|
||||
outcome.Count = record.anyErrorCount
|
||||
outcome.Cause = record.variedCause()
|
||||
return outcome
|
||||
}
|
||||
if hintable && record.count >= toolFailureHintAt && !record.hintShown {
|
||||
if record.count >= toolFailureHintAt && !record.hintShown {
|
||||
record.hintShown = true
|
||||
outcome.InjectHint = true
|
||||
}
|
||||
@@ -770,15 +603,3 @@ func (state *guardState) planReminder(turn int) string {
|
||||
|
||||
return ""
|
||||
}
|
||||
|
||||
// variedCause classifies a content-blind stop by what the whole streak held.
|
||||
func (record *toolFailureRecord) variedCause() toolFailureCause {
|
||||
switch {
|
||||
case record.sawRefused && record.sawExecuted:
|
||||
return toolFailureCauseVariedMixed
|
||||
case record.sawRefused:
|
||||
return toolFailureCauseVariedRefused
|
||||
default:
|
||||
return toolFailureCauseVariedExecuted
|
||||
}
|
||||
}
|
||||
@@ -2,8 +2,6 @@ package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
@@ -82,9 +80,6 @@ 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) {
|
||||
@@ -194,7 +189,7 @@ func TestUnknownExecSessionProbingTripsFailureHalt(t *testing.T) {
|
||||
var state guardState
|
||||
var stoppedAt int
|
||||
for i := 1; i <= toolFailureStopAt; i++ {
|
||||
out := state.observeToolResult(tools.WriteStdinToolName, true, true, tools.UnknownExecSessionError(i), "")
|
||||
out := state.observeToolResult(tools.WriteStdinToolName, true, tools.UnknownExecSessionError(i))
|
||||
if out.Stop {
|
||||
stoppedAt = i
|
||||
break
|
||||
@@ -205,223 +200,6 @@ func TestUnknownExecSessionProbingTripsFailureHalt(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// A permission denial repeats forever when the streak is keyed on the error
|
||||
// TEXT, because the denial message embeds the path or command that varies per
|
||||
// call.
|
||||
//
|
||||
// Observed, not theorised: a headless run made 384 denied calls over 26 minutes
|
||||
// without tripping a halt that stops at 6. Every denial carried the same typed
|
||||
// category and a different reason string, so errorSignature differed each time
|
||||
// and the record was rebuilt at count 1 on every call.
|
||||
//
|
||||
// TestUnknownExecSessionErrorSignatureIsIDInvariant fixed one message this way.
|
||||
// Keying on the category fixes the class, without needing every future denial
|
||||
// message to remember to be invariant.
|
||||
func TestPermissionDenialStreakSurvivesVaryingReasonText(t *testing.T) {
|
||||
var state guardState
|
||||
stoppedAt := 0
|
||||
for i := 1; i <= toolFailureStopAt; i++ {
|
||||
// The shape the loop actually produces: same tool, same category, a
|
||||
// different path every time.
|
||||
output := "Error: Permission denied for write_file: cannot write " +
|
||||
filepath.Join("C:", "ws", "pkg", "file"+strconv.Itoa(i)+".go")
|
||||
out := state.observeToolResult("write_file", true, true, output, DenialPermissionDenied)
|
||||
if out.Stop {
|
||||
stoppedAt = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if stoppedAt != toolFailureStopAt {
|
||||
t.Fatalf("denials with varying reason text stopped at %d, want %d", stoppedAt, toolFailureStopAt)
|
||||
}
|
||||
}
|
||||
|
||||
// The content-blind bound. A tool failing over and over with genuinely
|
||||
// DIFFERENT errors and no denial category is still a tool that is not working,
|
||||
// and the same-signature streak can never see it.
|
||||
//
|
||||
// Bounded well above toolFailureStopAt on purpose: a model iterating on a
|
||||
// tricky edit legitimately fails a few times with different errors while it
|
||||
// converges, and cutting that short is the regression
|
||||
// TestSuccessResetsBothFailureCounters guards.
|
||||
func TestToolFailingWithDifferentErrorsEveryTimeStillStops(t *testing.T) {
|
||||
var state guardState
|
||||
stoppedAt := 0
|
||||
for i := 1; i <= toolFailureAnyErrorStopAt; i++ {
|
||||
out := state.observeToolResult("bash", true, true, "distinct failure "+strconv.Itoa(i), "")
|
||||
if out.Stop {
|
||||
stoppedAt = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if stoppedAt != toolFailureAnyErrorStopAt {
|
||||
t.Fatalf("a tool failing with a new error each call stopped at %d, want %d", stoppedAt, toolFailureAnyErrorStopAt)
|
||||
}
|
||||
}
|
||||
|
||||
// Neither counter may outlive a success, or a long run that fails occasionally
|
||||
// and recovers would eventually halt for no reason. This is the property that
|
||||
// keeps the content-blind bound safe to add.
|
||||
func TestSuccessResetsBothFailureCounters(t *testing.T) {
|
||||
var state guardState
|
||||
for i := 1; i < toolFailureAnyErrorStopAt; i++ {
|
||||
if out := state.observeToolResult("bash", true, true, "distinct failure "+strconv.Itoa(i), ""); out.Stop {
|
||||
t.Fatalf("stopped at %d before the success that should reset it", i)
|
||||
}
|
||||
}
|
||||
state.observeToolResult("bash", false, false, "ok", "")
|
||||
|
||||
// Same again from zero. Reaching the bound a second time proves the counter
|
||||
// restarted rather than merely paused.
|
||||
stoppedAt := 0
|
||||
for i := 1; i <= toolFailureAnyErrorStopAt; i++ {
|
||||
if out := state.observeToolResult("bash", true, true, "later failure "+strconv.Itoa(i), ""); out.Stop {
|
||||
stoppedAt = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if stoppedAt != toolFailureAnyErrorStopAt {
|
||||
t.Fatalf("after a success the tool stopped at %d, want a full fresh %d", stoppedAt, toolFailureAnyErrorStopAt)
|
||||
}
|
||||
}
|
||||
|
||||
// Records are keyed per tool, so ANOTHER tool succeeding in between must not
|
||||
// clear the failing tool's streak.
|
||||
//
|
||||
// This is the realistic shape of the run that motivated the fix: the model kept
|
||||
// making progress elsewhere — reading files, updating its plan — while one tool
|
||||
// was refused over and over. A reset keyed on "something succeeded" rather than
|
||||
// "this tool succeeded" would make the halt unreachable in exactly the runs that
|
||||
// need it.
|
||||
func TestAnotherToolSucceedingDoesNotClearAFailingToolsStreak(t *testing.T) {
|
||||
var state guardState
|
||||
stoppedAt := 0
|
||||
for i := 1; i <= toolFailureStopAt; i++ {
|
||||
state.observeToolResult("read_file", false, false, "ok", "")
|
||||
out := state.observeToolResult("write_file", true, false,
|
||||
"Error: Permission denied for write_file: "+strconv.Itoa(i), DenialPermissionDenied)
|
||||
if out.Stop {
|
||||
stoppedAt = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if stoppedAt != toolFailureStopAt {
|
||||
t.Fatalf("denials interleaved with another tool's successes stopped at %d, want %d", stoppedAt, toolFailureStopAt)
|
||||
}
|
||||
}
|
||||
|
||||
// alwaysPromptingTool is never allowed to run: it exists so a Run-level test can
|
||||
// drive real permission denials through the loop.
|
||||
type alwaysPromptingTool struct{ ran int }
|
||||
|
||||
func (tool *alwaysPromptingTool) Name() string { return "bash" }
|
||||
func (tool *alwaysPromptingTool) Description() string { return "test shell tool" }
|
||||
func (tool *alwaysPromptingTool) Parameters() tools.Schema {
|
||||
return tools.Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]tools.PropertySchema{"command": {Type: "string"}},
|
||||
Required: []string{"command"},
|
||||
AdditionalProperties: false,
|
||||
}
|
||||
}
|
||||
func (tool *alwaysPromptingTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectShell, Permission: tools.PermissionPrompt, Reason: "runs shell commands"}
|
||||
}
|
||||
func (tool *alwaysPromptingTool) Run(context.Context, map[string]any) tools.Result {
|
||||
tool.ran++
|
||||
return tools.Result{Status: tools.StatusOK, Output: "should never run"}
|
||||
}
|
||||
|
||||
// The regression for the gap this whole change existed to close, driven through
|
||||
// Run rather than through the guard helper.
|
||||
//
|
||||
// The first version of this fix was a no-op in production and every one of its
|
||||
// unit tests passed, because they called observeToolResult directly with
|
||||
// failed=true. The loop passes isRetriableToolError, which returns false for a
|
||||
// categorized denial, so the guard took its success branch and deleted the
|
||||
// record before it could key on the category. A denied tool still looped to the
|
||||
// turn limit.
|
||||
//
|
||||
// This drives the real path: a tool that always prompts, an approver that always
|
||||
// denies, and a different command each turn so the denial REASON — and therefore
|
||||
// the error text — varies exactly as it does in a real run.
|
||||
func TestRunStopsARepeatedlyDeniedToolAtTheFailureBound(t *testing.T) {
|
||||
tool := &alwaysPromptingTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
|
||||
// Comfortably more turns than the bound, so reaching the bound is what stops
|
||||
// the run rather than exhausting the provider or MaxTurns.
|
||||
turns := make([][]zeroruntime.StreamEvent, 0, toolFailureStopAt+4)
|
||||
for i := range toolFailureStopAt + 4 {
|
||||
turns = append(turns, toolTurn("call-"+strconv.Itoa(i), "bash",
|
||||
`{"command":"touch /etc/file`+strconv.Itoa(i)+`"}`))
|
||||
}
|
||||
provider := &mockProvider{turns: turns}
|
||||
|
||||
denials := 0
|
||||
result, err := Run(context.Background(), "do the thing", provider, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAsk,
|
||||
MaxTurns: len(turns) + 5,
|
||||
OnPermissionRequest: func(_ context.Context, request PermissionRequest) (PermissionDecision, error) {
|
||||
denials++
|
||||
// The varying half. In production this is the refused path or command;
|
||||
// here it is the command, which lands in the denial message the guard
|
||||
// used to key on.
|
||||
return PermissionDecision{
|
||||
Action: PermissionDecisionDeny,
|
||||
Reason: "refused " + request.ToolName + " call " + strconv.Itoa(denials),
|
||||
}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if denials != toolFailureStopAt {
|
||||
t.Errorf("the run made %d denied calls, want it halted at %d", denials, toolFailureStopAt)
|
||||
}
|
||||
if tool.ran != 0 {
|
||||
t.Errorf("the denied tool executed %d times; the denial must precede execution", tool.ran)
|
||||
}
|
||||
want := toolFailureStopAnswer("bash", toolFailureStopAt, toolFailureCauseSameRefusal)
|
||||
if result.FinalAnswer != want {
|
||||
t.Errorf("final answer =\n %q\nwant\n %q", result.FinalAnswer, want)
|
||||
}
|
||||
}
|
||||
|
||||
// The stop message must describe the bound that actually tripped. When the
|
||||
// content-blind counter is what halts the run, the same-signature streak is
|
||||
// usually 1, and reporting that would tell the user a tool failed once after it
|
||||
// failed a dozen different ways.
|
||||
func TestVariedFailureStopAnswerReportsTheRightCounter(t *testing.T) {
|
||||
var state guardState
|
||||
var outcome toolFailureOutcome
|
||||
for i := 1; i <= toolFailureAnyErrorStopAt; i++ {
|
||||
outcome = state.observeToolResult("bash", true, true, "distinct failure "+strconv.Itoa(i), "")
|
||||
if outcome.Stop {
|
||||
break
|
||||
}
|
||||
}
|
||||
if !outcome.Stop {
|
||||
t.Fatal("never stopped")
|
||||
}
|
||||
if outcome.Cause != toolFailureCauseVariedExecuted {
|
||||
t.Errorf("Cause = %v for a stop driven by the content-blind counter over executed failures", outcome.Cause)
|
||||
}
|
||||
if outcome.Count != toolFailureAnyErrorStopAt {
|
||||
t.Errorf("Count = %d, want the counter that tripped (%d)", outcome.Count, toolFailureAnyErrorStopAt)
|
||||
}
|
||||
answer := toolFailureStopAnswer("bash", outcome.Count, outcome.Cause)
|
||||
if !strings.Contains(answer, "varying errors") {
|
||||
t.Errorf("stop answer describes the wrong cause: %q", answer)
|
||||
}
|
||||
if strings.Contains(answer, "with the same error") {
|
||||
t.Errorf("stop answer claims a same-error loop after distinct failures: %q", answer)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGuardStateResetsToolOnlyStreakOnEmptyNonToolTurn(t *testing.T) {
|
||||
var state guardState
|
||||
toolOnly := zeroruntime.CollectedStream{
|
||||
@@ -735,22 +513,6 @@ 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{})
|
||||
@@ -767,51 +529,6 @@ 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) {
|
||||
|
||||
@@ -84,12 +84,102 @@ func TestBlockedByHookResultFallsBackWhenReasonEmpty(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// THE BLOCKING HOOK'S OWN NOTICE IS A NOTICE, NOT PART OF THE VETO TEXT.
|
||||
//
|
||||
// It used to arrive folded into Reason, so it reached the model inside the error
|
||||
// message and never reached EnforcementNotices, which is what every surface
|
||||
// renders as the disclosure and what the session record keeps apart from the
|
||||
// body.
|
||||
func TestBlockedByHookResultCarriesTheBlockersNoticeOnTheTypedSlice(t *testing.T) {
|
||||
const notice = "denyRead is configured, so the write jail is not confining writes"
|
||||
out := blockedByHookResult(
|
||||
ToolCall{ID: "c4", Name: "bash"},
|
||||
hooks.DispatchOutcome{Blocked: true, BlockedBy: "policy", Reason: "policy violation", Notices: []string{notice}},
|
||||
)
|
||||
if len(out.EnforcementNotices) != 1 || out.EnforcementNotices[0] != notice {
|
||||
t.Fatalf("EnforcementNotices = %#v, want the blocking hook's notice", out.EnforcementNotices)
|
||||
}
|
||||
if strings.Contains(out.Output, notice) {
|
||||
t.Errorf("the notice is also inside the veto message, so it is said twice:\n%s", out.Output)
|
||||
}
|
||||
if got := strings.Count(out.ModelOutput(), notice); got != 1 {
|
||||
t.Errorf("the model sees the notice %d times, want once:\n%s", got, out.ModelOutput())
|
||||
}
|
||||
}
|
||||
|
||||
// AN EARLIER HOOK'S NOTICE SURVIVES A REASON THAT QUOTES IT.
|
||||
//
|
||||
// The veto path used to drop every notice whose text appeared anywhere in the
|
||||
// blocking hook's reason, on the theory that the reason already carried it. A
|
||||
// blocking hook that printed the same sentence, for example by relaying another
|
||||
// tool's stderr, therefore removed an earlier hook's disclosure from the typed
|
||||
// slice.
|
||||
func TestBlockedByHookResultKeepsAnEarlierNoticeTheReasonQuotes(t *testing.T) {
|
||||
const notice = "denyRead is configured, so the write jail is not confining writes"
|
||||
out := blockedByHookResult(
|
||||
ToolCall{ID: "c5", Name: "bash"},
|
||||
hooks.DispatchOutcome{
|
||||
Blocked: true,
|
||||
BlockedBy: "veto",
|
||||
Reason: "refusing: the last run reported \"" + notice + "\"",
|
||||
Notices: []string{notice},
|
||||
},
|
||||
)
|
||||
if len(out.EnforcementNotices) != 1 || out.EnforcementNotices[0] != notice {
|
||||
t.Fatalf("the earlier hook's notice was dropped because the veto quoted it: %#v", out.EnforcementNotices)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDispatchHelpersAreNoopWithoutDispatcher(t *testing.T) {
|
||||
options := Options{} // Hooks is nil
|
||||
if _, blocked := dispatchBeforeTool(context.Background(), options, ToolCall{Name: "bash"}, nil); blocked {
|
||||
t.Fatal("a nil dispatcher must never block a tool")
|
||||
}
|
||||
if feedback := dispatchAfterTool(context.Background(), options, ToolCall{Name: "bash"}, nil, tools.Result{}); feedback != "" {
|
||||
t.Fatalf("a nil dispatcher must yield no feedback, got %q", feedback)
|
||||
if feedback, notices := dispatchAfterTool(context.Background(), options, ToolCall{Name: "bash"}, nil, tools.Result{}); feedback != "" || notices != nil {
|
||||
t.Fatalf("a nil dispatcher must yield no feedback and no notices, got %q and %v", feedback, notices)
|
||||
}
|
||||
}
|
||||
|
||||
// A SUCCESSFUL beforeTool HOOK'S NOTICE MUST STAY TYPED.
|
||||
//
|
||||
// executeToolCall used to read the beforeTool outcome only when Blocked was true,
|
||||
// so a hook that ran fine and produced an enforcement notice — for instance that
|
||||
// it ran under the weakened DenyRead token — put that notice in the audit record
|
||||
// and nowhere anybody could see it. The first fix delivered it as prose appended
|
||||
// to the result output, which reached the model and no interactive surface,
|
||||
// because those build their enforcement furniture from the typed slice.
|
||||
//
|
||||
// So the delivery is the typed field, and this pins the merge: the hook's notice
|
||||
// first, the tool's own after it, each exactly once, blanks contributing nothing.
|
||||
func TestBeforeToolNoticesMergeIntoTheTypedField(t *testing.T) {
|
||||
const notice = "hook ran without WRITE_RESTRICTED because denyRead is configured"
|
||||
const toolOwned = "the sandbox dropped the network capability for this call"
|
||||
|
||||
result := withAppliedHookNotices(ToolResult{Output: "ok"}, []string{notice}, nil)
|
||||
if len(result.EnforcementNotices) != 1 || result.EnforcementNotices[0] != notice {
|
||||
t.Fatalf("a successful beforeTool notice did not reach the typed field: %v", result.EnforcementNotices)
|
||||
}
|
||||
// AND NOT THE OUTPUT AS WELL, or every surface that renders the slice shows
|
||||
// the disclosure twice.
|
||||
if strings.Contains(result.Output, notice) {
|
||||
t.Errorf("the notice was written into the output as well as the typed field: %q", result.Output)
|
||||
}
|
||||
|
||||
// Both arrive, hook first, when the tool carries its own.
|
||||
result = withAppliedHookNotices(ToolResult{EnforcementNotices: []string{toolOwned}}, []string{notice}, nil)
|
||||
if len(result.EnforcementNotices) != 2 || result.EnforcementNotices[0] != notice || result.EnforcementNotices[1] != toolOwned {
|
||||
t.Fatalf("the hook and tool notices did not merge in order: %v", result.EnforcementNotices)
|
||||
}
|
||||
|
||||
// The same disclosure from both sides is carried once.
|
||||
result = withAppliedHookNotices(ToolResult{EnforcementNotices: []string{notice}}, []string{notice}, nil)
|
||||
if len(result.EnforcementNotices) != 1 {
|
||||
t.Errorf("one disclosure reported by both the hook and the tool was carried %d times: %v", len(result.EnforcementNotices), result.EnforcementNotices)
|
||||
}
|
||||
|
||||
// Blank notices contribute nothing, so a run with no hook output stays silent.
|
||||
result = withAppliedHookNotices(ToolResult{}, []string{"", " "}, nil)
|
||||
if len(result.EnforcementNotices) != 0 {
|
||||
t.Errorf("blank hook notices produced %v, want nothing", result.EnforcementNotices)
|
||||
}
|
||||
}
|
||||
+210
-256
@@ -587,14 +587,6 @@ 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
|
||||
}
|
||||
@@ -719,60 +711,6 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
|
||||
// executed concurrently, consumed strictly in order below.
|
||||
var precomputed []precomputedToolResult
|
||||
precomputedStart, precomputedEnd := 0, 0
|
||||
// closeOutRemaining terminates every advertised call from `from` onward
|
||||
// TRUTHFULLY, and every early return in this loop goes through it.
|
||||
//
|
||||
// Read-ahead broke the assumption the abort placeholders were written
|
||||
// under. executeParallelReadBatch runs an entire eligible run before the
|
||||
// loop consumes any of it, so "not consumed yet" stopped meaning "not
|
||||
// executed": a sibling can already have run, committed its side effects
|
||||
// and any authorization credit they carry, and still sit past the index
|
||||
// where a terminal branch fires. Recording it as aborted threw that result
|
||||
// away, so the transcript disagreed with what had actually happened, and
|
||||
// its callbacks, trace counter, task observation and images went with it.
|
||||
//
|
||||
// So each remaining call is put in the one state that is true of it:
|
||||
// completed, and finalized exactly once with the same bookkeeping the main
|
||||
// path performs, or unstarted, and aborted. The guard is deliberately NOT
|
||||
// consulted for a drained sibling: these results cannot reverse a stop
|
||||
// decision that has already been made, they are only owed an honest record.
|
||||
closeOutRemaining := func(from int) {
|
||||
for next := from; next < len(collected.ToolCalls); next++ {
|
||||
sibling, ran := precomputedResultFor(precomputed, precomputedStart, precomputedEnd, next)
|
||||
if !ran {
|
||||
messages = appendAbortedToolResults(messages, collected.ToolCalls[next:next+1])
|
||||
continue
|
||||
}
|
||||
nextCall := collected.ToolCalls[next]
|
||||
if options.OnToolCall != nil {
|
||||
options.OnToolCall(nextCall)
|
||||
}
|
||||
options.Trace.Counter(trace.CounterToolCalls, 1)
|
||||
recordOutputBudgetTrace(options.Trace, sibling)
|
||||
task.observe(taskStateEvent{kind: taskStateEventToolResult, arguments: nextCall.Arguments, toolResult: sibling})
|
||||
if options.OnToolResult != nil {
|
||||
options.OnToolResult(sibling)
|
||||
}
|
||||
for _, name := range sibling.LoadedTools {
|
||||
loaded[name] = true
|
||||
}
|
||||
// The same fields the consumption path below records, provider call
|
||||
// ID included: a strict replay pairs the result to its tool_use by
|
||||
// that ID, and a finalized sibling without it was the one result in
|
||||
// the turn that could not be paired. Reported by @jatmn.
|
||||
messages = append(messages, zeroruntime.Message{
|
||||
Role: zeroruntime.MessageRoleTool,
|
||||
Content: sibling.ModelOutput(),
|
||||
ToolCallID: sibling.ToolCallID,
|
||||
ToolCallProviderID: nextCall.ProviderCallID,
|
||||
IsError: sibling.Status == tools.StatusError,
|
||||
ChangedFiles: append([]string(nil), sibling.ChangedFiles...),
|
||||
})
|
||||
if len(sibling.Images) > 0 {
|
||||
toolResultsWithImages = append(toolResultsWithImages, sibling)
|
||||
}
|
||||
}
|
||||
}
|
||||
for index, call := range collected.ToolCalls {
|
||||
// When this call starts a consecutive run of >= 2 auto-allowed read-only
|
||||
// calls, execute the whole run concurrently now (see parallel_tools.go).
|
||||
@@ -837,13 +775,13 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
|
||||
abortErr = ctx.Err()
|
||||
}
|
||||
if abortErr != nil {
|
||||
closeOutRemaining(index + 1)
|
||||
messages = appendAbortedToolResults(messages, collected.ToolCalls[index+1:])
|
||||
messages = append(messages, buildToolImageMessages()...)
|
||||
result.Messages = copyMessages(messages)
|
||||
return result, abortErr
|
||||
}
|
||||
if stopReason := stopReasonFromToolResult(toolResult); stopReason != "" {
|
||||
closeOutRemaining(index + 1)
|
||||
messages = appendAbortedToolResults(messages, collected.ToolCalls[index+1:])
|
||||
messages = append(messages, buildToolImageMessages()...)
|
||||
result.FinalAnswer = toolResult.ModelOutput()
|
||||
result.StopReason = stopReason
|
||||
@@ -858,29 +796,8 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
|
||||
// aren't fixed by reformatting the call, so a "match this schema" hint
|
||||
// would misdirect the model toward JSON shape or blocked behavior.
|
||||
retriableFailure := isRetriableToolError(toolResult)
|
||||
// A categorized denial is NOT retriable — retrying it verbatim is
|
||||
// pointless — but it is still a failure the streaks must count, or a
|
||||
// refused tool loops until the turn limit. Passing retriableFailure for
|
||||
// both is what let that happen: observeToolResult took its success
|
||||
// branch and deleted the record before it could key on the category.
|
||||
// isPolicyRefusal rather than DenialReason alone: the categories are
|
||||
// only attached where a typed denial is built, so a headless prompt
|
||||
// refusal and an uncategorized sandbox block were counted as SUCCESS
|
||||
// and cleared the record they were supposed to accumulate.
|
||||
policyRefusal := isPolicyRefusal(toolResult)
|
||||
countedFailure := retriableFailure || policyRefusal
|
||||
outcome := guards.observeToolResult(call.Name, countedFailure, retriableFailure, toolResult.ModelOutput(), toolResult.DenialReason)
|
||||
// The profile's failure-streak trigger is for RETRIABLE failures: it
|
||||
// restores the displaced turn budget and reasoning effort on the theory
|
||||
// that the tool is struggling and needs room. A policy refusal is not a
|
||||
// struggling tool, it is an answer, and spending the one-shot
|
||||
// escalation on one contradicts that trigger's own contract. Denials
|
||||
// still count for the halt above; they just do not buy more budget.
|
||||
if policyRefusal {
|
||||
posture.observeToolOutcome(toolFailureOutcome{}, toolResult)
|
||||
} else {
|
||||
posture.observeToolOutcome(outcome, toolResult)
|
||||
}
|
||||
outcome := guards.observeToolResult(call.Name, retriableFailure, toolResult.ModelOutput())
|
||||
posture.observeToolOutcome(outcome, toolResult)
|
||||
if outcome.Stop {
|
||||
// The assistant message advertised EVERY collected tool call, but
|
||||
// the guard halts mid-turn so the calls after this one never run.
|
||||
@@ -888,22 +805,10 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
|
||||
// every tool_use has a matching tool_result and the recorded
|
||||
// messages stay valid for a strict provider replay (Anthropic
|
||||
// rejects a tool_use with no answering tool_result).
|
||||
closeOutRemaining(index + 1)
|
||||
messages = appendAbortedToolResults(messages, collected.ToolCalls[index+1:])
|
||||
messages = append(messages, buildToolImageMessages()...)
|
||||
result.FinalAnswer = toolFailureStopAnswer(call.Name, outcome.Count, outcome.Cause)
|
||||
result.FinalAnswer = toolFailureStopAnswer(call.Name, outcome.Count)
|
||||
result.Messages = copyMessages(messages)
|
||||
// THE HALT HAS TO CROSS THE TERMINAL-STATUS BOUNDARY TOO.
|
||||
//
|
||||
// This branch returns straight out, so it never reached the
|
||||
// completion gate that the max-turns paths below go through. Under
|
||||
// RequireCompletionSignal `zero exec` treats only Incomplete as exit
|
||||
// 4 and otherwise emits run_end("success", 0), so a task that was
|
||||
// denied six times, or failed through the twelve-call bound, became
|
||||
// a successful automation result having done none of the work.
|
||||
if options.RequireCompletionSignal {
|
||||
result.Incomplete = true
|
||||
result.IncompleteReason = toolFailureIncompleteReason(call.Name, outcome.Cause)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
if outcome.InjectHint && failureHint == "" {
|
||||
@@ -1562,10 +1467,31 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
|
||||
}
|
||||
|
||||
// beforeTool hooks may veto the call before it runs (a non-zero exit blocks).
|
||||
//
|
||||
// A SUCCESSFUL beforeTool HOOK STILL HAS SOMETHING TO SAY, BUT ONLY ITS NOTICE.
|
||||
//
|
||||
// Reading the outcome only when Blocked left the enforcement disclosure in the
|
||||
// audit record and nowhere the model or the operator could see it: a hook could
|
||||
// run under the weakened DenyRead token and say so to nobody.
|
||||
//
|
||||
// Notices, NOT Messages. Messages is presentation text that hookMessage builds
|
||||
// by folding the notice together with the hook's ordinary stdout, so delivering
|
||||
// it would put every successful hook's routine logging, large diagnostics, and
|
||||
// whatever text a hook happened to process into the next model request. That is
|
||||
// a behaviour change nobody asked for and a standing input channel. main is
|
||||
// silent for successful hooks and stays silent here for everything except the
|
||||
// disclosure. Carried to the tool result below, the same surface afterTool
|
||||
// feedback already uses.
|
||||
var beforeToolNotices []string
|
||||
// Filled after the tool runs, and merged with beforeTool's at the one
|
||||
// finalization below so both phases land on the same contract.
|
||||
var afterToolNotices []string
|
||||
if toolFound {
|
||||
if outcome, blocked := dispatchBeforeTool(ctx, options, call, args); blocked {
|
||||
outcome, blocked := dispatchBeforeTool(ctx, options, call, args)
|
||||
if blocked {
|
||||
return blockedByHookResult(call, outcome), nil
|
||||
}
|
||||
beforeToolNotices = outcome.Notices
|
||||
}
|
||||
args = shellExecutionArgsForApproval(call.Name, args, decisionAction, options)
|
||||
|
||||
@@ -1611,7 +1537,10 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
|
||||
})
|
||||
if retryResult, directResult, retried, action, reason, prefix, abortErr := maybeRetryUnsandboxedAfterSandboxRestriction(ctx, registry, call, tool, args, result, permissionMode, options, progressCallback); retried || directResult != nil || abortErr != nil {
|
||||
if directResult != nil {
|
||||
return *directResult, abortErr
|
||||
// A denied, cancelled, or ungrantable retry still returns a result for a
|
||||
// call whose beforeTool hook already ran. Without this the disclosure is
|
||||
// produced and then dropped on the floor.
|
||||
return withAppliedHookNotices(*directResult, beforeToolNotices, nil), abortErr
|
||||
}
|
||||
result = retryResult
|
||||
permissionGranted = true
|
||||
@@ -1637,7 +1566,15 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
|
||||
// afterTool hooks run once the tool has executed; their output (e.g. a
|
||||
// formatter or vet result) is surfaced back to the model on the result.
|
||||
if toolFound {
|
||||
if feedback := dispatchAfterTool(ctx, options, call, args, result); feedback != "" {
|
||||
// NEITHER PHASE'S NOTICES ARE FOLDED IN HERE. They are typed enforcement
|
||||
// data on both sides now, and appending them as hook prose was what kept
|
||||
// beforeTool's out of every surface that renders the typed slice, and what
|
||||
// made afterTool's arrive twice once that slice was composed. Only the
|
||||
// hook's own output, which is what an afterTool validator asked to say,
|
||||
// still travels this way.
|
||||
feedback, notices := dispatchAfterTool(ctx, options, call, args, result)
|
||||
afterToolNotices = notices
|
||||
if strings.TrimSpace(feedback) != "" {
|
||||
var didRedact bool
|
||||
result.Output, didRedact = appendHookFeedback(result.Output, feedback)
|
||||
if didRedact {
|
||||
@@ -1664,60 +1601,32 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
|
||||
// Secret scrubbing happens at the registry boundary (the single point both
|
||||
// the agent loop and the MCP server pass through), so result.Output is
|
||||
// already redacted here and result.Redacted reflects whether it changed.
|
||||
return ToolResult{
|
||||
Risk: executedRisk,
|
||||
ToolCallID: call.ID,
|
||||
Name: call.Name,
|
||||
Status: result.Status,
|
||||
Output: result.ModelOutput(),
|
||||
Truncated: result.Truncated,
|
||||
Meta: result.Meta,
|
||||
Images: result.Images,
|
||||
Redacted: result.Redacted,
|
||||
ChangedFiles: result.ChangedFiles,
|
||||
FileDiffs: result.FileDiffs,
|
||||
ChangeSummaries: result.ChangeSummaries,
|
||||
Display: result.HumanDisplay(),
|
||||
Outcome: result.Outcome,
|
||||
LoadedTools: loadedToolsFromResult(result.Meta),
|
||||
//
|
||||
// Wrapped in the same finalization every other return from a hooked call
|
||||
// uses, so the normal path cannot drift away from the veto and retry paths.
|
||||
return withAppliedHookNotices(ToolResult{
|
||||
Risk: executedRisk,
|
||||
ToolCallID: call.ID,
|
||||
Name: call.Name,
|
||||
Status: result.Status,
|
||||
Output: result.BaseModelOutput(),
|
||||
Truncated: result.Truncated,
|
||||
Meta: result.Meta,
|
||||
EnforcementNotices: append([]string(nil), result.EnforcementNotices...),
|
||||
Images: result.Images,
|
||||
Redacted: result.Redacted,
|
||||
ChangedFiles: result.ChangedFiles,
|
||||
FileDiffs: result.FileDiffs,
|
||||
ChangeSummaries: result.ChangeSummaries,
|
||||
Display: result.BaseDisplay(),
|
||||
Outcome: result.Outcome,
|
||||
LoadedTools: loadedToolsFromResult(result.Meta),
|
||||
// A tool may signal a mid-run model escalation by carrying the target id
|
||||
// in Meta["escalate_to_model"]. Lift it into the typed loop-level field;
|
||||
// the Run turn loop performs the actual provider switch. Empty for every
|
||||
// ordinary tool result.
|
||||
RequestedModel: result.Meta["escalate_to_model"],
|
||||
// ONE REFUSAL IDENTITY, derived here so classification and streak
|
||||
// accounting cannot disagree.
|
||||
//
|
||||
// The registry marks its pre-execution refusals in metadata, and
|
||||
// isPolicyRefusal read that marker while observeToolResult keyed on
|
||||
// DenialReason, which those paths leave empty. The guard therefore fell
|
||||
// back to errorSignature(output), and two refusals of the same category
|
||||
// with different wording looked like two different failures. A model
|
||||
// alternating capture_artifact's browser_screenshot and browser_pdf against
|
||||
// a disabled driver is refused identically each time, and never tripped the
|
||||
// six-call refusal halt: only the generic twelve-error fallback stopped it,
|
||||
// reporting varied errors rather than a repeated refusal.
|
||||
DenialReason: denialCategoryForResult(result),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// denialCategoryForResult gives every pre-execution refusal a stable category,
|
||||
// whether it was built as a typed denial or only marked in metadata.
|
||||
//
|
||||
// The mapping lives here, at the one boundary tools.Result becomes ToolResult,
|
||||
// rather than at each producer. A marker without a category is the shape that
|
||||
// caused this: it classified as a refusal and keyed as an error signature.
|
||||
func denialCategoryForResult(result tools.Result) DenialCategory {
|
||||
switch result.Meta[tools.PolicyRefusalMeta] {
|
||||
case tools.PolicyRefusalToolNotEnabled:
|
||||
return DenialFiltered
|
||||
case tools.PolicyRefusalPermissionDenied, tools.PolicyRefusalPermissionRequired:
|
||||
return DenialPermissionDenied
|
||||
case tools.PolicyRefusalSandboxDenied, tools.PolicyRefusalSandboxApproval:
|
||||
return DenialSandboxBlock
|
||||
default:
|
||||
return DenialNone
|
||||
}
|
||||
}, beforeToolNotices, afterToolNotices), nil
|
||||
}
|
||||
|
||||
const sandboxNamespaceLimitedReason = "sandbox output is limited to the sandbox PID namespace; host/global state requires approval"
|
||||
@@ -2050,10 +1959,6 @@ func toolResultFromPrePermissionReject(call ToolCall, result tools.Result) ToolR
|
||||
Display: display,
|
||||
LoadedTools: loadedToolsFromResult(meta),
|
||||
RequestedModel: meta["escalate_to_model"],
|
||||
// The SAME identity the executed path derives. This is the route a
|
||||
// RejectBeforePermission refusal takes, so leaving it empty here is what
|
||||
// made the marker classify as a refusal and key as an error signature.
|
||||
DenialReason: denialCategoryForResult(tools.Result{Meta: meta}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2105,9 +2010,9 @@ func dispatchBeforeTool(ctx context.Context, options Options, call ToolCall, arg
|
||||
// dispatchAfterTool runs configured afterTool hooks once a tool has executed and
|
||||
// returns any advisory output (e.g. a formatter or vet result) to surface back
|
||||
// to the model. afterTool hooks never block. A nil dispatcher is a no-op.
|
||||
func dispatchAfterTool(ctx context.Context, options Options, call ToolCall, args map[string]any, result tools.Result) string {
|
||||
func dispatchAfterTool(ctx context.Context, options Options, call ToolCall, args map[string]any, result tools.Result) (string, []string) {
|
||||
if options.Hooks == nil || hooksSuppressed(options) {
|
||||
return ""
|
||||
return "", nil
|
||||
}
|
||||
outcome := options.Hooks.Dispatch(ctx, hooks.DispatchInput{
|
||||
Event: hooks.EventAfterTool,
|
||||
@@ -2123,7 +2028,13 @@ func dispatchAfterTool(ctx context.Context, options Options, call ToolCall, args
|
||||
"changedFiles": result.ChangedFiles,
|
||||
},
|
||||
})
|
||||
return strings.TrimSpace(strings.Join(outcome.Messages, "\n"))
|
||||
// Both halves, because they have different destinations. The hook's own output
|
||||
// is prose that belongs in the result body where an afterTool validator's
|
||||
// findings have always gone; its enforcement notices are typed data that the
|
||||
// caller merges into the result the same way beforeTool's are. Returning only
|
||||
// the joined messages was what left afterTool on the old channel while
|
||||
// beforeTool moved, and a notice reported by both then appeared twice.
|
||||
return strings.TrimSpace(strings.Join(outcome.Messages, "\n")), outcome.Notices
|
||||
}
|
||||
|
||||
// dispatchSessionStart runs configured sessionStart hooks once before the first
|
||||
@@ -2187,7 +2098,7 @@ func blockedByHookResult(call ToolCall, outcome hooks.DispatchOutcome) ToolResul
|
||||
reason = "blocked by a beforeTool hook"
|
||||
}
|
||||
message := fmt.Sprintf("Error: %q was blocked by hook %q: %s", call.Name, outcome.BlockedBy, reason)
|
||||
return ToolResult{
|
||||
result := ToolResult{
|
||||
ToolCallID: call.ID,
|
||||
Name: call.Name,
|
||||
Status: tools.StatusError,
|
||||
@@ -2195,12 +2106,124 @@ func blockedByHookResult(call ToolCall, outcome hooks.DispatchOutcome) ToolResul
|
||||
Redacted: redacted,
|
||||
DenialReason: DenialHookBlocked,
|
||||
}
|
||||
// Dispatch runs hooks in order and stops at the first veto, so an earlier hook
|
||||
// may already have run under a weakened token before this one said no. Its
|
||||
// notice describes something that happened and has to survive the veto, and so
|
||||
// does the blocking hook's own.
|
||||
//
|
||||
// ALL OF THEM, FROM THE TYPED SLICE. Reason carries the hook's words only, so
|
||||
// nothing here is said twice and nothing has to be filtered out. The filter
|
||||
// this replaced matched notices against the reason text, and dropped an
|
||||
// earlier hook's disclosure whenever the blocking hook's output quoted it.
|
||||
return withAppliedHookNotices(result, outcome.Notices, nil)
|
||||
}
|
||||
|
||||
// withAppliedHookNotices is the single place a hook enforcement notice
|
||||
// reaches a tool result on a path that does NOT run afterTool.
|
||||
//
|
||||
// The normal tail joins the notices with the afterTool feedback and delivers
|
||||
// both at once. Two other exits return a result for a call whose hook already
|
||||
// ran: a later hook's veto, and a denied, cancelled, or ungrantable unsandboxed
|
||||
// retry. Routing all three through one function is what keeps "the hook ran
|
||||
// under this token" from depending on which exit the call happened to take.
|
||||
//
|
||||
// NO REBUDGET HERE, AND THAT IS LOAD-BEARING ON WHAT MAY PASS THROUGH. The
|
||||
// normal tail appends and then calls Registry.RebudgetAfterHook, because what it
|
||||
// appends is afterTool feedback: hook stdout, which a hook can make arbitrarily
|
||||
// large. These notices cannot be: they are CommandPlan.Notes, which only this
|
||||
// package's own fixed sentences may fill (none does today; #1006 refuses the
|
||||
// denyRead trade the last one described), and nothing hook-authored reaches
|
||||
// this slice: DispatchOutcome.Messages is where
|
||||
// hook output lives, and the capture site deliberately does not read it.
|
||||
//
|
||||
// So if anything ever widens what is delivered here to include text a hook or a
|
||||
// tool can size, this needs the rebudget step as well, which means converting
|
||||
// through tools.Result the way the tail does rather than editing Output in
|
||||
// place. Do not widen it without that.
|
||||
// withAppliedHookNotices carries a hook's disclosures onto the result as TYPED
|
||||
// enforcement notices, and is the single finalization point every return path
|
||||
// from a hooked call goes through.
|
||||
//
|
||||
// BOTH PHASES, because half a symmetry is its own defect. beforeTool was moved
|
||||
// onto the typed slice while afterTool was left folding its notices into the
|
||||
// prose feedback, and since the two carry the identical fixed string, a
|
||||
// disclosure reported by both then reached the model twice: once from the typed
|
||||
// prepend every surface composes, once inside the hook feedback in the body. A
|
||||
// bash or exec card showed it in the furniture and again in the body.
|
||||
//
|
||||
// It used to fold them into result.Output as prose. That reached the provider,
|
||||
// which reads the output, and reached nothing else. Every interactive surface
|
||||
// builds its enforcement furniture from the typed slice, and for an edit or a
|
||||
// write the card shows Display.Preview instead of Output, so the disclosure was
|
||||
// absent from the live card in both its collapsed and expanded states and from
|
||||
// the session payload that restores them. A notice about a weakened token was
|
||||
// therefore shown to the model and hidden from the operator, on exactly the
|
||||
// results where something was written.
|
||||
//
|
||||
// Typed here, the canonical accessors compose it per surface, the same way a
|
||||
// tool's own notices already work. Nothing is written into Output as well:
|
||||
// decoration has one owner per surface or the disclosure appears twice.
|
||||
// Hook notices lead, in run order, with the tool's own after them: dedupe makes
|
||||
// the ordering moot for the duplicate case that motivated this, and for distinct
|
||||
// disclosures it keeps the shape the tool-owned notices already had.
|
||||
func withAppliedHookNotices(result ToolResult, before []string, after []string) ToolResult {
|
||||
hookNotices := before
|
||||
if len(after) > 0 {
|
||||
hookNotices = append(append([]string(nil), before...), after...)
|
||||
}
|
||||
merged, didRedact := mergeEnforcementNotices(hookNotices, result.EnforcementNotices)
|
||||
result.EnforcementNotices = merged
|
||||
if didRedact {
|
||||
result.Redacted = true
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// mergeEnforcementNotices puts the hook disclosures ahead of the
|
||||
// tool's own and drops exact repeats, so a surface rendering the slice shows
|
||||
// each disclosure exactly once.
|
||||
//
|
||||
// Hook notices are third-party text arriving on an intercepted path that bypasses
|
||||
// the registry's redaction boundary, so they are scrubbed here the way
|
||||
// appendHookFeedback scrubbed them while they travelled as prose. The bool
|
||||
// reports whether scrubbing changed anything, so Redacted keeps matching the
|
||||
// registry's contract.
|
||||
func mergeEnforcementNotices(before []string, own []string) ([]string, bool) {
|
||||
if len(before) == 0 && len(own) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
merged := make([]string, 0, len(before)+len(own))
|
||||
seen := make(map[string]struct{}, len(before)+len(own))
|
||||
redacted := false
|
||||
for index, notice := range append(append([]string(nil), before...), own...) {
|
||||
if index < len(before) {
|
||||
scrubbed := redaction.RedactString(notice, redaction.Options{})
|
||||
redacted = redacted || scrubbed != notice
|
||||
notice = scrubbed
|
||||
}
|
||||
if strings.TrimSpace(notice) == "" {
|
||||
continue
|
||||
}
|
||||
if _, already := seen[notice]; already {
|
||||
continue
|
||||
}
|
||||
seen[notice] = struct{}{}
|
||||
merged = append(merged, notice)
|
||||
}
|
||||
if len(merged) == 0 {
|
||||
return nil, redacted
|
||||
}
|
||||
return merged, redacted
|
||||
}
|
||||
|
||||
// appendHookFeedback appends afterTool hook output to a tool result's output,
|
||||
// scrubbed for secrets like every other string crossing the tool boundary. The
|
||||
// returned bool reports whether scrubbing changed the feedback, so the caller can
|
||||
// set ToolResult.Redacted to match the registry's redaction contract.
|
||||
// The joiner that folded beforeTool notices in with the afterTool feedback is
|
||||
// gone. Sending a typed enforcement notice out as hook prose was the defect, not
|
||||
// the delivery: see withAppliedHookNotices. afterTool feedback still arrives here
|
||||
// on its own, which is all this path was ever meant to carry.
|
||||
func appendHookFeedback(output string, feedback string) (string, bool) {
|
||||
scrubbed := redaction.RedactString(feedback, redaction.Options{})
|
||||
redacted := scrubbed != feedback
|
||||
@@ -2225,64 +2248,21 @@ func isRetriableToolError(result ToolResult) bool {
|
||||
// A categorized denial (filtered / permission / sandbox) is a policy decision,
|
||||
// not a transient failure — never retry it. This is robust to message wording
|
||||
// (the text checks below remain as a fallback for results lacking the field).
|
||||
return !isPolicyRefusal(result)
|
||||
}
|
||||
|
||||
// isPolicyRefusal reports a result the run refused on policy grounds: a
|
||||
// permission gate, a filter, a sandbox preflight, or a hook.
|
||||
//
|
||||
// Split out of isRetriableToolError so the two questions cannot drift apart,
|
||||
// which is exactly what they had done. The halt guard asked
|
||||
// `DenialReason != DenialNone`, and a category is only attached on the paths
|
||||
// that build a TYPED denial. A headless run leaves OnPermissionRequest nil, so
|
||||
// the loop never reaches its typed-denial branch and the registry returns a bare
|
||||
// `Error: Permission required ...` carrying no category. A sandbox preflight
|
||||
// denial on a non-shell tool loses its SandboxDecision converting to ToolResult
|
||||
// and arrives as an uncategorized `Sandbox block`.
|
||||
//
|
||||
// Both are refusals, and both were counted as SUCCESS, so observeToolResult
|
||||
// cleared the record and the same refused call could repeat to MaxTurns. That is
|
||||
// the loop this guard exists to stop. The text checks below already enumerated
|
||||
// these outcomes for the retriable question; they simply were not reachable from
|
||||
// the counting one.
|
||||
func isPolicyRefusal(result ToolResult) bool {
|
||||
// A REFUSAL IS A FAILURE. Everything below inspects metadata and then output
|
||||
// text, and neither is meaningful on a result the tool completed.
|
||||
//
|
||||
// isRetriableToolError gates on this before it calls here, so the boundary
|
||||
// held while that was the only caller. Extracting this helper and calling it
|
||||
// straight from the counting path dropped the gate: an allowed `bash`
|
||||
// printing "Sandbox block", or a read_file returning a document that quotes
|
||||
// one of these phrases, was counted as a denial. Six such successes tripped
|
||||
// the same-signature stop and ended a healthy run with a refusal answer.
|
||||
//
|
||||
// The gate belongs here rather than at each caller, because the next caller
|
||||
// will forget it too.
|
||||
if result.Status != tools.StatusError {
|
||||
if result.DenialReason != DenialNone {
|
||||
return false
|
||||
}
|
||||
if result.DenialReason != DenialNone {
|
||||
return true
|
||||
if result.Meta["permission_action"] == string(PermissionActionDeny) {
|
||||
return false
|
||||
}
|
||||
if result.Meta[tools.PermissionActionMeta] == string(PermissionActionDeny) {
|
||||
return true
|
||||
switch {
|
||||
case strings.Contains(result.Output, "is not enabled for this run"),
|
||||
strings.Contains(result.Output, "Permission denied for "),
|
||||
strings.Contains(result.Output, "Permission required for "),
|
||||
strings.Contains(result.Output, "Sandbox block"),
|
||||
strings.Contains(result.Output, "Sandbox approval required for "):
|
||||
return false
|
||||
}
|
||||
// STRUCTURED ONLY. This used to fall back to matching phrases in Output
|
||||
// ("Sandbox block", "Permission denied for ", and so on), which cannot work:
|
||||
// Output is tool-controlled. `bash` preserves arbitrary stdout and stderr on
|
||||
// a StatusError for any nonzero exit, so an ALLOWED command running
|
||||
// `printf 'Sandbox block\n' >&2; exit 1` has actually executed, carries no
|
||||
// DenialReason, and was still classified as refused. The loop then withheld
|
||||
// the schema hint and the profile failure-streak recovery, counted the
|
||||
// executed failure toward the refusal halt, and a later stop told the user a
|
||||
// tool had been refused when it had run.
|
||||
//
|
||||
// read_file returning a document that quotes one of those phrases did the
|
||||
// same thing, which is the more likely way a real session hits it.
|
||||
//
|
||||
// The registry now marks every path that returns BEFORE the tool runs, so
|
||||
// the question is answerable from provenance. Nothing below reads Output.
|
||||
return tools.IsPolicyRefusalResult(tools.Result{Meta: result.Meta})
|
||||
return true
|
||||
}
|
||||
|
||||
// scrubInterceptedOutput mirrors the registry's scrubResultSecrets boundary for
|
||||
@@ -2378,18 +2358,19 @@ func askUserFallbackResult(ctx context.Context, registry *tools.Registry, call T
|
||||
Cwd: options.Cwd,
|
||||
})
|
||||
return ToolResult{
|
||||
ToolCallID: call.ID,
|
||||
Name: call.Name,
|
||||
Status: result.Status,
|
||||
Output: result.ModelOutput(),
|
||||
Truncated: result.Truncated,
|
||||
Meta: result.Meta,
|
||||
Redacted: result.Redacted,
|
||||
ChangedFiles: result.ChangedFiles,
|
||||
FileDiffs: result.FileDiffs,
|
||||
ChangeSummaries: result.ChangeSummaries,
|
||||
Display: result.HumanDisplay(),
|
||||
Outcome: result.Outcome,
|
||||
ToolCallID: call.ID,
|
||||
Name: call.Name,
|
||||
Status: result.Status,
|
||||
Output: result.BaseModelOutput(),
|
||||
Truncated: result.Truncated,
|
||||
Meta: result.Meta,
|
||||
EnforcementNotices: append([]string(nil), result.EnforcementNotices...),
|
||||
Redacted: result.Redacted,
|
||||
ChangedFiles: result.ChangedFiles,
|
||||
FileDiffs: result.FileDiffs,
|
||||
ChangeSummaries: result.ChangeSummaries,
|
||||
Display: result.BaseDisplay(),
|
||||
Outcome: result.Outcome,
|
||||
}
|
||||
}
|
||||
return ToolResult{
|
||||
@@ -2848,7 +2829,7 @@ func deniedPermissionResult(call ToolCall, reason string, requestEvent Permissio
|
||||
Output: "Error: Permission denied for " + call.Name + ": " + reason,
|
||||
DenialReason: denial,
|
||||
Meta: map[string]string{
|
||||
tools.PermissionActionMeta: string(event.Action),
|
||||
"permission_action": string(event.Action),
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -2874,7 +2855,7 @@ func canceledPermissionResult(call ToolCall, reason string, requestEvent Permiss
|
||||
Output: "Error: Permission approval cancelled for " + call.Name + ": " + reason,
|
||||
DenialReason: DenialApprovalCanceled,
|
||||
Meta: map[string]string{
|
||||
tools.PermissionActionMeta: string(event.Action),
|
||||
"permission_action": string(event.Action),
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -3728,33 +3709,6 @@ func toolResultImageMessage(result ToolResult, options Options) (zeroruntime.Mes
|
||||
}, true
|
||||
}
|
||||
|
||||
// precomputedResultFor reports the read-ahead result for call index, and whether
|
||||
// that call actually produced one. A batch entry that never started is still
|
||||
// unstarted as far as the transcript is concerned.
|
||||
//
|
||||
// PRODUCING A RESULT AND ASKING THE RUN TO STOP ARE DIFFERENT FACTS. Treating a
|
||||
// non-nil abortErr as "unstarted" conflated them, and a cancelled permission
|
||||
// request is exactly the case where they disagree: executeToolCall builds the
|
||||
// cancellation result and returns it WITH ErrPermissionApprovalCanceled, so an
|
||||
// earlier sibling reaching a terminal branch discarded a real result and wrote
|
||||
// an aborted placeholder in its place. The transcript then denied that the call
|
||||
// had completed permission handling, and the cancellation was missing from
|
||||
// OnToolResult, the trace counters and the task observation even though the
|
||||
// permission event had already happened.
|
||||
//
|
||||
// The terminal decision is not affected: draining such a sibling only records
|
||||
// what occurred, it cannot reverse a stop already selected.
|
||||
func precomputedResultFor(precomputed []precomputedToolResult, start, end, index int) (ToolResult, bool) {
|
||||
if index < start || index >= end {
|
||||
return ToolResult{}, false
|
||||
}
|
||||
entry := precomputed[index-start]
|
||||
if !entry.completed {
|
||||
return ToolResult{}, false
|
||||
}
|
||||
return entry.result, true
|
||||
}
|
||||
|
||||
func modelAcceptsToolImages(options Options) bool {
|
||||
if options.SupportsVision != nil {
|
||||
return options.SupportsVision(options.Model)
|
||||
|
||||
@@ -685,23 +685,11 @@ func TestIsRetriableToolError(t *testing.T) {
|
||||
{"success", ToolResult{Status: tools.StatusOK}, false},
|
||||
{"bad arguments", ToolResult{Status: tools.StatusError, Output: "Error: Failed to parse arguments for x: bad json"}, true},
|
||||
{"execution failure", ToolResult{Status: tools.StatusError, Output: "Error: read foo.txt: no such file"}, true},
|
||||
// Each refusal carries the provenance its production path attaches: the
|
||||
// disabled-tool path sets a denial category, the registry gates set the
|
||||
// pre-execution marker. Text is never what decides.
|
||||
{"disabled tool", ToolResult{Status: tools.StatusError, Output: `Error: Tool "x" is not enabled for this run.`, DenialReason: DenialFiltered}, false},
|
||||
{"disabled tool", ToolResult{Status: tools.StatusError, Output: `Error: Tool "x" is not enabled for this run.`}, false},
|
||||
{"permission denied (meta)", ToolResult{Status: tools.StatusError, Output: "Error: Permission denied for x: needs approval", Meta: map[string]string{"permission_action": "deny"}}, false},
|
||||
{"permission required", ToolResult{Status: tools.StatusError, Output: "Error: Permission required for x: approve first", Meta: map[string]string{tools.PolicyRefusalMeta: tools.PolicyRefusalPermissionRequired}}, false},
|
||||
{"sandbox block", ToolResult{Status: tools.StatusError, Output: "Error: Sandbox block: outside_workspace", Meta: map[string]string{tools.PolicyRefusalMeta: tools.PolicyRefusalSandboxDenied}}, false},
|
||||
{"sandbox approval", ToolResult{Status: tools.StatusError, Output: "Error: Sandbox approval required for x: network", Meta: map[string]string{tools.PolicyRefusalMeta: tools.PolicyRefusalSandboxApproval}}, false},
|
||||
// THE INVERSE, which is the failure this classifier had. An allowed tool
|
||||
// that RAN, failed, and happened to print one of those phrases must stay
|
||||
// retriable: `bash` preserves arbitrary stdout and stderr on any nonzero
|
||||
// exit, and read_file returns whatever the file says.
|
||||
{"executed failure printing the sandbox phrase", ToolResult{Status: tools.StatusError, Output: "Sandbox block\n"}, true},
|
||||
{"executed failure printing a permission denial", ToolResult{Status: tools.StatusError, Output: "cp: Permission denied for /etc/hosts"}, true},
|
||||
{"executed failure printing a permission requirement", ToolResult{Status: tools.StatusError, Output: "Permission required for sudo"}, true},
|
||||
{"executed failure printing the approval phrase", ToolResult{Status: tools.StatusError, Output: "Sandbox approval required for curl"}, true},
|
||||
{"executed failure quoting the disabled-tool phrase", ToolResult{Status: tools.StatusError, Output: "grep: web_fetch is not enabled for this run"}, true},
|
||||
{"permission required", ToolResult{Status: tools.StatusError, Output: "Error: Permission required for x: approve first"}, false},
|
||||
{"sandbox block", ToolResult{Status: tools.StatusError, Output: "Error: Sandbox block: outside_workspace"}, false},
|
||||
{"sandbox approval", ToolResult{Status: tools.StatusError, Output: "Error: Sandbox approval required for x: network"}, false},
|
||||
// Structured denial categories are authoritative regardless of message text.
|
||||
{"denial: filtered", ToolResult{Status: tools.StatusError, Output: "anything", DenialReason: DenialFiltered}, false},
|
||||
{"denial: permission", ToolResult{Status: tools.StatusError, Output: "anything", DenialReason: DenialPermissionDenied}, false},
|
||||
|
||||
@@ -1,74 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
"github.com/Twigpine/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// mixedOutcomeTool needs permission on every call. The test denies every other
|
||||
// request, and a call that is allowed executes and fails with an error naming
|
||||
// its own item, so no single identity repeats and only the content-blind bound
|
||||
// can stop the run.
|
||||
type mixedOutcomeTool struct{ ran int }
|
||||
|
||||
func (tool *mixedOutcomeTool) Name() string { return "bash" }
|
||||
func (tool *mixedOutcomeTool) Description() string { return "probe" }
|
||||
func (tool *mixedOutcomeTool) Parameters() tools.Schema { return tools.Schema{Type: "object"} }
|
||||
func (tool *mixedOutcomeTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectShell, Permission: tools.PermissionPrompt, Reason: "runs a command"}
|
||||
}
|
||||
|
||||
func (tool *mixedOutcomeTool) Run(context.Context, map[string]any) tools.Result {
|
||||
tool.ran++
|
||||
return tools.Result{Status: tools.StatusError, Output: "Error: item " + strconv.Itoa(tool.ran) + " failed"}
|
||||
}
|
||||
|
||||
// A MIXED HALT READS THE SAME HEADLESS AS INTERACTIVE.
|
||||
//
|
||||
// The stop answer for the content-blind bound over executed failures and
|
||||
// refusals says the tool "failed or was refused". The headless reason for the
|
||||
// same halt said only that it "failed repeatedly", so zero exec and the
|
||||
// interactive answer disagreed about one stop. Reported by @jatmn.
|
||||
func TestAMixedGuardHaltIsReportedAsMixedUnderTheCompletionGate(t *testing.T) {
|
||||
tool := &mixedOutcomeTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
turns := make([][]zeroruntime.StreamEvent, 0, toolFailureAnyErrorStopAt+4)
|
||||
for i := range toolFailureAnyErrorStopAt + 4 {
|
||||
turns = append(turns, toolTurn("c"+strconv.Itoa(i), "bash", `{"item":`+strconv.Itoa(i)+`}`))
|
||||
}
|
||||
asked := 0
|
||||
result, err := Run(context.Background(), "run them", &mockProvider{turns: turns}, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAsk,
|
||||
MaxTurns: len(turns) + 5,
|
||||
RequireCompletionSignal: true,
|
||||
OnPermissionRequest: func(context.Context, PermissionRequest) (PermissionDecision, error) {
|
||||
asked++
|
||||
if asked%2 == 0 {
|
||||
return PermissionDecision{Action: PermissionDecisionAllow}, nil
|
||||
}
|
||||
return PermissionDecision{Action: PermissionDecisionDeny}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tool.ran == 0 || asked <= tool.ran {
|
||||
t.Fatalf("the run did not mix executed failures with refusals (ran %d of %d asked), so this proves nothing", tool.ran, asked)
|
||||
}
|
||||
if !strings.Contains(result.FinalAnswer, "failed or was refused") {
|
||||
t.Fatalf("the halt was not the mixed content-blind bound: %q", result.FinalAnswer)
|
||||
}
|
||||
if !result.Incomplete {
|
||||
t.Fatal("a guard halt under the completion gate must be Incomplete")
|
||||
}
|
||||
if !strings.Contains(result.IncompleteReason, "failed or was refused") || !strings.Contains(result.IncompleteReason, "bash") {
|
||||
t.Errorf("IncompleteReason = %q, want it to say the tool failed or was refused, as the stop answer does", result.IncompleteReason)
|
||||
}
|
||||
}
|
||||
@@ -1,232 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
"github.com/Twigpine/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// readAheadTool is read-only and thread-safe, so a consecutive run of its calls
|
||||
// is executed as one parallel batch before the loop consumes any of it.
|
||||
type readAheadTool struct {
|
||||
mu sync.Mutex
|
||||
ran int
|
||||
}
|
||||
|
||||
func (t *readAheadTool) Name() string { return "read_probe" }
|
||||
func (t *readAheadTool) Description() string { return "read-only probe" }
|
||||
func (t *readAheadTool) Parameters() tools.Schema {
|
||||
return tools.Schema{Type: "object", Properties: map[string]tools.PropertySchema{"id": {Type: "string"}}}
|
||||
}
|
||||
|
||||
func (t *readAheadTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow, Reason: "test"}
|
||||
}
|
||||
|
||||
func (t *readAheadTool) Capabilities() tools.ToolCapabilities {
|
||||
return tools.ToolCapabilities{Effect: tools.EffectReadOnly, ThreadSafe: true}
|
||||
}
|
||||
|
||||
func (t *readAheadTool) Run(_ context.Context, args map[string]any) tools.Result {
|
||||
t.mu.Lock()
|
||||
t.ran++
|
||||
t.mu.Unlock()
|
||||
id, _ := args["id"].(string)
|
||||
// Varying output, so no identity repeats and only the content-blind counter
|
||||
// can halt the run.
|
||||
return tools.Result{Status: tools.StatusError, Output: "failure for " + id}
|
||||
}
|
||||
|
||||
func readCall(id string) []zeroruntime.StreamEvent {
|
||||
return []zeroruntime.StreamEvent{
|
||||
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: id, ToolCallProviderID: "provider-" + id, ToolName: "read_probe"},
|
||||
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: id, ArgumentsFragment: `{"id":"` + id + `"}`},
|
||||
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: id},
|
||||
}
|
||||
}
|
||||
|
||||
// READ-AHEAD MEANS "NOT CONSUMED YET" NO LONGER MEANS "NOT EXECUTED".
|
||||
//
|
||||
// executeParallelReadBatch runs an entire eligible run of read calls before the
|
||||
// loop consumes any of them. Every terminal branch in the consumption loop used
|
||||
// to close out the calls after the current index with aborted placeholders, on
|
||||
// the assumption that they had not run. A sibling that had already executed was
|
||||
// therefore recorded as aborted: its real result was discarded, the transcript
|
||||
// disagreed with what had happened, and its callbacks, trace counter, task
|
||||
// observation and images were lost with it. Where such a sibling is a
|
||||
// successful read, execution may already have committed file-observation credit
|
||||
// for content the model never receives.
|
||||
func TestParallelSiblingsAreFinalizedWhenTheGuardHalts(t *testing.T) {
|
||||
tool := &readAheadTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
|
||||
// Eleven single-call turns bring the content-blind counter to one below its
|
||||
// bound, so the first call of the pair below trips it.
|
||||
var turns [][]zeroruntime.StreamEvent
|
||||
for i := 0; i < toolFailureAnyErrorStopAt-1; i++ {
|
||||
id := "s" + strconv.Itoa(i)
|
||||
turns = append(turns, append(readCall(id), zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}))
|
||||
}
|
||||
// Then one turn advertising two parallel-eligible calls. Both execute in the
|
||||
// batch; the halt fires while consuming the first.
|
||||
pair := append(readCall("pA"), readCall("pB")...)
|
||||
turns = append(turns, append(pair, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}))
|
||||
turns = append(turns, append(readCall("unreached"), zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}))
|
||||
|
||||
var reported []ToolResult
|
||||
var announced int
|
||||
result, err := Run(context.Background(), "read things", &mockProvider{turns: turns}, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAuto,
|
||||
MaxTurns: len(turns) + 5,
|
||||
OnToolCall: func(ToolCall) { announced++ },
|
||||
OnToolResult: func(r ToolResult) { reported = append(reported, r) },
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.HasPrefix(result.FinalAnswer, "Agent stopped:") {
|
||||
t.Fatalf("the guard did not halt this run: %q", result.FinalAnswer)
|
||||
}
|
||||
|
||||
// Every execution is reported exactly once, and the announcements match.
|
||||
if tool.ran != len(reported) {
|
||||
t.Errorf("the tool executed %d times but %d results were reported; a completed call was recorded as aborted",
|
||||
tool.ran, len(reported))
|
||||
}
|
||||
if announced != len(reported) {
|
||||
t.Errorf("OnToolCall fired %d times against %d results; the two callbacks disagree", announced, len(reported))
|
||||
}
|
||||
|
||||
// Strict pairing: one tool result per advertised tool call, and the drained
|
||||
// sibling carries its REAL output rather than the aborted placeholder.
|
||||
byID := map[string]int{}
|
||||
var siblingContent string
|
||||
for _, message := range result.Messages {
|
||||
if message.Role != zeroruntime.MessageRoleTool {
|
||||
continue
|
||||
}
|
||||
byID[message.ToolCallID]++
|
||||
if message.ToolCallID == "pB" {
|
||||
siblingContent = message.Content
|
||||
}
|
||||
// A strict replay pairs each result to its tool_use by the provider
|
||||
// call ID, so the drained sibling needs it as much as the consumed call.
|
||||
if (message.ToolCallID == "pA" || message.ToolCallID == "pB") && message.ToolCallProviderID != "provider-"+message.ToolCallID {
|
||||
t.Errorf("tool result %s carries provider call ID %q, want %q", message.ToolCallID, message.ToolCallProviderID, "provider-"+message.ToolCallID)
|
||||
}
|
||||
}
|
||||
for _, id := range []string{"pA", "pB"} {
|
||||
if byID[id] != 1 {
|
||||
t.Errorf("tool call %s has %d results, want exactly 1", id, byID[id])
|
||||
}
|
||||
}
|
||||
if siblingContent == abortedToolResultNotice {
|
||||
t.Error("the sibling that already ran was recorded as aborted")
|
||||
}
|
||||
if !strings.Contains(siblingContent, "failure for pB") {
|
||||
t.Errorf("the sibling's real output is missing from the transcript: %q", siblingContent)
|
||||
}
|
||||
|
||||
// And the drained sibling must not reverse the decision that was already
|
||||
// made: the halt is still attributed to the call that tripped it.
|
||||
if !strings.Contains(result.FinalAnswer, strconv.Itoa(toolFailureAnyErrorStopAt)) {
|
||||
t.Errorf("the stop answer changed after draining siblings: %q", result.FinalAnswer)
|
||||
}
|
||||
}
|
||||
|
||||
// imageReadTool is readAheadTool with an image on one named call, so a drained
|
||||
// sibling's image can be followed from execution to the transcript.
|
||||
type imageReadTool struct {
|
||||
readAheadTool
|
||||
imageFor string
|
||||
}
|
||||
|
||||
func (t *imageReadTool) Run(ctx context.Context, args map[string]any) tools.Result {
|
||||
result := t.readAheadTool.Run(ctx, args)
|
||||
if id, _ := args["id"].(string); id == t.imageFor {
|
||||
result.Images = []zeroruntime.ImageBlock{{MediaType: "image/png", Data: []byte("\x89PNG")}}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// The drained sibling's image is part of what ran, so it follows the tool
|
||||
// results like any other call's image. Images reach the transcript through the
|
||||
// turn's deferred builder, which the halt path calls after closing out the
|
||||
// siblings; a drain that skipped the sibling's image would lose it silently.
|
||||
func TestParallelSiblingImageSurvivesTheGuardHalt(t *testing.T) {
|
||||
tool := &imageReadTool{imageFor: "pB"}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
|
||||
var turns [][]zeroruntime.StreamEvent
|
||||
for i := 0; i < toolFailureAnyErrorStopAt-1; i++ {
|
||||
id := "s" + strconv.Itoa(i)
|
||||
turns = append(turns, append(readCall(id), zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}))
|
||||
}
|
||||
pair := append(readCall("pA"), readCall("pB")...)
|
||||
turns = append(turns, append(pair, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}))
|
||||
|
||||
result, err := Run(context.Background(), "read things", &mockProvider{turns: turns}, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAuto,
|
||||
MaxTurns: len(turns) + 5,
|
||||
SupportsVision: func(string) bool { return true },
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.HasPrefix(result.FinalAnswer, "Agent stopped:") {
|
||||
t.Fatalf("the guard did not halt this run: %q", result.FinalAnswer)
|
||||
}
|
||||
images := 0
|
||||
for _, message := range result.Messages {
|
||||
if message.Role == zeroruntime.MessageRoleUser {
|
||||
images += len(message.Images)
|
||||
}
|
||||
}
|
||||
if images != 1 {
|
||||
t.Errorf("the drained sibling's image reached the transcript %d times, want once", images)
|
||||
}
|
||||
}
|
||||
|
||||
// A call that genuinely never ran still gets an aborted placeholder, so the fix
|
||||
// did not simply stop aborting anything.
|
||||
func TestUnstartedCallsStillGetAbortedPlaceholders(t *testing.T) {
|
||||
tool := &readAheadTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
|
||||
var turns [][]zeroruntime.StreamEvent
|
||||
for i := 0; i < toolFailureAnyErrorStopAt-1; i++ {
|
||||
id := "s" + strconv.Itoa(i)
|
||||
turns = append(turns, append(readCall(id), zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}))
|
||||
}
|
||||
// One eligible pair, then a THIRD call outside the batch window is not what
|
||||
// happens here: the whole run is eligible, so instead advertise a single call
|
||||
// and let the halt fire on it, leaving nothing precomputed beyond it.
|
||||
turns = append(turns, append(readCall("solo"), zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}))
|
||||
turns = append(turns, append(readCall("unreached"), zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone}))
|
||||
|
||||
result, err := Run(context.Background(), "read things", &mockProvider{turns: turns}, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAuto,
|
||||
MaxTurns: len(turns) + 5,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.HasPrefix(result.FinalAnswer, "Agent stopped:") {
|
||||
t.Fatalf("the guard did not halt: %q", result.FinalAnswer)
|
||||
}
|
||||
if tool.ran != toolFailureAnyErrorStopAt {
|
||||
t.Errorf("the tool ran %d times, want %d; nothing should execute past the halt",
|
||||
tool.ran, toolFailureAnyErrorStopAt)
|
||||
}
|
||||
}
|
||||
@@ -2,7 +2,6 @@ package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
@@ -13,17 +12,8 @@ import (
|
||||
// them one after another serializes pure I/O waits. A consecutive run of
|
||||
// capability-safe read-only calls is executed concurrently instead; results are
|
||||
// then consumed in the original call order, so guard counters, message
|
||||
// ordering, abort semantics, and the surface's call/result event pairing match
|
||||
// sequential execution.
|
||||
//
|
||||
// What can differ is how many calls RAN. A run executes before the guard
|
||||
// consumes its first result, so a halt on an earlier call can leave up to
|
||||
// maxParallelReadTools-1 siblings already executed. The guard bounds count
|
||||
// consumed results, not executions, and the extra executions are reads by the
|
||||
// eligibility rule below; closeOutRemaining records them without feeding the
|
||||
// guard. Reported by @jatmn.
|
||||
//
|
||||
// Runs never span a mutating call: a
|
||||
// ordering, abort semantics, and the surface's call/result event pairing are
|
||||
// byte-identical to sequential execution. Runs never span a mutating call: a
|
||||
// read that follows a write must observe the write, so eligibility is decided
|
||||
// per consecutive run, not per batch.
|
||||
//
|
||||
@@ -46,15 +36,6 @@ const maxParallelReadTools = 8
|
||||
type precomputedToolResult struct {
|
||||
result ToolResult
|
||||
abortErr error
|
||||
// completed records whether this call PRODUCED a result, which is a
|
||||
// different question from whether it asks the enclosing run to stop.
|
||||
//
|
||||
// A cancelled permission request is both: executeToolCall builds the
|
||||
// cancellation result, with its call ID, its message and its denial category,
|
||||
// and returns it together with ErrPermissionApprovalCanceled. Reading the
|
||||
// error as "never started" throws that result away, and the transcript then
|
||||
// denies that permission handling happened at all.
|
||||
completed bool
|
||||
}
|
||||
|
||||
// parallelSafeToolCall reports whether call may run concurrently with its
|
||||
@@ -208,15 +189,7 @@ func executeParallelReadBatch(ctx context.Context, registry *tools.Registry, cal
|
||||
semaphore <- struct{}{}
|
||||
defer func() { <-semaphore }()
|
||||
result, abortErr := executeToolCall(ctx, registry, calls[index], permissionMode, batchOptions)
|
||||
// Recorded HERE, where both halves are in hand, rather than inferred later
|
||||
// from the error. A populated call ID is what every path that produced a
|
||||
// result sets and what a context cancellation before execution leaves
|
||||
// empty.
|
||||
results[index-start] = precomputedToolResult{
|
||||
result: result,
|
||||
abortErr: abortErr,
|
||||
completed: strings.TrimSpace(result.ToolCallID) != "",
|
||||
}
|
||||
results[index-start] = precomputedToolResult{result: result, abortErr: abortErr}
|
||||
}(index)
|
||||
}
|
||||
waitGroup.Wait()
|
||||
|
||||
@@ -1,116 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
"github.com/Twigpine/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// forgedMarkerTool executes, fails, and claims in its metadata that the registry
|
||||
// refused it before it ran.
|
||||
type forgedMarkerTool struct {
|
||||
ran int
|
||||
key string
|
||||
value string
|
||||
}
|
||||
|
||||
func (tool *forgedMarkerTool) Name() string { return "bash" }
|
||||
func (tool *forgedMarkerTool) Description() string { return "test shell tool" }
|
||||
func (tool *forgedMarkerTool) Parameters() tools.Schema {
|
||||
return tools.Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]tools.PropertySchema{"command": {Type: "string"}},
|
||||
Required: []string{"command"},
|
||||
AdditionalProperties: false,
|
||||
}
|
||||
}
|
||||
func (tool *forgedMarkerTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow, Reason: "reads files"}
|
||||
}
|
||||
func (tool *forgedMarkerTool) Run(context.Context, map[string]any) tools.Result {
|
||||
tool.ran++
|
||||
return tools.Result{
|
||||
Status: tools.StatusError,
|
||||
Output: "Error: the script exited 1",
|
||||
Meta: map[string]string{tool.key: tool.value},
|
||||
}
|
||||
}
|
||||
|
||||
// THE FORGED MARKER MUST NOT CHANGE THE RUN.
|
||||
//
|
||||
// Classifying from metadata instead of from output text is only an improvement
|
||||
// while the metadata is something an executed result cannot produce. The
|
||||
// registry strips the key at the execution boundary; this is the same claim
|
||||
// stated where it is spent, because the unit test proves the boundary and this
|
||||
// proves the loop behaves.
|
||||
//
|
||||
// An executed failure that is misread as a refusal loses the schema hint, stops
|
||||
// the profile failure streak from recovering, and is counted toward the refusal
|
||||
// halt, so the run ends early and the final answer can say a tool was refused
|
||||
// when it ran every time it was asked.
|
||||
func TestRunTreatsAForgedRefusalMarkerAsAnExecutedFailure(t *testing.T) {
|
||||
for _, forged := range []struct{ key, value string }{
|
||||
{tools.PolicyRefusalMeta, tools.PolicyRefusalSandboxDenied},
|
||||
{tools.PolicyRefusalMeta, tools.PolicyRefusalPermissionDenied},
|
||||
{tools.PolicyRefusalMeta, tools.PolicyRefusalToolNotEnabled},
|
||||
{tools.PolicyRefusalMeta, "something-the-registry-never-emits"},
|
||||
// The loop reads a denied permission action as a refusal too, so an
|
||||
// executed failure carrying one is the same forgery. Reported by @jatmn.
|
||||
{tools.PermissionActionMeta, string(PermissionActionDeny)},
|
||||
} {
|
||||
value := forged.value
|
||||
t.Run(forged.key+"="+value, func(t *testing.T) {
|
||||
tool := &forgedMarkerTool{key: forged.key, value: value}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
|
||||
// One more call than the hint threshold, so the hint has to have been
|
||||
// injected by the last one, then a final text turn to end the run.
|
||||
calls := toolFailureHintAt + 1
|
||||
turns := make([][]zeroruntime.StreamEvent, 0, calls+1)
|
||||
for i := range calls {
|
||||
turns = append(turns, toolTurn("call-"+strconv.Itoa(i), "bash", `{"command":"./flaky.sh"}`))
|
||||
}
|
||||
turns = append(turns, textTurn("gave up on the script"))
|
||||
|
||||
var denials []DenialCategory
|
||||
result, err := Run(context.Background(), "run the script", &mockProvider{turns: turns}, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAsk,
|
||||
MaxTurns: len(turns) + 5,
|
||||
OnToolResult: func(toolResult ToolResult) {
|
||||
denials = append(denials, toolResult.DenialReason)
|
||||
},
|
||||
OnPermissionRequest: func(context.Context, PermissionRequest) (PermissionDecision, error) {
|
||||
t.Error("permission was requested for an allow-safety tool; this test must exercise the executed-failure path")
|
||||
return PermissionDecision{Action: PermissionDecisionDeny}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tool.ran != calls {
|
||||
t.Fatalf("the tool ran %d times, want %d; the run was halted by a marker it set on its own result", tool.ran, calls)
|
||||
}
|
||||
for index, denial := range denials {
|
||||
if denial != DenialNone {
|
||||
t.Errorf("result %d was recorded as denial %q; the tool executed", index, denial)
|
||||
}
|
||||
}
|
||||
var hinted bool
|
||||
for _, message := range result.Messages {
|
||||
if strings.Contains(message.Content, toolFailureHintMarker) {
|
||||
hinted = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hinted {
|
||||
t.Errorf("no retry hint was injected for an executed failure carrying %q, so the model got no correction", value)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,196 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strconv"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
"github.com/Twigpine/zero/internal/trace"
|
||||
"github.com/Twigpine/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// THE UNCATEGORIZED REFUSALS ARE THE ONES THIS CHANGE EXISTS FOR, SO THEY ARE THE
|
||||
// ONES THAT MUST BE DRIVEN THROUGH Run.
|
||||
//
|
||||
// A categorized denial is easy to test and was never the gap. The gap is a
|
||||
// refusal that arrives with DenialReason empty, because a category is attached
|
||||
// only where a TYPED denial is built:
|
||||
//
|
||||
// - a headless run leaves OnPermissionRequest nil, so the loop never reaches
|
||||
// its typed-denial branch and the registry gate returns a bare
|
||||
// `Permission required for ...`
|
||||
// - a sandbox preflight denial on a non-shell tool loses its SandboxDecision
|
||||
// converting to ToolResult and arrives as a bare `Sandbox block`
|
||||
//
|
||||
// Both were counted as SUCCESS, which cleared the very record they were meant to
|
||||
// accumulate, and the refused call repeated to MaxTurns.
|
||||
//
|
||||
// Asserting that through the helper alone would prove nothing: the first version
|
||||
// of this fix passed every helper test while being a no-op in production,
|
||||
// because the loop asked a different question than the tests did. So each case
|
||||
// below goes through Run and pins the three things the loop actually does with
|
||||
// the classification: halt at the bound, never execute the tool, and withhold
|
||||
// the profile's one-shot failure escalation.
|
||||
|
||||
// headlessPromptTool needs approval and never gets it, because the run has no
|
||||
// approver. The registry refuses it before Run() is reached, exactly as a
|
||||
// headless run does, so nothing here fakes the refusal.
|
||||
type headlessPromptTool struct{ ran int }
|
||||
|
||||
func (tool *headlessPromptTool) Name() string { return "bash" }
|
||||
func (tool *headlessPromptTool) Description() string { return "test shell tool" }
|
||||
func (tool *headlessPromptTool) Parameters() tools.Schema {
|
||||
return tools.Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]tools.PropertySchema{"command": {Type: "string"}},
|
||||
Required: []string{"command"},
|
||||
AdditionalProperties: false,
|
||||
}
|
||||
}
|
||||
func (tool *headlessPromptTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectShell, Permission: tools.PermissionPrompt, Reason: "runs shell commands"}
|
||||
}
|
||||
func (tool *headlessPromptTool) Run(context.Context, map[string]any) tools.Result {
|
||||
tool.ran++
|
||||
return tools.Result{Status: tools.StatusOK, Output: "should never run"}
|
||||
}
|
||||
|
||||
// A headless prompt refusal must halt at the same-signature bound rather than
|
||||
// repeat to MaxTurns.
|
||||
//
|
||||
// It used to key on the refusal TEXT, which the registry happens to hold
|
||||
// constant here, and the stop then described a repeated failure. The registry
|
||||
// marks this path, so the category is now derived at the ToolResult boundary and
|
||||
// the streak keys on that instead. The bound is unchanged; the wording is more
|
||||
// honest, because a headless prompt refusal is a refusal and not a tool that
|
||||
// keeps failing.
|
||||
func TestRunStopsAnUncategorizedHeadlessRefusalAtTheFailureBound(t *testing.T) {
|
||||
tool := &headlessPromptTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
|
||||
// Comfortably more turns than the bound, so reaching the bound is what stops
|
||||
// the run rather than exhausting the provider or MaxTurns.
|
||||
turns := make([][]zeroruntime.StreamEvent, 0, toolFailureStopAt+4)
|
||||
for i := range toolFailureStopAt + 4 {
|
||||
turns = append(turns, toolTurn("call-"+strconv.Itoa(i), "bash",
|
||||
`{"command":"touch /etc/file`+strconv.Itoa(i)+`"}`))
|
||||
}
|
||||
provider := &mockProvider{turns: turns}
|
||||
|
||||
recorder := trace.NewRecorder("policy-refusal-session", "run-1", "fast")
|
||||
result, err := Run(context.Background(), "do the thing", provider, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAsk,
|
||||
MaxTurns: len(turns) + 5,
|
||||
Trace: recorder,
|
||||
// Armed well below the halt bound, so a refusal counted as a retriable
|
||||
// failure would visibly spend it.
|
||||
Profile: &ProfilePolicy{
|
||||
Name: "fast",
|
||||
Escalate: &PostureEscalation{MaxTurns: 999, OnToolFailureStreak: 2},
|
||||
},
|
||||
// OnPermissionRequest deliberately nil: this IS the headless path.
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if len(provider.requests) != toolFailureStopAt {
|
||||
t.Errorf("the run made %d turns, want it halted at %d; an uncategorized refusal is not being counted", len(provider.requests), toolFailureStopAt)
|
||||
}
|
||||
if tool.ran != 0 {
|
||||
t.Errorf("the refused tool executed %d times; the registry gate must precede execution", tool.ran)
|
||||
}
|
||||
// Categorized now, so the stop answer says refused. Previously the loop could
|
||||
// not tell and reported a repeated failure.
|
||||
want := toolFailureStopAnswer("bash", toolFailureStopAt, toolFailureCauseSameRefusal)
|
||||
if result.FinalAnswer != want {
|
||||
t.Errorf("final answer =\n %q\nwant\n %q", result.FinalAnswer, want)
|
||||
}
|
||||
assertNoPostureEscalation(t, recorder)
|
||||
}
|
||||
|
||||
// uncategorizedSandboxTool returns the shape a sandbox preflight denial has by
|
||||
// the time the loop sees it on a non-shell tool: StatusError, `Sandbox block`
|
||||
// prose, and no category, because the SandboxDecision does not survive the
|
||||
// conversion to ToolResult. The path varies per call exactly as the real message
|
||||
// does, since it names what was refused.
|
||||
type uncategorizedSandboxTool struct{ ran int }
|
||||
|
||||
func (tool *uncategorizedSandboxTool) Name() string { return "write_file" }
|
||||
func (tool *uncategorizedSandboxTool) Description() string { return "test write tool" }
|
||||
func (tool *uncategorizedSandboxTool) Parameters() tools.Schema {
|
||||
return tools.Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]tools.PropertySchema{"path": {Type: "string"}},
|
||||
Required: []string{"path"},
|
||||
AdditionalProperties: false,
|
||||
}
|
||||
}
|
||||
func (tool *uncategorizedSandboxTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectWrite, Permission: tools.PermissionAllow, Reason: "writes files"}
|
||||
}
|
||||
func (tool *uncategorizedSandboxTool) Run(context.Context, map[string]any) tools.Result {
|
||||
tool.ran++
|
||||
return tools.Result{
|
||||
Status: tools.StatusError,
|
||||
Output: "Sandbox block: write to /out-" + strconv.Itoa(tool.ran) + ".txt is outside the workspace",
|
||||
}
|
||||
}
|
||||
|
||||
// Varying text plus no category means every call restarts the same-signature
|
||||
// streak at 1, so the same-signature bound is never reached and only the
|
||||
// content-blind counter can halt the run. Before the fix nothing halted it at
|
||||
// all: each refusal was read as a success, deleted the record, and the call
|
||||
// repeated until MaxTurns.
|
||||
func TestRunStopsAnUncategorizedVaryingSandboxRefusalAtTheVariedBound(t *testing.T) {
|
||||
tool := &uncategorizedSandboxTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
|
||||
turns := make([][]zeroruntime.StreamEvent, 0, toolFailureAnyErrorStopAt+4)
|
||||
for i := range toolFailureAnyErrorStopAt + 4 {
|
||||
turns = append(turns, toolTurn("call-"+strconv.Itoa(i), "write_file",
|
||||
`{"path":"/out-`+strconv.Itoa(i)+`.txt"}`))
|
||||
}
|
||||
provider := &mockProvider{turns: turns}
|
||||
|
||||
recorder := trace.NewRecorder("policy-refusal-session", "run-2", "fast")
|
||||
result, err := Run(context.Background(), "write the files", provider, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAsk,
|
||||
MaxTurns: len(turns) + 5,
|
||||
Trace: recorder,
|
||||
Profile: &ProfilePolicy{
|
||||
Name: "fast",
|
||||
Escalate: &PostureEscalation{MaxTurns: 999, OnToolFailureStreak: 2},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if tool.ran != toolFailureAnyErrorStopAt {
|
||||
t.Errorf("the tool was called %d times, want the run halted at %d", tool.ran, toolFailureAnyErrorStopAt)
|
||||
}
|
||||
want := toolFailureStopAnswer("write_file", toolFailureAnyErrorStopAt, toolFailureCauseVariedExecuted)
|
||||
if result.FinalAnswer != want {
|
||||
t.Errorf("final answer =\n %q\nwant\n %q", result.FinalAnswer, want)
|
||||
}
|
||||
assertNoPostureEscalation(t, recorder)
|
||||
}
|
||||
|
||||
// The profile's failure-streak trigger restores turn budget and reasoning effort
|
||||
// on the theory that a struggling tool needs room. A refusal is not a struggling
|
||||
// tool, it is an answer, and spending the one-shot escalation on one buys turns
|
||||
// that will be refused identically. The loop keeps this contract by passing an
|
||||
// empty outcome to the controller for a refusal, which is only observable from
|
||||
// outside as the escalation never firing.
|
||||
func assertNoPostureEscalation(t *testing.T, recorder *trace.Recorder) {
|
||||
t.Helper()
|
||||
if got := recorder.Finish().Counter(trace.CounterPostureEscalations); got != 0 {
|
||||
t.Errorf("posture_escalations = %d, want 0; a policy refusal spent the one-shot failure escalation", got)
|
||||
}
|
||||
}
|
||||
@@ -1,265 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
"github.com/Twigpine/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// A REFUSAL IS A FAILURE, AND THE CLASSIFIER MUST SAY SO BEFORE IT READS ANYTHING
|
||||
// ELSE.
|
||||
//
|
||||
// isPolicyRefusal decides on denial category, then permission metadata, then
|
||||
// output text. None of those questions is meaningful about a call the tool
|
||||
// completed, and the last one is answered by content the model does not control:
|
||||
// a `bash` that greps the sandbox docs, a read_file returning this very source
|
||||
// file, any successful command whose output happens to carry one of the phrases.
|
||||
//
|
||||
// isRetriableToolError gated on StatusError before calling in, so the boundary
|
||||
// held while that was the only caller. Extracting the helper and calling it from
|
||||
// the counting path dropped the gate, and every successful result was put to a
|
||||
// question that only failures can answer.
|
||||
//
|
||||
// Each case below is a SUCCESS. None may be read as a refusal.
|
||||
func TestPolicyRefusalNeverClassifiesASuccessfulResult(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
result ToolResult
|
||||
// structured marks a result carrying real provenance: a denial category,
|
||||
// permission metadata, or the registry's pre-execution marker. Only those
|
||||
// may be read as a refusal when they fail. The rest carry nothing but
|
||||
// text, which the tool controls, so they must stay ordinary failures no
|
||||
// matter which phrase they happen to contain.
|
||||
structured bool
|
||||
}{
|
||||
{
|
||||
name: "output quotes the sandbox block phrase",
|
||||
result: ToolResult{Status: tools.StatusOK, Output: "grep: docs/sandbox.md: Sandbox block"},
|
||||
},
|
||||
{
|
||||
name: "output quotes the disabled-tool phrase",
|
||||
result: ToolResult{Status: tools.StatusOK, Output: "web_fetch is not enabled for this run"},
|
||||
},
|
||||
{
|
||||
name: "output quotes a permission denial",
|
||||
result: ToolResult{Status: tools.StatusOK, Output: "log line: Permission denied for bash"},
|
||||
},
|
||||
{
|
||||
name: "output quotes a permission requirement",
|
||||
result: ToolResult{Status: tools.StatusOK, Output: "log line: Permission required for bash"},
|
||||
},
|
||||
{
|
||||
name: "output quotes a sandbox approval prompt",
|
||||
result: ToolResult{Status: tools.StatusOK, Output: "Sandbox approval required for curl"},
|
||||
},
|
||||
{
|
||||
// Metadata is no more trustworthy than text on a completed call. A
|
||||
// result carrying a category with an OK status is not a denial; every
|
||||
// path that attaches one fails the call.
|
||||
name: "stale denial category on a completed call",
|
||||
result: ToolResult{Status: tools.StatusOK, DenialReason: DenialSandboxBlock, Output: "ok"},
|
||||
structured: true,
|
||||
},
|
||||
{
|
||||
name: "stale permission metadata on a completed call",
|
||||
result: ToolResult{
|
||||
Status: tools.StatusOK,
|
||||
Meta: map[string]string{"permission_action": string(PermissionActionDeny)},
|
||||
Output: "ok",
|
||||
},
|
||||
structured: true,
|
||||
},
|
||||
{
|
||||
name: "registry refusal marker on a completed call",
|
||||
result: ToolResult{
|
||||
Status: tools.StatusOK,
|
||||
Meta: map[string]string{tools.PolicyRefusalMeta: tools.PolicyRefusalSandboxDenied},
|
||||
Output: "ok",
|
||||
},
|
||||
structured: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, testCase := range cases {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
if isPolicyRefusal(testCase.result) {
|
||||
t.Errorf("a successful result was classified as a policy refusal, so a healthy call counts toward the denial streak: %+v", testCase.result)
|
||||
}
|
||||
failed := testCase.result
|
||||
failed.Status = tools.StatusError
|
||||
switch {
|
||||
case testCase.structured && !isPolicyRefusal(failed):
|
||||
t.Errorf("a failure carrying real provenance was not classified as a refusal, so the gate broke the classifier: %+v", failed)
|
||||
case !testCase.structured && isPolicyRefusal(failed):
|
||||
t.Errorf("a failure carrying only text was classified as a refusal; output is tool-controlled and cannot decide this: %+v", failed)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// blockTextShellTool runs freely and prints text the classifier keys on. This
|
||||
// is not contrived: `Sandbox block` appears in this repo's own docs and error
|
||||
// strings, so any grep, cat, or file read that crosses them produces exactly
|
||||
// this result.
|
||||
type blockTextShellTool struct{ ran int }
|
||||
|
||||
func (tool *blockTextShellTool) Name() string { return "bash" }
|
||||
func (tool *blockTextShellTool) Description() string { return "test shell tool" }
|
||||
func (tool *blockTextShellTool) Parameters() tools.Schema {
|
||||
return tools.Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]tools.PropertySchema{"command": {Type: "string"}},
|
||||
Required: []string{"command"},
|
||||
AdditionalProperties: false,
|
||||
}
|
||||
}
|
||||
func (tool *blockTextShellTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow, Reason: "reads files"}
|
||||
}
|
||||
func (tool *blockTextShellTool) Run(context.Context, map[string]any) tools.Result {
|
||||
tool.ran++
|
||||
return tools.Result{
|
||||
Status: tools.StatusOK,
|
||||
Output: "docs/windows-sandbox.md:41: Sandbox block (write jail) is reported as ...",
|
||||
}
|
||||
}
|
||||
|
||||
// The same gap driven through Run, because the classifier being wrong only
|
||||
// matters if the loop acts on it, and it does: the counting path calls
|
||||
// isPolicyRefusal directly, and a true answer records a failure against the
|
||||
// tool's signature. Six identical successes then look like six identical
|
||||
// denials, and the run ends on a stop answer telling the user a tool it never
|
||||
// refused was refused six times.
|
||||
//
|
||||
// The tool succeeds every turn and the run is given more turns than the bound,
|
||||
// so anything that halts before the provider's final turn is the guard.
|
||||
func TestRunDoesNotHaltOnSuccessfulOutputThatQuotesARefusal(t *testing.T) {
|
||||
tool := &blockTextShellTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
|
||||
calls := toolFailureStopAt + 4
|
||||
turns := make([][]zeroruntime.StreamEvent, 0, calls+1)
|
||||
for i := range calls {
|
||||
turns = append(turns, toolTurn("call-"+strconv.Itoa(i), "bash",
|
||||
`{"command":"grep -n 'Sandbox block' docs/windows-sandbox.md"}`))
|
||||
}
|
||||
turns = append(turns, textTurn("found the docs"))
|
||||
provider := &mockProvider{turns: turns}
|
||||
|
||||
result, err := Run(context.Background(), "search the docs", provider, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAsk,
|
||||
MaxTurns: len(turns) + 5,
|
||||
OnPermissionRequest: func(context.Context, PermissionRequest) (PermissionDecision, error) {
|
||||
t.Error("the loop asked permission for an allow-safety tool; this test must exercise the success path")
|
||||
return PermissionDecision{Action: PermissionDecisionDeny}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if tool.ran != calls {
|
||||
t.Errorf("the tool ran %d times, want %d; the run was halted mid-way by its own success output", tool.ran, calls)
|
||||
}
|
||||
if result.FinalAnswer != "found the docs" {
|
||||
t.Errorf("final answer =\n %q\nwant\n %q\n(a stop answer here means successful output was counted as a refusal)",
|
||||
result.FinalAnswer, "found the docs")
|
||||
}
|
||||
}
|
||||
|
||||
// failingPhraseTool RAN and failed. Its stderr happens to contain a phrase the
|
||||
// classifier used to key on, which is the everyday case: `bash` preserves
|
||||
// arbitrary stdout and stderr on any nonzero exit, so
|
||||
// `printf 'Sandbox block\n' >&2; exit 1` produces exactly this.
|
||||
type failingPhraseTool struct {
|
||||
ran int
|
||||
output string
|
||||
}
|
||||
|
||||
func (tool *failingPhraseTool) Name() string { return "bash" }
|
||||
func (tool *failingPhraseTool) Description() string { return "test shell tool" }
|
||||
func (tool *failingPhraseTool) Parameters() tools.Schema {
|
||||
return tools.Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]tools.PropertySchema{"command": {Type: "string"}},
|
||||
Required: []string{"command"},
|
||||
AdditionalProperties: false,
|
||||
}
|
||||
}
|
||||
func (tool *failingPhraseTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow, Reason: "reads files"}
|
||||
}
|
||||
func (tool *failingPhraseTool) Run(context.Context, map[string]any) tools.Result {
|
||||
tool.ran++
|
||||
return tools.Result{Status: tools.StatusError, Output: tool.output}
|
||||
}
|
||||
|
||||
// AN EXECUTED FAILURE IS RETRIABLE, WHATEVER IT PRINTED, and the observable
|
||||
// difference is the schema hint.
|
||||
//
|
||||
// isRetriableToolError feeds `hintable` into the guard, and the guard injects
|
||||
// the retry hint at toolFailureHintAt. Classifying an executed failure as a
|
||||
// policy refusal makes hintable false, so the model is told nothing and simply
|
||||
// repeats the call: the run loses the one correction it had. That is on top of
|
||||
// the failure being accumulated as a denial.
|
||||
//
|
||||
// Each phrase below is one the classifier used to match on. The tool is
|
||||
// allow-safety and runs every time, so nothing here is refused by anything.
|
||||
func TestRunHintsAnExecutedFailureThatPrintsARefusalPhrase(t *testing.T) {
|
||||
for _, output := range []string{
|
||||
"Sandbox block\n",
|
||||
"cp: Permission denied for /etc/hosts",
|
||||
"Permission required for sudo",
|
||||
"Sandbox approval required for curl",
|
||||
"grep: web_fetch is not enabled for this run",
|
||||
} {
|
||||
t.Run(output, func(t *testing.T) {
|
||||
tool := &failingPhraseTool{output: output}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
|
||||
// One more call than the hint threshold, so the hint has to have been
|
||||
// injected by the last one, then a final text turn to end the run.
|
||||
calls := toolFailureHintAt + 1
|
||||
turns := make([][]zeroruntime.StreamEvent, 0, calls+1)
|
||||
for i := range calls {
|
||||
turns = append(turns, toolTurn("call-"+strconv.Itoa(i), "bash", `{"command":"./flaky.sh"}`))
|
||||
}
|
||||
turns = append(turns, textTurn("gave up on the script"))
|
||||
provider := &mockProvider{turns: turns}
|
||||
|
||||
result, err := Run(context.Background(), "run the script", provider, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAsk,
|
||||
MaxTurns: len(turns) + 5,
|
||||
OnPermissionRequest: func(context.Context, PermissionRequest) (PermissionDecision, error) {
|
||||
t.Error("permission was requested for an allow-safety tool; this test must exercise the executed-failure path")
|
||||
return PermissionDecision{Action: PermissionDecisionDeny}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tool.ran != calls {
|
||||
t.Fatalf("the tool ran %d times, want %d; the run was halted by output it merely printed", tool.ran, calls)
|
||||
}
|
||||
|
||||
var hinted bool
|
||||
for _, message := range result.Messages {
|
||||
if strings.Contains(message.Content, toolFailureHintMarker) {
|
||||
hinted = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hinted {
|
||||
t.Errorf("no retry hint was injected for an executed failure printing %q; it was classified as a policy refusal, so the model got no correction", output)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,84 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// UNCATEGORIZED REFUSALS MUST STILL COUNT.
|
||||
//
|
||||
// The halt guard asked DenialReason != DenialNone, but a category is only
|
||||
// attached where a TYPED denial is built. A headless run leaves
|
||||
// OnPermissionRequest nil, so the loop never reaches that branch and the
|
||||
// registry returns a bare "Error: Permission required ..." with no category. A
|
||||
// sandbox preflight denial on a non-shell tool loses its SandboxDecision in the
|
||||
// conversion and arrives as an uncategorized "Sandbox block".
|
||||
//
|
||||
// Both were counted as SUCCESS, which cleared the very record the guard
|
||||
// accumulates, so the same refused call could repeat to MaxTurns. That is the
|
||||
// loop this PR exists to stop.
|
||||
func TestUncategorizedPolicyRefusalsAreCounted(t *testing.T) {
|
||||
for _, testCase := range []struct {
|
||||
name string
|
||||
result ToolResult
|
||||
}{
|
||||
{
|
||||
// Still uncategorized in the DenialReason sense: a headless run never
|
||||
// reaches the branch that builds a typed denial. It carries the
|
||||
// registry's pre-execution marker instead, which is what makes it
|
||||
// classifiable without reading the message.
|
||||
name: "headless prompt refusal carries no denial category",
|
||||
result: ToolResult{
|
||||
Status: tools.StatusError,
|
||||
Output: `Error: Permission required for bash: The tool is marked "prompt" and was not executed.`,
|
||||
Meta: map[string]string{tools.PolicyRefusalMeta: tools.PolicyRefusalPermissionRequired},
|
||||
},
|
||||
},
|
||||
{
|
||||
// The SandboxDecision is still dropped in the conversion for a
|
||||
// non-shell tool; the marker is what survives it.
|
||||
name: "sandbox preflight denial loses its decision in conversion",
|
||||
result: ToolResult{
|
||||
Status: tools.StatusError,
|
||||
Output: "Sandbox block: write outside the workspace",
|
||||
Meta: map[string]string{tools.PolicyRefusalMeta: tools.PolicyRefusalSandboxDenied},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "an explicitly denied permission action",
|
||||
result: ToolResult{Status: tools.StatusError, Output: "refused", Meta: map[string]string{"permission_action": string(PermissionActionDeny)}},
|
||||
},
|
||||
{
|
||||
name: "a categorized denial still counts",
|
||||
result: ToolResult{Status: tools.StatusError, Output: "denied", DenialReason: DenialFiltered},
|
||||
},
|
||||
} {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
if !isPolicyRefusal(testCase.result) {
|
||||
t.Fatal("refusal was treated as a success, so the guard's record is cleared and the call can repeat to MaxTurns")
|
||||
}
|
||||
if isRetriableToolError(testCase.result) {
|
||||
t.Error("a policy refusal is not retriable: repeating it verbatim cannot help")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// An ordinary failure is still retriable and still not a refusal, or the guard
|
||||
// would halt on things the model can legitimately fix by trying again.
|
||||
func TestOrdinaryFailuresAreNotPolicyRefusals(t *testing.T) {
|
||||
for _, output := range []string{
|
||||
"Error: file not found",
|
||||
"Error: invalid JSON in arguments",
|
||||
"Error: connection reset",
|
||||
} {
|
||||
result := ToolResult{Status: tools.StatusError, Output: output}
|
||||
if isPolicyRefusal(result) {
|
||||
t.Errorf("%q was classified as a policy refusal", output)
|
||||
}
|
||||
if !isRetriableToolError(result) {
|
||||
t.Errorf("%q should stay retriable", output)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,87 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// PRODUCING A RESULT AND ASKING THE RUN TO STOP ARE DIFFERENT FACTS.
|
||||
//
|
||||
// closeOutRemaining decides between finalizing a completed sibling and writing
|
||||
// an aborted placeholder, and it used to make that decision from abortErr. A
|
||||
// cancelled permission request is exactly where the two facts disagree:
|
||||
// executeToolCall builds the cancellation result, with its call ID, its message
|
||||
// and its denial category, and returns it TOGETHER with
|
||||
// ErrPermissionApprovalCanceled. Reading the error as "never started" discards a
|
||||
// real result and records the opposite of what happened, and the cancellation
|
||||
// then goes missing from OnToolResult, the trace counters and the task
|
||||
// observation even though the permission event already fired.
|
||||
func TestABatchEntryWithAResultAndAnAbortErrorIsDrained(t *testing.T) {
|
||||
canceled := canceledPermissionResult(
|
||||
ToolCall{ID: "call-b", Name: "read_probe"},
|
||||
"cancelled in TUI",
|
||||
PermissionEvent{ToolName: "read_probe"},
|
||||
)
|
||||
precomputed := []precomputedToolResult{
|
||||
{result: ToolResult{ToolCallID: "call-a", Name: "read_probe", Status: tools.StatusOK}, completed: true},
|
||||
{
|
||||
result: canceled,
|
||||
abortErr: fmt.Errorf("%w for read_probe", ErrPermissionApprovalCanceled),
|
||||
completed: true,
|
||||
},
|
||||
}
|
||||
|
||||
sibling, ran := precomputedResultFor(precomputed, 0, 2, 1)
|
||||
if !ran {
|
||||
t.Fatal("a cancelled permission result was treated as a call that never started")
|
||||
}
|
||||
if sibling.ToolCallID != "call-b" {
|
||||
t.Errorf("ToolCallID = %q, want the cancellation result rather than an empty placeholder", sibling.ToolCallID)
|
||||
}
|
||||
if sibling.DenialReason != DenialApprovalCanceled {
|
||||
t.Errorf("DenialReason = %q, want the cancellation preserved", sibling.DenialReason)
|
||||
}
|
||||
}
|
||||
|
||||
// An entry that never produced a result is still unstarted, or the fix would
|
||||
// have stopped aborting anything.
|
||||
func TestABatchEntryWithNoResultIsStillUnstarted(t *testing.T) {
|
||||
precomputed := []precomputedToolResult{
|
||||
{abortErr: errors.New("cancelled before the tool ran")},
|
||||
}
|
||||
if _, ran := precomputedResultFor(precomputed, 0, 1, 0); ran {
|
||||
t.Error("an entry that produced nothing was finalized as if it had run")
|
||||
}
|
||||
}
|
||||
|
||||
// A call outside the batch window never started either.
|
||||
func TestACallOutsideTheBatchWindowIsUnstarted(t *testing.T) {
|
||||
precomputed := []precomputedToolResult{{result: ToolResult{ToolCallID: "call-a"}, completed: true}}
|
||||
if _, ran := precomputedResultFor(precomputed, 0, 1, 1); ran {
|
||||
t.Error("a call past the batch window was reported as executed")
|
||||
}
|
||||
}
|
||||
|
||||
// And the flag is recorded where the truth is known rather than inferred later,
|
||||
// so a real batch marks what it ran.
|
||||
func TestTheBatchRecordsWhatItRan(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(&readAheadTool{})
|
||||
calls := []ToolCall{
|
||||
{ID: "call-a", Name: "read_probe", Arguments: `{"id":"a"}`},
|
||||
{ID: "call-b", Name: "read_probe", Arguments: `{"id":"b"}`},
|
||||
}
|
||||
results := executeParallelReadBatch(context.Background(), registry, calls, 0, 2, PermissionModeAuto, Options{Registry: registry})
|
||||
for index, entry := range results {
|
||||
if !entry.completed {
|
||||
t.Errorf("entry %d executed but was not recorded as completed: %#v", index, entry)
|
||||
}
|
||||
if entry.result.ToolCallID != calls[index].ID {
|
||||
t.Errorf("entry %d is keyed to %q, want %q", index, entry.result.ToolCallID, calls[index].ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,176 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
"github.com/Twigpine/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// legacyDenialSignaturePrefix is the string provenance used to be encoded in.
|
||||
// The identity was one namespace holding both a normalized error signature and
|
||||
// a synthetic "denial:<category>" key, and provenance was recovered afterwards
|
||||
// with strings.HasPrefix. It survives here, and only here, so the regression
|
||||
// names the exact spelling that used to collide.
|
||||
const legacyDenialSignaturePrefix = "denial:"
|
||||
|
||||
func runUntilGuardHalt(t *testing.T, gate bool) Result {
|
||||
t.Helper()
|
||||
tool := &uncategorizedSandboxTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
turns := make([][]zeroruntime.StreamEvent, 0, toolFailureAnyErrorStopAt+4)
|
||||
for i := range toolFailureAnyErrorStopAt + 4 {
|
||||
turns = append(turns, toolTurn("c"+strconv.Itoa(i), "write_file", `{"path":"/o-`+strconv.Itoa(i)+`.txt"}`))
|
||||
}
|
||||
result, err := Run(context.Background(), "write", &mockProvider{turns: turns}, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAsk,
|
||||
MaxTurns: len(turns) + 5,
|
||||
RequireCompletionSignal: gate,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.HasPrefix(result.FinalAnswer, "Agent stopped:") {
|
||||
t.Fatalf("the guard did not halt this run: %q", result.FinalAnswer)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// A GUARD HALT IS UNFINISHED WORK, AND HEADLESS HAS TO HEAR ABOUT IT.
|
||||
//
|
||||
// The halt returns straight out of the tool loop, so it never reached the
|
||||
// completion gate the max-turns paths go through. zero exec treats only
|
||||
// Result.Incomplete as exit 4 and otherwise reports run_end success with exit
|
||||
// 0, so a task denied six times, or failed through the twelve-call bound, came
|
||||
// back as a successful automation result having done none of the work asked.
|
||||
func TestGuardHaltIsIncompleteUnderTheCompletionGate(t *testing.T) {
|
||||
headless := runUntilGuardHalt(t, true)
|
||||
if !headless.Incomplete {
|
||||
t.Error("a guard halt reports the run as complete; zero exec would call a run that did nothing a success")
|
||||
}
|
||||
if !strings.Contains(headless.IncompleteReason, "write_file") {
|
||||
t.Errorf("the incomplete reason does not name the tool that halted the run: %q", headless.IncompleteReason)
|
||||
}
|
||||
// And the interactive default is untouched: Incomplete exists for the
|
||||
// headless gate, so setting it without one would be inventing a status.
|
||||
if interactive := runUntilGuardHalt(t, false); interactive.Incomplete {
|
||||
t.Error("Incomplete was set without RequireCompletionSignal, changing interactive behaviour")
|
||||
}
|
||||
}
|
||||
|
||||
// executedDenialLookalikeTool runs, fails, and prints exactly the string the
|
||||
// guard used to store for a real permission refusal.
|
||||
type executedDenialLookalikeTool struct{ ran int }
|
||||
|
||||
func (t *executedDenialLookalikeTool) Name() string { return "bash" }
|
||||
func (t *executedDenialLookalikeTool) Description() string { return "probe" }
|
||||
func (t *executedDenialLookalikeTool) Parameters() tools.Schema { return tools.Schema{Type: "object"} }
|
||||
func (t *executedDenialLookalikeTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectShell, Permission: tools.PermissionAllow}
|
||||
}
|
||||
|
||||
func (t *executedDenialLookalikeTool) Run(context.Context, map[string]any) tools.Result {
|
||||
t.ran++
|
||||
return tools.Result{Status: tools.StatusError, Output: legacyDenialSignaturePrefix + string(DenialPermissionDenied)}
|
||||
}
|
||||
|
||||
// TRUSTED PROVENANCE AND UNTRUSTED CONTENT MUST NOT SHARE A NAMESPACE.
|
||||
//
|
||||
// While the identity was a single string, a command that printed exactly
|
||||
// "denial:permission_denied" and exited non-zero took on the identity of a real
|
||||
// permission refusal: it merged streaks with one, and the run told the user the
|
||||
// tool had been refused when it had executed every time.
|
||||
func TestAnExecutedErrorCannotImpersonateARefusal(t *testing.T) {
|
||||
tool := &executedDenialLookalikeTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
turns := make([][]zeroruntime.StreamEvent, 0, toolFailureStopAt+4)
|
||||
for i := range toolFailureStopAt + 4 {
|
||||
turns = append(turns, toolTurn("c"+strconv.Itoa(i), "bash", `{}`))
|
||||
}
|
||||
result, err := Run(context.Background(), "run it", &mockProvider{turns: turns}, Options{
|
||||
Registry: registry, PermissionMode: PermissionModeAsk, MaxTurns: len(turns) + 5,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tool.ran == 0 {
|
||||
t.Fatal("the tool never ran, so this no longer covers an executed failure")
|
||||
}
|
||||
if strings.Contains(result.FinalAnswer, "was refused") {
|
||||
t.Errorf("the tool executed %d times and nothing refused it, but the run says it was refused: %q", tool.ran, result.FinalAnswer)
|
||||
}
|
||||
if !strings.Contains(result.FinalAnswer, "with the same error") {
|
||||
t.Errorf("an executed failure repeating identically should read as the same error: %q", result.FinalAnswer)
|
||||
}
|
||||
}
|
||||
|
||||
// AND THE AGGREGATE IS ITS OWN FACT.
|
||||
//
|
||||
// The content-blind bound fires precisely when no single identity repeated
|
||||
// enough, so the identity present at the end says nothing about the eleven
|
||||
// before it. Alternating two refusal categories reaches twelve without either
|
||||
// reaching six; the tool never ran once, and the answer used to describe that
|
||||
// as the tool failing with varying errors.
|
||||
func TestTheContentBlindBoundKeepsRefusalProvenance(t *testing.T) {
|
||||
for _, testCase := range []struct {
|
||||
name string
|
||||
kinds []DenialCategory
|
||||
wantCause toolFailureCause
|
||||
wantAnswer string
|
||||
}{
|
||||
{
|
||||
name: "alternating refusal categories",
|
||||
kinds: []DenialCategory{DenialPermissionDenied, DenialSandboxBlock},
|
||||
wantCause: toolFailureCauseVariedRefused,
|
||||
wantAnswer: "was refused",
|
||||
},
|
||||
{
|
||||
// Mixed is DECIDED, not inherited from whichever field a switch
|
||||
// tested first. Neither "failed" alone nor "refused" alone is true of
|
||||
// this sequence, so the wording says both.
|
||||
name: "executed failures mixed with refusals",
|
||||
kinds: []DenialCategory{DenialNone, DenialPermissionDenied},
|
||||
wantCause: toolFailureCauseVariedMixed,
|
||||
wantAnswer: "failed or was refused",
|
||||
},
|
||||
{
|
||||
name: "executed failures only",
|
||||
kinds: []DenialCategory{DenialNone, DenialNone},
|
||||
wantCause: toolFailureCauseVariedExecuted,
|
||||
wantAnswer: "with varying errors",
|
||||
},
|
||||
} {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
state := newGuardState()
|
||||
var outcome toolFailureOutcome
|
||||
for index := 0; index < toolFailureAnyErrorStopAt; index++ {
|
||||
// Distinct prose every call, exactly as a real refusal or error
|
||||
// naming what it touched would produce.
|
||||
outcome = state.observeToolResult("bash", true, false,
|
||||
"failure on item "+strconv.Itoa(index), testCase.kinds[index%len(testCase.kinds)])
|
||||
if outcome.Stop {
|
||||
break
|
||||
}
|
||||
}
|
||||
if !outcome.Stop {
|
||||
t.Fatal("twelve consecutive failures did not trip the content-blind bound")
|
||||
}
|
||||
if outcome.Count != toolFailureAnyErrorStopAt {
|
||||
t.Errorf("Count = %d, want the counter that tripped (%d)", outcome.Count, toolFailureAnyErrorStopAt)
|
||||
}
|
||||
if outcome.Cause != testCase.wantCause {
|
||||
t.Errorf("Cause = %v, want %v", outcome.Cause, testCase.wantCause)
|
||||
}
|
||||
answer := toolFailureStopAnswer("bash", outcome.Count, outcome.Cause)
|
||||
if !strings.Contains(answer, testCase.wantAnswer) {
|
||||
t.Errorf("stop answer = %q, want it to contain %q", answer, testCase.wantAnswer)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,103 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/sandbox"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
"github.com/Twigpine/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// deniedWriteTool records every execution. Its body must never run: the sandbox
|
||||
// refuses the call at the registry boundary, before the tool is reached.
|
||||
type deniedWriteTool struct{ ran int }
|
||||
|
||||
func (t *deniedWriteTool) Name() string { return "write_file" }
|
||||
func (t *deniedWriteTool) Description() string { return "test write tool" }
|
||||
func (t *deniedWriteTool) Parameters() tools.Schema {
|
||||
return tools.Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]tools.PropertySchema{"path": {Type: "string"}},
|
||||
Required: []string{"path"},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *deniedWriteTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectWrite, Permission: tools.PermissionAllow, Reason: "writes files"}
|
||||
}
|
||||
|
||||
func (t *deniedWriteTool) Run(context.Context, map[string]any) tools.Result {
|
||||
t.ran++
|
||||
return tools.Result{Status: tools.StatusOK, Output: "wrote"}
|
||||
}
|
||||
|
||||
// THROUGH THE REAL SANDBOX, NOT A FAKE STANDING WHERE IT WOULD BE.
|
||||
//
|
||||
// The neighbouring uncategorized test deliberately hand-builds the shape a
|
||||
// refusal has when no category survives, which is a fact about the LOOP. This
|
||||
// one covers the producer: a real sandbox.Engine evaluation denies the call,
|
||||
// the registry turns that into a categorized refusal, and the guard keys on the
|
||||
// category. Without it, reverting refusalResult(..., PolicyRefusalSandboxDenied)
|
||||
// in registry.go to a plain errorResult left both the agent and tools suites
|
||||
// green, so nothing depended on the marker at all.
|
||||
//
|
||||
// Every assertion here fails if a different link breaks: ran proves the body was
|
||||
// skipped, the request count proves the CATEGORY bound halted it rather than the
|
||||
// content-blind one, and the wording proves the provenance survived conversion.
|
||||
func TestRunHaltsARealSandboxDenialOnTheCategoryBound(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
forbidden := t.TempDir()
|
||||
policy := sandbox.DefaultPolicy()
|
||||
// Applies whenever the sandbox is enforcing, so this does not depend on the
|
||||
// workspace boundary being reachable from a synthetic request.
|
||||
policy.DenyWrite = []string{forbidden}
|
||||
|
||||
tool := &deniedWriteTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
|
||||
// A different path every call. Note that errorSignature truncates at 80
|
||||
// characters and the message opens with a fixed 40-character prefix followed
|
||||
// by a long temp root, so the signatures collide here anyway: the request
|
||||
// count below pins WHICH bound halted the run, but it does not by itself
|
||||
// prove the streak was held by the category. The wording assertion is the one
|
||||
// that discriminates, and it is the one that fails when the marker is gone.
|
||||
turns := make([][]zeroruntime.StreamEvent, 0, toolFailureStopAt+4)
|
||||
for i := range toolFailureStopAt + 4 {
|
||||
target := filepath.Join(forbidden, "escape-"+strconv.Itoa(i)+".txt")
|
||||
turns = append(turns, toolTurn("c"+strconv.Itoa(i), "write_file", `{"path":`+strconv.Quote(target)+`}`))
|
||||
}
|
||||
provider := &mockProvider{turns: turns}
|
||||
|
||||
result, err := Run(context.Background(), "write the files", provider, Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAsk,
|
||||
MaxTurns: len(turns) + 5,
|
||||
Cwd: workspace,
|
||||
Sandbox: sandbox.NewEngine(sandbox.EngineOptions{
|
||||
WorkspaceRoot: workspace,
|
||||
Policy: policy,
|
||||
}),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if tool.ran != 0 {
|
||||
t.Errorf("the tool body ran %d times; the sandbox must refuse before execution", tool.ran)
|
||||
}
|
||||
if len(provider.requests) != toolFailureStopAt {
|
||||
t.Errorf("the run made %d turns, want the category bound at %d; a categorized refusal whose prose varies must still halt on its category",
|
||||
len(provider.requests), toolFailureStopAt)
|
||||
}
|
||||
if want := toolFailureStopAnswer("write_file", toolFailureStopAt, toolFailureCauseSameRefusal); result.FinalAnswer != want {
|
||||
t.Errorf("final answer =\n %q\nwant\n %q", result.FinalAnswer, want)
|
||||
}
|
||||
if !strings.Contains(result.FinalAnswer, "was refused") {
|
||||
t.Errorf("the refusal provenance did not survive to the answer: %q", result.FinalAnswer)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// CardBody is the card body source: the rich card-only Display.Preview (a
|
||||
// code/diff preview) when present on a successful result, else the text the
|
||||
// model also saw. Error results keep their output so the failure shows.
|
||||
//
|
||||
// UNDECORATED, ALWAYS. The enforcement disclosure is carried separately as typed
|
||||
// notices and rendered by the card as its own furniture, so a body that already
|
||||
// had the notice composed into it drew the warning twice: once in the notice
|
||||
// lines and once at the top of the output. A rich preview never carried it, so
|
||||
// only the no-preview results (every bash and exec card, and every error) were
|
||||
// wrong, which is exactly the shape a preview-only test cannot see.
|
||||
//
|
||||
// One owner for composition: this returns the base text, and whoever presents it
|
||||
// decorates once. Provider-facing text still goes through ModelOutput.
|
||||
func (result ToolResult) CardBody() string {
|
||||
display := result.BaseDisplay()
|
||||
if strings.TrimSpace(display.Preview) != "" && (result.Status != tools.StatusError || result.Outcome.Finalized()) {
|
||||
return display.Preview
|
||||
}
|
||||
return result.BaseModelOutput()
|
||||
}
|
||||
|
||||
// ToolResultSessionPayload is THE serialization of a tool result into a session
|
||||
// event, shared by every writer: the TUI, headless exec, and the ACP agent.
|
||||
//
|
||||
// It lives here, beside ToolResult, because it has three callers and only one of
|
||||
// them is a terminal. The writers used to spell the payload separately, and each
|
||||
// time one of them was left behind the same thing went missing on reload. The
|
||||
// headless writer persisted only the decorated ModelOutput, so a CLI-written
|
||||
// result restored into the TUI arrived with no typed notices and no undecorated
|
||||
// body. The ACP writer persisted the raw Output field, which after the
|
||||
// output/notice split is the UNDECORATED text, so an ACP client that loaded a
|
||||
// session saw every sandboxed result with its disclosure gone, although the
|
||||
// live update for the same call had shown it. All of them write to the same
|
||||
// session store, so a session written by one surface is restored by another.
|
||||
// One owner for the contract means one place where a field can go missing, and a
|
||||
// test against this function covers every writer.
|
||||
//
|
||||
// output is the provider-facing text with the disclosure composed in.
|
||||
// displayPreview is the undecorated card body, written whenever it differs from
|
||||
// output. A result carrying enforcement notices ALWAYS differs, because output
|
||||
// is decorated and the card body is not, so the undecorated body is written
|
||||
// even when it is empty. Readers key on the field being PRESENT rather than
|
||||
// non-empty for exactly that case: a command that printed nothing under an
|
||||
// enforced profile has an empty body and a real notice, and falling back to
|
||||
// output there would restore the decorated text and draw the notice twice.
|
||||
func ToolResultSessionPayload(result ToolResult) map[string]any {
|
||||
output := result.ModelOutput()
|
||||
payload := map[string]any{
|
||||
"toolCallId": result.ToolCallID,
|
||||
"name": result.Name,
|
||||
"status": string(result.Status),
|
||||
"output": output,
|
||||
}
|
||||
if preview := result.CardBody(); preview != output {
|
||||
payload["displayPreview"] = preview
|
||||
}
|
||||
if result.Truncated {
|
||||
payload["truncated"] = true
|
||||
}
|
||||
if result.Redacted {
|
||||
payload["redacted"] = true
|
||||
}
|
||||
if len(result.Meta) > 0 {
|
||||
payload["meta"] = result.Meta
|
||||
}
|
||||
if len(result.EnforcementNotices) > 0 {
|
||||
payload["enforcementNotices"] = result.EnforcementNotices
|
||||
}
|
||||
if len(result.ChangedFiles) > 0 {
|
||||
payload["changedFiles"] = result.ChangedFiles
|
||||
}
|
||||
if len(result.ChangeSummaries) > 0 {
|
||||
payload["changeSummaries"] = result.ChangeSummaries
|
||||
}
|
||||
return payload
|
||||
}
|
||||
@@ -1,73 +0,0 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// THE ANSWER MUST CLAIM ONLY WHAT THE COUNTER ESTABLISHED.
|
||||
//
|
||||
// anyErrorCount establishes that the tool failed consecutively without a
|
||||
// success. It does NOT establish that every failure differed: five A, five B and
|
||||
// two C reach the content-blind bound of 12 without any signature repeating six
|
||||
// times, and three of those errors were shared. The old wording called that
|
||||
// "each with a different error".
|
||||
//
|
||||
// The signature bound has the inverse problem. A denial streak keys on the
|
||||
// CATEGORY, and the prose behind it embeds the path or command refused, so it
|
||||
// differs on every call. Calling that "the same error" is false in the other
|
||||
// direction.
|
||||
func TestStopAnswerDoesNotOverclaimTheFailurePattern(t *testing.T) {
|
||||
varied := toolFailureStopAnswer("bash", 12, toolFailureCauseVariedExecuted)
|
||||
if strings.Contains(varied, "each with a different error") {
|
||||
t.Errorf("the content-blind bound does not track that every failure differed: %q", varied)
|
||||
}
|
||||
if !strings.Contains(varied, "varying errors") {
|
||||
t.Errorf("the varied stop should still say the errors varied: %q", varied)
|
||||
}
|
||||
|
||||
refused := toolFailureStopAnswer("bash", 6, toolFailureCauseSameRefusal)
|
||||
if strings.Contains(refused, "same error") {
|
||||
t.Errorf("a denial streak covers refusals of different paths, so it is not the same error: %q", refused)
|
||||
}
|
||||
if !strings.Contains(refused, "refused") {
|
||||
t.Errorf("a denial streak should say it was refused: %q", refused)
|
||||
}
|
||||
|
||||
// The one claim that IS justified: a signature streak really did repeat the
|
||||
// same error signature, so that wording stays.
|
||||
same := toolFailureStopAnswer("bash", 6, toolFailureCauseSameError)
|
||||
if !strings.Contains(same, "same error") {
|
||||
t.Errorf("a signature streak may still be described as the same error: %q", same)
|
||||
}
|
||||
}
|
||||
|
||||
// The mixed-signature case jatmn asked for, driven through the real counter
|
||||
// rather than asserted about the wording in isolation: a run whose errors vary
|
||||
// must trip the content-blind bound and must not be described as all-different.
|
||||
func TestMixedSignatureStreakTripsTheContentBlindBound(t *testing.T) {
|
||||
var state guardState
|
||||
var outcome toolFailureOutcome
|
||||
// Five of one error, five of another, two of a third: twelve failures, no
|
||||
// signature repeating six times in a row.
|
||||
for index, output := range []string{
|
||||
"Error: A", "Error: A", "Error: A", "Error: A", "Error: A",
|
||||
"Error: B", "Error: B", "Error: B", "Error: B", "Error: B",
|
||||
"Error: C", "Error: C",
|
||||
} {
|
||||
outcome = state.observeToolResult("bash", true, true, output, DenialNone)
|
||||
if outcome.Stop && index < 11 {
|
||||
t.Fatalf("halted early at failure %d", index+1)
|
||||
}
|
||||
}
|
||||
if !outcome.Stop {
|
||||
t.Fatal("twelve consecutive failures did not trip the content-blind bound")
|
||||
}
|
||||
if outcome.Cause != toolFailureCauseVariedExecuted {
|
||||
t.Errorf("a mixed-signature streak of executed failures should report varied-executed, got %v", outcome.Cause)
|
||||
}
|
||||
answer := toolFailureStopAnswer("bash", outcome.Count, outcome.Cause)
|
||||
if strings.Contains(answer, "each with a different error") {
|
||||
t.Errorf("three of these errors were shared, so they were not each different: %q", answer)
|
||||
}
|
||||
}
|
||||
@@ -1,20 +1,14 @@
|
||||
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
|
||||
@@ -23,56 +17,17 @@ var isolatedUserDirsRoot string
|
||||
// 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") {
|
||||
|
||||
+28
-6
@@ -79,6 +79,10 @@ type ToolResult struct {
|
||||
// The full result may be recoverable through Meta["spill_path"].
|
||||
Truncated bool
|
||||
Meta map[string]string
|
||||
// EnforcementNotices mirrors tools.Result.EnforcementNotices so the
|
||||
// disclosure survives the conversion into the agent-facing result and
|
||||
// reaches the model, the transcript and the interactive display.
|
||||
EnforcementNotices []string
|
||||
// Images the tool produced, delivered to the model as a following user
|
||||
// message rather than on this result. See tools.Result.Images.
|
||||
Images []zeroruntime.ImageBlock
|
||||
@@ -110,24 +114,42 @@ type ToolResult struct {
|
||||
RequestedModel string
|
||||
}
|
||||
|
||||
// ModelOutput returns the bounded provider-facing result while preserving
|
||||
// compatibility with synthetic and restored results created before outcomes
|
||||
// were finalized.
|
||||
func (result ToolResult) ModelOutput() string {
|
||||
// BaseModelOutput is the bounded provider-facing result WITHOUT the enforcement
|
||||
// disclosure composed into it, mirroring tools.Result.BaseModelOutput.
|
||||
//
|
||||
// A surface that renders the typed EnforcementNotices itself must build its body
|
||||
// from here, or the disclosure appears twice. Decoration has exactly one owner
|
||||
// per surface: either the text carries it or the surface draws it, never both.
|
||||
func (result ToolResult) BaseModelOutput() string {
|
||||
if result.Outcome.Finalized() {
|
||||
return result.Outcome.ModelView
|
||||
}
|
||||
return result.Output
|
||||
}
|
||||
|
||||
// HumanDisplay returns the presentation intended for interactive surfaces.
|
||||
func (result ToolResult) HumanDisplay() tools.Display {
|
||||
// BaseDisplay is BaseModelOutput's presentation half, and carries no enforcement
|
||||
// notices for the same reason.
|
||||
func (result ToolResult) BaseDisplay() tools.Display {
|
||||
if result.Outcome.Finalized() {
|
||||
return result.Outcome.HumanView
|
||||
}
|
||||
return result.Display
|
||||
}
|
||||
|
||||
// ModelOutput returns the bounded provider-facing result while preserving
|
||||
// compatibility with synthetic and restored results created before outcomes
|
||||
// were finalized.
|
||||
func (result ToolResult) ModelOutput() string {
|
||||
return tools.WithEnforcementNotices(result.BaseModelOutput(), result.EnforcementNotices)
|
||||
}
|
||||
|
||||
// HumanDisplay returns the presentation intended for interactive surfaces.
|
||||
func (result ToolResult) HumanDisplay() tools.Display {
|
||||
display := result.BaseDisplay()
|
||||
display.Summary = tools.WithEnforcementNotices(display.Summary, result.EnforcementNotices)
|
||||
return display
|
||||
}
|
||||
|
||||
// DenialCategory classifies why a tool call was blocked before it executed.
|
||||
type DenialCategory string
|
||||
|
||||
|
||||
@@ -1,57 +0,0 @@
|
||||
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(""))
|
||||
}
|
||||
@@ -1,54 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
+19
-25
@@ -874,6 +874,19 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
|
||||
}
|
||||
fmt.Fprintf(stderr, "warning: MCP server %s unavailable, skipped: %s\n", skipped.Name, redaction.ErrorMessage(skipped.Err, redaction.Options{}))
|
||||
}
|
||||
// AND WHAT THE SERVERS THAT DID START RAN UNDER. A stdio MCP server prepared
|
||||
// with a weakened write jail serves the whole session from that process, so
|
||||
// the disclosure is about startup and no later tool result can carry it. Said
|
||||
// once, here, next to the skip warnings, rather than pasted onto every
|
||||
// response the server produces. Network servers launch no local process and
|
||||
// report nothing, which is why the optional background registration is not
|
||||
// asked: its only member is the built-in HTTP default, which starts no local
|
||||
// process. A stdio default would need this statement from that path too.
|
||||
// NOT deferred: stderr here is the bare terminal, and the TUI takes it over at
|
||||
// deps.runTUI below. Delivery stops before that hand-off, so a late launch can
|
||||
// never write raw text into the alt screen; see stopMCPDisclosures's call site.
|
||||
guardedStderr, stopMCPDisclosures := reportMCPStartupDisclosures(stderr, mcpRuntime)
|
||||
stderr = guardedStderr
|
||||
// Make local plugins live: register their declared tools into the registry and
|
||||
// collect their hooks + skill roots for the dispatcher and skill tool below.
|
||||
// Done after specialist + MCP registration so plugin tools are part of the
|
||||
@@ -964,12 +977,15 @@ 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.
|
||||
startupNotices := []string{sandboxEngine.DegradedNotice()}
|
||||
modelRPM := zeroruntime.NewModelRPMLimiter(resolved.ModelRPM)
|
||||
// The terminal stops being ours on the next line. Stop and join the disclosure
|
||||
// pump first: anything already queued is printed here, on this goroutine, and a
|
||||
// launch that resolves later is dropped rather than written raw over the TUI.
|
||||
stopMCPDisclosures()
|
||||
return deps.runTUI(context.Background(), tui.Options{
|
||||
ModelRPM: modelRPM,
|
||||
Cwd: workspaceRoot,
|
||||
@@ -1222,9 +1238,6 @@ 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
|
||||
@@ -1236,26 +1249,12 @@ 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
|
||||
}
|
||||
lifecycleHooks := &specialist.LifecycleHooks{}
|
||||
executor := specialist.Executor{Paths: paths, LifecycleHooks: lifecycleHooks}
|
||||
executor := specialist.Executor{Paths: paths}
|
||||
runtime, err := specialist.RegisterTools(registry, executor)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -1274,12 +1273,7 @@ func registerSpecialistTools(registry *tools.Registry, workspaceRoot string, max
|
||||
return nil, err
|
||||
}
|
||||
swarm.RegisterTools(registry, sw)
|
||||
return &agentToolRuntime{
|
||||
specialist: runtime,
|
||||
swarm: sw,
|
||||
specialists: specialistSummaries(paths),
|
||||
lifecycleHooks: lifecycleHooks,
|
||||
}, nil
|
||||
return &agentToolRuntime{specialist: runtime, swarm: sw, specialists: specialistSummaries(paths)}, nil
|
||||
}
|
||||
|
||||
// specialistSummaries loads the available specialists (built-ins + user/project
|
||||
|
||||
@@ -92,25 +92,13 @@ 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))
|
||||
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")
|
||||
}
|
||||
lines = append(lines, fmt.Sprintf(" hash: %s -> %s", shortHash(result.PreviousHash), shortHash(result.Hash)))
|
||||
} 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
|
||||
}
|
||||
|
||||
@@ -7,8 +7,6 @@ import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/skills"
|
||||
)
|
||||
|
||||
func writeSourceSkillDir(t *testing.T, dir string, content string) string {
|
||||
@@ -393,133 +391,3 @@ 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())
|
||||
}
|
||||
}
|
||||
+39
-21
@@ -365,6 +365,18 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
|
||||
return writeExecProviderError(stdout, stderr, options.outputFormat, "mcp_error", err.Error())
|
||||
}
|
||||
defer closeMCPRuntime(stderr, mcpRuntime)
|
||||
// Said HERE, before --list-tools and before the first result, because both
|
||||
// return early and the process this describes is already running by now. On
|
||||
// stderr, so text, JSON and stream-JSON framing on stdout are untouched:
|
||||
// this is the same channel the skipped-server and trust notices use.
|
||||
// Deferred AFTER closeMCPRuntime was deferred, so it runs BEFORE it: the
|
||||
// pump stops and joins while stderr is still ours, and only then are the
|
||||
// clients closed.
|
||||
// Adopt the guarded writer for the rest of startup: the pump is live from
|
||||
// here until stop, and everything below writes to this same stderr.
|
||||
guardedStderr, stopDisclosures := reportMCPStartupDisclosures(stderr, mcpRuntime)
|
||||
stderr = guardedStderr
|
||||
defer stopDisclosures()
|
||||
}
|
||||
pluginActivation = activatePlugins(workspaceRoot, registry, deps, stderr, trustRoot, executionRunner)
|
||||
registerLocalControlTools(registry, workspaceRoot, resolved.LocalControl)
|
||||
@@ -655,7 +667,6 @@ 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,
|
||||
@@ -713,7 +724,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, registry, call); ok {
|
||||
if checkpoint, ok := sessionRecorder.captureCheckpoint(workspaceRoot, call); ok {
|
||||
writer.checkpoint(checkpoint)
|
||||
}
|
||||
},
|
||||
@@ -723,25 +734,7 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
|
||||
},
|
||||
OnToolResult: func(result agent.ToolResult) {
|
||||
writer.toolResult(result)
|
||||
payload := map[string]any{
|
||||
"toolCallId": result.ToolCallID,
|
||||
"name": result.Name,
|
||||
"status": string(result.Status),
|
||||
"output": result.Output,
|
||||
}
|
||||
if len(result.Meta) > 0 {
|
||||
payload["meta"] = result.Meta
|
||||
}
|
||||
if result.Truncated {
|
||||
payload["truncated"] = true
|
||||
}
|
||||
if result.Redacted {
|
||||
payload["redacted"] = true
|
||||
}
|
||||
if len(result.ChangedFiles) > 0 {
|
||||
payload["changedFiles"] = result.ChangedFiles
|
||||
}
|
||||
sessionRecorder.append(sessions.EventToolResult, payload)
|
||||
sessionRecorder.append(sessions.EventToolResult, persistedToolResultPayload(result))
|
||||
},
|
||||
OnUsage: func(u agent.Usage) {
|
||||
writer.usage(u)
|
||||
@@ -1505,3 +1498,28 @@ func writeTraceSnapshot(snapshot *trace.TurnTrace, dest string, stderr io.Writer
|
||||
defer file.Close()
|
||||
return trace.WriteNDJSON(file, snapshot)
|
||||
}
|
||||
|
||||
// persistedToolResultPayload renders one tool result for the durable session
|
||||
// log.
|
||||
//
|
||||
// IT USES THE ACCESSOR, NOT THE RAW FIELD. agent.ToolResult stores the
|
||||
// undecorated model text alongside the typed enforcement notices, and
|
||||
// ModelOutput is what composes the two. Replay reads this payload straight back
|
||||
// into the transcript without reconstructing a ToolResult, so a disclosure that
|
||||
// is not rendered here is simply absent from resumed and compacted context even
|
||||
// though it was visible during the original run.
|
||||
//
|
||||
// Both headless writers go through this, because they previously spelled the
|
||||
// same payload separately and had already drifted: one persisted the raw field
|
||||
// while the stream writer used the accessor.
|
||||
func persistedToolResultPayload(result agent.ToolResult) map[string]any {
|
||||
// ONE CONTRACT WITH THE TUI. This used to build its own payload with the
|
||||
// decorated output only, and both writers append to the same default
|
||||
// session store the TUI resumes from. A CLI-written result restored into
|
||||
// the TUI therefore arrived without typed enforcement notices and without
|
||||
// the undecorated card body, so a long collapsed result rendered no body
|
||||
// and, with it, no disclosure. internal/agent owns the shape now, beside
|
||||
// ToolResult, and every writer calls the same function rather than keeping
|
||||
// a matching copy of it.
|
||||
return agent.ToolResultSessionPayload(result)
|
||||
}
|
||||
@@ -1,206 +0,0 @@
|
||||
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
|
||||
}
|
||||
@@ -1,90 +0,0 @@
|
||||
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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/agent"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// THE HEADLESS WRITER PERSISTS THE SAME SHAPE THE TUI DOES.
|
||||
//
|
||||
// The tui-side restore test proves the shared payload restores a disclosure
|
||||
// exactly once; this proves the CLI actually WRITES that shared payload rather
|
||||
// than its own. The two used to be spelled separately and the headless one had
|
||||
// already drifted to decorated output only. Reverting the delegation leaves
|
||||
// the tui test green and fails this one, which is the point of having both.
|
||||
func TestHeadlessPayloadCarriesTypedNoticesAndUndecoratedBody(t *testing.T) {
|
||||
const notice = "denyRead is configured, so the Windows sandbox uses the token shape without WRITE_RESTRICTED (#869)"
|
||||
result := agent.ToolResult{
|
||||
ToolCallID: "call-cli",
|
||||
Name: "bash",
|
||||
Status: tools.StatusOK,
|
||||
Output: "PROBE-BODY",
|
||||
Truncated: true,
|
||||
EnforcementNotices: []string{notice},
|
||||
}
|
||||
payload := persistedToolResultPayload(result)
|
||||
|
||||
notices, _ := payload["enforcementNotices"].([]string)
|
||||
if len(notices) != 1 || notices[0] != notice {
|
||||
t.Errorf("headless payload does not carry the typed notice: %#v", payload["enforcementNotices"])
|
||||
}
|
||||
preview, _ := payload["displayPreview"].(string)
|
||||
if preview != "PROBE-BODY" {
|
||||
t.Errorf("headless payload does not carry the undecorated body as displayPreview: %q", preview)
|
||||
}
|
||||
output, _ := payload["output"].(string)
|
||||
if !strings.Contains(output, notice) || !strings.Contains(output, "PROBE-BODY") {
|
||||
t.Errorf("provider-facing output is no longer the decorated text: %q", output)
|
||||
}
|
||||
// The one field the headless writer added on its own must survive the
|
||||
// delegation, or a truncation marker silently stops being persisted.
|
||||
if truncated, _ := payload["truncated"].(bool); !truncated {
|
||||
t.Errorf("truncated flag was lost in the shared payload: %#v", payload["truncated"])
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
@@ -18,11 +19,15 @@ 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, registry *tools.Registry, call agent.ToolCall) (sessions.Event, bool) {
|
||||
func (recorder *execSessionRecorder) captureCheckpoint(workspaceRoot string, call agent.ToolCall) (sessions.Event, bool) {
|
||||
if recorder.prepared.Store == nil || recorder.prepared.Session.SessionID == "" {
|
||||
return sessions.Event{}, false
|
||||
}
|
||||
targets := agent.CheckpointTargets(workspaceRoot, registry, call)
|
||||
var args map[string]any
|
||||
if call.Arguments != "" {
|
||||
_ = json.Unmarshal([]byte(call.Arguments), &args)
|
||||
}
|
||||
targets := tools.MutationTargets(workspaceRoot, call.Name, args)
|
||||
if len(targets) == 0 {
|
||||
return sessions.Event{}, false
|
||||
}
|
||||
|
||||
@@ -148,7 +148,7 @@ func runExecSpecDraft(run execSpecDraftRun) int {
|
||||
"name": call.Name,
|
||||
"arguments": call.Arguments,
|
||||
})
|
||||
if checkpoint, ok := sessionRecorder.captureCheckpoint(run.workspaceRoot, run.registry, call); ok {
|
||||
if checkpoint, ok := sessionRecorder.captureCheckpoint(run.workspaceRoot, call); ok {
|
||||
writer.checkpoint(checkpoint)
|
||||
}
|
||||
},
|
||||
@@ -161,22 +161,7 @@ func runExecSpecDraft(run execSpecDraftRun) int {
|
||||
if info, ok := execSpecDraftInfoFromToolResult(result); ok {
|
||||
draftInfo = info
|
||||
}
|
||||
payload := map[string]any{
|
||||
"toolCallId": result.ToolCallID,
|
||||
"name": result.Name,
|
||||
"status": string(result.Status),
|
||||
"output": result.Output,
|
||||
}
|
||||
if len(result.Meta) > 0 {
|
||||
payload["meta"] = result.Meta
|
||||
}
|
||||
if result.Redacted {
|
||||
payload["redacted"] = true
|
||||
}
|
||||
if len(result.ChangedFiles) > 0 {
|
||||
payload["changedFiles"] = result.ChangedFiles
|
||||
}
|
||||
sessionRecorder.append(sessions.EventToolResult, payload)
|
||||
sessionRecorder.append(sessions.EventToolResult, persistedToolResultPayload(result))
|
||||
},
|
||||
OnUsage: func(u agent.Usage) {
|
||||
writer.usage(u)
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/config"
|
||||
"github.com/Twigpine/zero/internal/mcp"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// disclosingExecRuntime is an MCP runtime that launched a process under reduced
|
||||
// enforcement.
|
||||
type disclosingExecRuntime struct {
|
||||
noopMCPRuntime
|
||||
disclosures []mcp.StartupDisclosure
|
||||
}
|
||||
|
||||
func (r disclosingExecRuntime) StartupDisclosures() []mcp.StartupDisclosure { return r.disclosures }
|
||||
|
||||
const execDisclosureNotice = "denyRead is configured, so the write jail is not confining writes"
|
||||
|
||||
// isolateConfigDirs points every config/cache root at test-owned storage.
|
||||
//
|
||||
// Without it these tests build a sandbox engine against the developer's REAL
|
||||
// config dir, trigger the one-time grant migration there, and the migration
|
||||
// notice then turns up on a LATER test's stderr, failing whichever test happens
|
||||
// to assert an empty one. The failure moves between runs, which is what makes it
|
||||
// look like flakiness rather than contamination.
|
||||
func isolateConfigDirs(t *testing.T) {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
t.Setenv("HOME", dir)
|
||||
t.Setenv("USERPROFILE", dir)
|
||||
t.Setenv("APPDATA", dir)
|
||||
t.Setenv("LOCALAPPDATA", dir)
|
||||
t.Setenv("XDG_CONFIG_HOME", dir)
|
||||
t.Setenv("XDG_CACHE_HOME", dir)
|
||||
}
|
||||
|
||||
func execDisclosureDeps(cwd string) appDeps {
|
||||
return appDeps{
|
||||
getwd: func() (string, error) { return cwd, nil },
|
||||
resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) { return execResolvedConfig(), nil },
|
||||
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) {
|
||||
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"docs": {Type: "stdio", Command: "docs-mcp"},
|
||||
}}, nil
|
||||
},
|
||||
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
|
||||
return disclosingExecRuntime{disclosures: []mcp.StartupDisclosure{
|
||||
{Name: "docs", Notices: []string{execDisclosureNotice}},
|
||||
}}, nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// A HEADLESS RUN IS A DISCLOSURE SURFACE TOO.
|
||||
//
|
||||
// `zero exec` registers workspace MCP servers through the same sandbox-backed
|
||||
// runner interactive startup uses, so a stdio server here can launch under the
|
||||
// weakened token and serve the whole run. Reporting the disclosure only from the
|
||||
// TUI meant every text, JSON, stream-JSON and --list-tools caller was told
|
||||
// nothing about the enforcement trade, for a process that was already running.
|
||||
func TestExecReportsMCPStartupDisclosures(t *testing.T) {
|
||||
isolateConfigDirs(t)
|
||||
for _, format := range []string{"", "--output-format=json", "--output-format=stream-json"} {
|
||||
name := format
|
||||
if name == "" {
|
||||
name = "text"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
args := []string{"exec", "--list-tools"}
|
||||
if format != "" {
|
||||
args = append(args, format)
|
||||
}
|
||||
var stdout, stderr bytes.Buffer
|
||||
if code := runWithDeps(args, &stdout, &stderr, execDisclosureDeps(t.TempDir())); code != exitSuccess {
|
||||
t.Fatalf("exit = %d, stderr = %s", code, stderr.String())
|
||||
}
|
||||
if !strings.Contains(stderr.String(), execDisclosureNotice) {
|
||||
t.Errorf("the headless run said nothing about the enforcement trade: %q", stderr.String())
|
||||
}
|
||||
if !strings.Contains(stderr.String(), "docs") {
|
||||
t.Errorf("the report does not name the server: %q", stderr.String())
|
||||
}
|
||||
// The machine-readable surfaces must stay parseable: the disclosure
|
||||
// belongs on stderr precisely so stdout framing is untouched.
|
||||
if format == "--output-format=json" {
|
||||
var any map[string]any
|
||||
if err := json.Unmarshal(bytes.TrimSpace(stdout.Bytes()), &any); err != nil {
|
||||
t.Errorf("stdout is no longer valid JSON: %v (%q)", err, stdout.String())
|
||||
}
|
||||
}
|
||||
if format == "--output-format=stream-json" {
|
||||
for _, line := range strings.Split(strings.TrimSpace(stdout.String()), "\n") {
|
||||
if strings.TrimSpace(line) == "" {
|
||||
continue
|
||||
}
|
||||
var any map[string]any
|
||||
if err := json.Unmarshal([]byte(line), &any); err != nil {
|
||||
t.Errorf("a stream-json line is not valid JSON: %v (%q)", err, line)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// A run whose servers launched nothing says nothing.
|
||||
func TestExecWithoutDisclosuresStaysQuiet(t *testing.T) {
|
||||
isolateConfigDirs(t)
|
||||
deps := execDisclosureDeps(t.TempDir())
|
||||
deps.registerMCPTools = func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
|
||||
return noopMCPRuntime{}, nil
|
||||
}
|
||||
var stdout, stderr bytes.Buffer
|
||||
if code := runWithDeps([]string{"exec", "--list-tools"}, &stdout, &stderr, deps); code != exitSuccess {
|
||||
t.Fatalf("exit = %d, stderr = %s", code, stderr.String())
|
||||
}
|
||||
if strings.Contains(stderr.String(), "reduced enforcement") {
|
||||
t.Errorf("a run with nothing to disclose printed one: %q", stderr.String())
|
||||
}
|
||||
}
|
||||
@@ -2,7 +2,6 @@ package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
@@ -83,105 +82,11 @@ 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`.")
|
||||
@@ -661,31 +566,10 @@ 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]
|
||||
|
||||
@@ -5,7 +5,6 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
@@ -334,234 +333,3 @@ 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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -319,7 +319,7 @@ Flags:
|
||||
--command <cmd> Command to run (required)
|
||||
--name <name> Human-readable hook name
|
||||
--description <text> Hook description
|
||||
--matcher <pattern> Tool/specialist matcher (beforeTool/afterTool/specialistStart/specialistStop)
|
||||
--matcher <pattern> Tool matcher (beforeTool/afterTool only)
|
||||
--arg <value> Command argument (repeatable)
|
||||
--user Write to user config instead of the project
|
||||
--json Print command result as JSON
|
||||
|
||||
@@ -741,13 +741,6 @@ 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
|
||||
@@ -807,29 +800,11 @@ 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
|
||||
@@ -881,19 +856,15 @@ func (cfg *mcpWritableConfig) setServerDisabled(name string, disabled bool) (boo
|
||||
server = map[string]json.RawMessage{}
|
||||
}
|
||||
|
||||
// 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
|
||||
current := false
|
||||
if rawDisabled, ok := server["disabled"]; ok && len(rawDisabled) > 0 && string(rawDisabled) != "null" {
|
||||
if err := json.Unmarshal(rawDisabled, ¤t); err != nil {
|
||||
return false, false, err
|
||||
}
|
||||
}
|
||||
changed := current != disabled
|
||||
if disabled || shippedDisabled {
|
||||
data, err := json.Marshal(disabled)
|
||||
if disabled {
|
||||
data, err := json.Marshal(true)
|
||||
if err != nil {
|
||||
return false, false, err
|
||||
}
|
||||
|
||||
@@ -2,18 +2,10 @@ 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"
|
||||
)
|
||||
|
||||
@@ -52,255 +44,3 @@ 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,163 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Twigpine/zero/internal/config"
|
||||
"github.com/Twigpine/zero/internal/mcp"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// A LAUNCH THAT COMPLETES AFTER THE REPORTER HAS RUN MUST STILL BE SAID, ONCE,
|
||||
// AND ONLY WHILE SOMEONE OWNS THE WRITER.
|
||||
//
|
||||
// Both production paths call reportMCPStartupDisclosures exactly once, right
|
||||
// after RegisterTools returns. A stdio attempt abandoned at the connect timeout
|
||||
// can still be inside cmd.Start at that moment; the process then starts under
|
||||
// the reduced write confinement, the reaper closes its client, and a reporter
|
||||
// that merely SAMPLED the runtime has already come and gone. The retained sink
|
||||
// held the fact and nobody read it again, so the operator saw the skipped
|
||||
// server and never the disclosure.
|
||||
//
|
||||
// This drives the REAL reporter against a REAL runtime, rather than polling
|
||||
// StartupDisclosures by hand, which is what an earlier regression did and which
|
||||
// is precisely how it masked the original bug: a test that re-reads on the
|
||||
// tester's behalf proves nothing about a production path that does not.
|
||||
//
|
||||
// It also never reads stderr while the pump could write. The buffer is examined
|
||||
// only after stop has joined the pump, which is the same discipline both
|
||||
// production callers follow, and is why this passes under -race.
|
||||
func TestLateMCPLaunchReachesTheStartupReporterExactlyOnce(t *testing.T) {
|
||||
const notice = "MCP server started without WRITE_RESTRICTED because denyRead is configured (#869)"
|
||||
released := make(chan struct{})
|
||||
published := make(chan struct{})
|
||||
|
||||
runtime, err := mcp.RegisterTools(context.Background(), tools.NewRegistry(),
|
||||
config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"slow": {Type: "stdio", Command: "slow-mcp"},
|
||||
}},
|
||||
mcp.RegisterOptions{
|
||||
ConnectTimeout: 50 * time.Millisecond,
|
||||
ClientFactory: func(ctx context.Context, server mcp.Server) (mcp.ToolClient, error) {
|
||||
// Held past the registration timeout AND past the settle grace, so
|
||||
// registration has already reaped this attempt and returned.
|
||||
<-released
|
||||
mcp.PublishLaunchForTest(ctx, []string{notice})
|
||||
// Publishing is synchronous into the stream, so by the time this
|
||||
// closes the disclosure is queued and stop cannot race past it.
|
||||
close(published)
|
||||
return nil, errors.New("initialize failed long after start")
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = runtime.Close() })
|
||||
|
||||
// The reporter runs ONCE, here, exactly as runExec and the interactive
|
||||
// startup path run it: before the launch has resolved.
|
||||
var stderr bytes.Buffer
|
||||
_, stop := reportMCPStartupDisclosures(&stderr, runtime)
|
||||
|
||||
close(released)
|
||||
<-published
|
||||
// The owner ends delivery and joins the pump. Every write to the buffer has
|
||||
// happened by the time this returns, so the reads below are unsynchronised
|
||||
// only because there is no longer anything to synchronise with.
|
||||
stop()
|
||||
|
||||
got := stderr.String()
|
||||
if n := strings.Count(got, notice); n != 1 {
|
||||
t.Fatalf("a launch that completed after the reporter ran was disclosed %d time(s), want exactly 1:\n%s", n, got)
|
||||
}
|
||||
if !strings.Contains(got, "MCP server slow started with reduced enforcement") {
|
||||
t.Errorf("the late disclosure does not name the server:\n%s", got)
|
||||
}
|
||||
if skipped := runtime.Skipped(); len(skipped) != 1 {
|
||||
t.Errorf("the server should still be recorded as skipped: %#v", skipped)
|
||||
}
|
||||
}
|
||||
|
||||
// And a server whose launch was already known when the reporter ran is said
|
||||
// once by it, and not again by the late path.
|
||||
func TestKnownMCPLaunchIsNotReportedTwice(t *testing.T) {
|
||||
const notice = "MCP server started under reduced enforcement"
|
||||
runtime, err := mcp.RegisterTools(context.Background(), tools.NewRegistry(),
|
||||
config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"fast": {Type: "stdio", Command: "fast-mcp"},
|
||||
}},
|
||||
mcp.RegisterOptions{
|
||||
ConnectTimeout: time.Second,
|
||||
ClientFactory: func(ctx context.Context, server mcp.Server) (mcp.ToolClient, error) {
|
||||
mcp.PublishLaunchForTest(ctx, []string{notice})
|
||||
return nil, errors.New("initialize failed after start")
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = runtime.Close() })
|
||||
|
||||
var stderr bytes.Buffer
|
||||
_, stop := reportMCPStartupDisclosures(&stderr, runtime)
|
||||
stop()
|
||||
if n := strings.Count(stderr.String(), notice); n != 1 {
|
||||
t.Fatalf("a launch known at registration was disclosed %d time(s), want exactly 1:\n%s", n, stderr.String())
|
||||
}
|
||||
}
|
||||
|
||||
// THE OWNERSHIP BOUNDARY ITSELF: once the caller has stopped delivery, nothing
|
||||
// may write to its writer again.
|
||||
//
|
||||
// This is the property that the retained presentation callback could not hold.
|
||||
// It invoked the CLI's print function from the abandoned connect goroutine
|
||||
// whenever the launch happened to resolve, so a write could land after runExec
|
||||
// had returned or after Bubble Tea had taken the alt screen. The interactive
|
||||
// path stops delivery on the line before it hands over the terminal, and this
|
||||
// pins what that buys: a launch resolving afterwards is dropped, not printed.
|
||||
func TestMCPDisclosureAfterStopIsDroppedNotWritten(t *testing.T) {
|
||||
const notice = "MCP server started under reduced enforcement"
|
||||
released := make(chan struct{})
|
||||
published := make(chan struct{})
|
||||
|
||||
runtime, err := mcp.RegisterTools(context.Background(), tools.NewRegistry(),
|
||||
config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"slow": {Type: "stdio", Command: "slow-mcp"},
|
||||
}},
|
||||
mcp.RegisterOptions{
|
||||
ConnectTimeout: 50 * time.Millisecond,
|
||||
ClientFactory: func(ctx context.Context, server mcp.Server) (mcp.ToolClient, error) {
|
||||
<-released
|
||||
mcp.PublishLaunchForTest(ctx, []string{notice})
|
||||
close(published)
|
||||
return nil, errors.New("initialize failed long after start")
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = runtime.Close() })
|
||||
|
||||
var stderr bytes.Buffer
|
||||
_, stop := reportMCPStartupDisclosures(&stderr, runtime)
|
||||
// The owner gives up the writer BEFORE the launch resolves, which is the
|
||||
// interactive hand-off to the TUI.
|
||||
stop()
|
||||
|
||||
close(released)
|
||||
<-published
|
||||
// Asserting an absence, so the wrong behaviour is given time to appear: once
|
||||
// publishing has returned the disclosure is queued, and a delivery path that
|
||||
// outlived stop would have this long to print it. With delivery ended and the
|
||||
// pump joined there is no writer left, so this window changes nothing.
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
if got := stderr.String(); got != "" {
|
||||
t.Fatalf("a launch that resolved after the owner stopped still wrote to its writer: %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/mcp"
|
||||
)
|
||||
|
||||
type disclosingRuntime struct {
|
||||
noopMCPRuntime
|
||||
disclosures []mcp.StartupDisclosure
|
||||
}
|
||||
|
||||
func (runtime disclosingRuntime) StartupDisclosures() []mcp.StartupDisclosure {
|
||||
return runtime.disclosures
|
||||
}
|
||||
|
||||
// SAID ONCE, WHERE THE USER IS ALREADY BEING TOLD WHAT STARTED.
|
||||
//
|
||||
// The disclosure describes a server PROCESS, which serves the whole session, so
|
||||
// it cannot ride on a tool result and must not be repeated on every one.
|
||||
func TestStartupDisclosuresAreReportedOnce(t *testing.T) {
|
||||
const notice = "denyRead is configured, so the write jail is not confining writes"
|
||||
var stderr bytes.Buffer
|
||||
reportMCPStartupDisclosures(&stderr, disclosingRuntime{
|
||||
disclosures: []mcp.StartupDisclosure{{Name: "docs", Notices: []string{notice}}},
|
||||
})
|
||||
output := stderr.String()
|
||||
if count := strings.Count(output, notice); count != 1 {
|
||||
t.Errorf("the disclosure appears %d times, want exactly 1: %q", count, output)
|
||||
}
|
||||
if !strings.Contains(output, "docs") {
|
||||
t.Errorf("the report does not name the server it is about: %q", output)
|
||||
}
|
||||
}
|
||||
|
||||
// A run with nothing to disclose prints nothing at all.
|
||||
func TestNoStartupDisclosuresPrintNothing(t *testing.T) {
|
||||
var stderr bytes.Buffer
|
||||
reportMCPStartupDisclosures(&stderr, disclosingRuntime{})
|
||||
if stderr.Len() != 0 {
|
||||
t.Errorf("a run with no disclosure wrote %q", stderr.String())
|
||||
}
|
||||
stderr.Reset()
|
||||
// And a runtime that launches nothing is not required to answer.
|
||||
reportMCPStartupDisclosures(&stderr, noopMCPRuntime{})
|
||||
if stderr.Len() != 0 {
|
||||
t.Errorf("a runtime that launches nothing wrote %q", stderr.String())
|
||||
}
|
||||
}
|
||||
@@ -3,9 +3,11 @@ package cli
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/Twigpine/zero/internal/config"
|
||||
"github.com/Twigpine/zero/internal/execution"
|
||||
@@ -190,3 +192,107 @@ func isSensitiveMCPDisplayKey(key string) bool {
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// mcpStartupDisclosing is the optional interface a runtime implements when it
|
||||
// can report what its launched server processes ran under. Optional rather than
|
||||
// part of mcpToolRuntime so a runtime that launches nothing, and every test
|
||||
// double, stays unchanged.
|
||||
type mcpStartupDisclosing interface {
|
||||
StartupDisclosures() []mcp.StartupDisclosure
|
||||
}
|
||||
|
||||
// mcpStartupStreaming is the push form: the runtime queues each disclosure as a
|
||||
// typed event, including a launch that completes after registration returned,
|
||||
// and this package drains it on the goroutine that owns stderr.
|
||||
type mcpStartupStreaming interface {
|
||||
StartupDisclosureStream() *mcp.StartupDisclosureStream
|
||||
}
|
||||
|
||||
// reportMCPStartupDisclosures states once what enforcement applied to the MCP
|
||||
// server processes this run launched.
|
||||
//
|
||||
// A PUSH, NOT A SAMPLE. This used to read StartupDisclosures once, here, and a
|
||||
// stdio attempt abandoned at the connect timeout could still be inside cmd.Start
|
||||
// at that moment. The process then started under the reduced write confinement,
|
||||
// the reaper closed its client, and nothing read the runtime again: the operator
|
||||
// saw the skipped-server warning and never the disclosure.
|
||||
//
|
||||
// THIS GOROUTINE OWNS THE WRITER. The runtime queues typed disclosures; every
|
||||
// write to stderr happens either on the caller's goroutine (the set already known
|
||||
// when this returns, in server order, so startup output keeps its order) or on
|
||||
// the single pump started here, never both at once and never after stop.
|
||||
//
|
||||
// The returned stop ends delivery and joins the pump, so no write to stderr can
|
||||
// outlive the caller's ownership of it. The caller must run it before handing the
|
||||
// terminal to anything else. A disclosure that arrives after stop is dropped: it
|
||||
// is worth printing while someone owns the writer, and worth losing rather than
|
||||
// writing into a screen that now belongs to Bubble Tea. Anything already queued
|
||||
// when stop runs is still printed, on the caller's goroutine, with the pump
|
||||
// already finished.
|
||||
//
|
||||
// ONE WRITER, ONE CALLER AT A TIME. Joining the pump stops writes after its
|
||||
// lifetime but does nothing about the overlap: startup keeps emitting plugin,
|
||||
// trust, peer, provider and validation output to the same writer while the pump
|
||||
// is live. That is unsafe for an ordinary bytes.Buffer and interleaves lines even
|
||||
// on a writer that tolerates concurrent calls. A mutex private to the pump would
|
||||
// not help, because the other writes do not go through it. So the returned writer
|
||||
// is a guarded view of the caller's, and the caller adopts it for the rest of
|
||||
// startup; both sides then take the same lock.
|
||||
//
|
||||
// The pull form is kept for a runtime that implements no stream, which today is
|
||||
// only test doubles; it has no late launches to deliver, so its stop is a no-op
|
||||
// and its writer is handed back unchanged.
|
||||
func reportMCPStartupDisclosures(stderr io.Writer, runtime mcpToolRuntime) (guarded io.Writer, stop func()) {
|
||||
serialized := &serializedWriter{writer: stderr}
|
||||
print := func(disclosure mcp.StartupDisclosure) {
|
||||
for _, notice := range disclosure.Notices {
|
||||
fmt.Fprintf(serialized, "notice: MCP server %s started with reduced enforcement: %s\n", disclosure.Name, notice)
|
||||
}
|
||||
}
|
||||
printAll := func(disclosures []mcp.StartupDisclosure) {
|
||||
for _, disclosure := range disclosures {
|
||||
print(disclosure)
|
||||
}
|
||||
}
|
||||
streaming, ok := runtime.(mcpStartupStreaming)
|
||||
if !ok {
|
||||
if disclosing, ok := runtime.(mcpStartupDisclosing); ok {
|
||||
printAll(disclosing.StartupDisclosures())
|
||||
}
|
||||
return stderr, func() {}
|
||||
}
|
||||
stream := streaming.StartupDisclosureStream()
|
||||
if stream == nil {
|
||||
return stderr, func() {}
|
||||
}
|
||||
printAll(stream.Drain())
|
||||
pumped := make(chan struct{})
|
||||
go func() {
|
||||
defer close(pumped)
|
||||
for stream.Wait() {
|
||||
printAll(stream.Drain())
|
||||
}
|
||||
}()
|
||||
var once sync.Once
|
||||
return serialized, func() {
|
||||
once.Do(func() {
|
||||
stream.Close()
|
||||
<-pumped
|
||||
printAll(stream.Drain())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// serializedWriter gives one underlying writer a single owner at a time, so the
|
||||
// late-disclosure pump and the foreground startup path cannot be inside it
|
||||
// together.
|
||||
type serializedWriter struct {
|
||||
mu sync.Mutex
|
||||
writer io.Writer
|
||||
}
|
||||
|
||||
func (w *serializedWriter) Write(data []byte) (int, error) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
return w.writer.Write(data)
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Twigpine/zero/internal/config"
|
||||
"github.com/Twigpine/zero/internal/mcp"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// blockingWriter makes the overlap deterministic instead of hoping to hit it.
|
||||
//
|
||||
// The first write parks inside Write until the test releases it, which is the
|
||||
// window the pump and the foreground startup path really share: startup keeps
|
||||
// emitting plugin, trust, peer and provider output to the same stderr while a
|
||||
// late MCP disclosure can arrive at any moment. It also records whether it was
|
||||
// ever entered twice at once, which is the property under test.
|
||||
type blockingWriter struct {
|
||||
mu sync.Mutex
|
||||
inside int
|
||||
overlaps int
|
||||
buf bytes.Buffer
|
||||
|
||||
block chan struct{}
|
||||
blockOne sync.Once
|
||||
entered chan struct{}
|
||||
}
|
||||
|
||||
func newBlockingWriter() *blockingWriter {
|
||||
return &blockingWriter{block: make(chan struct{}), entered: make(chan struct{})}
|
||||
}
|
||||
|
||||
func (w *blockingWriter) Write(data []byte) (int, error) {
|
||||
w.mu.Lock()
|
||||
w.inside++
|
||||
if w.inside > 1 {
|
||||
w.overlaps++
|
||||
}
|
||||
w.mu.Unlock()
|
||||
|
||||
// Only the first writer parks, and it announces that it is inside.
|
||||
first := false
|
||||
w.blockOne.Do(func() {
|
||||
first = true
|
||||
close(w.entered)
|
||||
})
|
||||
if first {
|
||||
<-w.block
|
||||
}
|
||||
|
||||
w.mu.Lock()
|
||||
n, err := w.buf.Write(data)
|
||||
w.inside--
|
||||
w.mu.Unlock()
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (w *blockingWriter) overlapCount() int {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
return w.overlaps
|
||||
}
|
||||
|
||||
func (w *blockingWriter) String() string {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
return w.buf.String()
|
||||
}
|
||||
|
||||
// ONE CALLER AT A TIME, FOR THE WHOLE OVERLAP.
|
||||
//
|
||||
// Joining the pump at stop bounds writes to the pump's lifetime but says nothing
|
||||
// about what happens DURING it. The startup paths keep writing to the same
|
||||
// io.Writer the whole time, and the caller may legitimately hand in a plain
|
||||
// bytes.Buffer, which corrupts under concurrent use. A mutex private to the pump
|
||||
// would not have helped, because the foreground writes do not go through it; the
|
||||
// reporter therefore hands back a guarded view of the caller's writer and the
|
||||
// caller adopts it, so both sides take the same lock.
|
||||
func TestLateDisclosureAndForegroundStartupNeverShareTheWriter(t *testing.T) {
|
||||
const notice = "MCP server started under reduced enforcement"
|
||||
released := make(chan struct{})
|
||||
published := make(chan struct{})
|
||||
|
||||
runtime, err := mcp.RegisterTools(context.Background(), tools.NewRegistry(),
|
||||
config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"slow": {Type: "stdio", Command: "slow-mcp"},
|
||||
}},
|
||||
mcp.RegisterOptions{
|
||||
ConnectTimeout: 50 * time.Millisecond,
|
||||
ClientFactory: func(ctx context.Context, server mcp.Server) (mcp.ToolClient, error) {
|
||||
<-released
|
||||
mcp.PublishLaunchForTest(ctx, []string{notice})
|
||||
close(published)
|
||||
return nil, errors.New("initialize failed long after start")
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = runtime.Close() })
|
||||
|
||||
writer := newBlockingWriter()
|
||||
guarded, stop := reportMCPStartupDisclosures(writer, runtime)
|
||||
if guarded == io.Writer(writer) {
|
||||
t.Fatal("SETUP INVALID: the reporter handed back the raw writer, so the caller cannot share its lock")
|
||||
}
|
||||
|
||||
// Foreground startup writes through the guarded writer and parks inside it,
|
||||
// exactly as a slow terminal would.
|
||||
foregroundDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(foregroundDone)
|
||||
fmt.Fprintln(guarded, "warning: MCP server other unavailable, skipped: dial tcp: refused")
|
||||
}()
|
||||
<-writer.entered
|
||||
|
||||
// While the foreground write is parked, the late launch resolves and the pump
|
||||
// tries to print. If the two did not share a lock, this would enter Write
|
||||
// concurrently.
|
||||
close(released)
|
||||
<-published
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
close(writer.block)
|
||||
<-foregroundDone
|
||||
stop()
|
||||
|
||||
if n := writer.overlapCount(); n != 0 {
|
||||
t.Fatalf("the pump and foreground startup were inside the writer together %d time(s)", n)
|
||||
}
|
||||
got := writer.String()
|
||||
if count := strings.Count(got, notice); count != 1 {
|
||||
t.Fatalf("the late disclosure was written %d time(s), want exactly 1 after stop drained it:\n%s", count, got)
|
||||
}
|
||||
if !strings.Contains(got, "unavailable, skipped") {
|
||||
t.Errorf("the foreground startup message was lost:\n%s", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/agent"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// A DISCLOSURE THAT IS NOT PERSISTED DID NOT SURVIVE THE RUN.
|
||||
//
|
||||
// The session log is what a resumed or compacted conversation is rebuilt from,
|
||||
// and replay reads this payload's "output" straight into the transcript without
|
||||
// reconstructing an agent.ToolResult. So persisting the raw undecorated field
|
||||
// makes a warning that was visible during the original run disappear the moment
|
||||
// the session is resumed, with nothing failing anywhere to say so.
|
||||
func TestPersistedToolResultKeepsTheEnforcementNotice(t *testing.T) {
|
||||
const notice = "least-privilege notice: read access was narrowed"
|
||||
payload := persistedToolResultPayload(agent.ToolResult{
|
||||
ToolCallID: "call-1",
|
||||
Name: "bash",
|
||||
Status: tools.StatusOK,
|
||||
Output: "the command output",
|
||||
EnforcementNotices: []string{notice},
|
||||
})
|
||||
|
||||
output, _ := payload["output"].(string)
|
||||
if count := strings.Count(output, notice); count != 1 {
|
||||
t.Errorf("persisted output carries the notice %d times, want exactly 1: %q", count, output)
|
||||
}
|
||||
if !strings.Contains(output, "the command output") {
|
||||
t.Errorf("persisted output lost the command output: %q", output)
|
||||
}
|
||||
}
|
||||
|
||||
// The other fields still round-trip, so the shared helper did not quietly drop
|
||||
// what the two writers used to record separately.
|
||||
func TestPersistedToolResultKeepsItsOtherFields(t *testing.T) {
|
||||
payload := persistedToolResultPayload(agent.ToolResult{
|
||||
ToolCallID: "call-2",
|
||||
Name: "write_file",
|
||||
Status: tools.StatusError,
|
||||
Output: "boom",
|
||||
Meta: map[string]string{"k": "v"},
|
||||
Truncated: true,
|
||||
Redacted: true,
|
||||
ChangedFiles: []string{"a.go"},
|
||||
})
|
||||
for _, field := range []string{"toolCallId", "name", "status", "output", "meta", "truncated", "redacted", "changedFiles"} {
|
||||
if _, ok := payload[field]; !ok {
|
||||
t.Errorf("payload is missing %q: %#v", field, payload)
|
||||
}
|
||||
}
|
||||
if payload["status"] != string(tools.StatusError) {
|
||||
t.Errorf("status = %v, want %q", payload["status"], tools.StatusError)
|
||||
}
|
||||
}
|
||||
|
||||
// An ordinary result records exactly what it did before, so the accessor is not
|
||||
// adding anything where there is nothing to add.
|
||||
func TestPersistedToolResultLeavesAnOrdinaryResultAlone(t *testing.T) {
|
||||
payload := persistedToolResultPayload(agent.ToolResult{
|
||||
ToolCallID: "call-3",
|
||||
Name: "bash",
|
||||
Status: tools.StatusOK,
|
||||
Output: "plain output",
|
||||
})
|
||||
if got := payload["output"]; got != "plain output" {
|
||||
t.Errorf("persisted output = %v, want it untouched", got)
|
||||
}
|
||||
}
|
||||
@@ -1,145 +0,0 @@
|
||||
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)
|
||||
})
|
||||
}
|
||||
@@ -1,168 +0,0 @@
|
||||
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
|
||||
}
|
||||
@@ -1,282 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -1,51 +0,0 @@
|
||||
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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -3,48 +3,26 @@ package config
|
||||
import (
|
||||
"reflect"
|
||||
"strings"
|
||||
|
||||
"github.com/Twigpine/zero/internal/aimlapi"
|
||||
)
|
||||
|
||||
// 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
|
||||
// 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
|
||||
// `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.
|
||||
|
||||
@@ -291,92 +291,3 @@ 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")
|
||||
}
|
||||
}
|
||||
@@ -455,7 +455,7 @@ func projectAPIKeyEnvAllowed(profile ProviderProfile, envName string) bool {
|
||||
if baseURL == "" {
|
||||
baseURL = descriptor.DefaultBaseURL
|
||||
}
|
||||
if !providercatalog.IsCatalogBaseURL(descriptor, baseURL) {
|
||||
if !sameBaseURL(baseURL, descriptor.DefaultBaseURL) {
|
||||
return false
|
||||
}
|
||||
return containsStringFold(descriptor.AuthEnvVars, envName)
|
||||
@@ -463,7 +463,7 @@ func projectAPIKeyEnvAllowed(profile ProviderProfile, envName string) bool {
|
||||
|
||||
func projectBaseURLAllowed(profile ProviderProfile) bool {
|
||||
descriptor, ok := catalogDescriptorForProfile(profile)
|
||||
return ok && providercatalog.IsCatalogBaseURL(descriptor, profile.BaseURL)
|
||||
return ok && sameBaseURL(profile.BaseURL, descriptor.DefaultBaseURL)
|
||||
}
|
||||
|
||||
func hasInheritedProviderCredentialMaterial(project ProviderProfile, candidate ProviderProfile) bool {
|
||||
@@ -1128,13 +1128,10 @@ func applyCatalogDescriptor(profile *ProviderProfile, descriptor providercatalog
|
||||
if profile.Model == "" {
|
||||
profile.Model = descriptor.DefaultModel
|
||||
}
|
||||
// A legacy catalog base URL (an earlier default that still serves the
|
||||
// provider, e.g. a profile saved before a host move) counts as the catalog
|
||||
// endpoint, so it keeps the catalog env var binding.
|
||||
canonicalCatalogEndpoint := !explicitBaseURL || providercatalog.IsCatalogBaseURL(descriptor, profile.BaseURL)
|
||||
if profile.APIKeyEnv == "" && len(descriptor.AuthEnvVars) > 0 && canonicalCatalogEndpoint {
|
||||
if profile.APIKeyEnv == "" && len(descriptor.AuthEnvVars) > 0 && (!explicitBaseURL || sameBaseURL(profile.BaseURL, descriptor.DefaultBaseURL)) {
|
||||
profile.APIKeyEnv = descriptor.AuthEnvVars[0]
|
||||
}
|
||||
canonicalCatalogEndpoint := !explicitBaseURL || sameBaseURL(profile.BaseURL, descriptor.DefaultBaseURL)
|
||||
if len(descriptor.CustomHeaders) > 0 && canonicalCatalogEndpoint {
|
||||
catalogHeaders := descriptor.CustomHeaders
|
||||
if strings.EqualFold(strings.TrimSpace(descriptor.ID), "aimlapi") {
|
||||
|
||||
@@ -1357,74 +1357,6 @@ func TestResolveAppliesProviderCatalogDefaults(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAppliesOpenGatewayCatalogDefaults(t *testing.T) {
|
||||
path := writeConfig(t, `{
|
||||
"activeProvider": "gitlawb-opengateway",
|
||||
"providers": [{
|
||||
"name": "gitlawb-opengateway",
|
||||
"catalogID": "gitlawb-opengateway"
|
||||
}]
|
||||
}`)
|
||||
|
||||
resolved, err := Resolve(ResolveOptions{
|
||||
UserConfigPath: path,
|
||||
Env: map[string]string{
|
||||
"GITLAWB_OPENGATEWAY_API_KEY": "ogw_live_test",
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Resolve() error = %v", err)
|
||||
}
|
||||
if resolved.Provider.BaseURL != "https://api.twigpine.com/v1" {
|
||||
t.Fatalf("BaseURL = %q, want https://api.twigpine.com/v1", resolved.Provider.BaseURL)
|
||||
}
|
||||
if resolved.Provider.Model != "xiaomi/mimo-v2.6-pro" {
|
||||
t.Fatalf("Model = %q, want xiaomi/mimo-v2.6-pro", resolved.Provider.Model)
|
||||
}
|
||||
if resolved.Provider.APIKeyEnv != "GITLAWB_OPENGATEWAY_API_KEY" || resolved.Provider.APIKey != "ogw_live_test" {
|
||||
t.Fatalf("credential = (%q, %q), want GITLAWB_OPENGATEWAY_API_KEY env-resolved key", resolved.Provider.APIKeyEnv, resolved.Provider.APIKey)
|
||||
}
|
||||
}
|
||||
|
||||
// Profiles saved before the Twigpine rename persist the catalog's old base URL
|
||||
// (https://opengateway.gitlawb.com/v1) explicitly. That host still serves the
|
||||
// gateway, so those configs must keep resolving exactly as before: same host,
|
||||
// same model, catalog env var still bound, credential still picked up.
|
||||
func TestResolveKeepsLegacyOpenGatewayHostProfiles(t *testing.T) {
|
||||
path := writeConfig(t, `{
|
||||
"activeProvider": "gitlawb-opengateway",
|
||||
"providers": [{
|
||||
"name": "gitlawb-opengateway",
|
||||
"catalogID": "gitlawb-opengateway",
|
||||
"providerKind": "openai-compatible",
|
||||
"baseURL": "https://opengateway.gitlawb.com/v1",
|
||||
"model": "mimo-v2.5-pro"
|
||||
}]
|
||||
}`)
|
||||
|
||||
resolved, err := Resolve(ResolveOptions{
|
||||
UserConfigPath: path,
|
||||
Env: map[string]string{
|
||||
"GITLAWB_OPENGATEWAY_API_KEY": "ogw_live_test",
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Resolve() error = %v", err)
|
||||
}
|
||||
if resolved.Provider.BaseURL != "https://opengateway.gitlawb.com/v1" {
|
||||
t.Fatalf("BaseURL = %q, want the saved legacy host kept", resolved.Provider.BaseURL)
|
||||
}
|
||||
if resolved.Provider.Model != "mimo-v2.5-pro" {
|
||||
t.Fatalf("Model = %q, want the saved model kept", resolved.Provider.Model)
|
||||
}
|
||||
if resolved.Provider.APIKeyEnv != "GITLAWB_OPENGATEWAY_API_KEY" {
|
||||
t.Fatalf("APIKeyEnv = %q, want GITLAWB_OPENGATEWAY_API_KEY", resolved.Provider.APIKeyEnv)
|
||||
}
|
||||
if resolved.Provider.APIKey != "ogw_live_test" {
|
||||
t.Fatalf("APIKey = %q, want env-resolved key", resolved.Provider.APIKey)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAppliesMiniMaxCNCatalogDefaults(t *testing.T) {
|
||||
path := writeConfig(t, `{
|
||||
"activeProvider": "mini-cn",
|
||||
|
||||
@@ -724,39 +724,6 @@ 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"`
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
package execution
|
||||
|
||||
import "sync"
|
||||
|
||||
// ChildLaunchTracker is the one monotonic answer to "did the requested child
|
||||
// run", shared by every consumer of a prepared command.
|
||||
//
|
||||
// THE REPORT HAS THREE STATES, NOT TWO. For a plan whose child is created inside
|
||||
// an adapter, the report file passes through: not settled yet, settled with no
|
||||
// child, and child created. An absent file, an empty file the adapter opened but
|
||||
// has not written, a half-written one, and a decode error are all the FIRST
|
||||
// state, and collapsing them to "no child" is a claim about a process that may be
|
||||
// running at that instant. Consumers used to make that collapse independently and
|
||||
// then compensate for the timing on their own, which is why fixing one
|
||||
// presentation site kept exposing the next.
|
||||
//
|
||||
// Two things make an answer terminal here:
|
||||
//
|
||||
// - Confirm, when the consumer has observed the child directly. A stdio MCP
|
||||
// server answering initialize is that observation: the adapter speaks no MCP,
|
||||
// so a well-formed response can only have come from the requested child.
|
||||
// - Settle, when the adapter process itself has exited. After that the report
|
||||
// will never change, so whatever it says, including nothing, is the truth.
|
||||
//
|
||||
// Settle runs from Cleanup, BEFORE the report file is deleted. Reading the
|
||||
// evidence and then destroying it in one step is what lets a consumer ask after
|
||||
// close and still get an answer, rather than racing a deletion it does not
|
||||
// control.
|
||||
//
|
||||
// The decision only ever moves from unknown to known and never back, so two
|
||||
// consumers of the same prepared command cannot disagree, and a retry cannot turn
|
||||
// a launch that happened into one that did not.
|
||||
type ChildLaunchTracker struct {
|
||||
mu sync.Mutex
|
||||
settled bool
|
||||
launched bool
|
||||
|
||||
ownedByAdapter bool
|
||||
report func() (AdapterReport, error)
|
||||
}
|
||||
|
||||
// NewChildLaunchTracker builds the tracker for a prepared command and returns it
|
||||
// with a Cleanup that settles before releasing the plan's resources.
|
||||
//
|
||||
// The returned cleanup is the one the caller must use. Calling the prepared
|
||||
// command's own Cleanup instead deletes the report while the decision is still
|
||||
// unknown, which is the ordering hazard this type exists to remove.
|
||||
func NewChildLaunchTracker(prepared PreparedCommand) (*ChildLaunchTracker, func()) {
|
||||
tracker := &ChildLaunchTracker{
|
||||
ownedByAdapter: prepared.ChildLaunchOwnedByAdapter,
|
||||
report: prepared.Report,
|
||||
}
|
||||
// A plan whose child is the started process needs no evidence: Start
|
||||
// succeeding IS the launch, so the answer is terminal from the beginning.
|
||||
if !tracker.ownedByAdapter {
|
||||
tracker.settled = true
|
||||
tracker.launched = true
|
||||
}
|
||||
inner := prepared.Cleanup
|
||||
return tracker, func() {
|
||||
tracker.Settle()
|
||||
if inner != nil {
|
||||
inner()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Confirm latches a launch the consumer observed for itself.
|
||||
//
|
||||
// Positive evidence is always terminal. Nothing that happens later can make a
|
||||
// child that answered not have run, so this needs no settlement and cannot be
|
||||
// undone by a subsequent Settle finding no report.
|
||||
func (tracker *ChildLaunchTracker) Confirm() {
|
||||
if tracker == nil {
|
||||
return
|
||||
}
|
||||
tracker.mu.Lock()
|
||||
defer tracker.mu.Unlock()
|
||||
tracker.settled = true
|
||||
tracker.launched = true
|
||||
}
|
||||
|
||||
// Settle freezes the answer because the adapter has finished.
|
||||
//
|
||||
// One last read first: the adapter may have published between the consumer's
|
||||
// last look and its exit, and this is the only remaining chance to see it.
|
||||
func (tracker *ChildLaunchTracker) Settle() {
|
||||
if tracker == nil {
|
||||
return
|
||||
}
|
||||
tracker.mu.Lock()
|
||||
defer tracker.mu.Unlock()
|
||||
if tracker.settled {
|
||||
return
|
||||
}
|
||||
tracker.launched = tracker.readReportLocked()
|
||||
tracker.settled = true
|
||||
}
|
||||
|
||||
// Launched reports the decision so far.
|
||||
//
|
||||
// Before settlement this reads the report and answers true only for a published
|
||||
// launch. A negative answer here is NOT remembered, because the adapter may still
|
||||
// be between creating the child and recording it; the next caller asks again.
|
||||
// Consumers that need a final answer settle first, which Cleanup does for them.
|
||||
func (tracker *ChildLaunchTracker) Launched() bool {
|
||||
if tracker == nil {
|
||||
return false
|
||||
}
|
||||
tracker.mu.Lock()
|
||||
defer tracker.mu.Unlock()
|
||||
if tracker.settled {
|
||||
return tracker.launched
|
||||
}
|
||||
if tracker.readReportLocked() {
|
||||
tracker.launched = true
|
||||
tracker.settled = true
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Settled reports whether the answer is final, for callers that would otherwise
|
||||
// present an unknown as a fact.
|
||||
func (tracker *ChildLaunchTracker) Settled() bool {
|
||||
if tracker == nil {
|
||||
return false
|
||||
}
|
||||
tracker.mu.Lock()
|
||||
defer tracker.mu.Unlock()
|
||||
return tracker.settled
|
||||
}
|
||||
|
||||
// readReportLocked answers only the positive case. A missing file, an empty one,
|
||||
// a partial write and a decode error are all "nothing published yet", which is
|
||||
// the state this type refuses to record as an outcome.
|
||||
func (tracker *ChildLaunchTracker) readReportLocked() bool {
|
||||
if tracker.report == nil {
|
||||
return false
|
||||
}
|
||||
report, err := tracker.report()
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return ResolveChildLaunched(true, tracker.ownedByAdapter, report)
|
||||
}
|
||||
@@ -0,0 +1,168 @@
|
||||
package execution
|
||||
|
||||
import "testing"
|
||||
|
||||
// scriptedReport is an adapter report whose answer changes over time, the way a
|
||||
// real one does: the file is opened empty, decodes to an error while nothing has
|
||||
// been written, and only later carries the fact.
|
||||
type scriptedReport struct {
|
||||
launched bool
|
||||
reads int
|
||||
err error
|
||||
}
|
||||
|
||||
func (script *scriptedReport) read() (AdapterReport, error) {
|
||||
script.reads++
|
||||
if script.err != nil {
|
||||
return AdapterReport{}, script.err
|
||||
}
|
||||
if !script.launched {
|
||||
// What an empty or half-written file decodes to. NOT childLaunched=false,
|
||||
// which would be the adapter stating an outcome.
|
||||
return AdapterReport{}, nil
|
||||
}
|
||||
launched := true
|
||||
return AdapterReport{ChildLaunched: &launched}, nil
|
||||
}
|
||||
|
||||
func trackerOver(script *scriptedReport) (*ChildLaunchTracker, func()) {
|
||||
return NewChildLaunchTracker(PreparedCommand{
|
||||
ChildLaunchOwnedByAdapter: true,
|
||||
Report: script.read,
|
||||
})
|
||||
}
|
||||
|
||||
// A NEGATIVE READ IS NOT AN OUTCOME UNTIL THE ADAPTER IS DONE.
|
||||
//
|
||||
// This is the whole point of the type. Reading the report between the adapter
|
||||
// creating the child and recording it answers "nothing published", and caching
|
||||
// that answer freezes "not yet" into "never" for the rest of the session.
|
||||
func TestAnUnsettledNegativeIsNotRemembered(t *testing.T) {
|
||||
script := &scriptedReport{}
|
||||
tracker, _ := trackerOver(script)
|
||||
|
||||
if tracker.Launched() {
|
||||
t.Fatal("SETUP INVALID: an unpublished report answered launched, so there is no negative here to cache")
|
||||
}
|
||||
if tracker.Settled() {
|
||||
t.Fatal("a read that found nothing settled the decision; the adapter may still be about to publish")
|
||||
}
|
||||
|
||||
// The adapter publishes, as it does a moment after CreateProcessAsUser.
|
||||
script.launched = true
|
||||
if !tracker.Launched() {
|
||||
t.Fatal("the launch published after the first read was never seen; the earlier negative was cached")
|
||||
}
|
||||
if !tracker.Settled() {
|
||||
t.Fatal("a confirmed launch left the decision open")
|
||||
}
|
||||
}
|
||||
|
||||
// And once it IS an outcome it stays one, so the answer is monotonic and two
|
||||
// consumers of the same prepared command cannot disagree.
|
||||
func TestASettledLaunchIsNeverWithdrawn(t *testing.T) {
|
||||
script := &scriptedReport{launched: true}
|
||||
tracker, cleanup := trackerOver(script)
|
||||
|
||||
if !tracker.Launched() {
|
||||
t.Fatal("SETUP INVALID: a published report did not answer launched")
|
||||
}
|
||||
// Cleanup deletes the report in production, which is what a later read would
|
||||
// find. The decision must not follow it.
|
||||
script.launched = false
|
||||
cleanup()
|
||||
if !tracker.Launched() {
|
||||
t.Fatal("the decision followed the report file into deletion; a server that ran became one that did not")
|
||||
}
|
||||
}
|
||||
|
||||
// CLEANUP SETTLES BEFORE IT DESTROYS THE EVIDENCE.
|
||||
//
|
||||
// The report is a file the plan's cleanup removes. A consumer that asks after
|
||||
// cleanup used to read an absent report and answer "no child" about a server that
|
||||
// really did run, which is why the decision was being taken early and hitting the
|
||||
// unsettled window instead.
|
||||
func TestCleanupSettlesBeforeReleasingThePlan(t *testing.T) {
|
||||
script := &scriptedReport{launched: true}
|
||||
var order []string
|
||||
tracker, cleanup := NewChildLaunchTracker(PreparedCommand{
|
||||
ChildLaunchOwnedByAdapter: true,
|
||||
Report: func() (AdapterReport, error) {
|
||||
order = append(order, "read")
|
||||
return script.read()
|
||||
},
|
||||
Cleanup: func() { order = append(order, "release") },
|
||||
})
|
||||
|
||||
cleanup()
|
||||
if len(order) != 2 || order[0] != "read" || order[1] != "release" {
|
||||
t.Fatalf("cleanup order = %v, want the report read before the plan is released", order)
|
||||
}
|
||||
if !tracker.Launched() {
|
||||
t.Fatal("the launch published just before cleanup was lost; nothing read the report on the way out")
|
||||
}
|
||||
}
|
||||
|
||||
// An adapter that exits having published nothing settles negative, and that
|
||||
// answer is final: this is the case the unsettled rule must not swallow.
|
||||
func TestAnAdapterThatPublishedNothingSettlesNegative(t *testing.T) {
|
||||
script := &scriptedReport{}
|
||||
tracker, cleanup := trackerOver(script)
|
||||
cleanup()
|
||||
|
||||
if !tracker.Settled() {
|
||||
t.Fatal("the adapter finished without the decision becoming final")
|
||||
}
|
||||
if tracker.Launched() {
|
||||
t.Fatal("an adapter that published nothing was credited with a launch")
|
||||
}
|
||||
// And a report that changes after the adapter is gone changes nothing.
|
||||
script.launched = true
|
||||
if tracker.Launched() {
|
||||
t.Fatal("a settled decision was reopened by a later read")
|
||||
}
|
||||
}
|
||||
|
||||
// Direct evidence outranks the report and needs no settlement, because nothing
|
||||
// that happens later can make a child that answered not have run.
|
||||
func TestConfirmIsTerminalWithoutAReport(t *testing.T) {
|
||||
script := &scriptedReport{}
|
||||
tracker, cleanup := trackerOver(script)
|
||||
|
||||
tracker.Confirm()
|
||||
if !tracker.Launched() || !tracker.Settled() {
|
||||
t.Fatal("an observed child did not settle the decision")
|
||||
}
|
||||
cleanup()
|
||||
if !tracker.Launched() {
|
||||
t.Fatal("settling from cleanup overwrote an observed launch with an absent report")
|
||||
}
|
||||
}
|
||||
|
||||
// A plan whose started process IS the requested command has nothing to decide,
|
||||
// and must not be made to wait on a report it will never have.
|
||||
func TestAPlanTheAdapterDoesNotOwnIsLaunchedFromTheStart(t *testing.T) {
|
||||
tracker, _ := NewChildLaunchTracker(PreparedCommand{ChildLaunchOwnedByAdapter: false})
|
||||
if !tracker.Launched() || !tracker.Settled() {
|
||||
t.Fatal("a directly started command was not treated as launched")
|
||||
}
|
||||
}
|
||||
|
||||
// A report that fails to read is the unsettled state, not a negative outcome.
|
||||
func TestADecodeErrorIsNotAnAnswer(t *testing.T) {
|
||||
script := &scriptedReport{err: errScriptedReportBroken}
|
||||
tracker, _ := trackerOver(script)
|
||||
|
||||
if tracker.Launched() {
|
||||
t.Fatal("a broken report was read as a launch")
|
||||
}
|
||||
if tracker.Settled() {
|
||||
t.Fatal("a broken report settled the decision; a half-written file is a moment, not an outcome")
|
||||
}
|
||||
}
|
||||
|
||||
var errScriptedReportBroken = errScriptedReport("unexpected end of JSON input")
|
||||
|
||||
type errScriptedReport string
|
||||
|
||||
func (e errScriptedReport) Error() string { return string(e) }
|
||||
@@ -173,6 +173,11 @@ type Enforcement struct {
|
||||
Level string `json:"level,omitempty"`
|
||||
Degraded bool `json:"degraded,omitempty"`
|
||||
DowngradeReason string `json:"downgradeReason,omitempty"`
|
||||
// Notices are least-privilege disclosures about the enforcement actually
|
||||
// applied to THIS command, as opposed to the diagnostic views produced by
|
||||
// `zero sandbox policy` and `zero sandbox check`. A trade an operator only
|
||||
// discovers by running a separate diagnostic command is not disclosed.
|
||||
Notices []string `json:"notices,omitempty"`
|
||||
}
|
||||
|
||||
type Outcome struct {
|
||||
@@ -182,7 +187,18 @@ type Outcome struct {
|
||||
Exit *Exit `json:"exit,omitempty"`
|
||||
Denial *Denial `json:"denial,omitempty"`
|
||||
Enforcement Enforcement `json:"enforcement"`
|
||||
Changes []Change `json:"changes,omitempty"`
|
||||
// Launched records whether an OS process was actually created, observed at
|
||||
// the boundary that calls Run rather than inferred afterwards.
|
||||
//
|
||||
// OutcomeKind is not a launch-state field, and reading it as one is wrong in
|
||||
// both directions. A child that ran and then produced an unreadable adapter
|
||||
// report is rewritten to a setup failure, so inference drops a disclosure that
|
||||
// did apply; a context already cancelled before os.StartProcess yields a
|
||||
// cancellation, so inference claims reduced enforcement for a child that never
|
||||
// existed. Report decoding can fail after launch without rewriting history,
|
||||
// and cancellation happens on either side of Start.
|
||||
Launched bool `json:"launched,omitempty"`
|
||||
Changes []Change `json:"changes,omitempty"`
|
||||
}
|
||||
|
||||
// AdapterReport is the structured, machine-readable result emitted by a
|
||||
@@ -190,6 +206,51 @@ type Outcome struct {
|
||||
// command text cannot impersonate a policy decision.
|
||||
type AdapterReport struct {
|
||||
Denial *Denial `json:"denial,omitempty"`
|
||||
// ChildLaunched is the adapter's authoritative statement that the REQUESTED
|
||||
// process started, for a plan where the command the runner starts is not that
|
||||
// process.
|
||||
//
|
||||
// A wrapped plan starts a helper, and the helper creates the sandboxed child
|
||||
// only after validating the setup marker, applying ACLs, checking the network
|
||||
// policy, building capability SIDs and minting the restricted token. Any of
|
||||
// those can fail with the helper already running, so the runner's own
|
||||
// exec.Cmd.Process tells it the WRAPPER started and nothing about the child.
|
||||
// Only the adapter sees that transition, so only the adapter may report it.
|
||||
//
|
||||
// nil means the adapter does not speak to this, and the runner keeps its own
|
||||
// observation. That is correct for every direct, unwrapped command, where the
|
||||
// process the runner starts IS the requested one.
|
||||
ChildLaunched *bool `json:"childLaunched,omitempty"`
|
||||
}
|
||||
|
||||
// ChildLaunched reports whether this outcome describes a process that actually
|
||||
// started, from the recorded fact rather than from the terminal outcome kind.
|
||||
func (outcome Outcome) ChildLaunched() bool {
|
||||
return outcome.Launched
|
||||
}
|
||||
|
||||
// AppliedEnforcementNotices returns the least-privilege disclosures that are
|
||||
// true of what actually happened.
|
||||
//
|
||||
// ONE DECISION, AT THE BOUNDARY WHERE THE OUTCOME IS KNOWN. Enforcement.Notices
|
||||
// is planned: it describes the shape the command was PREPARED to run under, and
|
||||
// planning is not proof that anything ran. Every consumer that copied the field
|
||||
// straight out therefore made the completed-enforcement claim for commands that
|
||||
// never launched, telling an operator the write jail had been traded away for a
|
||||
// child that failed before it existed.
|
||||
//
|
||||
// Keeping this on Outcome rather than repeating an outcome-kind switch in hooks,
|
||||
// plugins and tools is the point: a new pre-launch outcome kind has to be
|
||||
// classified once, here, instead of being silently disclosed by whichever
|
||||
// consumer was not updated.
|
||||
func (outcome Outcome) AppliedEnforcementNotices() []string {
|
||||
if !outcome.ChildLaunched() {
|
||||
return nil
|
||||
}
|
||||
if len(outcome.Enforcement.Notices) == 0 {
|
||||
return nil
|
||||
}
|
||||
return append([]string(nil), outcome.Enforcement.Notices...)
|
||||
}
|
||||
|
||||
func (outcome Outcome) Validate() error {
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
package execution
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os/exec"
|
||||
"runtime"
|
||||
"testing"
|
||||
)
|
||||
|
||||
const launchStateNotice = "denyRead is configured, so the write jail is not confining writes"
|
||||
|
||||
func launchStateShell(ctx context.Context, script string) *exec.Cmd {
|
||||
if runtime.GOOS == "windows" {
|
||||
return exec.CommandContext(ctx, "cmd.exe", "/c", script)
|
||||
}
|
||||
return exec.CommandContext(ctx, "/bin/sh", "-c", script)
|
||||
}
|
||||
|
||||
// launchStatePreparer plans a command carrying an enforcement notice, and can
|
||||
// make the adapter report fail after the child has already run.
|
||||
type launchStatePreparer struct {
|
||||
script string
|
||||
reportErr error
|
||||
missing bool
|
||||
}
|
||||
|
||||
func (p *launchStatePreparer) PrepareExecution(ctx context.Context, _ Request) (PreparedCommand, error) {
|
||||
command := launchStateShell(ctx, p.script)
|
||||
if p.missing {
|
||||
command = exec.CommandContext(ctx, "definitely-not-a-real-binary-zzz")
|
||||
}
|
||||
prepared := PreparedCommand{
|
||||
Command: command,
|
||||
Enforcement: Enforcement{Notices: []string{launchStateNotice}},
|
||||
}
|
||||
if p.reportErr != nil {
|
||||
prepared.Report = func() (AdapterReport, error) { return AdapterReport{}, p.reportErr }
|
||||
}
|
||||
return prepared, nil
|
||||
}
|
||||
|
||||
func captured(t *testing.T, ctx context.Context, p *launchStatePreparer) CapturedResult {
|
||||
t.Helper()
|
||||
return NewRunner(p).ExecuteCaptured(ctx, CapturedRequest{Request: Request{
|
||||
Origin: OriginHook,
|
||||
Mode: ModeCaptured,
|
||||
Command: Command{Name: "irrelevant"},
|
||||
WorkingDirectory: t.TempDir(),
|
||||
WorkspaceRoots: []string{t.TempDir()},
|
||||
Approval: ApprovalContext{PolicyVersion: PolicyVersion},
|
||||
}})
|
||||
}
|
||||
|
||||
// THE OUTCOME KIND IS NOT A LAUNCH-STATE FIELD, IN EITHER DIRECTION.
|
||||
//
|
||||
// Deriving launch from the terminal kind is wrong twice over. The adapter report
|
||||
// is read AFTER Run, so a child that really ran and then produced an unreadable
|
||||
// report is rewritten to a setup failure: inference drops a disclosure that did
|
||||
// apply. And a context already cancelled before os.StartProcess still selects a
|
||||
// cancellation, so inference claims reduced enforcement for a process that never
|
||||
// existed.
|
||||
func TestLaunchStateIsRecordedNotInferred(t *testing.T) {
|
||||
t.Run("ran, then the adapter report failed", func(t *testing.T) {
|
||||
result := captured(t, context.Background(), &launchStatePreparer{
|
||||
script: "exit 0",
|
||||
reportErr: errors.New("adapter report is unreadable"),
|
||||
})
|
||||
if result.Outcome.Kind != OutcomeSandboxSetupFailure {
|
||||
t.Fatalf("SETUP INVALID: kind = %q, want the report failure to rewrite it", result.Outcome.Kind)
|
||||
}
|
||||
if !result.Outcome.Launched {
|
||||
t.Error("a child that ran was recorded as never launched")
|
||||
}
|
||||
if got := result.Outcome.AppliedEnforcementNotices(); len(got) != 1 {
|
||||
t.Errorf("the disclosure was dropped for a child that did run: %#v", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("cancelled before the process started", func(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
result := captured(t, ctx, &launchStatePreparer{script: "exit 0"})
|
||||
if result.Outcome.Launched {
|
||||
t.Error("a process that never started was recorded as launched")
|
||||
}
|
||||
if got := result.Outcome.AppliedEnforcementNotices(); len(got) != 0 {
|
||||
t.Errorf("an enforcement trade was claimed for a process that never existed: %#v", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("never found", func(t *testing.T) {
|
||||
result := captured(t, context.Background(), &launchStatePreparer{missing: true})
|
||||
if result.Outcome.Launched {
|
||||
t.Error("a missing executable was recorded as launched")
|
||||
}
|
||||
if got := result.Outcome.AppliedEnforcementNotices(); len(got) != 0 {
|
||||
t.Errorf("a missing executable claimed an enforcement trade: %#v", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("ordinary success still discloses", func(t *testing.T) {
|
||||
result := captured(t, context.Background(), &launchStatePreparer{script: "exit 0"})
|
||||
if !result.Outcome.Launched {
|
||||
t.Fatal("an ordinary run was recorded as never launched")
|
||||
}
|
||||
if got := result.Outcome.AppliedEnforcementNotices(); len(got) != 1 {
|
||||
t.Errorf("an ordinary run lost its disclosure: %#v", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
package execution
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
const liveLaunchNotice = "denyRead is configured, so the write jail is not confining writes"
|
||||
|
||||
// liveReportReader mirrors sandbox.CommandPlan.ExecutionReport: read the file the
|
||||
// helper publishes, treat "not there yet" as nothing recorded.
|
||||
func liveReportReader(path string) func() (AdapterReport, error) {
|
||||
return func() (AdapterReport, error) {
|
||||
raw, err := os.ReadFile(path)
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
return AdapterReport{}, nil
|
||||
}
|
||||
if err != nil {
|
||||
return AdapterReport{}, err
|
||||
}
|
||||
var report AdapterReport
|
||||
if err := json.Unmarshal(raw, &report); err != nil {
|
||||
return AdapterReport{}, err
|
||||
}
|
||||
return report, nil
|
||||
}
|
||||
}
|
||||
|
||||
// liveHelperCommand stands in for the Windows helper: publish the launch fact the
|
||||
// way it does right after CreateProcessAsUser, then stay alive the way it does
|
||||
// while waiting on the child.
|
||||
func liveHelperCommand(t *testing.T, reportPath string, publish bool) *exec.Cmd {
|
||||
t.Helper()
|
||||
if publish {
|
||||
if err := os.WriteFile(reportPath, []byte(`{"childLaunched":true}`), 0o600); err != nil {
|
||||
t.Fatalf("publish the launch report: %v", err)
|
||||
}
|
||||
}
|
||||
if runtime.GOOS == "windows" {
|
||||
// A child that holds itself open without exiting.
|
||||
return exec.Command("cmd.exe", "/c", "pause")
|
||||
}
|
||||
return exec.Command("/bin/sh", "-c", "sleep 30")
|
||||
}
|
||||
|
||||
// liveRequest is a valid interactive request; ProcessManager.Start validates it
|
||||
// before anything under test runs.
|
||||
func liveRequest(t *testing.T) Request {
|
||||
t.Helper()
|
||||
return Request{
|
||||
Origin: OriginInteractiveCommand,
|
||||
Mode: ModeInteractive,
|
||||
Command: Command{Name: "helper"},
|
||||
WorkingDirectory: t.TempDir(),
|
||||
WorkspaceRoots: []string{t.TempDir()},
|
||||
Approval: ApprovalContext{PolicyVersion: PolicyVersion},
|
||||
}
|
||||
}
|
||||
|
||||
// A RETAINED SESSION HAS TO DISCLOSE WHILE IT IS STILL RUNNING.
|
||||
//
|
||||
// The helper publishes childLaunched immediately after it creates the restricted
|
||||
// child, and only then waits for it. The manager read that report exclusively in
|
||||
// the post-Wait goroutine, so for the whole live lifetime of a wrapped session the
|
||||
// report was the zero value: the first exec_command reply and every write_stdin
|
||||
// poll resolved Launched=false and disclosed nothing, while the fact sat readable
|
||||
// on disk. A watcher, or a retained session nobody polls to completion, would
|
||||
// never be told the write jail had been traded away.
|
||||
//
|
||||
// Driven through the real ProcessManager with a real child process, and asserted
|
||||
// on the ProcessResult the tool layer consumes.
|
||||
func TestALiveWrappedSessionCarriesTheLaunchFact(t *testing.T) {
|
||||
directory := t.TempDir()
|
||||
reportPath := filepath.Join(directory, "report.json")
|
||||
command := liveHelperCommand(t, reportPath, true)
|
||||
|
||||
manager := NewProcessManager(ProcessManagerOptions{})
|
||||
result, err := manager.Start(context.Background(), ProcessStart{
|
||||
Prepared: PreparedCommand{
|
||||
Command: command,
|
||||
ChildLaunchOwnedByAdapter: true,
|
||||
Enforcement: Enforcement{Notices: []string{liveLaunchNotice}},
|
||||
Report: liveReportReader(reportPath),
|
||||
},
|
||||
Request: liveRequest(t),
|
||||
}, 300*time.Millisecond)
|
||||
if err != nil {
|
||||
t.Fatalf("Start: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { manager.StopAll() })
|
||||
|
||||
// SETUP: this has to be the LIVE state, and the fact has to be on disk, or the
|
||||
// assertion below would be about a terminal read.
|
||||
if result.Exited {
|
||||
t.Fatal("SETUP INVALID: the stand-in helper exited, so the live lifecycle is not under test")
|
||||
}
|
||||
if _, statErr := os.Stat(reportPath); statErr != nil {
|
||||
t.Fatalf("SETUP INVALID: the launch report is not on disk, so there is nothing to observe: %v", statErr)
|
||||
}
|
||||
|
||||
if !ResolveChildLaunched(true, result.ChildLaunchOwnedByAdapter, result.Report) {
|
||||
t.Fatal("a live wrapped session resolved as not launched, so its enforcement disclosure is withheld while the command runs")
|
||||
}
|
||||
|
||||
// And again on a poll, which is the write_stdin leg.
|
||||
polled, err := manager.Continue(context.Background(), ProcessContinue{ProcessID: result.ProcessID, Wait: 200 * time.Millisecond})
|
||||
if err != nil {
|
||||
t.Fatalf("Continue: %v", err)
|
||||
}
|
||||
if polled.Exited {
|
||||
t.Fatal("SETUP INVALID: the helper exited before the poll, so the live poll is not under test")
|
||||
}
|
||||
if !ResolveChildLaunched(true, polled.ChildLaunchOwnedByAdapter, polled.Report) {
|
||||
t.Fatal("a live poll of a wrapped session resolved as not launched")
|
||||
}
|
||||
if got := polled.Enforcement.Notices; len(got) != 1 || got[0] != liveLaunchNotice {
|
||||
t.Fatalf("the live poll carries notices %v, want exactly the one planned notice", got)
|
||||
}
|
||||
}
|
||||
|
||||
// AND A HELPER THAT NEVER CREATED THE CHILD STAYS SILENT.
|
||||
//
|
||||
// This is the negative the live read must not destroy. A helper that starts and
|
||||
// then fails setup, ACL application, or CreateProcessAsUser has an outer process
|
||||
// running and no child, so nothing may be promoted from the fact that the
|
||||
// wrapper itself is alive.
|
||||
func TestALiveWrappedSessionWithNoReportedChildStaysSilent(t *testing.T) {
|
||||
directory := t.TempDir()
|
||||
reportPath := filepath.Join(directory, "report.json")
|
||||
command := liveHelperCommand(t, reportPath, false)
|
||||
|
||||
manager := NewProcessManager(ProcessManagerOptions{})
|
||||
result, err := manager.Start(context.Background(), ProcessStart{
|
||||
Prepared: PreparedCommand{
|
||||
Command: command,
|
||||
ChildLaunchOwnedByAdapter: true,
|
||||
Enforcement: Enforcement{Notices: []string{liveLaunchNotice}},
|
||||
Report: liveReportReader(reportPath),
|
||||
},
|
||||
Request: liveRequest(t),
|
||||
}, 300*time.Millisecond)
|
||||
if err != nil {
|
||||
t.Fatalf("Start: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { manager.StopAll() })
|
||||
|
||||
if result.Exited {
|
||||
t.Fatal("SETUP INVALID: the stand-in helper exited, so the live lifecycle is not under test")
|
||||
}
|
||||
if _, statErr := os.Stat(reportPath); statErr == nil {
|
||||
t.Fatal("SETUP INVALID: a report exists, so this is not the no-child case")
|
||||
}
|
||||
if ResolveChildLaunched(true, result.ChildLaunchOwnedByAdapter, result.Report) {
|
||||
t.Fatal("a helper that reported no child was promoted to a launch, so the operator is told a write jail was traded away for a child that never existed")
|
||||
}
|
||||
}
|
||||
|
||||
// A report that appears mid-flight is observed on the next poll, which is what
|
||||
// makes this a lifecycle transition rather than a start-time snapshot.
|
||||
func TestTheLaunchFactIsObservedWhenItAppearsMidFlight(t *testing.T) {
|
||||
directory := t.TempDir()
|
||||
reportPath := filepath.Join(directory, "report.json")
|
||||
command := liveHelperCommand(t, reportPath, false)
|
||||
|
||||
manager := NewProcessManager(ProcessManagerOptions{})
|
||||
result, err := manager.Start(context.Background(), ProcessStart{
|
||||
Prepared: PreparedCommand{
|
||||
Command: command,
|
||||
ChildLaunchOwnedByAdapter: true,
|
||||
Report: liveReportReader(reportPath),
|
||||
},
|
||||
Request: liveRequest(t),
|
||||
}, 200*time.Millisecond)
|
||||
if err != nil {
|
||||
t.Fatalf("Start: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { manager.StopAll() })
|
||||
if ResolveChildLaunched(true, result.ChildLaunchOwnedByAdapter, result.Report) {
|
||||
t.Fatal("SETUP INVALID: nothing was published yet, so the first result must not be a launch")
|
||||
}
|
||||
|
||||
if err := os.WriteFile(reportPath, []byte(`{"childLaunched":true}`), 0o600); err != nil {
|
||||
t.Fatalf("publish the launch report: %v", err)
|
||||
}
|
||||
polled, err := manager.Continue(context.Background(), ProcessContinue{ProcessID: result.ProcessID, Wait: 200 * time.Millisecond})
|
||||
if err != nil {
|
||||
t.Fatalf("Continue: %v", err)
|
||||
}
|
||||
if !ResolveChildLaunched(true, polled.ChildLaunchOwnedByAdapter, polled.Report) {
|
||||
t.Fatal("the launch published while the session was live was never observed")
|
||||
}
|
||||
}
|
||||
@@ -79,8 +79,12 @@ type ProcessResult struct {
|
||||
Enforcement Enforcement
|
||||
Report AdapterReport
|
||||
ReportErr error
|
||||
Changes []Change
|
||||
Metadata map[string]string
|
||||
// ChildLaunchOwnedByAdapter carries the prepared plan's ownership of the
|
||||
// requested-child launch fact through to the caller, which for a retained
|
||||
// session no longer has the plan.
|
||||
ChildLaunchOwnedByAdapter bool
|
||||
Changes []Change
|
||||
Metadata map[string]string
|
||||
}
|
||||
|
||||
type ProcessSnapshot struct {
|
||||
@@ -162,6 +166,7 @@ func (manager *ProcessManager) Start(ctx context.Context, input ProcessStart, wa
|
||||
command: command,
|
||||
request: request,
|
||||
enforcement: input.Prepared.Enforcement,
|
||||
ownedLaunch: input.Prepared.ChildLaunchOwnedByAdapter,
|
||||
report: input.Prepared.Report,
|
||||
cleanup: input.Prepared.Cleanup,
|
||||
stdin: stdin,
|
||||
@@ -423,31 +428,46 @@ func (manager *ProcessManager) removeCompletedLater(process *managedProcess) {
|
||||
}
|
||||
|
||||
type managedProcess struct {
|
||||
id int
|
||||
commandText string
|
||||
cwd string
|
||||
relativeCwd string
|
||||
startedAt time.Time
|
||||
lastUsedAt time.Time
|
||||
tty bool
|
||||
command *exec.Cmd
|
||||
request Request
|
||||
enforcement Enforcement
|
||||
report func() (AdapterReport, error)
|
||||
cleanup func()
|
||||
stdin io.WriteCloser
|
||||
output *processOutputBuffer
|
||||
reaped chan struct{}
|
||||
doneOnce sync.Once
|
||||
done chan struct{}
|
||||
kill func(int) error
|
||||
mu sync.Mutex
|
||||
exitCode *int
|
||||
waitErr error
|
||||
resultReport AdapterReport
|
||||
reportErr error
|
||||
changes []Change
|
||||
metadata map[string]string
|
||||
id int
|
||||
commandText string
|
||||
cwd string
|
||||
relativeCwd string
|
||||
startedAt time.Time
|
||||
lastUsedAt time.Time
|
||||
tty bool
|
||||
command *exec.Cmd
|
||||
request Request
|
||||
enforcement Enforcement
|
||||
ownedLaunch bool
|
||||
// launchObserved latches the adapter's launch transition the first time it is
|
||||
// seen, so a live result can report it. Guarded by mu.
|
||||
launchObserved bool
|
||||
report func() (AdapterReport, error)
|
||||
cleanup func()
|
||||
stdin io.WriteCloser
|
||||
output *processOutputBuffer
|
||||
reaped chan struct{}
|
||||
doneOnce sync.Once
|
||||
done chan struct{}
|
||||
kill func(int) error
|
||||
mu sync.Mutex
|
||||
exitCode *int
|
||||
waitErr error
|
||||
resultReport AdapterReport
|
||||
reportErr error
|
||||
changes []Change
|
||||
metadata map[string]string
|
||||
}
|
||||
|
||||
// launchedReportLocked returns the report to hand out, with a latched live
|
||||
// launch folded in. Caller holds mu.
|
||||
func (process *managedProcess) launchedReportLocked() AdapterReport {
|
||||
report := process.resultReport
|
||||
if report.ChildLaunched == nil && process.launchObserved {
|
||||
launched := true
|
||||
report.ChildLaunched = &launched
|
||||
}
|
||||
return report
|
||||
}
|
||||
|
||||
func (process *managedProcess) markDone(err error, exitCode int, report AdapterReport, reportErr error, changes []Change) {
|
||||
@@ -455,14 +475,76 @@ func (process *managedProcess) markDone(err error, exitCode int, report AdapterR
|
||||
process.waitErr = err
|
||||
process.exitCode = &exitCode
|
||||
process.resultReport = report
|
||||
// The plan's cleanup has already removed the report file by the time this
|
||||
// runs on some orderings, so a terminal read can answer "nothing recorded"
|
||||
// about a child that demonstrably started. A launch we already saw is not
|
||||
// un-seen by that.
|
||||
if report.ChildLaunched == nil && process.launchObserved {
|
||||
launched := true
|
||||
process.resultReport.ChildLaunched = &launched
|
||||
}
|
||||
process.reportErr = reportErr
|
||||
process.changes = append([]Change(nil), changes...)
|
||||
process.mu.Unlock()
|
||||
process.doneOnce.Do(func() { close(process.done) })
|
||||
}
|
||||
|
||||
// observeLaunch reads the adapter's launch report while the process is still
|
||||
// running, and latches a confirmed launch.
|
||||
//
|
||||
// THE LAUNCH FACT IS A LIFECYCLE TRANSITION, NOT TERMINAL DATA. The Windows
|
||||
// helper publishes childLaunched immediately after CreateProcessAsUser creates
|
||||
// the restricted child, and then waits for it. The manager used to read the
|
||||
// report only in the post-Wait goroutine, so for the entire live lifetime of a
|
||||
// retained session the report was the zero value: the first exec_command reply
|
||||
// and every write_stdin poll resolved Launched=false and disclosed nothing,
|
||||
// even though the fact was sitting readable on disk. A watcher or an abandoned
|
||||
// retained session could therefore never be told the write jail had been traded
|
||||
// away. The MCP launcher already reads the report while its server is live;
|
||||
// this is the same read, in the launcher that was left behind.
|
||||
//
|
||||
// SILENCE IS NOT A NEGATIVE, BUT AN EXPLICIT NEGATIVE IS. An absent, partial or
|
||||
// undecodable report, and a helper that failed before it ever created the child,
|
||||
// must all leave the live result exactly as it was: not confirmed, nothing
|
||||
// disclosed. Reading absence as false, or surfacing a read error or a denial
|
||||
// from here, would let a mid-flight poll rewrite a running command into a setup
|
||||
// failure, and absence is the normal state once the plan's cleanup has removed
|
||||
// the file.
|
||||
//
|
||||
// A report that SAYS false is different, and it revokes. The Windows helper
|
||||
// publishes the launch before it resumes the suspended child, so there is a
|
||||
// window where the fact is readable and the child has still executed nothing; if
|
||||
// the resume then fails, the helper retracts the record with an explicit false.
|
||||
// Without this, a poll that landed inside that window would hold a launch that
|
||||
// never happened, and hold it through completion, since the final read finds the
|
||||
// file cleaned away and restores what was latched.
|
||||
//
|
||||
// Which is why the observation is repeated while the command runs rather than
|
||||
// latched once. It costs one small read per poll on a wrapped plan, and every
|
||||
// unwrapped plan still does no extra work at all; a fact that can be withdrawn
|
||||
// is not one to cache.
|
||||
func (process *managedProcess) observeLaunch() {
|
||||
if process.report == nil {
|
||||
return
|
||||
}
|
||||
process.mu.Lock()
|
||||
skip := !process.ownedLaunch
|
||||
process.mu.Unlock()
|
||||
if skip || process.doneClosed() {
|
||||
return
|
||||
}
|
||||
report, err := process.report()
|
||||
if err != nil || report.ChildLaunched == nil {
|
||||
return
|
||||
}
|
||||
process.mu.Lock()
|
||||
process.launchObserved = *report.ChildLaunched
|
||||
process.mu.Unlock()
|
||||
}
|
||||
|
||||
func (process *managedProcess) collectResult(ctx context.Context, wait time.Duration, interrupted bool) ProcessResult {
|
||||
output, truncated := process.collect(ctx, wait)
|
||||
process.observeLaunch()
|
||||
process.mu.Lock()
|
||||
exitCode := 0
|
||||
exited := process.exitCode != nil
|
||||
@@ -473,8 +555,9 @@ func (process *managedProcess) collectResult(ctx context.Context, wait time.Dura
|
||||
ProcessID: process.id, CommandText: process.commandText, RelativeCwd: process.relativeCwd,
|
||||
TTY: process.tty, Output: output, OutputTruncated: truncated, Exited: exited,
|
||||
ExitCode: exitCode, Interrupted: interrupted, Request: process.request,
|
||||
Enforcement: process.enforcement, Report: process.resultReport, ReportErr: process.reportErr,
|
||||
Changes: append([]Change(nil), process.changes...), Metadata: cloneStringMap(process.metadata),
|
||||
Enforcement: process.enforcement, Report: process.launchedReportLocked(), ReportErr: process.reportErr,
|
||||
ChildLaunchOwnedByAdapter: process.ownedLaunch,
|
||||
Changes: append([]Change(nil), process.changes...), Metadata: cloneStringMap(process.metadata),
|
||||
}
|
||||
process.mu.Unlock()
|
||||
return result
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
package execution
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// A LAUNCH THE HELPER TAKES BACK MUST NOT SURVIVE IN EITHER RESULT.
|
||||
//
|
||||
// The Windows helper publishes childLaunched before it resumes the suspended
|
||||
// child, so the fact is readable while the child has still executed nothing. A
|
||||
// poll landing in that window latched it, and the latch outlived the file: the
|
||||
// terminal read finds the report cleaned away and restores what was observed, so
|
||||
// a child that never became runnable was reported as launched, and the write-jail
|
||||
// trade it never made was disclosed as if it had.
|
||||
//
|
||||
// Deleting the report on that path cannot fix it, because deletion is also what a
|
||||
// normal cleanup does. The helper retracts with an explicit false instead, and
|
||||
// this is the interleaving that pins it: publish, live read, retract, poll, then
|
||||
// finish and read terminally.
|
||||
func TestARetractedLaunchIsNotCommittedLiveOrTerminally(t *testing.T) {
|
||||
directory := t.TempDir()
|
||||
reportPath := filepath.Join(directory, "report.json")
|
||||
command := liveHelperCommand(t, reportPath, true)
|
||||
|
||||
manager := NewProcessManager(ProcessManagerOptions{})
|
||||
result, err := manager.Start(context.Background(), ProcessStart{
|
||||
Prepared: PreparedCommand{
|
||||
Command: command,
|
||||
ChildLaunchOwnedByAdapter: true,
|
||||
Enforcement: Enforcement{Notices: []string{liveLaunchNotice}},
|
||||
Report: liveReportReader(reportPath),
|
||||
},
|
||||
Request: liveRequest(t),
|
||||
}, 300*time.Millisecond)
|
||||
if err != nil {
|
||||
t.Fatalf("Start: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { manager.StopAll() })
|
||||
|
||||
// SETUP: the publication window. The live read has to have seen the positive,
|
||||
// or the revocation below would have nothing to revoke.
|
||||
if result.Exited {
|
||||
t.Fatal("SETUP INVALID: the stand-in helper exited, so the live lifecycle is not under test")
|
||||
}
|
||||
if !ResolveChildLaunched(true, result.ChildLaunchOwnedByAdapter, result.Report) {
|
||||
t.Fatal("SETUP INVALID: the published launch was not observed live, so this test cannot show it being taken back")
|
||||
}
|
||||
|
||||
// The resume fails: the helper retracts the record it published.
|
||||
if err := os.WriteFile(reportPath, []byte(`{"childLaunched":false}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
polled, err := manager.Continue(context.Background(), ProcessContinue{ProcessID: result.ProcessID, Wait: 200 * time.Millisecond})
|
||||
if err != nil {
|
||||
t.Fatalf("Continue: %v", err)
|
||||
}
|
||||
if polled.Exited {
|
||||
t.Fatal("SETUP INVALID: the helper exited before the poll, so the live poll is not under test")
|
||||
}
|
||||
if ResolveChildLaunched(true, polled.ChildLaunchOwnedByAdapter, polled.Report) {
|
||||
t.Error("a live poll still reports a launch the helper retracted, so the disclosure stands for a child that never ran")
|
||||
}
|
||||
|
||||
// Finish, with the plan's cleanup removing the report the way it always does.
|
||||
// The terminal read then finds absence, which is exactly the case the restore
|
||||
// exists for, and it must not resurrect the retracted launch.
|
||||
if !manager.Stop(result.ProcessID) {
|
||||
t.Fatal("Stop: process not found")
|
||||
}
|
||||
_ = os.Remove(reportPath)
|
||||
final, err := manager.Continue(context.Background(), ProcessContinue{ProcessID: result.ProcessID, Wait: time.Second})
|
||||
if err != nil {
|
||||
t.Fatalf("Continue after stop: %v", err)
|
||||
}
|
||||
if !final.Exited {
|
||||
t.Fatal("SETUP INVALID: the helper never exited, so there is no terminal result under test")
|
||||
}
|
||||
if ResolveChildLaunched(true, final.ChildLaunchOwnedByAdapter, final.Report) {
|
||||
t.Error("the final result reports a launch the helper retracted")
|
||||
}
|
||||
}
|
||||
|
||||
// AND A GENUINE LAUNCH STILL SURVIVES ITS OWN CLEANUP. The revocation must key
|
||||
// on an explicit denial and nothing else: a report removed by the plan's normal
|
||||
// cleanup is silence, and silence does not take back what was seen.
|
||||
func TestAConfirmedLaunchSurvivesTheReportCleanup(t *testing.T) {
|
||||
directory := t.TempDir()
|
||||
reportPath := filepath.Join(directory, "report.json")
|
||||
command := liveHelperCommand(t, reportPath, true)
|
||||
|
||||
manager := NewProcessManager(ProcessManagerOptions{})
|
||||
result, err := manager.Start(context.Background(), ProcessStart{
|
||||
Prepared: PreparedCommand{
|
||||
Command: command,
|
||||
ChildLaunchOwnedByAdapter: true,
|
||||
Enforcement: Enforcement{Notices: []string{liveLaunchNotice}},
|
||||
Report: liveReportReader(reportPath),
|
||||
},
|
||||
Request: liveRequest(t),
|
||||
}, 300*time.Millisecond)
|
||||
if err != nil {
|
||||
t.Fatalf("Start: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { manager.StopAll() })
|
||||
if !ResolveChildLaunched(true, result.ChildLaunchOwnedByAdapter, result.Report) {
|
||||
t.Fatal("SETUP INVALID: the published launch was not observed live")
|
||||
}
|
||||
|
||||
if !manager.Stop(result.ProcessID) {
|
||||
t.Fatal("Stop: process not found")
|
||||
}
|
||||
_ = os.Remove(reportPath)
|
||||
final, err := manager.Continue(context.Background(), ProcessContinue{ProcessID: result.ProcessID, Wait: time.Second})
|
||||
if err != nil {
|
||||
t.Fatalf("Continue after stop: %v", err)
|
||||
}
|
||||
if !final.Exited {
|
||||
t.Fatal("SETUP INVALID: the helper never exited")
|
||||
}
|
||||
if !ResolveChildLaunched(true, final.ChildLaunchOwnedByAdapter, final.Report) {
|
||||
t.Error("a confirmed launch was lost when the plan's cleanup removed its report")
|
||||
}
|
||||
}
|
||||
@@ -24,6 +24,12 @@ type PreparedCommand struct {
|
||||
Enforcement Enforcement
|
||||
Report func() (AdapterReport, error)
|
||||
Cleanup func()
|
||||
// ChildLaunchOwnedByAdapter marks a plan where Command is a WRAPPER and the
|
||||
// requested process is created inside it, so exec.Cmd.Process says nothing
|
||||
// about whether the sandboxed child ever existed. The adapter must state the
|
||||
// fact in its report; if it does not, the runner treats the child as not
|
||||
// launched rather than crediting the wrapper's start.
|
||||
ChildLaunchOwnedByAdapter bool
|
||||
}
|
||||
|
||||
type CapturedRequest struct {
|
||||
@@ -90,10 +96,25 @@ func (runner *Runner) ExecuteCaptured(ctx context.Context, input CapturedRequest
|
||||
prepared.Command.Stdout = stdout
|
||||
prepared.Command.Stderr = stderr
|
||||
runErr := prepared.Command.Run()
|
||||
// Observed HERE, from the only thing that knows: exec.Cmd sets Process only
|
||||
// once os.StartProcess has succeeded, so this is false for a missing
|
||||
// executable and for a context cancelled before Start, and true for anything
|
||||
// that ran, including a later timeout or cancellation.
|
||||
launched := prepared.Command.Process != nil
|
||||
report, reportErr := AdapterReport{}, error(nil)
|
||||
if prepared.Report != nil {
|
||||
report, reportErr = prepared.Report()
|
||||
}
|
||||
// A WRAPPED PLAN'S LAUNCH BIT BELONGS TO THE ADAPTER. The line above observes
|
||||
// the process THIS command started, which for a Windows restricted-token plan
|
||||
// is the helper, not the requested executable: the helper validates the setup
|
||||
// marker, applies ACLs, checks the network policy, builds capability SIDs and
|
||||
// mints the restricted token after it is already running, and any of those can
|
||||
// fail with no sandboxed child ever created. Believing the outer bit there
|
||||
// reports that reads were denied as requested when only the unsandboxed
|
||||
// adapter ran. An adapter that owns the inner transition overrides it; one
|
||||
// that stays silent leaves the direct-command observation alone.
|
||||
launched = ResolveChildLaunched(launched, prepared.ChildLaunchOwnedByAdapter, report)
|
||||
result := CapturedResult{
|
||||
Stdout: stdout.String(),
|
||||
Stderr: stderr.String(),
|
||||
@@ -101,6 +122,7 @@ func (runner *Runner) ExecuteCaptured(ctx context.Context, input CapturedRequest
|
||||
Err: runErr,
|
||||
Outcome: Outcome{
|
||||
Enforcement: prepared.Enforcement,
|
||||
Launched: launched,
|
||||
},
|
||||
}
|
||||
exitCode := commandExitCode(runErr)
|
||||
@@ -157,6 +179,28 @@ func (runner *Runner) Prepare(ctx context.Context, request Request) (PreparedCom
|
||||
return preparer.PrepareExecution(ctx, request)
|
||||
}
|
||||
|
||||
// ResolveChildLaunched decides whether the REQUESTED process launched.
|
||||
//
|
||||
// ONE IMPLEMENTATION, because every launcher needs the same answer and each one
|
||||
// that re-derived it got a different one. observed is what the caller saw of the
|
||||
// process IT started, which for a wrapped plan is the helper and not the
|
||||
// requested child.
|
||||
//
|
||||
// - the adapter stated the fact: believe the adapter, in both directions.
|
||||
// - the adapter owns the fact and stayed silent: not launched. An absent report
|
||||
// must not be read as proof that enforcement applied.
|
||||
// - nobody owns it but the caller: keep the direct observation, which is
|
||||
// correct for a direct command and for bwrap.
|
||||
func ResolveChildLaunched(observed bool, ownedByAdapter bool, report AdapterReport) bool {
|
||||
if report.ChildLaunched != nil {
|
||||
return *report.ChildLaunched
|
||||
}
|
||||
if ownedByAdapter {
|
||||
return false
|
||||
}
|
||||
return observed
|
||||
}
|
||||
|
||||
func capturedSetupFailure(message string, err error, enforcement Enforcement) CapturedResult {
|
||||
return CapturedResult{
|
||||
Stderr: message,
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
package execution
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// A WRAPPER'S START IS NOT THE REQUESTED CHILD'S START.
|
||||
//
|
||||
// For a Windows restricted-token plan the command the runner starts is the
|
||||
// sandbox helper, not the executable the caller asked for. Inside that helper,
|
||||
// setup-marker validation, unelevated ACL application, network-policy
|
||||
// validation, capability and offline SID construction, restricted-token creation
|
||||
// and the CreateProcessAsUser call all happen afterwards, and each of them can
|
||||
// return with no sandboxed child ever created. exec.Cmd.Process is already
|
||||
// non-nil by then, so reading the launch state off it reports that reads were
|
||||
// denied as requested when the only thing that ran was the unsandboxed adapter.
|
||||
//
|
||||
// The fact belongs to whoever sees the transition. These pin both directions of
|
||||
// that boundary.
|
||||
type wrappedPreparer struct {
|
||||
script string
|
||||
owned bool
|
||||
childLaunched *bool
|
||||
reportNothing bool
|
||||
}
|
||||
|
||||
func (p *wrappedPreparer) PrepareExecution(ctx context.Context, _ Request) (PreparedCommand, error) {
|
||||
prepared := PreparedCommand{
|
||||
Command: launchStateShell(ctx, p.script),
|
||||
Enforcement: Enforcement{Notices: []string{launchStateNotice}},
|
||||
ChildLaunchOwnedByAdapter: p.owned,
|
||||
}
|
||||
if !p.reportNothing {
|
||||
launched := p.childLaunched
|
||||
prepared.Report = func() (AdapterReport, error) {
|
||||
return AdapterReport{ChildLaunched: launched}, nil
|
||||
}
|
||||
}
|
||||
return prepared, nil
|
||||
}
|
||||
|
||||
func capturedWrapped(t *testing.T, p *wrappedPreparer) CapturedResult {
|
||||
t.Helper()
|
||||
return NewRunner(p).ExecuteCaptured(context.Background(), CapturedRequest{Request: Request{
|
||||
Origin: OriginHook,
|
||||
Mode: ModeCaptured,
|
||||
Command: Command{Name: "irrelevant"},
|
||||
WorkingDirectory: t.TempDir(),
|
||||
WorkspaceRoots: []string{t.TempDir()},
|
||||
Approval: ApprovalContext{PolicyVersion: PolicyVersion},
|
||||
}})
|
||||
}
|
||||
|
||||
func TestWrappedPlanDisclosesOnlyWhatTheAdapterConfirms(t *testing.T) {
|
||||
// The helper starts and then fails before it can create the restricted child:
|
||||
// a bad setup marker, an ACL it could not apply, a network policy it rejected,
|
||||
// a token it could not mint. The wrapper process exists; the sandboxed one
|
||||
// never did, so nothing may be claimed about enforcement.
|
||||
t.Run("helper ran but never created the child", func(t *testing.T) {
|
||||
no := false
|
||||
result := capturedWrapped(t, &wrappedPreparer{script: "exit 1", owned: true, childLaunched: &no})
|
||||
// The wrapper really did run, which is the whole point: its exit code is
|
||||
// the script's. Without this the test could pass because nothing executed.
|
||||
if result.Outcome.Exit == nil || result.Outcome.Exit.Code != 1 {
|
||||
t.Fatalf("SETUP INVALID: the wrapper itself must have run; outcome = %+v", result.Outcome)
|
||||
}
|
||||
if notices := result.Outcome.AppliedEnforcementNotices(); len(notices) != 0 {
|
||||
t.Fatalf("a helper that never created the restricted child disclosed %q", notices)
|
||||
}
|
||||
})
|
||||
|
||||
// Same shape, but the adapter says nothing at all. Silence from the owner of
|
||||
// the fact is not permission to fall back to the wrapper's own start.
|
||||
t.Run("adapter that owns the fact stayed silent", func(t *testing.T) {
|
||||
result := capturedWrapped(t, &wrappedPreparer{script: "exit 1", owned: true, reportNothing: true})
|
||||
if notices := result.Outcome.AppliedEnforcementNotices(); len(notices) != 0 {
|
||||
t.Fatalf("an unreported child launch was disclosed as applied enforcement: %q", notices)
|
||||
}
|
||||
})
|
||||
|
||||
// And the other side of the boundary: a restricted child that really started
|
||||
// and then exited non-zero DID run under the disclosed enforcement, so the
|
||||
// notice must still be made, exactly once.
|
||||
t.Run("restricted child started, then failed", func(t *testing.T) {
|
||||
yes := true
|
||||
result := capturedWrapped(t, &wrappedPreparer{script: "exit 3", owned: true, childLaunched: &yes})
|
||||
notices := result.Outcome.AppliedEnforcementNotices()
|
||||
if len(notices) != 1 || notices[0] != launchStateNotice {
|
||||
t.Fatalf("a child that ran and then failed disclosed %q, want exactly one %q", notices, launchStateNotice)
|
||||
}
|
||||
})
|
||||
|
||||
// A direct, unwrapped command is unchanged: the process the runner starts IS
|
||||
// the requested one, so its own observation still decides.
|
||||
t.Run("direct command keeps its own observation", func(t *testing.T) {
|
||||
result := capturedWrapped(t, &wrappedPreparer{script: "exit 0", owned: false, reportNothing: true})
|
||||
notices := result.Outcome.AppliedEnforcementNotices()
|
||||
if len(notices) != 1 || notices[0] != launchStateNotice {
|
||||
t.Fatalf("a direct command that ran disclosed %q, want exactly one %q", notices, launchStateNotice)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -36,7 +36,25 @@ type DispatchOutcome struct {
|
||||
// Messages collects the output (stdout, else stderr) of each hook that
|
||||
// produced any, in run order. afterTool validators use this to feed results
|
||||
// (e.g. a formatter diff or vet warning) back to the model on the tool result.
|
||||
//
|
||||
// PRESENTATION TEXT, NOT A NOTICE CHANNEL, and it no longer carries notices at
|
||||
// all. hookMessage used to fold them in so afterTool got both on one line;
|
||||
// that made the disclosure arrive twice once the typed slice was composed by
|
||||
// every surface. Enforcement disclosures are on Notices, for both hook phases.
|
||||
// A caller that only wants to know what the sandbox did must read that:
|
||||
// delivering this slice would put every successful hook's routine logging into
|
||||
// the model's context.
|
||||
Messages []string
|
||||
// Notices carries only the enforcement disclosures, one entry per notice, in
|
||||
// run order across every hook that ran.
|
||||
//
|
||||
// Separate from Messages because they answer different questions and have
|
||||
// different audiences. A notice says the hook ran under a weakened token,
|
||||
// which the model and the operator both need; the hook's own output is for
|
||||
// afterTool validators that asked to be heard. Appended as each hook runs,
|
||||
// rather than read off the final result, so a disclosure from a hook that
|
||||
// already ran survives a later hook's veto ending the chain.
|
||||
Notices []string
|
||||
}
|
||||
|
||||
type commandResult struct {
|
||||
@@ -45,6 +63,13 @@ type commandResult struct {
|
||||
Stderr string
|
||||
Err error // set when the command could not be executed (not a non-zero exit)
|
||||
TimedOut bool // the hook started but its deadline/cancellation fired before it returned
|
||||
// Notices carries the enforcement disclosures the execution runner attached.
|
||||
//
|
||||
// Same reason as the plugin path: the generic execution contract is not
|
||||
// transport-only. Enforcement.Notices says what the sandbox actually did, and
|
||||
// a projection that keeps only stdout, stderr and an exit code drops it, so a
|
||||
// hook ran under a weakened token with nothing said about it.
|
||||
Notices []string
|
||||
}
|
||||
|
||||
// commandRunner executes one hook command. It is injectable so the dispatch
|
||||
@@ -139,6 +164,10 @@ func executionCommandRunner(runner *execution.Runner) commandRunner {
|
||||
Stderr: stderr,
|
||||
Err: commandErr,
|
||||
TimedOut: result.Outcome.Kind == execution.OutcomeTimedOut,
|
||||
// One shared decision: see Outcome.AppliedEnforcementNotices. A setup
|
||||
// failure or a missing executable launched no hook child, so the notice
|
||||
// would describe a token trade nobody made.
|
||||
Notices: result.Outcome.AppliedEnforcementNotices(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -182,11 +211,19 @@ func (dispatcher *Dispatcher) Dispatch(ctx context.Context, input DispatchInput)
|
||||
if message := hookMessage(result); message != "" {
|
||||
outcome.Messages = append(outcome.Messages, message)
|
||||
}
|
||||
// Appended per hook rather than read off the final result, so hook A's
|
||||
// disclosure is not lost when hook B stops the chain. A notice describes
|
||||
// something that ALREADY happened.
|
||||
for _, notice := range result.Notices {
|
||||
if strings.TrimSpace(notice) != "" {
|
||||
outcome.Notices = append(outcome.Notices, notice)
|
||||
}
|
||||
}
|
||||
|
||||
if blocked {
|
||||
outcome.Blocked = true
|
||||
outcome.BlockedBy = hook.ID
|
||||
outcome.Reason = blockReason(result)
|
||||
outcome.Reason = blockCause(result)
|
||||
// Stop on the first veto: the action is already denied.
|
||||
return outcome
|
||||
}
|
||||
@@ -243,14 +280,38 @@ func classifyResult(event Event, result commandResult) (AuditStatus, bool) {
|
||||
|
||||
// hookMessage returns the output worth surfacing from a hook run: stdout when
|
||||
// present, else stderr. Empty when the hook produced no output.
|
||||
// STDOUT OR STDERR, AND NOTHING ELSE. The enforcement notices used to be
|
||||
// prepended here as well, so that afterTool got the disclosure and its own output
|
||||
// on one line. That was the only delivery available when this was written, and it
|
||||
// is not any more: the notices travel typed on DispatchOutcome.Notices, and the
|
||||
// agent loop merges them into ToolResult.EnforcementNotices for beforeTool and
|
||||
// afterTool alike.
|
||||
//
|
||||
// Keeping the fold as well made the same disclosure arrive twice, once from the
|
||||
// typed slice that every surface composes and once inside the hook feedback block
|
||||
// in the body. The two writers carry the identical fixed string, so the model saw
|
||||
// it doubled and a bash or exec card showed it in both the furniture and the body.
|
||||
// One owner per fact, and for a notice that owner is the typed slice.
|
||||
func hookMessage(result commandResult) string {
|
||||
if trimmed := strings.TrimSpace(result.Stdout); trimmed != "" {
|
||||
return trimmed
|
||||
message := strings.TrimSpace(result.Stdout)
|
||||
if message == "" {
|
||||
message = strings.TrimSpace(result.Stderr)
|
||||
}
|
||||
return strings.TrimSpace(result.Stderr)
|
||||
return message
|
||||
}
|
||||
|
||||
func blockReason(result commandResult) string {
|
||||
// blockCause explains a veto in the hook's own words, and says nothing about
|
||||
// enforcement.
|
||||
//
|
||||
// THE SAME RULE AS hookMessage. The blocking hook's notices used to be folded in
|
||||
// here as well, from before the typed slice existed, and Dispatch had already
|
||||
// appended them to DispatchOutcome.Notices, so a hook that vetoed while running
|
||||
// without write confinement put one disclosure on two channels. The one consumer
|
||||
// made up for it by dropping every notice whose text appeared anywhere in the
|
||||
// reason, which also dropped an EARLIER hook's disclosure whenever the blocking
|
||||
// hook's output happened to quote it. The notice now travels on the typed slice
|
||||
// only, whichever hook in the chain produced it. Reported by CodeRabbit.
|
||||
func blockCause(result commandResult) string {
|
||||
if result.TimedOut {
|
||||
if trimmed := strings.TrimSpace(result.Stderr); trimmed != "" {
|
||||
return "hook timed out: " + trimmed
|
||||
@@ -288,7 +349,14 @@ func (dispatcher *Dispatcher) recordCompleted(hook Definition, input DispatchInp
|
||||
Matcher: hook.Matcher,
|
||||
ToolCallID: input.ToolCallID,
|
||||
Status: status,
|
||||
Results: []AuditResult{{ExitCode: result.ExitCode, Stdout: result.Stdout, Stderr: result.Stderr}},
|
||||
Results: []AuditResult{{
|
||||
ExitCode: result.ExitCode,
|
||||
Stdout: result.Stdout,
|
||||
Stderr: result.Stderr,
|
||||
// The notice is not in stdout or stderr by design, so the durable record
|
||||
// has to carry it or the fact ends with the dispatch result.
|
||||
EnforcementNotices: append([]string(nil), result.Notices...),
|
||||
}},
|
||||
DurationMs: durationMs,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -244,8 +244,8 @@ func TestClassifyResultTimedOutFailsClosedForBeforeTool(t *testing.T) {
|
||||
if status, blocked := classifyResult(EventAfterTool, timedOut); blocked || status != AuditError {
|
||||
t.Fatalf("afterTool timeout classify = (%q, %v), want (error, false)", status, blocked)
|
||||
}
|
||||
if reason := blockReason(timedOut); !strings.Contains(reason, "timed out") {
|
||||
t.Fatalf("blockReason = %q, want a timeout message", reason)
|
||||
if reason := blockCause(timedOut); !strings.Contains(reason, "timed out") {
|
||||
t.Fatalf("blockCause = %q, want a timeout message", reason)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
package hooks
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// auditedDispatcher wires a real audit store to a dispatcher whose hook result
|
||||
// is whatever the caller wants, and returns the events that survived the write.
|
||||
func auditedDispatcher(t *testing.T, hook Definition, result commandResult) []AuditEvent {
|
||||
t.Helper()
|
||||
store, err := NewAuditStore(AuditStoreOptions{AuditPath: filepath.Join(t.TempDir(), "audit.jsonl")})
|
||||
if err != nil {
|
||||
t.Fatalf("NewAuditStore: %v", err)
|
||||
}
|
||||
dispatcher := NewDispatcher(DispatcherOptions{
|
||||
Config: beforeToolConfig(hook),
|
||||
Audit: store,
|
||||
run: func(context.Context, string, []string, []byte, string, []string) commandResult {
|
||||
return result
|
||||
},
|
||||
})
|
||||
dispatcher.Dispatch(context.Background(), DispatchInput{Event: EventBeforeTool, ToolName: "bash", ToolCallID: "call_1"})
|
||||
|
||||
// READ BACK FROM DISK, not from the in-memory event the append returned. The
|
||||
// durable reader is the consumer this field exists for.
|
||||
events, err := store.ReadEvents()
|
||||
if err != nil {
|
||||
t.Fatalf("ReadEvents: %v", err)
|
||||
}
|
||||
return events
|
||||
}
|
||||
|
||||
func completedResults(t *testing.T, events []AuditEvent) []AuditResult {
|
||||
t.Helper()
|
||||
for _, event := range events {
|
||||
if len(event.Results) > 0 {
|
||||
return event.Results
|
||||
}
|
||||
}
|
||||
t.Fatalf("no completed record was written: %#v", events)
|
||||
return nil
|
||||
}
|
||||
|
||||
// THE TRANSIENT DISPATCH RESULT IS NOT WHERE THIS FACT CAN LIVE.
|
||||
//
|
||||
// recordCompleted kept an exit code, stdout and stderr, and the notice is
|
||||
// deliberately in none of those. So once the dispatch result was gone, an audit
|
||||
// or recovery reader could not tell that a hook had run under the weakened
|
||||
// DenyRead token, whatever the hook did afterwards.
|
||||
func TestTheAuditRecordKeepsTheEnforcementNotice(t *testing.T) {
|
||||
const notice = "denyRead is configured, so the write jail is not confining writes"
|
||||
|
||||
for _, testCase := range []struct {
|
||||
name string
|
||||
result commandResult
|
||||
}{
|
||||
{"launched and succeeded", commandResult{ExitCode: 0, Stdout: "looks fine", Notices: []string{notice}}},
|
||||
{"vetoed the tool", commandResult{ExitCode: 2, Stderr: "policy violation", Notices: []string{notice}}},
|
||||
{"silent hook", commandResult{ExitCode: 0, Notices: []string{notice}}},
|
||||
} {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
results := completedResults(t, auditedDispatcher(t,
|
||||
Definition{ID: "policy", Event: EventBeforeTool, Command: "policy-check", Enabled: true},
|
||||
testCase.result))
|
||||
if len(results) != 1 {
|
||||
t.Fatalf("results = %#v, want one", results)
|
||||
}
|
||||
if len(results[0].EnforcementNotices) != 1 || results[0].EnforcementNotices[0] != notice {
|
||||
t.Errorf("the durable record lost the disclosure: %#v", results[0])
|
||||
}
|
||||
// The existing semantics are untouched.
|
||||
if results[0].ExitCode != testCase.result.ExitCode {
|
||||
t.Errorf("ExitCode = %d, want %d", results[0].ExitCode, testCase.result.ExitCode)
|
||||
}
|
||||
if results[0].Stdout != testCase.result.Stdout || results[0].Stderr != testCase.result.Stderr {
|
||||
t.Errorf("stdout/stderr changed: %#v", results[0])
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// A hook with nothing to disclose writes exactly what it wrote before, so a
|
||||
// reader of historical records sees no difference.
|
||||
func TestAnOrdinaryHookWritesNoEnforcementField(t *testing.T) {
|
||||
results := completedResults(t, auditedDispatcher(t,
|
||||
Definition{ID: "policy", Event: EventBeforeTool, Command: "policy-check", Enabled: true},
|
||||
commandResult{ExitCode: 0, Stdout: "looks fine"}))
|
||||
if len(results[0].EnforcementNotices) != 0 {
|
||||
t.Errorf("a hook with no disclosure recorded one: %#v", results[0])
|
||||
}
|
||||
}
|
||||
|
||||
// And the durable record inherits the launch-state rule rather than restating
|
||||
// it: a hook that never started records no enforcement claim.
|
||||
func TestTheAuditRecordMakesNoClaimForAHookThatNeverLaunched(t *testing.T) {
|
||||
store, err := NewAuditStore(AuditStoreOptions{AuditPath: filepath.Join(t.TempDir(), "audit.jsonl")})
|
||||
if err != nil {
|
||||
t.Fatalf("NewAuditStore: %v", err)
|
||||
}
|
||||
dispatcher := NewDispatcher(DispatcherOptions{
|
||||
Config: beforeToolConfig(Definition{ID: "policy", Event: EventBeforeTool, Command: "policy-check", Enabled: true}),
|
||||
Audit: store,
|
||||
Cwd: t.TempDir(),
|
||||
Execution: newRunnerFor(¬icePreparer{build: func() *exec.Cmd { return exec.Command("definitely-not-a-real-binary-zzz") }}),
|
||||
})
|
||||
dispatcher.Dispatch(context.Background(), DispatchInput{Event: EventBeforeTool, ToolName: "bash", ToolCallID: "call_1"})
|
||||
events, err := store.ReadEvents()
|
||||
if err != nil {
|
||||
t.Fatalf("ReadEvents: %v", err)
|
||||
}
|
||||
for _, result := range completedResults(t, events) {
|
||||
if len(result.EnforcementNotices) != 0 {
|
||||
t.Errorf("the durable record claims an enforcement trade for a hook that never started: %#v", result)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
//go:build !windows
|
||||
|
||||
package hooks
|
||||
|
||||
// sleepScript keeps a launched child alive long enough for a timeout to fire.
|
||||
const sleepScript = "sleep 2"
|
||||
@@ -0,0 +1,6 @@
|
||||
package hooks
|
||||
|
||||
// sleepScript keeps a launched child alive long enough for a timeout to fire.
|
||||
// Paired with the !windows file of the same name; both must exist or the
|
||||
// platform without one silently loses the timeout case.
|
||||
const sleepScript = "ping -n 3 127.0.0.1 > NUL"
|
||||
@@ -0,0 +1,184 @@
|
||||
package hooks
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os/exec"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Twigpine/zero/internal/execution"
|
||||
)
|
||||
|
||||
const launchStateNotice = "denyRead is configured, so the write jail is not confining writes"
|
||||
|
||||
// shellCommand builds a portable child that the platform can actually launch.
|
||||
func shellCommand(script string) *exec.Cmd {
|
||||
if runtime.GOOS == "windows" {
|
||||
return exec.Command("cmd.exe", "/c", script)
|
||||
}
|
||||
return exec.Command("/bin/sh", "-c", script)
|
||||
}
|
||||
|
||||
// noticePreparer plans a command carrying an enforcement notice, and can fail
|
||||
// the way the sandbox does before the child exists.
|
||||
type noticePreparer struct {
|
||||
prepareErr error
|
||||
build func() *exec.Cmd
|
||||
}
|
||||
|
||||
func (preparer *noticePreparer) PrepareExecution(_ context.Context, request execution.Request) (execution.PreparedCommand, error) {
|
||||
if preparer.prepareErr != nil {
|
||||
return execution.PreparedCommand{}, preparer.prepareErr
|
||||
}
|
||||
command := preparer.build
|
||||
if command == nil {
|
||||
command = func() *exec.Cmd { return exec.Command(request.Command.Name, request.Command.Args...) }
|
||||
}
|
||||
return execution.PreparedCommand{
|
||||
Command: command(),
|
||||
Enforcement: execution.Enforcement{Notices: []string{launchStateNotice}},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// PLANNING A WRAPPED COMMAND IS NOT PROOF THAT ANYTHING RAN.
|
||||
//
|
||||
// Enforcement.Notices describes the shape the command was PREPARED to run
|
||||
// under. Copying it straight out made the completed-enforcement claim for
|
||||
// commands that never existed: a sandbox setup failure and a missing executable
|
||||
// are both decided before the child launches, so the hook message told the
|
||||
// operator the write jail had been traded away for a process that never
|
||||
// started.
|
||||
//
|
||||
// Everything after launch keeps the disclosure, including a nonzero exit, a
|
||||
// timeout and a cancellation: those happened to a child that really did run
|
||||
// under that token.
|
||||
//
|
||||
// Driven through the execution runner rather than a hand-built commandResult,
|
||||
// because the projection is the thing under test.
|
||||
func TestTheHookRunnerOnlyDisclosesEnforcementForAChildThatLaunched(t *testing.T) {
|
||||
for _, testCase := range []struct {
|
||||
name string
|
||||
preparer *noticePreparer
|
||||
timeout time.Duration
|
||||
wantNotice bool
|
||||
wantTimedOut bool
|
||||
}{
|
||||
{
|
||||
name: "sandbox setup failed before the child existed",
|
||||
preparer: ¬icePreparer{prepareErr: errors.New("could not build the restricted token")},
|
||||
wantNotice: false,
|
||||
},
|
||||
{
|
||||
name: "the executable was never found",
|
||||
preparer: ¬icePreparer{build: func() *exec.Cmd {
|
||||
return exec.Command("definitely-not-a-real-binary-zzz")
|
||||
}},
|
||||
wantNotice: false,
|
||||
},
|
||||
{
|
||||
name: "the child launched and succeeded",
|
||||
preparer: ¬icePreparer{build: func() *exec.Cmd { return shellCommand("exit 0") }},
|
||||
wantNotice: true,
|
||||
},
|
||||
{
|
||||
name: "the child launched and exited nonzero",
|
||||
preparer: ¬icePreparer{build: func() *exec.Cmd { return shellCommand("exit 3") }},
|
||||
wantNotice: true,
|
||||
},
|
||||
{
|
||||
name: "the child launched and timed out",
|
||||
preparer: ¬icePreparer{build: func() *exec.Cmd { return shellCommand(sleepScript) }},
|
||||
timeout: 150 * time.Millisecond,
|
||||
wantNotice: true,
|
||||
wantTimedOut: true,
|
||||
},
|
||||
} {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
if testCase.timeout > 0 {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(ctx, testCase.timeout)
|
||||
defer cancel()
|
||||
}
|
||||
run := executionCommandRunner(execution.NewRunner(testCase.preparer))
|
||||
result := run(ctx, "hook-command", nil, nil, t.TempDir(), nil)
|
||||
|
||||
if result.TimedOut != testCase.wantTimedOut {
|
||||
t.Fatalf("TimedOut = %v, want %v: the case did not reach the outcome kind it is named for", result.TimedOut, testCase.wantTimedOut)
|
||||
}
|
||||
if got := len(result.Notices) > 0; got != testCase.wantNotice {
|
||||
t.Fatalf("notices present = %v, want %v: %#v", got, testCase.wantNotice, result.Notices)
|
||||
}
|
||||
// Asserted on Notices, which is where the disclosure travels. It used
|
||||
// to be folded into the hook message as well, and the agent loop now
|
||||
// merges Notices into the typed enforcement slice for both hook
|
||||
// phases, so carrying it in both places delivered it twice.
|
||||
notices := strings.Join(result.Notices, "\n")
|
||||
if testCase.wantNotice && !strings.Contains(notices, launchStateNotice) {
|
||||
t.Errorf("a launched child lost its disclosure:\n%s", notices)
|
||||
}
|
||||
if !testCase.wantNotice && strings.Contains(notices, launchStateNotice) {
|
||||
t.Errorf("a child that never launched claimed the token was traded away:\n%s", notices)
|
||||
}
|
||||
if message := hookMessage(result); strings.Contains(message, launchStateNotice) {
|
||||
t.Errorf("the disclosure also rode along in the hook message, so it reaches the model twice:\n%s", message)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// And the same rule has to hold on the veto path, which returns early and is what
|
||||
// the model actually sees. Both channels are checked: the notice travels on
|
||||
// Notices, and Reason must not grow a copy of it either.
|
||||
func TestAVetoingHookThatNeverLaunchedClaimsNoEnforcement(t *testing.T) {
|
||||
dispatcher := NewDispatcher(DispatcherOptions{
|
||||
Config: beforeToolConfig(Definition{ID: "policy", Event: EventBeforeTool, Command: "policy-check", Enabled: true}),
|
||||
Cwd: t.TempDir(),
|
||||
// A missing executable rather than a prepare error: a prepare error never
|
||||
// builds the PreparedCommand, so its outcome carries no planned notice and
|
||||
// the assertion below would hold with the launch gate deleted. This shape
|
||||
// plans the notice and then fails to launch.
|
||||
Execution: execution.NewRunner(¬icePreparer{build: func() *exec.Cmd {
|
||||
return exec.Command("definitely-not-a-real-binary-zzz")
|
||||
}}),
|
||||
})
|
||||
outcome := dispatcher.Dispatch(context.Background(), DispatchInput{Event: EventBeforeTool, ToolName: "bash"})
|
||||
if !outcome.Blocked {
|
||||
t.Fatal("SETUP INVALID: a beforeTool hook that could not run must fail closed, or the veto path is not exercised")
|
||||
}
|
||||
if len(outcome.Notices) != 0 {
|
||||
t.Errorf("the dispatch outcome claims an enforcement trade for a hook that never started: %#v", outcome.Notices)
|
||||
}
|
||||
if strings.Contains(outcome.Reason, launchStateNotice) {
|
||||
t.Errorf("the veto reason claims an enforcement trade for a hook that never started:\n%s", outcome.Reason)
|
||||
}
|
||||
}
|
||||
|
||||
// A launched hook still carries it all the way into the dispatch outcome, on the
|
||||
// typed slice and only there.
|
||||
func TestALaunchedHookCarriesTheNoticeIntoTheDispatchOutcome(t *testing.T) {
|
||||
dispatcher := NewDispatcher(DispatcherOptions{
|
||||
Config: beforeToolConfig(Definition{ID: "policy", Event: EventBeforeTool, Command: "policy-check", Enabled: true}),
|
||||
Cwd: t.TempDir(),
|
||||
Execution: execution.NewRunner(¬icePreparer{build: func() *exec.Cmd { return shellCommand("exit 2") }}),
|
||||
})
|
||||
outcome := dispatcher.Dispatch(context.Background(), DispatchInput{Event: EventBeforeTool, ToolName: "bash"})
|
||||
if !outcome.Blocked {
|
||||
t.Fatal("SETUP INVALID: the hook did not veto, so the reason path is not exercised")
|
||||
}
|
||||
if got := strings.Count(strings.Join(outcome.Notices, "\n"), launchStateNotice); got != 1 {
|
||||
t.Errorf("a hook that really ran under the weakened token disclosed it %d times on Notices, want once: %#v", got, outcome.Notices)
|
||||
}
|
||||
if strings.Contains(outcome.Reason, launchStateNotice) {
|
||||
t.Errorf("the veto reason carries a second copy of the notice:\n%s", outcome.Reason)
|
||||
}
|
||||
}
|
||||
|
||||
// newRunnerFor keeps the audit tests readable without importing the execution
|
||||
// package into every file that needs one.
|
||||
func newRunnerFor(preparer *noticePreparer) *execution.Runner {
|
||||
return execution.NewRunner(preparer)
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
package hooks
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// A HOOK THAT RAN UNDER THE WEAKENED TOKEN STILL SAYS SO, ON THE TYPED CHANNEL.
|
||||
//
|
||||
// The projection once kept stdout, stderr and an exit code and dropped the
|
||||
// notices, so such a hook ran silently. The first fix put them into the message
|
||||
// alongside the hook's own output, which delivered them but made Messages two
|
||||
// things at once. Now that the agent loop merges Notices into the typed
|
||||
// EnforcementNotices for beforeTool and afterTool alike, folding them into the
|
||||
// message as well delivered the same disclosure twice.
|
||||
//
|
||||
// So the property is unchanged and its carrier moved: whatever the hook printed,
|
||||
// the notice reaches Notices exactly once and never rides along in Messages.
|
||||
func TestAHookSurfacesTheEnforcementNoticeOnTheTypedChannel(t *testing.T) {
|
||||
const notice = "denyRead is configured, so the write jail is not confining writes"
|
||||
|
||||
for _, testCase := range []struct {
|
||||
name string
|
||||
result commandResult
|
||||
}{
|
||||
{"hook printed nothing", commandResult{ExitCode: 0, Notices: []string{notice}}},
|
||||
{"hook printed to stdout", commandResult{ExitCode: 0, Stdout: "looks fine", Notices: []string{notice}}},
|
||||
{"hook printed to stderr only", commandResult{ExitCode: 0, Stderr: "a warning", Notices: []string{notice}}},
|
||||
} {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
result := testCase.result
|
||||
dispatcher := NewDispatcher(DispatcherOptions{
|
||||
Config: beforeToolConfig(Definition{ID: "vet", Event: EventAfterTool, Command: "vet", Enabled: true}),
|
||||
run: func(context.Context, string, []string, []byte, string, []string) commandResult {
|
||||
return result
|
||||
},
|
||||
})
|
||||
outcome := dispatcher.Dispatch(context.Background(), DispatchInput{Event: EventAfterTool, ToolName: "bash"})
|
||||
|
||||
if got := strings.Count(strings.Join(outcome.Notices, "\n"), notice); got != 1 {
|
||||
t.Fatalf("the notice reached Notices %d times, want once: %v", got, outcome.Notices)
|
||||
}
|
||||
// AND NOT IN THE PROSE AS WELL, which is what made it arrive twice.
|
||||
if joined := strings.Join(outcome.Messages, "\n"); strings.Contains(joined, notice) {
|
||||
t.Errorf("the notice also rode along in Messages, so every surface that composes the typed slice shows it twice:\n%s", joined)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// A hook's own output is untouched, with a notice or without one.
|
||||
func TestAHookMessageCarriesOnlyTheHooksOwnOutput(t *testing.T) {
|
||||
const notice = "denyRead is configured, so the write jail is not confining writes"
|
||||
if message := hookMessage(commandResult{ExitCode: 0, Stdout: "looks fine"}); message != "looks fine" {
|
||||
t.Errorf("hookMessage = %q, want the hook's own output untouched", message)
|
||||
}
|
||||
if message := hookMessage(commandResult{ExitCode: 0, Stdout: "looks fine", Notices: []string{notice}}); message != "looks fine" {
|
||||
t.Errorf("hookMessage = %q, want the notice left to the typed channel", message)
|
||||
}
|
||||
if message := hookMessage(commandResult{ExitCode: 0}); message != "" {
|
||||
t.Errorf("a silent hook with no notice produced %q", message)
|
||||
}
|
||||
if message := hookMessage(commandResult{ExitCode: 0, Notices: []string{notice}}); message != "" {
|
||||
t.Errorf("a silent hook produced %q, want nothing: its disclosure travels typed", message)
|
||||
}
|
||||
}
|
||||
|
||||
// THROUGH Dispatch, NOT A HAND-BUILT commandResult.
|
||||
//
|
||||
// The blocking branch returns immediately, so it never touches hookMessage, and a
|
||||
// vetoing beforeTool hook that ran without write confinement once reported only
|
||||
// the veto. The disclosure still has to come out of that branch, and it has to
|
||||
// come out ONCE: the blocking hook's notices used to be folded into Reason as well
|
||||
// as appended to Notices, so the same fact travelled on two channels and the
|
||||
// consumer had to guess which copy to drop. Reason is the hook's words, and the
|
||||
// notice travels on the typed slice like every other hook's.
|
||||
func TestABlockedBeforeToolHookCarriesItsNoticeOnTheTypedSliceOnly(t *testing.T) {
|
||||
const notice = "denyRead is configured, so the write jail is not confining writes"
|
||||
|
||||
dispatcher := NewDispatcher(DispatcherOptions{
|
||||
Config: beforeToolConfig(Definition{ID: "policy", Event: EventBeforeTool, Command: "policy-check", Enabled: true}),
|
||||
run: func(context.Context, string, []string, []byte, string, []string) commandResult {
|
||||
return commandResult{ExitCode: 2, Stderr: "policy violation", Notices: []string{notice}}
|
||||
},
|
||||
})
|
||||
|
||||
outcome := dispatcher.Dispatch(context.Background(), DispatchInput{Event: EventBeforeTool, ToolName: "bash"})
|
||||
if !outcome.Blocked {
|
||||
t.Fatal("SETUP INVALID: the hook did not block, so the blocking branch was never taken")
|
||||
}
|
||||
if len(outcome.Notices) != 1 || outcome.Notices[0] != notice {
|
||||
t.Errorf("the veto lost the enforcement notice from the typed slice: %#v", outcome.Notices)
|
||||
}
|
||||
if outcome.Reason != "policy violation" {
|
||||
t.Errorf("Reason = %q, want the hook's own explanation and nothing else", outcome.Reason)
|
||||
}
|
||||
// Across BOTH channels, because a copy in each is the defect this pins.
|
||||
if got := strings.Count(outcome.Reason, notice) + strings.Count(strings.Join(outcome.Notices, "\n"), notice); got != 1 {
|
||||
t.Errorf("the notice appears %d times across Reason and Notices, want once", got)
|
||||
}
|
||||
}
|
||||
|
||||
// And a veto with no notice reads exactly as it did before.
|
||||
func TestABlockedHookWithoutANoticeIsUnchanged(t *testing.T) {
|
||||
dispatcher := NewDispatcher(DispatcherOptions{
|
||||
Config: beforeToolConfig(Definition{ID: "policy", Event: EventBeforeTool, Command: "policy-check", Enabled: true}),
|
||||
run: func(context.Context, string, []string, []byte, string, []string) commandResult {
|
||||
return commandResult{ExitCode: 2, Stderr: "policy violation"}
|
||||
},
|
||||
})
|
||||
outcome := dispatcher.Dispatch(context.Background(), DispatchInput{Event: EventBeforeTool, ToolName: "bash"})
|
||||
if outcome.Reason != "policy violation" {
|
||||
t.Errorf("Reason = %q, want the hook's own explanation untouched", outcome.Reason)
|
||||
}
|
||||
}
|
||||
+15
-2
@@ -118,6 +118,19 @@ type AuditResult struct {
|
||||
ExitCode int `json:"exitCode"`
|
||||
Stdout string `json:"stdout,omitempty"`
|
||||
Stderr string `json:"stderr,omitempty"`
|
||||
// EnforcementNotices are the least-privilege disclosures that were true of
|
||||
// this hook's execution.
|
||||
//
|
||||
// A DURABLE READER CANNOT RECOVER A FACT THAT WAS DROPPED IN CONVERSION. The
|
||||
// notice is deliberately not written into stdout or stderr, so recording only
|
||||
// those three fields meant that once the transient dispatch result was gone,
|
||||
// nothing could tell an audit or recovery reader that a successful, failing or
|
||||
// vetoing hook had run under the weakened DenyRead token.
|
||||
//
|
||||
// Typed rather than a rendered line, and omitempty, so historical records that
|
||||
// predate the field read back unchanged and an ordinary hook writes exactly
|
||||
// what it wrote before.
|
||||
EnforcementNotices []string `json:"enforcementNotices,omitempty"`
|
||||
}
|
||||
|
||||
type AuditEvent struct {
|
||||
@@ -779,7 +792,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, afterTool, specialistStart, and specialistStop hooks."}
|
||||
return Definition{}, manifestError{fieldPath: field + ".matcher", message: "matcher is only supported for beforeTool and afterTool hooks."}
|
||||
}
|
||||
command, err := requiredString(obj, field+".command")
|
||||
if err != nil {
|
||||
@@ -843,7 +856,7 @@ func matchesHookMatcher(matcher string, toolName string) bool {
|
||||
}
|
||||
|
||||
func eventSupportsMatcher(event Event) bool {
|
||||
return event == EventBeforeTool || event == EventAfterTool || event == EventSpecialistStart || event == EventSpecialistStop
|
||||
return event == EventBeforeTool || event == EventAfterTool
|
||||
}
|
||||
|
||||
func toDiagnostic(err error, source ConfigSource, path string) Diagnostic {
|
||||
|
||||
@@ -180,7 +180,7 @@ func TestLoadConfigPreservesUserDisabledStateWhenProjectOmitsEnabled(t *testing.
|
||||
}
|
||||
|
||||
func TestLoadConfigRejectsMatchersOnLifecycleHooks(t *testing.T) {
|
||||
for _, event := range []string{"sessionStart", "sessionEnd"} {
|
||||
for _, event := range []string{"sessionStart", "specialistStart"} {
|
||||
t.Run(event, func(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
projectConfigPath := filepath.Join(dir, "hooks.json")
|
||||
@@ -210,46 +210,6 @@ 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})
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/config"
|
||||
"github.com/Twigpine/zero/internal/execution"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// helperCommandName is a command that resolves on this platform, so registration
|
||||
// reaches connectStdio. What actually runs is whatever the preparer returns.
|
||||
func helperCommandName() string {
|
||||
if runtime.GOOS == "windows" {
|
||||
return "cmd.exe"
|
||||
}
|
||||
return "sh"
|
||||
}
|
||||
|
||||
const adapterLaunchNotice = "denyRead is configured, so the write jail is not confining writes"
|
||||
|
||||
// adapterHelperPreparer models the Windows wrapped plan: the command is the
|
||||
// HELPER, the launch fact is a report file, and the plan's cleanup removes that
|
||||
// file. The cleanup is the part that matters, because it is what runs inside
|
||||
// client.Close and destroys the evidence a later decision needs.
|
||||
type adapterHelperPreparer struct {
|
||||
reportPath string
|
||||
reportBody string
|
||||
writeReport bool
|
||||
called bool
|
||||
}
|
||||
|
||||
func (preparer *adapterHelperPreparer) PrepareExecution(_ context.Context, _ execution.Request) (execution.PreparedCommand, error) {
|
||||
preparer.called = true
|
||||
if preparer.writeReport {
|
||||
if err := os.WriteFile(preparer.reportPath, []byte(preparer.reportBody), 0o600); err != nil {
|
||||
return execution.PreparedCommand{}, err
|
||||
}
|
||||
}
|
||||
// A helper that starts, says nothing an MCP client understands, and exits, so
|
||||
// the handshake fails the way it does when the requested server never existed.
|
||||
var command *exec.Cmd
|
||||
if runtime.GOOS == "windows" {
|
||||
command = exec.Command("cmd.exe", "/c", "exit 0")
|
||||
} else {
|
||||
command = exec.Command("/bin/sh", "-c", "exit 0")
|
||||
}
|
||||
path := preparer.reportPath
|
||||
return execution.PreparedCommand{
|
||||
Command: command,
|
||||
ChildLaunchOwnedByAdapter: true,
|
||||
Enforcement: execution.Enforcement{Notices: []string{adapterLaunchNotice}},
|
||||
Report: func() (execution.AdapterReport, error) {
|
||||
raw, err := os.ReadFile(path)
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
return execution.AdapterReport{}, nil
|
||||
}
|
||||
if err != nil {
|
||||
return execution.AdapterReport{}, err
|
||||
}
|
||||
var report execution.AdapterReport
|
||||
if err := json.Unmarshal(raw, &report); err != nil {
|
||||
return execution.AdapterReport{}, err
|
||||
}
|
||||
return report, nil
|
||||
},
|
||||
Cleanup: func() { _ = os.Remove(path) },
|
||||
}, nil
|
||||
}
|
||||
|
||||
func registerWithAdapterHelper(t *testing.T, preparer *adapterHelperPreparer) *Runtime {
|
||||
t.Helper()
|
||||
runtime, err := RegisterTools(context.Background(), tools.NewRegistry(), config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"docs": {Type: "stdio", Command: helperCommandName()},
|
||||
}}, RegisterOptions{
|
||||
// No ClientFactory on purpose: an injected factory skips connectStdio
|
||||
// entirely, which is where the decision under test is made, and the test
|
||||
// would pass with the fix removed.
|
||||
Execution: execution.NewRunner(preparer),
|
||||
WorkspaceRoot: t.TempDir(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("RegisterTools: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = runtime.Close() })
|
||||
return runtime
|
||||
}
|
||||
|
||||
// A HELPER THAT NEVER CREATED THE SERVER HAS NOTHING TO DISCLOSE.
|
||||
//
|
||||
// For an adapter-owned launch, cmd.Start proves only that the sandbox helper
|
||||
// started. It can then fail setup-marker validation, ACL application, network
|
||||
// validation, token construction, or CreateProcessAsUser without ever creating
|
||||
// the requested MCP server. The launch sink already made that distinction; the
|
||||
// initialize-error path did not, and carried the planned notices out
|
||||
// unconditionally. The operator was told the server had run without write
|
||||
// confinement when no server had run at all.
|
||||
func TestAnMCPHelperThatReportedNoChildDisclosesNothing(t *testing.T) {
|
||||
directory := t.TempDir()
|
||||
preparer := &adapterHelperPreparer{
|
||||
reportPath: filepath.Join(directory, "report.json"),
|
||||
reportBody: `{"childLaunched":false}`,
|
||||
writeReport: true,
|
||||
}
|
||||
runtime := registerWithAdapterHelper(t, preparer)
|
||||
|
||||
// SETUP: the attempt really did fail, or there is no disclosure decision here.
|
||||
if len(runtime.Skipped()) == 0 {
|
||||
t.Fatal("SETUP INVALID: the server connected, so the initialize-failure path is not under test")
|
||||
}
|
||||
if got := runtime.StartupDisclosures(); len(got) != 0 {
|
||||
t.Fatalf("a helper that reported no child announced %v; no server ran, confined or otherwise", got)
|
||||
}
|
||||
}
|
||||
|
||||
// AND ONE THAT DID CREATE IT DISCLOSES ONCE.
|
||||
//
|
||||
// The companion case, and the one that keeps the assertion above from being
|
||||
// satisfied by a path that discloses nothing ever. The child ran under the
|
||||
// planned token and may have done filesystem work before the handshake failed,
|
||||
// so the disclosure has to survive the failure.
|
||||
func TestAnMCPHelperThatLaunchedTheChildDisclosesOnce(t *testing.T) {
|
||||
directory := t.TempDir()
|
||||
preparer := &adapterHelperPreparer{
|
||||
reportPath: filepath.Join(directory, "report.json"),
|
||||
reportBody: `{"childLaunched":true}`,
|
||||
writeReport: true,
|
||||
}
|
||||
runtime := registerWithAdapterHelper(t, preparer)
|
||||
|
||||
if !preparer.called {
|
||||
t.Fatal("SETUP INVALID: the preparer never ran, so connectStdio was never reached")
|
||||
}
|
||||
disclosures := runtime.StartupDisclosures()
|
||||
if len(disclosures) != 1 {
|
||||
t.Fatalf("a server that really ran under the weakened token produced %d disclosures, want exactly one: %v", len(disclosures), disclosures)
|
||||
}
|
||||
if len(disclosures[0].Notices) != 1 || disclosures[0].Notices[0] != adapterLaunchNotice {
|
||||
t.Fatalf("the disclosure carries %v, want the one planned notice", disclosures[0].Notices)
|
||||
}
|
||||
}
|
||||
|
||||
// A helper that wrote no report at all is the same answer as one that reported
|
||||
// no child: absence is not confirmation.
|
||||
func TestAnMCPHelperThatWroteNoReportDisclosesNothing(t *testing.T) {
|
||||
directory := t.TempDir()
|
||||
preparer := &adapterHelperPreparer{reportPath: filepath.Join(directory, "report.json")}
|
||||
runtime := registerWithAdapterHelper(t, preparer)
|
||||
|
||||
if len(runtime.Skipped()) == 0 {
|
||||
t.Fatal("SETUP INVALID: the server connected, so the initialize-failure path is not under test")
|
||||
}
|
||||
if got := runtime.StartupDisclosures(); len(got) != 0 {
|
||||
t.Fatalf("a helper that published nothing announced %v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,281 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Twigpine/zero/internal/config"
|
||||
"github.com/Twigpine/zero/internal/execution"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
const (
|
||||
adapterHelperModeEnv = "ZERO_TEST_MCP_ADAPTER_MODE"
|
||||
adapterHelperReportEnv = "ZERO_TEST_MCP_ADAPTER_REPORT"
|
||||
|
||||
// answerThenPublish models the ordering the fix is about: the child is
|
||||
// serving before the adapter has recorded that it exists.
|
||||
answerThenPublish = "answer-then-publish"
|
||||
// failThenPublish is the same ordering on the other exit: the handshake dies
|
||||
// first and the adapter publishes on its way out.
|
||||
failThenPublish = "fail-then-publish"
|
||||
// failAndNeverPublish is the companion that keeps the two above from being
|
||||
// satisfied by disclosing unconditionally.
|
||||
failAndNeverPublish = "fail-and-never-publish"
|
||||
)
|
||||
|
||||
// TestAdapterHelperProcess is the helper process body, not a test.
|
||||
//
|
||||
// It exists so the report can be published at a controlled point RELATIVE TO THE
|
||||
// HANDSHAKE. A fixture that writes the finished report during PrepareExecution,
|
||||
// as the older test does, publishes before the helper command has even started,
|
||||
// so every ordering this file is about is already over before the parent looks.
|
||||
func TestAdapterHelperProcess(t *testing.T) {
|
||||
mode := os.Getenv(adapterHelperModeEnv)
|
||||
if mode == "" {
|
||||
t.Skip("not the helper process")
|
||||
}
|
||||
reportPath := os.Getenv(adapterHelperReportEnv)
|
||||
publish := func() {
|
||||
// Same shape the Windows helper writes, into the file the preparer already
|
||||
// created empty.
|
||||
_ = os.WriteFile(reportPath, []byte(`{"childLaunched":true}`), 0o600)
|
||||
}
|
||||
|
||||
switch mode {
|
||||
case answerThenPublish:
|
||||
serveOneMCPSessionThenPublish(publish)
|
||||
case failThenPublish:
|
||||
// The handshake dies here, before anything is recorded. Closing stdout is
|
||||
// the child going away; the adapter is still running.
|
||||
_ = os.Stdout.Close()
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
publish()
|
||||
case failAndNeverPublish:
|
||||
_ = os.Stdout.Close()
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
}
|
||||
// Before the framework can write anything to the pipe the parent is reading.
|
||||
os.Exit(0)
|
||||
}
|
||||
|
||||
// serveOneMCPSessionThenPublish answers the handshake and only then records that
|
||||
// the child exists, which is the interleaving that used to lose the disclosure.
|
||||
func serveOneMCPSessionThenPublish(publish func()) {
|
||||
reader := bufio.NewReader(os.Stdin)
|
||||
out := bufio.NewWriter(os.Stdout)
|
||||
published := false
|
||||
for {
|
||||
line, err := reader.ReadString('\n')
|
||||
if strings.TrimSpace(line) != "" {
|
||||
var message struct {
|
||||
ID *int `json:"id"`
|
||||
Method string `json:"method"`
|
||||
}
|
||||
if json.Unmarshal([]byte(strings.TrimSpace(line)), &message) == nil && message.ID != nil {
|
||||
var result string
|
||||
switch message.Method {
|
||||
case "initialize":
|
||||
result = `{"protocolVersion":"2024-11-05"}`
|
||||
case "tools/list":
|
||||
result = `{"tools":[]}`
|
||||
default:
|
||||
result = `{}`
|
||||
}
|
||||
fmt.Fprintf(out, `{"jsonrpc":"2.0","id":%d,"result":%s}`+"\n", *message.ID, result)
|
||||
_ = out.Flush()
|
||||
if message.Method == "initialize" && !published {
|
||||
// AFTER the response is on the wire. The parent can act on a
|
||||
// handshake that succeeded before this line runs.
|
||||
time.Sleep(300 * time.Millisecond)
|
||||
publish()
|
||||
published = true
|
||||
}
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// livePublishPreparer models the Windows wrapped plan with the real publication
|
||||
// ordering: the report file is created EMPTY before the launch, exactly as
|
||||
// openWindowsExecutionReport does, and the helper fills it in later.
|
||||
type livePublishPreparer struct {
|
||||
mode string
|
||||
reportPath string
|
||||
called bool
|
||||
}
|
||||
|
||||
func (preparer *livePublishPreparer) PrepareExecution(_ context.Context, _ execution.Request) (execution.PreparedCommand, error) {
|
||||
preparer.called = true
|
||||
// The empty file the parent can see before anything is published. This is the
|
||||
// state that is neither "no child" nor "child created".
|
||||
if err := os.WriteFile(preparer.reportPath, nil, 0o600); err != nil {
|
||||
return execution.PreparedCommand{}, err
|
||||
}
|
||||
command := exec.Command(os.Args[0], "-test.run=^TestAdapterHelperProcess$")
|
||||
command.Env = append(os.Environ(),
|
||||
adapterHelperModeEnv+"="+preparer.mode,
|
||||
adapterHelperReportEnv+"="+preparer.reportPath,
|
||||
)
|
||||
path := preparer.reportPath
|
||||
return execution.PreparedCommand{
|
||||
Command: command,
|
||||
ChildLaunchOwnedByAdapter: true,
|
||||
Enforcement: execution.Enforcement{Notices: []string{adapterLaunchNotice}},
|
||||
Report: func() (execution.AdapterReport, error) {
|
||||
raw, err := os.ReadFile(path)
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
return execution.AdapterReport{}, nil
|
||||
}
|
||||
if err != nil {
|
||||
return execution.AdapterReport{}, err
|
||||
}
|
||||
var report execution.AdapterReport
|
||||
// An empty or half-written file decodes to an error, which is the
|
||||
// unsettled state and not an answer.
|
||||
if err := json.Unmarshal(raw, &report); err != nil {
|
||||
return execution.AdapterReport{}, err
|
||||
}
|
||||
return report, nil
|
||||
},
|
||||
Cleanup: func() { _ = os.Remove(path) },
|
||||
}, nil
|
||||
}
|
||||
|
||||
func registerWithLivePublisher(t *testing.T, mode string) (*Runtime, *livePublishPreparer) {
|
||||
t.Helper()
|
||||
preparer := &livePublishPreparer{mode: mode, reportPath: filepath.Join(t.TempDir(), "report.json")}
|
||||
runtime, err := RegisterTools(context.Background(), tools.NewRegistry(), config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"docs": {Type: "stdio", Command: helperCommandName()},
|
||||
}}, RegisterOptions{
|
||||
Execution: execution.NewRunner(preparer),
|
||||
WorkspaceRoot: t.TempDir(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("RegisterTools: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = runtime.Close() })
|
||||
if !preparer.called {
|
||||
t.Fatal("SETUP INVALID: the preparer never ran, so connectStdio was never reached")
|
||||
}
|
||||
return runtime, preparer
|
||||
}
|
||||
|
||||
// A HANDSHAKE THAT SUCCEEDED IS PROOF THE CHILD RAN, WHATEVER THE REPORT SAYS YET.
|
||||
//
|
||||
// The child is created before the adapter records it, and it inherits the MCP
|
||||
// pipes, so it can answer initialize while the report file is still the empty one
|
||||
// the adapter opened. Reading it at that instant and remembering the answer
|
||||
// turned "not yet" into "never" for the rest of the session, and the operator was
|
||||
// never told the server serving these tools ran without the write jail.
|
||||
//
|
||||
// The adapter speaks no MCP, so a well-formed response can only have come from
|
||||
// the requested child. That is terminal evidence and needs no report.
|
||||
func TestADisclosureSurvivesAReportPublishedAfterTheHandshake(t *testing.T) {
|
||||
runtime, _ := registerWithLivePublisher(t, answerThenPublish)
|
||||
|
||||
// SETUP: the server really connected, or this is the failure path instead.
|
||||
if skipped := runtime.Skipped(); len(skipped) != 0 {
|
||||
t.Fatalf("SETUP INVALID: the server did not connect (%v), so the successful-handshake path is not under test", skipped)
|
||||
}
|
||||
disclosures := runtime.StartupDisclosures()
|
||||
if len(disclosures) != 1 {
|
||||
t.Fatalf("a server that answered the handshake produced %d disclosures, want exactly one: %v", len(disclosures), disclosures)
|
||||
}
|
||||
if len(disclosures[0].Notices) != 1 || disclosures[0].Notices[0] != adapterLaunchNotice {
|
||||
t.Fatalf("the disclosure carries %v, want the one planned notice", disclosures[0].Notices)
|
||||
}
|
||||
}
|
||||
|
||||
// AND ON THE OTHER EXIT, THE DECISION WAITS FOR THE ADAPTER.
|
||||
//
|
||||
// Here the handshake dies first and the adapter publishes on its way out. The
|
||||
// decision used to be taken before Close, precisely because Close deletes the
|
||||
// report, so it read the empty file and answered "no child" about a server that
|
||||
// really had run. Settling from inside cleanup, with the file still there, is what
|
||||
// lets the answer be taken after the adapter is terminal instead of before it.
|
||||
func TestADisclosureSurvivesAReportPublishedAfterAFailedHandshake(t *testing.T) {
|
||||
runtime, _ := registerWithLivePublisher(t, failThenPublish)
|
||||
|
||||
// SETUP: the attempt really failed, or the successful path is being measured.
|
||||
if len(runtime.Skipped()) == 0 {
|
||||
t.Fatal("SETUP INVALID: the server connected, so the initialize-failure path is not under test")
|
||||
}
|
||||
disclosures := runtime.StartupDisclosures()
|
||||
if len(disclosures) != 1 {
|
||||
t.Fatalf("a helper that published its launch on the way out produced %d disclosures, want exactly one: %v", len(disclosures), disclosures)
|
||||
}
|
||||
|
||||
// THROUGH THE ERROR, not only through the sink. Registration merges the two,
|
||||
// so asserting the sink alone passes while the failure the operator actually
|
||||
// reads carries nothing. This is the carrier that exists because the client
|
||||
// holding the notices has already been closed and discarded by then.
|
||||
skipped := runtime.Skipped()
|
||||
if len(skipped) != 1 {
|
||||
t.Fatalf("skipped = %v, want exactly the one server", skipped)
|
||||
}
|
||||
carried := startupNoticesFromError(skipped[0].Err)
|
||||
if len(carried) != 1 || carried[0] != adapterLaunchNotice {
|
||||
t.Fatalf("the initialize failure carries %v, want the planned notice; the decision was taken before the adapter settled", carried)
|
||||
}
|
||||
}
|
||||
|
||||
// AND AN ADAPTER THAT NEVER PUBLISHED STILL DISCLOSES NOTHING.
|
||||
//
|
||||
// The companion that keeps both cases above from being satisfied by announcing
|
||||
// unconditionally. Settling reads the report one last time and finds nothing,
|
||||
// which after the adapter has exited is the answer rather than a race.
|
||||
func TestNoDisclosureWhenTheAdapterExitsWithoutPublishing(t *testing.T) {
|
||||
runtime, _ := registerWithLivePublisher(t, failAndNeverPublish)
|
||||
|
||||
if len(runtime.Skipped()) == 0 {
|
||||
t.Fatal("SETUP INVALID: the server connected, so the initialize-failure path is not under test")
|
||||
}
|
||||
if got := runtime.StartupDisclosures(); len(got) != 0 {
|
||||
t.Fatalf("a helper that published nothing announced %v; no server is known to have run", got)
|
||||
}
|
||||
}
|
||||
|
||||
// THE CARRIER ENFORCES THE RULE ITSELF, NOT ONLY ITS CALLERS.
|
||||
//
|
||||
// connectAndList reads StartupNotices on the success path and passes the result
|
||||
// straight through to the disclosure sources, so this method is a carrier of the
|
||||
// launch fact in its own right. Today the handshake confirmation makes the gate
|
||||
// redundant: every route that reaches here has already settled the decision
|
||||
// positive. It is kept because the redundancy is on the safe side. An edit that
|
||||
// moves or loses the confirmation makes this return nothing rather than announce
|
||||
// a confinement on the strength of the helper having started, which is the claim
|
||||
// this whole mechanism exists to stop making.
|
||||
//
|
||||
// Driven directly, because no path through connectStdio can reach it with a
|
||||
// negative decision, and a test that cannot construct the state it is about would
|
||||
// be asserting nothing.
|
||||
func TestStartupNoticesAreEmptyWhileTheLaunchIsUnknown(t *testing.T) {
|
||||
client := &Client{startupNotices: []string{adapterLaunchNotice}}
|
||||
|
||||
// SETUP: ungated, this client discloses, or the assertion below is vacuous.
|
||||
if len(client.StartupNotices()) != 1 {
|
||||
t.Fatal("SETUP INVALID: the client discloses nothing even before the gate, so the gate cannot be what is under test")
|
||||
}
|
||||
|
||||
client.launched = func() bool { return false }
|
||||
if got := client.StartupNotices(); len(got) != 0 {
|
||||
t.Fatalf("a client whose launch is not established carries %v", got)
|
||||
}
|
||||
client.launched = func() bool { return true }
|
||||
if got := client.StartupNotices(); len(got) != 1 || got[0] != adapterLaunchNotice {
|
||||
t.Fatalf("an established launch carries %v, want the planned notice", got)
|
||||
}
|
||||
}
|
||||
@@ -1,287 +0,0 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/Twigpine/zero/internal/aimlapi"
|
||||
"github.com/Twigpine/zero/internal/config"
|
||||
"github.com/Twigpine/zero/internal/redaction"
|
||||
)
|
||||
|
||||
// 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, and for the length of the
|
||||
// call in the collector its context carries (aimlapiSentKeys), so the reply can
|
||||
// be scrubbed: never in Server, its Identity (OAuth token binding), config.json,
|
||||
// list output, errors or tool output. Every key the call sent, redirect hops
|
||||
// included, and on the SSE transport the keys the stream request sent, is
|
||||
// scrubbed from the response's error text, method and result before anything
|
||||
// reads them (scrubAIMLAPIKeys) and from any error the call returns, a
|
||||
// transport error quoting a URL the server chose among them
|
||||
// (scrubAIMLAPIKeyError). An error status to such a request is also 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)
|
||||
if sent, ok := request.Context().Value(aimlapiSentKeysContextKey{}).(*aimlapiSentKeys); ok {
|
||||
sent.add(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
|
||||
}
|
||||
|
||||
// aimlapiSentKeys collects every key applyAIMLAPIRequestHeaders puts on the
|
||||
// requests of one call. The first request's context carries it, and net/http
|
||||
// gives each redirect hop that same context, so a key resolved again for a later
|
||||
// hop is collected too. The final hop's key alone is not enough: a same-origin
|
||||
// server answering the last hop has seen every earlier one.
|
||||
type aimlapiSentKeys struct {
|
||||
mu sync.Mutex
|
||||
keys []string
|
||||
}
|
||||
|
||||
type aimlapiSentKeysContextKey struct{}
|
||||
|
||||
// withAIMLAPISentKeys returns ctx carrying a fresh collector for one call.
|
||||
func withAIMLAPISentKeys(ctx context.Context) (context.Context, *aimlapiSentKeys) {
|
||||
sent := &aimlapiSentKeys{}
|
||||
return context.WithValue(ctx, aimlapiSentKeysContextKey{}, sent), sent
|
||||
}
|
||||
|
||||
func (sent *aimlapiSentKeys) add(key string) {
|
||||
sent.mu.Lock()
|
||||
defer sent.mu.Unlock()
|
||||
for _, existing := range sent.keys {
|
||||
if existing == key {
|
||||
return
|
||||
}
|
||||
}
|
||||
sent.keys = append(sent.keys, key)
|
||||
}
|
||||
|
||||
// list returns a copy of the keys collected so far.
|
||||
func (sent *aimlapiSentKeys) list() []string {
|
||||
sent.mu.Lock()
|
||||
defer sent.mu.Unlock()
|
||||
return append([]string(nil), sent.keys...)
|
||||
}
|
||||
|
||||
// scrubAIMLAPIKeys removes the injected keys from a response before any of it
|
||||
// is read: the error text, the method name and every string in the result,
|
||||
// object member names included. A server can echo the credential it was sent,
|
||||
// and MCP reports a failed tool call as a result rather than an error, so the
|
||||
// result needs it as much as the error text. Only the known values go, split
|
||||
// forms included (redaction.RedactSecretValues); a result that holds none of
|
||||
// them keeps its original bytes.
|
||||
func scrubAIMLAPIKeys(message rpcMessage, keys ...string) rpcMessage {
|
||||
var sent []string
|
||||
for _, key := range keys {
|
||||
if strings.TrimSpace(key) != "" {
|
||||
sent = append(sent, key)
|
||||
}
|
||||
}
|
||||
if len(sent) == 0 {
|
||||
return message
|
||||
}
|
||||
if message.Error != nil {
|
||||
scrubbed := *message.Error
|
||||
scrubbed.Message = redaction.RedactSecretValues(scrubbed.Message, sent...)
|
||||
message.Error = &scrubbed
|
||||
}
|
||||
message.Method = redaction.RedactSecretValues(message.Method, sent...)
|
||||
message.Result = scrubAIMLAPIKeysFromJSON(message.Result, sent)
|
||||
return message
|
||||
}
|
||||
|
||||
func scrubAIMLAPIKeysFromJSON(raw json.RawMessage, keys []string) json.RawMessage {
|
||||
if len(raw) == 0 {
|
||||
return raw
|
||||
}
|
||||
decoder := json.NewDecoder(bytes.NewReader(raw))
|
||||
decoder.UseNumber()
|
||||
var value any
|
||||
if err := decoder.Decode(&value); err != nil {
|
||||
// Decoding it into the caller's type fails the same way, and that
|
||||
// error does not quote the text.
|
||||
return raw
|
||||
}
|
||||
scrubbed, changed := scrubAIMLAPIKeysFromValue(value, keys)
|
||||
if !changed {
|
||||
return raw
|
||||
}
|
||||
encoded, err := json.Marshal(scrubbed)
|
||||
if err != nil {
|
||||
// Never hand back the unscrubbed bytes: an empty result reads as
|
||||
// nothing rather than as the key.
|
||||
return nil
|
||||
}
|
||||
return encoded
|
||||
}
|
||||
|
||||
func scrubAIMLAPIKeysFromValue(value any, keys []string) (any, bool) {
|
||||
switch typed := value.(type) {
|
||||
case string:
|
||||
scrubbed := redaction.RedactSecretValues(typed, keys...)
|
||||
return scrubbed, scrubbed != typed
|
||||
case []any:
|
||||
changed := false
|
||||
for index, item := range typed {
|
||||
if scrubbed, itemChanged := scrubAIMLAPIKeysFromValue(item, keys); itemChanged {
|
||||
typed[index] = scrubbed
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
return typed, changed
|
||||
case map[string]any:
|
||||
changed := false
|
||||
scrubbed := make(map[string]any, len(typed))
|
||||
for name, item := range typed {
|
||||
scrubbedName := redaction.RedactSecretValues(name, keys...)
|
||||
scrubbedItem, itemChanged := scrubAIMLAPIKeysFromValue(item, keys)
|
||||
if itemChanged || scrubbedName != name {
|
||||
changed = true
|
||||
}
|
||||
scrubbed[scrubbedName] = scrubbedItem
|
||||
}
|
||||
return scrubbed, changed
|
||||
}
|
||||
return value, false
|
||||
}
|
||||
|
||||
// scrubAIMLAPIKeyError returns err with the injected keys removed from its text.
|
||||
// An error the keys do not appear in keeps its chain. One they do appear in is
|
||||
// replaced by its scrubbed text, because its chain still holds the key (a
|
||||
// *url.Error keeps the URL the server chose); the replacement unwraps only to
|
||||
// the context sentinel the original matched, so a cancellation still reads as
|
||||
// one.
|
||||
func scrubAIMLAPIKeyError(err error, keys ...string) error {
|
||||
if err == nil || len(keys) == 0 {
|
||||
return err
|
||||
}
|
||||
text := err.Error()
|
||||
scrubbed := redaction.RedactSecretValues(text, keys...)
|
||||
if scrubbed == text {
|
||||
return err
|
||||
}
|
||||
replacement := &aimlapiScrubbedError{text: scrubbed}
|
||||
for _, sentinel := range []error{context.Canceled, context.DeadlineExceeded} {
|
||||
if errors.Is(err, sentinel) {
|
||||
replacement.cause = sentinel
|
||||
break
|
||||
}
|
||||
}
|
||||
return replacement
|
||||
}
|
||||
|
||||
// aimlapiScrubbedError is an error whose text had an injected key removed.
|
||||
type aimlapiScrubbedError struct {
|
||||
text string
|
||||
cause error
|
||||
}
|
||||
|
||||
func (scrubbed *aimlapiScrubbedError) Error() string { return scrubbed.text }
|
||||
|
||||
func (scrubbed *aimlapiScrubbedError) Unwrap() error { return scrubbed.cause }
|
||||
|
||||
// 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)
|
||||
}
|
||||
@@ -1,658 +0,0 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Twigpine/zero/internal/aimlapi"
|
||||
)
|
||||
|
||||
// scrubTestKey has no credential shape, so only the known-value scrub can find
|
||||
// it; a shape pattern would leave it alone.
|
||||
const scrubTestKey = "aiml0123456789abcdef0123456789abcd"
|
||||
|
||||
// aimlapiEcho stands in for mcp.aimlapi.com over streamable HTTP. It answers
|
||||
// each JSON-RPC request with reply(method, id), as a JSON body or as one SSE
|
||||
// event, and records the Authorization header every request carried.
|
||||
type aimlapiEcho struct {
|
||||
sse bool
|
||||
reply func(method string, id json.RawMessage) string
|
||||
mu sync.Mutex
|
||||
auth []string
|
||||
}
|
||||
|
||||
func (echo *aimlapiEcho) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
var body []byte
|
||||
if request.Body != nil {
|
||||
body, _ = io.ReadAll(request.Body)
|
||||
_ = request.Body.Close()
|
||||
}
|
||||
echo.mu.Lock()
|
||||
echo.auth = append(echo.auth, request.Header.Get("Authorization"))
|
||||
echo.mu.Unlock()
|
||||
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
|
||||
}
|
||||
reply := echo.reply(message.Method, message.ID)
|
||||
response := fakeHTTPResponse(request, http.StatusOK, reply)
|
||||
if echo.sse {
|
||||
response.Header.Set("Content-Type", "text/event-stream")
|
||||
response.Body = io.NopCloser(strings.NewReader("event: message\ndata: " + reply + "\n\n"))
|
||||
}
|
||||
return response, nil
|
||||
}
|
||||
|
||||
// assertSent fails unless every request carried want as its Authorization, so a
|
||||
// test cannot pass by never reaching the injected-key path.
|
||||
func (echo *aimlapiEcho) assertSent(t *testing.T, want string) {
|
||||
t.Helper()
|
||||
echo.mu.Lock()
|
||||
defer echo.mu.Unlock()
|
||||
if len(echo.auth) == 0 {
|
||||
t.Fatal("the fake server was never called")
|
||||
}
|
||||
for _, got := range echo.auth {
|
||||
if got != want {
|
||||
t.Fatalf("Authorization = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func rpcReply(id json.RawMessage, field string, value any) string {
|
||||
encoded, err := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": id, field: value})
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return string(encoded)
|
||||
}
|
||||
|
||||
// echoedKeys are the spellings of the key a server echoes back: whole, and
|
||||
// split by a NUL and by a zero-width space, which a plain replace would miss.
|
||||
func echoedKeys() []string {
|
||||
return []string{
|
||||
scrubTestKey,
|
||||
scrubTestKey[:10] + string(rune(0)) + scrubTestKey[10:],
|
||||
scrubTestKey[:20] + string(rune(0x200b)) + scrubTestKey[20:],
|
||||
}
|
||||
}
|
||||
|
||||
// assertNoPartOf fails if any eight-byte run of key survives in text.
|
||||
func assertNoPartOf(t *testing.T, text string, key string) {
|
||||
t.Helper()
|
||||
for start := 0; start+8 <= len(key); start++ {
|
||||
if strings.Contains(text, key[start:start+8]) {
|
||||
t.Fatalf("text still holds part of the key (%q): %q", key[start:start+8], text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// assertKeyScrubbed fails if any part of scrubTestKey survives in text, or if
|
||||
// text does not show the replacement.
|
||||
func assertKeyScrubbed(t *testing.T, text string) {
|
||||
t.Helper()
|
||||
assertNoPartOf(t, text, scrubTestKey)
|
||||
if !strings.Contains(text, "[REDACTED]") {
|
||||
t.Fatalf("text = %q, want the key replaced by [REDACTED]", text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScrubAIMLAPIKeysKeepsAResultWithoutTheKeyByteForByte(t *testing.T) {
|
||||
raw := json.RawMessage(`{"tools": [ {"name": "b", "description": "password=hunter2"} ], "a": 1.50}`)
|
||||
if got := scrubAIMLAPIKeysFromJSON(raw, []string{scrubTestKey}); string(got) != string(raw) {
|
||||
t.Fatalf("result without the key = %s, want the original bytes %s", got, raw)
|
||||
}
|
||||
|
||||
withKey := json.RawMessage(`{"content": [ {"type": "text", "text": "key ` + scrubTestKey + `"} ], "n": 1.50}`)
|
||||
got := scrubAIMLAPIKeysFromJSON(withKey, []string{scrubTestKey})
|
||||
assertKeyScrubbed(t, string(got))
|
||||
var decoded struct {
|
||||
Content []Content `json:"content"`
|
||||
N json.Number `json:"n"`
|
||||
}
|
||||
if err := json.Unmarshal(got, &decoded); err != nil || len(decoded.Content) != 1 || decoded.Content[0].Text != "key [REDACTED]" || decoded.N != "1.50" {
|
||||
t.Fatalf("scrubbed result = %s (%v), want valid JSON with only the key replaced and the number as written", got, err)
|
||||
}
|
||||
}
|
||||
|
||||
func newEchoNetworkClient(server Server, echo *aimlapiEcho) *networkClient {
|
||||
return &networkClient{server: server, client: mcpHTTPClient(server, echo), nextID: 1}
|
||||
}
|
||||
|
||||
func TestAIMLAPIMCPScrubsTheInjectedKeyFromResponses(t *testing.T) {
|
||||
isolateAIMLAPIKeyLookup(t, scrubTestKey)
|
||||
server := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL)
|
||||
for _, sse := range []bool{false, true} {
|
||||
transport := "json"
|
||||
if sse {
|
||||
transport = "sse"
|
||||
}
|
||||
t.Run(transport+"/error message", func(t *testing.T) {
|
||||
for _, echoed := range echoedKeys() {
|
||||
echo := &aimlapiEcho{sse: sse, reply: func(_ string, id json.RawMessage) string {
|
||||
return rpcReply(id, "error", map[string]any{"code": -32001, "message": "invalid credential " + echoed + " for this account"})
|
||||
}}
|
||||
err := newEchoNetworkClient(server, echo).initialize(context.Background())
|
||||
if err == nil {
|
||||
t.Fatal("initialize succeeded against a server that answered with an error")
|
||||
}
|
||||
echo.assertSent(t, "Bearer "+scrubTestKey)
|
||||
assertKeyScrubbed(t, err.Error())
|
||||
if !strings.Contains(err.Error(), "invalid credential [REDACTED] for this account") {
|
||||
t.Fatalf("error = %q, want the rest of the server text kept", err)
|
||||
}
|
||||
}
|
||||
})
|
||||
t.Run(transport+"/tool result", func(t *testing.T) {
|
||||
const image = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
|
||||
for _, echoed := range echoedKeys() {
|
||||
echo := &aimlapiEcho{sse: sse, reply: func(_ string, id json.RawMessage) string {
|
||||
return rpcReply(id, "result", map[string]any{
|
||||
"isError": true,
|
||||
"content": []any{
|
||||
map[string]any{"type": "text", "text": "upstream rejected key " + echoed},
|
||||
map[string]any{"type": "image", "mimeType": "image/png", "data": image},
|
||||
},
|
||||
})
|
||||
}}
|
||||
result, err := newEchoNetworkClient(server, echo).CallTool(context.Background(), "generate", map[string]any{"prompt": "x"})
|
||||
if err != nil {
|
||||
t.Fatalf("CallTool: %v", err)
|
||||
}
|
||||
echo.assertSent(t, "Bearer "+scrubTestKey)
|
||||
if len(result.Content) != 2 || !result.IsError {
|
||||
t.Fatalf("result = %+v, want the error result with both blocks", result)
|
||||
}
|
||||
assertKeyScrubbed(t, result.Content[0].Text)
|
||||
if result.Content[0].Text != "upstream rejected key [REDACTED]" {
|
||||
t.Fatalf("text block = %q, want only the key replaced", result.Content[0].Text)
|
||||
}
|
||||
if result.Content[1].Data != image {
|
||||
t.Fatalf("image data changed: %q", result.Content[1].Data)
|
||||
}
|
||||
}
|
||||
})
|
||||
t.Run(transport+"/result without the key", func(t *testing.T) {
|
||||
// Shapes stay as the server wrote them: only the known value goes.
|
||||
const description = "Set password=hunter2 or pass Authorization: Bearer other-token"
|
||||
echo := &aimlapiEcho{sse: sse, reply: func(_ string, id json.RawMessage) string {
|
||||
return rpcReply(id, "result", map[string]any{"tools": []any{map[string]any{
|
||||
"name": "generate",
|
||||
"description": description,
|
||||
"inputSchema": map[string]any{"type": "object", "properties": map[string]any{"api_key": map[string]any{"type": "string"}}},
|
||||
}}})
|
||||
}}
|
||||
tools, err := newEchoNetworkClient(server, echo).ListTools(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("ListTools: %v", err)
|
||||
}
|
||||
echo.assertSent(t, "Bearer "+scrubTestKey)
|
||||
if len(tools) != 1 || tools[0].Description != description {
|
||||
t.Fatalf("tools = %+v, want the description unchanged", tools)
|
||||
}
|
||||
properties, _ := tools[0].InputSchema["properties"].(map[string]any)
|
||||
if _, ok := properties["api_key"]; !ok {
|
||||
t.Fatalf("inputSchema = %v, want the api_key property kept", tools[0].InputSchema)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("json/unexpected method", func(t *testing.T) {
|
||||
echo := &aimlapiEcho{reply: func(string, json.RawMessage) string {
|
||||
return `{"jsonrpc":"2.0","method":"notifications/` + scrubTestKey + `"}`
|
||||
}}
|
||||
err := newEchoNetworkClient(server, echo).initialize(context.Background())
|
||||
if err == nil || !strings.Contains(err.Error(), "got method") {
|
||||
t.Fatalf("initialize error = %v, want the unexpected method reported", err)
|
||||
}
|
||||
echo.assertSent(t, "Bearer "+scrubTestKey)
|
||||
assertKeyScrubbed(t, err.Error())
|
||||
})
|
||||
}
|
||||
|
||||
func TestAIMLAPIMCPLeavesResponsesAloneWithoutTheInjectedKey(t *testing.T) {
|
||||
isolateAIMLAPIKeyLookup(t, scrubTestKey)
|
||||
// A user-configured Authorization header means nothing was injected, so the
|
||||
// server text is reported as written; this change scrubs only what Zero
|
||||
// itself attached.
|
||||
server := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL)
|
||||
server.Headers = map[string]string{"Authorization": "Bearer " + scrubTestKey}
|
||||
echo := &aimlapiEcho{reply: func(_ string, id json.RawMessage) string {
|
||||
return rpcReply(id, "error", map[string]any{"code": -32001, "message": "invalid credential " + scrubTestKey})
|
||||
}}
|
||||
err := newEchoNetworkClient(server, echo).initialize(context.Background())
|
||||
echo.assertSent(t, "Bearer "+scrubTestKey)
|
||||
if err == nil || !strings.Contains(err.Error(), "invalid credential "+scrubTestKey) {
|
||||
t.Fatalf("error = %v, want the server text unchanged when no key was injected", err)
|
||||
}
|
||||
}
|
||||
|
||||
// aimlapiSSEEcho stands in for mcp.aimlapi.com over the legacy SSE transport:
|
||||
// the GET opens a stream that starts with the endpoint event, and each POST is
|
||||
// answered on that stream with reply(method, id).
|
||||
type aimlapiSSEEcho struct {
|
||||
endpoint string
|
||||
reply func(method string, id json.RawMessage) string
|
||||
// postErr, when set, fails every POST at the transport.
|
||||
postErr error
|
||||
// rawEvent, when set, writes its whole event to the stream instead of a
|
||||
// message event built from reply.
|
||||
rawEvent func(method string, id json.RawMessage) string
|
||||
events chan string
|
||||
mu sync.Mutex
|
||||
auth []string
|
||||
}
|
||||
|
||||
func newAIMLAPISSEEcho(endpoint string, reply func(string, json.RawMessage) string) *aimlapiSSEEcho {
|
||||
return &aimlapiSSEEcho{endpoint: endpoint, reply: reply, events: make(chan string, 8)}
|
||||
}
|
||||
|
||||
func (echo *aimlapiSSEEcho) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
echo.mu.Lock()
|
||||
echo.auth = append(echo.auth, request.Header.Get("Authorization"))
|
||||
echo.mu.Unlock()
|
||||
if request.Method == http.MethodGet {
|
||||
reader, writer := io.Pipe()
|
||||
go func() {
|
||||
defer writer.Close()
|
||||
if _, err := io.WriteString(writer, "event: endpoint\ndata: "+echo.endpoint+"\n\n"); err != nil {
|
||||
return
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case event := <-echo.events:
|
||||
if _, err := io.WriteString(writer, event); err != nil {
|
||||
return
|
||||
}
|
||||
case <-request.Context().Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
response := fakeHTTPResponse(request, http.StatusOK, "")
|
||||
response.Header.Set("Content-Type", "text/event-stream")
|
||||
response.Body = reader
|
||||
return response, nil
|
||||
}
|
||||
if echo.postErr != nil {
|
||||
return nil, echo.postErr
|
||||
}
|
||||
var body []byte
|
||||
if request.Body != nil {
|
||||
body, _ = io.ReadAll(request.Body)
|
||||
_ = request.Body.Close()
|
||||
}
|
||||
var message struct {
|
||||
ID json.RawMessage `json:"id"`
|
||||
Method string `json:"method"`
|
||||
}
|
||||
_ = json.Unmarshal(body, &message)
|
||||
if len(message.ID) > 0 {
|
||||
if echo.rawEvent != nil {
|
||||
echo.events <- echo.rawEvent(message.Method, message.ID)
|
||||
} else {
|
||||
echo.events <- "event: message\ndata: " + echo.reply(message.Method, message.ID) + "\n\n"
|
||||
}
|
||||
}
|
||||
return fakeHTTPResponse(request, http.StatusAccepted, ""), nil
|
||||
}
|
||||
|
||||
// rotateAIMLAPIKeys makes the key resolver hand out keys in order, repeating the
|
||||
// last, the way a provider change between two requests of one call would.
|
||||
func rotateAIMLAPIKeys(t *testing.T, keys ...string) {
|
||||
t.Helper()
|
||||
var mu sync.Mutex
|
||||
previous := loadAIMLAPIServiceKey
|
||||
loadAIMLAPIServiceKey = func() (string, error) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
key := keys[0]
|
||||
if len(keys) > 1 {
|
||||
keys = keys[1:]
|
||||
}
|
||||
return key, nil
|
||||
}
|
||||
t.Cleanup(func() { loadAIMLAPIServiceKey = previous })
|
||||
}
|
||||
|
||||
func (echo *aimlapiSSEEcho) assertSent(t *testing.T, want string, requests int) {
|
||||
t.Helper()
|
||||
echo.mu.Lock()
|
||||
defer echo.mu.Unlock()
|
||||
if len(echo.auth) < requests {
|
||||
t.Fatalf("the fake server saw %d requests, want at least %d", len(echo.auth), requests)
|
||||
}
|
||||
for _, got := range echo.auth {
|
||||
if got != want {
|
||||
t.Fatalf("Authorization = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func newEchoSSEClient(echo *aimlapiSSEEcho) *remoteSSEClient {
|
||||
server := Server{Name: "aimlapi", Type: ServerTypeSSE, URL: aimlapi.MCPServerURL}
|
||||
return &remoteSSEClient{
|
||||
server: server,
|
||||
client: mcpHTTPClient(server, echo),
|
||||
nextID: 1,
|
||||
pending: map[string]chan ssePendingResponse{},
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIMLAPIMCPSSEScrubsTheInjectedKeyFromStreamAnswers(t *testing.T) {
|
||||
isolateAIMLAPIKeyLookup(t, scrubTestKey)
|
||||
for _, echoed := range echoedKeys() {
|
||||
echo := newAIMLAPISSEEcho("/messages", func(method string, id json.RawMessage) string {
|
||||
if method == "tools/call" {
|
||||
return rpcReply(id, "result", map[string]any{"isError": true, "content": []any{map[string]any{"type": "text", "text": "rejected " + echoed}}})
|
||||
}
|
||||
return rpcReply(id, "error", map[string]any{"code": -32001, "message": "invalid credential " + echoed})
|
||||
})
|
||||
client := newEchoSSEClient(echo)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
if err := client.openStream(ctx); err != nil {
|
||||
cancel()
|
||||
t.Fatalf("openStream: %v", err)
|
||||
}
|
||||
_, listErr := client.ListTools(ctx)
|
||||
result, callErr := client.CallTool(ctx, "generate", map[string]any{})
|
||||
_ = client.Close()
|
||||
cancel()
|
||||
|
||||
echo.assertSent(t, "Bearer "+scrubTestKey, 3)
|
||||
if listErr == nil {
|
||||
t.Fatal("ListTools succeeded against a server that answered with an error")
|
||||
}
|
||||
assertKeyScrubbed(t, listErr.Error())
|
||||
if callErr != nil || len(result.Content) != 1 {
|
||||
t.Fatalf("CallTool = %+v, %v, want the error result", result, callErr)
|
||||
}
|
||||
assertKeyScrubbed(t, result.Content[0].Text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIMLAPIMCPSSEScrubsTheStreamKeyAndThePostKey(t *testing.T) {
|
||||
// The key is resolved for every request, so the stream and a later POST can
|
||||
// carry different keys when the provider changes in between. An answer on
|
||||
// the stream comes from a server that was sent both.
|
||||
const streamKey = "strm9f8e7d6c5b4a39281706f5e4d3c2b1"
|
||||
const postKey = "post1a2b3c4d5e6f708192a3b4c5d6e7f8"
|
||||
rotateAIMLAPIKeys(t, streamKey, postKey)
|
||||
|
||||
echo := newAIMLAPISSEEcho("/messages", func(_ string, id json.RawMessage) string {
|
||||
return rpcReply(id, "error", map[string]any{"code": -32001, "message": "stream key " + streamKey + ", post key " + postKey})
|
||||
})
|
||||
client := newEchoSSEClient(echo)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
if err := client.openStream(ctx); err != nil {
|
||||
t.Fatalf("openStream: %v", err)
|
||||
}
|
||||
_, err := client.ListTools(ctx)
|
||||
_ = client.Close()
|
||||
|
||||
echo.mu.Lock()
|
||||
auth := append([]string(nil), echo.auth...)
|
||||
echo.mu.Unlock()
|
||||
if len(auth) != 2 || auth[0] != "Bearer "+streamKey || auth[1] != "Bearer "+postKey {
|
||||
t.Fatalf("Authorization headers = %q, want the stream key and then the post key", auth)
|
||||
}
|
||||
if err == nil {
|
||||
t.Fatal("ListTools succeeded against a server that answered with an error")
|
||||
}
|
||||
assertNoPartOf(t, err.Error(), streamKey)
|
||||
assertNoPartOf(t, err.Error(), postKey)
|
||||
if !strings.Contains(err.Error(), "stream key [REDACTED], post key [REDACTED]") {
|
||||
t.Fatalf("error = %q, want both keys replaced", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIMLAPIMCPSSEScrubsTheInjectedKeyFromAnEndpointError(t *testing.T) {
|
||||
isolateAIMLAPIKeyLookup(t, scrubTestKey)
|
||||
// An invalid escape makes the endpoint fail to parse, and the parse error
|
||||
// quotes the endpoint the server sent.
|
||||
echo := newAIMLAPISSEEcho("/messages/"+scrubTestKey+"%zz", func(string, json.RawMessage) string { return "" })
|
||||
client := newEchoSSEClient(echo)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
err := client.openStream(ctx)
|
||||
_ = client.Close()
|
||||
echo.assertSent(t, "Bearer "+scrubTestKey, 1)
|
||||
if err == nil || !strings.Contains(err.Error(), "endpoint") {
|
||||
t.Fatalf("openStream error = %v, want the endpoint parse failure", err)
|
||||
}
|
||||
assertKeyScrubbed(t, err.Error())
|
||||
}
|
||||
|
||||
// aimlapiHopEcho answers the first request with a 307 to location and hands the
|
||||
// second hop to second, recording the Authorization each hop carried.
|
||||
type aimlapiHopEcho struct {
|
||||
location string
|
||||
second func(request *http.Request, id json.RawMessage) (*http.Response, error)
|
||||
mu sync.Mutex
|
||||
auth []string
|
||||
}
|
||||
|
||||
func (echo *aimlapiHopEcho) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
var body []byte
|
||||
if request.Body != nil {
|
||||
body, _ = io.ReadAll(request.Body)
|
||||
_ = request.Body.Close()
|
||||
}
|
||||
echo.mu.Lock()
|
||||
echo.auth = append(echo.auth, request.Header.Get("Authorization"))
|
||||
hop := len(echo.auth)
|
||||
echo.mu.Unlock()
|
||||
if hop == 1 {
|
||||
response := fakeHTTPResponse(request, http.StatusTemporaryRedirect, "")
|
||||
response.Header.Set("Location", echo.location)
|
||||
return response, nil
|
||||
}
|
||||
var message struct {
|
||||
ID json.RawMessage `json:"id"`
|
||||
}
|
||||
_ = json.Unmarshal(body, &message)
|
||||
return echo.second(request, message.ID)
|
||||
}
|
||||
|
||||
func TestAIMLAPIMCPScrubsEveryKeySentAcrossRedirectHops(t *testing.T) {
|
||||
// The key is resolved again for a redirect hop, so the hops of one call can
|
||||
// carry different keys, and the same-origin server answering the last hop
|
||||
// has seen them all.
|
||||
const first = "aimlfirst0123456789abcdef01234567"
|
||||
const second = "aimlsecond0123456789abcdef0123456"
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
location string
|
||||
// hopKey is the key the redirect hop must carry, "" for none.
|
||||
hopKey string
|
||||
second func(*http.Request, json.RawMessage) (*http.Response, error)
|
||||
}{
|
||||
{
|
||||
name: "error echoes the first hop's key",
|
||||
location: "https://mcp.aimlapi.com/mcp/",
|
||||
hopKey: second,
|
||||
second: func(request *http.Request, id json.RawMessage) (*http.Response, error) {
|
||||
return fakeHTTPResponse(request, http.StatusOK, rpcReply(id, "error", map[string]any{"code": -32001, "message": "stale " + first})), nil
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "status body after a hop that dropped the key",
|
||||
location: "https://[email protected]/mcp",
|
||||
second: func(request *http.Request, _ json.RawMessage) (*http.Response, error) {
|
||||
return fakeHTTPResponse(request, http.StatusPaymentRequired, `{"error":"invalid credential `+first+`"}`), nil
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "transport error quoting the redirect target",
|
||||
location: "https://mcp.aimlapi.com/mcp?session=" + first,
|
||||
hopKey: second,
|
||||
second: func(*http.Request, json.RawMessage) (*http.Response, error) {
|
||||
return nil, errors.New("connection reset by peer")
|
||||
},
|
||||
},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
rotateAIMLAPIKeys(t, first, second)
|
||||
echo := &aimlapiHopEcho{location: tc.location, second: tc.second}
|
||||
server := aimlapiTestServer("aimlapi", aimlapi.MCPServerURL)
|
||||
client := &networkClient{server: server, client: mcpHTTPClient(server, echo), nextID: 1}
|
||||
err := client.initialize(context.Background())
|
||||
|
||||
echo.mu.Lock()
|
||||
auth := append([]string(nil), echo.auth...)
|
||||
echo.mu.Unlock()
|
||||
if len(auth) != 2 || auth[0] != "Bearer "+first {
|
||||
t.Fatalf("Authorization per hop = %q, want the first key on the first request and a redirect hop", auth)
|
||||
}
|
||||
if tc.hopKey != "" && auth[1] != "Bearer "+tc.hopKey {
|
||||
t.Fatalf("redirect hop Authorization = %q, want the key resolved again for it", auth[1])
|
||||
}
|
||||
// With no key, net/http fills Basic auth from the hop's own userinfo.
|
||||
if tc.hopKey == "" && (strings.Contains(auth[1], first) || strings.Contains(auth[1], second)) {
|
||||
t.Fatalf("redirect hop Authorization = %q, want no automatic key", auth[1])
|
||||
}
|
||||
if err == nil {
|
||||
t.Fatal("initialize succeeded against a failing server")
|
||||
}
|
||||
assertNoPartOf(t, err.Error(), first)
|
||||
assertNoPartOf(t, err.Error(), second)
|
||||
if !strings.Contains(err.Error(), "[REDACTED]") {
|
||||
t.Fatalf("error = %q, want the echoed key replaced", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIMLAPIMCPSSEScrubsTheStreamKeyFromAFailedPost(t *testing.T) {
|
||||
// The server chose an endpoint holding the key its stream request carried.
|
||||
// A POST to it that fails at the transport quotes that URL, and the POST
|
||||
// itself carried a different key, so only the stream's keys can scrub it.
|
||||
const streamKey = "strm0a1b2c3d4e5f60718293a4b5c6d7e8"
|
||||
const postKey = "post9e8d7c6b5a4f30291807f6e5d4c3b2"
|
||||
rotateAIMLAPIKeys(t, streamKey, postKey)
|
||||
echo := newAIMLAPISSEEcho("/messages/"+streamKey, func(string, json.RawMessage) string { return "" })
|
||||
echo.postErr = errors.New("connection reset by peer")
|
||||
client := newEchoSSEClient(echo)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
if err := client.openStream(ctx); err != nil {
|
||||
t.Fatalf("openStream: %v", err)
|
||||
}
|
||||
_, err := client.ListTools(ctx)
|
||||
_ = client.Close()
|
||||
|
||||
echo.mu.Lock()
|
||||
auth := append([]string(nil), echo.auth...)
|
||||
echo.mu.Unlock()
|
||||
if len(auth) != 2 || auth[0] != "Bearer "+streamKey || auth[1] != "Bearer "+postKey {
|
||||
t.Fatalf("Authorization headers = %q, want the stream key and then the post key", auth)
|
||||
}
|
||||
if err == nil || !strings.Contains(err.Error(), "post MCP SSE message") {
|
||||
t.Fatalf("ListTools error = %v, want the failed POST", err)
|
||||
}
|
||||
assertNoPartOf(t, err.Error(), streamKey)
|
||||
assertNoPartOf(t, err.Error(), postKey)
|
||||
if !strings.Contains(err.Error(), "/messages/[REDACTED]") {
|
||||
t.Fatalf("error = %q, want the endpoint quoted with the key replaced", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAIMLAPIMCPSSEScrubsTheKeyFromALaterEndpointError(t *testing.T) {
|
||||
// An endpoint event can arrive after the stream is open. Its parse error
|
||||
// then fails the pending requests directly rather than returning from
|
||||
// openStream, so it has to be scrubbed where it is made.
|
||||
isolateAIMLAPIKeyLookup(t, scrubTestKey)
|
||||
echo := newAIMLAPISSEEcho("/messages", nil)
|
||||
echo.rawEvent = func(string, json.RawMessage) string {
|
||||
return "event: endpoint\ndata: /messages/" + scrubTestKey + "%zz\n\n"
|
||||
}
|
||||
client := newEchoSSEClient(echo)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
if err := client.openStream(ctx); err != nil {
|
||||
t.Fatalf("openStream: %v", err)
|
||||
}
|
||||
_, err := client.ListTools(ctx)
|
||||
_ = client.Close()
|
||||
echo.assertSent(t, "Bearer "+scrubTestKey, 2)
|
||||
if err == nil || !strings.Contains(err.Error(), "endpoint") {
|
||||
t.Fatalf("ListTools error = %v, want the later endpoint's parse failure", err)
|
||||
}
|
||||
assertKeyScrubbed(t, err.Error())
|
||||
}
|
||||
|
||||
func TestAIMLAPIMCPSSEScrubsKeysFromAFailedStreamRedirect(t *testing.T) {
|
||||
// The stream request follows redirects like any other, so a hop that fails
|
||||
// at the transport quotes the target the server chose, which can hold the
|
||||
// key the first hop carried.
|
||||
const first = "strmfirst0123456789abcdef01234567"
|
||||
const second = "strmsecond0123456789abcdef0123456"
|
||||
rotateAIMLAPIKeys(t, first, second)
|
||||
echo := &aimlapiHopEcho{
|
||||
location: "https://mcp.aimlapi.com/mcp?session=" + first,
|
||||
second: func(*http.Request, json.RawMessage) (*http.Response, error) {
|
||||
return nil, errors.New("connection reset by peer")
|
||||
},
|
||||
}
|
||||
server := Server{Name: "aimlapi", Type: ServerTypeSSE, URL: aimlapi.MCPServerURL}
|
||||
client := &remoteSSEClient{server: server, client: mcpHTTPClient(server, echo), nextID: 1, pending: map[string]chan ssePendingResponse{}}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
err := client.openStream(ctx)
|
||||
_ = client.Close()
|
||||
|
||||
echo.mu.Lock()
|
||||
auth := append([]string(nil), echo.auth...)
|
||||
echo.mu.Unlock()
|
||||
if len(auth) != 2 || auth[0] != "Bearer "+first || auth[1] != "Bearer "+second {
|
||||
t.Fatalf("Authorization per hop = %q, want the first key and then the second", auth)
|
||||
}
|
||||
if err == nil || !strings.Contains(err.Error(), "open MCP SSE stream") {
|
||||
t.Fatalf("openStream error = %v, want the failed stream request", err)
|
||||
}
|
||||
assertNoPartOf(t, err.Error(), first)
|
||||
assertNoPartOf(t, err.Error(), second)
|
||||
if !strings.Contains(err.Error(), "session=[REDACTED]") {
|
||||
t.Fatalf("error = %q, want the redirect target quoted with the key replaced", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScrubAIMLAPIKeyErrorKeepsOnlyTheContextSentinel(t *testing.T) {
|
||||
for _, sentinel := range []error{context.Canceled, context.DeadlineExceeded} {
|
||||
original := fmt.Errorf("post MCP SSE message to aimlapi: %w", &url.Error{Op: "Post", URL: "https://mcp.aimlapi.com/messages/" + scrubTestKey, Err: sentinel})
|
||||
scrubbed := scrubAIMLAPIKeyError(original, scrubTestKey)
|
||||
assertKeyScrubbed(t, scrubbed.Error())
|
||||
if !errors.Is(scrubbed, sentinel) {
|
||||
t.Fatalf("scrubbed error %q no longer reads as %v", scrubbed, sentinel)
|
||||
}
|
||||
if unwrapped := errors.Unwrap(scrubbed); unwrapped != sentinel {
|
||||
t.Fatalf("scrubbed error unwraps to %v, want only %v", unwrapped, sentinel)
|
||||
}
|
||||
var urlErr *url.Error
|
||||
if errors.As(scrubbed, &urlErr) {
|
||||
t.Fatalf("scrubbed error still exposes the *url.Error holding the key: %v", urlErr.URL)
|
||||
}
|
||||
}
|
||||
|
||||
if scrubbed := scrubAIMLAPIKeyError(fmt.Errorf("decode: %s", scrubTestKey), scrubTestKey); errors.Unwrap(scrubbed) != nil {
|
||||
t.Fatalf("scrubbed error without a context sentinel unwraps to %v, want nothing", errors.Unwrap(scrubbed))
|
||||
}
|
||||
untouched := fmt.Errorf("call MCP: %w", context.Canceled)
|
||||
if got := scrubAIMLAPIKeyError(untouched, scrubTestKey); got != untouched {
|
||||
t.Fatalf("an error without the key was replaced: %v", got)
|
||||
}
|
||||
}
|
||||
@@ -1,537 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
+160
-3
@@ -54,6 +54,56 @@ type ToolClient interface {
|
||||
Close() error
|
||||
}
|
||||
|
||||
// startupDisclosureError carries a launch disclosure out through a failure.
|
||||
//
|
||||
// A launched process is a fact about the past: once Start has succeeded the
|
||||
// disclosure is true whatever the handshake does next. The client is the only
|
||||
// thing that holds it, and the failure paths close and discard the client, so
|
||||
// without this the fact dies with the connection it was attached to.
|
||||
type startupDisclosureError struct {
|
||||
err error
|
||||
notices []string
|
||||
}
|
||||
|
||||
func (e *startupDisclosureError) Error() string { return e.err.Error() }
|
||||
func (e *startupDisclosureError) Unwrap() error { return e.err }
|
||||
|
||||
// startupNoticesFromError recovers a launch disclosure from a failed connect.
|
||||
func startupNoticesFromError(err error) []string {
|
||||
var disclosure *startupDisclosureError
|
||||
if errors.As(err, &disclosure) {
|
||||
return disclosure.notices
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// startupDisclosing is the optional interface a client implements when its
|
||||
// LAUNCH carried a least-privilege disclosure. A network server launches no
|
||||
// local process, so it does not implement this and reports nothing, which is the
|
||||
// correct answer rather than an empty one.
|
||||
type startupDisclosing interface {
|
||||
StartupNotices() []string
|
||||
}
|
||||
|
||||
// StartupNotices reports the disclosures that applied to this server's launch.
|
||||
//
|
||||
// GATED ON THE SAME DECISION AS EVERY OTHER CARRIER. connectAndList reads this on
|
||||
// the success path and hands the result straight to the disclosure sources, so
|
||||
// for a while a successfully connected wrapped server disclosed on the strength
|
||||
// of cmd.Start returning: the HELPER starting, which is the one thing the report
|
||||
// exists because it does not prove. It happened to be right, since a completed
|
||||
// handshake does imply the child ran, but by coincidence rather than by rule, and
|
||||
// the failure path next to it was already asking the adapter.
|
||||
func (client *Client) StartupNotices() []string {
|
||||
if client == nil || len(client.startupNotices) == 0 {
|
||||
return nil
|
||||
}
|
||||
if client.launched != nil && !client.launched() {
|
||||
return nil
|
||||
}
|
||||
return append([]string(nil), client.startupNotices...)
|
||||
}
|
||||
|
||||
type Client struct {
|
||||
server Server
|
||||
cmd *exec.Cmd
|
||||
@@ -64,6 +114,20 @@ type Client struct {
|
||||
idMu sync.Mutex
|
||||
nextID int
|
||||
cleanup func()
|
||||
// startupNotices are the least-privilege disclosures that applied to THIS
|
||||
// server's launch.
|
||||
//
|
||||
// THE FACT DESCRIBES STARTUP, SO IT CANNOT BE RECOVERED FROM A TOOL RESULT.
|
||||
// A stdio server prepared under the weakened token runs for the whole session,
|
||||
// and connectStdio used to keep only the command and its cleanup, so nothing
|
||||
// downstream could tell the operator that the process serving these tools had
|
||||
// reduced write confinement. Kept typed here and rendered exactly once at
|
||||
// registration rather than pasted onto every later tool result.
|
||||
startupNotices []string
|
||||
// launched is the launch decision this server's disclosures are gated on, nil
|
||||
// for a plan whose started process IS the server and where there is nothing to
|
||||
// decide. See internal/execution/child_launch.go.
|
||||
launched func() bool
|
||||
|
||||
writeMu sync.Mutex
|
||||
writeQueue chan writeOp
|
||||
@@ -169,6 +233,11 @@ func (b *boundedBuffer) String() string {
|
||||
func connectStdio(ctx context.Context, server Server, options ConnectOptions) (*Client, error) {
|
||||
var cmd *exec.Cmd
|
||||
var cleanup func()
|
||||
var plannedEnforcement execution.Enforcement
|
||||
// Retained from the prepared plan rather than dropped: for a wrapped plan the
|
||||
// adapter, not cmd.Start, owns whether the requested server process exists.
|
||||
var launchTracker *execution.ChildLaunchTracker
|
||||
var ownedLaunch bool
|
||||
cleanupTransferred := false
|
||||
defer func() {
|
||||
if cleanup != nil && !cleanupTransferred {
|
||||
@@ -192,7 +261,11 @@ func connectStdio(ctx context.Context, server Server, options ConnectOptions) (*
|
||||
return nil, fmt.Errorf("start MCP server %s: %w", server.Name, err)
|
||||
}
|
||||
cmd = prepared.Command
|
||||
cleanup = prepared.Cleanup
|
||||
plannedEnforcement = prepared.Enforcement
|
||||
ownedLaunch = prepared.ChildLaunchOwnedByAdapter
|
||||
// The tracker's cleanup, not the plan's: it settles the launch decision
|
||||
// before the report file it was read from is deleted.
|
||||
launchTracker, cleanup = execution.NewChildLaunchTracker(prepared)
|
||||
} else {
|
||||
cmd = exec.CommandContext(ctx, server.Command, server.Args...)
|
||||
cmd.Env = mergeProcessEnv(server.Env)
|
||||
@@ -211,6 +284,55 @@ func connectStdio(ctx context.Context, server Server, options ConnectOptions) (*
|
||||
return nil, fmt.Errorf("start MCP server %s: %w", server.Name, err)
|
||||
}
|
||||
|
||||
// PUBLISHED AT START, not at return. Registration abandons a server that
|
||||
// exceeds the connect timeout, and everything below this line (initialize, and
|
||||
// tools/list above it in the caller) can hang past that. Announcing the launch
|
||||
// here is what lets an abandoned attempt still disclose the confinement its
|
||||
// process ran under.
|
||||
//
|
||||
// EXCEPT WHEN START IS NOT THE LAUNCH. For a wrapped plan the process started
|
||||
// above is the sandbox helper, which creates the MCP server only after marker,
|
||||
// ACL, network, SID and token setup, any of which can fail leaving no server at
|
||||
// all. Publishing here would make the durable delivery machinery reliably
|
||||
// announce a confinement nothing ever ran under. Those plans publish from
|
||||
// publishAdapterLaunch below, once the adapter has stated the fact.
|
||||
if !ownedLaunch {
|
||||
publishLaunch(ctx, plannedEnforcement.Notices)
|
||||
}
|
||||
// ONE ANSWER, AND IT IS ONLY CACHED ONCE IT CANNOT CHANGE.
|
||||
//
|
||||
// The adapter's report is a file the plan's cleanup deletes, so a decision made
|
||||
// after cleanup used to read an absent report and answer "no child" about a
|
||||
// server that really did run. Memoizing the first read fixed that ordering and
|
||||
// introduced another: the read can land while the adapter has created the child
|
||||
// and not yet recorded it, and "not yet" got frozen as "never".
|
||||
//
|
||||
// ChildLaunchTracker holds the three states apart. It caches a launch, never an
|
||||
// absence, until the adapter is terminal, and it settles from cleanup with the
|
||||
// report still on disk. See internal/execution/child_launch.go.
|
||||
//
|
||||
// It also gives the success, late and failed paths the same input. The failure
|
||||
// path used to carry client.StartupNotices() unconditionally while the sink was
|
||||
// gated on the adapter, which is two competing definitions of applied
|
||||
// enforcement: an operator was told a server ran without write confinement when
|
||||
// only the sandbox helper ran and the requested server never existed.
|
||||
launchedOnce := func() bool {
|
||||
if !ownedLaunch {
|
||||
return true
|
||||
}
|
||||
return launchTracker.Launched()
|
||||
}
|
||||
// publishAdapterLaunch announces a wrapped plan's launch, but only if the
|
||||
// adapter confirms the requested child was created. Called on both ways this
|
||||
// attempt can end, which is also where an attempt abandoned at the connect
|
||||
// timeout eventually arrives, so a late disclosure is still delivered once.
|
||||
publishAdapterLaunch := func() {
|
||||
if !ownedLaunch || !launchedOnce() {
|
||||
return
|
||||
}
|
||||
publishLaunch(ctx, plannedEnforcement.Notices)
|
||||
}
|
||||
|
||||
client := &Client{
|
||||
server: server,
|
||||
cmd: cmd,
|
||||
@@ -219,16 +341,51 @@ func connectStdio(ctx context.Context, server Server, options ConnectOptions) (*
|
||||
writer: newMessageWriter(stdin),
|
||||
nextID: 1,
|
||||
cleanup: cleanup,
|
||||
// Recorded only now, AFTER Start returned. Everything above returns early,
|
||||
// so a prepare failure or an executable that could not be launched records
|
||||
// nothing: same launch-state rule hooks and plugins use, expressed by where
|
||||
// this assignment sits rather than by another outcome-kind switch.
|
||||
startupNotices: append([]string(nil), plannedEnforcement.Notices...),
|
||||
launched: launchedOnce,
|
||||
}
|
||||
cleanupTransferred = true
|
||||
if err := client.initialize(ctx); err != nil {
|
||||
// THE LAUNCH ALREADY HAPPENED, so the fact has to leave through the error.
|
||||
// Start succeeded above, which means the process ran under the planned token
|
||||
// and may have done filesystem work before the handshake failed. Returning a
|
||||
// bare error discards the client, and with it the only carrier the notices
|
||||
// had, so the operator was told the server was unavailable and not that it
|
||||
// had already run without the write jail.
|
||||
// AFTER Close, which waits out the adapter and then settles the decision
|
||||
// with the report still on disk. Asking first would ask an adapter that may
|
||||
// be between creating the child and recording it, and read "not yet" as
|
||||
// "never": the connect timeout ends the attempt here while the helper is
|
||||
// still working.
|
||||
_ = client.Close()
|
||||
launched := launchedOnce()
|
||||
publishAdapterLaunch()
|
||||
message := strings.TrimSpace(stderr.String())
|
||||
failure := fmt.Errorf("initialize MCP server %s: %w", server.Name, err)
|
||||
if message != "" {
|
||||
return nil, fmt.Errorf("initialize MCP server %s: %w: %s", server.Name, err, message)
|
||||
failure = fmt.Errorf("initialize MCP server %s: %w: %s", server.Name, err, message)
|
||||
}
|
||||
return nil, fmt.Errorf("initialize MCP server %s: %w", server.Name, err)
|
||||
// Same decision as the sink above. For a wrapped plan whose helper started
|
||||
// and then failed before creating the requested server, there is nothing to
|
||||
// disclose: no server ran, confined or otherwise.
|
||||
var carried []string
|
||||
if launched {
|
||||
carried = client.StartupNotices()
|
||||
}
|
||||
return nil, &startupDisclosureError{err: failure, notices: carried}
|
||||
}
|
||||
// THE HANDSHAKE IS THE OBSERVATION. A wrapped plan's adapter speaks no MCP and
|
||||
// creates the child suspended, so a well-formed initialize response can only
|
||||
// have come from the requested server, already running. That is terminal
|
||||
// evidence and does not depend on the report having been read yet, which is
|
||||
// what keeps a long-lived session from having to wait for an adapter that will
|
||||
// not exit until the session ends.
|
||||
launchTracker.Confirm()
|
||||
publishAdapterLaunch()
|
||||
return client, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
// A MODEL-FACING PROTOCOL BOUNDARY IS A PRESENTATION CONSUMER.
|
||||
//
|
||||
// This branch changed the result contract: Result.Output holds the UNDECORATED
|
||||
// base text and ModelOutput is the sole model-facing projection that composes it
|
||||
// with the typed enforcement notices. tools/call serialized Output directly,
|
||||
// which was a complete value before and is not one now, so an affected Windows
|
||||
// command reached an MCP client with its ordinary output and no statement that
|
||||
// its DenyRead token shape left writes unconfined.
|
||||
//
|
||||
// Driven through Serve rather than the accessor, because the question is what
|
||||
// goes on the wire.
|
||||
func TestMCPToolsCallCarriesTheEnforcementNotice(t *testing.T) {
|
||||
const notice = "denyRead is configured, so the write jail is not confining writes"
|
||||
const output = "ran the command"
|
||||
|
||||
for _, testCase := range []struct {
|
||||
name string
|
||||
result tools.Result
|
||||
wantText string
|
||||
wantIsErr bool
|
||||
wantNotice bool
|
||||
}{
|
||||
{
|
||||
name: "successful command with a notice",
|
||||
result: tools.Result{Status: tools.StatusOK, Output: output, EnforcementNotices: []string{notice}},
|
||||
wantIsErr: false,
|
||||
wantNotice: true,
|
||||
},
|
||||
{
|
||||
name: "failed command with a notice",
|
||||
result: tools.Result{Status: tools.StatusError, Output: output, EnforcementNotices: []string{notice}},
|
||||
wantIsErr: true,
|
||||
wantNotice: true,
|
||||
},
|
||||
{
|
||||
name: "ordinary command with no notice",
|
||||
result: tools.Result{Status: tools.StatusOK, Output: output},
|
||||
wantIsErr: false,
|
||||
wantNotice: false,
|
||||
},
|
||||
} {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(serverFakeTool{
|
||||
name: "run_thing",
|
||||
description: "runs a thing",
|
||||
parameters: tools.Schema{Type: "object", AdditionalProperties: false},
|
||||
safety: tools.Safety{SideEffect: tools.SideEffectRead, Permission: tools.PermissionAllow, Reason: "test"},
|
||||
run: func(map[string]any) tools.Result { return testCase.result },
|
||||
})
|
||||
|
||||
var input bytes.Buffer
|
||||
writeServerTestMessage(t, &input, rpcMessage{ID: 1, Method: "initialize"})
|
||||
writeServerTestMessage(t, &input, rpcMessage{Method: "notifications/initialized"})
|
||||
writeServerTestMessage(t, &input, rpcMessage{
|
||||
ID: 2,
|
||||
Method: "tools/call",
|
||||
Params: mustRaw(map[string]any{"name": "run_thing", "arguments": map[string]any{}}),
|
||||
})
|
||||
|
||||
var out bytes.Buffer
|
||||
if err := Serve(context.Background(), &input, &out, registry, ServeOptions{Name: "zero-test", Version: "1.2.3"}); err != nil {
|
||||
t.Fatalf("Serve() error = %v", err)
|
||||
}
|
||||
reader := newMessageReader(&out)
|
||||
readServerTestMessage(t, reader) // initialize
|
||||
var call CallToolResult
|
||||
decodeServerTestResult(t, readServerTestMessage(t, reader), &call)
|
||||
|
||||
if len(call.Content) != 1 || call.Content[0].Type != "text" {
|
||||
t.Fatalf("content shape changed: %#v", call.Content)
|
||||
}
|
||||
text := call.Content[0].Text
|
||||
if call.IsError != testCase.wantIsErr {
|
||||
t.Errorf("IsError = %v, want %v", call.IsError, testCase.wantIsErr)
|
||||
}
|
||||
if count := strings.Count(text, output); count != 1 {
|
||||
t.Errorf("the command's own output appears %d times, want exactly 1: %q", count, text)
|
||||
}
|
||||
gotNotice := strings.Count(text, notice)
|
||||
if testCase.wantNotice && gotNotice != 1 {
|
||||
t.Errorf("the disclosure appears %d times, want exactly 1: %q", gotNotice, text)
|
||||
}
|
||||
if !testCase.wantNotice {
|
||||
if gotNotice != 0 {
|
||||
t.Errorf("a disclosure appeared for a command that had none: %q", text)
|
||||
}
|
||||
if text != output {
|
||||
t.Errorf("ordinary output was altered: got %q, want %q", text, output)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// launchSink carries the fact that a server's PROCESS STARTED out of the connect
|
||||
// attempt, without waiting for that attempt to finish.
|
||||
//
|
||||
// The startup notices used to leave connectStdio only on the returned client, or
|
||||
// on the returned error. Both require the attempt to return. Registration
|
||||
// abandons a server that exceeds the connect timeout, so a server that started
|
||||
// under the reduced write confinement and then hung in initialize or tools/list
|
||||
// was recorded as skipped with nothing said about the confinement it ran under.
|
||||
// The reaper that later collects the abandoned attempt runs after the serial
|
||||
// commit phase has finished, so it cannot contribute without breaking the
|
||||
// deterministic ordering that phase exists to provide.
|
||||
//
|
||||
// Publishing at Start splits the two facts apart, which is the point: launch and
|
||||
// connection usability have different lifetimes. A sink that was never published
|
||||
// to means Start never happened, so prepare, pipe, and Start failures stay silent
|
||||
// exactly as before.
|
||||
//
|
||||
// IT IS ALSO AN EVENT, NOT ONLY A VALUE. A retained sink that nobody re-reads is
|
||||
// still a lost disclosure: both production reporters sample once, immediately
|
||||
// after registration returns, and a Start that completes after that sample had
|
||||
// no way to reach them. onPublish lets a reporter subscribe; if the launch has
|
||||
// already happened by the time it subscribes, it is told at once, so the fact
|
||||
// reaches exactly one presentation regardless of which side won the race.
|
||||
type launchSink struct {
|
||||
mu sync.Mutex
|
||||
launched bool
|
||||
notices []string
|
||||
onPublish func(notices []string)
|
||||
delivered bool
|
||||
}
|
||||
|
||||
type launchSinkKey struct{}
|
||||
|
||||
// withLaunchSink attaches a sink to the context handed to the client factory.
|
||||
// Carried on the context rather than added to the factory signature so an
|
||||
// injected or third-party factory that knows nothing about it still works, and
|
||||
// simply discloses nothing.
|
||||
func withLaunchSink(ctx context.Context, sink *launchSink) context.Context {
|
||||
return context.WithValue(ctx, launchSinkKey{}, sink)
|
||||
}
|
||||
|
||||
// publishLaunch records that the process for this connect attempt has started,
|
||||
// along with the enforcement notices that applied to it. Safe on a context with
|
||||
// no sink, which is every caller outside registration.
|
||||
func publishLaunch(ctx context.Context, notices []string) {
|
||||
sink, _ := ctx.Value(launchSinkKey{}).(*launchSink)
|
||||
if sink == nil {
|
||||
return
|
||||
}
|
||||
sink.mu.Lock()
|
||||
sink.launched = true
|
||||
sink.notices = append([]string(nil), notices...)
|
||||
deliver := sink.pendingDeliveryLocked()
|
||||
sink.mu.Unlock()
|
||||
if deliver != nil {
|
||||
deliver()
|
||||
}
|
||||
}
|
||||
|
||||
// PublishLaunchForTest is publishLaunch for a test in another package that
|
||||
// injects a client factory and needs to mark its fake process as started. It
|
||||
// is the same function with the same context lookup, so a test exercises the
|
||||
// real sink rather than a stand-in, and it is inert on any context that did
|
||||
// not come through registration.
|
||||
func PublishLaunchForTest(ctx context.Context, notices []string) {
|
||||
publishLaunch(ctx, notices)
|
||||
}
|
||||
|
||||
// observe reports whether Start was reached and what applied to it. Read from
|
||||
// the registration goroutine while the connect goroutine may still be running,
|
||||
// hence the lock.
|
||||
func (sink *launchSink) observe() (bool, []string) {
|
||||
if sink == nil {
|
||||
return false, nil
|
||||
}
|
||||
sink.mu.Lock()
|
||||
defer sink.mu.Unlock()
|
||||
return sink.launched, append([]string(nil), sink.notices...)
|
||||
}
|
||||
|
||||
// subscribe registers the one presentation this launch should reach. If the
|
||||
// launch already happened, fn runs before subscribe returns; otherwise it runs
|
||||
// from publishLaunch. Either way it runs at most once, and a second subscriber
|
||||
// replaces nothing: the first delivery is the only delivery.
|
||||
func (sink *launchSink) subscribe(fn func(notices []string)) {
|
||||
if sink == nil || fn == nil {
|
||||
return
|
||||
}
|
||||
sink.mu.Lock()
|
||||
if sink.onPublish == nil {
|
||||
sink.onPublish = fn
|
||||
}
|
||||
deliver := sink.pendingDeliveryLocked()
|
||||
sink.mu.Unlock()
|
||||
if deliver != nil {
|
||||
deliver()
|
||||
}
|
||||
}
|
||||
|
||||
// pendingDeliveryLocked returns the delivery to perform, or nil, and marks it
|
||||
// done. Called with mu held; the returned closure must be invoked with mu
|
||||
// released, since a subscriber may itself take other locks.
|
||||
func (sink *launchSink) pendingDeliveryLocked() func() {
|
||||
if !sink.launched || sink.onPublish == nil || sink.delivered {
|
||||
return nil
|
||||
}
|
||||
sink.delivered = true
|
||||
fn := sink.onPublish
|
||||
notices := append([]string(nil), sink.notices...)
|
||||
return func() { fn(notices) }
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Twigpine/zero/internal/config"
|
||||
"github.com/Twigpine/zero/internal/tools"
|
||||
)
|
||||
|
||||
const launchNotice = "denyRead is configured, so the Windows sandbox uses the token shape without WRITE_RESTRICTED (#869)"
|
||||
|
||||
func registerWithFactory(t *testing.T, factory func(context.Context, Server) (ToolClient, error)) *Runtime {
|
||||
t.Helper()
|
||||
runtime, err := RegisterTools(context.Background(), tools.NewRegistry(), config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"slow": {Type: "stdio", Command: "slow-mcp"},
|
||||
}}, RegisterOptions{
|
||||
ConnectTimeout: 50 * time.Millisecond,
|
||||
ClientFactory: factory,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("RegisterTools error: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = runtime.Close() })
|
||||
return runtime
|
||||
}
|
||||
|
||||
// A SERVER THAT STARTED AND THEN HUNG STILL RAN UNDER THE REDUCED TOKEN.
|
||||
//
|
||||
// Registration abandons a server that exceeds the connect timeout and records it
|
||||
// as skipped. The startup notices used to leave connectStdio only on the returned
|
||||
// client or the returned error, and the abandoned attempt returns neither before
|
||||
// the serial commit phase is over, so the process ran with reduced write
|
||||
// confinement and startup said only that the server was skipped.
|
||||
func TestTimeoutAfterLaunchKeepsTheStartupDisclosure(t *testing.T) {
|
||||
runtime := registerWithFactory(t, func(ctx context.Context, server Server) (ToolClient, error) {
|
||||
// The process started under the reduced token, then initialize hangs.
|
||||
publishLaunch(ctx, []string{launchNotice})
|
||||
<-ctx.Done()
|
||||
return nil, ctx.Err()
|
||||
})
|
||||
|
||||
disclosures := runtime.StartupDisclosures()
|
||||
if len(disclosures) == 0 {
|
||||
t.Fatal("a server that started and then timed out disclosed nothing, so it ran under reduced write confinement unannounced")
|
||||
}
|
||||
if disclosures[0].Name != "slow" || len(disclosures[0].Notices) != 1 || disclosures[0].Notices[0] != launchNotice {
|
||||
t.Errorf("StartupDisclosures() = %#v, want one entry for slow carrying the launch notice", disclosures)
|
||||
}
|
||||
if skipped := runtime.Skipped(); len(skipped) != 1 || skipped[0].Name != "slow" {
|
||||
t.Errorf("the server should still be recorded as skipped: %#v", skipped)
|
||||
}
|
||||
}
|
||||
|
||||
// AND A TIMEOUT BEFORE LAUNCH STAYS SILENT.
|
||||
//
|
||||
// Without this, retaining the disclosure on timeout could be satisfied by
|
||||
// disclosing on every timeout, which would claim a token trade for a process
|
||||
// that was never created.
|
||||
func TestTimeoutBeforeLaunchDisclosesNothing(t *testing.T) {
|
||||
runtime := registerWithFactory(t, func(ctx context.Context, server Server) (ToolClient, error) {
|
||||
// Never reached Start: no publish.
|
||||
<-ctx.Done()
|
||||
return nil, ctx.Err()
|
||||
})
|
||||
|
||||
if disclosures := runtime.StartupDisclosures(); len(disclosures) != 0 {
|
||||
t.Errorf("a server that never started claimed a token trade: %#v", disclosures)
|
||||
}
|
||||
if skipped := runtime.Skipped(); len(skipped) != 1 {
|
||||
t.Errorf("the server should still be recorded as skipped: %#v", skipped)
|
||||
}
|
||||
}
|
||||
|
||||
// THE PUBLISH HAS TO SIT AFTER Start, AND ONLY A REAL LAUNCH PROVES IT.
|
||||
//
|
||||
// The two tests above inject a factory, so they exercise the registry's handling
|
||||
// of the sink and say nothing about where connectStdio publishes to it. Moving
|
||||
// the call one line up, above cmd.Start, leaves both of them green while every
|
||||
// failed launch starts claiming the token trade. This one drives the real
|
||||
// connectStdio with a command that cannot start.
|
||||
func TestAFailedStartPublishesNoLaunch(t *testing.T) {
|
||||
sink := &launchSink{}
|
||||
ctx := withLaunchSink(context.Background(), sink)
|
||||
|
||||
client, err := connectStdio(ctx, Server{
|
||||
Name: "missing",
|
||||
Type: "stdio",
|
||||
Command: "zero-nonexistent-mcp-binary-for-test",
|
||||
}, ConnectOptions{})
|
||||
if err == nil {
|
||||
if client != nil {
|
||||
_ = client.Close()
|
||||
}
|
||||
t.Fatal("expected a nonexistent executable to fail to start")
|
||||
}
|
||||
|
||||
if launched, notices := sink.observe(); launched {
|
||||
t.Errorf("a server whose process never started was published as launched (notices %#v)", notices)
|
||||
}
|
||||
}
|
||||
|
||||
// A START THAT COMPLETES JUST AFTER THE TIMEOUT MUST STILL BE DISCLOSED.
|
||||
//
|
||||
// connectStdio publishes only once cmd.Start has returned, and the timeout
|
||||
// branch used to sample the sink the instant it fired. Those interleave: the
|
||||
// sample reads empty, the result commits with no notice, and the reaper closes
|
||||
// the late client without being able to amend a commit that already happened.
|
||||
//
|
||||
// The real window is microseconds wide, so this drives the CONTRACT instead:
|
||||
// the attempt starts after the registration timeout has elapsed but inside the
|
||||
// settle grace, which is the case the synchronization exists to catch.
|
||||
func TestStartJustAfterTheTimeoutIsStillDisclosed(t *testing.T) {
|
||||
runtime := registerWithFactory(t, func(ctx context.Context, server Server) (ToolClient, error) {
|
||||
// 50ms registration timeout has fired; this lands inside launchSettleGrace.
|
||||
time.Sleep(120 * time.Millisecond)
|
||||
publishLaunch(ctx, []string{launchNotice})
|
||||
return nil, errors.New("initialize failed after start")
|
||||
})
|
||||
|
||||
disclosures := runtime.StartupDisclosures()
|
||||
if len(disclosures) != 1 || len(disclosures[0].Notices) != 1 || disclosures[0].Notices[0] != launchNotice {
|
||||
t.Fatalf("a process that started just after the timeout was not disclosed: %#v", disclosures)
|
||||
}
|
||||
if skipped := runtime.Skipped(); len(skipped) != 1 {
|
||||
t.Errorf("the server should still be recorded as skipped: %#v", skipped)
|
||||
}
|
||||
}
|
||||
|
||||
// And an attempt that never starts is not held for the grace, nor disclosed.
|
||||
func TestTimeoutBeforeStartIsNotDelayedOrDisclosed(t *testing.T) {
|
||||
start := time.Now()
|
||||
runtime := registerWithFactory(t, func(ctx context.Context, server Server) (ToolClient, error) {
|
||||
<-ctx.Done() // cancel arrives with the timeout; returns immediately
|
||||
return nil, ctx.Err()
|
||||
})
|
||||
elapsed := time.Since(start)
|
||||
|
||||
if disclosures := runtime.StartupDisclosures(); len(disclosures) != 0 {
|
||||
t.Errorf("a server that never started claimed a token trade: %#v", disclosures)
|
||||
}
|
||||
// The grace is 250ms; an attempt that returns on cancel must not pay it.
|
||||
if elapsed > 200*time.Millisecond {
|
||||
t.Errorf("registration waited %v for an attempt that never started", elapsed)
|
||||
}
|
||||
}
|
||||
|
||||
// AND A START THAT COMPLETES AFTER THE SETTLE GRACE MUST STILL BE DISCLOSED.
|
||||
//
|
||||
// The grace is only another timeout. Once it expires, registration reaps the
|
||||
// attempt in the background and returns; the process can still be inside
|
||||
// cmd.Start at that moment and start successfully afterwards. If Runtime were a
|
||||
// snapshot taken at commit time, that launch would have no owner: the reaper can
|
||||
// close the late client but cannot amend a value already returned, so the server
|
||||
// would have run under the reduced write confinement with startup reporting it
|
||||
// only as skipped.
|
||||
//
|
||||
// A larger grace changes the probability, not the contract, which is why this
|
||||
// releases the launch strictly AFTER the bound rather than inside it. The sink
|
||||
// outlives registration and StartupDisclosures reads through it.
|
||||
func TestStartAfterTheSettleGraceIsStillDisclosed(t *testing.T) {
|
||||
released := make(chan struct{})
|
||||
runtime := registerWithFactory(t, func(ctx context.Context, server Server) (ToolClient, error) {
|
||||
// Held past the 50ms registration timeout AND past launchSettleGrace, so
|
||||
// registration has already reaped this attempt and returned.
|
||||
<-released
|
||||
publishLaunch(ctx, []string{launchNotice})
|
||||
return nil, errors.New("initialize failed long after start")
|
||||
})
|
||||
|
||||
// Registration is done and the disclosure legitimately is not known yet.
|
||||
if disclosures := runtime.StartupDisclosures(); len(disclosures) != 0 {
|
||||
t.Fatalf("nothing had started yet, so nothing should be disclosed: %#v", disclosures)
|
||||
}
|
||||
|
||||
close(released)
|
||||
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
var disclosures []StartupDisclosure
|
||||
for time.Now().Before(deadline) {
|
||||
if disclosures = runtime.StartupDisclosures(); len(disclosures) > 0 {
|
||||
break
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
|
||||
if len(disclosures) != 1 || len(disclosures[0].Notices) != 1 || disclosures[0].Notices[0] != launchNotice {
|
||||
t.Fatalf("a process that started after the settle grace was never disclosed: %#v", disclosures)
|
||||
}
|
||||
// Reading again must not duplicate it.
|
||||
if again := runtime.StartupDisclosures(); len(again) != 1 {
|
||||
t.Errorf("a second read changed the disclosures: %#v", again)
|
||||
}
|
||||
if skipped := runtime.Skipped(); len(skipped) != 1 {
|
||||
t.Errorf("the server should still be recorded as skipped: %#v", skipped)
|
||||
}
|
||||
}
|
||||
@@ -35,11 +35,6 @@ type remoteSSEClient struct {
|
||||
pending map[string]chan ssePendingResponse
|
||||
streamErr error
|
||||
closed bool
|
||||
// streamKeys are the injected AI/ML API keys the stream request carried,
|
||||
// redirect hops included. openStream sets them once, before the reader
|
||||
// starts; the endpoint and the messages on the stream come from a server
|
||||
// that was sent them.
|
||||
streamKeys []string
|
||||
}
|
||||
|
||||
type ssePendingResponse struct {
|
||||
@@ -264,13 +259,12 @@ func (client *remoteSSEClient) request(ctx context.Context, method string, param
|
||||
client.pending[key] = responseChannel
|
||||
client.mu.Unlock()
|
||||
|
||||
sentKeys, err := client.post(ctx, rpcMessage{
|
||||
if err := client.post(ctx, rpcMessage{
|
||||
JSONRPC: "2.0",
|
||||
ID: id,
|
||||
Method: method,
|
||||
Params: rawParams,
|
||||
})
|
||||
if err != nil {
|
||||
}); err != nil {
|
||||
client.removePending(key)
|
||||
return err
|
||||
}
|
||||
@@ -280,15 +274,14 @@ func (client *remoteSSEClient) request(ctx context.Context, method string, param
|
||||
if response.err != nil {
|
||||
return response.err
|
||||
}
|
||||
message := scrubAIMLAPIKeys(response.message, append(append([]string(nil), client.streamKeys...), sentKeys...)...)
|
||||
if !rpcIDMatches(message.ID, id) {
|
||||
if !rpcIDMatches(response.message.ID, id) {
|
||||
return fmt.Errorf("MCP %s response id mismatch for server %s", method, client.server.Name)
|
||||
}
|
||||
if message.Error != nil {
|
||||
return fmt.Errorf("MCP %s failed: %s", method, message.Error.Message)
|
||||
if response.message.Error != nil {
|
||||
return fmt.Errorf("MCP %s failed: %s", method, response.message.Error.Message)
|
||||
}
|
||||
if target != nil && len(message.Result) > 0 {
|
||||
if err := json.Unmarshal(message.Result, target); err != nil {
|
||||
if target != nil && len(response.message.Result) > 0 {
|
||||
if err := json.Unmarshal(response.message.Result, target); err != nil {
|
||||
return fmt.Errorf("decode MCP %s result: %w", method, err)
|
||||
}
|
||||
}
|
||||
@@ -327,12 +320,11 @@ func (client *remoteSSEClient) notify(ctx context.Context, method string, params
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = client.post(ctx, rpcMessage{
|
||||
return client.post(ctx, rpcMessage{
|
||||
JSONRPC: "2.0",
|
||||
Method: method,
|
||||
Params: rawParams,
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (client *networkClient) post(ctx context.Context, message rpcMessage, expectResponse bool) (result rpcMessage, err error) {
|
||||
@@ -340,11 +332,6 @@ func (client *networkClient) post(ctx context.Context, message rpcMessage, expec
|
||||
if err != nil {
|
||||
return rpcMessage{}, err
|
||||
}
|
||||
ctx, sent := withAIMLAPISentKeys(ctx)
|
||||
// From here an error can quote the server or a URL it chose, a redirect
|
||||
// target among them, so every key the call sent, on any hop, is scrubbed
|
||||
// from it. This runs last, after the body is closed.
|
||||
defer func() { err = scrubAIMLAPIKeyError(err, sent.list()...) }()
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodPost, client.server.URL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return rpcMessage{}, fmt.Errorf("create MCP %s request for %s: %w", client.server.Type, client.server.Name, err)
|
||||
@@ -370,18 +357,12 @@ func (client *networkClient) post(ctx context.Context, message rpcMessage, expec
|
||||
if response.StatusCode == http.StatusNoContent {
|
||||
return rpcMessage{}, fmt.Errorf("MCP %s server %s returned no response body", client.server.Type, client.server.Name)
|
||||
}
|
||||
decoded, err := client.decodeResponse(response)
|
||||
if err != nil {
|
||||
return rpcMessage{}, err
|
||||
}
|
||||
return scrubAIMLAPIKeys(decoded, sent.list()...), nil
|
||||
return client.decodeResponse(response)
|
||||
}
|
||||
|
||||
func (client *remoteSSEClient) openStream(ctx context.Context) (err error) {
|
||||
func (client *remoteSSEClient) openStream(ctx context.Context) error {
|
||||
streamCtx, cancel := context.WithCancel(ctx)
|
||||
client.streamCancel = cancel
|
||||
streamCtx, sent := withAIMLAPISentKeys(streamCtx)
|
||||
defer func() { err = scrubAIMLAPIKeyError(err, sent.list()...) }()
|
||||
|
||||
request, err := http.NewRequestWithContext(streamCtx, http.MethodGet, client.server.URL, nil)
|
||||
if err != nil {
|
||||
@@ -399,7 +380,6 @@ func (client *remoteSSEClient) openStream(ctx context.Context) (err error) {
|
||||
return err
|
||||
}
|
||||
client.streamBody = response.Body
|
||||
client.streamKeys = sent.list()
|
||||
|
||||
ready := make(chan error, 1)
|
||||
go client.readStream(response.Body, ready)
|
||||
@@ -411,41 +391,31 @@ func (client *remoteSSEClient) openStream(ctx context.Context) (err error) {
|
||||
}
|
||||
}
|
||||
|
||||
// post sends message to the endpoint and returns the injected AI/ML API keys the
|
||||
// request carried, redirect hops included, since the answer arrives on the
|
||||
// stream from a server that was sent them. The endpoint is the server's own
|
||||
// choice and can hold one of the stream's keys, and a transport error quotes it,
|
||||
// so the stream's keys are scrubbed from any error as well as the POST's.
|
||||
func (client *remoteSSEClient) post(ctx context.Context, message rpcMessage) (sentKeys []string, err error) {
|
||||
ctx, sent := withAIMLAPISentKeys(ctx)
|
||||
defer func() {
|
||||
sentKeys = sent.list()
|
||||
err = scrubAIMLAPIKeyError(err, append(append([]string(nil), client.streamKeys...), sentKeys...)...)
|
||||
}()
|
||||
func (client *remoteSSEClient) post(ctx context.Context, message rpcMessage) (err error) {
|
||||
endpointURL, err := client.currentEndpoint()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return err
|
||||
}
|
||||
body, err := json.Marshal(message)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return err
|
||||
}
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create MCP SSE post for %s: %w", client.server.Name, err)
|
||||
return fmt.Errorf("create MCP SSE post for %s: %w", client.server.Name, err)
|
||||
}
|
||||
client.applyPostHeaders(request)
|
||||
|
||||
response, err := client.client.Do(request)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("post MCP SSE message to %s: %w", client.server.Name, err)
|
||||
return fmt.Errorf("post MCP SSE message to %s: %w", client.server.Name, err)
|
||||
}
|
||||
defer closeResponseBody(&err, client.server, response.Body)
|
||||
if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
|
||||
return nil, httpStatusError(client.server, response)
|
||||
return httpStatusError(client.server, response)
|
||||
}
|
||||
_, _ = io.Copy(io.Discard, response.Body)
|
||||
return nil, nil
|
||||
return nil
|
||||
}
|
||||
|
||||
func (client *networkClient) applyHeaders(request *http.Request) {
|
||||
@@ -457,7 +427,6 @@ 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) {
|
||||
@@ -465,7 +434,6 @@ 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) {
|
||||
@@ -474,7 +442,6 @@ 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) {
|
||||
@@ -556,8 +523,7 @@ func (client *remoteSSEClient) readStream(reader io.Reader, ready chan<- error)
|
||||
func (client *remoteSSEClient) setEndpoint(value string) error {
|
||||
endpointURL, err := resolveSSEEndpointURL(client.server.URL, strings.TrimSpace(value))
|
||||
if err != nil {
|
||||
// A URL parse error quotes the endpoint, which is the server's text.
|
||||
return scrubAIMLAPIKeyError(err, client.streamKeys...)
|
||||
return err
|
||||
}
|
||||
client.mu.Lock()
|
||||
client.endpointURL = endpointURL
|
||||
@@ -641,18 +607,12 @@ 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%s", server.Type, server.Name, response.StatusCode, hint)
|
||||
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%s", server.Type, server.Name, response.StatusCode, detail, hint)
|
||||
return fmt.Errorf("MCP %s server %s returned HTTP %d: %s", server.Type, server.Name, response.StatusCode, detail)
|
||||
}
|
||||
|
||||
func decodeSSERPCMessage(reader io.Reader) (rpcMessage, error) {
|
||||
|
||||
Loaded 100 of 242 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user