mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-08 17:56:14 -10:00
Compare commits
23
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
adb3350ea5 | ||
|
|
865bd5c22b | ||
|
|
2749447cdc | ||
|
|
ad34dc8d81 | ||
|
|
6edf9a8b78 | ||
|
|
90dcfd127e | ||
|
|
01c3498205 | ||
|
|
1ec7219aee | ||
|
|
1a81292d3f | ||
|
|
a78c6c357e | ||
|
|
8dd33ec39e | ||
|
|
d7ac85cbf9 | ||
|
|
d66ad71532 | ||
|
|
d065467cd3 | ||
|
|
0eab63c908 | ||
|
|
64a783b08c | ||
|
|
9c292092a2 | ||
|
|
060d38ce71 | ||
|
|
04fd3c01a7 | ||
|
|
dc15e82298 | ||
|
|
2d2450e9a7 | ||
|
|
cabfeefc4e | ||
|
|
0c1de25a04 |
No files matched your search
@@ -1 +1 @@
|
||||
{".":"0.7.0"}
|
||||
{".":"0.8.0"}
|
||||
@@ -91,6 +91,14 @@ These classes drive multi-round reviews. Fix them before requesting review:
|
||||
final-component-only no-follow is insufficient. On multi-step setup, roll back
|
||||
only what this run created; never destroy pre-existing resources you did
|
||||
not create; never report success when cleanup or unlock failed.
|
||||
- **Normalize before matching:** Any transform that REMOVES bytes without
|
||||
leaving a gap is also a reassembler, so it has to run before whatever matches
|
||||
on the result. Redacting by shape and then stripping control bytes lets a
|
||||
credential split by a NUL, an ESC or a C1 byte pass the patterns as two
|
||||
fragments and be rejoined on the way out; comparing a path against a form
|
||||
produced by the same resolver the kernel just used makes a redirect agree with
|
||||
itself. Strip, decode, or canonicalize first, then match. The unsplit value
|
||||
passing is not evidence, so the test needs a split case.
|
||||
- **Atomic shared state:** Write a complete temporary file, then atomically
|
||||
replace the destination so concurrent readers never see a partial write.
|
||||
Exclusive create or a write lock alone is not enough for readers. Serialize
|
||||
@@ -103,3 +111,42 @@ These classes drive multi-round reviews. Fix them before requesting review:
|
||||
shipped. Wire advertised entry points or shrink the claim. Do not bundle
|
||||
unrelated fixes. Do not widen security allowlists by name alone; check
|
||||
classification and side effects too.
|
||||
- **Host tool version floors:** Zero shells out to `git` and other host tools,
|
||||
and the project does not declare a minimum version for any of them. A flag,
|
||||
subcommand, or config key that is not long-established therefore raises that
|
||||
floor silently, and the failure on an older host is an unhelpful usage error
|
||||
rather than a clear message. Before adding one, check when it landed and say
|
||||
so where a caller would look. Record the floor for the exact subcommand you
|
||||
call, not for the flag in general: `--end-of-options` reached Git's shared
|
||||
parser in 2.24, but `git rev-parse` only accepted it from 2.30, and
|
||||
`git checkout` and `git reset` later still. Likewise `--show-current` is
|
||||
`git branch` 2.22, and `branch.autoSetupMerge=inherit` is 2.35. Then either
|
||||
keep a fallback for older versions or gate the call. A test that depends on a
|
||||
newer feature must skip below its floor instead of failing, and say which
|
||||
version it needs.
|
||||
- **Prove the test fails without the change:** A regression test earns its
|
||||
place only if it fails on the unfixed code, for the reason it claims. Run it
|
||||
both ways and quote the failure in the PR. The recurring miss is not a
|
||||
missing test but a test that never reaches the behavior it names, because an
|
||||
earlier guard rejects the input first, so it passes with and without the fix.
|
||||
When a test targets a specific layer, call that layer directly rather than
|
||||
relying on a public entry point that may refuse earlier.
|
||||
- **Hermetic tests:** Tests must not read or write the developer's real config,
|
||||
cache, or state directories. Redirect every root the code under test
|
||||
resolves, not just config. `os.UserConfigDir` and `os.UserCacheDir` read
|
||||
`XDG_CONFIG_HOME` and `XDG_CACHE_HOME` only on Linux and the BSDs. On macOS
|
||||
they derive from `HOME`, and on Windows from `%AppData%` and `%LocalAppData%`.
|
||||
Setting the XDG variables alone still leaves a macOS run writing to the real
|
||||
home, so set `HOME` as well, or take the directory as a parameter. Apply it to
|
||||
every test that reaches the same storage, not only the one where it was first
|
||||
needed; a test that passes on a clean machine and pollutes a real one is a
|
||||
blocker.
|
||||
- **Error codes vary by platform and by flag combination:** A check written
|
||||
against the code one target returns can be silently inert on another while
|
||||
still looking correct. `openat(O_DIRECTORY|O_NOFOLLOW)` on a symlink may
|
||||
report `ENOTDIR` or `ELOOP` depending on the platform, because the kernel can
|
||||
reject the directory mismatch before it ever resolves the link; accept both
|
||||
rather than assuming the one POSIX names. On Windows, `\\?\` and `\\.\`
|
||||
prefixes are not UNC paths, and an `NtCreateFile` information class must match
|
||||
the struct actually passed. Confirm the code each supported target returns
|
||||
instead of reasoning from the man page.
|
||||
@@ -6,6 +6,27 @@ 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.8.0](https://github.com/Gitlawb/zero/compare/v0.7.0...v0.8.0) (2026-08-21)
|
||||
|
||||
|
||||
### Features
|
||||
|
||||
* **security:** path-jail primitive, credential-store locking, worktree git hardening ([#891](https://github.com/Gitlawb/zero/issues/891)) ([060d38c](https://github.com/Gitlawb/zero/commit/060d38ce716feb7c5b6657b082cd35cb8dda3d75))
|
||||
* **update:** recognise Homebrew installs and add mise, go install docs ([#910](https://github.com/Gitlawb/zero/issues/910)) ([01c3498](https://github.com/Gitlawb/zero/commit/01c3498205e69d26e5f929225f9455b276444f39))
|
||||
* use Exa as default search provider ([#926](https://github.com/Gitlawb/zero/issues/926)) ([1ec7219](https://github.com/Gitlawb/zero/commit/1ec7219aee7b2840481ab0905f131644e27d2a4e))
|
||||
* **zerogit:** auto-create a conventional branch before push/pr on default branch ([#671](https://github.com/Gitlawb/zero/issues/671)) ([d065467](https://github.com/Gitlawb/zero/commit/d065467cd3ec3ed91c5dc60554c36fc4b4d6de4c))
|
||||
|
||||
|
||||
### Bug Fixes
|
||||
|
||||
* **acp:** three defects a desktop ACP client hits ([#915](https://github.com/Gitlawb/zero/issues/915)) ([a78c6c3](https://github.com/Gitlawb/zero/commit/a78c6c357e1d21146d7cd2e8030a754e06f7f1f5))
|
||||
* bump Go to 1.26.6 for stdlib vulnerability fixes ([#903](https://github.com/Gitlawb/zero/issues/903)) ([dc15e82](https://github.com/Gitlawb/zero/commit/dc15e822987d8d7de467c62c375518f435b74d8c))
|
||||
* **daemon:** clean up child after startup timeout ([#774](https://github.com/Gitlawb/zero/issues/774)) ([9c29209](https://github.com/Gitlawb/zero/commit/9c292092a26c9092fd873d072b0af1b0d9d39ef0))
|
||||
* **dictation:** redact API keys from streaming transcriber errors ([#852](https://github.com/Gitlawb/zero/issues/852)) ([64a783b](https://github.com/Gitlawb/zero/commit/64a783b08caa2c7b21f188033b9ed84a14e00e84))
|
||||
* **modelregistry:** expose reasoning effort for DeepSeek V4 models ([#931](https://github.com/Gitlawb/zero/issues/931)) ([90dcfd1](https://github.com/Gitlawb/zero/commit/90dcfd127e8a6d9902ec9f4e72f6d03fef1a0fc6))
|
||||
* **providers:** discover ChatGPT capabilities ([#890](https://github.com/Gitlawb/zero/issues/890)) ([2d2450e](https://github.com/Gitlawb/zero/commit/2d2450e9a744349f0d01b1d4e9ba29c24ba5650d))
|
||||
* **sandbox:** normalize launcher names before the command-prefix denylist ([#934](https://github.com/Gitlawb/zero/issues/934)) ([6edf9a8](https://github.com/Gitlawb/zero/commit/6edf9a8b78dc030dc44598919db1c7fa9d4f809a))
|
||||
|
||||
## [0.7.0](https://github.com/Gitlawb/zero/compare/v0.6.0...v0.7.0) (2026-08-10)
|
||||
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
|
||||
<p align="center">
|
||||
<a href="LICENSE"><img alt="license" src="https://img.shields.io/badge/license-MIT-blue"></a>
|
||||
<img alt="Go 1.26.5+" src="https://img.shields.io/badge/Go-1.26.5+-00ADD8?logo=go&logoColor=white">
|
||||
<img alt="Go 1.26.6+" src="https://img.shields.io/badge/Go-1.26.6+-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>
|
||||
<br>
|
||||
@@ -76,7 +76,7 @@ irm https://raw.githubusercontent.com/Gitlawb/zero/main/scripts/install.ps1 | ie
|
||||
|
||||
### From source
|
||||
|
||||
Source builds require Go 1.26.5+.
|
||||
Source builds require Go 1.26.6+.
|
||||
|
||||
```bash
|
||||
git clone https://github.com/Gitlawb/zero.git
|
||||
|
||||
+2
-2
@@ -6,7 +6,7 @@
|
||||
|
||||
<p align="center">
|
||||
<a href="LICENSE"><img alt="license" src="https://img.shields.io/badge/license-MIT-blue"></a>
|
||||
<img alt="Go 1.26.5+" src="https://img.shields.io/badge/Go-1.26.5+-00ADD8?logo=go&logoColor=white">
|
||||
<img alt="Go 1.26.6+" src="https://img.shields.io/badge/Go-1.26.6+-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>
|
||||
<br>
|
||||
@@ -57,7 +57,7 @@ irm https://raw.githubusercontent.com/Gitlawb/zero/main/scripts/install.ps1 | ie
|
||||
|
||||
### 从源码构建
|
||||
|
||||
源码构建需要 Go 1.26.5+。
|
||||
源码构建需要 Go 1.26.6+。
|
||||
|
||||
```bash
|
||||
git clone https://github.com/Gitlawb/zero.git
|
||||
|
||||
+1
-1
@@ -33,4 +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 the documented Firecrawl endpoint — see the README).
|
||||
keyless web search, to Exa's hosted MCP endpoint at `https://mcp.exa.ai/mcp`).
|
||||
+83
-8
@@ -1,14 +1,30 @@
|
||||
# Installing Zero
|
||||
|
||||
Zero is distributed as:
|
||||
Pick whichever fits how you already install things. None of these is the
|
||||
blessed one.
|
||||
|
||||
- an npm package, `@gitlawb/zero`
|
||||
- release archives on GitHub Releases
|
||||
- source builds with Go 1.26.5+
|
||||
| Method | Command | Self-update |
|
||||
| --- | --- | --- |
|
||||
| Install script (Linux, macOS) | `curl -fsSL https://raw.githubusercontent.com/Gitlawb/zero/main/scripts/install.sh \| bash` | `zero upgrade` |
|
||||
| Install script (Windows) | `irm https://raw.githubusercontent.com/Gitlawb/zero/main/scripts/install.ps1 \| iex` | `zero upgrade` |
|
||||
| npm | `npm install -g @gitlawb/zero` | `zero upgrade` (runs npm) |
|
||||
| Release archive | download from [Releases](https://github.com/Gitlawb/zero/releases) | `zero upgrade` |
|
||||
| mise | `mise use -g ubi:Gitlawb/zero` | `mise upgrade` |
|
||||
| `go install` | `go install github.com/Gitlawb/zero/cmd/zero@latest` | rerun the command |
|
||||
| Source | `go build -o zero ./cmd/zero` | rebuild |
|
||||
|
||||
The install scripts download a platform-specific release archive and require a
|
||||
published GitHub Release for the requested version. The npm package is
|
||||
self-contained: the platform binary installs from the npm registry.
|
||||
`zero upgrade` follows whatever owns the binary, with one exception it cannot
|
||||
detect. It delegates to npm for an npm install, refuses outright for a Homebrew
|
||||
keg and names `brew upgrade zero`, and replaces the binary in place for
|
||||
everything else — including a mise-managed one, which is the exception.
|
||||
|
||||
mise is the one case it cannot detect: a mise-managed binary is an ordinary
|
||||
standalone install on disk, so `zero upgrade` will replace it and leave mise
|
||||
describing a version that is no longer there. Use `mise upgrade` instead.
|
||||
|
||||
Release archives are the substrate for most of the above: the install scripts
|
||||
and the npm fallback both download a platform archive from a published GitHub
|
||||
Release and verify its checksum.
|
||||
|
||||
## npm
|
||||
|
||||
@@ -111,6 +127,45 @@ Defaults:
|
||||
- Version: latest GitHub release
|
||||
- Install path: `%LOCALAPPDATA%\zero\bin\zero.exe`
|
||||
|
||||
## mise
|
||||
|
||||
[mise](https://mise.jdx.dev/) installs Zero straight from the GitHub Release
|
||||
archives through its `ubi` backend, so there is nothing extra to publish and no
|
||||
registry in the middle:
|
||||
|
||||
```bash
|
||||
mise use -g ubi:Gitlawb/zero
|
||||
```
|
||||
|
||||
Pin a version the same way you would any other tool:
|
||||
|
||||
```bash
|
||||
mise use -g ubi:Gitlawb/[email protected]
|
||||
```
|
||||
|
||||
Updates come from `mise upgrade`. `zero upgrade` also works, because a
|
||||
mise-managed binary is an ordinary standalone install, but then mise's records
|
||||
describe a version that is no longer on disk. Prefer `mise upgrade`.
|
||||
|
||||
## go install
|
||||
|
||||
```bash
|
||||
go install github.com/Gitlawb/zero/cmd/zero@latest
|
||||
```
|
||||
|
||||
This builds from source, so it needs Go 1.26.6+ and it does not go through the
|
||||
release archives. Two consequences worth knowing before you pick it:
|
||||
|
||||
- On Linux you also need the sandbox helper, which is a separate binary and is
|
||||
not installed by this command. See
|
||||
[Sandbox Helpers For Source Builds](#sandbox-helpers-for-source-builds).
|
||||
Without it, native sandboxing is unavailable.
|
||||
- `zero upgrade` treats the result as a standalone install and will replace the
|
||||
binary with a release build rather than rebuilding from source. If you chose
|
||||
`go install` deliberately, rerun it instead.
|
||||
|
||||
macOS and Windows need no extra helper.
|
||||
|
||||
## From Source
|
||||
|
||||
```bash
|
||||
@@ -125,7 +180,7 @@ Build a local binary:
|
||||
go build -o zero ./cmd/zero
|
||||
```
|
||||
|
||||
Source builds require Go 1.26.5+.
|
||||
Source builds require Go 1.26.6+.
|
||||
|
||||
### Sandbox Helpers For Source Builds
|
||||
|
||||
@@ -242,3 +297,23 @@ zero upgrade
|
||||
```
|
||||
|
||||
See the [update guide](UPDATE.md) for update modes, flags, and platform details.
|
||||
|
||||
`zero upgrade` behaves differently depending on how Zero was installed, because
|
||||
overwriting a binary a package manager owns leaves that manager describing a
|
||||
version that is no longer there:
|
||||
|
||||
- **npm**: runs `npm install -g @gitlawb/zero@latest` for you.
|
||||
- **Homebrew**: refuses, and tells you to run `brew upgrade zero`. Replacing the
|
||||
keg binary directly would be reverted by the next `brew upgrade` or
|
||||
`brew reinstall`.
|
||||
- **everything else**: downloads the verified release archive and replaces the
|
||||
binary in place.
|
||||
|
||||
Homebrew is detected by the binary living inside a Cellar keg that carries
|
||||
Homebrew's own `INSTALL_RECEIPT.json`. The keg shape alone is not enough: a
|
||||
directory tree you built yourself can have exactly that shape, and refusing to
|
||||
update an install Homebrew never touched would be a real cost for a made-up
|
||||
reason. The receipt is the thing only Homebrew writes.
|
||||
|
||||
An ordinary install under `/usr/local/bin` is left alone even on an Intel Mac
|
||||
where that is also the Homebrew prefix.
|
||||
@@ -70,3 +70,43 @@ Endpoint resolution order:
|
||||
Installer scripts download the matching release asset for the local platform and
|
||||
verify its `.sha256` file. If Zero is already installed, run `zero upgrade`
|
||||
instead of reinstalling.
|
||||
|
||||
## Windows recovery state (standalone installs)
|
||||
|
||||
This section describes Windows only. On Linux and macOS a standalone update
|
||||
renames the staged file directly over the executable path through the
|
||||
installation directory's file descriptor, so the replacement is atomic and no
|
||||
aside copy, marker, or recovery record is ever created — a failed update leaves
|
||||
the previous binary in place and nothing to resolve.
|
||||
|
||||
On Windows, a running executable cannot be replaced in place, so the previous
|
||||
binary is moved aside to `<binary>.zero-update-<random>.old` first. The updater
|
||||
records that exact file — bound to its filesystem identity, in per-user state
|
||||
outside the installation directory — and deletes only that recorded copy after
|
||||
the next update is verified in place. Backups it did not create are never
|
||||
removed, so a file such as `zero.exe.before-manual-patch.old` is left alone. A
|
||||
recorded copy that something else holds open (an on-access scanner, an editor)
|
||||
stays recorded and is removed by a later update instead.
|
||||
|
||||
An update **refuses to run** while unresolved recovery state exists beside the
|
||||
binary:
|
||||
|
||||
| State on disk | Meaning |
|
||||
|---|---|
|
||||
| `<binary>.old` (or `<binary>.<suffix>.old`) plus a `.keep` marker | A previous update could not restore the original binary. The `.old` file may be the last binary the updater verified; the installed one may be unverified. |
|
||||
| `<binary>.…old.<suffix>.recovery` | A previous update could not even write the marker, so it moved the last verified binary to that name. |
|
||||
| The binary is missing and one or more `*.old` files exist | The previous attempt was interrupted between moves. |
|
||||
|
||||
The refusal names the paths involved and the two moves that end the state:
|
||||
either move the recovery binary back over the executable path, or — if the
|
||||
installed binary is the one you want — delete the `.keep` marker (or the
|
||||
`.recovery` copy, once you have verified the installed binary) and update again.
|
||||
|
||||
This is deliberately fail-closed and differs from older releases, which deleted
|
||||
`<binary>.old` and proceeded. The trade-off is that anyone who can write in the
|
||||
installation directory can plant `<binary>.old` and `<binary>.old.keep` there
|
||||
and make every subsequent update refuse until an operator clears them. That is
|
||||
the safe direction: the alternative is an update that overwrites the only
|
||||
verified copy of the previous binary. If updates start refusing, inspect those
|
||||
files before removing them — an installation directory that a lower-privileged
|
||||
account can write to is itself worth fixing.
|
||||
@@ -1,6 +1,6 @@
|
||||
module github.com/Gitlawb/zero
|
||||
|
||||
go 1.26.5
|
||||
go 1.26.6
|
||||
|
||||
require (
|
||||
charm.land/bubbles/v2 v2.1.1
|
||||
@@ -11,16 +11,17 @@ require (
|
||||
github.com/atotto/clipboard v0.1.4
|
||||
github.com/aymanbagabas/go-udiff v0.4.1
|
||||
github.com/charmbracelet/colorprofile v0.4.3
|
||||
github.com/charmbracelet/ultraviolet v0.0.0-20260703014108-f5a850f9c2b7
|
||||
github.com/charmbracelet/x/ansi v0.11.7
|
||||
github.com/charmbracelet/x/term v0.2.2
|
||||
github.com/coder/websocket v1.8.15
|
||||
github.com/ledongthuc/pdf v0.0.0-20250511090121-5959a4027728
|
||||
golang.org/x/image v0.45.0
|
||||
golang.org/x/sys v0.47.0
|
||||
mvdan.cc/sh/v3 v3.13.1
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/charmbracelet/ultraviolet v0.0.0-20260703014108-f5a850f9c2b7 // indirect
|
||||
github.com/charmbracelet/x/exp/golden v0.0.0-20260615092313-b57e5e6d29bb // indirect
|
||||
github.com/charmbracelet/x/termios v0.1.1 // indirect
|
||||
github.com/charmbracelet/x/windows v0.2.2 // indirect
|
||||
|
||||
@@ -64,6 +64,10 @@ github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavM
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
|
||||
golang.org/x/exp v0.0.0-20260611194520-c48552f49976 h1:X8Hz2ImujgbmetVuW+w2YkyZChE3cBpZi2P158rTG9M=
|
||||
golang.org/x/exp v0.0.0-20260611194520-c48552f49976/go.mod h1:vnf4pv9iKZXY58sQE1L86zmNWJ4159e1RkcWiLCkeEY=
|
||||
golang.org/x/image v0.44.0 h1:+tDekMZED9+LrtB3G5xzRggpVh9CARjZqROla3R3R+I=
|
||||
golang.org/x/image v0.44.0/go.mod h1:V8K3KE9KKKE+pLpQDOeN18w9oacNSvy1tDOirTu4xtY=
|
||||
golang.org/x/image v0.45.0 h1:FMb1nTbH5H9vF55SriQHgFw5GnNL9Jg6L25BwXKzhB0=
|
||||
golang.org/x/image v0.45.0/go.mod h1:n62x/7RqlwXDvGsSU4u6IUTUf6KghUZ9Bt7cG/T9Fx4=
|
||||
golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek=
|
||||
golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
|
||||
|
||||
+16
-2
@@ -296,6 +296,16 @@ func stopReasonFor(result agent.Result, err error) (string, error) {
|
||||
if errors.Is(err, context.Canceled) {
|
||||
return StopCancelled, nil
|
||||
}
|
||||
// A CLIENT THAT CANCELLED A PERMISSION PROMPT DID NOT FAIL. Only
|
||||
// context.Canceled was matched here, so dismissing a permission dialog —
|
||||
// a deliberate action, and for apply_patch the only refusal a client is
|
||||
// offered — came back as JSON-RPC -32603 carrying the internal sentinel
|
||||
// text. Editors and the desktop app both render that as a crashed turn,
|
||||
// so declining a tool looked like ZERO falling over. It is a
|
||||
// cancellation, and StopCancelled is what ACP has for saying so.
|
||||
if errors.Is(err, agent.ErrPermissionApprovalCanceled) {
|
||||
return StopCancelled, nil
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
if result.FinishReason == "length" {
|
||||
@@ -311,10 +321,14 @@ func stopReasonFor(result agent.Result, err error) (string, error) {
|
||||
// session/request_permission request and maps the outcome back. Failure to reach
|
||||
// the client fails closed to deny.
|
||||
func (a *Agent) requestPermission(ctx context.Context, sessionID string, req agent.PermissionRequest) (agent.PermissionDecision, error) {
|
||||
// Built ONCE and used for both halves. Sending one list and validating the
|
||||
// reply against another is the defect this whole path had; see
|
||||
// decisionFromOutcome.
|
||||
options := buildPermissionOptions(req)
|
||||
params := RequestPermissionParams{
|
||||
SessionID: sessionID,
|
||||
ToolCall: permissionToolCall(req),
|
||||
Options: buildPermissionOptions(req),
|
||||
Options: options,
|
||||
}
|
||||
var result RequestPermissionResult
|
||||
if err := a.conn.Call(ctx, MethodSessionRequestPerm, params, &result); err != nil {
|
||||
@@ -323,7 +337,7 @@ func (a *Agent) requestPermission(ctx context.Context, sessionID string, req age
|
||||
}
|
||||
return agent.PermissionDecision{Action: agent.PermissionDecisionDeny, Reason: "permission request failed: " + err.Error()}, nil
|
||||
}
|
||||
return decisionFromOutcome(result.Outcome, req.AvailableDecisions), nil
|
||||
return decisionFromOutcome(result.Outcome, options), nil
|
||||
}
|
||||
|
||||
func (a *Agent) emitPlan(registry *tools.Registry, note *notifier) {
|
||||
|
||||
@@ -0,0 +1,217 @@
|
||||
package acp
|
||||
|
||||
// desktopinterop_test.go covers three defects found by driving the real `zero
|
||||
// acp` binary from a desktop ACP client. Each is a case where ZERO and a
|
||||
// conforming client disagreed about the wire, and none of them failed loudly —
|
||||
// they produced a silent denial, a spurious crash, or two buttons the user
|
||||
// could not tell apart.
|
||||
//
|
||||
// A fourth, session/load restoring history without replaying it, is fixed by
|
||||
// #914 rather than here: that one also gives the replayed messages stable ids
|
||||
// and keeps session/resume replay-free, which this did not.
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/agent"
|
||||
)
|
||||
|
||||
// ---- an offered option is an acceptable answer ----
|
||||
|
||||
// The list the options were BUILT from and the list the answer was CHECKED
|
||||
// against were different whenever ZERO did not enumerate: the client was sent
|
||||
// Allow and Reject, and its reply was validated against an empty slice, so
|
||||
// every button it could possibly show failed closed to deny. The user clicked
|
||||
// Allow and ZERO recorded a denial, with nothing on screen to say so.
|
||||
func TestAnOfferedOptionIsAccepted(t *testing.T) {
|
||||
// No AvailableDecisions: the fallback path, which is every permission event
|
||||
// that is not a prompt.
|
||||
req := agent.PermissionRequest{ToolName: "bash"}
|
||||
|
||||
options := buildPermissionOptions(req)
|
||||
if len(options) == 0 {
|
||||
t.Fatal("no options were offered")
|
||||
}
|
||||
|
||||
for _, option := range options {
|
||||
decision := decisionFromOutcome(
|
||||
RequestPermissionOutcome{Outcome: OutcomeSelected, OptionID: option.OptionID},
|
||||
buildPermissionOptions(req),
|
||||
)
|
||||
if string(decision.Action) != option.OptionID {
|
||||
t.Fatalf("option %q was answered with %q (%s) — an offered option must be accepted",
|
||||
option.OptionID, decision.Action, decision.Reason)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The narrowing this validation exists for still holds: a client must not be
|
||||
// able to return a broader grant than it was shown.
|
||||
func TestAnOptionThatWasNotOfferedIsStillRefused(t *testing.T) {
|
||||
req := agent.PermissionRequest{ToolName: "bash"}
|
||||
|
||||
decision := decisionFromOutcome(
|
||||
RequestPermissionOutcome{Outcome: OutcomeSelected, OptionID: string(agent.PermissionDecisionAlwaysAllow)},
|
||||
buildPermissionOptions(req),
|
||||
)
|
||||
if decision.Action != agent.PermissionDecisionDeny {
|
||||
t.Fatalf("an unoffered broader grant was accepted as %q", decision.Action)
|
||||
}
|
||||
}
|
||||
|
||||
// The resolver is the single source of truth for both sides, which is what
|
||||
// makes the two agree by construction rather than by being kept in step.
|
||||
func TestOfferedDecisionsMatchWhatWasSent(t *testing.T) {
|
||||
for _, req := range []agent.PermissionRequest{
|
||||
{ToolName: "bash"},
|
||||
{ToolName: "bash", AvailableDecisions: []agent.PermissionDecisionAction{
|
||||
agent.PermissionDecisionAllow,
|
||||
agent.PermissionDecisionAllowForSession,
|
||||
agent.PermissionDecisionDeny,
|
||||
}},
|
||||
} {
|
||||
offered := buildPermissionOptions(req)
|
||||
for _, option := range buildPermissionOptions(req) {
|
||||
if !actionOffered(option.OptionID, offered) {
|
||||
t.Fatalf("option %q was sent but is not in the offered set %v", option.OptionID, offered)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---- two options the user can tell apart ----
|
||||
|
||||
// request_permissions offers plain allow and strict-review allow together, and
|
||||
// both were labelled "Allow". The label is the only thing an ACP client shows,
|
||||
// so the panel presented two identical buttons where one silently enabled
|
||||
// strict auto-review of what was granted.
|
||||
func TestEveryOfferedOptionHasADistinctLabel(t *testing.T) {
|
||||
req := agent.PermissionRequest{
|
||||
ToolName: "request_permissions",
|
||||
AvailableDecisions: []agent.PermissionDecisionAction{
|
||||
agent.PermissionDecisionAllow,
|
||||
agent.PermissionDecisionAllowStrict,
|
||||
agent.PermissionDecisionAllowForSession,
|
||||
agent.PermissionDecisionDeny,
|
||||
},
|
||||
}
|
||||
|
||||
seen := map[string]string{}
|
||||
for _, option := range buildPermissionOptions(req) {
|
||||
if previous, clash := seen[option.Name]; clash {
|
||||
t.Fatalf("options %q and %q share the label %q — a user cannot tell them apart",
|
||||
previous, option.OptionID, option.Name)
|
||||
}
|
||||
seen[option.Name] = option.OptionID
|
||||
}
|
||||
}
|
||||
|
||||
// The two remain distinct decisions: the label changed, the round trip did not.
|
||||
func TestStrictAllowStillRoundTrips(t *testing.T) {
|
||||
req := agent.PermissionRequest{
|
||||
ToolName: "request_permissions",
|
||||
AvailableDecisions: []agent.PermissionDecisionAction{
|
||||
agent.PermissionDecisionAllow,
|
||||
agent.PermissionDecisionAllowStrict,
|
||||
},
|
||||
}
|
||||
decision := decisionFromOutcome(
|
||||
RequestPermissionOutcome{Outcome: OutcomeSelected, OptionID: string(agent.PermissionDecisionAllowStrict)},
|
||||
buildPermissionOptions(req),
|
||||
)
|
||||
if decision.Action != agent.PermissionDecisionAllowStrict {
|
||||
t.Fatalf("strict allow round-tripped as %q", decision.Action)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- cancelling a permission is a cancellation, not a crash ----
|
||||
|
||||
// Only context.Canceled was recognised, so dismissing a permission dialog came
|
||||
// back as JSON-RPC -32603 carrying the internal sentinel text. Clients render
|
||||
// that as a failed turn, so declining a tool looked like ZERO falling over —
|
||||
// and for apply_patch, dismissing is the ONLY refusal a client is offered.
|
||||
func TestCancellingAPermissionEndsTheTurnAsCancelled(t *testing.T) {
|
||||
canceled := fmt.Errorf("%w for apply_patch", agent.ErrPermissionApprovalCanceled)
|
||||
|
||||
reason, err := stopReasonFor(agent.Result{}, canceled)
|
||||
if err != nil {
|
||||
t.Fatalf("stopReasonFor returned an error for a cancellation: %v", err)
|
||||
}
|
||||
if reason != StopCancelled {
|
||||
t.Fatalf("stop reason = %q, want %q", reason, StopCancelled)
|
||||
}
|
||||
}
|
||||
|
||||
// A genuine failure must still be an error. Mapping too much to cancelled would
|
||||
// hide real faults behind a clean ending, which is the opposite mistake.
|
||||
func TestARealFailureIsStillAnError(t *testing.T) {
|
||||
failure := errors.New("provider exploded")
|
||||
|
||||
if _, err := stopReasonFor(agent.Result{}, failure); !errors.Is(err, failure) {
|
||||
t.Fatalf("stopReasonFor(%v) = %v, want the failure preserved", failure, err)
|
||||
}
|
||||
}
|
||||
|
||||
// The sentinel is matched through errors.Is, so it survives the wrapping every
|
||||
// return site applies — each one adds the tool name.
|
||||
func TestPermissionCancellationSurvivesWrapping(t *testing.T) {
|
||||
if !errors.Is(fmt.Errorf("%w for bash", agent.ErrPermissionApprovalCanceled), agent.ErrPermissionApprovalCanceled) {
|
||||
t.Fatal("a wrapped permission cancellation was not recognised")
|
||||
}
|
||||
if errors.Is(errors.New("something else"), agent.ErrPermissionApprovalCanceled) {
|
||||
t.Fatal("an unrelated error was reported as a permission cancellation")
|
||||
}
|
||||
}
|
||||
|
||||
// A decision that never became an option cannot be selected.
|
||||
//
|
||||
// PermissionDecisionCancel is the case that exists today: optionKindFor drops
|
||||
// it because ACP expresses cancellation through the outcome, yet ZERO
|
||||
// enumerates it in AvailableDecisions for shell commands and apply_patch. While
|
||||
// the reply was validated against the DECISIONS rather than the sent options, a
|
||||
// client could return {"outcome":"selected","optionId":"cancel"} — an id it was
|
||||
// never shown — and abort the whole turn.
|
||||
//
|
||||
// Written over every option-less decision rather than over "cancel", so a
|
||||
// future action that optionKindFor declines to render is covered the day it is
|
||||
// added rather than the day someone remembers this test.
|
||||
func TestADecisionThatIsNotAnOptionCannotBeSelected(t *testing.T) {
|
||||
for _, toolName := range []string{"bash", "apply_patch"} {
|
||||
req := agent.PermissionRequest{
|
||||
ToolName: toolName,
|
||||
AvailableDecisions: []agent.PermissionDecisionAction{
|
||||
agent.PermissionDecisionAllow,
|
||||
agent.PermissionDecisionDeny,
|
||||
agent.PermissionDecisionCancel,
|
||||
},
|
||||
}
|
||||
options := buildPermissionOptions(req)
|
||||
|
||||
for _, action := range req.AvailableDecisions {
|
||||
if actionOffered(string(action), options) {
|
||||
continue // it was sent; selecting it is legitimate
|
||||
}
|
||||
decision := decisionFromOutcome(
|
||||
RequestPermissionOutcome{Outcome: OutcomeSelected, OptionID: string(action)},
|
||||
options,
|
||||
)
|
||||
if decision.Action != agent.PermissionDecisionDeny {
|
||||
t.Fatalf("%s: %q was never sent as an option but was accepted as %q",
|
||||
toolName, action, decision.Action)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Cancelling is still reachable — through the outcome, which is where ACP puts
|
||||
// it. Closing the selected-id route must not close the real one.
|
||||
func TestCancellingThroughTheOutcomeStillWorks(t *testing.T) {
|
||||
options := buildPermissionOptions(agent.PermissionRequest{ToolName: "apply_patch"})
|
||||
|
||||
decision := decisionFromOutcome(RequestPermissionOutcome{Outcome: OutcomeCancelled}, options)
|
||||
if decision.Action != agent.PermissionDecisionCancel {
|
||||
t.Fatalf("outcome=cancelled gave %q, want cancel", decision.Action)
|
||||
}
|
||||
}
|
||||
+57
-18
@@ -15,14 +15,7 @@ import (
|
||||
// PermissionOptions. Only the actions ZERO actually presented (AvailableDecisions)
|
||||
// are surfaced; the optionId is the ZERO action string for a clean round-trip.
|
||||
func buildPermissionOptions(req agent.PermissionRequest) []PermissionOption {
|
||||
actions := req.AvailableDecisions
|
||||
if len(actions) == 0 {
|
||||
// Sensible default if ZERO didn't enumerate: allow once / reject.
|
||||
actions = []agent.PermissionDecisionAction{
|
||||
agent.PermissionDecisionAllow,
|
||||
agent.PermissionDecisionDeny,
|
||||
}
|
||||
}
|
||||
actions := offeredDecisions(req)
|
||||
options := make([]PermissionOption, 0, len(actions))
|
||||
for _, action := range actions {
|
||||
kind, name := optionKindFor(action, req.PrefixApprovalEscalates)
|
||||
@@ -38,6 +31,32 @@ func buildPermissionOptions(req agent.PermissionRequest) []PermissionOption {
|
||||
return options
|
||||
}
|
||||
|
||||
// offeredDecisions resolves what this request actually offers the client.
|
||||
//
|
||||
// ONE RESOLVER, USED BY BOTH SIDES, and that is the entire point of it existing.
|
||||
// The fallback below used to live inside buildPermissionOptions while
|
||||
// requestPermission validated the client's answer against the raw
|
||||
// req.AvailableDecisions. When ZERO did not enumerate — which is every
|
||||
// permission event that is not a prompt, and includes a shell call whose sandbox
|
||||
// decision came back deny — the two disagreed completely: the client was sent
|
||||
// "Allow" and "Reject", and validation was performed against an EMPTY list, so
|
||||
// every option the client could possibly click failed closed to deny with
|
||||
// "permission option was not offered".
|
||||
//
|
||||
// The user clicked Allow and ZERO recorded a denial. Nothing surfaced it,
|
||||
// because a denial is a legitimate answer — the tool was simply refused and the
|
||||
// turn carried on as though the user had rejected it.
|
||||
func offeredDecisions(req agent.PermissionRequest) []agent.PermissionDecisionAction {
|
||||
if len(req.AvailableDecisions) > 0 {
|
||||
return req.AvailableDecisions
|
||||
}
|
||||
// Sensible default if ZERO didn't enumerate: allow once / reject.
|
||||
return []agent.PermissionDecisionAction{
|
||||
agent.PermissionDecisionAllow,
|
||||
agent.PermissionDecisionDeny,
|
||||
}
|
||||
}
|
||||
|
||||
// optionKindFor maps a ZERO decision action to an ACP PermissionOptionKind and a
|
||||
// human label. Returns an empty kind for actions that ACP expresses through the
|
||||
// outcome rather than an option (cancel).
|
||||
@@ -49,8 +68,16 @@ func buildPermissionOptions(req agent.PermissionRequest) []PermissionOption {
|
||||
// disclosure cannot be TUI-only.
|
||||
func optionKindFor(action agent.PermissionDecisionAction, escalates bool) (kind, name string) {
|
||||
switch action {
|
||||
case agent.PermissionDecisionAllow, agent.PermissionDecisionAllowStrict:
|
||||
case agent.PermissionDecisionAllow:
|
||||
return PermAllowOnce, "Allow"
|
||||
case agent.PermissionDecisionAllowStrict:
|
||||
// SPLIT FROM PLAIN ALLOW because the label is the only thing an ACP
|
||||
// client shows. Both were labelled "Allow", and request_permissions
|
||||
// offers them together — so the panel presented two identical buttons,
|
||||
// and whichever the user picked looked the same while one of them
|
||||
// silently turned on strict auto-review of what was granted. The option
|
||||
// id carries no room for the distinction, so it has to be in the name.
|
||||
return PermAllowOnce, "Allow with strict review"
|
||||
case agent.PermissionDecisionAllowForSession:
|
||||
return PermAllowAlways, "Allow for this session"
|
||||
case agent.PermissionDecisionAllowPrefix:
|
||||
@@ -70,19 +97,31 @@ func optionKindFor(action agent.PermissionDecisionAction, escalates bool) (kind,
|
||||
|
||||
// decisionFromOutcome maps the client's permission outcome back to a ZERO
|
||||
// decision. A cancelled outcome cancels the run; a selected option id is the ZERO
|
||||
// action verbatim (validated against what was offered); anything unrecognized
|
||||
// fails closed to deny.
|
||||
func decisionFromOutcome(outcome RequestPermissionOutcome, offered []agent.PermissionDecisionAction) agent.PermissionDecision {
|
||||
// action verbatim (validated against what was actually sent); anything
|
||||
// unrecognized fails closed to deny.
|
||||
//
|
||||
// IT TAKES THE OPTIONS THAT WERE SENT, not the decisions the request offered,
|
||||
// and the two are not the same list. Some actions never become options —
|
||||
// PermissionDecisionCancel is dropped by optionKindFor because ACP expresses
|
||||
// cancellation through the outcome instead — yet it IS enumerated in
|
||||
// AvailableDecisions for shell commands and apply_patch, the two commonest
|
||||
// prompts. Validating against the decisions therefore accepted
|
||||
// {"outcome":"selected","optionId":"cancel"}: an identifier no client was ever
|
||||
// shown, aborting the whole turn.
|
||||
//
|
||||
// Passing the sent options makes "what we offered" and "what we accept" the
|
||||
// same slice rather than two things kept in step, so anything optionKindFor
|
||||
// declines to render is excluded structurally instead of by name.
|
||||
func decisionFromOutcome(outcome RequestPermissionOutcome, offered []PermissionOption) agent.PermissionDecision {
|
||||
switch outcome.Outcome {
|
||||
case OutcomeCancelled:
|
||||
return agent.PermissionDecision{Action: agent.PermissionDecisionCancel, Reason: "client cancelled"}
|
||||
case OutcomeSelected:
|
||||
action := agent.PermissionDecisionAction(outcome.OptionID)
|
||||
// Bind to what was actually offered for THIS call: a client must not be able
|
||||
// to return a broader grant (always_allow / allow_for_session) that wasn't
|
||||
// presented. Anything not offered fails closed to deny.
|
||||
if actionOffered(action, offered) {
|
||||
return agent.PermissionDecision{Action: action}
|
||||
if actionOffered(outcome.OptionID, offered) {
|
||||
return agent.PermissionDecision{Action: agent.PermissionDecisionAction(outcome.OptionID)}
|
||||
}
|
||||
return agent.PermissionDecision{Action: agent.PermissionDecisionDeny, Reason: "permission option was not offered"}
|
||||
default:
|
||||
@@ -90,9 +129,9 @@ func decisionFromOutcome(outcome RequestPermissionOutcome, offered []agent.Permi
|
||||
}
|
||||
}
|
||||
|
||||
func actionOffered(action agent.PermissionDecisionAction, offered []agent.PermissionDecisionAction) bool {
|
||||
for _, a := range offered {
|
||||
if a == action {
|
||||
func actionOffered(optionID string, offered []PermissionOption) bool {
|
||||
for _, option := range offered {
|
||||
if option.OptionID == optionID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
@@ -56,10 +56,11 @@ func TestACPPrefixOptionsStaySilentWithoutEscalation(t *testing.T) {
|
||||
// what was offered, so decorating it would fail every prefix approval closed.
|
||||
func TestACPDisclosureDoesNotDisturbTheOptionIDRoundTrip(t *testing.T) {
|
||||
request := escalationRequest(true)
|
||||
for _, option := range buildPermissionOptions(request) {
|
||||
options := buildPermissionOptions(request)
|
||||
for _, option := range options {
|
||||
decision := decisionFromOutcome(
|
||||
RequestPermissionOutcome{Outcome: OutcomeSelected, OptionID: option.OptionID},
|
||||
request.AvailableDecisions,
|
||||
options,
|
||||
)
|
||||
if string(decision.Action) != option.OptionID {
|
||||
t.Errorf("option id %q round-tripped to %q (%s)", option.OptionID, decision.Action, decision.Reason)
|
||||
|
||||
@@ -39,11 +39,14 @@ func TestBuildPermissionOptionsDefault(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestDecisionFromOutcome(t *testing.T) {
|
||||
offered := []agent.PermissionDecisionAction{
|
||||
agent.PermissionDecisionAllow,
|
||||
agent.PermissionDecisionAlwaysAllow,
|
||||
agent.PermissionDecisionDeny,
|
||||
}
|
||||
offered := buildPermissionOptions(agent.PermissionRequest{
|
||||
ToolName: "bash",
|
||||
AvailableDecisions: []agent.PermissionDecisionAction{
|
||||
agent.PermissionDecisionAllow,
|
||||
agent.PermissionDecisionAlwaysAllow,
|
||||
agent.PermissionDecisionDeny,
|
||||
},
|
||||
})
|
||||
if d := decisionFromOutcome(RequestPermissionOutcome{Outcome: OutcomeCancelled}, offered); d.Action != agent.PermissionDecisionCancel {
|
||||
t.Errorf("cancelled -> %q, want cancel", d.Action)
|
||||
}
|
||||
|
||||
@@ -99,6 +99,43 @@ func TestProposedCommandPrefixRejectsUnsafeRequestedPrefix(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// The launcher denylist is only worth anything if it holds on the path a user
|
||||
// actually sees. A model naming its own prefix_rule is that path: before
|
||||
// launcher names were normalized, "python3.11" was offered as a one-token
|
||||
// grant, and hasStringPrefix then matched every later python3.11 invocation.
|
||||
func TestProposedCommandPrefixRejectsLauncherSpellingsInRequestedPrefix(t *testing.T) {
|
||||
for _, testCase := range []struct {
|
||||
command string
|
||||
rule []any
|
||||
}{
|
||||
{"python3.11 script.py", []any{"python3.11"}},
|
||||
{"python.exe script.py", []any{"python.exe"}},
|
||||
{"node.exe app.js", []any{"node.exe"}},
|
||||
{"cmd /c whoami", []any{"cmd"}},
|
||||
{"busybox sh -c id", []any{"busybox"}},
|
||||
{"uv run main.py", []any{"uv", "run"}},
|
||||
{"python3.11 -c import_os", []any{"python3.11", "-c"}},
|
||||
} {
|
||||
got := proposedCommandPrefix("bash", map[string]any{
|
||||
"command": testCase.command,
|
||||
"prefix_rule": testCase.rule,
|
||||
})
|
||||
if got != nil {
|
||||
t.Errorf("prefix_rule %v for %q was offered as %#v, want rejected", testCase.rule, testCase.command, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestProposedCommandPrefixStillOffersOrdinaryCommands(t *testing.T) {
|
||||
got := proposedCommandPrefix("bash", map[string]any{
|
||||
"command": "cargo build --release",
|
||||
"prefix_rule": []any{"cargo", "build"},
|
||||
})
|
||||
if len(got) != 2 || got[0] != "cargo" || got[1] != "build" {
|
||||
t.Fatalf("ordinary prefix rule = %#v, want [cargo build]", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProposedCommandPrefixRejectsUnsafeShellForms(t *testing.T) {
|
||||
cases := []string{
|
||||
"cat < in > out",
|
||||
|
||||
+172
-11
@@ -4,7 +4,9 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strconv"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/trace"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
@@ -390,6 +392,27 @@ type compactionState struct {
|
||||
// starts at 1.0 and converges via an EMA as each turn reports actual usage, so
|
||||
// later turns compact nearer to real capacity. Zero is treated as 1.0.
|
||||
calibrationRatio float64
|
||||
|
||||
// summarizerFactory lazily builds the dedicated (cheap) summarizer provider
|
||||
// (Options.Summarizer) for the current main model. summarizerProvider
|
||||
// memoizes the built provider and summarizerResolved records that the
|
||||
// factory ran (it may legitimately return no provider: the main model is
|
||||
// already the cheap one). summarizerBroken is the sticky fallback: once a
|
||||
// build or call fails, the run stays on the main provider until the next
|
||||
// model switch — a misconfigured cheap model must never break (or
|
||||
// repeatedly delay) compaction.
|
||||
summarizerFactory func(ctx context.Context, mainModelID string) (Provider, error)
|
||||
summarizerProvider Provider
|
||||
summarizerResolved bool
|
||||
summarizerBroken bool
|
||||
// window is the context window the threshold was derived from. The
|
||||
// reactive free-prune stage compares against it (the limit the provider
|
||||
// enforces), not against the 70% proactive threshold — a pruned history
|
||||
// between the two would still succeed on retry.
|
||||
window int
|
||||
// mainModel is the model the run currently streams on; the summarizer
|
||||
// choice is made against it and remade after a mid-run switch.
|
||||
mainModel string
|
||||
}
|
||||
|
||||
// calibrate folds one turn's (rawEstimate, actualPromptTokens) sample into the
|
||||
@@ -424,17 +447,88 @@ func (state *compactionState) calibratedTokens(raw int) int {
|
||||
|
||||
func newCompactionState(options Options, task *taskState) *compactionState {
|
||||
state := &compactionState{
|
||||
enabled: options.ContextWindow > 0,
|
||||
threshold: compactionThreshold(options.ContextWindow),
|
||||
preserveLast: options.CompactionPreserveLast,
|
||||
onUsage: options.OnUsage,
|
||||
task: task,
|
||||
planner: newContextPlanner(contextPlannerConfig{contextWindow: options.ContextWindow}),
|
||||
trace: options.Trace,
|
||||
enabled: options.ContextWindow > 0,
|
||||
threshold: compactionThreshold(options.ContextWindow),
|
||||
window: options.ContextWindow,
|
||||
mainModel: options.Model,
|
||||
preserveLast: options.CompactionPreserveLast,
|
||||
onUsage: options.OnUsage,
|
||||
task: task,
|
||||
planner: newContextPlanner(contextPlannerConfig{contextWindow: options.ContextWindow, serviceTier: options.ServiceTier}),
|
||||
trace: options.Trace,
|
||||
summarizerFactory: options.Summarizer,
|
||||
}
|
||||
return state
|
||||
}
|
||||
|
||||
// updateWindow re-derives the compaction threshold after a mid-run model
|
||||
// switch. The low-water mark is reset: it was measured against the old
|
||||
// threshold, and keeping it could suppress a compaction the smaller window
|
||||
// now requires. <= 0 (unknown model) leaves everything unchanged.
|
||||
func (state *compactionState) updateWindow(window int) {
|
||||
if window <= 0 {
|
||||
return
|
||||
}
|
||||
state.enabled = true
|
||||
state.window = window
|
||||
state.threshold = compactionThreshold(window)
|
||||
state.lowWaterMark = 0
|
||||
if state.planner != nil {
|
||||
state.planner.config.contextWindow = window
|
||||
}
|
||||
}
|
||||
|
||||
// switchModel records a mid-run model switch: the window is re-derived
|
||||
// (window <= 0 keeps the old one) and the summarizer choice is forgotten so
|
||||
// the next compaction re-resolves it against the new main model. A run that
|
||||
// started on the cheap model (no dedicated summarizer) and escalated to an
|
||||
// expensive one therefore regains cheap summaries — the exact workflow a
|
||||
// dedicated summarizer exists for.
|
||||
func (state *compactionState) switchModel(modelID string, window int) {
|
||||
state.mainModel = modelID
|
||||
state.updateWindow(window)
|
||||
state.summarizerProvider = nil
|
||||
state.summarizerResolved = false
|
||||
state.summarizerBroken = false
|
||||
}
|
||||
|
||||
// summarizeProvider returns the provider summarization calls should use and
|
||||
// whether it is the dedicated summarizer: that one when it applies to the
|
||||
// current main model and is healthy, else the main provider. The factory runs
|
||||
// at most once per main model; a build error marks the summarizer broken for
|
||||
// that model. Callers branch on the boolean rather than comparing Provider
|
||||
// values — interface comparison panics on a non-comparable dynamic type.
|
||||
func (state *compactionState) summarizeProvider(ctx context.Context, main Provider) (Provider, bool) {
|
||||
if state.summarizerFactory == nil || state.summarizerBroken {
|
||||
return main, false
|
||||
}
|
||||
if !state.summarizerResolved {
|
||||
state.summarizerResolved = true
|
||||
built, err := state.summarizerFactory(ctx, state.mainModel)
|
||||
if err != nil {
|
||||
state.summarizerBroken = true
|
||||
return main, false
|
||||
}
|
||||
state.summarizerProvider = built
|
||||
}
|
||||
if state.summarizerProvider == nil {
|
||||
return main, false
|
||||
}
|
||||
return state.summarizerProvider, true
|
||||
}
|
||||
|
||||
// fitsWindow reports whether a calibrated size is safely under the limit the
|
||||
// provider enforces. The proactive threshold (70% of the window) is a trigger
|
||||
// for the paid summarizer, not the hard limit; the reactive path compares
|
||||
// against the window itself, with a small margin for estimate error, so a
|
||||
// history pruned to, say, 85% of the window takes the free retry.
|
||||
func (state *compactionState) fitsWindow(size int) bool {
|
||||
if state.window <= 0 {
|
||||
return size <= state.threshold
|
||||
}
|
||||
return size <= state.window-state.window/20
|
||||
}
|
||||
|
||||
// maybeCompact runs proactive compaction at the top of a turn. It returns the
|
||||
// (possibly compacted) message slice. It is a no-op when compaction is disabled,
|
||||
// when the history is under threshold, or when the history has not grown past
|
||||
@@ -528,10 +622,34 @@ func (state *compactionState) recover(
|
||||
// context-limit error never triggers an unexpected summarization call.
|
||||
return messages, false, nil
|
||||
}
|
||||
if state.reactiveAttempted {
|
||||
if !isContextLimitError(errorMessage) {
|
||||
return messages, false, nil
|
||||
}
|
||||
if !isContextLimitError(errorMessage) {
|
||||
|
||||
// CHEAP FIRST STAGE (mirrors maybeCompact): a context-limit failure can
|
||||
// often be fixed for free by dropping superseded reads and pruning old
|
||||
// tool-result bodies. It runs regardless of the one-shot reactive budget
|
||||
// (it costs nothing and cannot loop: a second pass finds nothing left to
|
||||
// prune). The free retry is taken only when the pruned history fits the
|
||||
// window the provider enforces — otherwise a retry would fail the same
|
||||
// way, so fall through to the paid summarizer in this same call.
|
||||
toolTokens := estimateToolDefTokens(tools)
|
||||
prunedAny := false
|
||||
if pruned, reclaimed := pruneSupersededReadResults(messages); reclaimed > 0 {
|
||||
messages, prunedAny = pruned, true
|
||||
}
|
||||
if pruned, reclaimed := pruneStaleToolOutput(messages, state.preserveLast); reclaimed > 0 {
|
||||
messages, prunedAny = pruned, true
|
||||
}
|
||||
if prunedAny {
|
||||
size := state.calibratedTokens(estimateTokens(messages) + toolTokens)
|
||||
if state.fitsWindow(size) {
|
||||
state.lowWaterMark = size
|
||||
return messages, true, nil
|
||||
}
|
||||
}
|
||||
|
||||
if state.reactiveAttempted {
|
||||
return messages, false, nil
|
||||
}
|
||||
|
||||
@@ -571,11 +689,21 @@ func (state *compactionState) recover(
|
||||
// provider call. The summary stream intentionally does NOT forward OnText (so
|
||||
// compaction stays invisible on the user-facing surface), but it DOES forward
|
||||
// OnUsage so the summarizer's token cost is still counted by usage/budgeting.
|
||||
// The call goes to the dedicated (cheap) summarizer when one applies; if that
|
||||
// call fails the summarizer is marked broken for the current model and the
|
||||
// same slice is retried once on the main provider, so the caller sees at most
|
||||
// one (main-provider-quality) error.
|
||||
func (state *compactionState) summarizeClosure(ctx context.Context, provider Provider) func([]zeroruntime.Message) (string, error) {
|
||||
if state.planner == nil {
|
||||
state.planner = newContextPlanner(contextPlannerConfig{})
|
||||
}
|
||||
return func(toSummarize []zeroruntime.Message) (string, error) {
|
||||
chosen, dedicated := state.summarizeProvider(ctx, provider)
|
||||
summary, err := summarizeWithPlanner(ctx, chosen, toSummarize, state.onUsage, state.planner, state.trace)
|
||||
if err == nil || !dedicated {
|
||||
return summary, err
|
||||
}
|
||||
state.summarizerBroken = true
|
||||
return summarizeWithPlanner(ctx, provider, toSummarize, state.onUsage, state.planner, state.trace)
|
||||
}
|
||||
}
|
||||
@@ -659,8 +787,41 @@ func summarizeMessagesOnce(ctx context.Context, provider Provider, messages []ze
|
||||
return summary, nil
|
||||
}
|
||||
|
||||
// summaryToolResultClamp bounds each tool-result body in the summarizer input.
|
||||
// Tool outputs dominate the elided middle's bulk (build logs, file reads,
|
||||
// search results), yet the summary instructions only need their gist — the
|
||||
// model that ACTED on the full output already turned it into decisions the
|
||||
// surrounding messages record. 2,000 chars (~500 tokens) keeps the useful
|
||||
// head+tail of a result while cutting the summarizer input by 5-10x on
|
||||
// tool-heavy sessions; the summarizer's own cost is the price being cut.
|
||||
const summaryToolResultClamp = 2000
|
||||
|
||||
// summaryToolArgsClamp bounds each tool-call argument blob in the transcript:
|
||||
// a write_file call carries the whole file body in its arguments.
|
||||
const summaryToolArgsClamp = 500
|
||||
|
||||
// clampForSummary head+tail-truncates text for the summarizer transcript,
|
||||
// snapping to rune boundaries so the clamp never splits a UTF-8 sequence.
|
||||
func clampForSummary(text string, max int) string {
|
||||
if max <= 0 || len(text) <= max {
|
||||
return text
|
||||
}
|
||||
head := max * 3 / 4
|
||||
tail := max - head
|
||||
for head > 0 && !utf8.RuneStart(text[head]) {
|
||||
head--
|
||||
}
|
||||
start := len(text) - tail
|
||||
for start < len(text) && !utf8.RuneStart(text[start]) {
|
||||
start++
|
||||
}
|
||||
return text[:head] + "\n…[" + strconv.Itoa(len(text)-head-(len(text)-start)) + " chars omitted for summarization]…\n" + text[start:]
|
||||
}
|
||||
|
||||
// renderTranscript flattens messages into a plain-text transcript for the
|
||||
// summarizer. Secret scrubbing already happened upstream at the tool boundary.
|
||||
// Tool-result bodies and tool-call arguments are clamped (see the constants
|
||||
// above): the summary needs what was done and decided, not the raw payloads.
|
||||
func renderTranscript(messages []zeroruntime.Message) string {
|
||||
lines := make([]string, 0, len(messages))
|
||||
for _, message := range messages {
|
||||
@@ -670,13 +831,13 @@ func renderTranscript(messages []zeroruntime.Message) string {
|
||||
if len(message.ToolCalls) > 0 {
|
||||
calls := make([]string, 0, len(message.ToolCalls))
|
||||
for _, call := range message.ToolCalls {
|
||||
calls = append(calls, call.Name+"("+call.Arguments+")")
|
||||
calls = append(calls, call.Name+"("+clampForSummary(call.Arguments, summaryToolArgsClamp)+")")
|
||||
}
|
||||
line += "\n[tool calls: " + strings.Join(calls, "; ") + "]"
|
||||
}
|
||||
lines = append(lines, line)
|
||||
case zeroruntime.MessageRoleTool:
|
||||
lines = append(lines, "tool result: "+message.Content)
|
||||
lines = append(lines, "tool result: "+clampForSummary(message.Content, summaryToolResultClamp))
|
||||
default:
|
||||
lines = append(lines, string(message.Role)+": "+message.Content)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,469 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// compactableHistory returns a history whose middle is large enough to trip
|
||||
// the proactive threshold for a 1000-token window.
|
||||
func compactableHistory() []zeroruntime.Message {
|
||||
return []zeroruntime.Message{
|
||||
{Role: zeroruntime.MessageRoleSystem, Content: "sys"},
|
||||
{Role: zeroruntime.MessageRoleUser, Content: strings.Repeat("u", 4000)},
|
||||
{Role: zeroruntime.MessageRoleAssistant, Content: strings.Repeat("a", 4000)},
|
||||
{Role: zeroruntime.MessageRoleUser, Content: "u2"},
|
||||
{Role: zeroruntime.MessageRoleAssistant, Content: "a2"},
|
||||
{Role: zeroruntime.MessageRoleUser, Content: "u3"},
|
||||
}
|
||||
}
|
||||
|
||||
// Summarization must run on the dedicated (cheap) summarizer when one is
|
||||
// configured — the main provider sees no summarize traffic at all.
|
||||
func TestMaybeCompactUsesDedicatedSummarizer(t *testing.T) {
|
||||
cheap := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventText, Content: "CHEAP SUMMARY"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
main := &mockProvider{}
|
||||
factoryCalls := 0
|
||||
state := newCompactionState(Options{
|
||||
ContextWindow: 1000,
|
||||
CompactionPreserveLast: 2,
|
||||
Summarizer: func(context.Context, string) (Provider, error) {
|
||||
factoryCalls++
|
||||
return cheap, nil
|
||||
},
|
||||
}, nil)
|
||||
|
||||
compacted := state.maybeCompact(context.Background(), main, compactableHistory(), nil)
|
||||
|
||||
if len(cheap.requests) == 0 {
|
||||
t.Fatal("dedicated summarizer never received the summarize call")
|
||||
}
|
||||
if len(main.requests) != 0 {
|
||||
t.Fatalf("main provider must see no summarize traffic, got %d requests", len(main.requests))
|
||||
}
|
||||
if factoryCalls != 1 {
|
||||
t.Fatalf("summarizer factory calls = %d, want 1 (lazy, memoized)", factoryCalls)
|
||||
}
|
||||
if estimateTokens(compacted) >= estimateTokens(compactableHistory()) {
|
||||
t.Fatal("compaction did not shrink the history")
|
||||
}
|
||||
found := false
|
||||
for _, message := range compacted {
|
||||
if strings.Contains(message.Content, "CHEAP SUMMARY") {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatal("cheap summarizer's summary missing from the compacted history")
|
||||
}
|
||||
}
|
||||
|
||||
// A failing summarizer falls back to the main provider for the SAME slice and
|
||||
// stays broken for the rest of the run — no repeated failed cheap calls.
|
||||
func TestSummarizerFailureFallsBackToMainAndSticks(t *testing.T) {
|
||||
broken := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventError, Error: "boom: model not found"},
|
||||
}}}
|
||||
main := &mockProvider{turns: [][]zeroruntime.StreamEvent{
|
||||
{{Type: zeroruntime.StreamEventText, Content: "MAIN SUMMARY"}, {Type: zeroruntime.StreamEventDone}},
|
||||
{{Type: zeroruntime.StreamEventText, Content: "MAIN SUMMARY 2"}, {Type: zeroruntime.StreamEventDone}},
|
||||
}}
|
||||
state := newCompactionState(Options{
|
||||
ContextWindow: 1000,
|
||||
CompactionPreserveLast: 2,
|
||||
Summarizer: func(context.Context, string) (Provider, error) { return broken, nil },
|
||||
}, nil)
|
||||
summarize := state.summarizeClosure(context.Background(), main)
|
||||
input := []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "history"}}
|
||||
|
||||
summary, err := summarize(input)
|
||||
if err != nil {
|
||||
t.Fatalf("fallback must succeed via the main provider: %v", err)
|
||||
}
|
||||
if summary != "MAIN SUMMARY" {
|
||||
t.Fatalf("summary = %q, want the main provider's", summary)
|
||||
}
|
||||
if len(broken.requests) != 1 {
|
||||
t.Fatalf("broken summarizer requests = %d, want exactly 1", len(broken.requests))
|
||||
}
|
||||
if !state.summarizerBroken {
|
||||
t.Fatal("a failed summarizer must be marked broken for the run")
|
||||
}
|
||||
|
||||
// Second summarization: straight to main, no cheap retry.
|
||||
if _, err := summarize(input); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(broken.requests) != 1 {
|
||||
t.Fatalf("broken summarizer was retried (%d requests)", len(broken.requests))
|
||||
}
|
||||
if len(main.requests) != 2 {
|
||||
t.Fatalf("main provider requests = %d, want 2", len(main.requests))
|
||||
}
|
||||
}
|
||||
|
||||
// A factory that errors (misconfigured model) is equivalent to no summarizer:
|
||||
// main provider used, factory not retried.
|
||||
func TestSummarizerFactoryErrorFallsBack(t *testing.T) {
|
||||
main := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventText, Content: "MAIN SUMMARY"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
factoryCalls := 0
|
||||
state := newCompactionState(Options{
|
||||
ContextWindow: 1000,
|
||||
CompactionPreserveLast: 2,
|
||||
Summarizer: func(context.Context, string) (Provider, error) {
|
||||
factoryCalls++
|
||||
return nil, errors.New("unknown model")
|
||||
},
|
||||
}, nil)
|
||||
summarize := state.summarizeClosure(context.Background(), main)
|
||||
|
||||
summary, err := summarize([]zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "history"}})
|
||||
if err != nil || summary != "MAIN SUMMARY" {
|
||||
t.Fatalf("summary = %q, err = %v; want main provider fallback", summary, err)
|
||||
}
|
||||
if factoryCalls != 1 {
|
||||
t.Fatalf("factory calls = %d, want 1 (sticky failure)", factoryCalls)
|
||||
}
|
||||
if _, dedicated := state.summarizeProvider(context.Background(), main); dedicated {
|
||||
t.Fatal("after a factory failure the run must stay on the main provider")
|
||||
}
|
||||
if factoryCalls != 1 {
|
||||
t.Fatalf("factory retried after failure (%d calls)", factoryCalls)
|
||||
}
|
||||
}
|
||||
|
||||
// Tool payloads are clamped in the summarizer transcript: results to
|
||||
// summaryToolResultClamp, call arguments to summaryToolArgsClamp.
|
||||
func TestRenderTranscriptClampsToolPayloads(t *testing.T) {
|
||||
bigResult := strings.Repeat("r", 50_000)
|
||||
bigArgs := strings.Repeat("w", 20_000)
|
||||
transcript := renderTranscript([]zeroruntime.Message{
|
||||
{Role: zeroruntime.MessageRoleAssistant, Content: "editing", ToolCalls: []zeroruntime.ToolCall{
|
||||
{ID: "c1", Name: "write_file", Arguments: bigArgs},
|
||||
}},
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "c1", Content: bigResult},
|
||||
{Role: zeroruntime.MessageRoleUser, Content: "keep going"},
|
||||
})
|
||||
|
||||
if len(transcript) > summaryToolResultClamp+summaryToolArgsClamp+1024 {
|
||||
t.Fatalf("transcript is %d bytes; clamps did not apply", len(transcript))
|
||||
}
|
||||
if !strings.Contains(transcript, "chars omitted for summarization") {
|
||||
t.Fatal("clamped transcript must carry the omission marker")
|
||||
}
|
||||
if !strings.Contains(transcript, "keep going") {
|
||||
t.Fatal("non-tool content must survive unclamped")
|
||||
}
|
||||
// Head AND tail of the clamped payloads survive.
|
||||
if !strings.HasPrefix(transcript[strings.Index(transcript, "tool result: ")+len("tool result: "):], "r") ||
|
||||
!strings.Contains(transcript, "rrr\n\nuser: keep going") {
|
||||
t.Fatalf("head/tail of the tool result lost:\n%s", transcript[:200])
|
||||
}
|
||||
}
|
||||
|
||||
// bigToolHistory returns a history with n large tool results (~40k estimated
|
||||
// tokens each): the trailing protection window keeps the newest verbatim,
|
||||
// the older ones are prunable.
|
||||
func bigToolHistory(n int) []zeroruntime.Message {
|
||||
bigBody := strings.Repeat("x", 160_000)
|
||||
messages := []zeroruntime.Message{{Role: zeroruntime.MessageRoleSystem, Content: "sys"}}
|
||||
for i := 0; i < n; i++ {
|
||||
id := string(rune('a' + i))
|
||||
messages = append(messages,
|
||||
zeroruntime.Message{Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{{ID: id, Name: "read_file", Arguments: "{\"path\":\"" + id + "\"}"}}},
|
||||
zeroruntime.Message{Role: zeroruntime.MessageRoleTool, ToolCallID: id, Content: bigBody},
|
||||
)
|
||||
}
|
||||
return append(messages, zeroruntime.Message{Role: zeroruntime.MessageRoleUser, Content: "go"})
|
||||
}
|
||||
|
||||
func summaryProvider(text string) *mockProvider {
|
||||
return &mockProvider{turns: [][]zeroruntime.StreamEvent{
|
||||
{{Type: zeroruntime.StreamEventText, Content: text}, {Type: zeroruntime.StreamEventDone}},
|
||||
{{Type: zeroruntime.StreamEventText, Content: text}, {Type: zeroruntime.StreamEventDone}},
|
||||
}}
|
||||
}
|
||||
|
||||
// A context-limit error is first answered with the free prune stage; when the
|
||||
// pruned history fits the window the retry is free (no summarizer call) and
|
||||
// the one-shot reactive budget is untouched. The paid summarizer only runs
|
||||
// when a later error finds nothing left to prune.
|
||||
func TestRecoverPrunesBeforeSummarizing(t *testing.T) {
|
||||
provider := summaryProvider("SUMMARY")
|
||||
// The protected trailing tool output (~40k tokens) must fit the window.
|
||||
state := newCompactionState(Options{ContextWindow: 200_000, CompactionPreserveLast: 2}, nil)
|
||||
messages := bigToolHistory(5)
|
||||
|
||||
pruned, retried, err := state.recover(context.Background(), provider, messages, nil, "maximum context length exceeded")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !retried {
|
||||
t.Fatal("prunable history that fits must produce a free retry")
|
||||
}
|
||||
if len(provider.requests) != 0 {
|
||||
t.Fatalf("free prune stage must not call the summarizer (%d requests)", len(provider.requests))
|
||||
}
|
||||
if state.reactiveAttempted {
|
||||
t.Fatal("a free prune retry must not consume the reactive budget")
|
||||
}
|
||||
if estimateTokens(pruned) >= estimateTokens(messages) {
|
||||
t.Fatal("prune did not shrink the history")
|
||||
}
|
||||
|
||||
// Second context-limit error on the pruned history: nothing meaningfully
|
||||
// prunable remains, so the paid summarizer fires and consumes the budget.
|
||||
_, retried, err = state.recover(context.Background(), provider, pruned, nil, "maximum context length exceeded")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !retried {
|
||||
t.Fatal("second recover must compact and retry")
|
||||
}
|
||||
if len(provider.requests) == 0 {
|
||||
t.Fatal("second recover must reach the summarizer")
|
||||
}
|
||||
if !state.reactiveAttempted {
|
||||
t.Fatal("a paid recover must consume the reactive budget")
|
||||
}
|
||||
}
|
||||
|
||||
// When pruning reclaims something but the pruned history still exceeds the
|
||||
// window, recover falls through to the paid summarizer in the SAME call
|
||||
// instead of spending a retry that would fail the same way.
|
||||
func TestRecoverFallsThroughToSummarizerWhenPruneIsInsufficient(t *testing.T) {
|
||||
provider := summaryProvider("SUMMARY")
|
||||
// A 1,000-token window: even the protected trailing output overflows it.
|
||||
state := newCompactionState(Options{ContextWindow: 1000, CompactionPreserveLast: 2}, nil)
|
||||
messages := bigToolHistory(5)
|
||||
|
||||
compacted, retried, err := state.recover(context.Background(), provider, messages, nil, "maximum context length exceeded")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !retried {
|
||||
t.Fatal("insufficient prune must still end in a retry via summarization")
|
||||
}
|
||||
if len(provider.requests) == 0 {
|
||||
t.Fatal("insufficient prune must reach the summarizer in the same recover call")
|
||||
}
|
||||
if !state.reactiveAttempted {
|
||||
t.Fatal("the paid summarizer consumes the reactive budget")
|
||||
}
|
||||
if estimateTokens(compacted) >= estimateTokens(messages) {
|
||||
t.Fatal("recover did not shrink the history")
|
||||
}
|
||||
}
|
||||
|
||||
// The free prune stage is not gated behind the one-shot budget: after a paid
|
||||
// reactive compaction, a later context-limit error still gets a free retry
|
||||
// when pruning makes the history fit.
|
||||
func TestRecoverFreePruneRunsAfterBudgetIsSpent(t *testing.T) {
|
||||
provider := summaryProvider("SUMMARY")
|
||||
state := newCompactionState(Options{ContextWindow: 200_000, CompactionPreserveLast: 2}, nil)
|
||||
state.reactiveAttempted = true
|
||||
messages := bigToolHistory(5)
|
||||
|
||||
pruned, retried, err := state.recover(context.Background(), provider, messages, nil, "maximum context length exceeded")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !retried || len(provider.requests) != 0 {
|
||||
t.Fatalf("free prune must retry without the summarizer even when the budget is spent (retried=%v, requests=%d)", retried, len(provider.requests))
|
||||
}
|
||||
if estimateTokens(pruned) >= estimateTokens(messages) {
|
||||
t.Fatal("prune did not shrink the history")
|
||||
}
|
||||
// With nothing left to prune and the budget spent, the caller keeps its
|
||||
// original error — recover neither loops nor summarizes.
|
||||
_, retried, err = state.recover(context.Background(), provider, pruned, nil, "maximum context length exceeded")
|
||||
if err != nil || retried || len(provider.requests) != 0 {
|
||||
t.Fatalf("spent budget must not summarize again (retried=%v, requests=%d, err=%v)", retried, len(provider.requests), err)
|
||||
}
|
||||
}
|
||||
|
||||
// The free-retry gate compares against the window the provider enforces, not
|
||||
// the 70% proactive threshold: a history pruned to ~80% of the window still
|
||||
// takes the free retry.
|
||||
func TestRecoverFreeRetryGateUsesWindowNotThreshold(t *testing.T) {
|
||||
provider := summaryProvider("SUMMARY")
|
||||
messages := bigToolHistory(5)
|
||||
pruned, _ := pruneStaleToolOutput(messages, 2)
|
||||
prunedTokens := estimateTokens(pruned)
|
||||
// Choose a window so the pruned size sits between 70% and 95% of it.
|
||||
window := prunedTokens * 100 / 80
|
||||
state := newCompactionState(Options{ContextWindow: window, CompactionPreserveLast: 2}, nil)
|
||||
if prunedTokens <= state.threshold {
|
||||
t.Fatalf("test setup: pruned size %d must exceed the proactive threshold %d", prunedTokens, state.threshold)
|
||||
}
|
||||
|
||||
_, retried, err := state.recover(context.Background(), provider, messages, nil, "maximum context length exceeded")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !retried || len(provider.requests) != 0 {
|
||||
t.Fatalf("a history that fits the window must take the free retry (retried=%v, requests=%d)", retried, len(provider.requests))
|
||||
}
|
||||
}
|
||||
|
||||
// After a mid-run model switch the summarizer is re-resolved against the new
|
||||
// main model: a run that started on the cheap model (no dedicated summarizer)
|
||||
// and escalated gets a cheap summarizer, and a previously broken one is
|
||||
// retried for the new model.
|
||||
func TestSwitchModelReresolvesSummarizer(t *testing.T) {
|
||||
cheap := summaryProvider("CHEAP SUMMARY")
|
||||
main := summaryProvider("MAIN SUMMARY")
|
||||
var factoryModels []string
|
||||
state := newCompactionState(Options{
|
||||
ContextWindow: 1000,
|
||||
CompactionPreserveLast: 2,
|
||||
Model: "cheap-model",
|
||||
Summarizer: func(_ context.Context, mainModelID string) (Provider, error) {
|
||||
factoryModels = append(factoryModels, mainModelID)
|
||||
if mainModelID == "cheap-model" {
|
||||
return nil, nil // main already is the cheap model
|
||||
}
|
||||
return cheap, nil
|
||||
},
|
||||
}, nil)
|
||||
|
||||
if _, dedicated := state.summarizeProvider(context.Background(), main); dedicated {
|
||||
t.Fatal("on the cheap model the main provider must summarize")
|
||||
}
|
||||
if _, dedicated := state.summarizeProvider(context.Background(), main); dedicated || len(factoryModels) != 1 {
|
||||
t.Fatalf("factory must run once per main model, calls=%v", factoryModels)
|
||||
}
|
||||
|
||||
state.switchModel("big-model", 400_000)
|
||||
if state.threshold != compactionThreshold(400_000) || state.window != 400_000 {
|
||||
t.Fatalf("switchModel must re-derive the window, got threshold=%d window=%d", state.threshold, state.window)
|
||||
}
|
||||
if got, dedicated := state.summarizeProvider(context.Background(), main); !dedicated || got != Provider(cheap) {
|
||||
t.Fatal("after escalating to an expensive model the dedicated cheap summarizer must be used")
|
||||
}
|
||||
if len(factoryModels) != 2 || factoryModels[1] != "big-model" {
|
||||
t.Fatalf("factory must be re-invoked with the new main model, calls=%v", factoryModels)
|
||||
}
|
||||
|
||||
// A broken summarizer is retried after the next switch.
|
||||
state.summarizerBroken = true
|
||||
if _, dedicated := state.summarizeProvider(context.Background(), main); dedicated {
|
||||
t.Fatal("a broken summarizer must fall back to main")
|
||||
}
|
||||
state.switchModel("other-model", 0)
|
||||
if got, dedicated := state.summarizeProvider(context.Background(), main); !dedicated || got != Provider(cheap) {
|
||||
t.Fatal("a model switch must clear the broken flag and rebuild the summarizer")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactionUpdateWindow(t *testing.T) {
|
||||
state := newCompactionState(Options{ContextWindow: 100_000, CompactionPreserveLast: 2}, nil)
|
||||
state.lowWaterMark = 50_000
|
||||
before := state.threshold
|
||||
|
||||
state.updateWindow(0) // unknown model: unchanged
|
||||
if state.threshold != before || state.lowWaterMark != 50_000 {
|
||||
t.Fatal("updateWindow(0) must be a no-op")
|
||||
}
|
||||
|
||||
state.updateWindow(10_000)
|
||||
if state.threshold >= before {
|
||||
t.Fatalf("threshold = %d, want re-derived below %d", state.threshold, before)
|
||||
}
|
||||
if state.lowWaterMark != 0 {
|
||||
t.Fatal("low-water mark must reset on a window change")
|
||||
}
|
||||
if !state.enabled {
|
||||
t.Fatal("a positive window must keep compaction enabled")
|
||||
}
|
||||
}
|
||||
|
||||
// End-to-end: after escalate_model switches providers, the loop re-derives the
|
||||
// context window via ContextWindowFor, visible in the OnContext report.
|
||||
func TestRunEscalationRederivesContextWindow(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(escalatingTool{target: "big-model"})
|
||||
switched := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventText, Content: "done"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call-1", ToolName: "escalate"},
|
||||
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call-1", ArgumentsFragment: `{}`},
|
||||
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call-1"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
}}}
|
||||
|
||||
var windows []int
|
||||
result, err := Run(context.Background(), "go", provider, Options{
|
||||
Registry: registry,
|
||||
ContextWindow: 100_000,
|
||||
Model: "small-model",
|
||||
ModelSwitcher: func(context.Context, string) (Provider, error) { return switched, nil },
|
||||
ContextWindowFor: func(modelID string) int {
|
||||
if modelID == "big-model" {
|
||||
return 400_000
|
||||
}
|
||||
return 0
|
||||
},
|
||||
OnContext: func(breakdown ContextBreakdown) { windows = append(windows, breakdown.ContextWindow) },
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.FinalAnswer != "done" {
|
||||
t.Fatalf("final answer = %q", result.FinalAnswer)
|
||||
}
|
||||
if len(windows) < 2 {
|
||||
t.Fatalf("expected at least 2 OnContext reports, got %d", len(windows))
|
||||
}
|
||||
if windows[0] != 100_000 {
|
||||
t.Fatalf("turn 1 window = %d, want the original 100000", windows[0])
|
||||
}
|
||||
if windows[len(windows)-1] != 400_000 {
|
||||
t.Fatalf("post-escalation window = %d, want 400000", windows[len(windows)-1])
|
||||
}
|
||||
}
|
||||
|
||||
// nonComparableProvider has a slice field, so comparing two Provider interface
|
||||
// values holding it would panic; the summarizer path must never compare them.
|
||||
type nonComparableProvider struct {
|
||||
inner *mockProvider
|
||||
tags []string
|
||||
}
|
||||
|
||||
func (p nonComparableProvider) StreamCompletion(ctx context.Context, request zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) {
|
||||
return p.inner.StreamCompletion(ctx, request)
|
||||
}
|
||||
|
||||
func TestSummarizeClosureNeverComparesProviderValues(t *testing.T) {
|
||||
broken := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
|
||||
{Type: zeroruntime.StreamEventError, Error: "boom: model not found"},
|
||||
}}}
|
||||
main := nonComparableProvider{inner: summaryProvider("MAIN SUMMARY"), tags: []string{"x"}}
|
||||
state := newCompactionState(Options{
|
||||
ContextWindow: 1000,
|
||||
CompactionPreserveLast: 2,
|
||||
Summarizer: func(context.Context, string) (Provider, error) { return nonComparableProvider{inner: broken}, nil },
|
||||
}, nil)
|
||||
summary, err := state.summarizeClosure(context.Background(), main)([]zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "history"}})
|
||||
if err != nil || summary != "MAIN SUMMARY" {
|
||||
t.Fatalf("fallback with non-comparable providers must not panic and must use main, got %q err=%v", summary, err)
|
||||
}
|
||||
// And with no dedicated summarizer at all the closure still works.
|
||||
state = newCompactionState(Options{ContextWindow: 1000, CompactionPreserveLast: 2}, nil)
|
||||
if _, err := state.summarizeClosure(context.Background(), main)([]zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "history"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
type contextPlannerConfig struct {
|
||||
contextWindow int
|
||||
promptCacheKey string
|
||||
serviceTier string
|
||||
promptParts systemPromptParts
|
||||
}
|
||||
|
||||
@@ -62,6 +63,7 @@ func (planner *contextPlanner) plan(messages []zeroruntime.Message, toolDefs []z
|
||||
Messages: requestMessages,
|
||||
Tools: planner.snapshotTools(toolDefs, fingerprint),
|
||||
ReasoningEffort: reasoningEffort,
|
||||
ServiceTier: planner.config.serviceTier,
|
||||
PromptCacheKey: planner.config.promptCacheKey,
|
||||
}
|
||||
breakdown := MeasureContext(request.Messages, request.Tools, planner.config.contextWindow)
|
||||
|
||||
@@ -25,6 +25,7 @@ func TestContextPlannerPreservesProviderRequest(t *testing.T) {
|
||||
planner := newContextPlanner(contextPlannerConfig{
|
||||
contextWindow: 128_000,
|
||||
promptCacheKey: "session-1",
|
||||
serviceTier: "priority",
|
||||
promptParts: systemPromptParts{prompt: "system", baseInstructions: "system"},
|
||||
})
|
||||
|
||||
@@ -33,6 +34,7 @@ func TestContextPlannerPreservesProviderRequest(t *testing.T) {
|
||||
Messages: copyMessages(messages),
|
||||
Tools: toolDefs,
|
||||
ReasoningEffort: "medium",
|
||||
ServiceTier: "priority",
|
||||
PromptCacheKey: "session-1",
|
||||
}
|
||||
if !reflect.DeepEqual(plan.Request, want) {
|
||||
|
||||
@@ -144,9 +144,10 @@ func TestRunSwitchesProviderOnEscalationRequest(t *testing.T) {
|
||||
var switchedTo string
|
||||
switchCount := 0
|
||||
result, err := Run(context.Background(), "go", firstProvider, Options{
|
||||
Registry: registry,
|
||||
Model: "claude-sonnet-4.5",
|
||||
MaxTurns: 4,
|
||||
Registry: registry,
|
||||
Model: "claude-sonnet-4.5",
|
||||
ServiceTier: "priority",
|
||||
MaxTurns: 4,
|
||||
ModelSwitcher: func(_ context.Context, modelID string) (Provider, error) {
|
||||
switchCount++
|
||||
switchedTo = modelID
|
||||
@@ -167,9 +168,15 @@ func TestRunSwitchesProviderOnEscalationRequest(t *testing.T) {
|
||||
if len(firstProvider.requests) != 1 {
|
||||
t.Fatalf("expected first provider to handle exactly the escalation turn, got %d requests", len(firstProvider.requests))
|
||||
}
|
||||
if got := firstProvider.requests[0].ServiceTier; got != "priority" {
|
||||
t.Fatalf("initial request service tier = %q, want priority", got)
|
||||
}
|
||||
if len(secondProvider.requests) != 1 {
|
||||
t.Fatalf("expected second provider to handle the post-switch turn, got %d requests", len(secondProvider.requests))
|
||||
}
|
||||
if got := secondProvider.requests[0].ServiceTier; got != "" {
|
||||
t.Fatalf("post-switch request retained service tier %q", got)
|
||||
}
|
||||
if result.FinalAnswer != "answered on the upgraded model" {
|
||||
t.Fatalf("expected final answer from the swapped provider, got %q", result.FinalAnswer)
|
||||
}
|
||||
@@ -211,9 +218,10 @@ func TestRunEscalationSwitcherErrorIsNonFatal(t *testing.T) {
|
||||
provider := &mockProvider{turns: escalateThenAnswerTurns("recovered")}
|
||||
|
||||
result, err := Run(context.Background(), "go", provider, Options{
|
||||
Registry: registry,
|
||||
Model: "claude-sonnet-4.5",
|
||||
MaxTurns: 4,
|
||||
Registry: registry,
|
||||
Model: "claude-sonnet-4.5",
|
||||
ServiceTier: "priority",
|
||||
MaxTurns: 4,
|
||||
ModelSwitcher: func(_ context.Context, _ string) (Provider, error) {
|
||||
return nil, errors.New("provider build blew up")
|
||||
},
|
||||
@@ -225,6 +233,9 @@ func TestRunEscalationSwitcherErrorIsNonFatal(t *testing.T) {
|
||||
if len(provider.requests) != 2 {
|
||||
t.Fatalf("expected the run to continue on the original provider, got %d requests", len(provider.requests))
|
||||
}
|
||||
if got := provider.requests[1].ServiceTier; got != "priority" {
|
||||
t.Fatalf("failed-switch request service tier = %q, want priority", got)
|
||||
}
|
||||
if result.FinalAnswer != "recovered" {
|
||||
t.Fatalf("expected final answer after non-fatal switch error, got %q", result.FinalAnswer)
|
||||
}
|
||||
|
||||
+40
-12
@@ -40,7 +40,15 @@ const (
|
||||
toolResultControlSpecReview = "spec_review_required"
|
||||
)
|
||||
|
||||
var errPermissionApprovalCanceled = errors.New("permission approval cancelled")
|
||||
// ErrPermissionApprovalCanceled ends a run because the caller cancelled a
|
||||
// permission prompt, rather than because anything failed.
|
||||
//
|
||||
// EXPORTED FOR THE SURFACES. A cancel is a normal ending and each surface says
|
||||
// so its own way — the TUI unwinds, and ACP has a "cancelled" stop reason for
|
||||
// exactly this. While it was unexported ACP could not distinguish it from a
|
||||
// fault and reported it as an internal error, so declining a tool looked like
|
||||
// the agent crashing.
|
||||
var ErrPermissionApprovalCanceled = errors.New("permission approval cancelled")
|
||||
|
||||
// isImageRejectionError reports whether err is a provider 400 that rejects
|
||||
// image/multimodal content. This is checked BEFORE the compaction-retry path
|
||||
@@ -186,6 +194,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
|
||||
planner := newContextPlanner(contextPlannerConfig{
|
||||
contextWindow: options.ContextWindow,
|
||||
promptCacheKey: options.SessionID,
|
||||
serviceTier: options.ServiceTier,
|
||||
promptParts: promptParts,
|
||||
})
|
||||
messages := zeroruntime.SeedMessagesWithImages(promptParts.prompt, prompt, options.Images)
|
||||
@@ -829,12 +838,27 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
|
||||
_ = session.Prewarm(ctx)
|
||||
provider = sessionProvider{session: session}
|
||||
options.Model = turnRequestedModel
|
||||
// Service tiers are model capabilities. A successful escalation may
|
||||
// target a model that does not support the tier selected before the
|
||||
// switch, so let the destination use its default.
|
||||
options.ServiceTier = ""
|
||||
planner.config.serviceTier = ""
|
||||
options.Trace.Counter(trace.CounterModelSwitches, 1)
|
||||
// KNOWN LIMITATION (deferred): the compactor's context-window budget
|
||||
// is fixed at run start from options.ContextWindow and is NOT updated
|
||||
// here, so a switch to a model with a different window keeps compacting
|
||||
// against the original budget. Fixing it needs the switcher to also
|
||||
// report the new window — out of scope for this change.
|
||||
// Re-derive the context window for the new model so compaction and
|
||||
// the OnContext budget report track the model actually in force —
|
||||
// without this a switch to a smaller-window model can overflow the
|
||||
// target, and a switch to a larger one over-compacts. An unknown
|
||||
// model (<= 0) keeps the original window. The compactor also
|
||||
// re-resolves its dedicated summarizer against the new model.
|
||||
window := 0
|
||||
if options.ContextWindowFor != nil {
|
||||
window = options.ContextWindowFor(turnRequestedModel)
|
||||
}
|
||||
if window > 0 {
|
||||
options.ContextWindow = window
|
||||
planner.config.contextWindow = window
|
||||
}
|
||||
compactor.switchModel(turnRequestedModel, window)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1358,7 +1382,7 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
|
||||
}
|
||||
case PermissionDecisionCancel:
|
||||
emitCanceledPermission(options, call, requestEvent, decisionReason)
|
||||
return canceledPermissionResult(call, decisionReason, requestEvent), fmt.Errorf("%w for %s", errPermissionApprovalCanceled, call.Name)
|
||||
return canceledPermissionResult(call, decisionReason, requestEvent), fmt.Errorf("%w for %s", ErrPermissionApprovalCanceled, call.Name)
|
||||
default:
|
||||
emitDeniedPermission(options, call, requestEvent, decisionReason)
|
||||
return deniedPermissionResult(call, decisionReason, requestEvent), nil
|
||||
@@ -1542,7 +1566,7 @@ func maybeRetryUnsandboxedAfterSandboxRestriction(ctx context.Context, registry
|
||||
case PermissionDecisionCancel:
|
||||
emitCanceledPermission(options, call, requestEvent, reason)
|
||||
canceled := canceledPermissionResult(call, reason, requestEvent)
|
||||
return result, &canceled, true, decision.Action, reason, nil, fmt.Errorf("%w for %s", errPermissionApprovalCanceled, call.Name)
|
||||
return result, &canceled, true, decision.Action, reason, nil, fmt.Errorf("%w for %s", ErrPermissionApprovalCanceled, call.Name)
|
||||
default:
|
||||
emitDeniedPermission(options, call, requestEvent, reason)
|
||||
denied := deniedPermissionResult(call, reason, requestEvent)
|
||||
@@ -1590,7 +1614,7 @@ func maybeRetryWithNetworkAfterSandboxDenial(ctx context.Context, registry *tool
|
||||
case PermissionDecisionCancel:
|
||||
emitCanceledPermission(options, call, requestEvent, reason)
|
||||
canceled := canceledPermissionResult(call, reason, requestEvent)
|
||||
return result, &canceled, true, decision.Action, reason, fmt.Errorf("%w for %s", errPermissionApprovalCanceled, call.Name)
|
||||
return result, &canceled, true, decision.Action, reason, fmt.Errorf("%w for %s", ErrPermissionApprovalCanceled, call.Name)
|
||||
default:
|
||||
emitDeniedPermission(options, call, requestEvent, reason)
|
||||
denied := deniedPermissionResult(call, reason, requestEvent)
|
||||
@@ -3081,13 +3105,15 @@ func partitionToolsCached(registry *tools.Registry, permissionMode PermissionMod
|
||||
// Deferral may activate only when tool_search is actually runnable; otherwise
|
||||
// the loop would hide deferred tools behind a loader the dispatch gate rejects
|
||||
// — an inescapable dead-end. "Runnable" mirrors executeToolCall's gate:
|
||||
// registered, not in DisabledTools, and advertised in the current permission
|
||||
// mode (e.g. not spec-draft, where tool_search is not advertised). The
|
||||
// registered, model-visible, not in DisabledTools, and advertised in the
|
||||
// current permission mode (e.g. not spec-draft, where tool_search is not
|
||||
// advertised). The
|
||||
// EnabledTools allowlist is intentionally NOT checked here — tool_search is
|
||||
// exempt from the allowlist at dispatch, so an allowlist that omits it must
|
||||
// not disable deferral.
|
||||
loader, loaderFound := registry.Get(tools.ToolSearchToolName)
|
||||
loaderUsable := loaderFound &&
|
||||
tools.ModelVisible(loader) &&
|
||||
!containsToolName(options.DisabledTools, tools.ToolSearchToolName) &&
|
||||
ToolAdvertised(loader, permissionMode)
|
||||
|
||||
@@ -3202,7 +3228,9 @@ func runtimeToolDefinition(tool tools.Tool) zeroruntime.ToolDefinition {
|
||||
}
|
||||
|
||||
func ToolVisible(tool tools.Tool, permissionMode PermissionMode, enabledTools []string, disabledTools []string) bool {
|
||||
return ToolAllowedByFilters(tool.Name(), enabledTools, disabledTools) && ToolAdvertised(tool, permissionMode)
|
||||
return tools.ModelVisible(tool) &&
|
||||
ToolAllowedByFilters(tool.Name(), enabledTools, disabledTools) &&
|
||||
ToolAdvertised(tool, permissionMode)
|
||||
}
|
||||
|
||||
func ToolAllowedByFilters(name string, enabledTools []string, disabledTools []string) bool {
|
||||
|
||||
@@ -1853,7 +1853,7 @@ func TestRunAbortsWhenPermissionRequestCanceled(t *testing.T) {
|
||||
},
|
||||
})
|
||||
|
||||
if !errors.Is(err, errPermissionApprovalCanceled) {
|
||||
if !errors.Is(err, ErrPermissionApprovalCanceled) {
|
||||
t.Fatalf("expected permission approval cancel error, got %v", err)
|
||||
}
|
||||
if result.FinalAnswer != "" {
|
||||
@@ -4340,3 +4340,95 @@ func TestPlanModeHonorsBeforeToolVeto(t *testing.T) {
|
||||
t.Fatalf("expected tool result to mention the beforeTool veto, got %q", combined)
|
||||
}
|
||||
}
|
||||
|
||||
// A cancelled permission on either sandbox-retry path must abort the run with a
|
||||
// recognisable cancellation, and must not run the escalated retry.
|
||||
//
|
||||
// BOTH PATHS, because both wrap ErrPermissionApprovalCanceled and each is
|
||||
// reached by a different kind of denial: one by an unrestricted platform denial
|
||||
// asking for approval, the other by an external-network denial. Surfaces map
|
||||
// this sentinel onto their own idea of a cancel — ACP returns stopReason
|
||||
// "cancelled" for it — so a path that produced a bare error instead would show
|
||||
// the user a crash where they had simply declined.
|
||||
func TestCancellingASandboxRetryAbortsWithoutRetrying(t *testing.T) {
|
||||
denial := func(capability execution.CapabilityKind, scope string) tools.Result {
|
||||
return tools.Result{
|
||||
Status: tools.StatusError,
|
||||
ExecutionOutcome: &execution.Outcome{
|
||||
State: execution.StateDenied,
|
||||
Kind: execution.OutcomeEnforcementDenied,
|
||||
Denial: &execution.Denial{
|
||||
Capability: execution.Capability{Kind: capability, Scope: scope},
|
||||
Source: execution.DenialSourcePlatformSandbox,
|
||||
Reason: "denied by the sandbox",
|
||||
Recoverable: true,
|
||||
NextAction: execution.DenialNextActionRequestApproval,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
for _, testCase := range []struct {
|
||||
name string
|
||||
result tools.Result
|
||||
run func(context.Context, *tools.Registry, ToolCall, tools.Tool, map[string]any, tools.Result, Options) error
|
||||
}{
|
||||
{
|
||||
name: "unsandboxed retry",
|
||||
result: denial(execution.CapabilityUnrestricted, "host"),
|
||||
run: func(ctx context.Context, registry *tools.Registry, call ToolCall, tool tools.Tool, args map[string]any, result tools.Result, options Options) error {
|
||||
_, _, _, _, _, _, abortErr := maybeRetryUnsandboxedAfterSandboxRestriction(
|
||||
ctx, registry, call, tool, args, result, PermissionModeAsk, options, nil)
|
||||
return abortErr
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "network retry",
|
||||
result: denial(execution.CapabilityExternalNetwork, ""),
|
||||
run: func(ctx context.Context, registry *tools.Registry, call ToolCall, tool tools.Tool, args map[string]any, result tools.Result, options Options) error {
|
||||
_, _, _, _, _, abortErr := maybeRetryWithNetworkAfterSandboxDenial(
|
||||
ctx, registry, call, tool, args, result, PermissionModeAsk, options, nil)
|
||||
return abortErr
|
||||
},
|
||||
},
|
||||
} {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
retryTool := &sandboxNamespaceLimitedRetryTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(retryTool)
|
||||
|
||||
asked := 0
|
||||
options := Options{
|
||||
Registry: registry,
|
||||
PermissionMode: PermissionModeAsk,
|
||||
Autonomy: "medium",
|
||||
Sandbox: sandbox.NewEngine(sandbox.EngineOptions{
|
||||
WorkspaceRoot: t.TempDir(),
|
||||
Policy: sandbox.DefaultPolicy(),
|
||||
}),
|
||||
OnPermissionRequest: func(context.Context, PermissionRequest) (PermissionDecision, error) {
|
||||
asked++
|
||||
return PermissionDecision{Action: PermissionDecisionCancel, Reason: "client cancelled"}, nil
|
||||
},
|
||||
}
|
||||
|
||||
call := ToolCall{ID: "call-1", Name: "bash", Arguments: `{"command":"curl example.com"}`}
|
||||
args := map[string]any{"command": "curl example.com"}
|
||||
|
||||
abortErr := testCase.run(context.Background(), registry, call, retryTool, args, testCase.result, options)
|
||||
|
||||
if asked != 1 {
|
||||
t.Fatalf("permission was requested %d times, want exactly 1", asked)
|
||||
}
|
||||
if !errors.Is(abortErr, ErrPermissionApprovalCanceled) {
|
||||
t.Fatalf("abort error = %v, want it to match ErrPermissionApprovalCanceled", abortErr)
|
||||
}
|
||||
// The escalated retry is the thing the permission was gating. A
|
||||
// cancelled prompt that still ran it would have run the command the
|
||||
// user just declined.
|
||||
if len(retryTool.calls) != 0 {
|
||||
t.Fatalf("the escalated retry ran %d times after a cancel: %+v", len(retryTool.calls), retryTool.calls)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -67,3 +67,21 @@ func TestEagerToolSchemaTokenBudget(t *testing.T) {
|
||||
t.Fatalf("eager tool schemas are %d tokens, over the %d ceiling — trim a schema, drop a core tool, or raise the ceiling deliberately (deferring will NOT help: this test disables deferral)", got, maxEagerToolSchemaTokens)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentAdvertisesBothEditTools(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
for _, tool := range tools.CoreToolsScoped(t.TempDir(), nil) {
|
||||
registry.Register(tool)
|
||||
}
|
||||
exposed, _ := partitionTools(registry, PermissionModeAsk, Options{}, map[string]bool{})
|
||||
names := make(map[string]bool, len(exposed))
|
||||
for _, definition := range exposed {
|
||||
names[definition.Name] = true
|
||||
}
|
||||
if !names["apply_patch"] {
|
||||
t.Fatal("agent must retain apply_patch for multi-hunk changes")
|
||||
}
|
||||
if !names["edit_file"] {
|
||||
t.Fatal("agent must receive edit_file for targeted exact replacements")
|
||||
}
|
||||
}
|
||||
@@ -52,7 +52,7 @@ work.
|
||||
|
||||
- Choose the narrowest tool that safely accomplishes the step. Prefer native
|
||||
file tools - read_file, read_minified_file, list_directory, glob, grep,
|
||||
write_file, edit_file, apply_patch - over shelling out to
|
||||
edit_file, apply_patch, write_file - over shelling out to
|
||||
cat/sed/awk/python for file operations.
|
||||
They are safer, reviewable, and produce clean diffs.
|
||||
- Prefer read_minified_file when initially exploring source code; it preserves
|
||||
@@ -61,10 +61,18 @@ work.
|
||||
- Keep edits focused and reviewable. A single patch may update several related
|
||||
files when they form one coherent change; do not hide unrelated edits in a
|
||||
bulk shell or script rewrite.
|
||||
- For edits to existing files, prefer edit_file or apply_patch with minimal,
|
||||
targeted diffs. Match the existing indentation, imports, and idioms. Match the
|
||||
file's comment density: do not add explanatory comments unless the user asks or
|
||||
the code is already comment-dense.
|
||||
- For edits to existing files, use edit_file for a targeted change (old_string
|
||||
must match the file exactly and be unique, so include a few surrounding
|
||||
lines) and apply_patch when one coherent change spans several hunks or
|
||||
files. Use write_file only to create a file or when most of it changes; do
|
||||
not rewrite a whole file to change a few lines. Match the existing
|
||||
indentation, imports, and idioms. Match the file's comment density: do not
|
||||
add explanatory comments unless the user asks or the code is already
|
||||
comment-dense.
|
||||
- A successful edit result already confirms the change; do not re-read a file
|
||||
just to verify an edit that succeeded. If an edit fails, read the error, fix
|
||||
the old_string or hunk, and retry the same tool rather than switching to a
|
||||
full rewrite.
|
||||
- Solve the problem as posed, not a more general version of it. Add no
|
||||
speculative abstraction, configurability, or handling for cases that cannot
|
||||
occur, and nothing the user did not ask for. A small diff can still be
|
||||
@@ -104,7 +112,7 @@ work.
|
||||
you need to clean up a running foreground command yourself, use write_stdin.
|
||||
- write_stdin's session_id is only ever an id returned by a still-running
|
||||
exec_command; never guess or probe ids. If you have no such session, start one
|
||||
with exec_command, or use write_file/edit_file/apply_patch for file changes.
|
||||
with exec_command, or use edit_file/apply_patch/write_file for file changes.
|
||||
- write_stdin with empty input polls an existing exec_command session, and
|
||||
`\u0003` interrupts it. Sending other stdin bytes may require approval because
|
||||
it can drive the running process beyond the original command. Non-tty sessions
|
||||
|
||||
@@ -60,8 +60,12 @@ const openAIPromptAddendum = `<model_guidance>
|
||||
longer answers, fenced code blocks for code, and ` + "`inline code`" + ` for paths,
|
||||
commands, and symbols.
|
||||
- Strongly prefer the native file tools (read_file, list_directory, grep, glob,
|
||||
write_file, edit_file, apply_patch) over shelling out to cat/sed/awk/python for
|
||||
file work. Make one tool call per file; do not batch file writes into a script.
|
||||
edit_file, apply_patch, write_file) over shelling out to cat/sed/awk/python
|
||||
for file work. read_file, edit_file and write_file take one file per call;
|
||||
a single apply_patch may span several files when they form one coherent
|
||||
change. Do not batch file writes into a script. Independent tool calls
|
||||
(several reads, several edit_file calls to different files) belong in the
|
||||
same turn.
|
||||
- Persist until the task is fully handled this turn: gather context, implement,
|
||||
run the validators, and report — do not stop at a partial result.
|
||||
</model_guidance>`
|
||||
|
||||
@@ -36,7 +36,7 @@ func TestCoreSystemPromptIncludesCodingQualityRules(t *testing.T) {
|
||||
"inspect the target file",
|
||||
"plan then act",
|
||||
"choose the narrowest tool",
|
||||
"prefer edit_file or apply_patch",
|
||||
"for edits to existing files, use edit_file",
|
||||
"verify after edits",
|
||||
"honor the active permission mode",
|
||||
"avoid broad refactors",
|
||||
|
||||
+21
-4
@@ -325,6 +325,7 @@ type Options struct {
|
||||
ProviderName string
|
||||
Model string
|
||||
ReasoningEffort string
|
||||
ServiceTier string
|
||||
Cwd string
|
||||
SystemPrompt string
|
||||
// TransientSystemPrompt adds trusted runtime guidance for this run only.
|
||||
@@ -347,10 +348,26 @@ type Options struct {
|
||||
// CompactionPreserveLast is how many trailing messages compaction keeps
|
||||
// verbatim. <= 0 falls back to defaultCompactionPreserveLast.
|
||||
CompactionPreserveLast int
|
||||
Registry *tools.Registry
|
||||
PermissionMode PermissionMode
|
||||
Autonomy string
|
||||
Sandbox *sandbox.Engine
|
||||
// Summarizer, when set, lazily builds the provider used for compaction
|
||||
// summarization calls — typically a cheap/fast model, since summaries at
|
||||
// main-model prices are the single most expensive recurring event in a
|
||||
// long run. It receives the main model currently in force and may return
|
||||
// (nil, nil) when no dedicated summarizer applies to it (the main model is
|
||||
// already the cheap one). Built on the first paid compaction and again
|
||||
// after a mid-run model switch. Any failure (build or call) falls back to
|
||||
// the run's main provider until the next switch, so a misconfigured
|
||||
// summarizer can never break compaction. nil keeps today's behavior.
|
||||
Summarizer func(ctx context.Context, mainModelID string) (Provider, error)
|
||||
// ContextWindowFor, when set, resolves a model ID to its context window so
|
||||
// the compactor can re-derive its threshold after a mid-run model switch
|
||||
// (escalate_model). Without it a switch keeps compacting against the
|
||||
// original model's window — overflowing a smaller target or over-compacting
|
||||
// a larger one. Return <= 0 when the model is unknown (window unchanged).
|
||||
ContextWindowFor func(modelID string) int
|
||||
Registry *tools.Registry
|
||||
PermissionMode PermissionMode
|
||||
Autonomy string
|
||||
Sandbox *sandbox.Engine
|
||||
// FileTracker records per-session file read/write versions so the write tools
|
||||
// can detect a file changed on disk outside Zero since it was last read. nil
|
||||
// disables the check. Created once per session and threaded into every tool run.
|
||||
|
||||
@@ -29,10 +29,6 @@ func ParseTraceEventKeys(stdout string) []string {
|
||||
return keys
|
||||
}
|
||||
|
||||
func TraceEventKeys(stdout string) []string {
|
||||
return ParseTraceEventKeys(stdout)
|
||||
}
|
||||
|
||||
func MissingTraceEvents(required []string, stdout string) []string {
|
||||
seen := map[string]bool{}
|
||||
for _, key := range ParseTraceEventKeys(stdout) {
|
||||
|
||||
@@ -453,7 +453,12 @@ func (manager *Manager) normalizeLoadedTask(fileTaskID string, task Task) (Task,
|
||||
}
|
||||
outputFile, err := manager.outputFile(task.ID, task.OutputFile)
|
||||
if err != nil {
|
||||
return Task{}, false, err
|
||||
var migrated bool
|
||||
outputFile, migrated = manager.migrateLegacyOutputFile(task.ID, task.OutputFile)
|
||||
if !migrated {
|
||||
return Task{}, false, err
|
||||
}
|
||||
changed = true
|
||||
}
|
||||
if outputFile != task.OutputFile {
|
||||
changed = true
|
||||
@@ -562,6 +567,32 @@ func (manager *Manager) outputFile(taskID string, requested string) (string, err
|
||||
return path, nil
|
||||
}
|
||||
|
||||
// migrateLegacyOutputFile recovers task metadata written under a previous home
|
||||
// or XDG data root. It accepts only the canonical per-task output
|
||||
// name in the historic zero/background layout, and only when the corresponding
|
||||
// regular file already exists inside the active manager root.
|
||||
func (manager *Manager) migrateLegacyOutputFile(taskID string, requested string) (string, bool) {
|
||||
requested = strings.TrimSpace(requested)
|
||||
if !filepath.IsAbs(requested) {
|
||||
return "", false
|
||||
}
|
||||
requested = filepath.Clean(requested)
|
||||
expectedName := taskID + ".ndjson"
|
||||
if filepath.Base(requested) != expectedName {
|
||||
return "", false
|
||||
}
|
||||
legacyRoot := filepath.Dir(requested)
|
||||
if filepath.Base(legacyRoot) != "background" || filepath.Base(filepath.Dir(legacyRoot)) != "zero" {
|
||||
return "", false
|
||||
}
|
||||
canonical := filepath.Join(manager.rootDir, expectedName)
|
||||
info, err := os.Lstat(canonical)
|
||||
if err != nil || !info.Mode().IsRegular() {
|
||||
return "", false
|
||||
}
|
||||
return canonical, true
|
||||
}
|
||||
|
||||
func validTaskID(taskID string) bool {
|
||||
return taskIDPattern.MatchString(strings.TrimSpace(taskID))
|
||||
}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package background
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
@@ -191,6 +192,88 @@ func TestManagerSkipsInvalidPersistedMetadata(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerMigratesLegacyOutputPathWhenCanonicalOutputExists(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
const taskID = "legacy_task"
|
||||
canonicalOutput := filepath.Join(root, taskID+".ndjson")
|
||||
if err := os.WriteFile(canonicalOutput, []byte("legacy output\n"), 0o600); err != nil {
|
||||
t.Fatalf("write canonical output: %v", err)
|
||||
}
|
||||
legacyOutput := filepath.Join(t.TempDir(), "zero", "background", taskID+".ndjson")
|
||||
metadata, err := json.Marshal(Task{
|
||||
ID: taskID,
|
||||
Type: "specialist",
|
||||
Status: StatusCompleted,
|
||||
OutputFile: legacyOutput,
|
||||
StartedAt: time.Date(2026, 8, 1, 10, 0, 0, 0, time.UTC),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("marshal legacy metadata: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(root, taskID+".json"), append(metadata, '\n'), 0o600); err != nil {
|
||||
t.Fatalf("write legacy metadata: %v", err)
|
||||
}
|
||||
|
||||
manager, err := NewManager(root)
|
||||
if err != nil {
|
||||
t.Fatalf("NewManager returned error: %v", err)
|
||||
}
|
||||
task, ok := manager.Get(taskID)
|
||||
if !ok {
|
||||
t.Fatal("manager did not load migrated task")
|
||||
}
|
||||
if task.OutputFile != canonicalOutput {
|
||||
t.Fatalf("migrated output file = %q, want %q", task.OutputFile, canonicalOutput)
|
||||
}
|
||||
if warnings := manager.LoadWarnings(); len(warnings) != 0 {
|
||||
t.Fatalf("migration produced load warnings: %#v", warnings)
|
||||
}
|
||||
|
||||
persisted, err := os.ReadFile(filepath.Join(root, taskID+".json"))
|
||||
if err != nil {
|
||||
t.Fatalf("read migrated metadata: %v", err)
|
||||
}
|
||||
var saved Task
|
||||
if err := json.Unmarshal(persisted, &saved); err != nil {
|
||||
t.Fatalf("decode migrated metadata: %v", err)
|
||||
}
|
||||
if saved.OutputFile != canonicalOutput {
|
||||
t.Fatalf("persisted output file = %q, want %q", saved.OutputFile, canonicalOutput)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerRejectsUnrelatedAbsoluteOutputPathDuringMigration(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
const taskID = "legacy_task"
|
||||
if err := os.WriteFile(filepath.Join(root, taskID+".ndjson"), []byte("output\n"), 0o600); err != nil {
|
||||
t.Fatalf("write canonical output: %v", err)
|
||||
}
|
||||
metadata, err := json.Marshal(Task{
|
||||
ID: taskID,
|
||||
Type: "specialist",
|
||||
Status: StatusCompleted,
|
||||
OutputFile: filepath.Join(t.TempDir(), taskID+".ndjson"),
|
||||
StartedAt: time.Date(2026, 8, 1, 10, 0, 0, 0, time.UTC),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("marshal metadata: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(root, taskID+".json"), append(metadata, '\n'), 0o600); err != nil {
|
||||
t.Fatalf("write metadata: %v", err)
|
||||
}
|
||||
|
||||
manager, err := NewManager(root)
|
||||
if err != nil {
|
||||
t.Fatalf("NewManager returned error: %v", err)
|
||||
}
|
||||
if _, ok := manager.Get(taskID); ok {
|
||||
t.Fatal("manager loaded task with an unrelated absolute output path")
|
||||
}
|
||||
if warnings := strings.Join(manager.LoadWarnings(), "\n"); !strings.Contains(warnings, "output file must be inside") {
|
||||
t.Fatalf("load warnings = %q, want containment rejection", warnings)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerMarksReloadedRunningTaskNonKillable(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
manager, err := NewManager(root)
|
||||
|
||||
@@ -39,3 +39,21 @@ func ConfigureChildProcessGroup(cmd *exec.Cmd) {
|
||||
func terminateProcess(pid int) error {
|
||||
return execution.TerminateProcessTree(pid, terminationGracePeriod, terminationPollInterval)
|
||||
}
|
||||
|
||||
// terminateOwnedProcess terminates cmd's process group directly when its launch
|
||||
// configuration proves it was made a group leader. This avoids fragile Getpgid
|
||||
// rediscovery after an owned leader exits. Ordinary commands fall back to the
|
||||
// safe PID/tree path rather than assuming their PID is also a process-group ID.
|
||||
//
|
||||
// The Setpgid && Pgid == 0 guard only recognizes ConfigureChildProcessGroup's
|
||||
// own convention. A command made its own session leader via Setsid (as opposed
|
||||
// to Setpgid) is also its own process-group leader in practice, but takes the
|
||||
// slower rediscovery path here since Setsid isn't checked. Harmless today
|
||||
// because TerminateCommand has exactly one caller in this codebase; worth
|
||||
// covering explicitly if Setsid-configured commands start using this path too.
|
||||
func terminateOwnedProcess(cmd *exec.Cmd) (bool, error) {
|
||||
if cmd.SysProcAttr != nil && cmd.SysProcAttr.Setpgid && cmd.SysProcAttr.Pgid == 0 {
|
||||
return false, execution.TerminateProcessGroup(cmd.Process.Pid, terminationGracePeriod, terminationPollInterval)
|
||||
}
|
||||
return false, terminateProcess(cmd.Process.Pid)
|
||||
}
|
||||
@@ -147,11 +147,12 @@ func TestTerminateProcessKillsForkedChildren(t *testing.T) {
|
||||
t.Fatalf("terminateProcess: %v", err)
|
||||
}
|
||||
|
||||
// The forked child must be gone too (poll until reaped by init).
|
||||
// The forked child must no longer be running. An orphaned zombie is already
|
||||
// dead but may remain visible briefly until the platform's init reaps it.
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
// Only ESRCH proves the child is gone; any other error (e.g. EPERM) would
|
||||
// wrongly pass the test, so treat it as still-present and keep polling.
|
||||
for !errors.Is(syscall.Kill(childPID, syscall.Signal(0)), syscall.ESRCH) {
|
||||
// processStopped treats only ESRCH or an observed zombie state as stopped;
|
||||
// other errors (for example EPERM) remain live and keep polling.
|
||||
for !processStopped(childPID) {
|
||||
if time.Now().After(deadline) {
|
||||
_ = syscall.Kill(childPID, syscall.SIGKILL)
|
||||
t.Fatalf("forked child %d survived terminateProcess — group kill failed", childPID)
|
||||
@@ -159,3 +160,73 @@ func TestTerminateProcessKillsForkedChildren(t *testing.T) {
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTerminateCommandStopsUnconfiguredCommand(t *testing.T) {
|
||||
cmd := exec.Command("sleep", "30")
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.Fatalf("start: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
if cmd.ProcessState == nil {
|
||||
_ = cmd.Process.Kill()
|
||||
}
|
||||
})
|
||||
|
||||
if err := TerminateCommand(cmd); err != nil {
|
||||
t.Fatalf("TerminateCommand: %v", err)
|
||||
}
|
||||
if cmd.ProcessState == nil {
|
||||
t.Fatal("unconfigured command was not reaped")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTerminateCommandKillsChildAfterLeaderExits(t *testing.T) {
|
||||
grace, poll := terminationGracePeriod, terminationPollInterval
|
||||
terminationGracePeriod, terminationPollInterval = 2*time.Second, 20*time.Millisecond
|
||||
t.Cleanup(func() { terminationGracePeriod, terminationPollInterval = grace, poll })
|
||||
|
||||
// The leader exits immediately after launching the child. TerminateCommand
|
||||
// must capture and signal the process group before Wait reaps the leader and
|
||||
// makes its group identity unavailable.
|
||||
cmd := exec.Command("sh", "-c", "sleep 300 & echo $!; exit 0")
|
||||
ConfigureChildProcessGroup(cmd)
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
t.Fatalf("stdout pipe: %v", err)
|
||||
}
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.Fatalf("start: %v", err)
|
||||
}
|
||||
line, err := bufio.NewReader(stdout).ReadString('\n')
|
||||
if err != nil {
|
||||
t.Fatalf("read forked child pid: %v", err)
|
||||
}
|
||||
childPID, err := strconv.Atoi(strings.TrimSpace(line))
|
||||
if err != nil {
|
||||
t.Fatalf("parse forked child pid %q: %v", line, err)
|
||||
}
|
||||
|
||||
if err := TerminateCommand(cmd); err != nil {
|
||||
t.Fatalf("TerminateCommand: %v", err)
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for !processStopped(childPID) {
|
||||
if time.Now().After(deadline) {
|
||||
_ = syscall.Kill(childPID, syscall.SIGKILL)
|
||||
t.Fatalf("forked child %d survived TerminateCommand", childPID)
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
func processStopped(pid int) bool {
|
||||
if errors.Is(syscall.Kill(pid, syscall.Signal(0)), syscall.ESRCH) {
|
||||
return true
|
||||
}
|
||||
state, err := exec.Command("ps", "-o", "stat=", "-p", strconv.Itoa(pid)).Output()
|
||||
if err != nil {
|
||||
return errors.Is(syscall.Kill(pid, syscall.Signal(0)), syscall.ESRCH)
|
||||
}
|
||||
return strings.HasPrefix(strings.TrimSpace(string(state)), "Z")
|
||||
}
|
||||
@@ -3,11 +3,20 @@
|
||||
package background
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os/exec"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/execution"
|
||||
"golang.org/x/sys/windows"
|
||||
)
|
||||
|
||||
// stillActiveExitCode is the documented GetExitCodeProcess value for a process
|
||||
// that has not terminated. x/sys/windows does not expose Windows' STILL_ACTIVE
|
||||
// constant.
|
||||
const stillActiveExitCode uint32 = 259
|
||||
|
||||
var terminateProcessForTest = terminateProcess
|
||||
|
||||
// ConfigureChildProcessGroup is a no-op on Windows: process-tree termination is
|
||||
// delegated to execution.TerminateProcessTree, so no launch-time process-group
|
||||
// setup is required (the POSIX build sets Setpgid here instead).
|
||||
@@ -16,3 +25,28 @@ func ConfigureChildProcessGroup(cmd *exec.Cmd) { execution.ConfigureProcessGroup
|
||||
func terminateProcess(pid int) error {
|
||||
return execution.TerminateProcessTree(pid, 0, 0)
|
||||
}
|
||||
|
||||
// terminateOwnedProcess asks taskkill /T to terminate the tree rooted at cmd,
|
||||
// with KillProcessTree's direct Process.Kill fallback if taskkill fails. taskkill
|
||||
// can discover descendants only while the root PID still exists; unlike a POSIX
|
||||
// process group, a Windows tree has no independently addressable identity after
|
||||
// its root exits.
|
||||
func terminateOwnedProcess(cmd *exec.Cmd) (bool, error) {
|
||||
var (
|
||||
alreadyExited bool
|
||||
terminateErr error
|
||||
)
|
||||
err := cmd.Process.WithHandle(func(handle uintptr) {
|
||||
var exitCode uint32
|
||||
if windows.GetExitCodeProcess(windows.Handle(handle), &exitCode) == nil {
|
||||
alreadyExited = exitCode != stillActiveExitCode
|
||||
}
|
||||
// Keep the exact process identity pinned while the PID-based taskkill
|
||||
// operation runs, preventing the PID from being recycled underneath it.
|
||||
terminateErr = terminateProcessForTest(cmd.Process.Pid)
|
||||
})
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("pin process %d for termination: %w", cmd.Process.Pid, err)
|
||||
}
|
||||
return alreadyExited, terminateErr
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
//go:build windows
|
||||
|
||||
package background
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"errors"
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestTerminateCommandReturnsTreeFailureAfterRootFallbackReaps(t *testing.T) {
|
||||
cmd := exec.Command("powershell.exe", "-NoProfile", "-NonInteractive", "-Command", "Start-Sleep -Seconds 300")
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.Fatalf("start: %v", err)
|
||||
}
|
||||
treeErr := errors.New("taskkill /T failed")
|
||||
originalTerminateProcess := terminateProcessForTest
|
||||
terminateProcessForTest = func(pid int) error {
|
||||
if err := cmd.Process.Kill(); err != nil {
|
||||
return errors.Join(treeErr, err)
|
||||
}
|
||||
return treeErr
|
||||
}
|
||||
t.Cleanup(func() { terminateProcessForTest = originalTerminateProcess })
|
||||
|
||||
err := TerminateCommand(cmd)
|
||||
if !errors.Is(err, treeErr) {
|
||||
t.Fatalf("TerminateCommand error = %v, want tree failure", err)
|
||||
}
|
||||
if cmd.ProcessState == nil {
|
||||
t.Fatal("root process was not reaped")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTerminateCommandReapsExitedLeader(t *testing.T) {
|
||||
cmd := exec.Command("cmd.exe", "/c", "exit", "0")
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.Fatalf("start: %v", err)
|
||||
}
|
||||
// Let the child exit without calling Wait so cleanup exercises the Windows
|
||||
// taskkill/TerminateProcess race while it still owns the process handle.
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
|
||||
// jatmn's #774 finding: taskkill /T racing an already-exited PID (and the
|
||||
// Process.Kill fallback) both fail here, but the reap below still succeeds —
|
||||
// the process-handle exit-code check proves the leader was already gone, so
|
||||
// TerminateCommand must not surface the expected kill-attempt error.
|
||||
if err := TerminateCommand(cmd); err != nil {
|
||||
t.Fatalf("TerminateCommand: %v, want nil (leader was already gone and got reaped)", err)
|
||||
}
|
||||
if cmd.ProcessState == nil {
|
||||
t.Fatal("exited process was not reaped")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTerminateCommandDeadLeaderCannotIdentifyDescendantTree(t *testing.T) {
|
||||
cmd := exec.Command("powershell.exe", "-NoProfile", "-NonInteractive", "-Command",
|
||||
`$p = Start-Process powershell.exe -ArgumentList '-NoProfile','-Command','Start-Sleep -Seconds 300' -PassThru; $p.Id`)
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
t.Fatalf("stdout pipe: %v", err)
|
||||
}
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.Fatalf("start: %v", err)
|
||||
}
|
||||
line, err := bufio.NewReader(stdout).ReadString('\n')
|
||||
if err != nil {
|
||||
t.Fatalf("read descendant pid: %v", err)
|
||||
}
|
||||
descendantPID, err := strconv.Atoi(strings.TrimSpace(line))
|
||||
if err != nil {
|
||||
t.Fatalf("parse descendant pid %q: %v", line, err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_ = exec.Command("taskkill.exe", "/F", "/PID", strconv.Itoa(descendantPID)).Run()
|
||||
})
|
||||
|
||||
// Wait until the leader is observably exited but deliberately leave it
|
||||
// unreaped, matching the POSIX regression's ordering.
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for processExistsOnWindows(cmd.Process.Pid) {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("leader did not exit")
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
if err := TerminateCommand(cmd); err != nil {
|
||||
t.Fatalf("TerminateCommand: %v", err)
|
||||
}
|
||||
|
||||
// Windows has no persistent group ID analogous to POSIX: once the leader is
|
||||
// gone, taskkill /T cannot rediscover this descendant. Assert that platform
|
||||
// distinction rather than claiming TerminateCommand can provide it here.
|
||||
if !processExistsOnWindows(descendantPID) {
|
||||
t.Fatal("descendant unexpectedly stopped after its root exited")
|
||||
}
|
||||
}
|
||||
|
||||
func processExistsOnWindows(pid int) bool {
|
||||
return exec.Command("powershell.exe", "-NoProfile", "-NonInteractive", "-Command",
|
||||
"Get-Process -Id "+strconv.Itoa(pid)+" -ErrorAction Stop | Out-Null").Run() == nil
|
||||
}
|
||||
@@ -1,5 +1,14 @@
|
||||
package background
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"time"
|
||||
)
|
||||
|
||||
const commandReapTimeout = 3 * time.Second
|
||||
|
||||
// TerminateProcess stops a background process by PID — on Windows its process
|
||||
// tree; on POSIX its whole process group when the PID leads its own group (the
|
||||
// invariant ConfigureChildProcessGroup establishes for processes started through
|
||||
@@ -12,3 +21,69 @@ package background
|
||||
func TerminateProcess(pid int) error {
|
||||
return terminateProcess(pid)
|
||||
}
|
||||
|
||||
// TerminateCommand stops a started command and reaps its leader. The caller must
|
||||
// have exclusive ownership of cmd: it must not have previously called Wait or
|
||||
// Process.Release, and no goroutine may call either concurrently. On POSIX it
|
||||
// stops the whole group when the command was configured as its leader; ordinary
|
||||
// commands safely fall back to PID/tree discovery. On Windows, success for a
|
||||
// leader that was already dead confirms only that the leader was reaped;
|
||||
// descendants may survive because Windows cannot rediscover a tree from a dead
|
||||
// root. `zero daemon start` needs this operation when readiness times out: it
|
||||
// launched the child, so it must both stop the tree and collect the leader.
|
||||
//
|
||||
// The order matters: the tree is signalled first, because Wait releases the
|
||||
// leader's PID and a later group lookup could then resolve to nothing (or, worse,
|
||||
// to a recycled PID). Termination stays in the shared execution lifecycle
|
||||
// primitives rather than introducing a second platform-specific killer.
|
||||
//
|
||||
// If the bounded reap times out, the Wait call already started by this function
|
||||
// continues in its goroutine. From that point onward the caller must not access
|
||||
// cmd at all, including ProcessState, because that would race with Wait.
|
||||
func TerminateCommand(cmd *exec.Cmd) error {
|
||||
if cmd == nil || cmd.Process == nil {
|
||||
return errors.New("terminate command: process was never started")
|
||||
}
|
||||
leaderAlreadyExited, terminateErr := terminateOwnedProcess(cmd)
|
||||
reapErr := waitForTerminatedCommandWithin(cmd, commandReapTimeout)
|
||||
if reapErr != nil {
|
||||
if terminateErr != nil {
|
||||
return fmt.Errorf("%v (reap failed: %w)", terminateErr, reapErr)
|
||||
}
|
||||
return reapErr
|
||||
}
|
||||
if terminateErr != nil && !leaderAlreadyExited {
|
||||
return terminateErr
|
||||
}
|
||||
// Only discard a termination error when the platform demonstrated before
|
||||
// attempting termination that the leader had already exited. A successful
|
||||
// reap alone says nothing about whether a live descendant tree was stopped.
|
||||
return nil
|
||||
}
|
||||
|
||||
// classifyWaitError treats a non-zero exit status as success: a process being
|
||||
// terminated is expected to report one (or a signal), so the only interesting
|
||||
// failure is Wait itself not working.
|
||||
func classifyWaitError(waitErr error) error {
|
||||
var exitErr *exec.ExitError
|
||||
if waitErr != nil && !errors.As(waitErr, &exitErr) {
|
||||
return fmt.Errorf("reap process: %w", waitErr)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// waitForTerminatedCommandWithin bounds the reap so a child that somehow survives
|
||||
// termination cannot hang the caller — the daemon-start cleanup path runs while a
|
||||
// user waits at the CLI. On timeout its Wait goroutine is intentionally left to
|
||||
// finish: exec.Cmd permits only one Wait call, and abandoning it without a reaper
|
||||
// would leak the child once it eventually exits.
|
||||
func waitForTerminatedCommandWithin(cmd *exec.Cmd, timeout time.Duration) error {
|
||||
waitDone := make(chan error, 1)
|
||||
go func() { waitDone <- cmd.Wait() }()
|
||||
select {
|
||||
case waitErr := <-waitDone:
|
||||
return classifyWaitError(waitErr)
|
||||
case <-time.After(timeout):
|
||||
return fmt.Errorf("process %d did not reap after termination", cmd.Process.Pid)
|
||||
}
|
||||
}
|
||||
+135
-38
@@ -64,38 +64,53 @@ type appDeps struct {
|
||||
// getenv reads a process environment variable (production: os.Getenv, set in
|
||||
// defaultAppDeps — deliberately NOT filled by fillAppDeps, so tests are hermetic
|
||||
// against ambient vars like ZERO_PROVIDER unless they inject it). nil ⇒ empty.
|
||||
getenv func(string) string
|
||||
probeProviderHealth func(context.Context, providerhealth.Options) providerhealth.Result
|
||||
discoverProviderModels func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error)
|
||||
detectLocalRuntimes func(context.Context, provideronboarding.LocalDetectOptions) []provideronboarding.DetectedLocalRuntime
|
||||
openRouterLogin func(context.Context, provideroauth.OpenRouterOptions) (string, error)
|
||||
newSessionStore func() *sessions.Store
|
||||
loadPlugins func(plugins.LoadOptions) (plugins.LoadResult, error)
|
||||
loadHooks func(hooks.LoadOptions) (hooks.LoadResult, error)
|
||||
skillsDir func() string
|
||||
pluginsDir func() string
|
||||
toolsDir func() string
|
||||
newMCPStore func() (*mcp.PermissionStore, error)
|
||||
newMCPTokenStore func() (*mcp.TokenStore, error)
|
||||
newSandboxStore func() (*sandbox.GrantStore, error)
|
||||
selectSandboxBackend func(sandbox.BackendOptions) sandbox.Backend
|
||||
runSandboxSetupHelper func(path string, args []string, stdout io.Writer, stderr io.Writer) error
|
||||
registerMCPTools func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error)
|
||||
prepareWorktree func(context.Context, worktrees.Options) (worktrees.Result, error)
|
||||
releaseWorktree func(context.Context, worktrees.Options, string) error
|
||||
detectVerifyPlan func(string) (verify.Plan, error)
|
||||
runVerify func(context.Context, verify.Plan, verify.RunOptions) verify.Report
|
||||
runSelfVerify func(context.Context, verify.Plan, selfverify.Options) selfverify.Report
|
||||
runAgentEval func(context.Context, agentEvalOptions) (agentEvalReport, error)
|
||||
inspectChanges func(context.Context, zerogit.InspectOptions) (zerogit.ChangeSummary, error)
|
||||
commitChanges func(context.Context, zerogit.CommitOptions) (zerogit.CommitResult, error)
|
||||
pushChanges func(context.Context, zerogit.PushOptions) (zerogit.PushResult, error)
|
||||
createPR func(context.Context, zerogit.PROptions) (zerogit.PRResult, error)
|
||||
runTUI func(context.Context, tui.Options) int
|
||||
runEditor func(string) error
|
||||
checkUpdate func(context.Context, update.Options) (update.Result, error)
|
||||
applyUpdate func(context.Context, update.Options) (update.ApplyResult, error)
|
||||
now func() time.Time
|
||||
getenv func(string) string
|
||||
probeProviderHealth func(context.Context, providerhealth.Options) providerhealth.Result
|
||||
discoverProviderModels func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error)
|
||||
detectLocalRuntimes func(context.Context, provideronboarding.LocalDetectOptions) []provideronboarding.DetectedLocalRuntime
|
||||
openRouterLogin func(context.Context, provideroauth.OpenRouterOptions) (string, error)
|
||||
newSessionStore func() *sessions.Store
|
||||
loadPlugins func(plugins.LoadOptions) (plugins.LoadResult, error)
|
||||
loadHooks func(hooks.LoadOptions) (hooks.LoadResult, error)
|
||||
skillsDir func() string
|
||||
pluginsDir func() string
|
||||
toolsDir func() string
|
||||
newMCPStore func() (*mcp.PermissionStore, error)
|
||||
newMCPTokenStore func() (*mcp.TokenStore, error)
|
||||
newSandboxStore func() (*sandbox.GrantStore, error)
|
||||
selectSandboxBackend func(sandbox.BackendOptions) sandbox.Backend
|
||||
runSandboxSetupHelper func(path string, args []string, stdout io.Writer, stderr io.Writer) error
|
||||
registerMCPTools func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error)
|
||||
prepareWorktree func(context.Context, worktrees.Options) (worktrees.Result, error)
|
||||
releaseWorktree func(context.Context, worktrees.Options, string) error
|
||||
detectVerifyPlan func(string) (verify.Plan, error)
|
||||
runVerify func(context.Context, verify.Plan, verify.RunOptions) verify.Report
|
||||
runSelfVerify func(context.Context, verify.Plan, selfverify.Options) selfverify.Report
|
||||
runAgentEval func(context.Context, agentEvalOptions) (agentEvalReport, error)
|
||||
inspectChanges func(context.Context, zerogit.InspectOptions) (zerogit.ChangeSummary, error)
|
||||
commitChanges func(context.Context, zerogit.CommitOptions) (zerogit.CommitResult, error)
|
||||
pushChanges func(context.Context, zerogit.PushOptions) (zerogit.PushResult, error)
|
||||
createPR func(context.Context, zerogit.PROptions) (zerogit.PRResult, error)
|
||||
createBranch func(context.Context, zerogit.BranchOptions) (zerogit.BranchResult, error)
|
||||
isDefaultBranch func(context.Context, zerogit.DefaultBranchOptions) (bool, string, string, error)
|
||||
currentGitUser func(context.Context, string) string
|
||||
headCommitSubject func(context.Context, string) string
|
||||
commitsAhead func(context.Context, string, string, string) (int, error)
|
||||
isUnbornRemote func(context.Context, string, string) (bool, error)
|
||||
refreshTrackingRef func(context.Context, string, string, string) error
|
||||
branchUpstreamRemote func(context.Context, string, string) string
|
||||
branchUpstreamRemoteAndMerge func(context.Context, string, string) (string, string)
|
||||
resolveRemoteBranchTip func(context.Context, string, string, string) (string, error)
|
||||
remoteHasBranch func(context.Context, string, string, string) (bool, error)
|
||||
currentGitBranch func(context.Context, string) string
|
||||
currentBranchTip func(context.Context, string) string
|
||||
deleteBranch func(context.Context, string, string, string) error
|
||||
resetBranchRef func(context.Context, string, string, string, string) error
|
||||
runTUI func(context.Context, tui.Options) int
|
||||
runEditor func(string) error
|
||||
checkUpdate func(context.Context, update.Options) (update.Result, error)
|
||||
applyUpdate func(context.Context, update.Options) (update.ApplyResult, error)
|
||||
now func() time.Time
|
||||
}
|
||||
|
||||
type mcpToolRuntime interface {
|
||||
@@ -198,11 +213,52 @@ func defaultAppDeps() appDeps {
|
||||
commitChanges: zerogit.Commit,
|
||||
pushChanges: zerogit.Push,
|
||||
createPR: zerogit.CreatePR,
|
||||
runTUI: tui.Run,
|
||||
runEditor: openEditor,
|
||||
checkUpdate: update.Check,
|
||||
applyUpdate: update.Apply,
|
||||
now: time.Now,
|
||||
createBranch: zerogit.CreateBranch,
|
||||
isDefaultBranch: zerogit.IsDefaultBranch,
|
||||
currentGitUser: func(ctx context.Context, cwd string) string {
|
||||
return zerogit.CurrentGitUser(ctx, cwd, nil)
|
||||
},
|
||||
headCommitSubject: func(ctx context.Context, cwd string) string {
|
||||
return zerogit.HeadCommitSubject(ctx, cwd, nil)
|
||||
},
|
||||
commitsAhead: func(ctx context.Context, cwd, remote, branch string) (int, error) {
|
||||
return zerogit.CommitsAhead(ctx, cwd, remote, branch, nil)
|
||||
},
|
||||
isUnbornRemote: func(ctx context.Context, cwd, remote string) (bool, error) {
|
||||
return zerogit.IsUnbornRemote(ctx, cwd, remote, nil)
|
||||
},
|
||||
refreshTrackingRef: func(ctx context.Context, cwd, remote, branch string) error {
|
||||
return zerogit.RefreshTrackingRef(ctx, cwd, remote, branch, nil)
|
||||
},
|
||||
branchUpstreamRemote: func(ctx context.Context, cwd, branch string) string {
|
||||
return zerogit.UpstreamRemote(ctx, cwd, branch, nil)
|
||||
},
|
||||
branchUpstreamRemoteAndMerge: func(ctx context.Context, cwd, branch string) (string, string) {
|
||||
return zerogit.UpstreamRemoteAndMergeBranch(ctx, cwd, branch, nil)
|
||||
},
|
||||
resolveRemoteBranchTip: func(ctx context.Context, cwd, remote, branch string) (string, error) {
|
||||
return zerogit.ResolveRemoteBranchTip(ctx, cwd, remote, branch, nil)
|
||||
},
|
||||
remoteHasBranch: func(ctx context.Context, cwd, remote, branch string) (bool, error) {
|
||||
return zerogit.RemoteHasBranch(ctx, cwd, remote, branch, nil)
|
||||
},
|
||||
currentGitBranch: func(ctx context.Context, cwd string) string {
|
||||
return zerogit.CurrentBranch(ctx, cwd, nil)
|
||||
},
|
||||
currentBranchTip: func(ctx context.Context, cwd string) string {
|
||||
return zerogit.CurrentBranchTip(ctx, cwd, nil)
|
||||
},
|
||||
deleteBranch: func(ctx context.Context, cwd, fallbackBranch, branchToDelete string) error {
|
||||
return zerogit.DeleteBranch(ctx, cwd, fallbackBranch, branchToDelete, nil)
|
||||
},
|
||||
resetBranchRef: func(ctx context.Context, cwd, branch, newTip, expectedOld string) error {
|
||||
return zerogit.ResetBranchRef(ctx, cwd, branch, newTip, nil, expectedOld)
|
||||
},
|
||||
runTUI: tui.Run,
|
||||
runEditor: openEditor,
|
||||
checkUpdate: update.Check,
|
||||
applyUpdate: update.Apply,
|
||||
now: time.Now,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -562,6 +618,45 @@ func fillAppDeps(deps appDeps) appDeps {
|
||||
if deps.createPR == nil {
|
||||
deps.createPR = defaults.createPR
|
||||
}
|
||||
if deps.createBranch == nil {
|
||||
deps.createBranch = defaults.createBranch
|
||||
}
|
||||
if deps.isDefaultBranch == nil {
|
||||
deps.isDefaultBranch = defaults.isDefaultBranch
|
||||
}
|
||||
if deps.currentGitUser == nil {
|
||||
deps.currentGitUser = defaults.currentGitUser
|
||||
}
|
||||
if deps.headCommitSubject == nil {
|
||||
deps.headCommitSubject = defaults.headCommitSubject
|
||||
}
|
||||
if deps.commitsAhead == nil {
|
||||
deps.commitsAhead = defaults.commitsAhead
|
||||
}
|
||||
if deps.isUnbornRemote == nil {
|
||||
deps.isUnbornRemote = defaults.isUnbornRemote
|
||||
}
|
||||
if deps.refreshTrackingRef == nil {
|
||||
deps.refreshTrackingRef = defaults.refreshTrackingRef
|
||||
}
|
||||
if deps.branchUpstreamRemote == nil {
|
||||
deps.branchUpstreamRemote = defaults.branchUpstreamRemote
|
||||
}
|
||||
if deps.branchUpstreamRemoteAndMerge == nil {
|
||||
deps.branchUpstreamRemoteAndMerge = defaults.branchUpstreamRemoteAndMerge
|
||||
}
|
||||
if deps.remoteHasBranch == nil {
|
||||
deps.remoteHasBranch = defaults.remoteHasBranch
|
||||
}
|
||||
if deps.currentGitBranch == nil {
|
||||
deps.currentGitBranch = defaults.currentGitBranch
|
||||
}
|
||||
if deps.currentBranchTip == nil {
|
||||
deps.currentBranchTip = defaults.currentBranchTip
|
||||
}
|
||||
// resolveRemoteBranchTip, deleteBranch, and resetBranchRef stay nil when
|
||||
// unset so unit tests that mock createBranch without a real git tree do not
|
||||
// hit real git restore/delete commands.
|
||||
if deps.runTUI == nil {
|
||||
deps.runTUI = defaults.runTUI
|
||||
}
|
||||
@@ -761,7 +856,7 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
|
||||
// A server that could not be reached or validated is skipped, not fatal (one
|
||||
// bad MCP server must not abort startup) — surface each so a missing tool set is
|
||||
// explained rather than silently absent. A built-in default the user never
|
||||
// configured (e.g. keyless Firecrawl with no credentials) is the exception: it
|
||||
// configured (e.g. keyless Exa with no credentials) is the exception: it
|
||||
// was never asked for, so its failure is not worth a startup warning — only
|
||||
// servers the user actually configured warn on failure (issue #552).
|
||||
for _, skipped := range mcpRuntime.Skipped() {
|
||||
@@ -842,6 +937,7 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
|
||||
Version: version,
|
||||
Theme: theme,
|
||||
SavedTheme: resolved.Preferences.Theme,
|
||||
SavedPet: resolved.Preferences.Pet,
|
||||
UserConfigPath: userConfigPath,
|
||||
DoctorUserConfigPath: doctorUserConfigPath,
|
||||
ProjectConfigPath: projectConfigPath,
|
||||
@@ -852,6 +948,7 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
|
||||
FavoriteModels: resolved.Preferences.FavoriteModels,
|
||||
RecentModels: resolved.Preferences.RecentModels,
|
||||
RecapsEnabled: resolved.Preferences.RecapsEnabled(),
|
||||
CompactionModel: resolved.Preferences.CompactionModel,
|
||||
Provider: provider,
|
||||
NewProvider: deps.newProvider,
|
||||
ProbeProviderHealth: deps.probeProviderHealth,
|
||||
|
||||
@@ -211,12 +211,12 @@ func TestTUIStartupSuppressesWarningForUnconfiguredDefaultServer(t *testing.T) {
|
||||
},
|
||||
resolveMCPConfig: func(workspaceRoot string, excludeProject bool) (config.MCPConfig, error) {
|
||||
return config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"firecrawl": config.DefaultMCPServers()["firecrawl"],
|
||||
"exa": config.DefaultMCPServers()["exa"],
|
||||
}}, nil
|
||||
},
|
||||
registerMCPTools: func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
|
||||
return fakeMCPRuntimeWithSkips{skipped: []mcp.SkippedServer{
|
||||
{Name: "firecrawl", Err: errors.New("returned HTTP 401"), UnconfiguredDefault: true},
|
||||
{Name: "exa", Err: errors.New("returned HTTP 401"), UnconfiguredDefault: true},
|
||||
}}, nil
|
||||
},
|
||||
runTUI: func(ctx context.Context, options tui.Options) int {
|
||||
@@ -227,7 +227,7 @@ func TestTUIStartupSuppressesWarningForUnconfiguredDefaultServer(t *testing.T) {
|
||||
if exitCode != 0 {
|
||||
t.Fatalf("exitCode = %d stderr=%s", exitCode, stderr.String())
|
||||
}
|
||||
if strings.Contains(stderr.String(), "firecrawl") {
|
||||
if strings.Contains(stderr.String(), "exa") {
|
||||
t.Fatalf("stderr = %q, want no warning for an unconfigured default server", stderr.String())
|
||||
}
|
||||
}
|
||||
|
||||
+52
-10
@@ -1,6 +1,7 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
@@ -160,20 +161,61 @@ func runDaemonStartDetached(paths daemon.Paths, stdout io.Writer, stderr io.Writ
|
||||
cmd.Stdout = logFile
|
||||
cmd.Stderr = logFile
|
||||
background.ConfigureChildProcessGroup(cmd) // own process group: outlives this shell
|
||||
if err := cmd.Start(); err != nil {
|
||||
if err := startAndAwaitDaemonProcess(cmd, func() bool { return daemonReachable(paths) }, 5*time.Second, 25*time.Millisecond); err != nil {
|
||||
if errors.Is(err, errDaemonStartTimeout) {
|
||||
return writeAppError(stderr, err.Error()+"; see "+logPath, exitCrash)
|
||||
}
|
||||
return writeAppError(stderr, err.Error(), exitCrash)
|
||||
}
|
||||
_ = cmd.Process.Release()
|
||||
fmt.Fprintf(stdout, "zero daemon started (socket %s)\n", paths.Socket)
|
||||
return exitSuccess
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if daemonReachable(paths) {
|
||||
fmt.Fprintf(stdout, "zero daemon started (socket %s)\n", paths.Socket)
|
||||
return exitSuccess
|
||||
}
|
||||
time.Sleep(25 * time.Millisecond)
|
||||
var errDaemonStartTimeout = errors.New("daemon did not come up within timeout")
|
||||
|
||||
func startAndAwaitDaemonProcess(cmd *exec.Cmd, reachable func() bool, timeout time.Duration, pollInterval time.Duration) error {
|
||||
if err := cmd.Start(); err != nil {
|
||||
return err
|
||||
}
|
||||
return writeAppError(stderr, "daemon did not come up within timeout; see "+logPath, exitCrash)
|
||||
if waitForDaemonReadiness(reachable, timeout, pollInterval) {
|
||||
if err := cmd.Process.Release(); err != nil {
|
||||
cleanupErr := terminateAndReapDaemonProcess(cmd)
|
||||
if cleanupErr != nil {
|
||||
return fmt.Errorf("release daemon process: %w (cleanup failed: %v)", err, cleanupErr)
|
||||
}
|
||||
return fmt.Errorf("release daemon process: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if err := terminateAndReapDaemonProcess(cmd); err != nil {
|
||||
return fmt.Errorf("%w; cleanup failed: %v", errDaemonStartTimeout, err)
|
||||
}
|
||||
return errDaemonStartTimeout
|
||||
}
|
||||
|
||||
func waitForDaemonReadiness(reachable func() bool, timeout time.Duration, pollInterval time.Duration) bool {
|
||||
deadline := time.Now().Add(timeout)
|
||||
for {
|
||||
if reachable() {
|
||||
return true
|
||||
}
|
||||
remaining := time.Until(deadline)
|
||||
if remaining <= 0 {
|
||||
break
|
||||
}
|
||||
time.Sleep(min(pollInterval, remaining))
|
||||
}
|
||||
// The loop's last reachable() call already ran immediately before this one
|
||||
// (right after the final Sleep, with the same deadline check separating
|
||||
// them), so this adds no extra waiting time and is not a grace window for a
|
||||
// slow-starting daemon — it only re-checks the same instant in case
|
||||
// reachable() itself is non-deterministic. To tolerate a genuinely slow
|
||||
// start, raise the timeout passed to startAndAwaitDaemonProcess instead.
|
||||
return reachable()
|
||||
}
|
||||
|
||||
func terminateAndReapDaemonProcess(cmd *exec.Cmd) error {
|
||||
return background.TerminateCommand(cmd)
|
||||
}
|
||||
|
||||
func runDaemonStop(args []string, stdout io.Writer, stderr io.Writer) int {
|
||||
|
||||
@@ -2,8 +2,15 @@ package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/background"
|
||||
)
|
||||
|
||||
// isolateDaemonPaths points DefaultPaths at a temp dir so the test never touches
|
||||
@@ -121,3 +128,84 @@ func TestDaemonSubcommandsRejectExtraArgs(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitForDaemonReadinessChecksTimeoutBoundary(t *testing.T) {
|
||||
checks := 0
|
||||
ready := waitForDaemonReadiness(func() bool {
|
||||
checks++
|
||||
return checks == 2
|
||||
}, 0, time.Millisecond)
|
||||
if !ready {
|
||||
t.Fatal("daemon that became reachable at the timeout boundary was not detected")
|
||||
}
|
||||
if checks != 2 {
|
||||
t.Fatalf("reachability checks = %d, want initial and final checks", checks)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTerminateAndReapDaemonProcess(t *testing.T) {
|
||||
cmd := daemonDetachedChildCommand("hang")
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.Fatalf("start helper process: %v", err)
|
||||
}
|
||||
|
||||
if err := terminateAndReapDaemonProcess(cmd); err != nil {
|
||||
t.Fatalf("terminate and reap helper process: %v", err)
|
||||
}
|
||||
if cmd.ProcessState == nil {
|
||||
t.Fatalf("helper process was not reaped: state=%v", cmd.ProcessState)
|
||||
}
|
||||
if cmd.ProcessState.Success() {
|
||||
t.Fatalf("helper process unexpectedly exited successfully: %v", cmd.ProcessState)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStartAndAwaitDaemonProcessTimeoutTerminatesAndReaps(t *testing.T) {
|
||||
cmd := daemonDetachedChildCommand("hang")
|
||||
if err := startAndAwaitDaemonProcess(cmd, func() bool { return false }, 0, time.Millisecond); !errors.Is(err, errDaemonStartTimeout) {
|
||||
t.Fatalf("start and await error = %v, want timeout", err)
|
||||
}
|
||||
if cmd.ProcessState == nil {
|
||||
t.Fatal("timed-out helper process was not reaped")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStartAndAwaitDaemonProcessReleasesReadyChild(t *testing.T) {
|
||||
marker := filepath.Join(t.TempDir(), "child-finished")
|
||||
cmd := daemonDetachedChildCommand("mark", "ZERO_TEST_DAEMON_CHILD_MARKER="+marker)
|
||||
if err := startAndAwaitDaemonProcess(cmd, func() bool { return true }, time.Second, time.Millisecond); err != nil {
|
||||
t.Fatalf("start and await ready helper: %v", err)
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if _, err := os.Stat(marker); err == nil {
|
||||
return
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("released helper process did not remain alive to write its marker")
|
||||
}
|
||||
|
||||
func daemonDetachedChildCommand(mode string, extraEnv ...string) *exec.Cmd {
|
||||
cmd := exec.Command(os.Args[0], "-test.run=TestDaemonDetachedChildProcess")
|
||||
cmd.Env = append(os.Environ(), "ZERO_TEST_DAEMON_DETACHED_CHILD="+mode)
|
||||
cmd.Env = append(cmd.Env, extraEnv...)
|
||||
background.ConfigureChildProcessGroup(cmd)
|
||||
return cmd
|
||||
}
|
||||
|
||||
func TestDaemonDetachedChildProcess(t *testing.T) {
|
||||
switch os.Getenv("ZERO_TEST_DAEMON_DETACHED_CHILD") {
|
||||
case "":
|
||||
return
|
||||
case "mark":
|
||||
if err := os.WriteFile(os.Getenv("ZERO_TEST_DAEMON_CHILD_MARKER"), []byte("ready"), 0o600); err != nil {
|
||||
os.Exit(2)
|
||||
}
|
||||
os.Exit(0)
|
||||
}
|
||||
for {
|
||||
time.Sleep(time.Hour)
|
||||
}
|
||||
}
|
||||
+14
-10
@@ -682,16 +682,20 @@ func runExec(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) in
|
||||
ModelSwitcher: modelSwitcher,
|
||||
TurnSessionProvider: turnSessions,
|
||||
ModelSessionSwitcher: modelSessionSwitcher,
|
||||
ReasoningEffort: forwardEffort,
|
||||
Trace: traceRecorder,
|
||||
Cwd: workspaceRoot,
|
||||
Images: images,
|
||||
Registry: registry,
|
||||
PermissionMode: permissionMode,
|
||||
Autonomy: options.autonomy,
|
||||
SelfCorrect: selfCorrector,
|
||||
FileDiagnostics: fileDiagnostics,
|
||||
Profile: execProfile.Policy(displacedMaxTurns, execProfileFilledEffort),
|
||||
Summarizer: summarizerFactory(resolved, deps.newProvider),
|
||||
ContextWindowFor: func(modelID string) int {
|
||||
return modelregistry.AgentContextWindow(modelContextWindow(modelRegistry, modelID))
|
||||
},
|
||||
ReasoningEffort: forwardEffort,
|
||||
Trace: traceRecorder,
|
||||
Cwd: workspaceRoot,
|
||||
Images: images,
|
||||
Registry: registry,
|
||||
PermissionMode: permissionMode,
|
||||
Autonomy: options.autonomy,
|
||||
SelfCorrect: selfCorrector,
|
||||
FileDiagnostics: fileDiagnostics,
|
||||
Profile: execProfile.Policy(displacedMaxTurns, execProfileFilledEffort),
|
||||
// Headless exec: don't accept a no-tool-call turn as "done" while work
|
||||
// clearly remains (pending plan items / a mid-step continuation cue) —
|
||||
// nudge to continue, and finalize as INCOMPLETE rather than false success
|
||||
|
||||
@@ -9,15 +9,15 @@ import (
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
)
|
||||
|
||||
// `zero mcp disable firecrawl` must work even though firecrawl is a built-in
|
||||
// `zero mcp disable exa` must work even though exa is a built-in
|
||||
// default that is not written to the user's config file until overridden.
|
||||
func TestRunMCPDisableSeededFirecrawlDefault(t *testing.T) {
|
||||
func TestRunMCPDisableSeededExaDefault(t *testing.T) {
|
||||
configPath := filepath.Join(t.TempDir(), "zero", "config.json")
|
||||
// A config with no firecrawl entry — the default lives in code, not the file.
|
||||
// A config with no Exa entry — the default lives in code, not the file.
|
||||
writeMCPCommandRawConfig(t, configPath, `{"activeProvider":"fast"}`)
|
||||
|
||||
var out, errBuf bytes.Buffer
|
||||
code := runWithDeps([]string{"mcp", "disable", "firecrawl", "--json"}, &out, &errBuf, appDeps{
|
||||
code := runWithDeps([]string{"mcp", "disable", "exa", "--json"}, &out, &errBuf, appDeps{
|
||||
userConfigPath: func() (string, error) { return configPath, nil },
|
||||
})
|
||||
if code != exitSuccess {
|
||||
@@ -31,8 +31,8 @@ func TestRunMCPDisableSeededFirecrawlDefault(t *testing.T) {
|
||||
if err := json.Unmarshal(out.Bytes(), &payload); err != nil {
|
||||
t.Fatalf("decode disable JSON: %v\n%s", err, out.String())
|
||||
}
|
||||
if payload.ServerName != "firecrawl" || !payload.Disabled || !payload.Changed {
|
||||
t.Fatalf("disable payload = %#v, want firecrawl disabled+changed", payload)
|
||||
if payload.ServerName != "exa" || !payload.Disabled || !payload.Changed {
|
||||
t.Fatalf("disable payload = %#v, want exa disabled+changed", payload)
|
||||
}
|
||||
|
||||
// End-to-end: resolving that config now turns the default off.
|
||||
@@ -40,7 +40,7 @@ func TestRunMCPDisableSeededFirecrawlDefault(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
if !cfg.Servers["firecrawl"].Disabled {
|
||||
t.Fatal("expected `mcp disable firecrawl` to turn the seeded default off in the resolved config")
|
||||
if !cfg.Servers["exa"].Disabled {
|
||||
t.Fatal("expected `mcp disable exa` to turn the seeded default off in the resolved config")
|
||||
}
|
||||
}
|
||||
@@ -1,13 +1,8 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/oauth"
|
||||
"github.com/Gitlawb/zero/internal/providers"
|
||||
"github.com/Gitlawb/zero/internal/providers/providerio"
|
||||
)
|
||||
|
||||
@@ -28,46 +23,5 @@ import (
|
||||
// ProviderProfile.OAuthLoginCandidates so the runtime resolver, the Codex account
|
||||
// resolver, and the onboarding presence check never diverge.
|
||||
func oauthLoginForProfile(profile config.ProviderProfile) (providerio.TokenResolver, string) {
|
||||
candidates := profile.OAuthLoginCandidates()
|
||||
if len(candidates) == 0 {
|
||||
return nil, ""
|
||||
}
|
||||
store, err := oauth.NewStore(oauth.StoreOptions{})
|
||||
if err != nil {
|
||||
return nil, ""
|
||||
}
|
||||
_, key, ok := oauth.FirstStored(store, candidates)
|
||||
if !ok {
|
||||
// No login under any candidate (or unreadable/invalid keys) → API-key
|
||||
// auth, no resolver.
|
||||
return nil, ""
|
||||
}
|
||||
manager, err := oauth.NewManager(oauth.ManagerOptions{
|
||||
Store: store,
|
||||
HTTPClient: &http.Client{Timeout: 30 * time.Second},
|
||||
// Refreshing a token the user logged into (possibly a preset provider like
|
||||
// xAI) re-resolves that provider's OAuth config, which needs the preset.
|
||||
AllowPresets: true,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, ""
|
||||
}
|
||||
resolver := func(ctx context.Context, forceRefresh bool) (string, string, bool, error) {
|
||||
var token string
|
||||
var rerr error
|
||||
if forceRefresh {
|
||||
token, rerr = manager.Handle401(ctx, key)
|
||||
} else {
|
||||
token, rerr = manager.GetFresh(ctx, key)
|
||||
}
|
||||
if errors.Is(rerr, oauth.ErrNoToken) {
|
||||
// The login was removed since construction → fall back to the API key.
|
||||
return "", "", false, nil
|
||||
}
|
||||
if rerr != nil {
|
||||
return "", "", false, rerr
|
||||
}
|
||||
return "Authorization", "Bearer " + token, true, nil
|
||||
}
|
||||
return resolver, key
|
||||
return providers.OAuthLoginForProfile(profile)
|
||||
}
|
||||
@@ -73,6 +73,18 @@ func runProvidersModels(args []string, stdout io.Writer, stderr io.Writer, deps
|
||||
if description := strings.TrimSpace(model.Description); description != "" {
|
||||
entry["description"] = description
|
||||
}
|
||||
if len(model.ReasoningEfforts) > 0 {
|
||||
entry["reasoning_efforts"] = append([]string{}, model.ReasoningEfforts...)
|
||||
}
|
||||
if effort := strings.TrimSpace(model.DefaultReasoningEffort); effort != "" {
|
||||
entry["default_reasoning_effort"] = effort
|
||||
}
|
||||
if len(model.ServiceTiers) > 0 {
|
||||
entry["service_tiers"] = append([]string{}, model.ServiceTiers...)
|
||||
}
|
||||
if tier := strings.TrimSpace(model.DefaultServiceTier); tier != "" {
|
||||
entry["default_service_tier"] = tier
|
||||
}
|
||||
items = append(items, entry)
|
||||
}
|
||||
payload := map[string]any{
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/agent"
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/providers"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// summarizerFactory adapts the resolved profile and the authenticated provider
|
||||
// builder into agent.Options.Summarizer; the selection rules live in
|
||||
// providers.CompactionSummarizerFactory, shared with the TUI.
|
||||
func summarizerFactory(resolved config.ResolvedConfig, newProvider func(config.ProviderProfile) (zeroruntime.Provider, error)) func(context.Context, string) (agent.Provider, error) {
|
||||
return providers.CompactionSummarizerFactory(resolved.Provider, resolved.Preferences.CompactionModel, newProvider)
|
||||
}
|
||||
+2059
-15
File diff suppressed because it is too large.
Load diff
+538
-8
@@ -2,9 +2,12 @@ package cli
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -602,18 +605,20 @@ func parseChangesArgs(args []string, command string) (changesCommandOptions, boo
|
||||
return options, false, execUsageError{fmt.Sprintf("unexpected changes argument %q", arg)}
|
||||
}
|
||||
}
|
||||
if command != "commit" && options.message != "" {
|
||||
if command != "commit" && options.hasMessage {
|
||||
return options, false, execUsageError{"--message is only valid with `zero changes commit`"}
|
||||
}
|
||||
if command != "commit" && (options.hasMessage || options.dryRun || options.auto) {
|
||||
return options, false, execUsageError{"--message, --dry-run, and --auto are only valid with `zero changes commit`"}
|
||||
}
|
||||
if command == "commit" && options.hasMessage && options.auto {
|
||||
return options, false, execUsageError{"cannot specify both --message and --auto"}
|
||||
}
|
||||
if command != "commit" && command != "push" && options.dryRun {
|
||||
return options, false, execUsageError{"--dry-run is only valid with commit or push"}
|
||||
}
|
||||
// --auto on push/pr is the explicit opt-in for LLM branch naming (see
|
||||
// ensureFeatureBranch); on commit it opts into the LLM commit message.
|
||||
if command != "commit" && command != "push" && command != "pr" && options.auto {
|
||||
return options, false, execUsageError{"--auto is only valid with commit, push, or pr"}
|
||||
}
|
||||
if command != "inspect" && options.baseRef != "" {
|
||||
return options, false, execUsageError{"--base is only valid with `zero changes inspect`"}
|
||||
}
|
||||
@@ -934,8 +939,7 @@ Flags:
|
||||
--fill Automatically populate PR title and body from commits
|
||||
--draft Create PR as a draft
|
||||
--yes Confirm pushing to a default/protected branch
|
||||
-a, --auto Auto-generate commit message using LLM (use --dry-run to preview)
|
||||
--dry-run Preview commit metadata without mutating git state
|
||||
-a, --auto Use the LLM: commit generates the message, push/pr name the auto-created branch (sends the diff to the provider)
|
||||
--json Print JSON output
|
||||
-h, --help Show this help
|
||||
`)
|
||||
@@ -958,11 +962,31 @@ func runChangesPush(args []string, stdout io.Writer, stderr io.Writer, deps appD
|
||||
return writeExecUsageError(stderr, err.Error())
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
branch, remote, created, err := ensureFeatureBranch(ctx, stdout, workspaceRoot, options.remote, featureBranchOptions{
|
||||
JSONMode: options.json,
|
||||
AllowDefaultBranch: options.yes,
|
||||
DryRun: options.dryRun,
|
||||
AutoNaming: options.auto,
|
||||
MaxDiffBytes: options.maxDiffBytes,
|
||||
}, deps)
|
||||
if err != nil {
|
||||
return writeExecUsageError(stderr, err.Error())
|
||||
}
|
||||
|
||||
// context.Background() is deliberate: the 2-minute ctx above is the
|
||||
// feature-branch preflight (inspection, remote probes, branch creation).
|
||||
// The push itself must not be cancelled mid-transfer just because the
|
||||
// preflight budget elapsed; a large push can legitimately take longer.
|
||||
result, err := deps.pushChanges(context.Background(), zerogit.PushOptions{
|
||||
Cwd: workspaceRoot,
|
||||
Remote: options.remote,
|
||||
Remote: firstNonEmptyString(options.remote, remote),
|
||||
Branch: branch,
|
||||
Force: options.force,
|
||||
DryRun: options.dryRun,
|
||||
RequireNewRemoteBranch: created,
|
||||
AllowPushDefaultBranch: options.yes,
|
||||
})
|
||||
if err != nil {
|
||||
@@ -1010,15 +1034,62 @@ func runChangesPR(args []string, stdout io.Writer, stderr io.Writer, deps appDep
|
||||
return writeExecUsageError(stderr, err.Error())
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
||||
defer cancel()
|
||||
|
||||
branch, remote, created, err := ensureFeatureBranch(ctx, stdout, workspaceRoot, options.remote, featureBranchOptions{
|
||||
JSONMode: options.json,
|
||||
AllowDefaultBranch: options.yes,
|
||||
DryRun: false,
|
||||
AutoNaming: options.auto,
|
||||
MaxDiffBytes: options.maxDiffBytes,
|
||||
}, deps)
|
||||
if err != nil {
|
||||
return writeExecUsageError(stderr, err.Error())
|
||||
}
|
||||
|
||||
// Resolve the target remote for the unborn-remote check. Explicit --remote wins,
|
||||
// followed by the remote resolved by ensureFeatureBranch. When --yes is set,
|
||||
// fall back to the branch upstream or "origin".
|
||||
targetRemote := strings.TrimSpace(options.remote)
|
||||
if targetRemote == "" {
|
||||
targetRemote = remote
|
||||
}
|
||||
if targetRemote == "" {
|
||||
currentBranch := ""
|
||||
if deps.currentGitBranch != nil {
|
||||
currentBranch = deps.currentGitBranch(context.Background(), workspaceRoot)
|
||||
}
|
||||
if deps.branchUpstreamRemote != nil && currentBranch != "" {
|
||||
targetRemote = deps.branchUpstreamRemote(context.Background(), workspaceRoot, currentBranch)
|
||||
}
|
||||
if targetRemote == "" {
|
||||
targetRemote = "origin"
|
||||
}
|
||||
}
|
||||
|
||||
// Probe errors are intentionally ignored so the flow can continue; if the
|
||||
// probe fails, the subsequent push or PR operation provides fail-closed behavior.
|
||||
if deps.isUnbornRemote != nil {
|
||||
if unborn, unbornErr := deps.isUnbornRemote(context.Background(), workspaceRoot, targetRemote); unbornErr == nil && unborn {
|
||||
return writeExecUsageError(stderr, fmt.Sprintf("cannot create pull request on unborn remote %s: push the initial default branch first", targetRemote))
|
||||
}
|
||||
}
|
||||
|
||||
if !options.json {
|
||||
if _, err := fmt.Fprintln(stdout, "Pushing current branch to set upstream..."); err != nil {
|
||||
return exitCrash
|
||||
}
|
||||
}
|
||||
// context.Background() is deliberate here too: runChangesPR reuses the
|
||||
// 2-minute ctx for preflight (branch creation and remote probes), but the
|
||||
// push must not be cancelled mid-transfer if that budget elapsed.
|
||||
pushResult, err := deps.pushChanges(context.Background(), zerogit.PushOptions{
|
||||
Cwd: workspaceRoot,
|
||||
Remote: options.remote,
|
||||
Remote: firstNonEmptyString(options.remote, remote),
|
||||
Branch: branch,
|
||||
Force: options.force,
|
||||
RequireNewRemoteBranch: created,
|
||||
AllowPushDefaultBranch: options.yes,
|
||||
})
|
||||
if err != nil {
|
||||
@@ -1100,3 +1171,462 @@ func generateAutoCommitMessage(ctx context.Context, provider zeroruntime.Provide
|
||||
}
|
||||
return msg, nil
|
||||
}
|
||||
|
||||
type featureBranchOptions struct {
|
||||
JSONMode bool
|
||||
AllowDefaultBranch bool
|
||||
DryRun bool
|
||||
AutoNaming bool
|
||||
MaxDiffBytes int
|
||||
}
|
||||
|
||||
// branchPushPlan holds the resolved parameters for a push or PR workflow
|
||||
// before any branches are created or refs mutated. Separating plan resolution
|
||||
// from mutation allows tests to assert the resolution matrix directly without
|
||||
// touching a git repository.
|
||||
type branchPushPlan struct {
|
||||
ShouldBranch bool
|
||||
CurrentBranch string
|
||||
PushRemote string
|
||||
RequireNewRemoteBranch bool
|
||||
TrackingRemote string
|
||||
TrackingBranch string
|
||||
RestoreTip string
|
||||
PublishabilityBase string
|
||||
Summary zerogit.ChangeSummary
|
||||
Slug string
|
||||
}
|
||||
|
||||
// resolveBranchPushPlan resolves the branch, remotes, restore tip, and naming
|
||||
// slug for changes push / pr workflows. It performs all checks (clean working
|
||||
// tree, unborn remote detection, ahead count, restore tip resolution) but
|
||||
// executes no mutation (no branch creation, ref resets, or deletions).
|
||||
func resolveBranchPushPlan(ctx context.Context, stdout io.Writer, workspaceRoot string, requestedRemote string, opts featureBranchOptions, deps appDeps) (branchPushPlan, error) {
|
||||
if opts.AllowDefaultBranch || opts.DryRun {
|
||||
return branchPushPlan{
|
||||
ShouldBranch: false,
|
||||
PushRemote: strings.TrimSpace(requestedRemote),
|
||||
RequireNewRemoteBranch: false,
|
||||
}, nil
|
||||
}
|
||||
|
||||
if deps.isDefaultBranch == nil {
|
||||
return branchPushPlan{}, fmt.Errorf("isDefaultBranch dependency missing")
|
||||
}
|
||||
isDefault, currentBranch, remote, err := deps.isDefaultBranch(ctx, zerogit.DefaultBranchOptions{Cwd: workspaceRoot, Remote: requestedRemote})
|
||||
if err != nil {
|
||||
return branchPushPlan{}, err
|
||||
}
|
||||
if !isDefault {
|
||||
// Whether this push needs the nonexistence lease is decided by one
|
||||
// live check against the destination remote, not by inference from
|
||||
// local git state (upstream config, a generated-branch marker):
|
||||
// every local signal that stood in for that check turned out to be
|
||||
// wrong in some ordering (inherited config from checkout -b, a
|
||||
// config write that failed after a successful push, a push aimed at
|
||||
// a different remote than the one local config describes). A branch
|
||||
// missing on the remote gets the nonexistence lease, protecting a
|
||||
// concurrent creator of the same name; a branch already there gets a
|
||||
// plain push, and git's own fast-forward check is the safety net —
|
||||
// if the remote history isn't ours, the push fails with a clear
|
||||
// non-fast-forward error instead of an attempted silent recovery.
|
||||
if deps.remoteHasBranch == nil {
|
||||
return branchPushPlan{
|
||||
ShouldBranch: false,
|
||||
CurrentBranch: currentBranch,
|
||||
PushRemote: remote,
|
||||
RequireNewRemoteBranch: true,
|
||||
}, nil
|
||||
}
|
||||
targetRemote := firstNonEmptyString(requestedRemote, remote)
|
||||
exists, existsErr := deps.remoteHasBranch(ctx, workspaceRoot, targetRemote, currentBranch)
|
||||
if existsErr != nil {
|
||||
return branchPushPlan{}, fmt.Errorf("cannot check whether %s already exists on remote %s: %w", currentBranch, targetRemote, existsErr)
|
||||
}
|
||||
return branchPushPlan{
|
||||
ShouldBranch: false,
|
||||
CurrentBranch: currentBranch,
|
||||
PushRemote: remote,
|
||||
RequireNewRemoteBranch: !exists,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// CreateBranch and Push publish commits only. A dirty working tree would
|
||||
// leave uncommitted edits behind under a branch/PR that does not include
|
||||
// them, so refuse until the tree is clean (commit or stash first).
|
||||
workingTree, err := deps.inspectChanges(ctx, zerogit.InspectOptions{Cwd: workspaceRoot})
|
||||
if err != nil {
|
||||
return branchPushPlan{}, fmt.Errorf("failed to inspect working tree: %w", err)
|
||||
}
|
||||
if !workingTree.Clean {
|
||||
return branchPushPlan{}, fmt.Errorf("working tree has uncommitted changes; commit or stash them before pushing from the default branch")
|
||||
}
|
||||
|
||||
// Capture the source default branch's own upstream before refreshing or
|
||||
// leaving it. --remote selects the push destination; restore must not
|
||||
// rewrite local main to upstream/main when main still tracks origin/main
|
||||
// (fork setups), and refreshing must update the exact tracking ref that
|
||||
// will supply both the ahead count and the restore target.
|
||||
trackingRemote := remote
|
||||
trackingBranch := currentBranch
|
||||
if deps.branchUpstreamRemoteAndMerge != nil {
|
||||
if r, m := deps.branchUpstreamRemoteAndMerge(ctx, workspaceRoot, currentBranch); r != "" && m != "" {
|
||||
trackingRemote = r
|
||||
trackingBranch = m
|
||||
}
|
||||
}
|
||||
|
||||
// Refresh the local remote-tracking ref before trusting it: IsDefaultBranch
|
||||
// already contacted the remote for its symref check, but the tracking ref
|
||||
// commitsAhead reads from is only a local cache (written at clone or the
|
||||
// last fetch) and can sit behind the remote's live tip. Left stale, a real
|
||||
// publishable range could look like zero commits ahead, and the diff
|
||||
// derived below (which reuses this same ref as its base) would be stale
|
||||
// too. A failure here (offline, or a genuinely unborn remote with no such
|
||||
// ref to fetch yet) is folded into the same fail-closed path as an
|
||||
// unresolvable ahead count below.
|
||||
var fetchErr error
|
||||
if deps.refreshTrackingRef != nil {
|
||||
fetchErr = deps.refreshTrackingRef(ctx, workspaceRoot, trackingRemote, trackingBranch)
|
||||
}
|
||||
|
||||
// The ahead check gates on the source upstream (trackingRemote), but the
|
||||
// feature branch is pushed to `pushRemote`, which can be a different
|
||||
// destination (--remote on a fork). Refuse to auto-branch onto an unborn
|
||||
// destination: without an initial default branch the new feature branch
|
||||
// becomes the destination's first ref and its HEAD is left pointing at it.
|
||||
// This is checked once, here, against the push destination, before resolving
|
||||
// commit tips or creating any branches.
|
||||
pushRemote := firstNonEmptyString(requestedRemote, remote)
|
||||
if deps.isUnbornRemote != nil {
|
||||
if unbornDest, unbornDestErr := deps.isUnbornRemote(ctx, workspaceRoot, pushRemote); unbornDestErr == nil && unbornDest {
|
||||
return branchPushPlan{}, fmt.Errorf("remote %s has no branches yet; push the initial default branch first with --yes (`zero changes push --yes`), then use auto-branching for subsequent work", pushRemote)
|
||||
}
|
||||
}
|
||||
|
||||
// Branching off the default branch only makes sense when HEAD carries a
|
||||
// commit that is not already on the remote default branch. A clean,
|
||||
// up-to-date default branch would otherwise publish a feature branch at the
|
||||
// exact default tip. If the ahead count cannot be determined (for example
|
||||
// the remote-tracking ref was never fetched), fail rather than guess. When
|
||||
// the remote is confirmed unborn, refuse auto-branching entirely: the
|
||||
// initial default branch must be established first (changes push --yes)
|
||||
// so the remote has a real HEAD before any feature branch is published.
|
||||
ahead, aheadErr := deps.commitsAhead(ctx, workspaceRoot, trackingRemote, trackingBranch)
|
||||
if fetchErr != nil || aheadErr != nil {
|
||||
var unbornRemote bool
|
||||
var unbornErr error
|
||||
if deps.isUnbornRemote != nil {
|
||||
unbornRemote, unbornErr = deps.isUnbornRemote(ctx, workspaceRoot, trackingRemote)
|
||||
}
|
||||
if unbornErr == nil && unbornRemote {
|
||||
return branchPushPlan{}, fmt.Errorf("remote %s has no branches yet; push the initial default branch first with --yes (`zero changes push --yes`), then use auto-branching for subsequent work", trackingRemote)
|
||||
}
|
||||
cause := aheadErr
|
||||
if cause == nil {
|
||||
cause = fetchErr
|
||||
}
|
||||
return branchPushPlan{}, fmt.Errorf("cannot determine whether HEAD is ahead of %s/%s: %w; fetch the remote tracking branch first", trackingRemote, trackingBranch, cause)
|
||||
}
|
||||
if ahead == 0 {
|
||||
return branchPushPlan{}, fmt.Errorf("no changes to publish: HEAD is not ahead of %s/%s; commit your work before pushing", trackingRemote, trackingBranch)
|
||||
}
|
||||
|
||||
// Resolve the restore tip to a concrete commit before creating any branch.
|
||||
// It must never be a bare local branch name: resolving "main" yields the
|
||||
// very commit the feature branch is about to own, so restore becomes a
|
||||
// silent no-op and the publishable commits stay on the default branch. A
|
||||
// named remote resolves through its just-refreshed remote-tracking ref; a
|
||||
// direct URL or path asks the remote itself. Fail closed rather than
|
||||
// substitute a local branch name.
|
||||
restoreTip := ""
|
||||
if deps.resolveRemoteBranchTip != nil {
|
||||
tip, tipErr := deps.resolveRemoteBranchTip(ctx, workspaceRoot, trackingRemote, trackingBranch)
|
||||
if tipErr != nil {
|
||||
return branchPushPlan{}, fmt.Errorf("cannot resolve restore tip for %s/%s: %w", trackingRemote, trackingBranch, tipErr)
|
||||
}
|
||||
restoreTip = tip
|
||||
}
|
||||
if restoreTip == "" {
|
||||
restoreTip = trackingRemote + "/" + trackingBranch
|
||||
}
|
||||
|
||||
// Name the branch (and, with --auto, send the provider) from what HEAD is
|
||||
// actually ahead of the resolved remote branch by, using the same ref
|
||||
// commitsAhead just checked. A working-tree snapshot can describe edits a
|
||||
// commit-only push will never include.
|
||||
baseRef := restoreTip
|
||||
summary, err := deps.inspectChanges(ctx, zerogit.InspectOptions{Cwd: workspaceRoot, BaseRef: baseRef, MaxDiffBytes: opts.MaxDiffBytes})
|
||||
if err != nil {
|
||||
return branchPushPlan{}, fmt.Errorf("failed to inspect changes: %w", err)
|
||||
}
|
||||
|
||||
slug := fallbackBranchSlug(summary)
|
||||
if len(summary.Files) == 0 {
|
||||
// An empty commit (or one whose only content the diff omits) leaves
|
||||
// nothing to derive a slug from; name the branch from the commit
|
||||
// subject instead.
|
||||
if subject := deps.headCommitSubject(ctx, workspaceRoot); subject != "" {
|
||||
slug = zerogit.SlugifyBranchComponent(subject)
|
||||
}
|
||||
}
|
||||
if opts.AutoNaming && strings.TrimSpace(summary.Diff) != "" {
|
||||
llmSlug := ""
|
||||
if resolved, cfgErr := deps.resolveConfig(workspaceRoot, config.Overrides{}); cfgErr == nil && config.HasProviderProfile(resolved.Provider) {
|
||||
if provider, provErr := deps.newProvider(resolved.Provider); provErr == nil {
|
||||
if !opts.JSONMode {
|
||||
fmt.Fprintln(stdout, "Generating branch name using LLM...")
|
||||
}
|
||||
genCtx, cancel := context.WithTimeout(ctx, 60*time.Second)
|
||||
generated, genErr := generateAutoBranchSlug(genCtx, provider, resolved.Provider.Model, redactChangeSummary(summary))
|
||||
cancel()
|
||||
if genErr == nil && generated != "" {
|
||||
llmSlug = generated
|
||||
}
|
||||
}
|
||||
}
|
||||
if llmSlug != "" {
|
||||
slug = llmSlug
|
||||
} else if !opts.JSONMode {
|
||||
// Cover resolveConfig failures, missing provider profiles,
|
||||
// newProvider errors, generation errors, and empty results so the
|
||||
// user is not left guessing after "Generating..." or silence.
|
||||
fmt.Fprintln(stdout, "LLM branch naming unavailable; using deterministic name.")
|
||||
}
|
||||
}
|
||||
|
||||
return branchPushPlan{
|
||||
ShouldBranch: true,
|
||||
CurrentBranch: currentBranch,
|
||||
PushRemote: pushRemote,
|
||||
RequireNewRemoteBranch: true,
|
||||
TrackingRemote: trackingRemote,
|
||||
TrackingBranch: trackingBranch,
|
||||
RestoreTip: restoreTip,
|
||||
PublishabilityBase: baseRef,
|
||||
Summary: summary,
|
||||
Slug: slug,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ensureFeatureBranch is the branch-naming step `zero changes push`/`pr` run
|
||||
// before pushing: pushing straight to the default branch is refused deeper in
|
||||
// zerogit.Push, so rather than surface that as a dead end, create and switch
|
||||
// to a conventionally named "<user>/<slug>" branch first. It returns the
|
||||
// branch push/pr should target, or "" to mean "current HEAD branch, unchanged"
|
||||
// (zerogit.Push already treats an empty Branch that way), plus the remote the
|
||||
// preflight resolved (requestedRemote, then the original branch's configured
|
||||
// upstream, then "origin"), plus whether Push should require that branch not
|
||||
// already exist on that remote. Callers must pass the remote to Push: a
|
||||
// freshly created branch has no tracking configuration, so Push's own
|
||||
// fallback would silently retarget "origin" even when the work came from a
|
||||
// branch tracking a different remote. Callers must also pass that third value
|
||||
// through as Push's RequireNewRemoteBranch: CreateBranch's own
|
||||
// remote-collision probe runs before this returns, and closing that race
|
||||
// requires Push's push itself to assert the destination is still new. That
|
||||
// requirement also carries across a retry: if the current branch is already
|
||||
// non-default (this call didn't create it) but has not been successfully
|
||||
// published to the exact target remote under the same branch name, keep the
|
||||
// nonexistence lease. branch.<name>.remote alone is not enough: with
|
||||
// branch.autoSetupMerge=inherit, checkout -b copies remote=origin from the
|
||||
// source branch before any push. After a successful create, the original
|
||||
// default branch ref is moved back to its own upstream tip (not necessarily
|
||||
// the push destination) so the feature branch exclusively owns the
|
||||
// publishable commits; a failed restore rolls the feature branch back.
|
||||
// allowDefaultBranch (the --yes flag) and dryRun both opt out via the ""
|
||||
// branch / false return, leaving Push's own guard/preview behavior on the
|
||||
// default branch unaffected.
|
||||
//
|
||||
// autoNaming gates the LLM naming path (--auto): these commands were
|
||||
// git-only, and sending the change diff to a configured provider on every
|
||||
// default-branch push would silently export source code nobody asked to
|
||||
// share. Without the opt-in the name comes from deterministic local
|
||||
// information only. maxDiffBytes caps the committed-range diff Inspect
|
||||
// returns, so a user who passed --diff-bytes to bound the proprietary source
|
||||
// sent for LLM naming has that cap honored here just as the commit path does.
|
||||
// The working tree must be clean and HEAD must be ahead of the resolved
|
||||
// remote branch; both are verified before any branch is created. The
|
||||
// inspectChanges, commitsAhead, headCommitSubject, currentGitUser, and
|
||||
// createBranch dependencies are required on the default-branch path and are
|
||||
// called without nil guards; fillAppDeps populates all of them. The
|
||||
// refreshTrackingRef, resolveRemoteBranchTip, isUnbornRemote, resetBranchRef,
|
||||
// and deleteBranch dependencies are optional and are skipped when nil.
|
||||
func ensureFeatureBranch(ctx context.Context, stdout io.Writer, workspaceRoot string, requestedRemote string, opts featureBranchOptions, deps appDeps) (string, string, bool, error) {
|
||||
plan, err := resolveBranchPushPlan(ctx, stdout, workspaceRoot, requestedRemote, opts, deps)
|
||||
if err != nil {
|
||||
return "", "", false, err
|
||||
}
|
||||
if !plan.ShouldBranch {
|
||||
return plan.CurrentBranch, plan.PushRemote, plan.RequireNewRemoteBranch, nil
|
||||
}
|
||||
|
||||
expectedRestoreTip := ""
|
||||
if deps.currentBranchTip != nil {
|
||||
expectedRestoreTip = deps.currentBranchTip(ctx, workspaceRoot)
|
||||
}
|
||||
|
||||
name := zerogit.BuildBranchName(deps.currentGitUser(ctx, workspaceRoot), plan.Slug)
|
||||
result, err := deps.createBranch(ctx, zerogit.BranchOptions{Cwd: workspaceRoot, Name: name, Remote: plan.PushRemote})
|
||||
if err != nil {
|
||||
return "", "", false, fmt.Errorf("failed to create branch: %w", err)
|
||||
}
|
||||
// The normal flow commits on the default branch first, then creates the
|
||||
// feature branch at the same HEAD. Without moving the original default
|
||||
// ref back, local main keeps the pre-squash commits and diverges from
|
||||
// its upstream after a squash-merge (and a later pull/push can re-publish
|
||||
// them). Once the feature branch owns those commits, restore the default
|
||||
// branch to the source branch's own upstream tip captured above.
|
||||
// Failure-safe: if the restore fails, delete the feature branch and
|
||||
// check the default branch back out so the tree is not left half-moved.
|
||||
if deps.resetBranchRef != nil {
|
||||
if err := deps.resetBranchRef(ctx, workspaceRoot, plan.CurrentBranch, plan.RestoreTip, expectedRestoreTip); err != nil {
|
||||
// A compare-and-swap conflict means the default branch moved
|
||||
// while we were branching, so restoreTip is no longer the tip we
|
||||
// meant to restore. The rollback below deletes the branch that
|
||||
// now exclusively owns the user's commits, which would destroy
|
||||
// work to recover from a race we can simply report. Keep the
|
||||
// branch and let the user reconcile the default branch.
|
||||
if errors.Is(err, zerogit.ErrCompareAndSwapConflict) {
|
||||
return "", "", false, fmt.Errorf("failed to restore default branch %s to %s after auto-branching: %w; generated branch %s was preserved because the default branch changed concurrently", plan.CurrentBranch, plan.RestoreTip, err, result.Branch)
|
||||
}
|
||||
var rollbackErr error
|
||||
if deps.deleteBranch != nil {
|
||||
rollbackErr = deps.deleteBranch(ctx, workspaceRoot, plan.CurrentBranch, result.Branch)
|
||||
}
|
||||
if rollbackErr != nil {
|
||||
return "", "", false, fmt.Errorf("failed to restore default branch %s to %s after auto-branching: %w; rollback of %s also failed: %v", plan.CurrentBranch, plan.RestoreTip, err, result.Branch, rollbackErr)
|
||||
}
|
||||
return "", "", false, fmt.Errorf("failed to restore default branch %s to %s after auto-branching: %w", plan.CurrentBranch, plan.RestoreTip, err)
|
||||
}
|
||||
}
|
||||
if !opts.JSONMode {
|
||||
fmt.Fprintf(stdout, "Created branch %s (was on %s)\n", result.Branch, plan.CurrentBranch)
|
||||
}
|
||||
return result.Branch, plan.PushRemote, true, nil
|
||||
}
|
||||
|
||||
// fallbackBranchSlug derives a deterministic branch-name slug from a change
|
||||
// summary without calling an LLM, so ensureFeatureBranch still works when no
|
||||
// provider is configured.
|
||||
func fallbackBranchSlug(summary zerogit.ChangeSummary) string {
|
||||
switch len(summary.Files) {
|
||||
case 0:
|
||||
return "changes"
|
||||
case 1:
|
||||
// Git paths are always slash-separated regardless of platform, so
|
||||
// path.Base (not filepath.Base) extracts the filename portably.
|
||||
return zerogit.SlugifyBranchComponent(path.Base(summary.Files[0].Path))
|
||||
default:
|
||||
return fmt.Sprintf("update-%d-files", len(summary.Files))
|
||||
}
|
||||
}
|
||||
|
||||
// generateAutoBranchSlug asks the model for a short kebab-case slug
|
||||
// describing the diff, mirroring generateAutoCommitMessage's prompt shape.
|
||||
func generateAutoBranchSlug(ctx context.Context, provider zeroruntime.Provider, model string, summary zerogit.ChangeSummary) (string, error) {
|
||||
var promptBuilder strings.Builder
|
||||
promptBuilder.WriteString("Analyze the following git diff and generate a short git branch name slug for it.\n")
|
||||
promptBuilder.WriteString("The slug must be 2 to 5 lowercase words separated by hyphens (kebab-case), using only letters, digits, and hyphens, with no prefix like \"feature/\" or \"fix/\" and no surrounding quotes.\n")
|
||||
promptBuilder.WriteString("Output ONLY the raw slug text, nothing else.\n\n")
|
||||
promptBuilder.WriteString("Git Diff:\n")
|
||||
promptBuilder.WriteString(summary.Diff)
|
||||
|
||||
request := zeroruntime.CompletionRequest{
|
||||
Messages: []zeroruntime.Message{
|
||||
{Role: zeroruntime.MessageRoleUser, Content: promptBuilder.String()},
|
||||
},
|
||||
}
|
||||
stream, err := provider.StreamCompletion(ctx, request)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
collected := zeroruntime.CollectStream(ctx, stream)
|
||||
if collected.Error != "" {
|
||||
return "", fmt.Errorf("%s", collected.Error)
|
||||
}
|
||||
|
||||
slug := zerogit.SlugifyBranchComponent(extractBranchSlug(collected.Text))
|
||||
if slug == "" {
|
||||
return "", fmt.Errorf("provider returned empty branch slug")
|
||||
}
|
||||
return slug, nil
|
||||
}
|
||||
|
||||
// slugLineRe matches a line that already reads as a kebab-case slug (letters,
|
||||
// digits, and internal single hyphens only). extractBranchSlug prefers such a
|
||||
// line so a preamble sentence is never mistaken for the slug.
|
||||
var slugLineRe = regexp.MustCompile(`^[a-zA-Z0-9]+(-[a-zA-Z0-9]+)*$`)
|
||||
|
||||
func isPreambleText(s string) bool {
|
||||
lower := strings.ToLower(strings.TrimSpace(s))
|
||||
if strings.HasSuffix(lower, ":") || strings.HasSuffix(lower, ".") || strings.HasSuffix(lower, "!") || strings.HasSuffix(lower, "?") {
|
||||
return true
|
||||
}
|
||||
preambles := []string{
|
||||
"here is", "here's", "below is", "below are", "sure", "certainly",
|
||||
"i suggest", "suggested branch", "branch name", "recommended branch",
|
||||
"how about", "you could use",
|
||||
}
|
||||
for _, p := range preambles {
|
||||
if strings.HasPrefix(lower, p) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// extractBranchSlug pulls the intended slug out of a model response that didn't
|
||||
// follow the "output only the raw slug" instruction exactly. It drops Markdown
|
||||
// code-fence lines, then prefers a line that already looks like a kebab-case
|
||||
// slug: that skips a leading preamble such as "Here is a suggested branch
|
||||
// name:" in favor of the "add-login-page" line that follows it. One-word
|
||||
// acknowledgements that happen to match the slug regex ("Sure", "Certainly")
|
||||
// are filtered as preamble before the early return, so "Sure\nadd login page"
|
||||
// yields the real suggestion. Inline labeled forms ("Branch name: add-login-page")
|
||||
// still work: the preamble check runs on the extracted value after the label
|
||||
// is stripped. If no line is strictly kebab-cased, it prefers a plausible
|
||||
// non-preamble line so plain multi-word replies still slugify correctly.
|
||||
func extractBranchSlug(text string) string {
|
||||
var firstLine string
|
||||
|
||||
for _, line := range strings.Split(text, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if strings.HasPrefix(line, "```") {
|
||||
continue
|
||||
}
|
||||
line = strings.TrimSpace(strings.Trim(line, `"'`))
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
candidate := line
|
||||
if idx := strings.Index(line, ":"); idx != -1 {
|
||||
after := strings.TrimSpace(strings.Trim(line[idx+1:], `"'`))
|
||||
if after != "" {
|
||||
candidate = after
|
||||
}
|
||||
}
|
||||
// Models often terminate a suggestion with sentence punctuation;
|
||||
// strip it so "add login page." is recovered as the slug rather than
|
||||
// rejected as preamble (isPreambleText treats a trailing '.' as a
|
||||
// full sentence). True preamble lines still match the prefix list.
|
||||
candidate = strings.TrimSpace(strings.TrimRight(candidate, ".!?"))
|
||||
if candidate == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
// Preamble classification runs before accepting a slug-shaped line so
|
||||
// "Sure" does not win over a later "add-login-page". The check is on
|
||||
// candidate (post label-strip) so "Branch name: add-login-page" still
|
||||
// returns the extracted slug.
|
||||
if slugLineRe.MatchString(candidate) && !isPreambleText(candidate) {
|
||||
return candidate
|
||||
}
|
||||
|
||||
if firstLine == "" && !isPreambleText(candidate) {
|
||||
firstLine = candidate
|
||||
}
|
||||
}
|
||||
|
||||
return firstLine
|
||||
}
|
||||
@@ -105,7 +105,17 @@ func TestLoadProviderCommandTimeout(t *testing.T) {
|
||||
//
|
||||
// This is a hang-breaker, not a tuned margin. Nothing in a healthy run waits on
|
||||
// it, since both tests return about one WaitDelay after the fixture exits.
|
||||
const providerCommandTestBudget = 60 * time.Second
|
||||
const (
|
||||
providerCommandTestBudget = 90 * time.Second
|
||||
providerCommandReadyWaitAttempts = 45
|
||||
)
|
||||
|
||||
func TestProviderCommandBudgetLeavesCleanupMargin(t *testing.T) {
|
||||
readyWaitCeiling := time.Duration(providerCommandReadyWaitAttempts) * time.Second
|
||||
if providerCommandTestBudget < 2*readyWaitCeiling {
|
||||
t.Fatalf("provider command test budget %s leaves less than one readiness window for cleanup after %s", providerCommandTestBudget, readyWaitCeiling)
|
||||
}
|
||||
}
|
||||
|
||||
// backgroundChildLifetime must exceed providerCommandTestBudget so that a child
|
||||
// dying of old age is unobservable: any run where runProviderCommand has not
|
||||
@@ -341,15 +351,15 @@ func writeCommand(t *testing.T, script commandScript) string {
|
||||
// assertion fail is far more diagnosable than blocking until the
|
||||
// timeout fires and reporting an unrelated error. Each iteration is
|
||||
// about a second, since `ping -n 2` sleeps between its two echoes,
|
||||
// so the ceiling is roughly 15s: far longer than any cold PowerShell
|
||||
// start seen on a runner, and far shorter than
|
||||
// providerCommandTestBudget, so giving up surfaces a legible
|
||||
// so the ceiling is roughly 45s: enough for a heavily loaded Windows
|
||||
// runner while leaving the other half of providerCommandTestBudget for
|
||||
// process exit, PID-file verification, and cleanup. Giving up surfaces a legible
|
||||
// PID-file failure rather than a timeout that hides the real cause.
|
||||
lines = append(lines,
|
||||
":zeroWaitReady",
|
||||
"if exist \""+readyFile+"\" goto zeroReady",
|
||||
"set /a ZERO_READY_TRIES+=1",
|
||||
"if %ZERO_READY_TRIES% GEQ 15 goto zeroReady",
|
||||
"if %ZERO_READY_TRIES% GEQ "+itoa(providerCommandReadyWaitAttempts)+" goto zeroReady",
|
||||
"ping -n 2 127.0.0.1 >nul",
|
||||
"goto zeroWaitReady",
|
||||
":zeroReady",
|
||||
|
||||
@@ -0,0 +1,294 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
type jsonObjectSpan struct {
|
||||
start int
|
||||
end int
|
||||
members []jsonMemberSpan
|
||||
}
|
||||
|
||||
type jsonMemberSpan struct {
|
||||
key string
|
||||
leadingStart int
|
||||
keyStart int
|
||||
valueStart int
|
||||
valueEnd int
|
||||
commaAfter int
|
||||
}
|
||||
|
||||
// setPetPreferenceJSON changes only preferences.pet. It keeps the original
|
||||
// bytes for every unrelated member so choosing a companion does not reorder or
|
||||
// reformat the user's configuration file.
|
||||
func setPetPreferenceJSON(data []byte, pet string) ([]byte, error) {
|
||||
rootStart := skipJSONSpace(data, 0)
|
||||
root, err := parseJSONObject(data, rootStart)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
preferencesIndex := lastJSONMember(root.members, "preferences")
|
||||
if preferencesIndex < 0 {
|
||||
if pet == "" {
|
||||
return data, nil
|
||||
}
|
||||
encodedPet, err := json.Marshal(pet)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("encode pet preference: %w", err)
|
||||
}
|
||||
return insertJSONMember(data, root, "preferences", petJSONObject(encodedPet)), nil
|
||||
}
|
||||
|
||||
preferenceMember := root.members[preferencesIndex]
|
||||
preferenceValue := bytes.TrimSpace(data[preferenceMember.valueStart:preferenceMember.valueEnd])
|
||||
if bytes.Equal(preferenceValue, []byte("null")) {
|
||||
if pet == "" {
|
||||
if lastJSONMember(root.members[:preferencesIndex], "preferences") >= 0 {
|
||||
return data, nil
|
||||
}
|
||||
return removeJSONMember(data, root, preferencesIndex), nil
|
||||
}
|
||||
encodedPet, err := json.Marshal(pet)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("encode pet preference: %w", err)
|
||||
}
|
||||
return replaceJSONRange(data, preferenceMember.valueStart, preferenceMember.valueEnd, petJSONObject(encodedPet)), nil
|
||||
}
|
||||
preferences, err := parseJSONObject(data, preferenceMember.valueStart)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("preferences must be a JSON object: %w", err)
|
||||
}
|
||||
petIndex := lastJSONMember(preferences.members, "pet")
|
||||
if pet != "" {
|
||||
encodedPet, err := json.Marshal(pet)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("encode pet preference: %w", err)
|
||||
}
|
||||
if petIndex >= 0 {
|
||||
member := preferences.members[petIndex]
|
||||
return replaceJSONRange(data, member.valueStart, member.valueEnd, encodedPet), nil
|
||||
}
|
||||
return insertJSONMember(data, preferences, "pet", encodedPet), nil
|
||||
}
|
||||
|
||||
for petIndex >= 0 {
|
||||
data = removeJSONMember(data, preferences, petIndex)
|
||||
root, err = parseJSONObject(data, rootStart)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
preferencesIndex = lastJSONMember(root.members, "preferences")
|
||||
if preferencesIndex < 0 {
|
||||
return data, nil
|
||||
}
|
||||
preferenceMember = root.members[preferencesIndex]
|
||||
preferences, err = parseJSONObject(data, preferenceMember.valueStart)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
petIndex = lastJSONMember(preferences.members, "pet")
|
||||
}
|
||||
if len(preferences.members) == 0 {
|
||||
if lastJSONMember(root.members[:preferencesIndex], "preferences") >= 0 {
|
||||
return data, nil
|
||||
}
|
||||
return removeJSONMember(data, root, preferencesIndex), nil
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
func parseJSONObject(data []byte, start int) (jsonObjectSpan, error) {
|
||||
if start < 0 || start >= len(data) || data[start] != '{' {
|
||||
return jsonObjectSpan{}, fmt.Errorf("expected JSON object")
|
||||
}
|
||||
object := jsonObjectSpan{start: start}
|
||||
leadingStart := start + 1
|
||||
for {
|
||||
i := skipJSONSpace(data, leadingStart)
|
||||
if i >= len(data) {
|
||||
return jsonObjectSpan{}, fmt.Errorf("unterminated JSON object")
|
||||
}
|
||||
if data[i] == '}' {
|
||||
object.end = i
|
||||
return object, nil
|
||||
}
|
||||
if data[i] != '"' {
|
||||
return jsonObjectSpan{}, fmt.Errorf("expected JSON object key")
|
||||
}
|
||||
keyStart := i
|
||||
keyEnd, err := scanJSONString(data, i)
|
||||
if err != nil {
|
||||
return jsonObjectSpan{}, err
|
||||
}
|
||||
var key string
|
||||
if err := json.Unmarshal(data[keyStart:keyEnd], &key); err != nil {
|
||||
return jsonObjectSpan{}, err
|
||||
}
|
||||
i = skipJSONSpace(data, keyEnd)
|
||||
if i >= len(data) || data[i] != ':' {
|
||||
return jsonObjectSpan{}, fmt.Errorf("expected colon after JSON object key")
|
||||
}
|
||||
valueStart := skipJSONSpace(data, i+1)
|
||||
valueEnd, err := scanJSONValue(data, valueStart)
|
||||
if err != nil {
|
||||
return jsonObjectSpan{}, err
|
||||
}
|
||||
i = skipJSONSpace(data, valueEnd)
|
||||
member := jsonMemberSpan{
|
||||
key: key,
|
||||
leadingStart: leadingStart,
|
||||
keyStart: keyStart,
|
||||
valueStart: valueStart,
|
||||
valueEnd: valueEnd,
|
||||
commaAfter: -1,
|
||||
}
|
||||
if i < len(data) && data[i] == ',' {
|
||||
member.commaAfter = i
|
||||
object.members = append(object.members, member)
|
||||
leadingStart = i + 1
|
||||
continue
|
||||
}
|
||||
if i >= len(data) || data[i] != '}' {
|
||||
return jsonObjectSpan{}, fmt.Errorf("expected comma or end of JSON object")
|
||||
}
|
||||
object.members = append(object.members, member)
|
||||
object.end = i
|
||||
return object, nil
|
||||
}
|
||||
}
|
||||
|
||||
func scanJSONString(data []byte, start int) (int, error) {
|
||||
for i := start + 1; i < len(data); i++ {
|
||||
switch data[i] {
|
||||
case '\\':
|
||||
i++
|
||||
case '"':
|
||||
return i + 1, nil
|
||||
}
|
||||
}
|
||||
return 0, fmt.Errorf("unterminated JSON string")
|
||||
}
|
||||
|
||||
func scanJSONValue(data []byte, start int) (int, error) {
|
||||
if start >= len(data) {
|
||||
return 0, fmt.Errorf("missing JSON value")
|
||||
}
|
||||
if data[start] == '"' {
|
||||
return scanJSONString(data, start)
|
||||
}
|
||||
if data[start] == '{' || data[start] == '[' {
|
||||
stack := []byte{data[start]}
|
||||
for i := start + 1; i < len(data); i++ {
|
||||
if data[i] == '"' {
|
||||
end, err := scanJSONString(data, i)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
i = end - 1
|
||||
continue
|
||||
}
|
||||
switch data[i] {
|
||||
case '{', '[':
|
||||
stack = append(stack, data[i])
|
||||
case '}', ']':
|
||||
open := stack[len(stack)-1]
|
||||
if (open == '{' && data[i] != '}') || (open == '[' && data[i] != ']') {
|
||||
return 0, fmt.Errorf("mismatched JSON delimiter")
|
||||
}
|
||||
stack = stack[:len(stack)-1]
|
||||
if len(stack) == 0 {
|
||||
return i + 1, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0, fmt.Errorf("unterminated JSON value")
|
||||
}
|
||||
for i := start; i < len(data); i++ {
|
||||
switch data[i] {
|
||||
case ',', '}', ']', ' ', '\t', '\r', '\n':
|
||||
if i == start {
|
||||
return 0, fmt.Errorf("missing JSON value")
|
||||
}
|
||||
return i, nil
|
||||
}
|
||||
}
|
||||
return len(data), nil
|
||||
}
|
||||
|
||||
func skipJSONSpace(data []byte, start int) int {
|
||||
for start < len(data) {
|
||||
switch data[start] {
|
||||
case ' ', '\t', '\r', '\n':
|
||||
start++
|
||||
default:
|
||||
return start
|
||||
}
|
||||
}
|
||||
return start
|
||||
}
|
||||
|
||||
func lastJSONMember(members []jsonMemberSpan, key string) int {
|
||||
for i := len(members) - 1; i >= 0; i-- {
|
||||
if members[i].key == key {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
func petJSONObject(encodedPet []byte) []byte {
|
||||
result := append([]byte(`{"pet":`), encodedPet...)
|
||||
return append(result, '}')
|
||||
}
|
||||
|
||||
func replaceJSONRange(data []byte, start, end int, replacement []byte) []byte {
|
||||
result := make([]byte, 0, len(data)-(end-start)+len(replacement))
|
||||
result = append(result, data[:start]...)
|
||||
result = append(result, replacement...)
|
||||
return append(result, data[end:]...)
|
||||
}
|
||||
|
||||
func removeJSONMember(data []byte, object jsonObjectSpan, index int) []byte {
|
||||
member := object.members[index]
|
||||
start, end := member.leadingStart, member.valueEnd
|
||||
if member.commaAfter >= 0 {
|
||||
end = member.commaAfter + 1
|
||||
} else if index > 0 {
|
||||
start = object.members[index-1].commaAfter
|
||||
}
|
||||
return replaceJSONRange(data, start, end, nil)
|
||||
}
|
||||
|
||||
func insertJSONMember(data []byte, object jsonObjectSpan, key string, value []byte) []byte {
|
||||
encodedKey, _ := json.Marshal(key)
|
||||
entry := append(append(encodedKey, ':', ' '), value...)
|
||||
if len(object.members) == 0 {
|
||||
interior := data[object.start+1 : object.end]
|
||||
if newline := bytes.LastIndexByte(interior, '\n'); newline >= 0 {
|
||||
closingIndent := interior[newline+1:]
|
||||
indentUnit := []byte(" ")
|
||||
if bytes.Contains(closingIndent, []byte("\t")) {
|
||||
indentUnit = []byte("\t")
|
||||
}
|
||||
lineBreak := []byte("\n")
|
||||
if newline > 0 && interior[newline-1] == '\r' {
|
||||
lineBreak = []byte("\r\n")
|
||||
}
|
||||
insert := append(append(append(lineBreak, closingIndent...), indentUnit...), entry...)
|
||||
return replaceJSONRange(data, object.start+1, object.start+1, insert)
|
||||
}
|
||||
return replaceJSONRange(data, object.end, object.end, entry)
|
||||
}
|
||||
last := object.members[len(object.members)-1]
|
||||
separator := []byte(", ")
|
||||
if bytes.Contains(data[last.valueEnd:object.end], []byte("\n")) {
|
||||
lineStart := bytes.LastIndexByte(data[:last.keyStart], '\n') + 1
|
||||
indent := data[lineStart:last.keyStart]
|
||||
separator = append([]byte(",\n"), indent...)
|
||||
}
|
||||
insert := append(separator, entry...)
|
||||
return replaceJSONRange(data, last.valueEnd, last.valueEnd, insert)
|
||||
}
|
||||
@@ -5,22 +5,20 @@ import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
// DefaultMCPServers returns the MCP servers Zero ships ENABLED by default so web
|
||||
// search and scraping 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 point firecrawl at a self-hosted instance, or
|
||||
// add an API-key header to lift the free-tier 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`).
|
||||
//
|
||||
// Keyless Firecrawl routes requests through firecrawl.dev (1,000 free credits per
|
||||
// month, no account). Self-host Firecrawl (AGPL-3.0) for unlimited and private
|
||||
// use. Zero only calls it over the network, so Firecrawl's license never reaches
|
||||
// into Zero's own code.
|
||||
// Exa's hosted MCP server works anonymously with rate limits. Users can add an
|
||||
// Exa API key for higher limits.
|
||||
func DefaultMCPServers() map[string]MCPServerConfig {
|
||||
return map[string]MCPServerConfig{
|
||||
"firecrawl": {
|
||||
"exa": {
|
||||
Type: "http",
|
||||
URL: "https://mcp.firecrawl.dev/v2/mcp",
|
||||
URL: "https://mcp.exa.ai/mcp",
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -35,13 +33,13 @@ func IsDefaultMCPServer(name string) bool {
|
||||
|
||||
// IsUnconfiguredDefault reports whether server is one of Zero's built-in
|
||||
// defaults that the user never wrote an entry for in their config — i.e. it is
|
||||
// running with whatever Zero ships (e.g. keyless Firecrawl, no credentials).
|
||||
// running with whatever Zero ships (e.g. keyless Exa, no credentials).
|
||||
//
|
||||
// Both conditions below must hold:
|
||||
// - !server.configured: the user's JSON never declared an object for this
|
||||
// server key at all (set by MCPServerConfig.UnmarshalJSON only when it
|
||||
// actually ran for this key). Any explicit action — including a
|
||||
// disable/enable toggle like `zero mcp enable firecrawl` that leaves the
|
||||
// disable/enable toggle like `zero mcp enable exa` that leaves the
|
||||
// resolved value unchanged — sets configured, so it always counts as
|
||||
// user-configured, even though the value comparison below could not tell
|
||||
// the difference on its own.
|
||||
@@ -58,3 +56,90 @@ func IsUnconfiguredDefault(name string, server MCPServerConfig) bool {
|
||||
def, ok := DefaultMCPServers()[strings.TrimSpace(name)]
|
||||
return ok && !server.configured && reflect.DeepEqual(def, server)
|
||||
}
|
||||
|
||||
// retiredDefaultMCPServer describes a built-in default Zero no longer ships.
|
||||
type retiredDefaultMCPServer struct {
|
||||
// successor is the default that replaced it.
|
||||
successor string
|
||||
// shipped is the value Zero used to seed for this name. A user entry
|
||||
// written while the default was still shipped may name only an override
|
||||
// (a header, say) and rely on these fields for the rest, so they have to
|
||||
// outlive the default itself.
|
||||
shipped MCPServerConfig
|
||||
}
|
||||
|
||||
// retiredDefaultMCPServers maps a retired built-in default to what replaced it
|
||||
// and to the value Zero used to ship for it.
|
||||
var retiredDefaultMCPServers = map[string]retiredDefaultMCPServer{
|
||||
"firecrawl": {
|
||||
successor: "exa",
|
||||
shipped: MCPServerConfig{
|
||||
Type: "http",
|
||||
URL: "https://mcp.firecrawl.dev/v2/mcp",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// migrateRetiredDefaultMCPServers keeps an upgrade across a default-provider
|
||||
// swap from changing what the user's config means. It handles the two ways a
|
||||
// rename can break a config written against the old default:
|
||||
//
|
||||
// - A user who ran `zero mcp disable <old default>` made an explicit choice
|
||||
// not to open that outbound connection. Without carrying that decision to
|
||||
// the successor, the upgrade re-opens it under the new name — and because
|
||||
// the replacement looks like an untouched default, even the startup
|
||||
// warning stays quiet (see IsUnconfiguredDefault and issue #552).
|
||||
// - A user who customized the old default could name only the field they
|
||||
// were changing, because the seeded default supplied the transport. Once
|
||||
// the default stops being seeded that entry resolves to a server with no
|
||||
// type, url, or command, which fails NormalizeConfig and takes the whole
|
||||
// startup down rather than just that one server.
|
||||
//
|
||||
// It runs immediately after the user layer merges (see ResolveMCP), which
|
||||
// scopes it to user-level decisions — the only scope whose disable is sticky.
|
||||
// The carried-over disable applies only when the user never declared the
|
||||
// successor themselves: an explicit `exa` entry wins whether it enables or
|
||||
// disables. Because the disable is recorded as a user-level decision, the
|
||||
// lower-trust project layer cannot lift it, while `zero mcp enable exa` still
|
||||
// can (the CLI override scope merges with canReenable=true).
|
||||
func migrateRetiredDefaultMCPServers(cfg *MCPConfig) {
|
||||
for name, retired := range retiredDefaultMCPServers {
|
||||
entry, ok := cfg.Servers[name]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if inheritsRetiredTransport(entry) {
|
||||
// Re-seed what Zero used to ship underneath the user's fields, so a
|
||||
// partial entry still resolves to the complete server it named when
|
||||
// it was written. Their own values still win — this only fills the
|
||||
// gap the retired default used to cover.
|
||||
entry = mergeMCPServer(retired.shipped, entry, true)
|
||||
cfg.Servers[name] = entry
|
||||
}
|
||||
if !entry.disabledSet || !entry.Disabled {
|
||||
continue
|
||||
}
|
||||
replacement, ok := cfg.Servers[retired.successor]
|
||||
if !ok || replacement.configured {
|
||||
continue
|
||||
}
|
||||
replacement.Disabled = true
|
||||
replacement.disabledSet = true
|
||||
cfg.Servers[retired.successor] = replacement
|
||||
}
|
||||
}
|
||||
|
||||
// inheritsRetiredTransport reports whether entry named no transport of its own
|
||||
// — no type, url, command, or args — and so relied entirely on the retired
|
||||
// default to supply one.
|
||||
//
|
||||
// Anything else is a complete definition the user owns (a self-hosted url, a
|
||||
// stdio command) and must be left alone: grafting the retired default's http
|
||||
// type onto an entry that names a command produces a server that is rejected
|
||||
// outright, turning a working config into a startup failure.
|
||||
func inheritsRetiredTransport(entry MCPServerConfig) bool {
|
||||
return strings.TrimSpace(entry.Type) == "" &&
|
||||
strings.TrimSpace(entry.URL) == "" &&
|
||||
strings.TrimSpace(entry.Command) == "" &&
|
||||
len(entry.Args) == 0
|
||||
}
|
||||
@@ -7,10 +7,10 @@ import (
|
||||
)
|
||||
|
||||
func TestIsDefaultMCPServer(t *testing.T) {
|
||||
if !IsDefaultMCPServer("firecrawl") {
|
||||
t.Fatal("firecrawl should be a built-in default")
|
||||
if !IsDefaultMCPServer("exa") {
|
||||
t.Fatal("exa should be a built-in default")
|
||||
}
|
||||
if IsDefaultMCPServer(" firecrawl ") == false {
|
||||
if IsDefaultMCPServer(" exa ") == false {
|
||||
t.Fatal("IsDefaultMCPServer should trim whitespace")
|
||||
}
|
||||
if IsDefaultMCPServer("not-a-default") {
|
||||
@@ -18,45 +18,45 @@ func TestIsDefaultMCPServer(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMCPSeedsEnabledFirecrawlDefault(t *testing.T) {
|
||||
func TestResolveMCPSeedsEnabledExaDefault(t *testing.T) {
|
||||
cfg, err := ResolveMCP(ResolveOptions{})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
firecrawl, ok := cfg.Servers["firecrawl"]
|
||||
exa, ok := cfg.Servers["exa"]
|
||||
if !ok {
|
||||
t.Fatal("expected the firecrawl default to be seeded with no user config")
|
||||
t.Fatal("expected the exa default to be seeded with no user config")
|
||||
}
|
||||
if firecrawl.Type != "http" || firecrawl.URL != "https://mcp.firecrawl.dev/v2/mcp" {
|
||||
t.Fatalf("unexpected firecrawl default: %#v", firecrawl)
|
||||
if exa.Type != "http" || exa.URL != "https://mcp.exa.ai/mcp" {
|
||||
t.Fatalf("unexpected exa default: %#v", exa)
|
||||
}
|
||||
if firecrawl.Disabled {
|
||||
t.Fatal("the firecrawl default must be enabled out of the box")
|
||||
if exa.Disabled {
|
||||
t.Fatal("the exa default must be enabled out of the box")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMCPUserCanDisableDefault(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"disabled":true}}}}`), 0o600); err != nil {
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"exa":{"disabled":true}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := ResolveMCP(ResolveOptions{UserConfigPath: path})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
if !cfg.Servers["firecrawl"].Disabled {
|
||||
if !cfg.Servers["exa"].Disabled {
|
||||
t.Fatal("a user must be able to disable the default by writing over it")
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsUnconfiguredDefault(t *testing.T) {
|
||||
if !IsUnconfiguredDefault("firecrawl", DefaultMCPServers()["firecrawl"]) {
|
||||
t.Fatal("an untouched firecrawl default should be reported as unconfigured")
|
||||
if !IsUnconfiguredDefault("exa", DefaultMCPServers()["exa"]) {
|
||||
t.Fatal("an untouched exa default should be reported as unconfigured")
|
||||
}
|
||||
if IsUnconfiguredDefault("firecrawl", MCPServerConfig{Type: "http", URL: "http://localhost:3002/mcp"}) {
|
||||
if IsUnconfiguredDefault("exa", MCPServerConfig{Type: "http", URL: "https://example.com/mcp"}) {
|
||||
t.Fatal("a server overriding the default URL is no longer unconfigured")
|
||||
}
|
||||
if IsUnconfiguredDefault("firecrawl", MCPServerConfig{Type: "http", URL: "https://mcp.firecrawl.dev/v2/mcp", Auth: "bearer"}) {
|
||||
if IsUnconfiguredDefault("exa", MCPServerConfig{Type: "http", URL: "https://mcp.exa.ai/mcp", Auth: "bearer"}) {
|
||||
t.Fatal("a server with credentials added is no longer unconfigured")
|
||||
}
|
||||
if IsUnconfiguredDefault("not-a-default", MCPServerConfig{}) {
|
||||
@@ -65,57 +65,173 @@ func TestIsUnconfiguredDefault(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestResolveMCPExplicitReenableIsNotUnconfiguredDefault(t *testing.T) {
|
||||
// `zero mcp enable firecrawl` after a prior disable writes {"disabled":false}
|
||||
// `zero mcp enable exa` after a prior disable writes {"disabled":false}
|
||||
// explicitly. The resolved value is identical to the untouched default (both
|
||||
// enabled, no credentials), but the user DID take an explicit action here, so
|
||||
// IsUnconfiguredDefault must not treat it as untouched (issue #563 review).
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"disabled":false}}}}`), 0o600); err != nil {
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"exa":{"disabled":false}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := ResolveMCP(ResolveOptions{UserConfigPath: path})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
firecrawl := cfg.Servers["firecrawl"]
|
||||
if firecrawl.Disabled {
|
||||
t.Fatalf("explicit re-enable should leave the server enabled: %#v", firecrawl)
|
||||
exa := cfg.Servers["exa"]
|
||||
if exa.Disabled {
|
||||
t.Fatalf("explicit re-enable should leave the server enabled: %#v", exa)
|
||||
}
|
||||
if IsUnconfiguredDefault("firecrawl", firecrawl) {
|
||||
if IsUnconfiguredDefault("exa", exa) {
|
||||
t.Fatal("an explicit enable/disable toggle must count as user-configured, even though the resolved value matches the default")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMCPExplicitRedeclareOfDefaultValuesIsNotUnconfiguredDefault(t *testing.T) {
|
||||
// A user who copies firecrawl's exact default type/url into their config
|
||||
// A user who copies Exa's exact default type/url into their config
|
||||
// (e.g. from an example file, planning to add credentials later) produces a
|
||||
// resolved value byte-identical to DefaultMCPServers()["firecrawl"] — the
|
||||
// resolved value byte-identical to DefaultMCPServers()["exa"] — the
|
||||
// same trap TestResolveMCPExplicitReenableIsNotUnconfiguredDefault covers for
|
||||
// the disabled toggle. IsUnconfiguredDefault must still treat this as
|
||||
// user-configured because the user's JSON declared an entry for it, even
|
||||
// though a plain resolved-value comparison could not tell the difference.
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"type":"http","url":"https://mcp.firecrawl.dev/v2/mcp"}}}}`), 0o600); err != nil {
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"exa":{"type":"http","url":"https://mcp.exa.ai/mcp"}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := ResolveMCP(ResolveOptions{UserConfigPath: path})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
firecrawl := cfg.Servers["firecrawl"]
|
||||
want := DefaultMCPServers()["firecrawl"]
|
||||
if firecrawl.Type != want.Type || firecrawl.URL != want.URL || firecrawl.Disabled != want.Disabled {
|
||||
t.Fatalf("expected the resolved value to match the default's fields exactly: %#v", firecrawl)
|
||||
exa := cfg.Servers["exa"]
|
||||
want := DefaultMCPServers()["exa"]
|
||||
if exa.Type != want.Type || exa.URL != want.URL || exa.Disabled != want.Disabled {
|
||||
t.Fatalf("expected the resolved value to match the default's fields exactly: %#v", exa)
|
||||
}
|
||||
if IsUnconfiguredDefault("firecrawl", firecrawl) {
|
||||
if IsUnconfiguredDefault("exa", exa) {
|
||||
t.Fatal("redeclaring the default's exact values is still an explicit user configuration, not an untouched default")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMCPUserCanOverrideDefaultURLKeepingOtherFields(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
// Point firecrawl at a self-hosted instance; the default's Type must survive.
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"url":"http://localhost:3002/mcp"}}}}`), 0o600); err != nil {
|
||||
// Point Exa at a proxy; the default's Type must survive.
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"exa":{"url":"https://example.com/mcp"}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := ResolveMCP(ResolveOptions{UserConfigPath: path})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
exa := cfg.Servers["exa"]
|
||||
if exa.URL != "https://example.com/mcp" {
|
||||
t.Fatalf("user override of the default URL did not apply: %#v", exa)
|
||||
}
|
||||
if exa.Type != "http" {
|
||||
t.Fatalf("override should keep the default's other fields (type), got %#v", exa)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMCPCarriesLegacyDefaultDisableToSuccessor(t *testing.T) {
|
||||
// Upgrade path: the user disabled the firecrawl default Zero used to ship.
|
||||
// Swapping the default to exa must not re-open an outbound connection they
|
||||
// explicitly switched off.
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"disabled":true}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := ResolveMCP(ResolveOptions{UserConfigPath: path})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
if !cfg.Servers["exa"].Disabled {
|
||||
t.Fatalf("a prior disable of the retired default must carry to its replacement: %#v", cfg.Servers["exa"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMCPExplicitSuccessorBeatsLegacyDefaultDisable(t *testing.T) {
|
||||
// The user disabled the old default but has since explicitly configured exa.
|
||||
// Their newer, explicit choice wins — the migration must not override it.
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"disabled":true},"exa":{"disabled":false}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := ResolveMCP(ResolveOptions{UserConfigPath: path})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
exa := cfg.Servers["exa"]
|
||||
if exa.Disabled {
|
||||
t.Fatalf("an explicit exa entry must survive the legacy-disable migration: %#v", exa)
|
||||
}
|
||||
if IsUnconfiguredDefault("exa", exa) {
|
||||
t.Fatal("an explicitly configured exa is not an untouched default")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMCPLegacyDefaultLeftEnabledDoesNotDisableSuccessor(t *testing.T) {
|
||||
// A user who configured firecrawl without disabling it (e.g. added a header)
|
||||
// never opted out of the default search server, so exa stays enabled.
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"headers":{"Authorization":"Bearer k"}}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := ResolveMCP(ResolveOptions{UserConfigPath: path})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
if cfg.Servers["exa"].Disabled {
|
||||
t.Fatal("only an explicit disable of the retired default should carry over")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMCPLegacyDefaultDisableIsLiftableByOverride(t *testing.T) {
|
||||
// The carried-over disable is a user-level decision, not a permanent one:
|
||||
// `zero mcp enable exa` merges through the CLI override scope, which is the
|
||||
// one layer allowed to re-enable.
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"disabled":true}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
enabled := MCPServerConfig{Disabled: false, disabledSet: true}
|
||||
cfg, err := ResolveMCP(ResolveOptions{
|
||||
UserConfigPath: path,
|
||||
Overrides: Overrides{MCP: MCPConfig{Servers: map[string]MCPServerConfig{"exa": enabled}}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
if cfg.Servers["exa"].Disabled {
|
||||
t.Fatalf("an explicit enable must lift the carried-over disable: %#v", cfg.Servers["exa"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMCPLegacyDisableSurvivesProjectReenable(t *testing.T) {
|
||||
// The project layer is lower-trust and must not lift the carried-over
|
||||
// user-level disable — the same guard a direct `zero mcp disable exa` gets.
|
||||
dir := t.TempDir()
|
||||
userPath := filepath.Join(dir, "config.json")
|
||||
if err := os.WriteFile(userPath, []byte(`{"mcp":{"servers":{"firecrawl":{"disabled":true}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
projectPath := filepath.Join(dir, "project.json")
|
||||
if err := os.WriteFile(projectPath, []byte(`{"mcp":{"servers":{"exa":{"disabled":false}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := ResolveMCP(ResolveOptions{UserConfigPath: userPath, ProjectConfigPath: projectPath})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
if !cfg.Servers["exa"].Disabled {
|
||||
t.Fatal("project config must not re-enable a server the user disabled")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMCPKeepsRetiredDefaultTransportForPartialEntry(t *testing.T) {
|
||||
// A user who added an API key to the firecrawl default wrote only the header
|
||||
// — the seeded default supplied the http type and URL. Retiring the default
|
||||
// must not strip those out from under them and leave an unusable server.
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"headers":{"Authorization":"Bearer k"}}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := ResolveMCP(ResolveOptions{UserConfigPath: path})
|
||||
@@ -123,10 +239,55 @@ func TestResolveMCPUserCanOverrideDefaultURLKeepingOtherFields(t *testing.T) {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
firecrawl := cfg.Servers["firecrawl"]
|
||||
if firecrawl.URL != "http://localhost:3002/mcp" {
|
||||
t.Fatalf("user override of the default URL did not apply: %#v", firecrawl)
|
||||
if firecrawl.Type != "http" || firecrawl.URL != "https://mcp.firecrawl.dev/v2/mcp" {
|
||||
t.Fatalf("a partial entry must keep the retired default's transport: %#v", firecrawl)
|
||||
}
|
||||
if firecrawl.Type != "http" {
|
||||
t.Fatalf("override should keep the default's other fields (type), got %#v", firecrawl)
|
||||
if firecrawl.Headers["Authorization"] != "Bearer k" {
|
||||
t.Fatalf("the user's own field must still win: %#v", firecrawl)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMCPDoesNotGraftRetiredDefaultOntoOwnTransport(t *testing.T) {
|
||||
// A self-hosted firecrawl over stdio is a complete definition the user owns.
|
||||
// Re-seeding the retired http default under it would produce an http server
|
||||
// carrying a command, which NormalizeConfig rejects outright.
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"command":"firecrawl-mcp"}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := ResolveMCP(ResolveOptions{UserConfigPath: path})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
firecrawl := cfg.Servers["firecrawl"]
|
||||
if firecrawl.Type != "" || firecrawl.URL != "" {
|
||||
t.Fatalf("an entry naming its own transport must be left alone: %#v", firecrawl)
|
||||
}
|
||||
if firecrawl.Command != "firecrawl-mcp" {
|
||||
t.Fatalf("the user's command must survive: %#v", firecrawl)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveMCPDisabledRetiredEntryStaysReEnableable(t *testing.T) {
|
||||
// The disable carries to exa, and the retired entry itself keeps a usable
|
||||
// transport so a later `zero mcp enable firecrawl` does not resolve to a
|
||||
// server with no type, url, or command.
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"disabled":true}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := ResolveMCP(ResolveOptions{UserConfigPath: path})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
firecrawl := cfg.Servers["firecrawl"]
|
||||
if !firecrawl.Disabled {
|
||||
t.Fatalf("the retired entry must stay disabled: %#v", firecrawl)
|
||||
}
|
||||
if firecrawl.Type != "http" || firecrawl.URL != "https://mcp.firecrawl.dev/v2/mcp" {
|
||||
t.Fatalf("a disabled retired entry should still carry a usable transport: %#v", firecrawl)
|
||||
}
|
||||
if !cfg.Servers["exa"].Disabled {
|
||||
t.Fatal("the disable must still carry to the successor")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,252 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSetPetPreservesUnrelatedConfig(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
original := FileConfig{
|
||||
ActiveProvider: "test",
|
||||
Providers: []ProviderProfile{{Name: "test", Model: "example"}},
|
||||
Preferences: PreferencesConfig{Theme: "dracula"},
|
||||
}
|
||||
data, err := json.Marshal(original)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(path, data, 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
updated, err := SetPet(path, "boba")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if updated.Preferences.Pet != "boba" || updated.Preferences.Theme != "dracula" || updated.ActiveProvider != "test" || len(updated.Providers) != 1 {
|
||||
t.Fatalf("SetPet lost config fields: %#v", updated)
|
||||
}
|
||||
reloadedData, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var reloaded FileConfig
|
||||
if err := json.Unmarshal(reloadedData, &reloaded); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if reloaded.Preferences.Pet != "boba" || reloaded.Preferences.Theme != "dracula" || reloaded.ActiveProvider != "test" || len(reloaded.Providers) != 1 {
|
||||
t.Fatalf("persisted config lost fields: %#v", reloaded)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetPetRejectsInvalidInputs(t *testing.T) {
|
||||
if _, err := SetPet(" ", "boba"); err == nil {
|
||||
t.Fatal("SetPet accepted a blank config path")
|
||||
}
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"preferences":`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := SetPet(path, "boba"); err == nil || !strings.Contains(err.Error(), "invalid config JSON") {
|
||||
t.Fatalf("SetPet malformed JSON error = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetPetCreatesNestedConfigAndClearsEmptyPreferences(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "nested", "zero", "config.json")
|
||||
if _, err := SetPet(path, " boba "); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var root map[string]json.RawMessage
|
||||
if err := json.Unmarshal(data, &root); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var preferences struct {
|
||||
Pet string `json:"pet"`
|
||||
}
|
||||
if err := json.Unmarshal(root["preferences"], &preferences); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if preferences.Pet != "boba" {
|
||||
t.Fatalf("persisted pet = %q, want %q", preferences.Pet, "boba")
|
||||
}
|
||||
root["future"] = json.RawMessage(`{"enabled":true}`)
|
||||
encoded, err := json.Marshal(root)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(path, encoded, 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := SetPet(path, " "); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
data, err = os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
root = nil
|
||||
if err := json.Unmarshal(data, &root); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, ok := root["preferences"]; ok {
|
||||
t.Fatalf("empty preferences were retained: %s", data)
|
||||
}
|
||||
var future struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
if err := json.Unmarshal(root["future"], &future); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !future.Enabled {
|
||||
t.Fatalf("unrelated future member was lost: %s", data)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetPetPreservesUnknownJSONMembers(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
original := []byte(`{
|
||||
"activeProvider": "test",
|
||||
"futureTopLevel": {"enabled": true},
|
||||
"preferences": {
|
||||
"theme": "dracula",
|
||||
"futurePreference": {"mode": "extended"}
|
||||
}
|
||||
}`)
|
||||
if err := os.WriteFile(path, original, 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := SetPet(path, "boba"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var root map[string]json.RawMessage
|
||||
if err := json.Unmarshal(data, &root); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var futureTopLevel struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
}
|
||||
if err := json.Unmarshal(root["futureTopLevel"], &futureTopLevel); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !futureTopLevel.Enabled {
|
||||
t.Fatalf("unknown top-level member = %s", root["futureTopLevel"])
|
||||
}
|
||||
var preferences map[string]json.RawMessage
|
||||
if err := json.Unmarshal(root["preferences"], &preferences); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var futurePreference struct {
|
||||
Mode string `json:"mode"`
|
||||
}
|
||||
if err := json.Unmarshal(preferences["futurePreference"], &futurePreference); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if futurePreference.Mode != "extended" {
|
||||
t.Fatalf("unknown preference member = %s", preferences["futurePreference"])
|
||||
}
|
||||
if got := string(preferences["pet"]); got != `"boba"` {
|
||||
t.Fatalf("pet preference = %s, want boba", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetPetPreservesExistingJSONLayoutAndMemberOrder(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
original := []byte("{\n\t\"zFuture\": {\"enabled\":true},\n\t\"preferences\": {\n\t\t\"theme\" : \"dracula\",\n\t\t\"pet\" : \"old\",\n\t\t\"futurePreference\": {\"mode\":\"extended\"}\n\t},\n\t\"activeProvider\": \"test\"\n}\n")
|
||||
if err := os.WriteFile(path, original, 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := SetPet(path, "boba"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := strings.Replace(string(original), `"old"`, `"boba"`, 1)
|
||||
if string(got) != want {
|
||||
t.Fatalf("SetPet reformatted or reordered config:\n--- got ---\n%s--- want ---\n%s", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetPetMinimallyEditsExistingJSON(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
original string
|
||||
pet string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "add pet to existing preferences",
|
||||
original: "{\n \"zFuture\": {\"text\":\"comma, brace } and \\\"quote\\\"\"},\n \"preferences\": {\n \"theme\": \"dracula\"\n },\n \"activeProvider\": \"test\"\n}\n",
|
||||
pet: "boba",
|
||||
want: "{\n \"zFuture\": {\"text\":\"comma, brace } and \\\"quote\\\"\"},\n \"preferences\": {\n \"theme\": \"dracula\",\n \"pet\": \"boba\"\n },\n \"activeProvider\": \"test\"\n}\n",
|
||||
},
|
||||
{
|
||||
name: "add preferences without reordering root",
|
||||
original: "{\n \"zFuture\": true,\n \"activeProvider\": \"test\"\n}\n",
|
||||
pet: "boba",
|
||||
want: "{\n \"zFuture\": true,\n \"activeProvider\": \"test\",\n \"preferences\": {\"pet\":\"boba\"}\n}\n",
|
||||
},
|
||||
{
|
||||
name: "remove pet but retain preferences",
|
||||
original: "{\n \"preferences\": {\n \"theme\": \"dracula\",\n \"pet\": \"boba\",\n \"future\": true\n },\n \"zFuture\": true\n}\n",
|
||||
pet: "",
|
||||
want: "{\n \"preferences\": {\n \"theme\": \"dracula\",\n \"future\": true\n },\n \"zFuture\": true\n}\n",
|
||||
},
|
||||
{
|
||||
name: "remove empty preferences member",
|
||||
original: "{\n \"zFuture\": true,\n \"preferences\": {\n \"pet\": \"boba\"\n },\n \"activeProvider\": \"test\"\n}\n",
|
||||
pet: "",
|
||||
want: "{\n \"zFuture\": true,\n \"activeProvider\": \"test\"\n}\n",
|
||||
},
|
||||
{
|
||||
name: "do not expose an older duplicate preference",
|
||||
original: "{\"preferences\":{\"pet\":\"cat\"},\"preferences\":{\"pet\":\"dog\"}}\n",
|
||||
pet: "",
|
||||
want: "{\"preferences\":{\"pet\":\"cat\"},\"preferences\":{}}\n",
|
||||
},
|
||||
{
|
||||
name: "do not expose an older duplicate behind null",
|
||||
original: "{\"preferences\":{\"pet\":\"cat\"},\"preferences\":null}\n",
|
||||
pet: "",
|
||||
want: "{\"preferences\":{\"pet\":\"cat\"},\"preferences\":null}\n",
|
||||
},
|
||||
{
|
||||
name: "preserve multiline empty preferences layout",
|
||||
original: "{\n \"preferences\": {\n },\n \"activeProvider\": \"test\"\n}\n",
|
||||
pet: "boba",
|
||||
want: "{\n \"preferences\": {\n \"pet\": \"boba\"\n },\n \"activeProvider\": \"test\"\n}\n",
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(test.original), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := SetPet(path, test.pet); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(got) != test.want {
|
||||
t.Fatalf("SetPet output:\n--- got ---\n%s--- want ---\n%s", got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -34,7 +34,7 @@ func TestResolveMCPExcludeProjectDropsProjectServers(t *testing.T) {
|
||||
if _, ok := cfg.Servers["user-srv"]; !ok {
|
||||
t.Fatal("the user server must always be present")
|
||||
}
|
||||
if _, ok := cfg.Servers["firecrawl"]; !ok {
|
||||
if _, ok := cfg.Servers["exa"]; !ok {
|
||||
t.Fatal("the built-in default must always be present")
|
||||
}
|
||||
})
|
||||
@@ -54,7 +54,7 @@ func TestResolveMCPExcludeProjectDropsProjectServers(t *testing.T) {
|
||||
if _, ok := cfg.Servers["user-srv"]; !ok {
|
||||
t.Fatal("the user server must survive when the project layer is dropped")
|
||||
}
|
||||
if _, ok := cfg.Servers["firecrawl"]; !ok {
|
||||
if _, ok := cfg.Servers["exa"]; !ok {
|
||||
t.Fatal("the built-in default must survive when the project layer is dropped")
|
||||
}
|
||||
})
|
||||
|
||||
@@ -180,7 +180,7 @@ func Resolve(options ResolveOptions) (ResolvedConfig, error) {
|
||||
}
|
||||
|
||||
func ResolveMCP(options ResolveOptions) (MCPConfig, error) {
|
||||
// Seed Zero's built-in default MCP servers (e.g. keyless Firecrawl for free,
|
||||
// Seed Zero's built-in default MCP servers (e.g. keyless Exa for free,
|
||||
// no-setup web search/scrape) BEFORE merging user/project config, so the user
|
||||
// can override any field or disable a default by writing over it.
|
||||
cfg := FileConfig{MCP: MCPConfig{Servers: DefaultMCPServers()}}
|
||||
@@ -194,6 +194,10 @@ func ResolveMCP(options ResolveOptions) (MCPConfig, error) {
|
||||
// server the user disabled (a user-level disable is sticky, but the user
|
||||
// scope itself may lift it).
|
||||
mergeMCPConfig(&cfg.MCP, fileConfig.MCP, true)
|
||||
// Reconcile config written against a default Zero has since retired,
|
||||
// here while the user layer is the only one merged — so a carried-over
|
||||
// disable counts as a user-level decision the project layer cannot lift.
|
||||
migrateRetiredDefaultMCPServers(&cfg.MCP)
|
||||
}
|
||||
// Drop the project layer when the workspace is untrusted, so a cloned repo's
|
||||
// ./.zero/config.json cannot register (and spawn) MCP servers. Fail-closed:
|
||||
@@ -273,6 +277,9 @@ func mergeConfig(dst *FileConfig, src FileConfig) {
|
||||
if strings.TrimSpace(src.Preferences.Theme) != "" {
|
||||
dst.Preferences.Theme = strings.TrimSpace(src.Preferences.Theme)
|
||||
}
|
||||
if strings.TrimSpace(src.Preferences.Pet) != "" {
|
||||
dst.Preferences.Pet = strings.TrimSpace(src.Preferences.Pet)
|
||||
}
|
||||
mergeLocalControlConfig(&dst.LocalControl, src.LocalControl)
|
||||
mergeKeyBindings(&dst.KeyBindings, src.KeyBindings)
|
||||
mergeSTTConfig(&dst.STT, src.STT)
|
||||
|
||||
@@ -168,6 +168,22 @@ func TestResolveLoadsRecentModelsFromUserConfigOnly(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveLoadsPetFromUserConfigOnly(t *testing.T) {
|
||||
userPath := writeConfig(t, `{"preferences":{"pet":" boba "}}`)
|
||||
projectPath := writeConfig(t, `{"preferences":{"pet":"project-pet"}}`)
|
||||
resolved, err := Resolve(ResolveOptions{
|
||||
UserConfigPath: userPath,
|
||||
ProjectConfigPath: projectPath,
|
||||
Env: map[string]string{},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resolved.Preferences.Pet != "boba" {
|
||||
t.Fatalf("Preferences.Pet = %q, want user preference boba", resolved.Preferences.Pet)
|
||||
}
|
||||
}
|
||||
|
||||
// recentModels must never be merged from project config (same posture as
|
||||
// favoriteModels): a project setting it with no corresponding user setting
|
||||
// should leave the resolved value empty, not pick up the project's list.
|
||||
|
||||
@@ -112,10 +112,19 @@ type PreferencesConfig struct {
|
||||
// name (e.g. "dracula"). Applied at startup below the --theme flag and
|
||||
// ZERO_THEME, so a /theme choice survives restart. Empty = unset (defaults auto).
|
||||
Theme string `json:"theme,omitempty"`
|
||||
// Pet is the terminal companion selected through /pets. Empty leaves pets
|
||||
// off until the user chooses one; "disabled" is the explicit off state.
|
||||
Pet string `json:"pet,omitempty"`
|
||||
// Recaps is a tri-state: nil (unset) defaults to ON; an explicit false means
|
||||
// the user turned idle recaps off. A *bool is its own tri-state, so no
|
||||
// custom unmarshal is needed (unlike ToolsConfig.DeferThreshold's int).
|
||||
Recaps *bool `json:"recaps,omitempty"`
|
||||
// CompactionModel routes compaction summarization calls to this model
|
||||
// instead of the session's main model (summaries at main-model prices are
|
||||
// the most expensive recurring event in long runs). Empty = automatic: a
|
||||
// curated cheap model on official endpoints, the main model elsewhere.
|
||||
// "main" forces the main model. ZERO_COMPACTION_MODEL overrides.
|
||||
CompactionModel string `json:"compactionModel,omitempty"`
|
||||
}
|
||||
|
||||
// RecentModelEntry is one provider-qualified model selection recorded in
|
||||
@@ -148,7 +157,8 @@ type KeyBindingsConfig struct {
|
||||
ToggleMouse KeyBindingDef `json:"toggleMouse,omitempty"`
|
||||
// CycleReasoning cycles through reasoning effort levels (default: ctrl+t).
|
||||
CycleReasoning KeyBindingDef `json:"cycleReasoning,omitempty"`
|
||||
// TogglePlan toggles the plan panel expansion (default: ctrl+p).
|
||||
// TogglePlan is retained for configuration compatibility. Plan updates render
|
||||
// in the transcript and no longer have a persistent panel to toggle.
|
||||
TogglePlan KeyBindingDef `json:"togglePlan,omitempty"`
|
||||
// ToggleSidebar toggles the right context sidebar (default: ctrl+b).
|
||||
ToggleSidebar KeyBindingDef `json:"toggleSidebar,omitempty"`
|
||||
@@ -461,7 +471,7 @@ type MCPServerConfig struct {
|
||||
// fields it set or what values they hold. A built-in default seeded by
|
||||
// DefaultMCPServers() is never unmarshaled from JSON, so it starts false;
|
||||
// any explicit entry in the user/project file — even one that happens to
|
||||
// repeat a default's exact field values (e.g. re-declaring firecrawl's
|
||||
// repeat a default's exact field values (e.g. re-declaring Exa's
|
||||
// default URL) — sets it true. IsUnconfiguredDefault checks this alongside
|
||||
// a resolved-value comparison, so redeclaring default values verbatim still
|
||||
// counts as user-configured.
|
||||
|
||||
@@ -591,6 +591,35 @@ func SetTheme(path string, theme string) (FileConfig, error) {
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
// SetPet persists only the terminal-pet preference while preserving every
|
||||
// unrelated user setting through the config writer's atomic replace path.
|
||||
func SetPet(path string, pet string) (FileConfig, error) {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return FileConfig{}, fmt.Errorf("config path is required")
|
||||
}
|
||||
cfg := FileConfig{}
|
||||
data := []byte("{}")
|
||||
if existing, err := os.ReadFile(path); err == nil {
|
||||
if err := json.Unmarshal(existing, &cfg); err != nil {
|
||||
return FileConfig{}, fmt.Errorf("invalid config JSON %s: %w", path, err)
|
||||
}
|
||||
data = existing
|
||||
} else if !os.IsNotExist(err) {
|
||||
return FileConfig{}, fmt.Errorf("read config %s: %w", path, err)
|
||||
}
|
||||
pet = strings.TrimSpace(pet)
|
||||
cfg.Preferences.Pet = pet
|
||||
data, err := setPetPreferenceJSON(data, pet)
|
||||
if err != nil {
|
||||
return FileConfig{}, fmt.Errorf("invalid config JSON %s: %w", path, err)
|
||||
}
|
||||
if err := writeConfigData(path, data); err != nil {
|
||||
return FileConfig{}, err
|
||||
}
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
// SetSTTModel persists the dictation model and its provider, mirroring
|
||||
// SetTheme (read-modify-atomic-write). provider must be one of the known STT
|
||||
// provider kinds; a local provider stores the model as stt.localModelPath,
|
||||
@@ -736,17 +765,23 @@ func NormalizeRecentModels(entries []RecentModelEntry) []RecentModelEntry {
|
||||
}
|
||||
|
||||
func writeConfigFile(path string, cfg FileConfig) error {
|
||||
data, err := json.MarshalIndent(cfg, "", " ")
|
||||
if err != nil {
|
||||
return fmt.Errorf("encode config JSON: %w", err)
|
||||
}
|
||||
return writeConfigData(path, data)
|
||||
}
|
||||
|
||||
func writeConfigData(path string, data []byte) error {
|
||||
dir := filepath.Dir(path)
|
||||
if dir != "." && dir != "" {
|
||||
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||
return fmt.Errorf("create config directory %s: %w", dir, err)
|
||||
}
|
||||
}
|
||||
data, err := json.MarshalIndent(cfg, "", " ")
|
||||
if err != nil {
|
||||
return fmt.Errorf("encode config JSON: %w", err)
|
||||
if len(data) == 0 || data[len(data)-1] != '\n' {
|
||||
data = append(data, '\n')
|
||||
}
|
||||
data = append(data, '\n')
|
||||
// Write-to-temp + rename: an in-place write interrupted mid-way (crash,
|
||||
// disk full) would leave the user's only config truncated or corrupt.
|
||||
tmp, err := os.CreateTemp(dir, ".zero-config-*.tmp")
|
||||
|
||||
@@ -0,0 +1,246 @@
|
||||
package credstore
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func fileStore(t *testing.T, dir string) *Store {
|
||||
t.Helper()
|
||||
store, err := New(Options{Dir: dir, Storage: "file"})
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
return store
|
||||
}
|
||||
|
||||
// THE REPRODUCTION, kept as the regression. Before the lock this reliably kept
|
||||
// 1 of 100: every writer read the same map, added its own provider, and the
|
||||
// last rename published a file missing all the others.
|
||||
func TestConcurrentSetKeepsEveryKey(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
store := fileStore(t, dir)
|
||||
|
||||
const n = 100
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, n)
|
||||
for i := 0; i < n; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
if err := store.Set(fmt.Sprintf("p%03d", i), fmt.Sprintf("k%03d", i)); err != nil {
|
||||
errs <- err
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
t.Fatalf("Set: %v", err)
|
||||
}
|
||||
|
||||
got, err := store.Providers()
|
||||
if err != nil {
|
||||
t.Fatalf("Providers: %v", err)
|
||||
}
|
||||
if len(got) != n {
|
||||
t.Fatalf("kept %d of %d concurrent writes", len(got), n)
|
||||
}
|
||||
// ...and every value is the one its own writer wrote, not a neighbour's.
|
||||
for i := 0; i < n; i++ {
|
||||
key, ok, err := store.Get(fmt.Sprintf("p%03d", i))
|
||||
if err != nil {
|
||||
t.Fatalf("Get: %v", err)
|
||||
}
|
||||
if !ok || key != fmt.Sprintf("k%03d", i) {
|
||||
t.Fatalf("p%03d = %q (present=%v); a writer's own value must survive", i, key, ok)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// DELETE RACES SET, and an unlocked delete loses the other writer's key.
|
||||
//
|
||||
// A first version of this only checked that keys nobody deleted survived — and
|
||||
// every reader sees those, so it passed with the lock removed from Delete. The
|
||||
// case that actually breaks is a Delete whose stale read predates a concurrent
|
||||
// Set: it writes back a map that never contained the new key, and the Set is
|
||||
// gone. So the assertion is on the SETS surviving, not on the bystanders.
|
||||
func TestADeleteCannotClobberAConcurrentSet(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
store := fileStore(t, dir)
|
||||
|
||||
// Throwaway keys for the deleters to remove, written first so a delete
|
||||
// always has something to do.
|
||||
const churn = 40
|
||||
for i := 0; i < churn; i++ {
|
||||
if err := store.Set(fmt.Sprintf("churn%02d", i), "v"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
const adds = 60
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < adds; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
_ = store.Set(fmt.Sprintf("added%02d", i), "v")
|
||||
}(i)
|
||||
}
|
||||
for i := 0; i < churn; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
_, _ = store.Delete(fmt.Sprintf("churn%02d", i))
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
got, err := store.Providers()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
survivors := map[string]bool{}
|
||||
for _, name := range got {
|
||||
survivors[name] = true
|
||||
}
|
||||
missing := []string{}
|
||||
for i := 0; i < adds; i++ {
|
||||
name := fmt.Sprintf("added%02d", i)
|
||||
if !survivors[name] {
|
||||
missing = append(missing, name)
|
||||
}
|
||||
}
|
||||
if len(missing) > 0 {
|
||||
t.Fatalf("%d of %d concurrent Set calls were clobbered by a racing Delete: %v",
|
||||
len(missing), adds, missing)
|
||||
}
|
||||
// ...and every churn key really was removed, so the deletes were not
|
||||
// no-ops that would make the check above vacuous.
|
||||
for i := 0; i < churn; i++ {
|
||||
if survivors[fmt.Sprintf("churn%02d", i)] {
|
||||
t.Fatalf("churn%02d survived; the deletes did nothing and this test proves nothing", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ACROSS PROCESSES, which is the case that matters: a plan running with
|
||||
// max_workers > 1 spawns child PROCESSES, and an in-process mutex would not
|
||||
// have helped any of them. The race detector cannot see this one either — it is
|
||||
// two OS processes — so it is driven for real.
|
||||
func TestConcurrentSetAcrossProcesses(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("spawns processes")
|
||||
}
|
||||
helper := os.Getenv("ZERO_CREDSTORE_HELPER_DIR")
|
||||
if helper != "" {
|
||||
// Child mode: write our slice of the keys and exit.
|
||||
store, err := New(Options{Dir: helper, Storage: "file"})
|
||||
if err != nil {
|
||||
os.Exit(3)
|
||||
}
|
||||
prefix := os.Getenv("ZERO_CREDSTORE_HELPER_PREFIX")
|
||||
for i := 0; i < 25; i++ {
|
||||
if err := store.Set(fmt.Sprintf("%s%02d", prefix, i), "v"); err != nil {
|
||||
os.Exit(4)
|
||||
}
|
||||
}
|
||||
os.Exit(0)
|
||||
}
|
||||
|
||||
dir := t.TempDir()
|
||||
const children = 4
|
||||
var wg sync.WaitGroup
|
||||
failures := make(chan string, children)
|
||||
for c := 0; c < children; c++ {
|
||||
wg.Add(1)
|
||||
go func(c int) {
|
||||
defer wg.Done()
|
||||
cmd := exec.Command(os.Args[0], "-test.run=TestConcurrentSetAcrossProcesses", "-test.v")
|
||||
cmd.Env = append(os.Environ(),
|
||||
"ZERO_CREDSTORE_HELPER_DIR="+dir,
|
||||
fmt.Sprintf("ZERO_CREDSTORE_HELPER_PREFIX=c%d_", c),
|
||||
)
|
||||
if out, err := cmd.CombinedOutput(); err != nil {
|
||||
failures <- fmt.Sprintf("child %d: %v\n%s", c, err, out)
|
||||
}
|
||||
}(c)
|
||||
}
|
||||
wg.Wait()
|
||||
close(failures)
|
||||
for failure := range failures {
|
||||
t.Fatal(failure)
|
||||
}
|
||||
|
||||
store := fileStore(t, dir)
|
||||
got, err := store.Providers()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got) != children*25 {
|
||||
t.Fatalf("kept %d of %d keys written by %d concurrent processes", len(got), children*25, children)
|
||||
}
|
||||
}
|
||||
|
||||
// The lock lives BESIDE the data file, never on it. write publishes by rename,
|
||||
// so a lock taken on the data file would be attached to an inode the next
|
||||
// writer has already replaced — every writer would appear to hold it.
|
||||
func TestTheLockIsNotTheDataFile(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
store := fileStore(t, dir)
|
||||
if store.lockPath() == store.file {
|
||||
t.Fatal("the lock is the data file; rename would carry it away")
|
||||
}
|
||||
if !strings.HasPrefix(filepath.Base(store.lockPath()), filepath.Base(store.file)) {
|
||||
t.Fatalf("lock %q is not beside the data file %q", store.lockPath(), store.file)
|
||||
}
|
||||
if err := store.Set("p", "k"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := os.Stat(store.lockPath()); err != nil {
|
||||
t.Fatalf("the lock file was not created: %v", err)
|
||||
}
|
||||
// ...and it survives a write, since the rename replaces only the data file.
|
||||
if err := store.Set("q", "k"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := os.Stat(store.lockPath()); err != nil {
|
||||
t.Fatalf("the lock file did not survive a write: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// The lock is a security boundary, so the path where it CANNOT be taken needs
|
||||
// coverage too: an operation that could not lock must report that, never proceed
|
||||
// unserialized and never report success. Replacing the credential directory with
|
||||
// a regular file makes the lock's MkdirAll fail on every platform.
|
||||
func TestOperationsFailWhenTheLockCannotBeAcquired(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
dir := filepath.Join(base, "creds")
|
||||
store := fileStore(t, dir)
|
||||
// Where the store expects its directory, put a regular file.
|
||||
if err := os.WriteFile(dir, []byte("not a directory"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if err := store.Set("openai", "sk-test"); err == nil {
|
||||
t.Error("Set reported success although the lock could not be acquired")
|
||||
}
|
||||
if _, err := store.Delete("openai"); err == nil {
|
||||
t.Error("Delete reported success although the lock could not be acquired")
|
||||
}
|
||||
if _, _, err := store.Get("openai"); err == nil {
|
||||
t.Error("Get reported success although the lock could not be acquired")
|
||||
}
|
||||
if _, err := store.Providers(); err == nil {
|
||||
t.Error("Providers reported success although the lock could not be acquired")
|
||||
}
|
||||
// And nothing may have been written through the failed path.
|
||||
if raw, err := os.ReadFile(dir); err != nil || string(raw) != "not a directory" {
|
||||
t.Errorf("the blocking file was modified: %q (err %v)", raw, err)
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,7 @@ package credstore
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
@@ -125,14 +126,27 @@ func (s *Store) Backend() string { return s.backend }
|
||||
func (s *Store) Encrypted() bool { return s.backend == "keyring" || s.backend == "encrypted-file" }
|
||||
|
||||
// Set stores (or overwrites) the API key for provider.
|
||||
func (s *Store) Set(provider, key string) error {
|
||||
func (s *Store) Set(provider, key string) (err error) {
|
||||
provider = normalizeProvider(provider)
|
||||
if provider == "" {
|
||||
return fmt.Errorf("credstore: provider is required")
|
||||
}
|
||||
if s.backend == "keyring" {
|
||||
// The OS keyring serializes its own writes; the file lock guards the
|
||||
// file backends only.
|
||||
return s.kr.Set(keyringService, keyringPrefix+provider, key)
|
||||
}
|
||||
// READ AND WRITE UNDER ONE LOCK. Separately locking each half would still
|
||||
// lose keys: two writers would each read the same map, each add their own
|
||||
// provider, and the second rename would publish a file missing the first.
|
||||
// Measured before this existed: 1 of 100 concurrent Set calls survived.
|
||||
release, err := s.acquireFileLock(true)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Joined, not chosen between: a release failure annotates the result instead
|
||||
// of masking the write error that actually explains what went wrong.
|
||||
defer func() { err = errors.Join(err, release()) }()
|
||||
data, err := s.read()
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -142,7 +156,7 @@ func (s *Store) Set(provider, key string) error {
|
||||
}
|
||||
|
||||
// Get returns the API key for provider and whether one is stored.
|
||||
func (s *Store) Get(provider string) (string, bool, error) {
|
||||
func (s *Store) Get(provider string) (value string, found bool, err error) {
|
||||
provider = normalizeProvider(provider)
|
||||
if provider == "" {
|
||||
return "", false, nil
|
||||
@@ -150,6 +164,15 @@ func (s *Store) Get(provider string) (string, bool, error) {
|
||||
if s.backend == "keyring" {
|
||||
return s.kr.Get(keyringService, keyringPrefix+provider)
|
||||
}
|
||||
// A shared lock, so a read serializes against a writer's publish. On unix
|
||||
// the rename is atomic and this only prevents reading an empty file mid-swap;
|
||||
// on Windows an unsynchronized reader holding the file open would block the
|
||||
// writer's rename outright, so readers must coordinate through the same lock.
|
||||
release, err := s.acquireFileLock(false)
|
||||
if err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
defer func() { err = errors.Join(err, release()) }()
|
||||
data, err := s.read()
|
||||
if err != nil {
|
||||
return "", false, err
|
||||
@@ -159,7 +182,7 @@ func (s *Store) Get(provider string) (string, bool, error) {
|
||||
}
|
||||
|
||||
// Delete removes the API key for provider, reporting whether one existed.
|
||||
func (s *Store) Delete(provider string) (bool, error) {
|
||||
func (s *Store) Delete(provider string) (existed bool, err error) {
|
||||
provider = normalizeProvider(provider)
|
||||
if provider == "" {
|
||||
return false, nil
|
||||
@@ -167,6 +190,13 @@ func (s *Store) Delete(provider string) (bool, error) {
|
||||
if s.backend == "keyring" {
|
||||
return s.kr.Delete(keyringService, keyringPrefix+provider)
|
||||
}
|
||||
// The same lock as Set: a delete is a read-modify-write too, and one racing
|
||||
// a set would otherwise resurrect or erase an unrelated provider.
|
||||
release, err := s.acquireFileLock(true)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
defer func() { err = errors.Join(err, release()) }()
|
||||
data, err := s.read()
|
||||
if err != nil {
|
||||
return false, err
|
||||
@@ -179,17 +209,22 @@ func (s *Store) Delete(provider string) (bool, error) {
|
||||
}
|
||||
|
||||
// Providers lists providers with a stored key (sorted), for display/audit.
|
||||
func (s *Store) Providers() ([]string, error) {
|
||||
func (s *Store) Providers() (names []string, err error) {
|
||||
if s.backend == "keyring" {
|
||||
// The OS keyring has no enumerate-by-prefix in our minimal surface; callers
|
||||
// that need a list pass known provider names to Get. Return empty here.
|
||||
return nil, nil
|
||||
}
|
||||
release, err := s.acquireFileLock(false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { err = errors.Join(err, release()) }()
|
||||
data, err := s.read()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
names := make([]string, 0, len(data))
|
||||
names = make([]string, 0, len(data))
|
||||
for name := range data {
|
||||
names = append(names, name)
|
||||
}
|
||||
@@ -264,6 +299,15 @@ func (s *Store) write(data map[string]string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// lockPath is the advisory lock guarding a read-modify-write of the credential
|
||||
// file. Beside the data file, never the data file itself — write publishes by
|
||||
// rename and would carry the lock away with the old inode.
|
||||
func (s *Store) lockPath() string { return s.file + ".lock" }
|
||||
|
||||
// filepathDir is filepath.Dir, named locally so the two platform lock files can
|
||||
// share it without either importing path/filepath for one call.
|
||||
func filepathDir(path string) string { return filepath.Dir(path) }
|
||||
|
||||
func normalizeProvider(provider string) string {
|
||||
return strings.ToLower(strings.TrimSpace(provider))
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
//go:build !windows
|
||||
|
||||
package credstore
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
"golang.org/x/sys/unix"
|
||||
)
|
||||
|
||||
// acquireFileLock takes an exclusive advisory lock (flock) so a read-modify-write
|
||||
// of the credential file is serialized against every other one — across
|
||||
// processes AND across goroutines, since flock is held per open file
|
||||
// description and two opens in one process contend exactly as two processes do.
|
||||
//
|
||||
// THE LOCK FILE IS SEPARATE FROM THE DATA FILE, and that is not tidiness. write
|
||||
// publishes by os.Rename, which replaces the inode; a lock taken on the data
|
||||
// file would be attached to an inode that the next writer has already replaced,
|
||||
// so every writer would appear to hold it. The lock lives on a file nothing
|
||||
// renames.
|
||||
func (s *Store) acquireFileLock(exclusive bool) (func() error, error) {
|
||||
path := s.lockPath()
|
||||
if err := os.MkdirAll(filepathDir(path), 0o700); err != nil {
|
||||
return nil, fmt.Errorf("credstore: lock dir: %w", err)
|
||||
}
|
||||
file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0o600)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("credstore: open lock: %w", err)
|
||||
}
|
||||
// Writers take LOCK_EX; readers take LOCK_SH so they run concurrently with
|
||||
// each other but still serialize against a writer's publish (see Get).
|
||||
how := unix.LOCK_SH
|
||||
if exclusive {
|
||||
how = unix.LOCK_EX
|
||||
}
|
||||
if err := unix.Flock(int(file.Fd()), how); err != nil {
|
||||
_ = file.Close()
|
||||
return nil, fmt.Errorf("credstore: lock: %w", err)
|
||||
}
|
||||
return func() error {
|
||||
// Close alone drops the flock, so the explicit unlock is belt-and-braces —
|
||||
// but its failure is reported rather than swallowed, because a cleanup that
|
||||
// did not complete must not be indistinguishable from one that did. The
|
||||
// callers join this into their result, so it annotates a successful write
|
||||
// instead of masking it.
|
||||
unlockErr := unix.Flock(int(file.Fd()), unix.LOCK_UN)
|
||||
closeErr := file.Close()
|
||||
if err := errors.Join(unlockErr, closeErr); err != nil {
|
||||
return fmt.Errorf("credstore: release lock: %w", err)
|
||||
}
|
||||
return nil
|
||||
}, nil
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
//go:build windows
|
||||
|
||||
package credstore
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
|
||||
"golang.org/x/sys/windows"
|
||||
)
|
||||
|
||||
// acquireFileLock takes an exclusive OS lock (LockFileEx) so a read-modify-write
|
||||
// of the credential file is serialized, matching the flock behaviour on unix.
|
||||
//
|
||||
// The lock file is SEPARATE from the data file for the same reason: write
|
||||
// publishes by rename, and a lock on the renamed file would be attached to
|
||||
// something the next writer has already replaced.
|
||||
func (s *Store) acquireFileLock(exclusive bool) (func() error, error) {
|
||||
path := s.lockPath()
|
||||
if err := os.MkdirAll(filepathDir(path), 0o700); err != nil {
|
||||
return nil, fmt.Errorf("credstore: lock dir: %w", err)
|
||||
}
|
||||
file, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0o600)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("credstore: open lock: %w", err)
|
||||
}
|
||||
handle := windows.Handle(file.Fd())
|
||||
overlapped := new(windows.Overlapped)
|
||||
// A fixed 1-byte region, blocking (no LOCKFILE_FAIL_IMMEDIATELY) so a
|
||||
// waiter queues rather than failing the operation. Writers pass
|
||||
// LOCKFILE_EXCLUSIVE_LOCK; readers omit it (a shared lock) so they still
|
||||
// serialize against a writer's publish — which matters more on Windows,
|
||||
// where an unsynchronized reader holding the file open blocks the rename.
|
||||
flags := uint32(0)
|
||||
if exclusive {
|
||||
flags = windows.LOCKFILE_EXCLUSIVE_LOCK
|
||||
}
|
||||
if err := windows.LockFileEx(handle, flags, 0, 1, 0, overlapped); err != nil {
|
||||
_ = file.Close()
|
||||
return nil, fmt.Errorf("credstore: lock: %w", err)
|
||||
}
|
||||
return func() error {
|
||||
// Reported rather than swallowed, for the same reason as the unix side: a
|
||||
// cleanup that did not complete must not look identical to one that did.
|
||||
// It matters more here — on Windows the handle staying open is what blocks
|
||||
// the next writer's rename, so a failed release has a visible consequence.
|
||||
unlockErr := windows.UnlockFileEx(handle, 0, 1, 0, overlapped)
|
||||
closeErr := file.Close()
|
||||
if err := errors.Join(unlockErr, closeErr); err != nil {
|
||||
return fmt.Errorf("credstore: release lock: %w", err)
|
||||
}
|
||||
return nil
|
||||
}, nil
|
||||
}
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/providers/providerio"
|
||||
"github.com/coder/websocket"
|
||||
)
|
||||
|
||||
@@ -60,7 +61,15 @@ func (d *deepgramTranscriber) StreamTranscribe(ctx context.Context, chunks <-cha
|
||||
HTTPHeader: http.Header{"Authorization": {"Token " + d.cfg.APIKey}},
|
||||
})
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("connecting to Deepgram: %w", err)
|
||||
// Preserve the context.Canceled sentinel (it carries no key) so an
|
||||
// Esc-abort during dial still matches the UI's errors.Is check. Key
|
||||
// on ctx.Err() itself rather than unwrapping err: a cancellation can
|
||||
// surface as a plain transport error (e.g. "closed network
|
||||
// connection") rather than context.Canceled directly.
|
||||
if ctx.Err() != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
return "", fmt.Errorf("connecting to Deepgram: %s", providerio.Redact(err.Error(), d.cfg.APIKey))
|
||||
}
|
||||
defer conn.CloseNow()
|
||||
|
||||
@@ -68,15 +77,26 @@ func (d *deepgramTranscriber) StreamTranscribe(ctx context.Context, chunks <-cha
|
||||
// so no conversion (unlike sherpa-onnx, which wants float32).
|
||||
writeErrCh := make(chan error, 1)
|
||||
go func() {
|
||||
for chunk := range chunks {
|
||||
if err := conn.Write(ctx, websocket.MessageBinary, chunk); err != nil {
|
||||
writeErrCh <- err
|
||||
// Select on ctx.Done() while waiting for the next chunk so a cancelled
|
||||
// session does not leave this goroutine blocked on an open idle channel.
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
writeErrCh <- ctx.Err()
|
||||
return
|
||||
case chunk, ok := <-chunks:
|
||||
if !ok {
|
||||
// CloseStream flushes any buffered audio and returns final
|
||||
// results before Deepgram closes the socket.
|
||||
writeErrCh <- conn.Write(ctx, websocket.MessageText, []byte(`{"type":"CloseStream"}`))
|
||||
return
|
||||
}
|
||||
if err := conn.Write(ctx, websocket.MessageBinary, chunk); err != nil {
|
||||
writeErrCh <- err
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
// CloseStream flushes any buffered audio and returns final results before
|
||||
// Deepgram closes the socket.
|
||||
writeErrCh <- conn.Write(ctx, websocket.MessageText, []byte(`{"type":"CloseStream"}`))
|
||||
}()
|
||||
|
||||
// Deepgram results are per-utterance segments, not cumulative: accumulate the
|
||||
@@ -100,7 +120,18 @@ func (d *deepgramTranscriber) StreamTranscribe(ctx context.Context, chunks <-cha
|
||||
}
|
||||
default:
|
||||
}
|
||||
return compose(), fmt.Errorf("Deepgram stream error: %w", err) //nolint:staticcheck // Preserve established user-facing error text.
|
||||
// A user abort cancels the streaming context; the UI matches it
|
||||
// with errors.Is(err, context.Canceled), so return the sentinel
|
||||
// itself (it carries no key) instead of a flat redacted string.
|
||||
// Key on ctx.Err() rather than unwrapping err: the writeErrCh
|
||||
// swap above can replace err with a plain transport error (e.g.
|
||||
// "closed network connection") that doesn't itself unwrap to
|
||||
// context.Canceled even though the cancellation is what caused it.
|
||||
if ctx.Err() != nil {
|
||||
return compose(), ctx.Err()
|
||||
}
|
||||
//nolint:staticcheck // Preserve established user-facing error text.
|
||||
return compose(), fmt.Errorf("Deepgram stream error: %s", providerio.Redact(err.Error(), d.cfg.APIKey))
|
||||
}
|
||||
if typ != websocket.MessageText {
|
||||
continue
|
||||
|
||||
@@ -0,0 +1,278 @@
|
||||
package dictation
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/coder/websocket"
|
||||
)
|
||||
|
||||
// Non-cancellation dial failure: handshake returns 101 with an Upgrade header
|
||||
// equal to the API key so the dial error embeds the key and redaction must strip it.
|
||||
func TestDeepgramDialFailureRedactsKey(t *testing.T) {
|
||||
const key = "sk-test-dial-key-abcdef"
|
||||
baseURL := dialFailServerWithUpgrade(t, key)
|
||||
|
||||
tr, err := NewDeepgramTranscriber(DeepgramConfig{APIKey: key, BaseURL: baseURL})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
chunks := make(chan []byte)
|
||||
close(chunks)
|
||||
_, ferr := tr.StreamTranscribe(context.Background(), chunks, nil)
|
||||
if ferr == nil {
|
||||
t.Fatal("expected dial failure")
|
||||
}
|
||||
if errors.Is(ferr, context.Canceled) {
|
||||
t.Fatalf("dial failure should not be context.Canceled: %v", ferr)
|
||||
}
|
||||
if strings.Contains(ferr.Error(), key) {
|
||||
t.Errorf("API key leaked from dial failure: %v", ferr)
|
||||
}
|
||||
if !strings.Contains(ferr.Error(), "connecting to Deepgram") {
|
||||
t.Errorf("expected dial-path error prefix, got: %v", ferr)
|
||||
}
|
||||
}
|
||||
|
||||
// dialFailServerWithUpgrade returns a ws URL whose handshake responds 101 with
|
||||
// Upgrade set to upgradeValue (not "websocket"), so websocket.Dial fails with
|
||||
// that value in the error string. Hijacks and closes immediately to avoid the
|
||||
// library's 3s body-drain timeout on failed dials.
|
||||
func dialFailServerWithUpgrade(t *testing.T, upgradeValue string) string {
|
||||
t.Helper()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
hj, ok := w.(http.Hijacker)
|
||||
if !ok {
|
||||
http.Error(w, "hijack not supported", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
conn, bufrw, err := hj.Hijack()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer conn.Close()
|
||||
_, _ = bufrw.WriteString("HTTP/1.1 101 Switching Protocols\r\n")
|
||||
_, _ = bufrw.WriteString("Connection: Upgrade\r\n")
|
||||
_, _ = bufrw.WriteString("Upgrade: " + upgradeValue + "\r\n\r\n")
|
||||
_ = bufrw.Flush()
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return "ws" + strings.TrimPrefix(srv.URL, "http")
|
||||
}
|
||||
|
||||
func TestDeepgramStreamTranscribeErrorRedaction(t *testing.T) {
|
||||
url := wsTestServer(t, func(ctx context.Context, c *websocket.Conn) {
|
||||
// Drain the client's audio frames, then wait for CloseStream so the
|
||||
// client has flushed its writes before we reject the connection. Closing
|
||||
// immediately (before CloseStream) risks the client failing on a write
|
||||
// error and never observing this API-key-bearing close reason.
|
||||
for {
|
||||
typ, data, err := c.Read(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if typ == websocket.MessageText && strings.Contains(string(data), "CloseStream") {
|
||||
break
|
||||
}
|
||||
}
|
||||
c.Close(websocket.StatusPolicyViolation, "invalid key sk-test-key")
|
||||
})
|
||||
|
||||
tr, err := NewDeepgramTranscriber(DeepgramConfig{APIKey: "sk-test-key", BaseURL: url})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
chunks := make(chan []byte, 1)
|
||||
chunks <- make([]byte, 320)
|
||||
close(chunks)
|
||||
_, ferr := tr.StreamTranscribe(context.Background(), chunks, func(string, bool) {})
|
||||
if ferr == nil {
|
||||
t.Fatal("expected error")
|
||||
}
|
||||
if strings.Contains(ferr.Error(), "sk-test-key") {
|
||||
t.Errorf("API key leaked: %v", ferr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeepgramStreamTranscribeCancelKeepsSentinel(t *testing.T) {
|
||||
firstFrame := make(chan struct{})
|
||||
url := wsTestServer(t, func(ctx context.Context, c *websocket.Conn) {
|
||||
// Hold the connection open, never answering, so the client blocks in
|
||||
// Read until its context is cancelled (the Esc-abort path).
|
||||
var once sync.Once
|
||||
for {
|
||||
if _, _, err := c.Read(ctx); err != nil {
|
||||
return
|
||||
}
|
||||
once.Do(func() { close(firstFrame) })
|
||||
}
|
||||
})
|
||||
|
||||
tr, err := NewDeepgramTranscriber(DeepgramConfig{APIKey: "sk-test-key", BaseURL: url})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
chunks := make(chan []byte, 1)
|
||||
defer close(chunks)
|
||||
chunks <- make([]byte, 320)
|
||||
// The channel stays open: the session is live when the user aborts.
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
_, ferr := tr.StreamTranscribe(ctx, chunks, nil)
|
||||
errCh <- ferr
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-firstFrame:
|
||||
cancel()
|
||||
ferr := <-errCh
|
||||
if !errors.Is(ferr, context.Canceled) {
|
||||
t.Fatalf("cancelled stream error lost the context.Canceled sentinel: %v", ferr)
|
||||
}
|
||||
case ferr := <-errCh:
|
||||
t.Fatalf("StreamTranscribe failed early instead of blocking: %v", ferr)
|
||||
}
|
||||
}
|
||||
|
||||
// Esc can cancel while the WebSocket dial is still pending. The dial error is
|
||||
// redacted with %s, so we must return ctx.Err() rather than a flat string.
|
||||
func TestDeepgramStreamTranscribeDialCancelKeepsSentinel(t *testing.T) {
|
||||
tr, err := NewDeepgramTranscriber(DeepgramConfig{
|
||||
APIKey: "sk-test-key",
|
||||
BaseURL: "ws://127.0.0.1:1", // never reached; ctx is already cancelled
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
chunks := make(chan []byte)
|
||||
defer close(chunks)
|
||||
|
||||
_, ferr := tr.StreamTranscribe(ctx, chunks, nil)
|
||||
if !errors.Is(ferr, context.Canceled) {
|
||||
t.Fatalf("dial-cancel lost the context.Canceled sentinel: %v", ferr)
|
||||
}
|
||||
}
|
||||
|
||||
// Esc right after accept can race the first write or the first Read; both
|
||||
// setup/write redaction sites must still yield context.Canceled.
|
||||
func TestDeepgramStreamTranscribeStartupCancelKeepsSentinel(t *testing.T) {
|
||||
accepted := make(chan struct{})
|
||||
url := wsTestServer(t, func(ctx context.Context, c *websocket.Conn) {
|
||||
close(accepted)
|
||||
for {
|
||||
if _, _, err := c.Read(ctx); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
tr, err := NewDeepgramTranscriber(DeepgramConfig{APIKey: "sk-test-key", BaseURL: url})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
chunks := make(chan []byte, 1)
|
||||
defer close(chunks)
|
||||
// No frames yet: the writer blocks on chunks or the reader blocks on Read.
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
_, ferr := tr.StreamTranscribe(ctx, chunks, nil)
|
||||
errCh <- ferr
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-accepted:
|
||||
cancel()
|
||||
ferr := <-errCh
|
||||
if !errors.Is(ferr, context.Canceled) {
|
||||
t.Fatalf("startup-cancel lost the context.Canceled sentinel: %v", ferr)
|
||||
}
|
||||
case ferr := <-errCh:
|
||||
t.Fatalf("StreamTranscribe failed before accept: %v", ferr)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDeepgramStreamErrorRedactsHeaderShapedKey exercises the same stream-error
|
||||
// redaction path as TestDeepgramStreamTranscribeErrorRedaction, with a close
|
||||
// reason that embeds the key next to an X-Api-Key-shaped header label.
|
||||
func TestDeepgramStreamErrorRedactsHeaderShapedKey(t *testing.T) {
|
||||
url := wsTestServer(t, func(ctx context.Context, c *websocket.Conn) {
|
||||
for {
|
||||
typ, data, err := c.Read(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if typ == websocket.MessageText && strings.Contains(string(data), "CloseStream") {
|
||||
break
|
||||
}
|
||||
}
|
||||
c.Close(websocket.StatusPolicyViolation, "X-Api-Key: sk-custom-secret-key-1234567890 failed")
|
||||
})
|
||||
|
||||
tr, err := NewDeepgramTranscriber(DeepgramConfig{APIKey: "sk-custom-secret-key-1234567890", BaseURL: url})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
chunks := make(chan []byte, 1)
|
||||
chunks <- make([]byte, 320)
|
||||
close(chunks)
|
||||
_, ferr := tr.StreamTranscribe(context.Background(), chunks, func(string, bool) {})
|
||||
if ferr == nil {
|
||||
t.Fatal("expected error")
|
||||
}
|
||||
if strings.Contains(ferr.Error(), "sk-custom-secret-key-1234567890") {
|
||||
t.Errorf("Custom header API key leaked: %v", ferr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeepgramSinglePassRedaction(t *testing.T) {
|
||||
url := wsTestServer(t, func(ctx context.Context, c *websocket.Conn) {
|
||||
for {
|
||||
typ, data, err := c.Read(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if typ == websocket.MessageText && strings.Contains(string(data), "CloseStream") {
|
||||
break
|
||||
}
|
||||
}
|
||||
c.Close(websocket.StatusPolicyViolation, "error with sk-key1234567890abcdef1234")
|
||||
})
|
||||
|
||||
tr, err := NewDeepgramTranscriber(DeepgramConfig{APIKey: "sk-key1234567890abcdef1234", BaseURL: url})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
chunks := make(chan []byte, 1)
|
||||
chunks <- make([]byte, 320)
|
||||
close(chunks)
|
||||
_, ferr := tr.StreamTranscribe(context.Background(), chunks, func(string, bool) {})
|
||||
if ferr == nil {
|
||||
t.Fatal("expected error")
|
||||
}
|
||||
got := ferr.Error()
|
||||
if strings.Contains(got, "sk-key1234567890abcdef1234") {
|
||||
t.Fatalf("API key leaked into stream error: %q", got)
|
||||
}
|
||||
// One key occurrence must produce exactly one redaction marker (no double wrap).
|
||||
if strings.Count(got, "[REDACTED]") != 1 {
|
||||
t.Fatalf("redaction markers in %q: want exactly 1 [REDACTED]", got)
|
||||
}
|
||||
}
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/providers/providerio"
|
||||
"github.com/coder/websocket"
|
||||
)
|
||||
|
||||
@@ -22,6 +23,10 @@ type OpenAIRealtimeConfig struct {
|
||||
Model string // default "gpt-4o-transcribe"
|
||||
// BaseURL overrides the wss endpoint (tests point it at a fake server).
|
||||
BaseURL string
|
||||
// writeErrInjector, when set, is consulted before each JSON websocket write.
|
||||
// Used by same-package tests to force controllable write failures that embed
|
||||
// secrets; production callers leave it nil.
|
||||
writeErrInjector func() error
|
||||
}
|
||||
|
||||
type openAIRealtimeTranscriber struct {
|
||||
@@ -59,7 +64,15 @@ func (o *openAIRealtimeTranscriber) StreamTranscribe(ctx context.Context, chunks
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("connecting to OpenAI Realtime: %w", err)
|
||||
// Preserve the context.Canceled sentinel (it carries no key) so an
|
||||
// Esc-abort during dial still matches the UI's errors.Is check. Key
|
||||
// on ctx.Err() itself rather than unwrapping err: a cancellation can
|
||||
// surface as a plain transport error (e.g. "closed network
|
||||
// connection") rather than context.Canceled directly.
|
||||
if ctx.Err() != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
return "", fmt.Errorf("connecting to OpenAI Realtime: %s", providerio.Redact(err.Error(), o.cfg.APIKey))
|
||||
}
|
||||
defer conn.CloseNow()
|
||||
|
||||
@@ -73,24 +86,53 @@ func (o *openAIRealtimeTranscriber) StreamTranscribe(ctx context.Context, chunks
|
||||
},
|
||||
},
|
||||
}
|
||||
if err := writeJSON(ctx, conn, sessionUpdate); err != nil {
|
||||
return "", fmt.Errorf("configuring OpenAI Realtime session: %w", err)
|
||||
if err := o.writeJSON(ctx, conn, sessionUpdate); err != nil {
|
||||
// Same cancellation short-circuit as the dial above: an Esc-abort
|
||||
// while the session update is in flight must stay context.Canceled.
|
||||
if ctx.Err() != nil {
|
||||
return "", ctx.Err()
|
||||
}
|
||||
return "", fmt.Errorf("configuring OpenAI Realtime session: %s", providerio.Redact(err.Error(), o.cfg.APIKey))
|
||||
}
|
||||
|
||||
writeErrCh := make(chan error, 1)
|
||||
go func() {
|
||||
for chunk := range chunks {
|
||||
appendMsg := map[string]any{
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": base64.StdEncoding.EncodeToString(chunk),
|
||||
}
|
||||
if err := writeJSON(ctx, conn, appendMsg); err != nil {
|
||||
writeErrCh <- err
|
||||
// Select on ctx.Done() while waiting for the next chunk so a cancelled
|
||||
// session does not leave this goroutine blocked on an open idle channel.
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
writeErrCh <- ctx.Err()
|
||||
return
|
||||
case chunk, ok := <-chunks:
|
||||
if !ok {
|
||||
// Commit the buffered audio to force final transcription.
|
||||
if err := o.writeJSON(ctx, conn, map[string]any{"type": "input_audio_buffer.commit"}); err != nil {
|
||||
writeErrCh <- err
|
||||
// Unblock the reader: without a close, conn.Read stays
|
||||
// parked on a silent peer and the error is never drained
|
||||
// from writeErrCh. CloseNow aborts instead of waiting on
|
||||
// the close handshake, which a silent peer never answers.
|
||||
conn.CloseNow()
|
||||
return
|
||||
}
|
||||
writeErrCh <- nil
|
||||
return
|
||||
}
|
||||
appendMsg := map[string]any{
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": base64.StdEncoding.EncodeToString(chunk),
|
||||
}
|
||||
if err := o.writeJSON(ctx, conn, appendMsg); err != nil {
|
||||
writeErrCh <- err
|
||||
// Same unblocking close as the commit path: a server that
|
||||
// goes silent after session setup would otherwise leave the
|
||||
// reader blocked in conn.Read with the error undrained.
|
||||
conn.CloseNow()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
// Commit the buffered audio to force final transcription of the utterance.
|
||||
writeErrCh <- writeJSON(ctx, conn, map[string]any{"type": "input_audio_buffer.commit"})
|
||||
}()
|
||||
|
||||
// OpenAI deltas are incremental (append), unlike Deepgram/sherpa's cumulative
|
||||
@@ -113,7 +155,17 @@ func (o *openAIRealtimeTranscriber) StreamTranscribe(ctx context.Context, chunks
|
||||
}
|
||||
default:
|
||||
}
|
||||
return compose(), fmt.Errorf("OpenAI Realtime stream error: %w", err)
|
||||
// A user abort cancels the streaming context; the UI matches it
|
||||
// with errors.Is(err, context.Canceled), so return the sentinel
|
||||
// itself (it carries no key) instead of a flat redacted string.
|
||||
// Key on ctx.Err() rather than unwrapping err: the writeErrCh
|
||||
// swap above can replace err with a plain transport error (e.g.
|
||||
// "closed network connection") that doesn't itself unwrap to
|
||||
// context.Canceled even though the cancellation is what caused it.
|
||||
if ctx.Err() != nil {
|
||||
return compose(), ctx.Err()
|
||||
}
|
||||
return compose(), fmt.Errorf("OpenAI Realtime stream error: %s", providerio.Redact(err.Error(), o.cfg.APIKey))
|
||||
}
|
||||
if typ != websocket.MessageText {
|
||||
continue
|
||||
@@ -125,6 +177,12 @@ func (o *openAIRealtimeTranscriber) StreamTranscribe(ctx context.Context, chunks
|
||||
if onPartial != nil {
|
||||
onPartial(compose(), false)
|
||||
}
|
||||
// onPartial may deliver the partial that makes the TUI cancel
|
||||
// (Esc). Observe that before the next Read, which can already
|
||||
// have a racing OpenAI error frame buffered.
|
||||
if ctx.Err() != nil {
|
||||
return compose(), ctx.Err()
|
||||
}
|
||||
case realtimeCompleted:
|
||||
// A completed item replaces the in-progress delta buffer with the
|
||||
// server's authoritative transcript for that segment.
|
||||
@@ -138,6 +196,9 @@ func (o *openAIRealtimeTranscriber) StreamTranscribe(ctx context.Context, chunks
|
||||
if onPartial != nil {
|
||||
onPartial(compose(), true)
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return compose(), ctx.Err()
|
||||
}
|
||||
// Once we've committed the buffer (user stopped) and the server has
|
||||
// returned a completed transcription, the utterance is done.
|
||||
if committed {
|
||||
@@ -147,7 +208,15 @@ func (o *openAIRealtimeTranscriber) StreamTranscribe(ctx context.Context, chunks
|
||||
case realtimeCommitted:
|
||||
committed = true
|
||||
case realtimeError:
|
||||
return compose(), fmt.Errorf("OpenAI Realtime error: %s", evt.text)
|
||||
// Esc can race an incoming error event: conn.Read may already have
|
||||
// the OpenAI error in hand by the time cancellation lands. Key on
|
||||
// ctx.Err() first, same as every other return point in this loop,
|
||||
// so a cancel that wins the race still surfaces as context.Canceled
|
||||
// instead of a spurious redacted failure alongside it.
|
||||
if ctx.Err() != nil {
|
||||
return compose(), ctx.Err()
|
||||
}
|
||||
return compose(), fmt.Errorf("OpenAI Realtime error: %s", providerio.Redact(evt.text, o.cfg.APIKey))
|
||||
}
|
||||
// The writer signals commit completion out of band; observe it so a
|
||||
// stop with no further audio still flips `committed`.
|
||||
@@ -155,6 +224,10 @@ func (o *openAIRealtimeTranscriber) StreamTranscribe(ctx context.Context, chunks
|
||||
case werr := <-writeErrCh:
|
||||
if werr == nil {
|
||||
committed = true
|
||||
} else if ctx.Err() != nil {
|
||||
return compose(), ctx.Err()
|
||||
} else {
|
||||
return compose(), fmt.Errorf("OpenAI Realtime stream error: %s", providerio.Redact(werr.Error(), o.cfg.APIKey))
|
||||
}
|
||||
default:
|
||||
}
|
||||
@@ -201,6 +274,15 @@ func parseRealtimeEvent(data []byte) realtimeEvent {
|
||||
return realtimeEvent{kind: realtimeOther}
|
||||
}
|
||||
|
||||
func (o *openAIRealtimeTranscriber) writeJSON(ctx context.Context, conn *websocket.Conn, v any) error {
|
||||
if o.cfg.writeErrInjector != nil {
|
||||
if err := o.cfg.writeErrInjector(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return writeJSON(ctx, conn, v)
|
||||
}
|
||||
|
||||
func writeJSON(ctx context.Context, conn *websocket.Conn, v any) error {
|
||||
data, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
|
||||
@@ -0,0 +1,437 @@
|
||||
package dictation
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/coder/websocket"
|
||||
)
|
||||
|
||||
// Non-cancellation dial failure: handshake returns 101 with an Upgrade header
|
||||
// equal to the API key so the dial error embeds the key and redaction must strip it.
|
||||
func TestOpenAIRealtimeDialFailureRedactsKey(t *testing.T) {
|
||||
const key = "sk-test-dial-key-abcdef"
|
||||
baseURL := dialFailServerWithUpgrade(t, key)
|
||||
|
||||
tr, err := NewOpenAIRealtimeTranscriber(OpenAIRealtimeConfig{APIKey: key, BaseURL: baseURL})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
chunks := make(chan []byte)
|
||||
close(chunks)
|
||||
_, ferr := tr.StreamTranscribe(context.Background(), chunks, nil)
|
||||
if ferr == nil {
|
||||
t.Fatal("expected dial failure")
|
||||
}
|
||||
if errors.Is(ferr, context.Canceled) {
|
||||
t.Fatalf("dial failure should not be context.Canceled: %v", ferr)
|
||||
}
|
||||
if strings.Contains(ferr.Error(), key) {
|
||||
t.Errorf("API key leaked from dial failure: %v", ferr)
|
||||
}
|
||||
if !strings.Contains(ferr.Error(), "connecting to OpenAI Realtime") {
|
||||
t.Errorf("expected dial-path error prefix, got: %v", ferr)
|
||||
}
|
||||
}
|
||||
|
||||
// Session-configuration write fails after dial: inject a write error that embeds
|
||||
// the API key (peer close cannot reach this path reliably — the session write
|
||||
// always lands in the kernel buffer first) and assert redaction.
|
||||
func TestOpenAIRealtimeSessionConfigFailureRedactsKey(t *testing.T) {
|
||||
const key = "sk-test-session-key-abcdef"
|
||||
url := wsTestServer(t, func(ctx context.Context, c *websocket.Conn) {
|
||||
// Hold open; the client fails on the injected session write before reading.
|
||||
<-ctx.Done()
|
||||
})
|
||||
|
||||
tr, err := NewOpenAIRealtimeTranscriber(OpenAIRealtimeConfig{
|
||||
APIKey: key,
|
||||
BaseURL: url,
|
||||
writeErrInjector: func() error {
|
||||
return errors.New("invalid API key " + key)
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
chunks := make(chan []byte)
|
||||
close(chunks)
|
||||
_, ferr := tr.StreamTranscribe(context.Background(), chunks, nil)
|
||||
if ferr == nil {
|
||||
t.Fatal("expected session-config failure")
|
||||
}
|
||||
if errors.Is(ferr, context.Canceled) {
|
||||
t.Fatalf("session-config failure should not be context.Canceled: %v", ferr)
|
||||
}
|
||||
if strings.Contains(ferr.Error(), key) {
|
||||
t.Errorf("API key leaked from session-config failure: %v", ferr)
|
||||
}
|
||||
if !strings.Contains(ferr.Error(), "configuring OpenAI Realtime session") {
|
||||
t.Errorf("expected session-config error prefix, got: %v", ferr)
|
||||
}
|
||||
}
|
||||
|
||||
// Asynchronous writer failure: session update succeeds, then the audio writer
|
||||
// injects a key-bearing write error. The server keeps the connection open and
|
||||
// emits a delta so the reader loop observes writeErrCh (not a peer close).
|
||||
func TestOpenAIRealtimeWriterFailureRedactsKey(t *testing.T) {
|
||||
const key = "sk-test-writer-key-abcdef"
|
||||
sessionDone := make(chan struct{})
|
||||
url := wsTestServer(t, func(ctx context.Context, c *websocket.Conn) {
|
||||
if _, _, err := c.Read(ctx); err != nil {
|
||||
return
|
||||
}
|
||||
close(sessionDone)
|
||||
// Keep reading and emit a delta so the client processes an event and
|
||||
// drains writeErrCh while the connection is still open.
|
||||
_ = c.Write(ctx, websocket.MessageText, []byte(
|
||||
`{"type":"conversation.item.input_audio_transcription.delta","delta":"hi"}`,
|
||||
))
|
||||
for {
|
||||
if _, _, err := c.Read(ctx); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
var writes atomic.Int32
|
||||
tr, err := NewOpenAIRealtimeTranscriber(OpenAIRealtimeConfig{
|
||||
APIKey: key,
|
||||
BaseURL: url,
|
||||
writeErrInjector: func() error {
|
||||
// First write is session.update; fail subsequent audio/commit writes.
|
||||
if writes.Add(1) == 1 {
|
||||
return nil
|
||||
}
|
||||
return errors.New("invalid API key " + key)
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
chunks := make(chan []byte, 1)
|
||||
chunks <- make([]byte, 480)
|
||||
// Leave open after the first frame so the writer hits the injected error
|
||||
// on that append rather than a clean commit.
|
||||
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
_, ferr := tr.StreamTranscribe(context.Background(), chunks, func(string, bool) {})
|
||||
done <- ferr
|
||||
}()
|
||||
|
||||
var ferr error
|
||||
select {
|
||||
case ferr = <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
close(chunks)
|
||||
t.Fatal("StreamTranscribe did not complete after writer failure")
|
||||
}
|
||||
close(chunks)
|
||||
|
||||
if ferr == nil {
|
||||
t.Fatal("expected writer failure")
|
||||
}
|
||||
if errors.Is(ferr, context.Canceled) {
|
||||
t.Fatalf("writer failure should not be context.Canceled: %v", ferr)
|
||||
}
|
||||
if strings.Contains(ferr.Error(), key) {
|
||||
t.Errorf("API key leaked from writer failure: %v", ferr)
|
||||
}
|
||||
if !strings.Contains(ferr.Error(), "OpenAI Realtime stream error") {
|
||||
t.Errorf("expected stream/writer error prefix, got: %v", ferr)
|
||||
}
|
||||
select {
|
||||
case <-sessionDone:
|
||||
default:
|
||||
// Session may race the injected audio write; prefix assertion above is enough.
|
||||
}
|
||||
}
|
||||
|
||||
// Writer failure with a silent server: the peer never sends a frame after
|
||||
// session setup, so the reader is parked in conn.Read. The writer must close
|
||||
// the connection on failure or the writeErrCh error is never observed and
|
||||
// StreamTranscribe blocks forever.
|
||||
func TestOpenAIRealtimeWriterFailureUnblocksSilentServer(t *testing.T) {
|
||||
const key = "sk-test-silent-writer-key-abcdef"
|
||||
url := wsTestServer(t, func(ctx context.Context, c *websocket.Conn) {
|
||||
// Hold the connection open without ever writing a frame.
|
||||
<-ctx.Done()
|
||||
})
|
||||
|
||||
var writes atomic.Int32
|
||||
tr, err := NewOpenAIRealtimeTranscriber(OpenAIRealtimeConfig{
|
||||
APIKey: key,
|
||||
BaseURL: url,
|
||||
writeErrInjector: func() error {
|
||||
// First write is session.update; fail the subsequent append write.
|
||||
if writes.Add(1) == 1 {
|
||||
return nil
|
||||
}
|
||||
return errors.New("invalid API key " + key)
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
chunks := make(chan []byte, 1)
|
||||
chunks <- make([]byte, 480)
|
||||
// Leave the channel open: the writer must fail on the append, not the commit.
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
_, ferr := tr.StreamTranscribe(ctx, chunks, func(string, bool) {})
|
||||
done <- ferr
|
||||
}()
|
||||
|
||||
select {
|
||||
case ferr := <-done:
|
||||
if ferr == nil {
|
||||
t.Fatal("expected writer failure")
|
||||
}
|
||||
if errors.Is(ferr, context.Canceled) {
|
||||
t.Fatalf("writer failure should not be context.Canceled: %v", ferr)
|
||||
}
|
||||
if strings.Contains(ferr.Error(), key) {
|
||||
t.Errorf("API key leaked from writer failure: %v", ferr)
|
||||
}
|
||||
if !strings.Contains(ferr.Error(), "OpenAI Realtime stream error") {
|
||||
t.Errorf("expected stream/writer error prefix, got: %v", ferr)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
close(chunks)
|
||||
t.Fatal("StreamTranscribe did not unblock after writer failure with a silent server")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIRealtimeStreamTranscribeErrorRedaction(t *testing.T) {
|
||||
url := wsTestServer(t, func(ctx context.Context, c *websocket.Conn) {
|
||||
defer c.Close(websocket.StatusNormalClosure, "")
|
||||
for {
|
||||
typ, _, err := c.Read(ctx)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if typ != websocket.MessageText {
|
||||
continue
|
||||
}
|
||||
_ = c.Write(ctx, websocket.MessageText, []byte(`{"type":"error","error":{"message":"invalid API key sk-test-key"}}`))
|
||||
}
|
||||
})
|
||||
|
||||
tr, err := NewOpenAIRealtimeTranscriber(OpenAIRealtimeConfig{APIKey: "sk-test-key", BaseURL: url})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
chunks := make(chan []byte, 1)
|
||||
chunks <- make([]byte, 480)
|
||||
close(chunks)
|
||||
_, ferr := tr.StreamTranscribe(context.Background(), chunks, func(string, bool) {})
|
||||
if ferr == nil {
|
||||
t.Fatal("expected error")
|
||||
}
|
||||
if strings.Contains(ferr.Error(), "sk-test-key") {
|
||||
t.Errorf("API key leaked: %v", ferr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIRealtimeStreamTranscribeCancelKeepsSentinel(t *testing.T) {
|
||||
firstFrame := make(chan struct{})
|
||||
url := wsTestServer(t, func(ctx context.Context, c *websocket.Conn) {
|
||||
// Hold the connection open, never answering, so the client blocks in
|
||||
// Read until its context is cancelled (the Esc-abort path).
|
||||
var once sync.Once
|
||||
for {
|
||||
if _, _, err := c.Read(ctx); err != nil {
|
||||
return
|
||||
}
|
||||
once.Do(func() { close(firstFrame) })
|
||||
}
|
||||
})
|
||||
|
||||
tr, err := NewOpenAIRealtimeTranscriber(OpenAIRealtimeConfig{APIKey: "sk-test-key", BaseURL: url})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
chunks := make(chan []byte, 1)
|
||||
defer close(chunks)
|
||||
chunks <- make([]byte, 480)
|
||||
// The channel stays open: the session is live when the user aborts.
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
_, ferr := tr.StreamTranscribe(ctx, chunks, nil)
|
||||
errCh <- ferr
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-firstFrame:
|
||||
cancel()
|
||||
ferr := <-errCh
|
||||
if !errors.Is(ferr, context.Canceled) {
|
||||
t.Fatalf("cancelled stream error lost the context.Canceled sentinel: %v", ferr)
|
||||
}
|
||||
case ferr := <-errCh:
|
||||
t.Fatalf("StreamTranscribe failed early instead of blocking: %v", ferr)
|
||||
}
|
||||
}
|
||||
|
||||
// Esc can cancel while the WebSocket dial is still pending. The dial error is
|
||||
// redacted with %s, so we must return ctx.Err() rather than a flat string.
|
||||
func TestOpenAIRealtimeStreamTranscribeDialCancelKeepsSentinel(t *testing.T) {
|
||||
tr, err := NewOpenAIRealtimeTranscriber(OpenAIRealtimeConfig{
|
||||
APIKey: "sk-test-key",
|
||||
BaseURL: "ws://127.0.0.1:1", // never reached; ctx is already cancelled
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
chunks := make(chan []byte)
|
||||
defer close(chunks)
|
||||
|
||||
_, ferr := tr.StreamTranscribe(ctx, chunks, nil)
|
||||
if !errors.Is(ferr, context.Canceled) {
|
||||
t.Fatalf("dial-cancel lost the context.Canceled sentinel: %v", ferr)
|
||||
}
|
||||
}
|
||||
|
||||
// Esc right after accept can hit the session-update write or the first Read;
|
||||
// both redaction sites must still yield context.Canceled.
|
||||
func TestOpenAIRealtimeStreamTranscribeStartupCancelKeepsSentinel(t *testing.T) {
|
||||
accepted := make(chan struct{})
|
||||
url := wsTestServer(t, func(ctx context.Context, c *websocket.Conn) {
|
||||
close(accepted)
|
||||
// Do not read: leave the client in session-update write or first Read.
|
||||
<-ctx.Done()
|
||||
})
|
||||
|
||||
tr, err := NewOpenAIRealtimeTranscriber(OpenAIRealtimeConfig{APIKey: "sk-test-key", BaseURL: url})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
chunks := make(chan []byte, 1)
|
||||
defer close(chunks)
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
_, ferr := tr.StreamTranscribe(ctx, chunks, nil)
|
||||
errCh <- ferr
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-accepted:
|
||||
cancel()
|
||||
ferr := <-errCh
|
||||
if !errors.Is(ferr, context.Canceled) {
|
||||
t.Fatalf("startup-cancel lost the context.Canceled sentinel: %v", ferr)
|
||||
}
|
||||
case ferr := <-errCh:
|
||||
t.Fatalf("StreamTranscribe failed before accept: %v", ferr)
|
||||
}
|
||||
}
|
||||
|
||||
// Cancel while the writer is blocked on more chunks. Whichever path observes
|
||||
// the cancellation first (the blocked conn.Read or the writeErrCh drain) must
|
||||
// return context.Canceled rather than a redacted flat string.
|
||||
func TestOpenAIRealtimeStreamTranscribeWriteCancelKeepsSentinel(t *testing.T) {
|
||||
sessionReceived := make(chan struct{})
|
||||
url := wsTestServer(t, func(ctx context.Context, c *websocket.Conn) {
|
||||
if _, _, err := c.Read(ctx); err != nil {
|
||||
return
|
||||
}
|
||||
close(sessionReceived)
|
||||
for {
|
||||
if _, _, err := c.Read(ctx); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
tr, err := NewOpenAIRealtimeTranscriber(OpenAIRealtimeConfig{APIKey: "sk-test-key", BaseURL: url})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
chunks := make(chan []byte, 1)
|
||||
chunks <- make([]byte, 480)
|
||||
defer close(chunks)
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
_, ferr := tr.StreamTranscribe(ctx, chunks, nil)
|
||||
errCh <- ferr
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-sessionReceived:
|
||||
cancel()
|
||||
ferr := <-errCh
|
||||
if !errors.Is(ferr, context.Canceled) {
|
||||
t.Fatalf("write-path cancel lost the context.Canceled sentinel: %v", ferr)
|
||||
}
|
||||
case ferr := <-errCh:
|
||||
t.Fatalf("StreamTranscribe failed early instead of blocking: %v", ferr)
|
||||
}
|
||||
}
|
||||
|
||||
// Esc can race an incoming OpenAI error event. conn.Read may already have
|
||||
// the error frame in hand by the time the cancellation lands. The server
|
||||
// fires a delta immediately followed by the error, back to back. Cancel from
|
||||
// inside the delta callback (same path the TUI uses when a partial triggers
|
||||
// Esc handling) so the cancel is observed before the next event is processed.
|
||||
// The returned error must still be context.Canceled, not the redacted OpenAI
|
||||
// error.
|
||||
func TestOpenAIRealtimeStreamTranscribeErrorRaceKeepsSentinel(t *testing.T) {
|
||||
url := wsTestServer(t, func(ctx context.Context, c *websocket.Conn) {
|
||||
if _, _, err := c.Read(ctx); err != nil {
|
||||
return
|
||||
}
|
||||
_ = c.Write(ctx, websocket.MessageText, []byte(
|
||||
`{"type":"conversation.item.input_audio_transcription.delta","delta":"hi"}`,
|
||||
))
|
||||
_ = c.Write(ctx, websocket.MessageText, []byte(
|
||||
`{"type":"error","error":{"message":"invalid API key sk-test-key"}}`,
|
||||
))
|
||||
<-ctx.Done()
|
||||
})
|
||||
|
||||
tr, err := NewOpenAIRealtimeTranscriber(OpenAIRealtimeConfig{APIKey: "sk-test-key", BaseURL: url})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
chunks := make(chan []byte, 1)
|
||||
defer close(chunks)
|
||||
chunks <- make([]byte, 480)
|
||||
var once sync.Once
|
||||
_, ferr := tr.StreamTranscribe(ctx, chunks, func(string, bool) {
|
||||
// Cancel synchronously from the partial callback, before the
|
||||
// stream loop continues to the already-buffered error frame.
|
||||
once.Do(cancel)
|
||||
})
|
||||
if !errors.Is(ferr, context.Canceled) {
|
||||
t.Fatalf("StreamTranscribe error = %v, want context.Canceled (cancel must win over a racing OpenAI error event)", ferr)
|
||||
}
|
||||
}
|
||||
@@ -9,7 +9,16 @@ import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
const PolicyVersion = 1
|
||||
// PolicyVersion is the version of the approval policy that recorded grants were
|
||||
// evaluated under. Bump it whenever a change narrows what may be approved: the
|
||||
// grant store then backs the file up, keeps deny grants, invalidates prior
|
||||
// approvals, and tells the user once, instead of silently honoring approvals
|
||||
// that the current rules would have refused.
|
||||
//
|
||||
// v2 tightened command-prefix validation so a launcher spelled with a version
|
||||
// or an executable extension ("python3.11", "python.exe") can no longer be
|
||||
// approved as an ordinary command.
|
||||
const PolicyVersion = 2
|
||||
|
||||
type Origin string
|
||||
|
||||
|
||||
@@ -6,6 +6,8 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
)
|
||||
@@ -38,12 +40,45 @@ func KillProcessTree(pid int) error {
|
||||
// TerminateProcessTree requests graceful termination, then force-kills the
|
||||
// process tree after grace. Callers retain their distinct persistence models;
|
||||
// this function owns only the OS lifecycle primitive.
|
||||
//
|
||||
// The signal target is REDISCOVERED here via processSignalTarget(pid), which
|
||||
// asks the OS for pid's current group. A caller that already knows pid was
|
||||
// configured as its own group leader at launch (e.g. via ConfigureProcessGroup)
|
||||
// should use TerminateProcessGroup instead — see its doc comment for why this
|
||||
// rediscovery is fragile once the leader may have already exited.
|
||||
func TerminateProcessTree(pid int, grace, poll time.Duration) error {
|
||||
target, err := processSignalTarget(pid)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
alive := func() bool { return syscall.Kill(target, syscall.Signal(0)) == nil }
|
||||
return terminateTarget(pid, target, grace, poll)
|
||||
}
|
||||
|
||||
// TerminateProcessGroup is like TerminateProcessTree, but for a caller that
|
||||
// already knows pid is (or was, at launch) its own process-group leader — e.g.
|
||||
// a command started after ConfigureProcessGroup(cmd), which sets Setpgid so
|
||||
// cmd.Process.Pid always leads its own group at the moment it is signalled.
|
||||
//
|
||||
// TerminateProcessTree instead rediscovers the group via syscall.Getpgid(pid)
|
||||
// at call time. On Darwin, that lookup can return ESRCH once an unreaped group
|
||||
// leader has already exited, even though live descendants remain in the group
|
||||
// it configured — TerminateProcessTree would then silently fall back to
|
||||
// signalling only the (already dead) leader PID and leave those descendants
|
||||
// running. Skipping rediscovery in favor of the launch-time invariant avoids
|
||||
// that: the negative PID is used directly, regardless of whether the leader is
|
||||
// still alive to be looked up.
|
||||
func TerminateProcessGroup(pid int, grace, poll time.Duration) error {
|
||||
if pid <= 1 {
|
||||
return fmt.Errorf("refusing to signal invalid pid %d", pid)
|
||||
}
|
||||
return terminateTarget(pid, -pid, grace, poll)
|
||||
}
|
||||
|
||||
// terminateTarget signals target (an individual PID or, negative, a process
|
||||
// group) with SIGTERM then, after grace, SIGKILL. reportPid identifies the
|
||||
// original process in error messages regardless of which form target takes.
|
||||
func terminateTarget(reportPid, target int, grace, poll time.Duration) error {
|
||||
alive := func() bool { return signalTargetRunning(target) }
|
||||
if err := syscall.Kill(target, syscall.SIGTERM); err != nil {
|
||||
if errors.Is(err, syscall.ESRCH) {
|
||||
return nil
|
||||
@@ -77,11 +112,91 @@ func TerminateProcessTree(pid int, grace, poll time.Duration) error {
|
||||
time.Sleep(poll)
|
||||
}
|
||||
if alive() {
|
||||
return fmt.Errorf("process %d did not exit after SIGKILL", pid)
|
||||
return fmt.Errorf("process %d did not exit after SIGKILL", reportPid)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// signalTargetRunning reports whether the signal target still has a process that
|
||||
// is actually RUNNING, as opposed to one waiting to be reaped.
|
||||
//
|
||||
// kill(2) with signal 0 succeeds for a zombie: the PID still exists until its
|
||||
// parent collects it. That matters because the target here is usually a process
|
||||
// GROUP, and a terminated leader's child is briefly a zombie before init/launchd
|
||||
// reaps it. Treating that window as "still alive" made termination sit through
|
||||
// both grace periods, SIGKILL an already-dead group, and then report
|
||||
// "did not exit after SIGKILL" for a tree that had in fact stopped — turning a
|
||||
// successful cleanup into a spurious failure (and a flaky one, since it depends on
|
||||
// reap timing).
|
||||
//
|
||||
// Zombie detection goes through ps, which is available on every Unix we target and
|
||||
// avoids a /proc dependency Darwin does not have. It is only consulted when the
|
||||
// cheap kill check says something is still there, so the common path costs nothing
|
||||
// extra. If ps cannot be run or its output cannot be parsed, the conservative
|
||||
// pre-existing answer (the kill result) stands.
|
||||
func signalTargetRunning(target int) bool {
|
||||
if syscall.Kill(target, syscall.Signal(0)) != nil {
|
||||
return false
|
||||
}
|
||||
// A negative target is a process GROUP: check whether any member is
|
||||
// non-zombie. A positive target is an individual PID that is NOT its own
|
||||
// group leader (processSignalTarget's fallback) — its actual PGID differs
|
||||
// from its own PID, so matching rows by PGID here would look for a group
|
||||
// that doesn't exist and always report "unknown" (conservatively alive).
|
||||
// Check its own PID's state instead.
|
||||
var running, ok bool
|
||||
if target < 0 {
|
||||
running, ok = processGroupHasRunningMember(-target)
|
||||
} else {
|
||||
running, ok = processIsRunning(target)
|
||||
}
|
||||
if !ok {
|
||||
return true
|
||||
}
|
||||
return running
|
||||
}
|
||||
|
||||
// processGroupHasRunningMember reports whether any member of pgid is in a state
|
||||
// other than zombie. ok is false when the group's states could not be determined.
|
||||
func processGroupHasRunningMember(pgid int) (running bool, ok bool) {
|
||||
return processTableStateMatches(func(_, rowPgid int) bool { return rowPgid == pgid })
|
||||
}
|
||||
|
||||
// processIsRunning reports whether pid itself is in a state other than zombie.
|
||||
// ok is false when pid's state could not be determined.
|
||||
func processIsRunning(pid int) (running bool, ok bool) {
|
||||
return processTableStateMatches(func(rowPid, _ int) bool { return rowPid == pid })
|
||||
}
|
||||
|
||||
// processTableStateMatches scans the process table once, via ps, and reports
|
||||
// whether any row satisfying match is non-zombie, and whether at least one
|
||||
// matching row was found at all (ok). No rows found means either the process
|
||||
// (or group) genuinely doesn't exist (a race with exit) or ps's output was
|
||||
// unparseable/restricted; either way, the caller should not claim knowledge.
|
||||
func processTableStateMatches(match func(pid, pgid int) bool) (running bool, ok bool) {
|
||||
output, err := exec.Command("ps", "-A", "-o", "pid=,pgid=,stat=").Output()
|
||||
if err != nil {
|
||||
return false, false
|
||||
}
|
||||
found := false
|
||||
for line := range strings.SplitSeq(string(output), "\n") {
|
||||
fields := strings.Fields(line)
|
||||
if len(fields) < 3 {
|
||||
continue
|
||||
}
|
||||
rowPid, errPid := strconv.Atoi(fields[0])
|
||||
rowPgid, errPgid := strconv.Atoi(fields[1])
|
||||
if errPid != nil || errPgid != nil || !match(rowPid, rowPgid) {
|
||||
continue
|
||||
}
|
||||
found = true
|
||||
if !strings.HasPrefix(fields[2], "Z") {
|
||||
return true, true
|
||||
}
|
||||
}
|
||||
return false, found
|
||||
}
|
||||
|
||||
func processSignalTarget(pid int) (int, error) {
|
||||
if pid <= 1 {
|
||||
return 0, fmt.Errorf("refusing to signal invalid pid %d", pid)
|
||||
@@ -91,10 +206,14 @@ func processSignalTarget(pid int) (int, error) {
|
||||
if pgid == pid {
|
||||
target = -pid
|
||||
}
|
||||
} else if errors.Is(err, syscall.ESRCH) {
|
||||
// Preserve the individual target; the signal call below treats ESRCH as
|
||||
// already gone, which is a successful lifecycle outcome.
|
||||
return pid, nil
|
||||
} else {
|
||||
if errors.Is(err, syscall.ESRCH) {
|
||||
// Preserve the individual target; the signal call below treats ESRCH as
|
||||
// already gone, which is a successful lifecycle outcome.
|
||||
return pid, nil
|
||||
}
|
||||
// Conservatively retain the individual PID after other lookup failures;
|
||||
// guessing and signalling a process group could affect unrelated processes.
|
||||
}
|
||||
return target, nil
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
//go:build !windows
|
||||
|
||||
package execution
|
||||
|
||||
import (
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"syscall"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestTerminateProcessTreeTreatsZombieGroupAsExited is the regression test for
|
||||
// counting unreaped processes as alive: kill(2) with signal 0 succeeds for a
|
||||
// zombie, so a group whose only remaining member is waiting to be reaped used to
|
||||
// sit through both grace periods, take a pointless SIGKILL, and then be reported
|
||||
// as "did not exit after SIGKILL" — a spurious cleanup failure whose timing
|
||||
// depended on when the reaper ran.
|
||||
func TestTerminateProcessTreeTreatsZombieGroupAsExited(t *testing.T) {
|
||||
cmd := exec.Command("sh", "-c", "exit 0")
|
||||
ConfigureProcessGroup(cmd)
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.Fatalf("start: %v", err)
|
||||
}
|
||||
pid := cmd.Process.Pid
|
||||
// Deliberately do NOT Wait: the child stays a zombie in its own group, which
|
||||
// is exactly the state a terminated tree passes through.
|
||||
t.Cleanup(func() { _ = cmd.Wait() })
|
||||
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for {
|
||||
if zombie, ok := processIsZombie(pid); ok && zombie {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Skip("child did not reach the zombie state; ps output is unavailable in this environment")
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
if syscall.Kill(-pid, syscall.Signal(0)) != nil {
|
||||
t.Skip("the zombie group is no longer signalable; cannot exercise the case")
|
||||
}
|
||||
|
||||
if signalTargetRunning(-pid) {
|
||||
t.Fatal("a group holding only a zombie must not count as running")
|
||||
}
|
||||
if err := TerminateProcessTree(pid, 50*time.Millisecond, 10*time.Millisecond); err != nil {
|
||||
t.Fatalf("TerminateProcessTree on an already-exited tree: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTerminateProcessTreeStopsRunningGroup is the positive control: a group with
|
||||
// a live member must still be seen as running and then actually stopped, so the
|
||||
// zombie tolerance above cannot degrade into ignoring real processes.
|
||||
func TestTerminateProcessTreeStopsRunningGroup(t *testing.T) {
|
||||
cmd := exec.Command("sh", "-c", "sleep 30")
|
||||
ConfigureProcessGroup(cmd)
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.Fatalf("start: %v", err)
|
||||
}
|
||||
pid := cmd.Process.Pid
|
||||
defer func() { _ = cmd.Wait() }()
|
||||
|
||||
if !signalTargetRunning(-pid) {
|
||||
t.Fatal("a group with a live member must count as running")
|
||||
}
|
||||
if err := TerminateProcessTree(pid, time.Second, 10*time.Millisecond); err != nil {
|
||||
t.Fatalf("TerminateProcessTree: %v", err)
|
||||
}
|
||||
if running, ok := processGroupHasRunningMember(pid); ok && running {
|
||||
t.Fatal("the group still has a running member after termination")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSignalTargetRunningTreatsZombieIndividualPIDAsExited is the regression
|
||||
// test for jatmn's second #774 finding: signalTargetRunning's zombie check
|
||||
// matched the individual-PID fallback target (a process that is NOT its own
|
||||
// group leader — processSignalTarget's positive-PID case) against process-table
|
||||
// rows by PGID. That target's actual PGID differs from its own PID by
|
||||
// definition (that's exactly why processSignalTarget fell back to the
|
||||
// individual PID instead of a negative group target), so no row ever matched,
|
||||
// "unknown" resulted every time, and signalTargetRunning conservatively (and
|
||||
// incorrectly) treated a zombie individual-PID target as still running.
|
||||
func TestSignalTargetRunningTreatsZombieIndividualPIDAsExited(t *testing.T) {
|
||||
cmd := exec.Command("sh", "-c", "exit 0")
|
||||
// Deliberately do NOT call ConfigureProcessGroup: the child inherits this
|
||||
// test process's group, so it is not its own leader and processSignalTarget
|
||||
// falls back to the individual (positive) PID.
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.Fatalf("start: %v", err)
|
||||
}
|
||||
pid := cmd.Process.Pid
|
||||
t.Cleanup(func() { _ = cmd.Wait() })
|
||||
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for {
|
||||
if zombie, ok := processIsZombie(pid); ok && zombie {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Skip("child did not reach the zombie state; ps output is unavailable in this environment")
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
|
||||
target, err := processSignalTarget(pid)
|
||||
if err != nil {
|
||||
t.Fatalf("processSignalTarget: %v", err)
|
||||
}
|
||||
if target != pid {
|
||||
t.Skip("child unexpectedly became its own group leader; cannot exercise the individual-PID fallback")
|
||||
}
|
||||
if signalTargetRunning(target) {
|
||||
t.Fatal("a zombie individual-PID target must not count as running")
|
||||
}
|
||||
}
|
||||
|
||||
// processIsZombie reports a process's zombie state via ps. ok is false when the
|
||||
// state could not be read.
|
||||
func processIsZombie(pid int) (zombie bool, ok bool) {
|
||||
output, err := exec.Command("ps", "-o", "stat=", "-p", strconv.Itoa(pid)).Output()
|
||||
if err != nil {
|
||||
return false, false
|
||||
}
|
||||
state := string(output)
|
||||
for len(state) > 0 && (state[0] == ' ' || state[0] == '\n' || state[0] == '\t') {
|
||||
state = state[1:]
|
||||
}
|
||||
if state == "" {
|
||||
return false, false
|
||||
}
|
||||
return state[0] == 'Z', true
|
||||
}
|
||||
@@ -3,6 +3,8 @@
|
||||
package execution
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
@@ -12,10 +14,34 @@ import (
|
||||
|
||||
func ConfigureProcessGroup(cmd *exec.Cmd) {}
|
||||
|
||||
// KillProcessTree asks taskkill /T to kill pid and descendants it can still
|
||||
// discover from that live root. If taskkill fails, it falls back to killing pid
|
||||
// alone but still returns the tree error because descendants may have survived.
|
||||
// Neither path can recover descendants after the root has exited.
|
||||
func KillProcessTree(pid int) error {
|
||||
if err := exec.Command(taskkillPath(), "/T", "/F", "/PID", strconv.Itoa(pid)).Run(); err == nil {
|
||||
return killProcessTree(pid, runTaskkill, killRootProcess)
|
||||
}
|
||||
|
||||
func killProcessTree(pid int, killTree, killRoot func(int) error) error {
|
||||
treeErr := killTree(pid)
|
||||
if treeErr == nil {
|
||||
return nil
|
||||
}
|
||||
rootErr := killRoot(pid)
|
||||
if rootErr != nil {
|
||||
return errors.Join(
|
||||
fmt.Errorf("kill process tree rooted at %d: %w", pid, treeErr),
|
||||
fmt.Errorf("kill root process %d: %w", pid, rootErr),
|
||||
)
|
||||
}
|
||||
return fmt.Errorf("kill process tree rooted at %d (root process killed directly): %w", pid, treeErr)
|
||||
}
|
||||
|
||||
func runTaskkill(pid int) error {
|
||||
return exec.Command(taskkillPath(), "/T", "/F", "/PID", strconv.Itoa(pid)).Run()
|
||||
}
|
||||
|
||||
func killRootProcess(pid int) error {
|
||||
process, err := os.FindProcess(pid)
|
||||
if err != nil {
|
||||
return err
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
//go:build windows
|
||||
|
||||
package execution
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestKillProcessTreeRetainsTreeErrorWhenRootFallbackSucceeds(t *testing.T) {
|
||||
treeErr := errors.New("taskkill failed")
|
||||
err := killProcessTree(42, func(int) error { return treeErr }, func(int) error { return nil })
|
||||
if !errors.Is(err, treeErr) {
|
||||
t.Fatalf("killProcessTree error = %v, want tree error", err)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "root process killed directly") {
|
||||
t.Fatalf("killProcessTree error = %q, want fallback context", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKillProcessTreeJoinsTreeAndRootErrors(t *testing.T) {
|
||||
treeErr := errors.New("taskkill failed")
|
||||
rootErr := errors.New("root kill failed")
|
||||
err := killProcessTree(42, func(int) error { return treeErr }, func(int) error { return rootErr })
|
||||
if !errors.Is(err, treeErr) || !errors.Is(err, rootErr) {
|
||||
t.Fatalf("killProcessTree error = %v, want both failures", err)
|
||||
}
|
||||
}
|
||||
@@ -35,7 +35,7 @@ type Server struct {
|
||||
// credential lookup uses it to avoid reusing legacy user tokens by name.
|
||||
ProjectConfigured bool
|
||||
// UnconfiguredDefault is true when this server is one of Zero's built-in
|
||||
// defaults (e.g. keyless Firecrawl) that the user never touched in their
|
||||
// defaults (e.g. keyless Exa) that the user never touched in their
|
||||
// config — no credentials, no overrides. Callers use it to avoid warning
|
||||
// loudly when a server nobody configured fails to connect.
|
||||
UnconfiguredDefault bool
|
||||
|
||||
+118
-3
@@ -1,6 +1,8 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
@@ -90,7 +92,7 @@ func TestNormalizeConfigValidatesTransportBoundaries(t *testing.T) {
|
||||
|
||||
func TestNormalizeConfigFlagsUnconfiguredDefault(t *testing.T) {
|
||||
cfg := config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"firecrawl": config.DefaultMCPServers()["firecrawl"],
|
||||
"exa": config.DefaultMCPServers()["exa"],
|
||||
"web": {
|
||||
Type: "http",
|
||||
URL: "https://example.com/mcp",
|
||||
@@ -107,8 +109,8 @@ func TestNormalizeConfigFlagsUnconfiguredDefault(t *testing.T) {
|
||||
byName[server.Name] = server
|
||||
}
|
||||
|
||||
if !byName["firecrawl"].UnconfiguredDefault {
|
||||
t.Fatal("an untouched firecrawl default should be flagged UnconfiguredDefault")
|
||||
if !byName["exa"].UnconfiguredDefault {
|
||||
t.Fatal("an untouched exa default should be flagged UnconfiguredDefault")
|
||||
}
|
||||
if byName["web"].UnconfiguredDefault {
|
||||
t.Fatal("a server the user configured must not be flagged UnconfiguredDefault")
|
||||
@@ -240,3 +242,116 @@ func TestCopyStringMapTrimsKeysAndPreservesValues(t *testing.T) {
|
||||
t.Fatalf("copied[TOKEN] = %q, want value preserved verbatim", copied["TOKEN"])
|
||||
}
|
||||
}
|
||||
|
||||
// Startup behavior for the firecrawl -> exa default migration: a user who
|
||||
// disabled the retired default must end up with no Exa server to connect to,
|
||||
// not merely a disabled entry in the resolved config.
|
||||
func TestNormalizeConfigSkipsExaAfterLegacyFirecrawlDisable(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"disabled":true}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := config.ResolveMCP(config.ResolveOptions{UserConfigPath: path})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
servers, err := NormalizeConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NormalizeConfig: %v", err)
|
||||
}
|
||||
for _, server := range servers {
|
||||
if server.Name == "exa" {
|
||||
t.Fatalf("exa must not be started after the user disabled the default it replaced: %#v", server)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The same path with no legacy disable still starts Exa, so the migration
|
||||
// cannot quietly switch off the default for everyone else.
|
||||
func TestNormalizeConfigStartsExaWithoutLegacyDisable(t *testing.T) {
|
||||
cfg, err := config.ResolveMCP(config.ResolveOptions{})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
servers, err := NormalizeConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("NormalizeConfig: %v", err)
|
||||
}
|
||||
for _, server := range servers {
|
||||
if server.Name == "exa" {
|
||||
if !server.UnconfiguredDefault {
|
||||
t.Fatalf("an untouched exa default should be flagged unconfigured: %#v", server)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Fatal("expected the exa default to be started out of the box")
|
||||
}
|
||||
|
||||
// The failure path CodeRabbit flagged: a firecrawl entry that named only a
|
||||
// header relied on the retired default for its transport. Resolving it must
|
||||
// still produce a complete server — an incomplete one fails NormalizeConfig,
|
||||
// which takes down every other server's startup with it, not just this one.
|
||||
func TestNormalizeConfigResolvesPartialRetiredDefaultEntry(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"headers":{"Authorization":"Bearer k"}}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := config.ResolveMCP(config.ResolveOptions{UserConfigPath: path})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
servers, err := NormalizeConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("a header-only legacy entry must not break startup: %v", err)
|
||||
}
|
||||
var firecrawl *Server
|
||||
var sawExa bool
|
||||
for i := range servers {
|
||||
switch servers[i].Name {
|
||||
case "firecrawl":
|
||||
firecrawl = &servers[i]
|
||||
case "exa":
|
||||
sawExa = true
|
||||
}
|
||||
}
|
||||
if firecrawl == nil {
|
||||
t.Fatalf("expected the customized firecrawl server to survive: %#v", servers)
|
||||
}
|
||||
if firecrawl.Type != ServerTypeHTTP || firecrawl.URL != "https://mcp.firecrawl.dev/v2/mcp" {
|
||||
t.Fatalf("firecrawl lost the transport its entry relied on: %#v", *firecrawl)
|
||||
}
|
||||
if firecrawl.Headers["Authorization"] != "Bearer k" {
|
||||
t.Fatalf("firecrawl lost the user's header: %#v", *firecrawl)
|
||||
}
|
||||
if !sawExa {
|
||||
t.Fatal("the new exa default should still start alongside it")
|
||||
}
|
||||
}
|
||||
|
||||
// The mirror case: an entry that names its own stdio transport must not have
|
||||
// the retired http default grafted onto it, which would be rejected outright.
|
||||
func TestNormalizeConfigLeavesRetiredEntryWithOwnTransportAlone(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "config.json")
|
||||
if err := os.WriteFile(path, []byte(`{"mcp":{"servers":{"firecrawl":{"command":"firecrawl-mcp"}}}}`), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := config.ResolveMCP(config.ResolveOptions{UserConfigPath: path})
|
||||
if err != nil {
|
||||
t.Fatalf("ResolveMCP: %v", err)
|
||||
}
|
||||
servers, err := NormalizeConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("a self-hosted stdio entry must keep working: %v", err)
|
||||
}
|
||||
for _, server := range servers {
|
||||
if server.Name != "firecrawl" {
|
||||
continue
|
||||
}
|
||||
if server.Type != ServerTypeStdio || server.Command != "firecrawl-mcp" {
|
||||
t.Fatalf("the user's own stdio transport was corrupted: %#v", server)
|
||||
}
|
||||
return
|
||||
}
|
||||
t.Fatalf("expected the self-hosted firecrawl server to survive: %#v", servers)
|
||||
}
|
||||
@@ -153,16 +153,16 @@ func TestRegisterToolsSkipsUnreachableServerAndKeepsOthers(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRegisterToolsFlagsUnconfiguredDefaultOnSkip(t *testing.T) {
|
||||
// firecrawl is seeded by config.DefaultMCPServers with no credentials; when a
|
||||
// user never configures it, a connect failure (e.g. the real HTTP 401 from
|
||||
// issue #552) must be recorded as UnconfiguredDefault so callers can avoid
|
||||
// Exa is seeded by config.DefaultMCPServers with no credentials; when a
|
||||
// user never configures it, a connect failure must be recorded as
|
||||
// UnconfiguredDefault so callers can avoid
|
||||
// warning about a server the user never asked for. "custom" is configured by
|
||||
// the user and must NOT be flagged even though it also fails to connect.
|
||||
registry := tools.NewRegistry()
|
||||
|
||||
runtime, err := RegisterTools(context.Background(), registry, config.MCPConfig{Servers: map[string]config.MCPServerConfig{
|
||||
"firecrawl": config.DefaultMCPServers()["firecrawl"],
|
||||
"custom": {Type: "stdio", Command: "custom-mcp"},
|
||||
"exa": config.DefaultMCPServers()["exa"],
|
||||
"custom": {Type: "stdio", Command: "custom-mcp"},
|
||||
}}, RegisterOptions{
|
||||
ClientFactory: func(_ context.Context, server Server) (ToolClient, error) {
|
||||
return nil, errors.New(server.Name + " connect failed")
|
||||
@@ -177,12 +177,12 @@ func TestRegisterToolsFlagsUnconfiguredDefaultOnSkip(t *testing.T) {
|
||||
for _, skipped := range runtime.Skipped() {
|
||||
byName[skipped.Name] = skipped
|
||||
}
|
||||
firecrawl, ok := byName["firecrawl"]
|
||||
exa, ok := byName["exa"]
|
||||
if !ok {
|
||||
t.Fatalf("Skipped() = %#v, want an entry for firecrawl", runtime.Skipped())
|
||||
t.Fatalf("Skipped() = %#v, want an entry for exa", runtime.Skipped())
|
||||
}
|
||||
if !firecrawl.UnconfiguredDefault {
|
||||
t.Fatalf("Skipped()[firecrawl] = %#v, want UnconfiguredDefault", firecrawl)
|
||||
if !exa.UnconfiguredDefault {
|
||||
t.Fatalf("Skipped()[exa] = %#v, want UnconfiguredDefault", exa)
|
||||
}
|
||||
custom, ok := byName["custom"]
|
||||
if !ok {
|
||||
|
||||
@@ -231,6 +231,12 @@ func reasoningEffortsForModelName(name string) []ReasoningEffort {
|
||||
// the gateway's translation concern — unknown fields are ignored, so the
|
||||
// worst case is a silent no-op rather than a 400.
|
||||
return []ReasoningEffort{ReasoningEffortLow, ReasoningEffortMedium, ReasoningEffortHigh}
|
||||
case strings.Contains(n, "deepseek"):
|
||||
// DeepSeek V4 (deepseek-v4-flash / deepseek-v4-pro) exposes thinking-mode
|
||||
// effort control through the OpenAI-compatible reasoning_effort parameter.
|
||||
// The API accepts low/high/max and maps medium/xhigh onto high, so only
|
||||
// the three distinct tiers are advertised.
|
||||
return []ReasoningEffort{ReasoningEffortLow, ReasoningEffortHigh, ReasoningEffortMax}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -20,6 +20,10 @@ func TestReasoningEffortsFallbackForGPT5AndOSeries(t *testing.T) {
|
||||
{"hy3", []ReasoningEffort{ReasoningEffortLow, ReasoningEffortMedium, ReasoningEffortHigh}},
|
||||
{"hunyuan-t1", []ReasoningEffort{ReasoningEffortLow, ReasoningEffortMedium, ReasoningEffortHigh}},
|
||||
{"openai/o3-mini", []ReasoningEffort{ReasoningEffortLow, ReasoningEffortMedium, ReasoningEffortHigh}}, // github-style vendor/ id
|
||||
{"deepseek-v4-flash", []ReasoningEffort{ReasoningEffortLow, ReasoningEffortHigh, ReasoningEffortMax}},
|
||||
{"deepseek-v4-pro", []ReasoningEffort{ReasoningEffortLow, ReasoningEffortHigh, ReasoningEffortMax}},
|
||||
{"deepseek-reasoner", []ReasoningEffort{ReasoningEffortLow, ReasoningEffortHigh, ReasoningEffortMax}},
|
||||
{"deepseek-chat", []ReasoningEffort{ReasoningEffortLow, ReasoningEffortHigh, ReasoningEffortMax}},
|
||||
{"gpt-4.1", nil}, // non-reasoning, registered: stays empty
|
||||
{"gpt-4o-mini", nil},
|
||||
{"ollama/llama3.1", nil},
|
||||
|
||||
@@ -29,6 +29,9 @@ const (
|
||||
// newer Anthropic models). Defined so a capability catalog that lists "max"
|
||||
// has a constant to map to; no curated model lists it yet.
|
||||
ReasoningEffortMax ReasoningEffort = "max"
|
||||
// ReasoningEffortUltra is advertised by models whose provider exposes an
|
||||
// additional tier above max.
|
||||
ReasoningEffortUltra ReasoningEffort = "ultra"
|
||||
)
|
||||
|
||||
type ModelCapability string
|
||||
@@ -410,7 +413,7 @@ func ValidRuntimeProviderKind(provider ProviderKind) bool {
|
||||
|
||||
func ValidReasoningEffort(effort ReasoningEffort) bool {
|
||||
switch effort {
|
||||
case ReasoningEffortNone, ReasoningEffortMinimal, ReasoningEffortLow, ReasoningEffortMedium, ReasoningEffortHigh, ReasoningEffortXHigh, ReasoningEffortMax:
|
||||
case ReasoningEffortNone, ReasoningEffortMinimal, ReasoningEffortLow, ReasoningEffortMedium, ReasoningEffortHigh, ReasoningEffortXHigh, ReasoningEffortMax, ReasoningEffortUltra:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
|
||||
@@ -0,0 +1,177 @@
|
||||
// Package pathjail confines file operations to a directory tree.
|
||||
//
|
||||
// It exists because the obvious guard does not work. Checking a path with
|
||||
// Lstat and then handing the same string to MkdirAll, CreateTemp, Rename or
|
||||
// Remove is two independent resolutions with a gap between them, and every
|
||||
// operation after the check re-resolves the ANCESTORS as well. A store that
|
||||
// verified only its own directory and file was therefore writable outside its
|
||||
// tree whenever any component above them was a link: the check passed, and the
|
||||
// syscall that followed went somewhere else.
|
||||
//
|
||||
// On Windows the same guard failed for a second, independent reason. A junction
|
||||
// is a reparse point but NOT a symlink, so os.Lstat reports it as ModeIrregular
|
||||
// with ModeSymlink clear, and a guard testing only ModeSymlink returns "not a
|
||||
// link" when pointed straight at one. Junctions also need no privilege to
|
||||
// create, unlike Windows symlinks, so this is reachable by exactly the
|
||||
// unprivileged user a workspace guard is meant to contain.
|
||||
//
|
||||
// The fix for both is to stop naming things by path. Open a handle on the root
|
||||
// once and address everything below it relative to that handle, so the kernel
|
||||
// resolves each component against an object that cannot be swapped underneath.
|
||||
// os.Root is that handle: it refuses any component that would escape, including
|
||||
// a Windows reparse point, and it is the same idea the sandbox uses for ACL
|
||||
// materialization.
|
||||
package pathjail
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io/fs"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ErrEscapes reports a target that is not inside the root it was confined to.
|
||||
// Returned before anything is opened, so a caller cannot act on the path first.
|
||||
var ErrEscapes = errors.New("path escapes its root")
|
||||
|
||||
// ErrReparse reports a link or reparse point where a real file or directory was
|
||||
// required.
|
||||
var ErrReparse = errors.New("refusing to write through a link")
|
||||
|
||||
// Open confines dir to root and returns a handle on root plus dir expressed
|
||||
// relative to it. The caller closes the handle.
|
||||
//
|
||||
// root is resolved normally, because it is the boundary the caller chose and
|
||||
// trusts (a workspace the operator opened, say, which may legitimately sit under
|
||||
// a link). Everything BELOW it is resolved against the returned handle and
|
||||
// cannot escape.
|
||||
func Open(root, dir string) (*os.Root, string, error) {
|
||||
// Only a genuinely EMPTY value is treated as missing; every other spelling is
|
||||
// used verbatim. Testing TrimSpace(...) == "" instead would retarget the very
|
||||
// paths this guards: " " is a legal directory name on unix, so trimming would
|
||||
// reject it as "no root" or silently swap dir for root, which is the same
|
||||
// retargeting bug as trimming "root " down to "root".
|
||||
if root == "" {
|
||||
return nil, "", fmt.Errorf("%w: no root to confine %q to", ErrEscapes, dir)
|
||||
}
|
||||
if dir == "" {
|
||||
dir = root
|
||||
}
|
||||
// Rel before Open, so an escaping target is refused without opening
|
||||
// anything. filepath.Rel is lexical, which is the right test here: it says
|
||||
// what the caller ASKED for, and the handle enforces what actually happens.
|
||||
relative, err := filepath.Rel(root, dir)
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("%w: %q against %q: %v", ErrEscapes, dir, root, err)
|
||||
}
|
||||
if relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) {
|
||||
return nil, "", fmt.Errorf("%w: %q is outside %q", ErrEscapes, dir, root)
|
||||
}
|
||||
// The root is created if it is missing, by ordinary pathname, because it is
|
||||
// the boundary the caller declared rather than anything being confined. The
|
||||
// stores this replaced called MkdirAll on the whole tree, so this creates
|
||||
// strictly less by pathname than before, and everything BELOW root is opened
|
||||
// relative to the handle.
|
||||
if err := os.MkdirAll(root, 0o700); err != nil {
|
||||
return nil, "", fmt.Errorf("create %s: %w", root, err)
|
||||
}
|
||||
handle, err := os.OpenRoot(root)
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("open %s: %w", root, err)
|
||||
}
|
||||
return handle, relative, nil
|
||||
}
|
||||
|
||||
// RefuseReparse fails when relative names a link or reparse point.
|
||||
//
|
||||
// Absent is fine: a name that does not exist yet is what a create is for.
|
||||
//
|
||||
// This is belt to os.Root's braces. The handle already refuses to traverse a
|
||||
// reparse point, so this exists to give the FINAL component a clear error
|
||||
// instead of an "escapes from parent" one, and to keep the refusal explicit at
|
||||
// the point a reader looks for it.
|
||||
//
|
||||
// ModeIrregular is tested alongside ModeSymlink deliberately: that is what a
|
||||
// Windows junction reports, and a guard checking only ModeSymlink lets one
|
||||
// straight through.
|
||||
func RefuseReparse(handle *os.Root, relative string) error {
|
||||
// A trailing separator makes os.Root.Lstat resolve the terminal component on
|
||||
// Windows (doInRootAlwaysResolveTerminalSlash), so a junction named with a
|
||||
// trailing slash would be stat'd as its target and wave through as "not a
|
||||
// link". Strip it so the final component itself is the thing inspected.
|
||||
//
|
||||
// os.IsPathSeparator rather than a literal `/\` set, because backslash is a
|
||||
// separator only on Windows: on unix it is an ordinary filename character, and
|
||||
// trimming it would inspect "linked" when asked about a symlink genuinely
|
||||
// named `linked\` — reporting "not a link" for one that is.
|
||||
if trimmed := trimTrailingSeparators(relative); trimmed != "" {
|
||||
relative = trimmed
|
||||
}
|
||||
info, err := handle.Lstat(relative)
|
||||
if err != nil {
|
||||
if errors.Is(err, fs.ErrNotExist) {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("inspect %s: %w", relative, err)
|
||||
}
|
||||
if info.Mode()&(os.ModeSymlink|os.ModeIrregular) != 0 {
|
||||
return fmt.Errorf("%w: %s", ErrReparse, relative)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// trimTrailingSeparators drops the path separators at the end of a name, using
|
||||
// the ACTIVE platform's definition of one: `/` and `\` on Windows, `/` only on
|
||||
// unix, where a backslash is a legal character in a filename.
|
||||
func trimTrailingSeparators(path string) string {
|
||||
end := len(path)
|
||||
for end > 0 && os.IsPathSeparator(path[end-1]) {
|
||||
end--
|
||||
}
|
||||
return path[:end]
|
||||
}
|
||||
|
||||
// randomBytes fills b with cryptographic randomness. A variable rather than a
|
||||
// direct crypto/rand call so a test can force the name collision that exercises
|
||||
// the O_EXCL retry below — the branch that matters, and one a real random source
|
||||
// will not reach on demand.
|
||||
var randomBytes = rand.Read
|
||||
|
||||
// CreateTemp makes an O_EXCL temp file named prefix.<random>suffix inside dir,
|
||||
// relative to handle, and returns it with its relative path.
|
||||
//
|
||||
// O_EXCL rather than a fixed name: a link planted at a predictable temp path
|
||||
// between the check and the create is the same race this package exists to
|
||||
// close, and O_EXCL makes the kernel refuse it rather than us.
|
||||
func CreateTemp(handle *os.Root, dir, prefix, suffix string) (*os.File, string, error) {
|
||||
// prefix and suffix are NAME fragments, not path fragments. A separator or a
|
||||
// "."/".." in either would relocate the file out of dir — a prefix of
|
||||
// "../escaped", or a suffix of ".tmp/../evil" that also erases the random
|
||||
// component and with it the O_EXCL unpredictability the file relies on.
|
||||
// Reject them up front so the file lands where dir promises.
|
||||
for _, fragment := range []string{prefix, suffix} {
|
||||
if strings.ContainsAny(fragment, `/\`) || fragment == "." || fragment == ".." {
|
||||
return nil, "", fmt.Errorf("%w: temporary name fragment %q must be a single path component", ErrEscapes, fragment)
|
||||
}
|
||||
}
|
||||
for attempt := 0; attempt < 10; attempt++ {
|
||||
noise := make([]byte, 8)
|
||||
if _, err := randomBytes(noise); err != nil {
|
||||
return nil, "", fmt.Errorf("generate a temporary name: %w", err)
|
||||
}
|
||||
relative := filepath.Join(dir, prefix+"."+hex.EncodeToString(noise)+suffix)
|
||||
file, err := handle.OpenFile(relative, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0o600)
|
||||
if err == nil {
|
||||
return file, relative, nil
|
||||
}
|
||||
if errors.Is(err, fs.ErrExist) {
|
||||
continue
|
||||
}
|
||||
return nil, "", err
|
||||
}
|
||||
return nil, "", errors.New("could not create a temporary file with an unused name")
|
||||
}
|
||||
@@ -0,0 +1,360 @@
|
||||
package pathjail
|
||||
|
||||
import (
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// linkDir points link at target, using whatever kind of link this platform lets
|
||||
// an unprivileged process create.
|
||||
//
|
||||
// On Windows that is a JUNCTION, not a symlink, and the difference is the whole
|
||||
// reason this package exists twice over. A symlink needs
|
||||
// SeCreateSymbolicLinkPrivilege, which an ordinary account does not hold, so a
|
||||
// test written with os.Symlink silently skips on Windows and proves nothing
|
||||
// there. A junction needs no privilege at all, so it is both the reachable
|
||||
// attack and the only one that can be tested.
|
||||
func linkDir(t *testing.T, target, link string) {
|
||||
t.Helper()
|
||||
if runtime.GOOS == "windows" {
|
||||
if out, err := exec.Command("cmd", "/c", "mklink", "/J", link, target).CombinedOutput(); err != nil {
|
||||
t.Skipf("cannot create a junction: %v %s", err, out)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err := os.Symlink(target, link); err != nil {
|
||||
t.Skipf("cannot create a symlink: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// The reported defect: a link in an ANCESTOR of the target, not at the target
|
||||
// itself. Every guard that checked only the final directory and file passed
|
||||
// this, because the components it checked were genuinely not links; the ones
|
||||
// the syscall then traversed were.
|
||||
func TestAnAncestorLinkCannotRedirectAWrite(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
outside := filepath.Join(base, "outside")
|
||||
workspace := filepath.Join(base, "workspace")
|
||||
for _, dir := range []string{outside, workspace} {
|
||||
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
linkDir(t, outside, filepath.Join(workspace, ".zero"))
|
||||
|
||||
handle, relative, err := Open(workspace, filepath.Join(workspace, ".zero", "store"))
|
||||
if err != nil {
|
||||
t.Fatalf("Open should confine, not refuse outright: %v", err)
|
||||
}
|
||||
defer handle.Close()
|
||||
|
||||
if err := handle.MkdirAll(relative, 0o700); err == nil {
|
||||
t.Fatal("created a directory through a linked ancestor, so a write can still leave the workspace")
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(outside, "store")); !os.IsNotExist(err) {
|
||||
t.Errorf("something was created outside the workspace, stat error = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// RefuseReparse has to catch a junction, which is the case os.ModeSymlink alone
|
||||
// misses. Asserted through the exported helper rather than by reading the mode,
|
||||
// so the test fails if the predicate is ever narrowed back.
|
||||
func TestRefuseReparseCatchesAJunctionNotJustASymlink(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
target := filepath.Join(base, "target")
|
||||
inside := filepath.Join(base, "inside")
|
||||
for _, dir := range []string{target, inside} {
|
||||
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
linkDir(t, target, filepath.Join(inside, "linked"))
|
||||
|
||||
handle, relative, err := Open(inside, inside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer handle.Close()
|
||||
|
||||
err = RefuseReparse(handle, filepath.Join(relative, "linked"))
|
||||
if !errors.Is(err, ErrReparse) {
|
||||
t.Fatalf("RefuseReparse(link) = %v, want ErrReparse; on Windows a junction reports ModeIrregular, so a ModeSymlink-only check returns nil here", err)
|
||||
}
|
||||
}
|
||||
|
||||
// A name that does not exist yet is what a create is for, so it must not be
|
||||
// refused. Without this the guard above would be satisfied by refusing
|
||||
// everything.
|
||||
func TestRefuseReparseAllowsAnAbsentName(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
handle, relative, err := Open(base, base)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer handle.Close()
|
||||
if err := RefuseReparse(handle, filepath.Join(relative, "not-there.md")); err != nil {
|
||||
t.Errorf("an absent name was refused: %v", err)
|
||||
}
|
||||
if err := RefuseReparse(handle, relative); err != nil {
|
||||
t.Errorf("an ordinary directory was refused: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Open refuses a target outside its root before opening anything, so a caller
|
||||
// cannot act on the path first.
|
||||
func TestOpenRefusesATargetOutsideTheRoot(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
inside := filepath.Join(base, "inside")
|
||||
if err := os.MkdirAll(inside, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, target := range []string{
|
||||
filepath.Join(base, "sibling"),
|
||||
filepath.Join(inside, "..", "..", "escape"),
|
||||
} {
|
||||
if _, _, err := Open(inside, target); !errors.Is(err, ErrEscapes) {
|
||||
t.Errorf("Open(%q, %q) = %v, want ErrEscapes", inside, target, err)
|
||||
}
|
||||
}
|
||||
if _, _, err := Open("", filepath.Join(base, "anything")); !errors.Is(err, ErrEscapes) {
|
||||
t.Error("an empty root was accepted; a boundary is not optional")
|
||||
}
|
||||
}
|
||||
|
||||
// CreateTemp must use O_EXCL and an unpredictable name, so a link planted at a
|
||||
// guessable temp path is refused by the kernel rather than followed.
|
||||
func TestCreateTempIsExclusiveAndUnpredictable(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
handle, relative, err := Open(base, base)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer handle.Close()
|
||||
|
||||
seen := map[string]bool{}
|
||||
for i := 0; i < 8; i++ {
|
||||
file, name, err := CreateTemp(handle, relative, "note", ".tmp")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
file.Close()
|
||||
if seen[name] {
|
||||
t.Fatalf("CreateTemp reused the name %q, so two concurrent writes would stomp each other", name)
|
||||
}
|
||||
seen[name] = true
|
||||
if filepath.Ext(name) != ".tmp" {
|
||||
t.Errorf("temp name %q lost its suffix, so a listing may pick it up as real content", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A legitimate directory whose name carries a space must be confined as spelled.
|
||||
// Trimming it would open a different filesystem object.
|
||||
//
|
||||
// A LEADING space, deliberately: Win32 normalizes a component's trailing spaces
|
||||
// and dots away, so "root " and "root" are one object there and a trailing-space
|
||||
// case cannot express the bug on Windows. A leading space survives on every
|
||||
// platform, so this runs everywhere rather than being gated to unix.
|
||||
// The path is RELATIVE, under t.Chdir, so the leading space sits at the start of
|
||||
// the whole string. An absolute filepath.Join(base, " root") would put the space
|
||||
// in the middle, where TrimSpace cannot reach it — the test would then pass
|
||||
// against the very guard it exists to catch, on every platform.
|
||||
func TestOpenPreservesSpaceInPaths(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
t.Chdir(base)
|
||||
const spaced = " root" // the leading space is part of the name
|
||||
|
||||
handle, relative, err := Open(spaced, spaced)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer handle.Close()
|
||||
|
||||
// Write through the handle, then confirm the bytes landed under the spelling
|
||||
// the caller asked for — the property that actually matters, and one
|
||||
// handle.Name() alone would not prove.
|
||||
file, _, err := CreateTemp(handle, relative, "probe", ".tmp")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
name := file.Name()
|
||||
if err := file.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(base, spaced, filepath.Base(name))); err != nil {
|
||||
t.Errorf("a file created through the handle is not under %q: %v — the leading space was trimmed away", spaced, err)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(base, "root")); err == nil {
|
||||
t.Error(`a trimmed "root" was created instead of the requested " root"`)
|
||||
}
|
||||
}
|
||||
|
||||
// A whitespace-only name is a legal directory on unix, so it must be confined as
|
||||
// given rather than treated as "no root" or silently swapped for the root.
|
||||
func TestOpenAcceptsAWhitespaceOnlyComponent(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("Win32 normalizes a whitespace-only component away; the name is not expressible there")
|
||||
}
|
||||
// The path must be whitespace-only IN FULL for this to bite, so it has to be
|
||||
// relative: an absolute path ending in " " still has a non-empty TrimSpace and
|
||||
// would sail past the old guard, making the test prove nothing.
|
||||
t.Chdir(t.TempDir())
|
||||
handle, _, err := Open(" ", " ")
|
||||
if err != nil {
|
||||
t.Fatalf("a whitespace-only path was rejected as missing: %v", err)
|
||||
}
|
||||
defer handle.Close()
|
||||
if _, err := os.Stat(" "); err != nil {
|
||||
t.Errorf("the whitespace-only directory was not created as spelled: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// prefix and suffix are name fragments. A separator or traversal in either must
|
||||
// be refused, or the file lands outside dir and (for a suffix erasing the random
|
||||
// component) loses the O_EXCL unpredictability it depends on.
|
||||
func TestCreateTempRejectsPathFragmentsThatEscapeTheDir(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
inside := filepath.Join(base, "inside")
|
||||
if err := os.MkdirAll(inside, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// root is base, dir is the subdirectory: a "../escaped" fragment lands in
|
||||
// base — still inside the os.Root, so os.Root does NOT catch it — but outside
|
||||
// the dir CreateTemp promised. This is the case the reviewers reproduced.
|
||||
handle, relative, err := Open(base, inside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer handle.Close()
|
||||
cases := []struct{ prefix, suffix string }{
|
||||
{"../escaped", ".tmp"},
|
||||
{`..\sibling`, ".tmp"},
|
||||
{"note", ".tmp/../evil"},
|
||||
{"note", ".."},
|
||||
{"..", ".tmp"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if _, _, err := CreateTemp(handle, relative, c.prefix, c.suffix); !errors.Is(err, ErrEscapes) {
|
||||
t.Errorf("CreateTemp(prefix=%q, suffix=%q) = %v, want ErrEscapes", c.prefix, c.suffix, err)
|
||||
}
|
||||
}
|
||||
// Nothing may have leaked into base above the promised dir.
|
||||
entries, _ := os.ReadDir(base)
|
||||
for _, entry := range entries {
|
||||
if entry.Name() != "inside" {
|
||||
t.Errorf("a temp fragment escaped into the parent: %q", entry.Name())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A junction named with a trailing separator must still be refused. On Windows
|
||||
// os.Root.Lstat resolves the terminal component when the name ends in a
|
||||
// separator, so without trimming it the link is stat'd as its target and waved
|
||||
// through.
|
||||
func TestRefuseReparseCatchesALinkNamedWithATrailingSeparator(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
target := filepath.Join(base, "target")
|
||||
inside := filepath.Join(base, "inside")
|
||||
for _, dir := range []string{target, inside} {
|
||||
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
linkDir(t, target, filepath.Join(inside, "linked"))
|
||||
handle, relative, err := Open(inside, inside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer handle.Close()
|
||||
name := filepath.Join(relative, "linked") + string(filepath.Separator)
|
||||
if err := RefuseReparse(handle, name); !errors.Is(err, ErrReparse) {
|
||||
t.Fatalf("RefuseReparse(%q) = %v, want ErrReparse", name, err)
|
||||
}
|
||||
}
|
||||
|
||||
// The O_EXCL retry is the branch that protects an existing file, and a real
|
||||
// random source will not collide on demand — so force it. The first two draws
|
||||
// return the same bytes: the name is pre-created, CreateTemp must refuse to
|
||||
// reuse it, leave its contents alone, and come back with a different name.
|
||||
func TestCreateTempRetriesWithoutClobberingAnExistingName(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
handle, relative, err := Open(base, base)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer handle.Close()
|
||||
|
||||
draws := [][]byte{
|
||||
{1, 1, 1, 1, 1, 1, 1, 1}, // collides with the pre-created file
|
||||
{1, 1, 1, 1, 1, 1, 1, 1}, // and again, so the retry must keep trying
|
||||
{2, 2, 2, 2, 2, 2, 2, 2}, // finally a free name
|
||||
}
|
||||
previous := randomBytes
|
||||
t.Cleanup(func() { randomBytes = previous })
|
||||
call := 0
|
||||
randomBytes = func(b []byte) (int, error) {
|
||||
copy(b, draws[min(call, len(draws)-1)])
|
||||
call++
|
||||
return len(b), nil
|
||||
}
|
||||
|
||||
taken := filepath.Join(base, "note."+hex.EncodeToString(draws[0])+".tmp")
|
||||
if err := os.WriteFile(taken, []byte("do not clobber me"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
file, name, err := CreateTemp(handle, relative, "note", ".tmp")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
if filepath.Base(name) == filepath.Base(taken) {
|
||||
t.Fatalf("CreateTemp returned the name that was already taken (%q), so O_EXCL did not protect it", name)
|
||||
}
|
||||
kept, err := os.ReadFile(taken)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(kept) != "do not clobber me" {
|
||||
t.Errorf("the pre-existing file was overwritten: %q", kept)
|
||||
}
|
||||
if call < 2 {
|
||||
t.Errorf("the random source was drawn %d time(s); the collision path never ran", call)
|
||||
}
|
||||
}
|
||||
|
||||
// On unix a backslash is an ordinary filename character, not a separator, so a
|
||||
// link genuinely named `linked\` must still be refused. Trimming backslashes
|
||||
// unconditionally would inspect "linked" instead and report "not a link" for one
|
||||
// that is.
|
||||
func TestRefuseReparseKeepsABackslashInAUnixName(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip(`backslash is a separator on Windows, so this name is not expressible there`)
|
||||
}
|
||||
base := t.TempDir()
|
||||
target := filepath.Join(base, "target")
|
||||
inside := filepath.Join(base, "inside")
|
||||
for _, dir := range []string{target, inside} {
|
||||
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
linkDir(t, target, filepath.Join(inside, `linked\`))
|
||||
|
||||
handle, relative, err := Open(inside, inside)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer handle.Close()
|
||||
|
||||
name := filepath.Join(relative, `linked\`)
|
||||
if err := RefuseReparse(handle, name); !errors.Is(err, ErrReparse) {
|
||||
t.Fatalf("RefuseReparse(%q) = %v, want ErrReparse — the backslash is part of the name here, not a separator", name, err)
|
||||
}
|
||||
}
|
||||
@@ -104,7 +104,7 @@ func windowsTokenOwner(token windows.Token) (*windows.SID, error) {
|
||||
}
|
||||
owner := (*windowsTokenOwnerInfo)(unsafe.Pointer(&buffer[0])).owner
|
||||
if owner == nil {
|
||||
return nil, errors.New("Windows access token has no default owner")
|
||||
return nil, errors.New("windows access token has no default owner")
|
||||
}
|
||||
return owner.Copy()
|
||||
}
|
||||
|
||||
@@ -8,17 +8,21 @@ import (
|
||||
)
|
||||
|
||||
type Model struct {
|
||||
ID string
|
||||
Description string
|
||||
ContextWindow int
|
||||
ToolCall bool
|
||||
Reasoning bool
|
||||
InputModalities []string
|
||||
OutputModalities []string
|
||||
InputCost float64
|
||||
OutputCost float64
|
||||
Tags []string
|
||||
Source string
|
||||
ID string
|
||||
Description string
|
||||
ContextWindow int
|
||||
ToolCall bool
|
||||
Reasoning bool
|
||||
ReasoningEfforts []string
|
||||
DefaultReasoningEffort string
|
||||
ServiceTiers []string
|
||||
DefaultServiceTier string
|
||||
InputModalities []string
|
||||
OutputModalities []string
|
||||
InputCost float64
|
||||
OutputCost float64
|
||||
Tags []string
|
||||
Source string
|
||||
}
|
||||
|
||||
const minimaxModelSource = "https://platform.minimax.io/docs/api-reference/api-overview"
|
||||
@@ -270,6 +274,8 @@ func dedupeModels(defaultModel string, models []Model) []Model {
|
||||
}
|
||||
model.InputModalities = append([]string{}, model.InputModalities...)
|
||||
model.OutputModalities = append([]string{}, model.OutputModalities...)
|
||||
model.ReasoningEfforts = append([]string{}, model.ReasoningEfforts...)
|
||||
model.ServiceTiers = append([]string{}, model.ServiceTiers...)
|
||||
model.Tags = append([]string{}, model.Tags...)
|
||||
seen[model.ID] = true
|
||||
result = append(result, model)
|
||||
|
||||
@@ -122,6 +122,25 @@ func TestModelsDoNotAliasMutableCatalogState(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestDedupeModelsDoesNotAliasCapabilitySlices(t *testing.T) {
|
||||
models := []Model{{
|
||||
ID: "reasoner",
|
||||
ReasoningEfforts: []string{"low", "high"},
|
||||
ServiceTiers: []string{"priority"},
|
||||
}}
|
||||
first := dedupeModels("reasoner", models)
|
||||
first[0].ReasoningEfforts[0] = "mutated"
|
||||
first[0].ServiceTiers[0] = "mutated"
|
||||
|
||||
second := dedupeModels("reasoner", models)
|
||||
if got, want := second[0].ReasoningEfforts[0], "low"; got != want {
|
||||
t.Fatalf("reasoning efforts alias source state: got %q, want %q", got, want)
|
||||
}
|
||||
if got, want := second[0].ServiceTiers[0], "priority"; got != want {
|
||||
t.Fatalf("service tiers alias source state: got %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func modelIDs(models []Model) []string {
|
||||
ids := make([]string, 0, len(models))
|
||||
for _, model := range models {
|
||||
|
||||
@@ -184,6 +184,10 @@ func parseOpenAIStyleModelList(body []byte) ([]remoteModel, error) {
|
||||
|
||||
func ModelsDevProviderID(provider providercatalog.Descriptor) string {
|
||||
switch strings.TrimSpace(provider.ID) {
|
||||
case "chatgpt":
|
||||
// ChatGPT authenticates against the Codex backend, while its public
|
||||
// model metadata is published under the OpenAI provider on models.dev.
|
||||
return "openai"
|
||||
case "dashscope":
|
||||
return "alibaba"
|
||||
case "github":
|
||||
@@ -230,6 +234,11 @@ type remoteModel struct {
|
||||
Tools bool `json:"tools"`
|
||||
// ReasoningRaw accepts both a boolean (models.dev) and OpenRouter's object form.
|
||||
ReasoningRaw json.RawMessage `json:"reasoning"`
|
||||
ReasoningEfforts []string `json:"reasoning_efforts"`
|
||||
DefaultReasoning string `json:"default_reasoning_effort"`
|
||||
ServiceTiers []string `json:"service_tiers"`
|
||||
AdditionalSpeed []string `json:"additional_speed_tiers"`
|
||||
DefaultServiceTier string `json:"default_service_tier"`
|
||||
Free bool `json:"free"`
|
||||
IsFree bool `json:"is_free"`
|
||||
InputCost float64 `json:"input_cost"`
|
||||
@@ -309,6 +318,8 @@ func (model remoteModel) toModel(key string, source string) Model {
|
||||
containsFold(model.SupportedParameters, "reasoning") ||
|
||||
containsFold(model.SupportedParameters, "reasoning_effort") ||
|
||||
containsFold(model.SupportedParameters, "include_reasoning")
|
||||
reasoningEfforts := model.supportedReasoningEfforts()
|
||||
serviceTiers := normalizedServiceTiers(model.ServiceTiers, model.AdditionalSpeed)
|
||||
|
||||
// Prefer the short display name for aggregator catalogs (OpenRouter's
|
||||
// description field is multi-sentence marketing copy). models.dev and
|
||||
@@ -319,20 +330,56 @@ func (model remoteModel) toModel(key string, source string) Model {
|
||||
}
|
||||
|
||||
return Model{
|
||||
ID: strings.TrimSpace(id),
|
||||
Description: description,
|
||||
ContextWindow: contextWindow,
|
||||
ToolCall: toolCall,
|
||||
Reasoning: reasoning,
|
||||
InputModalities: inputModalities,
|
||||
OutputModalities: outputModalities,
|
||||
InputCost: inputCost,
|
||||
OutputCost: outputCost,
|
||||
Tags: cleanStrings(model.Tags),
|
||||
Source: source,
|
||||
ID: strings.TrimSpace(id),
|
||||
Description: description,
|
||||
ContextWindow: contextWindow,
|
||||
ToolCall: toolCall,
|
||||
Reasoning: reasoning,
|
||||
ReasoningEfforts: reasoningEfforts,
|
||||
DefaultReasoningEffort: strings.TrimSpace(model.DefaultReasoning),
|
||||
ServiceTiers: serviceTiers,
|
||||
DefaultServiceTier: normalizeServiceTier(model.DefaultServiceTier),
|
||||
InputModalities: inputModalities,
|
||||
OutputModalities: outputModalities,
|
||||
InputCost: inputCost,
|
||||
OutputCost: outputCost,
|
||||
Tags: cleanStrings(model.Tags),
|
||||
Source: source,
|
||||
}
|
||||
}
|
||||
|
||||
func normalizedServiceTiers(tiers []string, legacy []string) []string {
|
||||
values := append(append([]string{}, tiers...), legacy...)
|
||||
seen := map[string]bool{}
|
||||
result := make([]string, 0, len(values))
|
||||
for _, value := range values {
|
||||
value = normalizeServiceTier(value)
|
||||
if value == "" || seen[value] {
|
||||
continue
|
||||
}
|
||||
seen[value] = true
|
||||
result = append(result, value)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func normalizeServiceTier(value string) string {
|
||||
value = strings.ToLower(strings.TrimSpace(value))
|
||||
if value == "fast" {
|
||||
return "priority"
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func (model remoteModel) supportedReasoningEfforts() []string {
|
||||
efforts := append([]string{}, model.ReasoningEfforts...)
|
||||
var info remoteReasoningInfo
|
||||
if err := json.Unmarshal(bytes.TrimSpace(model.ReasoningRaw), &info); err == nil {
|
||||
efforts = append(efforts, info.Supported...)
|
||||
}
|
||||
return cleanStrings(efforts)
|
||||
}
|
||||
|
||||
func (model remoteModel) supportsReasoning() bool {
|
||||
raw := bytes.TrimSpace(model.ReasoningRaw)
|
||||
if len(raw) == 0 || string(raw) == "null" {
|
||||
|
||||
@@ -241,6 +241,7 @@ func TestPublicLiveCatalog(t *testing.T) {
|
||||
|
||||
func TestModelsDevProviderIDMapsZeroAliases(t *testing.T) {
|
||||
tests := map[string]string{
|
||||
"chatgpt": "openai",
|
||||
"github": "github-models",
|
||||
"moonshot": "moonshotai",
|
||||
"nvidia-nim": "nvidia",
|
||||
|
||||
@@ -19,20 +19,27 @@ import (
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
)
|
||||
|
||||
const anthropicVersion = "2023-06-01"
|
||||
const (
|
||||
anthropicVersion = "2023-06-01"
|
||||
chatGPTModelsProtocolVersion = "0.146.0"
|
||||
)
|
||||
|
||||
type Model struct {
|
||||
ID string
|
||||
Description string
|
||||
ContextWindow int
|
||||
ToolCall bool
|
||||
Reasoning bool
|
||||
InputModalities []string
|
||||
OutputModalities []string
|
||||
InputCost float64
|
||||
OutputCost float64
|
||||
Tags []string
|
||||
Source string
|
||||
ID string
|
||||
Description string
|
||||
ContextWindow int
|
||||
ToolCall bool
|
||||
Reasoning bool
|
||||
ReasoningEfforts []string
|
||||
DefaultReasoningEffort string
|
||||
ServiceTiers []string
|
||||
DefaultServiceTier string
|
||||
InputModalities []string
|
||||
OutputModalities []string
|
||||
InputCost float64
|
||||
OutputCost float64
|
||||
Tags []string
|
||||
Source string
|
||||
}
|
||||
|
||||
type Options struct {
|
||||
@@ -43,6 +50,7 @@ type Options struct {
|
||||
OAuthResolver providerio.TokenResolver
|
||||
CodexAccountResolver openai.CodexAccountResolver
|
||||
UserAgent string
|
||||
ClientVersion string
|
||||
}
|
||||
|
||||
func DiscoverCatalog(ctx context.Context, provider providercatalog.Descriptor, profile config.ProviderProfile, options Options) ([]Model, error) {
|
||||
@@ -51,6 +59,7 @@ func DiscoverCatalog(ctx context.Context, provider providercatalog.Descriptor, p
|
||||
// without credentials so the picker stays current before a key is entered.
|
||||
canProbeProvider := modelDiscoveryAllowed(profile) && (!provider.RequiresAuth ||
|
||||
discoveryHasCredential(profile) ||
|
||||
options.OAuthResolver != nil ||
|
||||
publicLiveCatalogProvider(provider, profile))
|
||||
if canProbeProvider {
|
||||
liveModels, liveErr := Discover(ctx, profile, options)
|
||||
@@ -204,6 +213,18 @@ func discoverOpenAIModels(ctx context.Context, profile config.ProviderProfile, o
|
||||
}
|
||||
var configure func(*http.Request)
|
||||
if providercatalog.NormalizeID(profile.CatalogID) == "chatgpt" {
|
||||
parsed, parseErr := url.Parse(endpoint)
|
||||
if parseErr != nil {
|
||||
return nil, fmt.Errorf("parse ChatGPT models endpoint: %w", parseErr)
|
||||
}
|
||||
clientVersion := strings.TrimSpace(options.ClientVersion)
|
||||
if clientVersion == "" {
|
||||
clientVersion = chatGPTModelsProtocolVersion
|
||||
}
|
||||
query := parsed.Query()
|
||||
query.Set("client_version", clientVersion)
|
||||
parsed.RawQuery = query.Encode()
|
||||
endpoint = parsed.String()
|
||||
configure = func(request *http.Request) {
|
||||
account := ""
|
||||
if options.CodexAccountResolver != nil {
|
||||
@@ -326,24 +347,39 @@ func anthropicModelsEndpoint(baseURL string) (string, error) {
|
||||
return parsed.String(), nil
|
||||
}
|
||||
|
||||
type modelsResponseItem struct {
|
||||
ID string `json:"id"`
|
||||
Slug string `json:"slug"`
|
||||
DisplayName string `json:"display_name"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Visibility string `json:"visibility"`
|
||||
ContextWindow int `json:"context_window"`
|
||||
ContextWindowAlt int `json:"contextWindow"`
|
||||
ContextLength int `json:"context_length"`
|
||||
MaxContextLength int `json:"max_context_length"`
|
||||
Free bool `json:"free"`
|
||||
IsFree bool `json:"is_free"`
|
||||
ToolCall bool `json:"tool_call"`
|
||||
ToolCallCamel bool `json:"toolCall"`
|
||||
Tools bool `json:"tools"`
|
||||
Reasoning bool `json:"reasoning"`
|
||||
ReasoningEfforts []string `json:"reasoning_efforts"`
|
||||
DefaultReasoning string `json:"default_reasoning_level"`
|
||||
SupportedReasoning []struct {
|
||||
Effort string `json:"effort"`
|
||||
} `json:"supported_reasoning_levels"`
|
||||
ServiceTiers []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"service_tiers"`
|
||||
AdditionalSpeedTiers []string `json:"additional_speed_tiers"`
|
||||
DefaultServiceTier string `json:"default_service_tier"`
|
||||
SupportedParameters []string `json:"supported_parameters"`
|
||||
}
|
||||
|
||||
type modelsResponse struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
DisplayName string `json:"display_name"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
ContextWindow int `json:"context_window"`
|
||||
ContextWindowAlt int `json:"contextWindow"`
|
||||
ContextLength int `json:"context_length"`
|
||||
MaxContextLength int `json:"max_context_length"`
|
||||
Free bool `json:"free"`
|
||||
IsFree bool `json:"is_free"`
|
||||
ToolCall bool `json:"tool_call"`
|
||||
ToolCallCamel bool `json:"toolCall"`
|
||||
Tools bool `json:"tools"`
|
||||
Reasoning bool `json:"reasoning"`
|
||||
SupportedParameters []string `json:"supported_parameters"`
|
||||
} `json:"data"`
|
||||
Data []modelsResponseItem `json:"data"`
|
||||
Models []modelsResponseItem `json:"models"`
|
||||
}
|
||||
|
||||
func parseModelsResponse(body []byte) ([]Model, error) {
|
||||
@@ -352,12 +388,21 @@ func parseModelsResponse(body []byte) ([]Model, error) {
|
||||
return nil, fmt.Errorf("decode models response: %w", err)
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
models := make([]Model, 0, len(payload.Data))
|
||||
for _, item := range payload.Data {
|
||||
id := strings.TrimSpace(item.ID)
|
||||
items := payload.Data
|
||||
if len(items) == 0 {
|
||||
items = payload.Models
|
||||
}
|
||||
models := make([]Model, 0, len(items))
|
||||
for _, item := range items {
|
||||
id := firstNonEmptyDiscovery(item.ID, item.Slug)
|
||||
if id == "" || seen[id] {
|
||||
continue
|
||||
}
|
||||
// The ChatGPT catalog includes hidden/internal entries alongside picker
|
||||
// models. Generic OpenAI-compatible responses omit visibility entirely.
|
||||
if visibility := strings.ToLower(strings.TrimSpace(item.Visibility)); visibility != "" && visibility != "list" {
|
||||
continue
|
||||
}
|
||||
seen[id] = true
|
||||
description := firstNonEmptyDiscovery(
|
||||
item.DisplayName,
|
||||
@@ -381,13 +426,19 @@ func parseModelsResponse(body []byte) ([]Model, error) {
|
||||
discoveryContainsFold(item.SupportedParameters, "reasoning") ||
|
||||
discoveryContainsFold(item.SupportedParameters, "reasoning_effort") ||
|
||||
discoveryContainsFold(item.SupportedParameters, "include_reasoning")
|
||||
efforts := reasoningEfforts(item)
|
||||
serviceTiers := modelServiceTiers(item)
|
||||
models = append(models, Model{
|
||||
ID: id,
|
||||
Description: description,
|
||||
ContextWindow: contextWindow,
|
||||
ToolCall: toolCall,
|
||||
Reasoning: reasoning,
|
||||
Source: "live",
|
||||
ID: id,
|
||||
Description: description,
|
||||
ContextWindow: contextWindow,
|
||||
ToolCall: toolCall,
|
||||
Reasoning: reasoning || len(efforts) > 0,
|
||||
ReasoningEfforts: efforts,
|
||||
DefaultReasoningEffort: strings.TrimSpace(item.DefaultReasoning),
|
||||
ServiceTiers: serviceTiers,
|
||||
DefaultServiceTier: normalizeServiceTier(item.DefaultServiceTier),
|
||||
Source: "live",
|
||||
})
|
||||
}
|
||||
sort.SliceStable(models, func(i, j int) bool {
|
||||
@@ -399,6 +450,50 @@ func parseModelsResponse(body []byte) ([]Model, error) {
|
||||
return models, nil
|
||||
}
|
||||
|
||||
func modelServiceTiers(item modelsResponseItem) []string {
|
||||
values := append([]string{}, item.AdditionalSpeedTiers...)
|
||||
for _, tier := range item.ServiceTiers {
|
||||
values = append(values, tier.ID)
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
result := make([]string, 0, len(values))
|
||||
for _, value := range values {
|
||||
value = normalizeServiceTier(value)
|
||||
if value == "" || seen[value] {
|
||||
continue
|
||||
}
|
||||
seen[value] = true
|
||||
result = append(result, value)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func normalizeServiceTier(value string) string {
|
||||
value = strings.ToLower(strings.TrimSpace(value))
|
||||
if value == "fast" {
|
||||
return "priority"
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func reasoningEfforts(item modelsResponseItem) []string {
|
||||
efforts := append([]string{}, item.ReasoningEfforts...)
|
||||
for _, level := range item.SupportedReasoning {
|
||||
efforts = append(efforts, level.Effort)
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
result := make([]string, 0, len(efforts))
|
||||
for _, effort := range efforts {
|
||||
effort = strings.ToLower(strings.TrimSpace(effort))
|
||||
if effort == "" || seen[effort] {
|
||||
continue
|
||||
}
|
||||
seen[effort] = true
|
||||
result = append(result, effort)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func fetchCatalogModels(ctx context.Context, provider providercatalog.Descriptor, options Options) ([]Model, error) {
|
||||
models, err := providermodelcatalog.FetchRemote(ctx, provider, providermodelcatalog.FetchOptions{
|
||||
HTTPClient: options.HTTPClient,
|
||||
@@ -416,17 +511,21 @@ func modelsFromCatalog(models []providermodelcatalog.Model) []Model {
|
||||
result := make([]Model, 0, len(models))
|
||||
for _, model := range models {
|
||||
result = append(result, Model{
|
||||
ID: model.ID,
|
||||
Description: model.Description,
|
||||
ContextWindow: model.ContextWindow,
|
||||
ToolCall: model.ToolCall,
|
||||
Reasoning: model.Reasoning,
|
||||
InputModalities: append([]string{}, model.InputModalities...),
|
||||
OutputModalities: append([]string{}, model.OutputModalities...),
|
||||
InputCost: model.InputCost,
|
||||
OutputCost: model.OutputCost,
|
||||
Tags: append([]string{}, model.Tags...),
|
||||
Source: model.Source,
|
||||
ID: model.ID,
|
||||
Description: model.Description,
|
||||
ContextWindow: model.ContextWindow,
|
||||
ToolCall: model.ToolCall,
|
||||
Reasoning: model.Reasoning,
|
||||
ReasoningEfforts: append([]string{}, model.ReasoningEfforts...),
|
||||
DefaultReasoningEffort: model.DefaultReasoningEffort,
|
||||
ServiceTiers: append([]string{}, model.ServiceTiers...),
|
||||
DefaultServiceTier: model.DefaultServiceTier,
|
||||
InputModalities: append([]string{}, model.InputModalities...),
|
||||
OutputModalities: append([]string{}, model.OutputModalities...),
|
||||
InputCost: model.InputCost,
|
||||
OutputCost: model.OutputCost,
|
||||
Tags: append([]string{}, model.Tags...),
|
||||
Source: model.Source,
|
||||
})
|
||||
}
|
||||
return result
|
||||
@@ -440,7 +539,9 @@ func mergeLiveModels(provider providercatalog.Descriptor, liveModels []Model, ca
|
||||
hasCatalog := len(byID) > 0
|
||||
// Aggregators publish the live list as the source of truth. Keep live-only
|
||||
// ids even when a remote catalog also loaded, instead of intersecting.
|
||||
preferLive := providermodelcatalog.PublicLiveCatalog(provider.ID)
|
||||
preferLive := provider.Local ||
|
||||
providermodelcatalog.PublicLiveCatalog(provider.ID) ||
|
||||
providercatalog.NormalizeID(provider.ID) == "chatgpt"
|
||||
result := make([]Model, 0, len(liveModels))
|
||||
for _, live := range liveModels {
|
||||
if catalog, ok := byID[live.ID]; ok {
|
||||
@@ -458,6 +559,21 @@ func mergeLiveModels(provider providercatalog.Descriptor, liveModels []Model, ca
|
||||
if strings.TrimSpace(catalog.Description) == "" && strings.TrimSpace(live.Description) != "" {
|
||||
catalog.Description = live.Description
|
||||
}
|
||||
if len(live.ReasoningEfforts) > 0 {
|
||||
catalog.ReasoningEfforts = append([]string{}, live.ReasoningEfforts...)
|
||||
}
|
||||
if live.DefaultReasoningEffort != "" {
|
||||
catalog.DefaultReasoningEffort = live.DefaultReasoningEffort
|
||||
}
|
||||
if live.Reasoning || len(live.ReasoningEfforts) > 0 || live.DefaultReasoningEffort != "" {
|
||||
catalog.Reasoning = true
|
||||
}
|
||||
if len(live.ServiceTiers) > 0 {
|
||||
catalog.ServiceTiers = append([]string{}, live.ServiceTiers...)
|
||||
}
|
||||
if live.DefaultServiceTier != "" {
|
||||
catalog.DefaultServiceTier = live.DefaultServiceTier
|
||||
}
|
||||
catalog.Source = firstDiscoverySource(catalog.Source, "live")
|
||||
result = append(result, catalog)
|
||||
continue
|
||||
@@ -526,17 +642,21 @@ func liveModelAllowedWithoutCatalog(provider providercatalog.Descriptor, id stri
|
||||
|
||||
func catalogModelFromDiscovery(model Model) providermodelcatalog.Model {
|
||||
return providermodelcatalog.Model{
|
||||
ID: model.ID,
|
||||
Description: model.Description,
|
||||
ContextWindow: model.ContextWindow,
|
||||
ToolCall: model.ToolCall,
|
||||
Reasoning: model.Reasoning,
|
||||
InputModalities: append([]string{}, model.InputModalities...),
|
||||
OutputModalities: append([]string{}, model.OutputModalities...),
|
||||
InputCost: model.InputCost,
|
||||
OutputCost: model.OutputCost,
|
||||
Tags: append([]string{}, model.Tags...),
|
||||
Source: model.Source,
|
||||
ID: model.ID,
|
||||
Description: model.Description,
|
||||
ContextWindow: model.ContextWindow,
|
||||
ToolCall: model.ToolCall,
|
||||
Reasoning: model.Reasoning,
|
||||
ReasoningEfforts: append([]string{}, model.ReasoningEfforts...),
|
||||
DefaultReasoningEffort: model.DefaultReasoningEffort,
|
||||
ServiceTiers: append([]string{}, model.ServiceTiers...),
|
||||
DefaultServiceTier: model.DefaultServiceTier,
|
||||
InputModalities: append([]string{}, model.InputModalities...),
|
||||
OutputModalities: append([]string{}, model.OutputModalities...),
|
||||
InputCost: model.InputCost,
|
||||
OutputCost: model.OutputCost,
|
||||
Tags: append([]string{}, model.Tags...),
|
||||
Source: model.Source,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -80,6 +80,175 @@ func TestParseModelsResponseCapturesContextAndFree(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseModelsResponseSupportsChatGPTCatalog(t *testing.T) {
|
||||
models, err := parseModelsResponse([]byte(`{
|
||||
"models": [
|
||||
{"slug":"gpt-5.6-sol","display_name":"GPT-5.6 Sol","description":"Frontier coding model","visibility":"list","context_window":400000,"default_reasoning_level":"high","supported_reasoning_levels":[{"effort":"low"},{"effort":"medium"},{"effort":"high"},{"effort":"xhigh"}],"service_tiers":[{"id":"priority"}],"default_service_tier":"standard"},
|
||||
{"slug":"gpt-5.6-terra","display_name":"GPT-5.6 Terra","visibility":"list"},
|
||||
{"slug":"internal-router","display_name":"Internal","visibility":"hide"},
|
||||
{"slug":"server-only","display_name":"Server only","visibility":"none"}
|
||||
]
|
||||
}`))
|
||||
if err != nil {
|
||||
t.Fatalf("parseModelsResponse: %v", err)
|
||||
}
|
||||
if got, want := modelIDs(models), []string{"gpt-5.6-sol", "gpt-5.6-terra"}; strings.Join(got, ",") != strings.Join(want, ",") {
|
||||
t.Fatalf("models = %#v, want visible ChatGPT models %#v", got, want)
|
||||
}
|
||||
if models[0].Description != "GPT-5.6 Sol" || models[0].ContextWindow != 400000 {
|
||||
t.Fatalf("ChatGPT metadata = %#v", models[0])
|
||||
}
|
||||
if got := strings.Join(models[0].ReasoningEfforts, ","); got != "low,medium,high,xhigh" {
|
||||
t.Fatalf("reasoning efforts = %q", got)
|
||||
}
|
||||
if models[0].DefaultReasoningEffort != "high" || !models[0].Reasoning {
|
||||
t.Fatalf("reasoning metadata = %#v", models[0])
|
||||
}
|
||||
if got := strings.Join(models[0].ServiceTiers, ","); got != "priority" || models[0].DefaultServiceTier != "standard" {
|
||||
t.Fatalf("service tier metadata = %#v", models[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseModelsResponseNormalizesLegacyFastTier(t *testing.T) {
|
||||
models, err := parseModelsResponse([]byte(`{"data":[{"id":"gpt-test","additional_speed_tiers":["fast","priority"]}]}`))
|
||||
if err != nil {
|
||||
t.Fatalf("parseModelsResponse: %v", err)
|
||||
}
|
||||
if got := strings.Join(models[0].ServiceTiers, ","); got != "priority" {
|
||||
t.Fatalf("service tiers = %q, want priority", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeChatGPTModelsKeepsLiveOnlyEntries(t *testing.T) {
|
||||
models := mergeLiveModels(
|
||||
providercatalog.Descriptor{ID: "chatgpt"},
|
||||
[]Model{
|
||||
{ID: "gpt-5.5", ReasoningEfforts: []string{"low", "high"}, ServiceTiers: []string{"priority"}},
|
||||
{ID: "gpt-5.6-sol", ReasoningEfforts: []string{"low", "ultra"}},
|
||||
},
|
||||
[]Model{{ID: "gpt-5.5", Description: "fallback"}},
|
||||
)
|
||||
if got, want := strings.Join(modelIDs(models), ","), "gpt-5.5,gpt-5.6-sol"; got != want {
|
||||
t.Fatalf("models = %q, want %q", got, want)
|
||||
}
|
||||
if got := strings.Join(models[0].ReasoningEfforts, ","); got != "low,high" {
|
||||
t.Fatalf("merged reasoning efforts = %q", got)
|
||||
}
|
||||
if got := strings.Join(models[0].ServiceTiers, ","); got != "priority" {
|
||||
t.Fatalf("merged service tiers = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverCatalogChatGPTMergesOpenAIModelsDevMetadata(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/api.json":
|
||||
_, _ = w.Write([]byte(`{
|
||||
"openai": {"models": {
|
||||
"gpt-5.6-sol": {
|
||||
"id": "gpt-5.6-sol",
|
||||
"name": "GPT-5.6 Sol",
|
||||
"description": "Frontier coding model",
|
||||
"tool_call": true,
|
||||
"reasoning": true,
|
||||
"reasoning_efforts": ["low", "high", "max"],
|
||||
"limit": {"context": 1050000},
|
||||
"modalities": {"input": ["text", "image", "pdf"], "output": ["text"]}
|
||||
}
|
||||
}}
|
||||
}`))
|
||||
case "/backend-api/codex/models":
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"gpt-5.6-sol"}]}`))
|
||||
default:
|
||||
t.Errorf("unexpected request path %q", r.URL.Path)
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
provider, ok := providercatalog.Get("chatgpt")
|
||||
if !ok {
|
||||
t.Fatal("ChatGPT provider missing from catalog")
|
||||
}
|
||||
models, err := DiscoverCatalog(context.Background(), provider, config.ProviderProfile{
|
||||
Name: "chatgpt",
|
||||
CatalogID: "chatgpt",
|
||||
ProviderKind: config.ProviderKindOpenAICompatible,
|
||||
BaseURL: server.URL + "/backend-api/codex",
|
||||
}, Options{
|
||||
HTTPClient: server.Client(),
|
||||
ModelsDevURL: server.URL + "/api.json",
|
||||
OAuthResolver: func(context.Context, bool) (string, string, bool, error) {
|
||||
return "Authorization", "Bearer test", true, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("DiscoverCatalog returned error: %v", err)
|
||||
}
|
||||
if len(models) != 1 {
|
||||
t.Fatalf("models = %#v, want one ChatGPT live model", models)
|
||||
}
|
||||
model := models[0]
|
||||
if model.ID != "gpt-5.6-sol" || model.ContextWindow != 1050000 || !model.ToolCall || !model.Reasoning {
|
||||
t.Fatalf("merged ChatGPT model = %#v, want models.dev metadata", model)
|
||||
}
|
||||
if got := strings.Join(model.InputModalities, ","); got != "text,image,pdf" {
|
||||
t.Fatalf("input modalities = %q, want text,image,pdf", got)
|
||||
}
|
||||
if got := strings.Join(model.ReasoningEfforts, ","); got != "low,high,max" {
|
||||
t.Fatalf("reasoning efforts = %q, want low,high,max", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeLiveModelsUsesLiveDefaultsAndReasoningCapability(t *testing.T) {
|
||||
models := mergeLiveModels(
|
||||
providercatalog.Descriptor{ID: "chatgpt"},
|
||||
[]Model{{
|
||||
ID: "gpt-live",
|
||||
Reasoning: true,
|
||||
DefaultReasoningEffort: "high",
|
||||
DefaultServiceTier: "priority",
|
||||
}},
|
||||
[]Model{{
|
||||
ID: "gpt-live",
|
||||
Reasoning: false,
|
||||
DefaultReasoningEffort: "low",
|
||||
DefaultServiceTier: "standard",
|
||||
}},
|
||||
)
|
||||
if len(models) != 1 {
|
||||
t.Fatalf("merged models = %#v, want one model", models)
|
||||
}
|
||||
got := models[0]
|
||||
if !got.Reasoning || got.DefaultReasoningEffort != "high" || got.DefaultServiceTier != "priority" {
|
||||
t.Fatalf("merged live metadata = %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeLocalModelsKeepsLiveOnlyEntries(t *testing.T) {
|
||||
models := mergeLiveModels(
|
||||
providercatalog.Descriptor{ID: "lmstudio", Local: true},
|
||||
[]Model{
|
||||
{ID: "qwen3-coder-30b", ContextWindow: 32768},
|
||||
{ID: "text-embedding-local"},
|
||||
{ID: "qwen2.5-coder-32b-instruct", Description: "live description"},
|
||||
},
|
||||
[]Model{{ID: "qwen2.5-coder-32b-instruct", Description: "catalog description", ToolCall: true}},
|
||||
)
|
||||
|
||||
if got, want := strings.Join(modelIDs(models), ","), "qwen3-coder-30b,qwen2.5-coder-32b-instruct"; got != want {
|
||||
t.Fatalf("models = %q, want %q", got, want)
|
||||
}
|
||||
for _, model := range models {
|
||||
if model.ID == "qwen2.5-coder-32b-instruct" {
|
||||
if model.Description != "catalog description" || !model.ToolCall {
|
||||
t.Fatalf("merged catalog metadata = %#v", model)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverCatalogOpenGatewayUsesLiveListWithoutKey(t *testing.T) {
|
||||
// Catalog and live endpoints return distinct payloads so the merge must keep
|
||||
// live-only ids that are absent from the remote catalog response.
|
||||
@@ -206,6 +375,9 @@ func TestDiscoverChatGPTModelsUsesOAuthAndCodexHeaders(t *testing.T) {
|
||||
if got := r.URL.Path; got != "/backend-api/codex/models" {
|
||||
t.Errorf("path = %q, want Codex models endpoint", got)
|
||||
}
|
||||
if got := r.URL.Query().Get("client_version"); got != chatGPTModelsProtocolVersion {
|
||||
t.Errorf("client_version = %q, want %q", got, chatGPTModelsProtocolVersion)
|
||||
}
|
||||
wantToken := "Bearer old-token"
|
||||
wantAccount := "old-account"
|
||||
if requests == 2 {
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
// Compaction summarization model selection.
|
||||
//
|
||||
// Compaction summaries used to run on the session's main model — near a 200k
|
||||
// window one summarization is a ~100k-token input call at full price, the
|
||||
// single most expensive recurring event in a long run. A cheap model handles
|
||||
// the "dense factual summary" task fine, so summarization is routed to one
|
||||
// whenever the target is knowable.
|
||||
//
|
||||
// Resolution order: ZERO_COMPACTION_MODEL env > preferences.compactionModel
|
||||
// config > a curated per-kind default (OFFICIAL endpoints only) > "" (use the
|
||||
// main model). The literal value "main" at any level forces the main model.
|
||||
// A wrong choice is never fatal: the agent falls back to the main provider on
|
||||
// the first summarizer failure for the rest of the run.
|
||||
|
||||
// CompactionModelEnv overrides the compaction summarization model.
|
||||
const CompactionModelEnv = "ZERO_COMPACTION_MODEL"
|
||||
|
||||
// compactionMainSentinel disables the dedicated summarizer explicitly.
|
||||
const compactionMainSentinel = "main"
|
||||
|
||||
// Curated cheap summarizer defaults, applied only for provider kinds whose
|
||||
// model catalog is guaranteed (official endpoints). Custom / *-compatible
|
||||
// endpoints get NO default — their catalogs are unknowable, and a guaranteed
|
||||
// failed call per run is worse than main-model prices.
|
||||
const (
|
||||
defaultAnthropicCompactionModel = "claude-haiku-4-5-20251001"
|
||||
defaultGoogleCompactionModel = "gemini-2.5-flash-lite"
|
||||
defaultOpenAICompactionModel = "gpt-4.1-mini"
|
||||
)
|
||||
|
||||
// CompactionModelID returns the model that compaction summarization calls
|
||||
// should use for the given profile, or "" to use the session's main model.
|
||||
// preference is the resolved config value (preferences.compactionModel).
|
||||
func CompactionModelID(profile config.ProviderProfile, preference string) string {
|
||||
value := strings.TrimSpace(os.Getenv(CompactionModelEnv))
|
||||
if value == "" {
|
||||
value = strings.TrimSpace(preference)
|
||||
}
|
||||
if value != "" {
|
||||
if strings.EqualFold(value, compactionMainSentinel) {
|
||||
return ""
|
||||
}
|
||||
return value
|
||||
}
|
||||
return defaultCompactionModel(profile)
|
||||
}
|
||||
|
||||
func defaultCompactionModel(profile config.ProviderProfile) string {
|
||||
metadata, err := ResolveRuntimeMetadata(profile, Options{})
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
var candidate string
|
||||
switch metadata.ProviderKind {
|
||||
case config.ProviderKindAnthropic:
|
||||
candidate = defaultAnthropicCompactionModel
|
||||
case config.ProviderKindGoogle:
|
||||
candidate = defaultGoogleCompactionModel
|
||||
case config.ProviderKindOpenAI:
|
||||
candidate = defaultOpenAICompactionModel
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
// Session already runs on the cheap model: a dedicated summarizer buys
|
||||
// nothing and would only add a second provider instance.
|
||||
if strings.EqualFold(strings.TrimSpace(metadata.APIModel), candidate) {
|
||||
return ""
|
||||
}
|
||||
return candidate
|
||||
}
|
||||
|
||||
// CompactionSummarizerFactory returns the lazy constructor the agent loop
|
||||
// calls for its dedicated compaction summarizer (agent.Options.Summarizer).
|
||||
// The compactor passes the main model in force at that moment — on the first
|
||||
// paid compaction and again after a mid-run model switch — so the choice
|
||||
// tracks escalations: a run that starts on the cheap model gets no dedicated
|
||||
// summarizer, and gains one once it escalates to a pricier main model.
|
||||
// Returning (nil, nil) means "summarize on the main provider". The factory
|
||||
// itself is nil only when no provider builder exists. Shared by the exec and
|
||||
// TUI front ends so the selection rules cannot drift between them.
|
||||
func CompactionSummarizerFactory(base config.ProviderProfile, preference string, newProvider func(config.ProviderProfile) (zeroruntime.Provider, error)) func(context.Context, string) (zeroruntime.Provider, error) {
|
||||
if newProvider == nil {
|
||||
return nil
|
||||
}
|
||||
return func(_ context.Context, mainModelID string) (zeroruntime.Provider, error) {
|
||||
profile := base
|
||||
if strings.TrimSpace(mainModelID) != "" {
|
||||
profile.Model = mainModelID
|
||||
}
|
||||
modelID := CompactionModelID(profile, preference)
|
||||
if modelID == "" {
|
||||
return nil, nil
|
||||
}
|
||||
profile.Model = modelID
|
||||
return newProvider(profile)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/zeroruntime"
|
||||
)
|
||||
|
||||
func TestCompactionModelIDResolutionOrder(t *testing.T) {
|
||||
anthropic := config.ProviderProfile{
|
||||
Name: "anthropic",
|
||||
ProviderKind: config.ProviderKindAnthropic,
|
||||
Model: "claude-sonnet-4.5",
|
||||
}
|
||||
|
||||
t.Setenv(CompactionModelEnv, "my-env-model")
|
||||
if got := CompactionModelID(anthropic, "my-config-model"); got != "my-env-model" {
|
||||
t.Fatalf("env must win, got %q", got)
|
||||
}
|
||||
|
||||
t.Setenv(CompactionModelEnv, "main")
|
||||
if got := CompactionModelID(anthropic, "my-config-model"); got != "" {
|
||||
t.Fatalf("env 'main' must force the main model, got %q", got)
|
||||
}
|
||||
|
||||
t.Setenv(CompactionModelEnv, "")
|
||||
if got := CompactionModelID(anthropic, "my-config-model"); got != "my-config-model" {
|
||||
t.Fatalf("config preference must apply when env unset, got %q", got)
|
||||
}
|
||||
if got := CompactionModelID(anthropic, "main"); got != "" {
|
||||
t.Fatalf("config 'main' must force the main model, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactionModelIDCuratedDefaults(t *testing.T) {
|
||||
t.Setenv(CompactionModelEnv, "")
|
||||
|
||||
anthropic := config.ProviderProfile{Name: "anthropic", ProviderKind: config.ProviderKindAnthropic, Model: "claude-sonnet-4.5"}
|
||||
if got := CompactionModelID(anthropic, ""); got != defaultAnthropicCompactionModel {
|
||||
t.Fatalf("anthropic default = %q, want %q", got, defaultAnthropicCompactionModel)
|
||||
}
|
||||
|
||||
google := config.ProviderProfile{Name: "google", ProviderKind: config.ProviderKindGoogle, Model: "gemini-2.5-pro"}
|
||||
if got := CompactionModelID(google, ""); got != defaultGoogleCompactionModel {
|
||||
t.Fatalf("google default = %q, want %q", got, defaultGoogleCompactionModel)
|
||||
}
|
||||
|
||||
openai := config.ProviderProfile{Name: "openai", ProviderKind: config.ProviderKindOpenAI, Model: "gpt-4.1"}
|
||||
if got := CompactionModelID(openai, ""); got != defaultOpenAICompactionModel {
|
||||
t.Fatalf("openai default = %q, want %q", got, defaultOpenAICompactionModel)
|
||||
}
|
||||
|
||||
// Custom/compatible endpoints: catalog unknowable, no default.
|
||||
compatible := config.ProviderProfile{
|
||||
Name: "opengateway",
|
||||
ProviderKind: config.ProviderKindOpenAICompatible,
|
||||
BaseURL: "https://opengateway.gitlawb.com/v1",
|
||||
Model: "some-model",
|
||||
}
|
||||
if got := CompactionModelID(compatible, ""); got != "" {
|
||||
t.Fatalf("openai-compatible must have no default, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactionModelIDSkipsWhenMainIsAlreadyCheap(t *testing.T) {
|
||||
t.Setenv(CompactionModelEnv, "")
|
||||
haiku := config.ProviderProfile{
|
||||
Name: "anthropic",
|
||||
ProviderKind: config.ProviderKindAnthropic,
|
||||
Model: defaultAnthropicCompactionModel,
|
||||
}
|
||||
if got := CompactionModelID(haiku, ""); got != "" {
|
||||
t.Fatalf("already-cheap session must not get a dedicated summarizer, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactionSummarizerFactoryTracksMainModel(t *testing.T) {
|
||||
t.Setenv(CompactionModelEnv, "")
|
||||
base := config.ProviderProfile{Name: "anthropic", ProviderKind: config.ProviderKindAnthropic, Model: defaultAnthropicCompactionModel}
|
||||
var built []string
|
||||
newProvider := func(profile config.ProviderProfile) (zeroruntime.Provider, error) {
|
||||
built = append(built, profile.Model)
|
||||
return nil, errors.New("not needed")
|
||||
}
|
||||
factory := CompactionSummarizerFactory(base, "", newProvider)
|
||||
if factory == nil {
|
||||
t.Fatal("factory must exist when a provider builder exists")
|
||||
}
|
||||
// Main model is the cheap default: no dedicated summarizer, nothing built.
|
||||
if provider, err := factory(context.Background(), defaultAnthropicCompactionModel); provider != nil || err != nil || len(built) != 0 {
|
||||
t.Fatalf("cheap main model must yield (nil, nil) without building, got %v %v built=%v", provider, err, built)
|
||||
}
|
||||
// After escalating to an expensive model the cheap default is built.
|
||||
if _, err := factory(context.Background(), "claude-sonnet-4-5"); err == nil || len(built) != 1 || built[0] != defaultAnthropicCompactionModel {
|
||||
t.Fatalf("escalated main model must build the cheap summarizer, built=%v err=%v", built, err)
|
||||
}
|
||||
// An empty main model falls back to the base profile's model — the cheap
|
||||
// default here, so again nothing is built.
|
||||
built = nil
|
||||
if provider, err := factory(context.Background(), ""); provider != nil || err != nil || len(built) != 0 {
|
||||
t.Fatalf("empty main model must use the base profile, got %v %v built=%v", provider, err, built)
|
||||
}
|
||||
// The explicit "main" preference disables the dedicated summarizer.
|
||||
if provider, err := CompactionSummarizerFactory(base, "main", newProvider)(context.Background(), "claude-sonnet-4-5"); provider != nil || err != nil {
|
||||
t.Fatalf("preference main must disable the summarizer, got %v %v", provider, err)
|
||||
}
|
||||
if CompactionSummarizerFactory(base, "", nil) != nil {
|
||||
t.Fatal("no provider builder means no factory")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/oauth"
|
||||
"github.com/Gitlawb/zero/internal/providers/providerio"
|
||||
)
|
||||
|
||||
// OAuthLoginForProfile resolves the OAuth login used by profile once and
|
||||
// returns both its bearer resolver and credential-store key. Keeping the key
|
||||
// alongside the resolver lets auxiliary requests derive account-scoped headers
|
||||
// from the same login that supplies the bearer.
|
||||
func OAuthLoginForProfile(profile config.ProviderProfile) (providerio.TokenResolver, string) {
|
||||
candidates := profile.OAuthLoginCandidates()
|
||||
if len(candidates) == 0 {
|
||||
return nil, ""
|
||||
}
|
||||
store, err := oauth.NewStore(oauth.StoreOptions{})
|
||||
if err != nil {
|
||||
return nil, ""
|
||||
}
|
||||
_, key, ok := oauth.FirstStored(store, candidates)
|
||||
if !ok {
|
||||
return nil, ""
|
||||
}
|
||||
manager, err := oauth.NewManager(oauth.ManagerOptions{
|
||||
Store: store,
|
||||
HTTPClient: &http.Client{Timeout: 30 * time.Second},
|
||||
AllowPresets: true,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, ""
|
||||
}
|
||||
resolver := func(ctx context.Context, forceRefresh bool) (string, string, bool, error) {
|
||||
var token string
|
||||
var resolveErr error
|
||||
if forceRefresh {
|
||||
token, resolveErr = manager.Handle401(ctx, key)
|
||||
} else {
|
||||
token, resolveErr = manager.GetFresh(ctx, key)
|
||||
}
|
||||
if errors.Is(resolveErr, oauth.ErrNoToken) {
|
||||
return "", "", false, nil
|
||||
}
|
||||
if resolveErr != nil {
|
||||
return "", "", false, resolveErr
|
||||
}
|
||||
return "Authorization", "Bearer " + token, true, nil
|
||||
}
|
||||
return resolver, key
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package providers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/oauth"
|
||||
)
|
||||
|
||||
func TestOAuthLoginForProfileBindsBearerAndAccountToSameLogin(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "oauth-tokens.json")
|
||||
t.Setenv("ZERO_OAUTH_STORAGE", "file")
|
||||
t.Setenv("ZERO_OAUTH_TOKENS_PATH", path)
|
||||
store, err := oauth.NewStore(oauth.StoreOptions{FilePath: path})
|
||||
if err != nil {
|
||||
t.Fatalf("oauth store: %v", err)
|
||||
}
|
||||
key := oauth.ProviderKey("chatgpt")
|
||||
if err := store.Save(key, oauth.Token{
|
||||
AccessToken: "subscription-token",
|
||||
Account: "account-42",
|
||||
ExpiresAt: time.Now().Add(time.Hour),
|
||||
}); err != nil {
|
||||
t.Fatalf("save OAuth token: %v", err)
|
||||
}
|
||||
|
||||
resolver, loginKey := OAuthLoginForProfile(config.ProviderProfile{Name: "codex", CatalogID: "chatgpt"})
|
||||
if resolver == nil || loginKey != key {
|
||||
t.Fatalf("OAuth login = (%v, %q), want resolver bound to %q", resolver != nil, loginKey, key)
|
||||
}
|
||||
header, value, ok, err := resolver(context.Background(), false)
|
||||
if err != nil || !ok || header != "Authorization" || value != "Bearer subscription-token" {
|
||||
t.Fatalf("bearer resolution = (%q, %q, %v, %v)", header, value, ok, err)
|
||||
}
|
||||
account, ok, err := CodexAccountResolverForLogin(loginKey)(context.Background())
|
||||
if err != nil || !ok || account != "account-42" {
|
||||
t.Fatalf("account resolution = (%q, %v, %v)", account, ok, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOAuthLoginForProfileReturnsNoResolverWithoutUsableLogin(t *testing.T) {
|
||||
t.Setenv("ZERO_OAUTH_STORAGE", "file")
|
||||
t.Setenv("ZERO_OAUTH_TOKENS_PATH", filepath.Join(t.TempDir(), "oauth-tokens.json"))
|
||||
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
profile config.ProviderProfile
|
||||
}{
|
||||
{
|
||||
name: "profile with its own credential has no OAuth candidates",
|
||||
profile: config.ProviderProfile{Name: "openai", CatalogID: "openai", APIKey: "sk-configured"},
|
||||
},
|
||||
{
|
||||
name: "candidate has no stored login",
|
||||
profile: config.ProviderProfile{Name: "chatgpt", CatalogID: "chatgpt"},
|
||||
},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
resolver, key := OAuthLoginForProfile(test.profile)
|
||||
if resolver != nil || key != "" {
|
||||
t.Fatalf("OAuthLoginForProfile() = (%v, %q), want (nil, empty)", resolver != nil, key)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -80,6 +80,7 @@ type responsesRequest struct {
|
||||
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
|
||||
Tools []responsesTool `json:"tools,omitempty"`
|
||||
Reasoning *responsesReasoning `json:"reasoning,omitempty"`
|
||||
ServiceTier string `json:"service_tier,omitempty"`
|
||||
}
|
||||
|
||||
// responsesReasoning carries the reasoning controls for the Responses API. The
|
||||
@@ -265,6 +266,9 @@ func (p *CodexProvider) buildResponsesRequest(request zeroruntime.CompletionRequ
|
||||
// long thinking phase shows live progress instead of looking hung.
|
||||
req.Reasoning = &responsesReasoning{Effort: effort, Summary: "auto"}
|
||||
}
|
||||
if tier := openAIServiceTier(request.ServiceTier); tier != "" {
|
||||
req.ServiceTier = tier
|
||||
}
|
||||
return req, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -492,6 +492,49 @@ func TestCodexProviderForwardsReasoningEffort(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestCodexProviderNormalizesServiceTier(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
input string
|
||||
want string
|
||||
}{
|
||||
{name: "priority", input: "priority", want: "priority"},
|
||||
{name: "flex", input: "flex", want: "flex"},
|
||||
{name: "unsupported omitted", input: "standard", want: ""},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
var rec codexRequest
|
||||
srv := newCodexTestServer(t, &rec)
|
||||
defer srv.Close()
|
||||
|
||||
provider, err := NewCodexProvider(CodexOptions{
|
||||
Options: Options{APIKey: "sk-test", BaseURL: srv.URL, Model: "gpt-5-codex"},
|
||||
AccountID: "acc-x",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewCodexProvider: %v", err)
|
||||
}
|
||||
stream, err := provider.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{
|
||||
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hi"}},
|
||||
ServiceTier: test.input,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion: %v", err)
|
||||
}
|
||||
drainCodexEvents(t, stream)
|
||||
if test.want == "" {
|
||||
if _, ok := rec.body["service_tier"]; ok {
|
||||
t.Fatalf("body.service_tier = %#v, want omitted", rec.body["service_tier"])
|
||||
}
|
||||
return
|
||||
}
|
||||
if got := rec.body["service_tier"]; got != test.want {
|
||||
t.Fatalf("body.service_tier = %#v, want %q", got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCodexProviderStreamsReasoningSummaryDeltas(t *testing.T) {
|
||||
// reasoning_summary_text deltas must surface as StreamEventReasoning (live
|
||||
// "thinking"), in order, alongside the normal text output. Without this a long
|
||||
|
||||
@@ -467,6 +467,9 @@ func (provider *Provider) openAIRequest(request zeroruntime.CompletionRequest) c
|
||||
if effort := openAIReasoningEffort(request.ReasoningEffort); effort != "" {
|
||||
mapped.ReasoningEffort = effort
|
||||
}
|
||||
if tier := openAIServiceTier(request.ServiceTier); tier != "" {
|
||||
mapped.ServiceTier = tier
|
||||
}
|
||||
// prompt_cache_key is a documented OpenAI parameter for server-side prefix
|
||||
// cache routing. Official OpenAI accepts it; many openai-compatible
|
||||
// gateways (NVIDIA NIM, strict local proxies) reject unknown fields with a
|
||||
@@ -483,7 +486,7 @@ func (provider *Provider) openAIRequest(request zeroruntime.CompletionRequest) c
|
||||
Function: toolFunction{
|
||||
Name: tool.Name,
|
||||
Description: tool.Description,
|
||||
Parameters: tool.Parameters,
|
||||
Parameters: normalizeToolParameters(tool.Parameters),
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -491,6 +494,21 @@ func (provider *Provider) openAIRequest(request zeroruntime.CompletionRequest) c
|
||||
return mapped
|
||||
}
|
||||
|
||||
// normalizeToolParameters keeps schemas accepted by strict OpenAI-compatible
|
||||
// servers such as LM Studio. No-argument tools still need an object-valued
|
||||
// parameters.properties field instead of an omitted field.
|
||||
func normalizeToolParameters(parameters map[string]any) map[string]any {
|
||||
normalized := make(map[string]any, len(parameters)+1)
|
||||
for key, value := range parameters {
|
||||
normalized[key] = value
|
||||
}
|
||||
properties, ok := normalized["properties"].(map[string]any)
|
||||
if !ok || properties == nil {
|
||||
normalized["properties"] = map[string]any{}
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
|
||||
// promptCacheKeyDisabled reports whether the ZERO_DISABLE_PROMPT_CACHE_KEY
|
||||
// kill switch is set to a truthy value. "0" and "false" (any case) are
|
||||
// no-ops, matching how ZERO_FORMAT_ON_WRITE parses boolean flags, so an
|
||||
@@ -505,7 +523,16 @@ func promptCacheKeyDisabled() bool {
|
||||
// is dropped rather than risking a 400 on an unrecognized enum.
|
||||
func openAIReasoningEffort(requested string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(requested)) {
|
||||
case "minimal", "low", "medium", "high":
|
||||
case "minimal", "low", "medium", "high", "xhigh", "max", "ultra":
|
||||
return strings.ToLower(strings.TrimSpace(requested))
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func openAIServiceTier(requested string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(requested)) {
|
||||
case "priority", "flex":
|
||||
return strings.ToLower(strings.TrimSpace(requested))
|
||||
default:
|
||||
return ""
|
||||
|
||||
@@ -3,6 +3,7 @@ package openai
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"reflect"
|
||||
@@ -105,6 +106,72 @@ func TestStreamCompletionPostsChatCompletionRequest(t *testing.T) {
|
||||
if tool["type"] != "function" {
|
||||
t.Fatalf("unexpected tool wrapper: %#v", tool)
|
||||
}
|
||||
functionDefinition := tool["function"].(map[string]any)
|
||||
parameters := functionDefinition["parameters"].(map[string]any)
|
||||
properties, ok := parameters["properties"].(map[string]any)
|
||||
if !ok || len(properties) != 0 {
|
||||
t.Fatalf("parameters.properties = %#v, want empty object", parameters["properties"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestStreamCompletionSerializesTypedNilPropertiesAsEmptyObject locks in the
|
||||
// normalizeToolParameters fix: a tool whose parameters carry a typed-nil
|
||||
// properties map must serialize "properties":{} on the wire, never
|
||||
// "properties":null, because strict OpenAI-compatible servers (LM Studio)
|
||||
// reject the null form.
|
||||
func TestStreamCompletionSerializesTypedNilPropertiesAsEmptyObject(t *testing.T) {
|
||||
var gotBody map[string]any
|
||||
var gotRaw []byte
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
raw, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
t.Fatalf("read request body: %v", err)
|
||||
}
|
||||
gotRaw = raw
|
||||
if err := json.Unmarshal(raw, &gotBody); err != nil {
|
||||
t.Fatalf("decode request body: %v", err)
|
||||
}
|
||||
writeSSE(w, `{"choices":[]}`)
|
||||
writeSSE(w, `[DONE]`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
provider, err := New(Options{
|
||||
APIKey: "sk-secret",
|
||||
BaseURL: server.URL + "/",
|
||||
Model: "gpt-test",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New returned error: %v", err)
|
||||
}
|
||||
|
||||
stream, err := provider.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{
|
||||
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hi"}},
|
||||
Tools: []zeroruntime.ToolDefinition{{
|
||||
Name: "no_args",
|
||||
Description: "Takes no arguments",
|
||||
Parameters: map[string]any{"type": "object", "properties": map[string]any(nil)},
|
||||
}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion returned error: %v", err)
|
||||
}
|
||||
drain(stream)
|
||||
|
||||
if !strings.Contains(string(gotRaw), `"properties":{}`) {
|
||||
t.Fatalf("request body must serialize properties as an empty object, got: %s", gotRaw)
|
||||
}
|
||||
if strings.Contains(string(gotRaw), `"properties":null`) {
|
||||
t.Fatalf("request body must not serialize properties as null, got: %s", gotRaw)
|
||||
}
|
||||
tools := gotBody["tools"].([]any)
|
||||
tool := tools[0].(map[string]any)
|
||||
functionDefinition := tool["function"].(map[string]any)
|
||||
parameters := functionDefinition["parameters"].(map[string]any)
|
||||
properties, ok := parameters["properties"].(map[string]any)
|
||||
if !ok || len(properties) != 0 {
|
||||
t.Fatalf("parameters.properties = %#v, want empty object", parameters["properties"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewRequiresModelButNotAPIKey(t *testing.T) {
|
||||
@@ -860,6 +927,51 @@ func TestStreamCompletionSendsReasoningEffort(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestStreamCompletionNormalizesServiceTier(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
input string
|
||||
want string
|
||||
}{
|
||||
{name: "priority", input: "priority", want: "priority"},
|
||||
{name: "flex with surrounding whitespace", input: " FlEx ", want: "flex"},
|
||||
{name: "unsupported omitted", input: "standard", want: ""},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
var gotBody map[string]any
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil {
|
||||
t.Fatalf("decode request body: %v", err)
|
||||
}
|
||||
writeSSE(w, `[DONE]`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
provider, err := New(Options{BaseURL: server.URL + "/", Model: "gpt-test"})
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
stream, err := provider.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{
|
||||
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hi"}},
|
||||
ServiceTier: test.input,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion: %v", err)
|
||||
}
|
||||
drain(stream)
|
||||
if test.want == "" {
|
||||
if _, ok := gotBody["service_tier"]; ok {
|
||||
t.Fatalf("service_tier = %#v, want omitted", gotBody["service_tier"])
|
||||
}
|
||||
return
|
||||
}
|
||||
if got := gotBody["service_tier"]; got != test.want {
|
||||
t.Fatalf("service_tier = %#v, want %q", got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStreamCompletionOmitsReasoningEffortWhenUnsetOrInvalid(t *testing.T) {
|
||||
for _, effort := range []string{"", "none", "bogus"} {
|
||||
var gotBody map[string]any
|
||||
@@ -1434,3 +1546,66 @@ func TestOpenAIRequestPreservesCacheablePrefixAcrossTurns(t *testing.T) {
|
||||
t.Fatalf("wire prompt cache key must remain stable: first=%#v second=%#v", first["prompt_cache_key"], second["prompt_cache_key"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestStreamCompletionSerializesNonMapPropertiesAsEmptyObject locks in the
|
||||
// normalizeToolParameters type-assert branch: when a tool's parameters carry a
|
||||
// `properties` field whose value is a non-map (e.g., an empty slice), it must
|
||||
// be normalized to "properties":{} on the wire, never "properties":[] or
|
||||
// "properties":null.
|
||||
func TestStreamCompletionSerializesNonMapPropertiesAsEmptyObject(t *testing.T) {
|
||||
var gotBody map[string]any
|
||||
var gotRaw []byte
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
raw, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
t.Fatalf("read request body: %v", err)
|
||||
}
|
||||
gotRaw = raw
|
||||
if err := json.Unmarshal(raw, &gotBody); err != nil {
|
||||
t.Fatalf("decode request body: %v", err)
|
||||
}
|
||||
writeSSE(w, `{"choices":[]}`)
|
||||
writeSSE(w, `[DONE]`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
provider, err := New(Options{
|
||||
APIKey: "sk-secret",
|
||||
BaseURL: server.URL + "/",
|
||||
Model: "gpt-test",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("New returned error: %v", err)
|
||||
}
|
||||
|
||||
stream, err := provider.StreamCompletion(context.Background(), zeroruntime.CompletionRequest{
|
||||
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "hi"}},
|
||||
Tools: []zeroruntime.ToolDefinition{{
|
||||
Name: "no_args",
|
||||
Description: "Takes no arguments",
|
||||
Parameters: map[string]any{"type": "object", "properties": []any{}},
|
||||
}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion returned error: %v", err)
|
||||
}
|
||||
drain(stream)
|
||||
|
||||
if !strings.Contains(string(gotRaw), `"properties":{}`) {
|
||||
t.Fatalf("request body must serialize properties as an empty object, got: %s", gotRaw)
|
||||
}
|
||||
if strings.Contains(string(gotRaw), `"properties":null`) {
|
||||
t.Fatalf("request body must not serialize properties as null, got: %s", gotRaw)
|
||||
}
|
||||
if strings.Contains(string(gotRaw), `"properties":[]`) {
|
||||
t.Fatalf("request body must not serialize properties as an empty array, got: %s", gotRaw)
|
||||
}
|
||||
tools := gotBody["tools"].([]any)
|
||||
tool := tools[0].(map[string]any)
|
||||
functionDefinition := tool["function"].(map[string]any)
|
||||
parameters := functionDefinition["parameters"].(map[string]any)
|
||||
properties, ok := parameters["properties"].(map[string]any)
|
||||
if !ok || len(properties) != 0 {
|
||||
t.Fatalf("parameters.properties = %#v, want empty object", parameters["properties"])
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package openai
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestOpenAIReasoningEffortPreservesKnownProviderTiers(t *testing.T) {
|
||||
for _, effort := range []string{"minimal", "low", "medium", "high", "xhigh", "max", "ultra"} {
|
||||
if got := openAIReasoningEffort(effort); got != effort {
|
||||
t.Fatalf("openAIReasoningEffort(%q) = %q", effort, got)
|
||||
}
|
||||
}
|
||||
if got := openAIReasoningEffort("unknown"); got != "" {
|
||||
t.Fatalf("unknown effort = %q, want omitted", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIServiceTierNormalizesKnownValues(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
input string
|
||||
want string
|
||||
}{
|
||||
{input: "priority", want: "priority"},
|
||||
{input: "flex", want: "flex"},
|
||||
{input: " PRIORITY ", want: "priority"},
|
||||
{input: "FlEx", want: "flex"},
|
||||
{input: "", want: ""},
|
||||
{input: "standard", want: ""},
|
||||
} {
|
||||
t.Run(test.input, func(t *testing.T) {
|
||||
if got := openAIServiceTier(test.input); got != test.want {
|
||||
t.Fatalf("openAIServiceTier(%q) = %q, want %q", test.input, got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -153,8 +153,8 @@ func (s *turnSession) observePrefix(ctx context.Context, request zeroruntime.Com
|
||||
// computeFingerprint digests the wire-affecting request parameters: the
|
||||
// session's model and max-tokens (fixed at construction — CompletionRequest
|
||||
// carries no model field), the normalized reasoning effort as it would appear
|
||||
// on the wire, the prompt-cache key (constant within a session, and part of
|
||||
// the request the provider serializes), and an order-preserving digest of the
|
||||
// on the wire, the normalized service tier, the prompt-cache key (constant
|
||||
// within a session, and part of the request the provider serializes), and an order-preserving digest of the
|
||||
// advertised tools (request serialization preserves tool order, so a reorder
|
||||
// changes the wire bytes and counts as drift). Messages are excluded — they
|
||||
// grow every turn; this fingerprints request parameters, not conversation
|
||||
@@ -175,6 +175,8 @@ func (s *turnSession) computeFingerprint(request zeroruntime.CompletionRequest)
|
||||
builder.WriteByte('|')
|
||||
builder.WriteString(openAIReasoningEffort(request.ReasoningEffort))
|
||||
builder.WriteByte('|')
|
||||
builder.WriteString(openAIServiceTier(request.ServiceTier))
|
||||
builder.WriteByte('|')
|
||||
builder.WriteString(request.PromptCacheKey)
|
||||
builder.WriteByte('|')
|
||||
builder.WriteString(wireToolsDigest(request.Tools))
|
||||
|
||||
@@ -231,11 +231,18 @@ func TestTurnSessionFingerprintDriftOnToolsAndEffort(t *testing.T) {
|
||||
// Unrecognized efforts normalize to omitted — same as empty.
|
||||
droppedEffort := session.computeFingerprint(zeroruntime.CompletionRequest{
|
||||
Tools: []zeroruntime.ToolDefinition{{Name: "read_file", Parameters: map[string]any{"type": "object"}}},
|
||||
ReasoningEffort: "xhigh",
|
||||
ReasoningEffort: "bogus",
|
||||
})
|
||||
if base != droppedEffort {
|
||||
t.Fatal("an effort the wire would omit must fingerprint like no effort")
|
||||
}
|
||||
changedTier := session.computeFingerprint(zeroruntime.CompletionRequest{
|
||||
Tools: []zeroruntime.ToolDefinition{{Name: "read_file", Parameters: map[string]any{"type": "object"}}},
|
||||
ServiceTier: "priority",
|
||||
})
|
||||
if base == changedTier {
|
||||
t.Fatal("fingerprint did not drift on a changed service tier")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTurnSessionFingerprintDriftOnToolDescription(t *testing.T) {
|
||||
|
||||
@@ -6,6 +6,7 @@ type chatCompletionRequest struct {
|
||||
Tools []toolDefinition `json:"tools,omitempty"`
|
||||
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
|
||||
ReasoningEffort string `json:"reasoning_effort,omitempty"`
|
||||
ServiceTier string `json:"service_tier,omitempty"`
|
||||
Stream bool `json:"stream"`
|
||||
StreamOptions *streamOptions `json:"stream_options,omitempty"`
|
||||
// PromptCacheKey asks the backend to route the request to a replica that
|
||||
|
||||
@@ -1,14 +0,0 @@
|
||||
// Test seams: helpers only test code uses, kept out of the production binary.
|
||||
package providerio
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"io"
|
||||
)
|
||||
|
||||
// ScanSSEData parses Server-Sent Event data fields from a streaming response.
|
||||
func ScanSSEData(reader io.Reader, handle func(data string) bool) error {
|
||||
scanner := bufio.NewScanner(reader)
|
||||
scanner.Buffer(make([]byte, 0, 4096), maxSSELineBytes)
|
||||
return scanSSEPayloads(scanner, handle, nil)
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package sandbox
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestApplyPatchPathBlockOnlyRejectsRelativeTraversal(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
inside := filepath.Join(root, "main.js")
|
||||
structured := func(path string) string {
|
||||
return strings.Join([]string{"*** Begin Patch", "*** Update File: " + path, "@@", "-a", "+b", "*** End Patch"}, "\n")
|
||||
}
|
||||
for name, patch := range map[string]string{
|
||||
"absolute inside workspace": structured(inside),
|
||||
"relative": structured("main.js"),
|
||||
"decorated markers": "*** Begin Patch ***\n*** Update File: main.js\n@@\n-a\n+b\n*** End Patch ***",
|
||||
"no-space marker": "***Begin Patch\n*** Update File: main.js\n@@\n-a\n+b\n***End Patch",
|
||||
} {
|
||||
request := Request{ToolName: "apply_patch", WorkspaceRoot: root, SideEffect: SideEffectWrite, Args: map[string]any{"patch": patch}}
|
||||
if block := applyPatchPathBlock(request); block != nil {
|
||||
t.Fatalf("%s: unexpected block %+v", name, block)
|
||||
}
|
||||
}
|
||||
for _, path := range []string{"../escape.js", ".."} {
|
||||
for name, patch := range map[string]string{"canonical": structured(path), "no-space": "***Begin Patch\n*** Update File: " + path + "\n@@\n-a\n+b\n***End Patch"} {
|
||||
request := Request{ToolName: "apply_patch", WorkspaceRoot: root, SideEffect: SideEffectWrite, Args: map[string]any{"patch": patch}}
|
||||
block := applyPatchPathBlock(request)
|
||||
if block == nil || block.Code != BlockOutsideWorkspace {
|
||||
t.Fatalf("%s %q must be blocked as traversal, got %+v", name, path, block)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Every marker spelling the tool applies must be classified as structured at
|
||||
// the sandbox boundary too; otherwise the boundary scans the patch as a unified
|
||||
// diff, extracts no targets, and validates nothing (fail-open).
|
||||
func TestStructuredPatchClassifierMatchesToolSpellings(t *testing.T) {
|
||||
for _, header := range []string{"*** Begin Patch", "*** Begin Patch ***", "***Begin Patch", " *** Begin Patch ", "\ufeff*** Begin Patch"} {
|
||||
patch := header + "\n*** Update File: main.js\n@@\n-a\n+b\n*** End Patch"
|
||||
if !IsStructuredPatch(patch) {
|
||||
t.Fatalf("%q must classify as a structured patch", header)
|
||||
}
|
||||
if paths := applyPatchPaths(patch); len(paths) != 1 || paths[0] != "main.js" {
|
||||
t.Fatalf("%q: sandbox must extract the structured target, got %v", header, paths)
|
||||
}
|
||||
}
|
||||
for _, header := range []string{"--- a/x", "Begin Patch", "*** Begin Patchwork", "*** Update File: x"} {
|
||||
if IsStructuredPatch(header + "\n-a\n+b") {
|
||||
t.Fatalf("%q must not classify as a structured patch", header)
|
||||
}
|
||||
}
|
||||
if StructuredPatchMarker("*** End Patch ***") != "end" || StructuredPatchMarker("***End Patch") != "end" {
|
||||
t.Fatal("decorated end markers must classify as end")
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyPatchRequestPathsCarryAbsolutePathsToScopeValidation(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
// NewScope also grants the system temp dir, so a sibling t.TempDir() is
|
||||
// legitimately in scope; pick a path under the filesystem root instead.
|
||||
outside, err := filepath.Abs(filepath.Join(string(filepath.Separator), "zero-outside-workspace-test", "escape.js"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
scope, err := NewScope(root, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
inside := filepath.Join(root, "main.js")
|
||||
structured := func(header, footer, path string) string {
|
||||
return strings.Join([]string{header, "*** Update File: " + path, "@@", "-a", "+b", footer}, "\n")
|
||||
}
|
||||
for name, spelling := range map[string][2]string{"canonical": {"*** Begin Patch", "*** End Patch"}, "no-space": {"***Begin Patch", "***End Patch"}} {
|
||||
// Failure path: the exact path the boundary parsed must be denied by the
|
||||
// scope. structuredPatchHeaderPaths normalises separators to "/", so the
|
||||
// parsed form is compared in slash form and then validated as-is.
|
||||
paths := applyPatchRequestPaths(map[string]any{"patch": structured(spelling[0], spelling[1], outside)})
|
||||
if len(paths) != 1 || paths[0] != filepath.ToSlash(outside) {
|
||||
t.Fatalf("%s: absolute patch path must reach scope validation unchanged, got %v", name, paths)
|
||||
}
|
||||
if block := scope.validate(paths[0]); block == nil || block.Code != BlockOutsideWorkspace {
|
||||
t.Fatalf("%s: scope must deny the parsed outside path %q, got %+v", name, paths[0], block)
|
||||
}
|
||||
// Success path: the parsed inside path must be accepted by the scope.
|
||||
paths = applyPatchRequestPaths(map[string]any{"patch": structured(spelling[0], spelling[1], inside)})
|
||||
if len(paths) != 1 || paths[0] != filepath.ToSlash(inside) {
|
||||
t.Fatalf("%s: inside patch path must reach scope validation unchanged, got %v", name, paths)
|
||||
}
|
||||
if block := scope.validate(paths[0]); block != nil {
|
||||
t.Fatalf("%s: scope must accept the parsed inside path %q, got %+v", name, paths[0], block)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -148,8 +148,10 @@ func unsafeCommandPrefix(prefix []string) bool {
|
||||
return true
|
||||
}
|
||||
}
|
||||
normalized := append([]string(nil), prefix...)
|
||||
normalized[0] = normalizeLauncherName(normalized[0])
|
||||
for _, banned := range bannedCommandPrefixSuggestions {
|
||||
if sameStringSlice(prefix, banned) {
|
||||
if sameStringSlice(prefix, banned) || sameStringSlice(normalized, banned) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
@@ -165,25 +167,118 @@ func unsafeCommandPrefixPart(part string) bool {
|
||||
}
|
||||
|
||||
func unsafeCommandPrefixLauncher(program string) bool {
|
||||
program = strings.ToLower(strings.TrimSpace(program))
|
||||
if strings.ContainsAny(program, `/\`) {
|
||||
// A path is already refused. ":" and "~" are refused for the same reason:
|
||||
// on Windows they introduce names that resolve to a different file than the
|
||||
// one matched here, and no list can keep up with them. "PYTHON~1.EXE" is
|
||||
// python.exe under an 8.3 short name, and 8.3 truncates the stem, so
|
||||
// "POWERS~1.EXE" cannot be recognized by name at all; "python.exe::$DATA"
|
||||
// names the same executable through its default stream. Refusing the shape
|
||||
// costs a permission prompt on names no real command uses.
|
||||
if strings.ContainsAny(program, `/\:~`) {
|
||||
return true
|
||||
}
|
||||
switch program {
|
||||
case "bash", "sh", "zsh", "/bin/bash", "/bin/zsh",
|
||||
"pwsh", "powershell", "powershell.exe",
|
||||
"env", "sudo", "doas", "su", "run0", "osascript",
|
||||
name := normalizeLauncherName(program)
|
||||
return bannedLauncherName(name) || bannedLauncherName(strings.TrimRight(name, "0123456789"))
|
||||
}
|
||||
|
||||
func bannedLauncherName(name string) bool {
|
||||
switch name {
|
||||
case "bash", "sh", "zsh", "dash", "ash", "ksh", "csh", "tcsh", "fish", "busybox",
|
||||
"pwsh", "powershell", "cmd", "wsl",
|
||||
"env", "sudo", "sudoedit", "doas", "su", "run0", "osascript",
|
||||
"command", "eval", "exec", "time",
|
||||
"find", "xargs", "timeout", "nice", "nohup", "watch", "setsid", "stdbuf", "ionice",
|
||||
"ssh", "make", "npm", "npx",
|
||||
"python", "python3", "py", "pythonw", "pyw", "pypy", "pypy3",
|
||||
"node", "perl", "ruby", "php", "lua", "deno", "bun":
|
||||
"ssh", "make", "npm", "npx", "pnpm", "yarn", "bunx",
|
||||
"python", "py", "pythonw", "pyw", "pypy", "uv", "uvx",
|
||||
"node", "nodejs", "perl", "ruby", "php", "lua", "deno", "bun":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// normalizeLauncherName reduces a program name to the launcher it actually runs
|
||||
// so the list above cannot be stepped around by spelling the same interpreter
|
||||
// differently: "python3.11", "python.exe" and "PYTHON3" all reach it as a name
|
||||
// the list matches. internal/agent's commandName already normalizes case and
|
||||
// Windows executable extensions on the allow side; this side matched raw, so
|
||||
// every versioned or .exe-suffixed launcher validated as an ordinary command
|
||||
// and persisted as a prefix grant that auto-allows every later invocation.
|
||||
//
|
||||
// A name is only ever narrowed toward an existing entry, so the failure
|
||||
// direction is an extra permission prompt, never a silent grant.
|
||||
func normalizeLauncherName(program string) string {
|
||||
name := strings.ToLower(strings.TrimSpace(program))
|
||||
// Windows discards trailing dots and spaces when resolving a filename, so
|
||||
// "python." starts python.exe. Drop them before anything matches on the
|
||||
// result, on every platform: the deny side must not depend on which OS is
|
||||
// reading the grant, since the grants file travels with a synced home
|
||||
// directory.
|
||||
name = strings.TrimRight(name, ". ")
|
||||
for _, suffix := range []string{".exe", ".cmd", ".bat", ".com", ".ps1"} {
|
||||
if strings.HasSuffix(name, suffix) {
|
||||
name = strings.TrimSuffix(name, suffix)
|
||||
name = strings.TrimRight(name, ". ")
|
||||
break
|
||||
}
|
||||
}
|
||||
return trimLauncherVersion(trimLauncherABISuffix(trimLauncherBuildSuffix(name)))
|
||||
}
|
||||
|
||||
// trimLauncherBuildSuffix drops the build-channel suffixes distributions and
|
||||
// upstreams append to an otherwise unchanged interpreter, so "python3.11-dbg",
|
||||
// "bash-static" and "pwsh-preview" are the launchers they say they are. The set
|
||||
// is closed on purpose: a general "-word" strip would also swallow
|
||||
// "node-gyp", "python3-config" and "ruby-lsp", which are ordinary tools a user
|
||||
// should still be able to approve.
|
||||
func trimLauncherBuildSuffix(name string) string {
|
||||
for _, suffix := range []string{"-dbg", "-debug", "-static", "-nightly", "-preview", "-beta"} {
|
||||
if strings.HasSuffix(name, suffix) {
|
||||
return strings.TrimSuffix(name, suffix)
|
||||
}
|
||||
}
|
||||
return name
|
||||
}
|
||||
|
||||
// trimLauncherABISuffix drops CPython's ABI flags from a versioned interpreter
|
||||
// name, so "python3.7m" and "python3.6dm" reduce to their version, and the
|
||||
// free-threaded builds "python3.13t" and "python3.13td" reduce to theirs. The
|
||||
// flags are only removed when a digit sits underneath them, which leaves
|
||||
// ordinary names that happen to end in those letters alone: "sha256sum" keeps
|
||||
// its "sum", "zstd" its "td", "cat" its "t".
|
||||
func trimLauncherABISuffix(name string) string {
|
||||
trimmed := strings.TrimRight(name, "dmut")
|
||||
if trimmed == name || trimmed == "" {
|
||||
return name
|
||||
}
|
||||
if last := trimmed[len(trimmed)-1]; last < '0' || last > '9' {
|
||||
return name
|
||||
}
|
||||
return trimmed
|
||||
}
|
||||
|
||||
// trimLauncherVersion drops a trailing "<separator><digits>" version from a
|
||||
// program name, so "python3.11" becomes "python3" and "python2.7" becomes
|
||||
// "python2". Digits without a separator are part of the name, which keeps
|
||||
// "base64", "7z" and "sha256sum" intact; the caller handles the "python3" shape
|
||||
// by also testing the digit-stripped form.
|
||||
func trimLauncherVersion(name string) string {
|
||||
for {
|
||||
trimmed := strings.TrimRight(name, "0123456789")
|
||||
if trimmed == name || trimmed == "" {
|
||||
return name
|
||||
}
|
||||
if !strings.HasSuffix(trimmed, ".") && !strings.HasSuffix(trimmed, "-") {
|
||||
return name
|
||||
}
|
||||
next := trimmed[:len(trimmed)-1]
|
||||
if next == "" {
|
||||
return name
|
||||
}
|
||||
name = next
|
||||
}
|
||||
}
|
||||
|
||||
func sameStringSlice(left []string, right []string) bool {
|
||||
if len(left) != len(right) {
|
||||
return false
|
||||
|
||||
Loaded 100 of 264 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user