Compare commits

...
Author SHA1 Message Date
Vasanthdev2004 6cac6c2abd test(sandbox): make a denied write distinguishable from a runner failure
The denial assertions expected exit 1, which the runner also returns for
its own errors -- a failed marker validation, a bad argument, a token it
could not build. Written that way the test passes both when the sandbox
denies the write and when nothing ever ran, and the second case proves
nothing while being indistinguishable from success.

That is a poor property for any test and a bad one for this test, whose
entire job is to notice a confinement regression. The leading
"granted root is writable" assertion covered part of it, but said nothing
about whether the two later commands launched.

The denial commands now report a distinctive exit code the runner cannot
produce for itself, so the code proves cmd.exe actually ran and its
redirect was refused.

Re-checked in both directions rather than assumed: the test passes against
the fix, and restoring the World SID unconditionally fails it with
"exit code = 0, want 77" -- still catching the bug it was written for.

Raised by CodeRabbit on #865.
2026-08-03 18:24:14 +05:30
Vasanthdev2004 3c3a466f51 fix(sandbox): stop the Windows write jail honouring Everyone-granted paths
The restricted token put the World SID (S-1-1-0, Everyone) in its
restricted-SID list unconditionally, from the original sandbox baseline
and with no comment saying why.

A WRITE_RESTRICTED token runs two checks for a write and needs both: the
normal one against its enabled SIDs, and a second against the restricted
list. The write jail rests entirely on that second check only succeeding
where Zero has explicitly ACL'd a capability SID. Everyone is a SID every
principal carries, so on any path whose DACL grants Everyone write, the
restricted half passed for free and confinement fell back to the ordinary
user's own permissions -- the exact boundary this token exists to be
stricter than.

Reachable with no privilege, no symlink and no race: an Everyone-writable
directory is enough, and share roots opened with Everyone:F and loose
third-party installer ACLs supply them. Ruled out as the common case:
C:\Users\Public\Documents grants BATCH, not Everyone.

The runner already stated the rule this broke -- it declines to add the
user SID, Administrators or SYSTEM because "each has write access nearly
everywhere" -- while listing Everyone among the restricted SIDs.

Removed only from the WRITE_RESTRICTED token, because without that flag it
is load-bearing rather than gratuitous. The restricted-SID check then
applies to reads too, and default Windows DACLs grant BUILTIN\Users rather
than anything in this list, so a token with no Everyone cannot open cmd.exe
and dies at launch with STATUS_ACCESS_DENIED (0xC0000022). Dropping it
there was tried first and takes TestWindowsUnelevatedRealSandboxSmoke with
it. That path is taken only when a profile configures DenyRead, which is
already the posture trading capability for read-deny enforcement (#612), so
the bypass survives there deliberately and narrowly; closing it needs a
read-side grant that is not a universal group.

Verified against real tokens rather than reasoned about. The new smoke test
drives the real runner and asserts three outcomes that only agree when the
jail is intact: a granted root is writable, an ordinary outside path is
denied, and an Everyone-writable outside path is denied too. Before the
change its third case wrote the file with exit 0; the first two were
already correct, which is what made the gap invisible. Restoring the SID
unconditionally turns it red again.
2026-08-03 18:14:44 +05:30
8e266797a8 perf: improve harness context efficiency (#838)
* Improve harness context efficiency

* Address harness efficiency review findings

* Make recaps idle and cancellable

Delay recap generation until the session has been inactive, orient it around the active goal and next task, cancel stale work on activity, and render the result ephemerally above the composer.

Also strengthen the formatter safety regression to exercise formatter-modified edit output.

* Track every replace-all edit span

Translate each replacement into its post-edit line range, including cumulative shifts from earlier replacements, so follow-up edits remain authorized for every model-supplied change.

* Make review regressions cross-platform

Verify replace-all tracking through the canonical edit tool path instead of a platform-sensitive temp path, and cover cancellation when recaps are disabled during generation.

* fix(tools): derive SeenWhole from the recorded ranges

SeenWhole was a flag set only by a SINGLE read covering the file, while
SeenRange judged coverage from the accumulated ranges. The two then
disagreed about identical reads: after reading lines 1-50 and 51-100 of a
100-line file, SeenRange(1,100) was true and SeenWhole was false.

That matters because write_file gates overwrites on SeenWhole. A file past
the read byte budget cannot be read whole in one call — the read is
truncated and RecordSeenRange is correctly skipped, since the model did not
see those lines — so no sequence of reads could ever set the flag, and the
overwrite was refused permanently with:

    the intended change depends on content that has not been read exactly
    in this session. Read the affected lines with read_file, then retry.

Advice that could not change the answer. edit_file was unaffected because it
gates on SeenRange, so the escape was to edit rather than rewrite, but a
model following the error text would loop.

SeenWhole is now computed from the ranges against the file's recorded line
count, so scoped reads that together cover a file make it writable and the
error text becomes actionable. The observation carries that line count, and
a read reporting a different one clears the accumulated ranges: they
describe a shape the file no longer has.

Coverage is judged by merging sorted ranges rather than walking every line
against every range, which the previous SeenRange did — same answer, and it
no longer scans line-by-line over a large file.

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

* test(tools): key the large-file test on the resolved path

read_file resolves the workspace root through EvalSymlinks before recording
anything, so the tracker is keyed on the resolved path. The test asserted
against the raw t.TempDir() path, which is the same string on Windows and a
different one on macOS, where the temp dir sits under /var — a symlink to
/private/var. The SeenWhole assertions would have looked up a key that never
existed: passing here, failing on macOS CI.

Also switches the line generator to fmt.Fprintf (staticcheck QF1012).

Both from CodeRabbit on #842.

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

* test(tools): actually switch the line generator to fmt.Fprintf

The previous commit's message claimed this was done. It was not: the edit
silently did not apply and I reported it fixed without re-reading the file,
so builder.WriteString(fmt.Sprintf(...)) survived at the head — caught by
@anandh8x, not by me.

Verified this time by running the linter rather than by reading the diff:
golangci-lint --enable-only unused,ineffassign,staticcheck ./internal/tools/
reports 0 issues.

* fix: keep exact-read workflows recoverable

* test: cover permission normalization wiring

* fix: commit file observations after output budgeting

* fix: defer file observations across retries

* test: normalize patch line endings

---------

Co-authored-by: Vasanthdev2004 <[email protected]>
Co-authored-by: Claude Opus 5 <[email protected]>
2026-07-31 23:30:03 +05:30
Anandan d37de9214b fix: use OAuth credentials in provider health checks (#828)
* fix: use OAuth credentials in provider health checks

* fix(providerhealth): verify OAuth login for all probes

* fix(providerhealth): retain API key redaction

* fix(providerhealth): redact bare OAuth tokens

* test(cli): isolate OAuth doctor config

* test(cli): isolate doctor config path
2026-07-30 22:32:58 +05:30
Anandan 588635487f feat(tui): replace /retitle with local /rename (#826)
* Improve session naming and resume picker

* Address session rename review feedback
2026-07-30 19:58:02 +05:30
Anandan 845df6c7d2 fix(tui): hide unavailable sidebar shortcut (#821)
* fix(tui): hide unavailable sidebar shortcut

* test(tui): cover medium sidebar footer hint
2026-07-30 17:33:08 +05:30
Anandan 15eb16f736 fix(tui): exclude permission waits from turn time (#820)
* fix(tui): exclude permission waits from turn time

* fix(tui): exclude permission waits from TTFT
2026-07-30 17:03:27 +05:30
github-actions[bot] 097c2652cb chore(main): release 0.6.0 (#786)
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
2026-07-29 09:56:37 +08:00
15b7a5475c fix(acp): interrupt idle reads on cancellation (#782)
* fix(acp): interrupt idle reads on cancellation

* fix(acp): make idle-read interruption per-Read, not per-Serve-call

Addresses the P2 on #782.

The previous design raced a single, whole-connection Once (endRead) between the
read loop's own "I have a terminal result" claim and a watcher goroutine's
"ctx is done, treat this as interrupted" claim. This has a real, deterministic
failure mode that has nothing to do with goroutine-scheduling luck: bufio.Reader
can obtain data AND a terminal error from ONE underlying Read call (io.Reader
explicitly permits returning (n>0, err) together). If a line delimiter is found
within that data, ReadBytes returns the line with a nil error for that call, but
bufio caches the error internally and returns it — without ever touching the
underlying reader again — on the very next ReadBytes call. If cancellation
happens in the gap between that first, successful read and the next one (e.g.
while a synchronously-dispatched handler is still writing its response), the
watcher's ctx.Done() case fires and claims the Once BEFORE the read loop ever
gets back to a ReadBytes call — so when the cached error then surfaces (with
zero contact with the underlying reader, let alone anything closing it), it is
misattributed as a cancellation-driven interruption and silently swallowed as a
clean shutdown, hiding a genuine broken transport.

The fix moves interruption-detection down to the Read call itself instead of
bracketing the whole Serve loop. interruptibleReader wraps the connection's
reader: each Read runs the underlying call on its own goroutine and races it
against ctx.Done(). If the underlying call finishes first, its result is
returned untouched — cancellation a moment later has no bearing on a call that
already had its answer. If ctx.Done() fires first, the reader is closed (when
closable) to nudge a genuinely blocked read, a generation counter is bumped, and
the call still waits for and returns the underlying call's own result rather
than a synthesized one.

Serve then brackets each ReadBytes call (which may invoke the wrapped Read zero
or more times — zero when bufio answers from its own buffer, including a
previously-cached error) with a snapshot of that counter before and after. A
read whose bracket shows no change was never raced against cancellation — either
answered from the buffer or genuinely completed before ctx.Done() could matter —
so its outcome is reported truthfully regardless of the ctx state by the time it
surfaces. Only a bracket that changed proves this specific call actually took
the cancellation branch.

Verified the fix has teeth: reverted to the original implementation and
confirmed the new regression test (which forces the exact bufio-caching
sequence deterministically, with a panic guard proving the underlying reader is
never called a second time) fails 100% of the time against it, then restored
the fix and confirmed it passes.

Also applied the two open review threads on the existing tests: readStarted and
writeStarted are now buffered channels with non-blocking sends rather than
close(), so a future change causing an extra read or write can't panic on a
double-close.

Validated with go build/vet/test on Windows, plus go vet/build and
`go test -race ./internal/acp/... -count=50` (and count=100 focused on the
cancellation/error tests) under WSL/Linux — this session's Windows host has no
C toolchain for -race directly.

* fix(acp): preserve decided read results and require explicit reader ownership

Addresses jatmn's two #782 review findings:

- interruptibleReader.Read's select between resultCh and ctx.Done()
  chooses pseudo-randomly when both are ready, so a real terminal read
  error decided just before cancellation could be misattributed to
  shutdown and silently swallowed. Re-check resultCh non-blockingly
  before claiming a call for cancellation.

- NewConn inferred reader ownership from an io.Closer assertion and
  closed any reader that happened to implement it, including one a
  caller merely borrowed. Split into NewConn (never closes r) and
  NewOwnedConn (explicit ownership, closes r on cancellation), and
  switch the sole production call site (process stdin) to the latter.

Also applies CodeRabbit's atomic.Int64 and leading-ctx-parameter nits.

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

* fix(acp): synchronize read completion with cancellation

Amp-Thread-ID: https://ampcode.com/threads/T-019fa018-9e4e-7506-85ab-e151888bfd71

Co-authored-by: Pierre Bruno <[email protected]>

* fix(acp): preserve read completion across cancellation

Amp-Thread-ID: https://ampcode.com/threads/T-019fa55a-483a-7119-92aa-4953a72f2a76

Co-authored-by: Pierre Bruno <[email protected]>

* test(acp): bound buffered error synchronization

Amp-Thread-ID: https://ampcode.com/threads/T-019fa55a-483a-7119-92aa-4953a72f2a76

Co-authored-by: Pierre Bruno <[email protected]>

* fix(acp): report shutdown transport errors

Amp-Thread-ID: https://ampcode.com/threads/T-019fa7a0-1223-701d-9529-48ba5d7cf8c8

Co-authored-by: Pierre Bruno <[email protected]>

* test(acp): inject cancellation context

Amp-Thread-ID: https://ampcode.com/threads/T-019fa7a0-1223-701d-9529-48ba5d7cf8c8

Co-authored-by: Pierre Bruno <[email protected]>

---------

Co-authored-by: Claude Sonnet 5 <[email protected]>
Co-authored-by: Amp <[email protected]>
2026-07-29 09:46:55 +08:00
bdbda89f3e fix(swarm): close lifecycle admission on shutdown (#776)
* fix(swarm): close lifecycle admission on shutdown

* fix(swarm): admit lifecycle work without holding the lock across Launch

Addresses the P1 on #776 plus the two documentation P3s.

[P1] beginLifecycleAdmission returned with lifecycleMu.RLock() held for the
caller's entire duration — including the call into the injected
MemberLauncher.Launch. Close needs the write lock before it can call s.cancel(),
so a launcher that blocks in Launch waiting on its context (the ordinary shape of
a launcher that waits for a slot, a daemon connection, or a sandbox handshake)
deadlocked every Close caller: the launch held the read lock while Close waited
for the write lock to reach the very cancellation the launch needed to unblock.

beginLifecycleAdmission is now a counted ticket: it holds lifecycleMu only across
the closed check and a lifecycleWork.Add(1), then returns a release func
(lifecycleWork.Done) the caller defers. Close flips closed and releases the lock
BEFORE calling cancel, then waits out lifecycleWork (admitted-but-unfinished
work) followed by watchers (goroutines admitted work may still be starting) — so
it remains a true completion barrier without holding any lock across external
launcher code. Every call site (Spawn, Handoff, AdoptOrphans, the retry path in
watch, and afterExit) moved to the ticket pattern.

[P3] beginLifecycleAdmission's contract and the Scheduler/Close handoff are now
both documented at the point that matters: the func-returning signature makes
the release contract explicit in the type itself, and Scheduler's comment spells
out why the read/write lock ordering guarantees exactly one of Scheduler or
Close closes whatever scheduler ends up installed.

Tests: TestCloseCancelsLauncherBlockedInLaunch reproduces the exact deadlock
(spawn parked in Launch, Close in a second goroutine, both must complete once
cancellation reaches the launcher) — verified it hangs under the old held-lock
pattern by temporarily restoring it. TestCloseWaitsForAdmittedSpawn is the
converse: Close must not return while an admitted spawn is still running.

Validated with `go test -race ./internal/swarm/... -count=1` under WSL/Linux
(this session's Windows host has no C toolchain for -race directly), plus
`go vet ./...` under WSL and go build/vet/test on Windows.

* fix(swarm): fail late queue admissions and dequeues after shutdown

Addresses jatmn's two #776 review findings:

- Close swept (clearQueue) every team's queue before waiting for
  admitted-but-in-flight lifecycle operations to finish. A Spawn,
  Handoff, or AdoptOrphans call that got its ticket right before
  closed flipped could still append to a team's queue after the
  sweep ran, stranding that spec forever (never launched, since
  shutdown began; never failed, since the sweep won't run again).
  Reorder Close to wait for lifecycleWork before sweeping queues, so
  every queue append from already-admitted work is visible to the
  sweep.

- afterExitAdmitted dequeued a queued spec (t.onExit) and launched it
  unconditionally once its own ticket was granted, even though Close
  can flip closed in the gap between that dequeue (which makes the
  spec invisible to the clearQueue sweep) and the launch. Re-check
  closed after dequeuing so a spec pulled off the queue after
  shutdown began fails with ErrSwarmClosed instead of starting a new
  worker.

Both regression tests are verified to fail against the pre-fix code
(confirmed by temporarily reverting each fix and rerunning).

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

* fix(swarm): snapshot teams after lifecycle work stops

Amp-Thread-ID: https://ampcode.com/threads/T-019fa018-a4d5-7078-a80b-802bab23607e

Co-authored-by: Pierre Bruno <[email protected]>

* fix(swarm): reject late admitted launches during shutdown

Re-check the closed state at the final launcher boundary so lifecycle work admitted before Close cannot start or retry a member after shutdown has already won the race.

Add deterministic late-dispatch and scheduler-handoff coverage, stress concurrent Spawn and Close under the race detector, and keep the late-team snapshot regression isolated from launcher behavior.

Amp-Thread-ID: https://ampcode.com/threads/T-019fa55a-4f7f-766e-a312-a7db232b9987

Co-authored-by: Pierre Bruno <[email protected]>

* fix(swarm): reject post-close launch commits

Amp-Thread-ID: https://ampcode.com/threads/T-019fa7a0-1223-701d-9529-48ba5d7cf8c8

Co-authored-by: Pierre Bruno <[email protected]>

* fix(swarm): gate retry relaunch on the post-close commit

The post-Launch shutdown gate added for spawns and queue drains was missing
from the bounded-retry path, which called the launcher and swapped in the new
handle unguarded. Since FuncLauncher returns a live handle immediately without
consulting its context, a Close landing between the pre-check and Launch's
return left the watcher supervising a member relaunched after shutdown began —
and recording that member's outcome on the task.

Extract the commit gate as commitLaunch and route the relaunch through it via
relaunchAdmitted, so an initial launch and a retry pass the same three gates:
admission, the pre-Launch re-check, and the post-Launch commit. When shutdown
wins, the handle is reaped and the watcher records the terminal failure instead
of adopting it.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>

---------

Co-authored-by: Claude Sonnet 5 <[email protected]>
Co-authored-by: Amp <[email protected]>
2026-07-29 09:46:25 +08:00
Vasanth T 8acd42a3bd fix(providers): retry provably pre-send transport failures (#447) (#750)
* fix(providers): retry provably pre-send transport failures (#447)

A flaky link that failed the streaming completion POST at connect time
stopped the turn and forced a manual continue, because the shared retry
layer never retries transport errors: a completion POST is
non-idempotent, so a replay could duplicate billable work.

That rationale does not apply to a PROVABLY pre-send failure, where no
request bytes ever left this host: a replay cannot duplicate anything,
exactly like a 429/503 where the server told us it did not accept the
request. SendWithRetry now retries, with the existing backoff, only that
narrow class: DNS resolution failure, a refused or unreachable dial, and
a TLS handshake timeout. Every ambiguous or post-send failure stays
unretried and is checked first so a mixed message can never be misread:
connection reset, broken pipe, EOF, a bare i/o timeout, context
deadline. All three providers route their completion POST through this
layer, so one change covers them uniformly.

This deliberately narrows the issue's ask (which listed connection
reset) to the provably-safe subset, and keeps the reset-stays-unretried
test green. Tests cover the retry-on-pre-send, no-retry-on-ambiguous,
and predicate-classification paths.

NOTE FOR REVIEW: this refines a deliberately conservative, documented
no-retry stance. Opening it for the owner's sign-off on the stance, not
as a settled call.

* fix(providers): classify pre-send dials by errno so Windows retries too

CodeRabbit flagged that the pre-send predicate matched only POSIX dial
wording, so a refused or unreachable dial on Windows (whose text differs
entirely: 'No connection could be made because the target machine
actively refused it') would not be retried.

Match the syscall errno via errors.Is instead: Go maps ECONNREFUSED /
ENETUNREACH / EHOSTUNREACH / ECONNRESET to the host's real values,
including the WSA codes on Windows, so classification is identical on
every platform. Errno matching is also authoritative where the string
markers were only heuristic, closing the theoretical substring gap. The
string markers stay as a fallback for errors already flattened to text.
ECONNRESET is added to the post-send exclusions by identity so a Windows
reset is excluded too.

Tests build the production error shape (net.OpError -> os.SyscallError ->
errno) so the Windows path is exercised portably on every platform, plus
an end-to-end SendWithRetry case for an errno-wrapped refused dial. Cross
-compiled and vetted for linux and darwin.

* fix(providers): scope pre-send retry to the dial phase and fail fast on permanent errors

Addresses gnanam's two [Major] review findings on this PR.

[Major] Pre-send errno classification was not scoped to the dial, so a post-send
read/write carrying ECONNREFUSED/ENETUNREACH/EHOSTUNREACH (a route dropping
mid-generation) was misclassified "provably pre-send" and could replay a
non-idempotent completion POST up to 6 times. errors.Is matches an errno
anywhere in the chain regardless of which operation raised it, so
isPreSendTransportError now requires an Op=="dial" *net.OpError before admitting
those errnos or their POSIX wording. That makes "no request bytes were sent"
structural: a post-send failure can no longer replay a POST. The dial errno
match still covers Windows (Go maps the WSA* codes), so the platform coverage
this PR added is preserved.

[Major] Permanent failures hung the agent ~60s on the 429-tuned schedule.
NXDOMAIN (net.DNSError with IsNotFound) is now treated as permanent and not
retried, and the pre-send path gets its own short bound and sub-second schedule
(preSendMaxAttempts=3, 500ms/1s) instead of reusing the 2/4/8/16/30s rate-limit
backoff, so a mistyped host or a dead local daemon fails in about 1.5s.

[Minor] Fixed the stale SendWithRetry doc that still claimed transport errors are
never replayed.

Tests: post-send read/write errnos and NXDOMAIN now assert exactly 1 call (no
replay); dial errnos plus the string fallback assert the pre-send cap; added the
pre-send backoff schedule. Race detector and all provider packages green.

* fix(providers): match Windows WSA dial errnos so refused dials retry on Windows too

CodeRabbit caught a real Windows bug, which I verified with a real dial on
Windows: errors.Is(err, syscall.ECONNREFUSED) returns FALSE for a refused dial
there, because a real Windows dial carries the raw WSAECONNREFUSED (10061) while
Go's syscall.ECONNREFUSED is a distinct value the net package never produces. The
POSIX string fallback ("connection refused") does not match the Windows wording
("No connection could be made because the target machine actively refused it")
either. So the pre-send dial classification silently never fired on Windows, the
exact platform this PR set out to fix.

Move the dial errnos into a platform-specific list (dialPreSendErrnos): POSIX
everywhere, plus the WSA* codes on Windows. The existing wrapDialErrno fixtures
wrap the Go syscall constants, which is why they passed on both platforms and
missed this. Adds TestIsPreSendTransportErrorRealRefusedDial, which does a REAL
refused dial and asserts pre-send classification, so the platform's true error
shape is exercised on every OS.

* fix(providers): exclude Plan 9 from the POSIX dial-errno build tag

* fix(providers): keep the pre-send errno classification off Plan 9

retry.go referenced syscall.ECONNRESET and dialPreSendErrnos directly, but the
errno definitions live in build-tagged files that exclude Plan 9 (which has no
BSD-style errno constants), so the package stopped compiling for GOOS=plan9 on
this branch even though it builds on the merge base.

Move the two syscall touch-points behind per-platform helpers: retry.go now
calls isConnResetErrno and ranges dialPreSendErrnos, both provided by
retry_dialerrno.go, retry_dialerrno_windows.go, and a new
retry_dialerrno_plan9.go. The Plan 9 variant supplies an empty errno list and a
false isConnResetErrno, so retry.go falls back to its existing string-marker
heuristics there. Behavior on linux/darwin/windows is unchanged (isConnResetErrno
matches the same syscall.ECONNRESET as before). retry_test.go is tagged !plan9
because it constructs real syscall.Errno dial shapes.

Verified: GOOS=plan9 build and vet now pass; windows and linux build, vet, gofmt,
and the full providerio test suite stay green.

* fix(providers): stop pre-send retry from replaying a redirected POST

Addresses jatmn's review on #750.

[P1] client.Do could follow a 307/308 after already transmitting the original
POST, and a dial failure to the redirect target then satisfied
isPreSendTransportError and replayed the request, re-billing a completion the
first host already received. SendWithRetry now disables redirect following
(CheckRedirect returns http.ErrUseLastResponse on a cloned client), so any
dial/DNS/TLS error from Do can only come from the single initial request and
the pre-send classification is sound. A 3xx surfaces as the response for the
caller's existing non-2xx handling. Regression test asserts the POST is sent
once and the 307 is not followed.

[P2] The pre-send retry used the loop-wide attempt counter shared with 429/503
retries, so a couple of status retries exhausted the pre-send budget and the
first dial/DNS/TLS blip after them got no retry. Track pre-send failures and
their backoff ordinal in a separate counter. Regression test covers a
status-then-transport-failure sequence.

[P3] Correct the Plan 9 and wrapDialErrno comments: the Plan 9 string-marker
fallback keeps the pre-send inclusions and the separate reset exclusion, and
the errno fixtures do not reproduce a real Windows dial (that carries distinct
WSA errnos), which the real-dial test covers.

* fix(providers): decouple the status retry budget from pre-send retries

Addresses jatmn's follow-up on #750.

The previous change gave pre-send failures their own counter, but every retry
still advanced the loop-wide attempt counter that the 429/503 branch used. So
the coupling survived in the opposite direction: two safely replayed refused
dials shortened the later status retry window and started its backoff at a
later rung.

Give the status path its own statusAttempts counter and backoff ordinal. With
both budgets tracked separately the shared loop counter has no remaining
reader, so the loop no longer carries one. Behaviour is unchanged when only one
kind of failure occurs: the first retryable status still waits on rung 1 and
still gets maxAttempts requests.

Adds TestSendWithRetryStatusBudgetSurvivesPreSendRetries, the inverse of the
existing mixed-sequence test: two refused dials followed by rate limiting must
still leave the status path its full window.

* fix(providers): keep following redirects, but never replay a redirected POST

Addresses jatmn's follow-up on #750.

Blocking redirects outright was too blunt. SendWithRetry is the shared
completion path, so replacing every caller's CheckRedirect with
ErrUseLastResponse turned a legitimate 307/308 from a redirecting
OpenAI-compatible gateway into a terminal non-2xx that providers report as an
upstream error, and it discarded an injected client's redirect policy.

Observe the redirect instead of preventing it: a per-attempt flag is set from
CheckRedirect, which then delegates to the caller's policy when one is set and
otherwise mirrors net/http's default ten-hop cap. The pre-send retry is gated on
that flag, so an attempt that entered a redirect is never replayed while
redirecting endpoints keep working. The safety property is unchanged, since a
dial failure on a redirect hop follows a POST the origin already received.

Adds three regressions: a redirect is followed through to the final response, a
dial failure after a redirect is not replayed, and a caller-supplied
CheckRedirect still runs and still decides.

* test(providerio): pin the redirect asymmetry for retryable statuses

A reviewer read the redirect guard as incomplete, because it stops a redirected
POST being replayed after a transport error but not after a 429/503/529. The
asymmetry is intentional, and nothing recorded that, so the next reader would
reasonably reach the same conclusion.

The two cases differ in what is known about the request. A transport failure
after a redirect leaves it unknown whether the redirect target received and
began processing the POST, so replaying could duplicate billable work. A
retryable status means every hop answered: the 307/308 is the origin declining
and delegating, and the 429 is the target declining as well, so nothing accepted
the request and a replay duplicates nothing.

The exposure is also the same as a retryable status with no redirect involved,
which this package has always retried, so guarding here would not close a gap.
It would make a rate-limited completion behind a redirecting gateway fail
outright rather than back off, which is the regression class already raised
against disabling redirects wholesale.

Adds the case as a test asserting the intended behaviour, plus a comment at the
branch explaining why it is not guarded.

* fix(providers): retry connect timeouts and proxy CONNECT failures

Three classification gaps against #447, all raised in review.

A connect-phase timeout was never retried. It is provably pre-send for the same
structural reason a refused dial is, since the connection was never established
and no request bytes left this host, but it arrives carrying "i/o timeout" and
the blanket wording exclusion rejected it before the connect gate could admit
it. Intermittent connect timeouts are exactly the transient failure the issue
is about, so the most common case was the one not covered. It is now admitted
ahead of the wording exclusions, gated on the connect phase so a post-send read
timeout, which is also Timeout(), still cannot replay a POST.

A failed CONNECT through an HTTP proxy was never retried either. net tags it Op
"proxyconnect" rather than "dial", so it missed the gate entirely. Same
guarantee applies: the tunnel to the model host was never built. The gate now
accepts both, which matters most in the environment where transient connect
failures are most common.

ETIMEDOUT and WSAETIMEDOUT join the errno lists so a connect timeout still
classifies if it reaches us with the errno but without the timeout flag.

The pre-send budget also now resets after a hop that reaches the server. It is
meant to bound consecutive connect failures, not the lifetime of the call, and
without the reset an early run of blips permanently spent it, so a later blip
surfaced as a transport error while the status budget still had room. That is
the coupling the two counters were split apart to remove, in the one ordering
the earlier tests missed. Total work stays bounded, since every reset costs a
status attempt and those are capped.

Also corrects the SendWithAuthRetry comment, which still claimed the inner path
retried only 429/503/529.

Tests cover dial and proxyconnect timeouts as retryable and read/write timeouts
as not, a refused proxy CONNECT, and the interleaved pre-send/status/pre-send
sequence. Each fix fails its test when reverted.
2026-07-28 22:42:01 +08:00
Vasanth T 1c38e8cba5 fix(config): make the provider-command timeout an upper bound (#811) (#813)
providerCommandTimeout bounded only the wait, so the call could take arbitrarily
longer than the 5s its error text promised. Two phases sat outside it.

Starting the process ran before the clock began. On Windows that is
CREATE_SUSPENDED, job-object creation and assignment, and a system-wide thread
snapshot, none of it fast on a cold or contended machine. The deadline is now
armed first and the start races it, so process creation is inside the budget.

Draining after termination was unbounded too: the timer path called Terminate
and then received on the wait with no limit. Terminate normally makes that
immediate, which is exactly why it survived unnoticed, but a tree that resists
termination held the call open long past the deadline that had just expired.
That receive is now bounded, and a call that gives up leaves cmd.Wait running,
so the output buffers are synchronised: the I/O pump keeps writing into buffers
the caller already holds, which a plain bytes.Buffer cannot survive.

The timeout error also carried nothing actionable. It was a bare string with the
underlying error and the command's stderr both dropped, so a user whose provider
command was slow got no clue which phase ran out or what the command had said.
It now names the phase, wraps the cause, and includes the redacted output the
failure path already had.

Both bounds are asserted on extracted helpers rather than end to end, because
end to end they are untestable: a terminated command stops writing and its wait
returns at once, so an unbounded drain and an unsynchronised buffer both pass a
whole-call test. The first version of these tests did exactly that and proved
nothing. Removing either bound now fails, as does arming the deadline after the
start, which reports the wrong phase.
2026-07-28 22:41:09 +08:00
Felix-Ayush 58daea7a48 feat(providers): add Fireworks AI provider preset (#772)
* feat(providers): add Fireworks AI provider preset

Add an OpenAI-compatible Fireworks catalog entry with curated coding
models, models.dev alias mapping, onboarding tests, and README setup.

* fix(provider): use correct Fireworks DeepSeek V3 model ID

Update default and catalog model IDs to deepseek-v3-0324 per Fireworks API.
2026-07-28 22:40:41 +08:00
Vasanth T 3e7692d8f0 fix(sandbox): scrub the daemon token file pointer (#677) (#818)
scrubSensitiveEnv removed ZERO_DAEMON_REMOTE_TOKEN but not its file
counterpart. TokenFromEnv accepts either form, so a sandboxed command
still inherited the path to the bridge token, and the default posture is
read-all, so it could then read the token itself.

Scrubbing the pointer is the whole fix rather than half of one: the path
comes from this variable alone and TokenFromEnv has no default location
to fall back on, so a command that never sees the variable has nothing to
guess at. This mirrors GOOGLE_APPLICATION_CREDENTIALS, which is already
scrubbed for the same reason.

Extends the existing scrub table, which compares the surviving
environment exactly, so the case fails if the entry is dropped. Verified:
without the fix the assertion reports ZERO_DAEMON_REMOTE_TOKEN_FILE
surviving into the scrubbed environment.
2026-07-28 21:56:28 +08:00
Vasanth T 33c94ec7a3 fix(sandbox): deny reads of git's credential stores (#815) (#816)
* fix(sandbox): deny reads of git's credential stores (#815)

git's credential-store backend keeps host passwords and personal access tokens
in cleartext, at ~/.git-credentials or ~/.config/git/credentials depending on
whether the user is on the XDG layout. Neither was on the deny list, so any
command the agent runs could read them.

Scoped to those two files deliberately. #815 also covers SSH private keys and
the GPG keyring, and those are a harder trade: denying ~/.ssh stops a sandboxed
git push over SSH from working, and excluding that directory is not even
sufficient, since an IdentityFile directive can point key material anywhere.
These two cost nothing by comparison, because git reads them to authenticate
rather than to know who the user is, so a sandboxed git keeps working and
simply cannot act as the user.

The XDG one denies the file rather than the directory around it, which is the
opposite of the choice made for ~/.config/zero and is deliberate: ~/.config/git
also holds the global git config, which userGitConfigReadPaths grants on
purpose so a sandboxed git can read user.name and aliases instead of failing
outright. Denying the directory would take that away.

Tests assert both stores are denied, that the git config beside the XDG one
stays readable, and that an explicit read grant still re-includes them like
every other entry. Dropping either entry fails the first; widening the XDG
entry to the directory fails the config assertion.

* test(sandbox): assert the credential deny exists before the grant lifts it

The allow-read case only checked that the credential store was absent from
the deny list once an explicit grant covered it. A deny rule that was never
added satisfies that just as well, so the test passed with the production
change reverted and proved nothing about the interaction it was named for.

Check both runs: denied with no grant, absent with one. The first leg is
what establishes there is something for the grant to override.
2026-07-28 18:31:33 +08:00
Vasanth T c9f9dc02ef test(tools): stop internal/tools failing on stock Windows (#814)
* test(tools): stop internal/tools failing on stock Windows

This package fails on every run on a Windows box without PowerShell 7, which
is what a stock install is: PowerShell 7 is a separate download. CI never sees
any of it, because the runners have 7. Three distinct causes, all surfaced by
the move to PowerShell in #804.

The stop-intent test waited a fixed second for a session to die and asserted on
whatever was true by then. Under cmd.exe that was enough; under PowerShell the
tree is slower to stand up and tear down, and the FIRST stop in a fresh process
pays the interpreter's cold start, so it routinely returned with termination
under way but unfinished. The assertion read that as a failed stop. It now
re-sends the stop until the session reports it has exited. Repeating is safe,
since Terminate is a no-op against an already-dead tree, and it keeps the
interrupt flag that a poll with empty chars would drop.

Separately, and pre-existing, every test that starts a session hit a teardown
race: a terminated process does not release its handles the instant the session
reports exited, and Windows refuses to remove a directory with an open handle
where POSIX does not. t.TempDir removes once and fails the test on the sharing
violation, so ordinary teardown timing was reported as a broken test. Those
roots now clean up with a bounded retry, and give up quietly rather than
failing: the process is already terminated, so the worst case is a stale
directory under the OS temp root.

Last, the MSYS guard test asserted a windows_shell_syntax block using a command
containing `&&`. On PowerShell 5.1 the windows_powershell_version preflight
added by #804 fires first and shadows it. The `&&` was incidental to what that
case is about, so it is gone. This one is worth noting as a shape rather than a
typo: a new preflight silently changed which block an older assertion got, and
only on the shell CI does not run.

Verified by running the package repeatedly rather than once: previously failing
on every run, now passing 3 of 3, and the stop-intent test passing 8 of 8 alone
and 6 of 6 under CPU contention. Disabling the terminate path still fails it,
so the assertion continues to bind.

* test(tools): assert the stop intent killed the session

The retry loop added to survive a slow PowerShell teardown also made the
test unfalsifiable. The helper sleeps 5s, well inside the 30s deadline, so
a terminate that did nothing at all would let the sleep end on its own and
every existing assertion would still hold: the session reports exited, it
reports an exit code, and the interrupted flag echoes the request rather
than the outcome. Neutering KillProcessTree confirmed it, the test passed.

Assert instead that the helper's completion line never appears, which
distinguishes a killed session from a finished one without depending on
how long the loop took. Output is accumulated across every poll rather
than read off the last result, because each continuation returns only what
the session produced since the previous collection.
2026-07-28 18:00:41 +08:00
Vasanth T ec3b822505 fix(tui): paste clipboard images with Ctrl+V (#534) (#807)
* fix(tui): paste clipboard images with Ctrl+V

Ctrl+V did nothing when the clipboard held a screenshot, while right-click
paste attached it. The image plumbing was already there: routePaste probes the
OS clipboard for an image whenever a paste arrives with empty text.

The gap is that a clipboard holding only an image produces no bracketed paste
at all, since the terminal has no text to send, so routePaste never runs and its
probe never fires. Right-click reached the probe only because that path always
delivers a clipboardReadMsg, empty or not. Bubble's own Ctrl+V binding stays
disabled, so nothing else was listening for the key.

Handle Ctrl+V by issuing the image probe directly. It reads image content only,
never text, so bracketed paste remains the single text path and a Ctrl+V with
text on the clipboard is unaffected. The probe yields no message when there is
no image, keeping it a silent no-op, and it is gated on noBlockingModal so a
permission prompt or picker still swallows the key.

readClipboardImage becomes a var so the tests stub it instead of depending on
the machine's real clipboard.

Fixes #534

* fix(tui): also probe the clipboard on Cmd+V, and test the real key route

Two things raised in review, both fair.

The handler matched keyCtrl only, which is ModCtrl. macOS reports Command as
ModSuper, as the rest of the TUI already models it, so Cmd+V never reached the
image probe even though the change claimed to cover it. It now matches both, via
a keySuper helper alongside the existing modifier helpers.

Worth being precise about the limit: this only helps where the terminal delivers
the key to the application. A terminal that handles Cmd+V itself pastes the
clipboard TEXT and sends no key event at all, so an image-only clipboard still
produces nothing for this to react to. That is a separate problem from the one
being fixed here.

The tests also did not exercise the route they were named for. One asserted that
Update returned some non-nil command and then ran a freshly built probe; the
other never called Update at all. Both now stub the reader BEFORE Update and run
the command Update actually returned.

That matters more than it sounds. Making the branch return a different non-nil
command leaves the old assertion passing and fails the new one, which is exactly
the hole: the old test proved a command was issued, not that it was the image
probe. Removing the route entirely, or reverting to ModCtrl only, also fails now.

* fix(imageinput): report an unreadable clipboard image instead of silence

On macOS the clipboard reader shells to pngpaste or a PyObjC one-liner, and
neither ships with the OS: pngpaste is Homebrew-only and Apple dropped the
bundled PyObjC. A stock Mac therefore fails every time, and the failure was
returned as (nil, nil), which the caller reads as "no image on the clipboard"
and treats as a silent no-op. Pressing Ctrl+V produced nothing at all: no
image, no error, no hint that anything was missing.

The darwin path already knows the difference. It probes with `clipboard info`
first and only reaches the extraction step once the clipboard is confirmed to
hold an image, so a failure there is not absence, it is "there is an image and
nothing here can read it". That now returns ErrClipboardImageUnreadable, which
the existing error path turns into a transcript notice naming what to install.

The defect predates this branch, but this branch is what makes the path
reachable from Ctrl+V, so shipping without it would close #534 with a feature
that is inert on one of the three platforms. Raised by gnanam, who measured it
on a real Mac rather than inferring it.

Deliberately narrow. The Linux arm cannot tell "no tool" from "no image"
without probing separately, and Windows uses PowerShell, which is always
present, so neither is touched.

The test drives the whole path: stubbed reader returns the error, Ctrl+V
issues the probe, and the assertion is that the failure reaches the transcript
as an error row rather than stopping earlier. Swallowing the error again makes
it fail.

* test(imageinput): cover the darwin unreadable-clipboard decision itself

The test added with the fix asserted around the change rather than on it. It
substituted ReadClipboardImage with a stub returning the error, which proves
the TUI surfaces an error handed to it, but never reaches
readClipboardImageDarwin, where the decision lives. Reverting both production
return sites to nil,nil left every package green, which is the definition of a
test that is not testing the change. Caught by gnanam.

readClipboardImageDarwin shells out twice and had no seam, so the two steps are
now package-level vars, the same shape as the readClipboardImage var this
change already relied on. That makes the case that matters reachable: the info
probe succeeds and reports an image, and the extractor then fails. It cannot be
arranged from outside the function, because from there both steps are one
opaque call.

Five cases, and the negative ones carry as much weight as the positive. An
image with no working extractor, and an extractor that exits 0 with no bytes,
must both be actionable. A clipboard offering no image, and a failing info
probe, must both stay silent, because turning those into a notice would make
every stray Ctrl+V nag the user. The extractor must not run at all in either
silent case.

Repeating gnanam's experiment now fails: reverting both return sites to nil,nil
fails the two actionable cases and leaves the silent ones passing, which is
exactly the discrimination the fix is supposed to add.
2026-07-28 17:51:28 +08:00
110 changed files with 7185 additions and 665 deletions

No files matched your search

+1 -1
View File
@@ -1 +1 @@
{".":"0.5.0"}
{".":"0.6.0"}
+31
View File
@@ -6,6 +6,37 @@ 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.6.0](https://github.com/Gitlawb/zero/compare/v0.5.0...v0.6.0) (2026-07-29)
### Features
* add persistent session goals ([#803](https://github.com/Gitlawb/zero/issues/803)) ([5d1869e](https://github.com/Gitlawb/zero/commit/5d1869e7bc193c97e07418b57b823905c1151d5d))
* **providers:** add Fireworks AI provider preset ([#772](https://github.com/Gitlawb/zero/issues/772)) ([58daea7](https://github.com/Gitlawb/zero/commit/58daea7a482927e5ab35afaf97761c60bb35e113))
* **skills:** report hash drift in skill info ([#795](https://github.com/Gitlawb/zero/issues/795)) ([9fbdc34](https://github.com/Gitlawb/zero/commit/9fbdc341551601c13434cf7b9b6f5e4ea4c483b1))
### Bug Fixes
* **acp:** interrupt idle reads on cancellation ([#782](https://github.com/Gitlawb/zero/issues/782)) ([15b7a54](https://github.com/Gitlawb/zero/commit/15b7a5475cfeefddbc59634551451a8872b996b7))
* **cli:** include plugins info in shell completions ([#794](https://github.com/Gitlawb/zero/issues/794)) ([fa8734f](https://github.com/Gitlawb/zero/commit/fa8734f17927de8832dfe041e043cfb86f93d12a))
* **config:** kill provider-command process tree via job object on Windows ([#690](https://github.com/Gitlawb/zero/issues/690)) ([5c58256](https://github.com/Gitlawb/zero/commit/5c58256abc29c185b449f8ce79647eb3fae61882))
* **config:** make the provider-command timeout an upper bound ([#811](https://github.com/Gitlawb/zero/issues/811)) ([#813](https://github.com/Gitlawb/zero/issues/813)) ([1c38e8c](https://github.com/Gitlawb/zero/commit/1c38e8cba50db255f3cda8b7cf4298712b2933e1))
* **config:** stop billing provider-command fixtures to the 5s timeout ([#810](https://github.com/Gitlawb/zero/issues/810)) ([f5fa4ff](https://github.com/Gitlawb/zero/commit/f5fa4ffc64537897aa78a56edca36d8f1dacd021))
* **cron:** reserve job IDs atomically ([#686](https://github.com/Gitlawb/zero/issues/686)) ([d9b882e](https://github.com/Gitlawb/zero/commit/d9b882eb0a94f38de8447cb54b57860ccd8ed053))
* **lint:** resolve static analysis baseline ([#769](https://github.com/Gitlawb/zero/issues/769)) ([583e653](https://github.com/Gitlawb/zero/commit/583e6530bd849a8efcb2b88bbe8c297e8592e62c))
* **providers:** retry provably pre-send transport failures ([#447](https://github.com/Gitlawb/zero/issues/447)) ([#750](https://github.com/Gitlawb/zero/issues/750)) ([8acd42a](https://github.com/Gitlawb/zero/commit/8acd42a3bd743a35807d762f23a33fb61c171839))
* **sandbox:** add platform-specific write lock around grant state file ([#752](https://github.com/Gitlawb/zero/issues/752)) ([#755](https://github.com/Gitlawb/zero/issues/755)) ([88fb7b6](https://github.com/Gitlawb/zero/commit/88fb7b6192c1ac8173089967d5e9119f92d7a101))
* **sandbox:** clarify scoped permission prompts ([#798](https://github.com/Gitlawb/zero/issues/798)) ([9fe02ae](https://github.com/Gitlawb/zero/commit/9fe02ae52e3234cb9dae76e0c5594a287ccbb52a))
* **sandbox:** deny reads of git's credential stores ([#815](https://github.com/Gitlawb/zero/issues/815)) ([#816](https://github.com/Gitlawb/zero/issues/816)) ([33c94ec](https://github.com/Gitlawb/zero/commit/33c94ec7a3e7987e6e36b5fd5e037ff4e6e84c26))
* **sandbox:** preserve user config and retry denied writes ([#801](https://github.com/Gitlawb/zero/issues/801)) ([81a5e6c](https://github.com/Gitlawb/zero/commit/81a5e6cb716ba4b7ecf7ff3ff261dfa51fd01098))
* **sandbox:** scrub the daemon token file pointer ([#677](https://github.com/Gitlawb/zero/issues/677)) ([#818](https://github.com/Gitlawb/zero/issues/818)) ([3e7692d](https://github.com/Gitlawb/zero/commit/3e7692d8f0885d346809869935c81a9cf055d89e))
* **swarm:** close lifecycle admission on shutdown ([#776](https://github.com/Gitlawb/zero/issues/776)) ([bdbda89](https://github.com/Gitlawb/zero/commit/bdbda89f3e42d93e8058ab3e1e9ef402659530cd))
* **tui:** fix rendering corruption over multipass + Windows Terminal ([#709](https://github.com/Gitlawb/zero/issues/709)) ([5bb1bc1](https://github.com/Gitlawb/zero/commit/5bb1bc11c3188dcea7709ff07c8b25dda3541de8))
* **tui:** paste clipboard images with Ctrl+V ([#534](https://github.com/Gitlawb/zero/issues/534)) ([#807](https://github.com/Gitlawb/zero/issues/807)) ([ec3b822](https://github.com/Gitlawb/zero/commit/ec3b8225050eb62f0f41a9582ed2881d2f9dcc8e))
* **tui:** reset quiet clock for new runs ([#799](https://github.com/Gitlawb/zero/issues/799)) ([a4a805a](https://github.com/Gitlawb/zero/commit/a4a805a6c82c51cb255e909bd87ca183c8805cf2))
* **windows:** use PowerShell for command execution ([#804](https://github.com/Gitlawb/zero/issues/804)) ([518d2e5](https://github.com/Gitlawb/zero/commit/518d2e53ba3fa073090a9b41c278face85b59b64))
## [0.5.0](https://github.com/Gitlawb/zero/compare/v0.4.0...v0.5.0) (2026-07-22)
+7
View File
@@ -131,6 +131,7 @@ export ANTHROPIC_API_KEY=...
export GEMINI_API_KEY=...
export AIMLAPI_API_KEY=...
export LONGCAT_API_KEY=...
export FIREWORKS_API_KEY=...
export MINIMAX_API_KEY=...
export MINIMAXI_API_KEY=...
```
@@ -147,6 +148,12 @@ To configure Meituan LongCat (LongCat-2.0) directly, run:
zero providers setup longcat --set-active
```
To configure Fireworks AI directly, run:
```bash
zero providers setup fireworks --set-active
```
MiniMax presets use the Anthropic-compatible endpoints for the global and China
regions:
+138 -6
View File
@@ -17,6 +17,7 @@ import (
"fmt"
"io"
"sync"
"sync/atomic"
)
// JSON-RPC 2.0 standard error codes.
@@ -80,8 +81,9 @@ type NotifyFunc func(ctx context.Context, params json.RawMessage)
// requests/notifications — needed because ACP is bidirectional (the agent calls
// the client for session/request_permission, fs/*, terminal/*).
type Conn struct {
reader *bufio.Reader
w io.Writer
rawReader io.Reader // wrapped lazily in Serve, once ctx is known — see interruptibleReader
readerCloser io.Closer
w io.Writer
writeMu sync.Mutex // serializes all writes to w
@@ -96,10 +98,15 @@ type Conn struct {
wg sync.WaitGroup // tracks in-flight inbound handlers
}
// NewConn builds a peer reading ndjson from r and writing ndjson to w.
// NewConn builds a peer reading ndjson from r and writing ndjson to w. This
// Conn does not own r: cancelling Serve's context cannot forcibly close it, so
// a Read on r that's genuinely idle (never returns on its own) leaves Serve
// blocked until r itself produces something — e.g. a test that closes r's own
// write side to unblock it. Use NewOwnedConn when this Conn should own r's
// lifetime.
func NewConn(r io.Reader, w io.Writer) *Conn {
return &Conn{
reader: bufio.NewReader(r),
rawReader: r,
w: w,
handlers: make(map[string]HandlerFunc),
notifiers: make(map[string]NotifyFunc),
@@ -107,6 +114,18 @@ func NewConn(r io.Reader, w io.Writer) *Conn {
}
}
// NewOwnedConn is like NewConn but declares that this Conn exclusively owns r:
// if r implements io.Closer, Serve closes it on context cancellation to
// interrupt an otherwise-idle blocking Read. Only pass a reader whose lifetime
// this Conn should control — e.g. the process's own stdin — never a reader a
// caller retains or shares elsewhere, since cancellation tears it down
// permanently for every consumer, not just this Conn.
func NewOwnedConn(r io.Reader, w io.Writer) *Conn {
c := NewConn(r, w)
c.readerCloser, _ = r.(io.Closer)
return c
}
// Handle registers a request handler for method.
func (c *Conn) Handle(method string, fn HandlerFunc) { c.handlers[method] = fn }
@@ -128,17 +147,36 @@ func (c *Conn) Serve(ctx context.Context) error {
c.wg.Wait()
}()
interruptible := newInterruptibleReader(ctx, c.rawReader, c.readerCloser)
reader := bufio.NewReader(interruptible)
for {
// One ndjson value per line. Framing per line (rather than streaming the
// decoder) keeps a single malformed line from making the whole connection
// unrecoverable — we report -32700 and continue.
line, err := c.reader.ReadBytes('\n')
//
// generation brackets this SPECIFIC ReadBytes call: bufio may invoke
// interruptible.Read zero or more times underneath it (zero when the
// answer is already sitting in bufio's own buffer — including a
// previously-buffered, not-yet-surfaced error: bufio can return a
// complete line with a nil error while quietly caching an error it
// received ALONGSIDE that line's bytes, which then surfaces on the VERY
// NEXT call with no further contact with interruptible.Read at all). The
// generation only advances inside interruptible.Read's own cancellation
// branch, so "unchanged across this call" is proof this call's outcome
// didn't pass through that branch — whether because it was answered from
// the buffer or because it raced ctx.Done() and won. Only that proof
// makes it safe to call the outcome genuine.
before := interruptible.generation()
line, err := reader.ReadBytes('\n')
interrupted := interruptible.generation() != before
if len(bytes.TrimSpace(line)) > 0 {
c.handleLine(ctx, line)
}
if err != nil {
c.failAllPending(err)
if errors.Is(err, io.EOF) || ctx.Err() != nil {
if errors.Is(err, io.EOF) || interrupted {
return nil
}
return err
@@ -146,6 +184,100 @@ func (c *Conn) Serve(ctx context.Context) error {
}
}
// interruptibleReader wraps a reader so a context cancellation can unblock a
// currently in-flight Read without ever discarding — or misattributing — the
// real outcome of a call that wasn't actually blocked.
//
// Each Read runs the underlying call on its own goroutine and races it against
// ctx.Done(). If the underlying call finishes first, its result is returned
// exactly as received; ctx being cancelled a moment later has no bearing on a
// call that already had its answer. If ctx.Done() fires first, closer (when
// non-nil) is closed to nudge a genuinely blocked read, the generation counter
// is bumped, and the call STILL waits for and returns the underlying call's own
// result — never a synthesized one — since closing a reader's outcome is
// whatever the reader itself reports for it, not ours to invent.
//
// The counter is what lets Serve tell a call that was actually raced against
// cancellation apart from one merely answered while cancellation happened to be
// pending — the distinction the read loop's own doc comment explains.
type interruptibleReader struct {
inner io.Reader
closer io.Closer
ctx context.Context
closeOnce sync.Once
gen atomic.Int64
}
func newInterruptibleReader(ctx context.Context, inner io.Reader, closer io.Closer) *interruptibleReader {
return &interruptibleReader{inner: inner, closer: closer, ctx: ctx}
}
func (r *interruptibleReader) generation() int64 { return r.gen.Load() }
type interruptibleReadResult struct {
n int
err error
}
// interruptibleReadDecision arbitrates whether an underlying Read completed or
// cancellation claimed it first. Completion is recorded before the result is
// published so cancellation cannot misclassify a finished read merely because
// its goroutine has not sent the result yet.
type interruptibleReadDecision struct {
once sync.Once
}
func (d *interruptibleReadDecision) complete() { d.once.Do(func() {}) }
func (d *interruptibleReadDecision) interrupt() (claimed bool) {
d.once.Do(func() { claimed = true })
return claimed
}
// readInner records completion in the Read call's deferred epilogue. It also
// converts a reader panic into a terminal error so the helper goroutine always
// publishes a result for Read to join.
func (r *interruptibleReader) readInner(p []byte, decision *interruptibleReadDecision) (res interruptibleReadResult) {
defer func() {
if recovered := recover(); recovered != nil {
res = interruptibleReadResult{err: fmt.Errorf("acp: reader panicked: %v", recovered)}
}
decision.complete()
}()
res.n, res.err = r.inner.Read(p)
return res
}
func (r *interruptibleReader) Read(p []byte) (int, error) {
resultCh := make(chan interruptibleReadResult, 1)
var decision interruptibleReadDecision
go func() {
resultCh <- r.readInner(p, &decision)
}()
select {
case res := <-resultCh:
return res.n, res.err
case <-r.ctx.Done():
// resultCh may not be readable yet even when inner.Read has returned: its
// goroutine can be descheduled between recording completion and publishing
// the result. Only interrupt a call that cancellation actually claims.
if decision.interrupt() {
r.gen.Add(1)
if r.closer != nil {
r.closeOnce.Do(func() { _ = r.closer.Close() })
}
}
// Wait for the real result rather than returning synthetically: p is
// still being written by the goroutine above until it does, so returning
// early would race that write, and the actual error (whatever closing
// the reader produced, or whatever the reader was already about to
// report) is what the caller should see.
res := <-resultCh
return res.n, res.err
}
}
// handleLine decodes and dispatches one ndjson frame. A parse failure replies
// with -32700 (id null) and keeps the connection alive.
func (c *Conn) handleLine(ctx context.Context, line []byte) {
+288
View File
@@ -3,11 +3,30 @@ package acp
import (
"context"
"encoding/json"
"errors"
"io"
"sync"
"sync/atomic"
"testing"
"time"
)
type testReadCloser struct {
read func([]byte) (int, error)
close func() error
}
func (r testReadCloser) Read(p []byte) (int, error) { return r.read(p) }
func (r testReadCloser) Close() error { return r.close() }
type testReader func([]byte) (int, error)
func (r testReader) Read(p []byte) (int, error) { return r(p) }
type testWriter func([]byte) (int, error)
func (w testWriter) Write(p []byte) (int, error) { return w(p) }
// connPair wires two Conns together over in-memory pipes and serves both.
func connPair(t *testing.T) (a, b *Conn, stop func()) {
t.Helper()
@@ -25,6 +44,275 @@ func connPair(t *testing.T) (a, b *Conn, stop func()) {
}
}
func TestConnServeCancellationInterruptsIdleRead(t *testing.T) {
pipeReader, writer := io.Pipe()
t.Cleanup(func() { _ = writer.Close() })
// Buffered + non-blocking send, not close: interruptibleReader retries each
// underlying Read on its own goroutine, so if bufio ever calls this more than
// once, a close() here would panic on the second call.
readStarted := make(chan struct{}, 1)
closeCalls := make(chan struct{}, 2)
reader := testReadCloser{
read: func(p []byte) (int, error) {
select {
case readStarted <- struct{}{}:
default:
}
return pipeReader.Read(p)
},
close: func() error {
closeCalls <- struct{}{}
return pipeReader.Close()
},
}
conn := NewOwnedConn(reader, io.Discard)
ctx, cancel := context.WithCancel(context.Background())
done := make(chan error, 1)
go func() { done <- conn.Serve(ctx) }()
<-readStarted
cancel()
select {
case err := <-done:
if err != nil {
t.Fatalf("Serve returned %v after cancellation, want nil", err)
}
case <-time.After(time.Second):
t.Fatal("Serve did not return after cancelling an idle connection")
}
if got := len(closeCalls); got != 1 {
t.Fatalf("reader Close calls = %d, want 1", got)
}
}
func TestConnServePreservesTerminalReadErrorDuringCancellation(t *testing.T) {
for _, closable := range []bool{false, true} {
name := "non-closable"
if closable {
name = "closable"
}
t.Run(name, func(t *testing.T) {
wantErr := errors.New("read failed")
read := testReader(func(p []byte) (int, error) {
return copy(p, "not json"), wantErr
})
closeCalls := make(chan struct{}, 1)
var reader io.Reader = read
if closable {
reader = testReadCloser{
read: read,
close: func() error {
closeCalls <- struct{}{}
return nil
},
}
}
// Buffered + non-blocking send, not close: a later change adding a
// second write (e.g. an additional error frame) must not panic on a
// repeated close of an already-closed channel.
writeStarted := make(chan struct{}, 1)
releaseWrite := make(chan struct{})
writer := testWriter(func(p []byte) (int, error) {
select {
case writeStarted <- struct{}{}:
default:
}
<-releaseWrite
return len(p), nil
})
newConn := NewConn
if closable {
// Ownership matters here: this case exists to prove that even a
// Conn entitled to close its reader on cancellation still
// preserves a terminal error that was already decided.
newConn = NewOwnedConn
}
conn := newConn(reader, writer)
ctx, cancel := context.WithCancel(context.Background())
done := make(chan error, 1)
go func() { done <- conn.Serve(ctx) }()
<-writeStarted
cancel()
close(releaseWrite)
select {
case err := <-done:
if !errors.Is(err, wantErr) {
t.Fatalf("Serve returned %v, want terminal read error %v", err, wantErr)
}
case <-time.After(time.Second):
t.Fatal("Serve did not return after terminal read error")
}
if got := len(closeCalls); got != 0 {
t.Fatalf("reader Close calls = %d, want 0 after terminal read", got)
}
})
}
}
// TestConnServePreservesReadErrorBufferedAlongsideAPriorLine is the deterministic
// regression test for jatmn's #782 finding: a real, non-EOF ReadBytes failure can
// be DECIDED before cancellation ever happens, yet only SURFACE afterward, and
// must still be reported rather than swallowed as a clean shutdown.
//
// bufio.Reader can obtain data AND a terminal error from ONE underlying Read call
// (io.Reader explicitly permits returning (n>0, err) together). If a delimiter is
// found within that data, ReadBytes returns the line with a nil error for THIS
// call, but bufio caches the error internally and returns it — WITHOUT calling the
// underlying reader again — on the very next ReadBytes call. This test forces that
// exact sequence: the underlying reader hands back one complete, validly-framed
// line together with wantErr in a SINGLE call (asserted below to be its ONLY
// call), the loop dispatches that line through a synchronous path that blocks the
// read loop, cancellation is asserted to have happened before the loop attempts
// its next read, and only then is the block released. By the time the read loop
// asks for more input, cancellation is already in effect and the underlying
// reader is never touched again — so if Serve attributed the resulting error to
// cancellation merely because ctx was already done, it would incorrectly return
// nil. It must instead recognize that this read was answered entirely from
// bufio's own cache, never raced against ctx.Done(), and report wantErr.
func TestConnServePreservesReadErrorBufferedAlongsideAPriorLine(t *testing.T) {
wantErr := errors.New("read failed")
var readCalls atomic.Int32
read := testReader(func(p []byte) (int, error) {
readCalls.Add(1)
// An unsupported jsonrpc version with an id takes the SYNCHRONOUS
// writeError path in handleLine (not a dispatched goroutine), giving this
// test a reliable blocking point between processing this line and the
// loop's next ReadBytes call.
return copy(p, `{"jsonrpc":"1.0","id":1}`+"\n"), wantErr
})
writeStarted := make(chan struct{}, 1)
releaseWrite := make(chan struct{})
var releaseWriteOnce sync.Once
release := func() { releaseWriteOnce.Do(func() { close(releaseWrite) }) }
t.Cleanup(release)
writer := testWriter(func(p []byte) (int, error) {
select {
case writeStarted <- struct{}{}:
default:
}
<-releaseWrite
return len(p), nil
})
conn := NewConn(read, writer)
ctx, cancel := context.WithCancel(context.Background())
done := make(chan error, 1)
go func() { done <- conn.Serve(ctx) }()
// The synchronous write confirms the first (and, per the panic guard above,
// only) read has already returned and is being processed — cancelling now
// lands squarely in the gap between that read and the loop's next one, never
// racing an in-flight interruptibleReader.Read call.
select {
case <-writeStarted:
case <-time.After(time.Second):
t.Fatal("Serve did not reach the synchronous write")
}
cancel()
release()
select {
case err := <-done:
if !errors.Is(err, wantErr) {
t.Fatalf("Serve returned %v, want the buffered terminal read error %v", err, wantErr)
}
case <-time.After(time.Second):
t.Fatal("Serve did not return after the buffered terminal read error")
}
if got := readCalls.Load(); got != 1 {
t.Fatalf("underlying Read calls = %d, want exactly 1 (bufio must have served the second ReadBytes from its cache)", got)
}
}
// TestInterruptibleReaderPrefersDecidedResultOverCancellation is the regression
// test for jatmn's second #782 finding: interruptibleReader.Read's select
// between resultCh and ctx.Done() chooses pseudo-randomly when both are ready,
// so a real, already-decided outcome (e.g. a terminal transport error) could be
// misattributed to cancellation purely because ctx also happened to be done by
// the time the select ran. Each iteration lets the underlying Read return (so
// resultCh already holds the real outcome) before cancelling, maximizing the
// odds of landing exactly in that ambiguous window; the generation counter must
// never advance when the real, already-decided result is what gets returned.
func TestInterruptibleReaderPrefersDecidedResultOverCancellation(t *testing.T) {
wantErr := errors.New("read failed")
for i := 0; i < 100; i++ {
resultReady := make(chan struct{})
read := testReader(func(p []byte) (int, error) {
defer close(resultReady)
return copy(p, "x"), wantErr
})
ctx, cancel := context.WithCancel(context.Background())
r := newInterruptibleReader(ctx, read, nil)
readDone := make(chan interruptibleReadResult, 1)
go func() {
n, err := r.Read(make([]byte, 8))
readDone <- interruptibleReadResult{n, err}
}()
<-resultReady
time.Sleep(time.Millisecond) // bias resultCh's send ahead of cancel
cancel()
res := <-readDone
if !errors.Is(res.err, wantErr) {
t.Fatalf("iteration %d: Read returned %v, want %v", i, res.err, wantErr)
}
if got := r.generation(); got != 0 {
t.Fatalf("iteration %d: generation = %d, want 0 (a decided result must never be attributed to cancellation)", i, got)
}
cancel()
}
}
// TestInterruptibleReadCompletionWinsBeforeResultPublication exercises the
// actual inner-read wrapper at the handoff between the underlying Read and the
// helper goroutine's result publication. Cancellation must not reclassify a
// result returned by this wrapper as interrupted.
func TestInterruptibleReadCompletionWinsBeforeResultPublication(t *testing.T) {
wantErr := errors.New("read failed")
read := testReader(func(p []byte) (int, error) {
return copy(p, "x"), wantErr
})
r := newInterruptibleReader(context.Background(), read, nil)
var decision interruptibleReadDecision
readComplete := make(chan struct{})
resultCh := make(chan interruptibleReadResult)
go func() {
res := r.readInner(make([]byte, 8), &decision)
close(readComplete)
resultCh <- res
}()
<-readComplete
if decision.interrupt() {
t.Fatal("cancellation claimed a read that completed before result publication")
}
res := <-resultCh
if res.n != 1 || !errors.Is(res.err, wantErr) {
t.Fatalf("inner read result = (%d, %v), want (1, %v)", res.n, res.err, wantErr)
}
}
func TestConnServeReportsReaderPanicAsTerminalError(t *testing.T) {
read := testReader(func([]byte) (int, error) {
panic("boom")
})
conn := NewConn(read, io.Discard)
done := make(chan error, 1)
go func() { done <- conn.Serve(context.Background()) }()
select {
case err := <-done:
const want = "acp: reader panicked: boom"
if err == nil || err.Error() != want {
t.Fatalf("Serve returned %v, want %q", err, want)
}
case <-time.After(time.Second):
t.Fatal("Serve did not return after reader panic")
}
}
func TestConnRequestResponse(t *testing.T) {
a, b, stop := connPair(t)
defer stop()
@@ -5,9 +5,29 @@ import (
"strings"
"testing"
"github.com/Gitlawb/zero/internal/sandbox"
"github.com/Gitlawb/zero/internal/tools"
)
type normalizationRecordingShellTool struct {
calls []map[string]any
}
func (tool *normalizationRecordingShellTool) Name() string { return "bash" }
func (tool *normalizationRecordingShellTool) Description() string {
return "record normalized shell arguments"
}
func (tool *normalizationRecordingShellTool) Parameters() tools.Schema {
return tools.Schema{Type: "object", AdditionalProperties: true}
}
func (tool *normalizationRecordingShellTool) Safety() tools.Safety {
return tools.Safety{SideEffect: tools.SideEffectShell, Permission: tools.PermissionAllow}
}
func (tool *normalizationRecordingShellTool) Run(_ context.Context, args map[string]any) tools.Result {
tool.calls = append(tool.calls, cloneArgs(args))
return tools.Result{Status: tools.StatusOK}
}
// A shell call that sets sandbox_permissions: with_additional_permissions but
// omits (or malforms) additional_permissions can never succeed: the same
// validation runs again at grant time regardless of the user's decision. It
@@ -144,3 +164,120 @@ func TestExecuteToolCallStillPromptsForValidAdditionalPermissions(t *testing.T)
t.Fatalf("OnPermissionRequest called %d times, want exactly 1 for a valid elevation request", promptCalls)
}
}
func TestNormalizeProviderExpandedEmptyAdditionalPermissions(t *testing.T) {
for _, mode := range []string{
string(tools.SandboxPermissionsUseDefault),
string(tools.SandboxPermissionsWithAdditionalPermissions),
} {
args := map[string]any{
"command": "./zero --version",
"sandbox_permissions": mode,
"additional_permissions": map[string]any{
"file_system": map[string]any{
"deny_read": []any{},
"entries": []any{},
"read": []any{},
"write": []any{},
},
"network": map[string]any{"enabled": false},
},
}
normalizeNoopAdditionalPermissions(args, t.TempDir(), nil)
if _, exists := args["additional_permissions"]; exists {
t.Fatalf("mode %q retained empty additional_permissions: %#v", mode, args)
}
if got := args["sandbox_permissions"]; got != string(tools.SandboxPermissionsUseDefault) {
t.Fatalf("mode %q normalized to %q, want use_default", mode, got)
}
}
}
func TestExecuteToolCallNormalizesProviderExpandedEmptyAdditionalPermissions(t *testing.T) {
tool := &normalizationRecordingShellTool{}
registry := tools.NewRegistry()
registry.Register(tool)
result, err := executeToolCall(
context.Background(),
registry,
ToolCall{
ID: "c1",
Name: "bash",
Arguments: `{"command":"./zero --version","sandbox_permissions":"with_additional_permissions",` +
`"additional_permissions":{"file_system":{"deny_read":[],"entries":[],"read":[],"write":[]},"network":{"enabled":false}}}`,
},
PermissionModeAuto,
Options{Cwd: t.TempDir()},
)
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusOK {
t.Fatalf("provider-expanded empty permissions were rejected: %s", result.Output)
}
if len(tool.calls) != 1 {
t.Fatalf("shell tool calls = %d, want 1", len(tool.calls))
}
args := tool.calls[0]
if _, exists := args["additional_permissions"]; exists {
t.Fatalf("executeToolCall retained empty additional_permissions: %#v", args)
}
if got := args["sandbox_permissions"]; got != string(tools.SandboxPermissionsUseDefault) {
t.Fatalf("sandbox mode = %q, want use_default", got)
}
}
func TestNormalizeAdditionalPermissionsPreservesRealRequest(t *testing.T) {
args := map[string]any{
"sandbox_permissions": string(tools.SandboxPermissionsWithAdditionalPermissions),
"additional_permissions": map[string]any{
"network": map[string]any{"enabled": true},
},
}
normalizeNoopAdditionalPermissions(args, t.TempDir(), nil)
if _, exists := args["additional_permissions"]; !exists {
t.Fatal("non-empty additional_permissions was removed")
}
if got := args["sandbox_permissions"]; got != string(tools.SandboxPermissionsWithAdditionalPermissions) {
t.Fatalf("sandbox mode changed to %q", got)
}
}
func TestNormalizeAdditionalPermissionsRemovesWorkspaceReadAlreadyCoveredBySandbox(t *testing.T) {
root := t.TempDir()
engine := sandbox.NewEngine(sandbox.EngineOptions{
WorkspaceRoot: root,
Policy: sandbox.DefaultPolicy(),
})
args := map[string]any{
"command": "./zero --version",
"sandbox_permissions": string(tools.SandboxPermissionsWithAdditionalPermissions),
"additional_permissions": map[string]any{
"file_system": map[string]any{
"entries": []any{
map[string]any{
"access": "read",
"path": map[string]any{
"type": "path",
"path": ".",
},
},
},
"deny_read": []any{},
"read": []any{},
"write": []any{},
},
"network": map[string]any{"enabled": false},
},
}
normalizeNoopAdditionalPermissions(args, root, engine)
if _, exists := args["additional_permissions"]; exists {
t.Fatalf("already-covered additional_permissions was retained: %#v", args)
}
if got := args["sandbox_permissions"]; got != string(tools.SandboxPermissionsUseDefault) {
t.Fatalf("sandbox mode changed to %q", got)
}
}
+8
View File
@@ -413,6 +413,14 @@ func (state *compactionState) maybeCompact(
// the bodies of old, large tool results (the model has already acted on
// them). If that brings us back under threshold, skip the paid summarizer
// entirely and preserve recent turns verbatim.
if pruned, reclaimed := pruneSupersededReadResults(messages); reclaimed > 0 {
messages = pruned
size = state.calibratedTokens(estimateTokens(messages) + toolTokens)
if size <= state.threshold {
state.lowWaterMark = size
return messages
}
}
if pruned, reclaimed := pruneStaleToolOutput(messages, state.preserveLast); reclaimed > 0 {
messages = pruned
size = state.calibratedTokens(estimateTokens(messages) + toolTokens)
+37
View File
@@ -379,6 +379,43 @@ func TestRunNoCompactionWhenContextWindowZero(t *testing.T) {
}
}
func TestMaybeCompactStopsAfterSupersededReadsReachThreshold(t *testing.T) {
body := strings.Repeat("same source line\n", 1000)
messages := []zeroruntime.Message{
{Role: zeroruntime.MessageRoleSystem, Content: "system"},
toolCallWithArgs("old", "read_file", `{"path":"a.go"}`),
{Role: zeroruntime.MessageRoleTool, ToolCallID: "old", Content: body},
toolCallWithArgs("new", "read_file", `{"path":"a.go"}`),
{Role: zeroruntime.MessageRoleTool, ToolCallID: "new", Content: body},
}
pruned, reclaimed := pruneSupersededReadResults(messages)
if reclaimed <= 0 {
t.Fatal("test setup did not reclaim superseded read output")
}
prunedSize := estimateTokens(pruned)
if estimateTokens(messages) <= prunedSize {
t.Fatal("test setup does not cross the compaction threshold")
}
provider := &summarizeRecordingProvider{}
state := &compactionState{
enabled: true,
threshold: prunedSize,
preserveLast: 2,
}
got := state.maybeCompact(context.Background(), provider, messages, nil)
if provider.summarizeCalls != 0 {
t.Fatalf("cheap pruning should avoid the summarizer, got %d calls", provider.summarizeCalls)
}
if state.lowWaterMark != prunedSize {
t.Fatalf("lowWaterMark = %d, want %d", state.lowWaterMark, prunedSize)
}
if !strings.Contains(got[2].Content, "superseded identical") || got[4].Content != body {
t.Fatalf("unexpected pruned messages: %#v", got)
}
}
// reactiveProvider builds up history with a tool call on turn 1, then errors
// with a context-limit message on turn 2 (once the history is large enough that
// compaction can shrink it), then succeeds on the same-turn retry. Summary
+5 -6
View File
@@ -341,15 +341,15 @@ func TestPartitionToolsActiveHidesUnloadedExposesLoaded(t *testing.T) {
if search.Description != discovery || !strings.Contains(search.Description, "mcp") {
t.Fatalf("tool_search description must carry discovery source, got %q", search.Description)
}
if strings.Contains(search.Description, "mcp__srv__alpha") {
t.Fatalf("tool_search description must not list already-loaded tool names, got %q", search.Description)
if !strings.Contains(search.Description, "mcp__srv__alpha") {
t.Fatalf("stable tool_search catalog must retain already-loaded tool names, got %q", search.Description)
}
if !strings.Contains(search.Description, "mcp__srv__beta") {
t.Fatalf("tool_search description must list exact hidden tool names, got %q", search.Description)
}
}
func TestPartitionToolsActiveNothingHiddenEmptyReminder(t *testing.T) {
func TestPartitionToolsActiveNothingHiddenKeepsStableCatalog(t *testing.T) {
registry := tools.NewRegistry()
registry.Register(fakeDeferredTool{name: "mcp__srv__alpha", desc: "alpha"})
registry.Register(fakeDeferredTool{name: "mcp__srv__beta", desc: "beta"})
@@ -368,9 +368,8 @@ func TestPartitionToolsActiveNothingHiddenEmptyReminder(t *testing.T) {
if !exposedNames["tool_search"] {
t.Fatalf("expected tool_search exposed on active path, got %#v", exposed)
}
// No hidden tools means no dynamic discovery text is needed.
if reminder != "" {
t.Fatalf("expected empty discovery text when nothing is hidden, got %q", reminder)
if !strings.Contains(reminder, "mcp__srv__alpha") || !strings.Contains(reminder, "mcp__srv__beta") {
t.Fatalf("stable catalog must retain all deferred names after loading, got %q", reminder)
}
}
+85 -40
View File
@@ -1076,6 +1076,9 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
}, nil
}
}
if isShellCommandTool(call.Name) {
normalizeNoopAdditionalPermissions(args, options.Cwd, options.Sandbox)
}
// tool_search is the gateway to the allowlisted deferred tools, so a non-empty
// EnabledTools allowlist that omits it must NOT reject the call — otherwise the
// discovery tool exposed to the model rejects at dispatch time (an inescapable
@@ -1354,7 +1357,8 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
Cwd: options.Cwd,
// Per-session file version tracker so write_file/edit_file refuse to clobber
// a file that changed on disk outside Zero since it was last read.
FileTracker: options.FileTracker,
FileTracker: options.FileTracker,
DeferFileObservationCommit: true,
// Post-edit diagnostics are NOT forwarded to the tools: they used to run
// synchronously inside edit_file/write_file and block every mutating tool
// call on the language server (≥300ms debounce floor, 10s cap). The loop
@@ -1406,6 +1410,7 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
result = registry.RebudgetAfterHook(call.Name, args, result)
}
}
result = registry.CommitFileObservation(result, options.FileTracker)
// Stamp the sandbox risk classification for this EXECUTED call so run-policy
// observers can see the risk of an allowed mutation. Prefer the preflight
// decision's classification when the sandbox evaluated the call; otherwise
@@ -1690,20 +1695,21 @@ func sandboxRestrictionRetryReason(result tools.Result) string {
func runToolForNetworkRetry(ctx context.Context, registry *tools.Registry, name string, toolCallID string, args map[string]any, permissionMode PermissionMode, options Options, progressCallback func(streamjson.Event)) tools.Result {
return registry.RunWithOptions(ctx, name, cloneArgs(args), tools.RunOptions{
PermissionGranted: true,
PermissionMode: string(permissionMode),
Autonomy: options.Autonomy,
Sandbox: options.Sandbox,
ToolCallID: toolCallID,
SessionID: options.SessionID,
Model: options.Model,
ReasoningEffort: options.ReasoningEffort,
Depth: options.Depth,
Cwd: options.Cwd,
FileTracker: options.FileTracker,
EnabledTools: options.EnabledTools,
DisabledTools: options.DisabledTools,
Progress: progressCallback,
PermissionGranted: true,
PermissionMode: string(permissionMode),
Autonomy: options.Autonomy,
Sandbox: options.Sandbox,
ToolCallID: toolCallID,
SessionID: options.SessionID,
Model: options.Model,
ReasoningEffort: options.ReasoningEffort,
Depth: options.Depth,
Cwd: options.Cwd,
FileTracker: options.FileTracker,
DeferFileObservationCommit: true,
EnabledTools: options.EnabledTools,
DisabledTools: options.DisabledTools,
Progress: progressCallback,
})
}
@@ -1718,20 +1724,21 @@ func unsandboxedRetryArgs(args map[string]any) map[string]any {
func runToolForUnsandboxedRetry(ctx context.Context, registry *tools.Registry, name string, toolCallID string, args map[string]any, permissionMode PermissionMode, options Options, progressCallback func(streamjson.Event)) tools.Result {
return registry.RunWithOptions(ctx, name, args, tools.RunOptions{
PermissionGranted: true,
PermissionMode: string(permissionMode),
Autonomy: options.Autonomy,
Sandbox: options.Sandbox,
ToolCallID: toolCallID,
SessionID: options.SessionID,
Model: options.Model,
ReasoningEffort: options.ReasoningEffort,
Depth: options.Depth,
Cwd: options.Cwd,
FileTracker: options.FileTracker,
EnabledTools: options.EnabledTools,
DisabledTools: options.DisabledTools,
Progress: progressCallback,
PermissionGranted: true,
PermissionMode: string(permissionMode),
Autonomy: options.Autonomy,
Sandbox: options.Sandbox,
ToolCallID: toolCallID,
SessionID: options.SessionID,
Model: options.Model,
ReasoningEffort: options.ReasoningEffort,
Depth: options.Depth,
Cwd: options.Cwd,
FileTracker: options.FileTracker,
DeferFileObservationCommit: true,
EnabledTools: options.EnabledTools,
DisabledTools: options.DisabledTools,
Progress: progressCallback,
})
}
@@ -2748,6 +2755,45 @@ func shellCommandAdditionalPermissionsRequested(args map[string]any) bool {
return strings.TrimSpace(value) == string(tools.SandboxPermissionsWithAdditionalPermissions)
}
// normalizeNoopAdditionalPermissions treats a provider-expanded empty or
// already-covered additional_permissions object as if it had been omitted.
// Some strict-schema providers materialize every optional property and may add
// a workspace read that the default sandbox already grants; rejecting either
// no-op makes the model retry an otherwise valid command. Requests that broaden
// access and malformed payloads are left untouched so validation stays closed.
func normalizeNoopAdditionalPermissions(args map[string]any, basePath string, engine *sandbox.Engine) {
raw, exists := args["additional_permissions"]
if !exists || raw == nil {
return
}
data, err := json.Marshal(raw)
if err != nil {
return
}
var profile sandbox.RequestPermissionProfile
if err := json.Unmarshal(data, &profile); err != nil {
return
}
normalized, err := sandbox.NormalizeRequestPermissionProfile(profile, basePath)
if err != nil {
return
}
noop := normalized.Empty()
if !noop && engine != nil {
if grantProfile, grantErr := sandbox.RequestPermissionGrantProfile(normalized); grantErr == nil {
noop = engine.CoversRequestPermissions(grantProfile)
}
}
if !noop {
return
}
delete(args, "additional_permissions")
if rawMode, ok := args["sandbox_permissions"].(string); ok &&
strings.TrimSpace(rawMode) == string(tools.SandboxPermissionsWithAdditionalPermissions) {
args["sandbox_permissions"] = string(tools.SandboxPermissionsUseDefault)
}
}
func inlineAdditionalPermissionsProfile(args map[string]any, basePath string) (sandbox.RequestPermissionProfile, bool, error) {
if !shellCommandAdditionalPermissionsRequested(args) {
if _, exists := args["additional_permissions"]; exists {
@@ -2938,22 +2984,20 @@ func partitionToolsCached(registry *tools.Registry, permissionMode PermissionMod
// 2. tool_search — always present, right after the stable block;
// 3. loaded deferred tools, alpha-sorted — APPENDED, so an unloaded->loaded
// transition grows the tail instead of inserting into the eager block.
// (tool_search's discovery text still shrinks as tools load, so keeping it and
// the loaded tail AFTER the eager block preserves the eager tools' cache across a
// load; fully stabilizing the loader's own description is a scoped follow-up.)
// tool_search's compact catalog also stays byte-stable as tools load, so only
// the appended full-schema tail changes.
eager := make([]zeroruntime.ToolDefinition, 0, len(visible))
loadedTail := make([]zeroruntime.ToolDefinition, 0)
var hiddenTools []tools.Tool
var deferredCatalog []tools.Tool
for _, tool := range visible {
name := tool.Name()
if name == tools.ToolSearchToolName {
continue // added explicitly between the eager block and the loaded tail
}
if tools.IsDeferred(tool) {
deferredCatalog = append(deferredCatalog, tool)
if loaded[name] {
loadedTail = append(loadedTail, cachedRuntimeToolDefinition(defCache, tool))
} else {
hiddenTools = append(hiddenTools, tool)
}
continue
}
@@ -2966,15 +3010,16 @@ func partitionToolsCached(registry *tools.Registry, permissionMode PermissionMod
return loadedTail[left].Name < loadedTail[right].Name
})
discovery := ""
if len(hiddenTools) > 0 {
discovery = tools.BuildToolSearchDescription(hiddenTools)
}
// Keep the catalog byte-stable as tools load. tool_search already redirects
// requests for an exposed tool back to the current tool list, so retaining
// loaded names costs only compact catalog lines and avoids invalidating the
// provider cache at the loader definition on every load.
discovery := tools.BuildToolSearchDescription(deferredCatalog)
// tool_search is guaranteed runnable on the active path, so it is ALWAYS exposed
// — even when a non-empty EnabledTools allowlist omits it (the operator
// allowlisted the deferred tools, not the loader). It sits right after the stable
// eager block and carries the discovery list for still-hidden tools.
// eager block and carries the stable deferred-tool catalog.
description := loader.Description()
if discovery != "" {
description = discovery
+16 -3
View File
@@ -332,7 +332,8 @@ type sandboxDeniedExecCommandRetryTool struct {
func (tool *sandboxDeniedExecCommandRetryTool) Name() string { return "exec_command" }
type sandboxNetworkDeniedRetryTool struct {
calls []map[string]any
calls []map[string]any
deferFileObservationCommit []bool
}
func (tool *sandboxNetworkDeniedRetryTool) Name() string { return "bash" }
@@ -355,6 +356,7 @@ func (tool *sandboxNetworkDeniedRetryTool) Run(ctx context.Context, args map[str
}
func (tool *sandboxNetworkDeniedRetryTool) RunWithOptions(ctx context.Context, args map[string]any, options tools.RunOptions) tools.Result {
tool.calls = append(tool.calls, cloneArgs(args))
tool.deferFileObservationCommit = append(tool.deferFileObservationCommit, options.DeferFileObservationCommit)
if shellNetworkAllowed(ctx, options.Sandbox, args) {
return tools.Result{Status: tools.StatusOK, Output: "server started"}
}
@@ -397,7 +399,8 @@ func agentNativeBackendStub() sandbox.Backend {
}
type sandboxNamespaceLimitedRetryTool struct {
calls []map[string]any
calls []map[string]any
deferFileObservationCommit []bool
}
func (tool *sandboxNamespaceLimitedRetryTool) Name() string { return "bash" }
@@ -416,8 +419,12 @@ func (tool *sandboxNamespaceLimitedRetryTool) Parameters() tools.Schema {
func (tool *sandboxNamespaceLimitedRetryTool) Safety() tools.Safety {
return tools.Safety{SideEffect: tools.SideEffectShell, Permission: tools.PermissionPrompt, Reason: "runs shell commands"}
}
func (tool *sandboxNamespaceLimitedRetryTool) Run(_ context.Context, args map[string]any) tools.Result {
func (tool *sandboxNamespaceLimitedRetryTool) Run(ctx context.Context, args map[string]any) tools.Result {
return tool.RunWithOptions(ctx, args, tools.RunOptions{})
}
func (tool *sandboxNamespaceLimitedRetryTool) RunWithOptions(_ context.Context, args map[string]any, options tools.RunOptions) tools.Result {
tool.calls = append(tool.calls, cloneArgs(args))
tool.deferFileObservationCommit = append(tool.deferFileObservationCommit, options.DeferFileObservationCommit)
if args["sandbox_permissions"] == string(tools.SandboxPermissionsRequireEscalated) {
return tools.Result{Status: tools.StatusOK, Output: "stdout:\nUSER PID COMMAND\nanaxy 42 firefox\nanaxy 43 Discord"}
}
@@ -483,6 +490,9 @@ func TestRunRetriesShellUnsandboxedAfterSandboxNamespaceLimitedOutput(t *testing
if retryTool.calls[1]["sandbox_permissions"] != string(tools.SandboxPermissionsRequireEscalated) {
t.Fatalf("retry args = %#v, want require_escalated", retryTool.calls[1])
}
if len(retryTool.deferFileObservationCommit) != 2 || !retryTool.deferFileObservationCommit[1] {
t.Fatalf("retry defer flags = %#v, want retry observation deferred", retryTool.deferFileObservationCommit)
}
if len(requests) != 1 {
t.Fatalf("permission requests = %#v, want one unsandboxed retry approval", requests)
}
@@ -634,6 +644,9 @@ func TestRunRetriesNetworkDeniedShellWithNetworkGrant(t *testing.T) {
if _, escalated := retryTool.calls[1]["sandbox_permissions"]; escalated {
t.Fatalf("network retry must stay sandboxed, got retry args %#v", retryTool.calls[1])
}
if len(retryTool.deferFileObservationCommit) != 2 || !retryTool.deferFileObservationCommit[1] {
t.Fatalf("retry defer flags = %#v, want retry observation deferred", retryTool.deferFileObservationCommit)
}
if len(requests) != 1 || requests[0].Reason != sandbox.ReasonNetworkBlocked {
t.Fatalf("permission requests = %#v, want one network approval", requests)
}
@@ -3,6 +3,7 @@ package agent
import (
"context"
"os"
"path/filepath"
"strconv"
"strings"
"testing"
@@ -19,6 +20,16 @@ type propagationOutputTool struct {
func TestOutputBudgetHookHelperProcess(t *testing.T) {
for index, arg := range os.Args {
if arg == "--zero-output-budget-hook-repeat" && index+1 < len(os.Args) {
repeats, err := strconv.Atoi(os.Args[index+1])
if err != nil {
os.Exit(2)
}
if _, err := os.Stdout.WriteString(strings.Repeat("hook feedback ", repeats)); err != nil {
os.Exit(2)
}
os.Exit(0)
}
if arg != "--zero-output-budget-hook" || index+1 >= len(os.Args) {
continue
}
@@ -30,25 +41,102 @@ func TestOutputBudgetHookHelperProcess(t *testing.T) {
}
func largeOutputBudgetHookDispatcher() *hooks.Dispatcher {
feedback := strings.Repeat("hook feedback ", 200)
return outputBudgetHookDispatcherFor("propagation_output", 200)
}
func outputBudgetHookDispatcherFor(toolName string, repeats int) *hooks.Dispatcher {
feedback := strings.Repeat("hook feedback ", repeats)
hookArgs := []string{
"-test.run=TestOutputBudgetHookHelperProcess",
"--",
"--zero-output-budget-hook",
feedback,
}
if repeats > 1_000 {
hookArgs = []string{
"-test.run=TestOutputBudgetHookHelperProcess",
"--",
"--zero-output-budget-hook-repeat",
strconv.Itoa(repeats),
}
}
return hooks.NewDispatcher(hooks.DispatcherOptions{Config: hooks.Config{
Enabled: true,
Hooks: []hooks.Definition{{
ID: "large-feedback",
Event: hooks.EventAfterTool,
Matcher: "propagation_output",
Matcher: toolName,
Command: os.Args[0],
Args: []string{
"-test.run=TestOutputBudgetHookHelperProcess",
"--",
"--zero-output-budget-hook",
feedback,
},
Args: hookArgs,
Enabled: true,
}},
}})
}
func TestExecuteToolCallCommitsReadObservationAfterFinalBoundary(t *testing.T) {
root := t.TempDir()
targetPath := filepath.Join(root, "target.txt")
if err := os.WriteFile(targetPath, []byte("known content\n"), 0o600); err != nil {
t.Fatal(err)
}
linkPath := filepath.Join(root, "small.txt")
if err := os.Symlink(targetPath, linkPath); err != nil {
t.Skipf("symlink unavailable: %v", err)
}
resolvedPath, err := filepath.EvalSymlinks(linkPath)
if err != nil {
t.Fatal(err)
}
registry := tools.NewRegistry()
registry.Register(tools.NewReadFileTool(root))
tracker := tools.NewFileTracker()
result, abortErr := executeToolCall(context.Background(), registry, ToolCall{
ID: "read-small", Name: "read_file", Arguments: `{"path":"small.txt"}`,
}, PermissionModeAuto, Options{Cwd: root, FileTracker: tracker})
if abortErr != nil || result.Status != tools.StatusOK {
t.Fatalf("executeToolCall: result=%#v err=%v", result, abortErr)
}
if !tracker.SeenWhole(resolvedPath) {
t.Fatal("unmodified final read output was not committed as observed")
}
}
func TestExecuteToolCallDoesNotCommitReadObservationBeforeHookBudget(t *testing.T) {
t.Setenv("ZERO_TOOL_OUTPUT_CEILING_TOKENS", "80")
root := t.TempDir()
targetPath := filepath.Join(root, "target.txt")
if err := os.WriteFile(targetPath, []byte("known content\n"), 0o600); err != nil {
t.Fatal(err)
}
linkPath := filepath.Join(root, "small.txt")
if err := os.Symlink(targetPath, linkPath); err != nil {
t.Skipf("symlink unavailable: %v", err)
}
resolvedPath, err := filepath.EvalSymlinks(linkPath)
if err != nil {
t.Fatal(err)
}
registry := tools.NewRegistry()
registry.Register(tools.NewReadFileTool(root))
tracker := tools.NewFileTracker()
result, abortErr := executeToolCall(context.Background(), registry, ToolCall{
ID: "read-hook", Name: "read_file", Arguments: `{"path":"small.txt"}`,
}, PermissionModeAuto, Options{
Cwd: root, FileTracker: tracker, Hooks: outputBudgetHookDispatcherFor("read_file", 20_000),
})
if abortErr != nil {
t.Fatalf("executeToolCall: %v", abortErr)
}
if !result.Truncated {
t.Fatalf("precondition: hook feedback was not truncated: %#v", result)
}
if tracker.SeenWhole(resolvedPath) {
t.Fatal("read observation committed before hook output was budgeted")
}
}
func (tool propagationOutputTool) Name() string { return "propagation_output" }
func (tool propagationOutputTool) Description() string { return "returns output for propagation tests" }
func (tool propagationOutputTool) Parameters() tools.Schema {
+64
View File
@@ -121,3 +121,67 @@ func toolNameForResult(messages []zeroruntime.Message, index int) string {
func isPrunedPlaceholder(content string) bool {
return strings.HasPrefix(strings.TrimSpace(content), "[pruned ")
}
var supersedableReadTools = map[string]bool{
"read_file": true,
"read_minified_file": true,
"grep": true,
"glob": true,
"list_directory": true,
}
// pruneSupersededReadResults removes only older byte-identical results of the
// same read call. It runs when context pressure already requires a rewrite, so
// it does not trade a warm provider-cache prefix for a small apparent saving.
func pruneSupersededReadResults(messages []zeroruntime.Message) ([]zeroruntime.Message, int) {
type seenResultKey struct {
call string
content string
}
seen := map[seenResultKey]struct{}{}
out := messages
reclaimed := 0
copied := false
for index := len(messages) - 1; index >= 0; index-- {
message := messages[index]
if message.Role != zeroruntime.MessageRoleTool || isPrunedPlaceholder(message.Content) {
continue
}
call, ok := toolCallForResult(messages, index)
if !ok || !supersedableReadTools[call.Name] {
continue
}
key := seenResultKey{
call: call.Name + "\x00" + strings.TrimSpace(call.Arguments),
content: message.Content,
}
if _, exists := seen[key]; exists {
placeholder := fmt.Sprintf("[superseded identical %s result — the newer result remains in context]", call.Name)
if !copied {
out = append([]zeroruntime.Message(nil), messages...)
copied = true
}
out[index].Content = placeholder
reclaimed += max(0, ApproxTextTokens(message.Content)-ApproxTextTokens(placeholder))
continue
}
seen[key] = struct{}{}
}
return out, reclaimed
}
func toolCallForResult(messages []zeroruntime.Message, resultIndex int) (zeroruntime.ToolCall, bool) {
id := messages[resultIndex].ToolCallID
if id == "" {
return zeroruntime.ToolCall{}, false
}
for index := resultIndex - 1; index >= 0; index-- {
for _, call := range messages[index].ToolCalls {
if call.ID == id {
return call, true
}
}
}
return zeroruntime.ToolCall{}, false
}
+73
View File
@@ -25,6 +25,13 @@ func toolCallMsg(id, name string) zeroruntime.Message {
}
}
func toolCallWithArgs(id, name, args string) zeroruntime.Message {
return zeroruntime.Message{
Role: zeroruntime.MessageRoleAssistant,
ToolCalls: []zeroruntime.ToolCall{{ID: id, Name: name, Arguments: args}},
}
}
func TestPruneSkipsSmallSessions(t *testing.T) {
// Total reclaimable is under the gate → no pruning.
msgs := []zeroruntime.Message{
@@ -120,3 +127,69 @@ func TestPruneDoesNotMutateInput(t *testing.T) {
t.Fatal("pruneStaleToolOutput must not mutate the caller's slice")
}
}
func TestPruneSupersededReadResultsOnlyDropsIdenticalOlderReads(t *testing.T) {
body := strings.Repeat("same source line\n", 1000)
messages := []zeroruntime.Message{
toolCallWithArgs("old", "read_file", `{"path":"a.go"}`),
{Role: zeroruntime.MessageRoleTool, ToolCallID: "old", Content: body},
toolCallWithArgs("different", "read_file", `{"path":"b.go"}`),
{Role: zeroruntime.MessageRoleTool, ToolCallID: "different", Content: body},
toolCallWithArgs("new", "read_file", `{"path":"a.go"}`),
{Role: zeroruntime.MessageRoleTool, ToolCallID: "new", Content: body},
}
out, reclaimed := pruneSupersededReadResults(messages)
if reclaimed <= 0 || !strings.Contains(out[1].Content, "superseded identical") {
t.Fatalf("older duplicate was not pruned: reclaimed=%d output=%q", reclaimed, out[1].Content)
}
if out[3].Content != body || out[5].Content != body {
t.Fatal("different call and newest result must remain exact")
}
if messages[1].Content != body {
t.Fatal("input slice was mutated")
}
}
func TestPruneSupersededReadResultsKeepsChangedAndMutatingResults(t *testing.T) {
messages := []zeroruntime.Message{
toolCallWithArgs("r1", "read_file", `{"path":"a.go"}`),
{Role: zeroruntime.MessageRoleTool, ToolCallID: "r1", Content: "version one"},
toolCallWithArgs("r2", "read_file", `{"path":"a.go"}`),
{Role: zeroruntime.MessageRoleTool, ToolCallID: "r2", Content: "version two"},
toolCallWithArgs("e1", "edit_file", `{"path":"a.go"}`),
{Role: zeroruntime.MessageRoleTool, ToolCallID: "e1", Content: "edited"},
toolCallWithArgs("e2", "edit_file", `{"path":"a.go"}`),
{Role: zeroruntime.MessageRoleTool, ToolCallID: "e2", Content: "edited"},
}
out, reclaimed := pruneSupersededReadResults(messages)
if reclaimed != 0 {
t.Fatalf("unsafe results were pruned: %d", reclaimed)
}
for index := range messages {
if out[index].Content != messages[index].Content {
t.Fatalf("message %d changed unexpectedly", index)
}
}
}
func TestPruneSupersededReadResultsFindsMatchingBodyAcrossChangedRead(t *testing.T) {
bodyA := strings.Repeat("version a\n", 100)
bodyB := strings.Repeat("version b\n", 100)
messages := []zeroruntime.Message{
toolCallWithArgs("old-a", "read_file", `{"path":"a.go"}`),
{Role: zeroruntime.MessageRoleTool, ToolCallID: "old-a", Content: bodyA},
toolCallWithArgs("middle-b", "read_file", `{"path":"a.go"}`),
{Role: zeroruntime.MessageRoleTool, ToolCallID: "middle-b", Content: bodyB},
toolCallWithArgs("new-a", "read_file", `{"path":"a.go"}`),
{Role: zeroruntime.MessageRoleTool, ToolCallID: "new-a", Content: bodyA},
}
out, reclaimed := pruneSupersededReadResults(messages)
if reclaimed <= 0 || !strings.Contains(out[1].Content, "superseded identical") {
t.Fatalf("older matching body was not pruned across a changed result: reclaimed=%d output=%q", reclaimed, out[1].Content)
}
if out[3].Content != bodyB || out[5].Content != bodyA {
t.Fatal("changed middle result and newest matching result must remain exact")
}
}
@@ -26,6 +26,16 @@ func TestSpecialistDelegationContext(t *testing.T) {
if strings.Contains(got, "nameless") {
t.Fatalf("nameless entry should be skipped: %q", got)
}
for _, want := range []string{"reduce total context", "Do not delegate small direct work", "bounded assignment"} {
if !strings.Contains(got, want) {
t.Fatalf("missing evidence-driven delegation guidance %q: %q", want, got)
}
}
for _, unwanted := range []string{"delegate to it proactively", "launch several specialists in parallel"} {
if strings.Contains(got, unwanted) {
t.Fatalf("retained broad delegation guidance %q: %q", unwanted, got)
}
}
}
func TestSystemPromptIncludesDelegationOnlyWithSpecialists(t *testing.T) {
+7 -9
View File
@@ -179,21 +179,19 @@ func approvedCommandPrefixContext(options Options) string {
return "## Approved Command Prefixes\n\nThe following command prefixes have already been approved and do not need another permission prompt:\n" + strings.Join(lines, "\n")
}
// specialistDelegationContext nudges the orchestrator to offload read-heavy or
// parallelizable work to a specialist sub-agent via the Task tool, keeping large
// tool outputs out of the main context. It renders only when specialists are
// known (which is only where the Task tool is actually registered), so a run with
// no delegatable specialists produces the previous prompt unchanged.
// specialistDelegationContext describes when a specialist earns its extra
// context and coordination cost. It renders only when specialists are known
// (which is only where the Task tool is actually registered), so a run with no
// delegatable specialists produces the previous prompt unchanged.
func specialistDelegationContext(options Options) string {
if len(options.Specialists) == 0 {
return ""
}
var b strings.Builder
b.WriteString("<specialists>\n")
b.WriteString("Delegate focused or read-heavy work to a specialist sub-agent with the Task tool instead of doing it inline. ")
b.WriteString("When a request matches a specialist's purpose, delegate to it proactively — you do not need the user to ask first. ")
b.WriteString("This keeps large tool outputs — searches, file dumps, multi-step exploration — out of your own context, so you stay fast and token-efficient. ")
b.WriteString("Prefer delegating codebase search and exploration; for independent subtasks, launch several specialists in parallel. Handle small, direct edits yourself.\n")
b.WriteString("Use the Task tool when a specialist's focused context, expertise, or an independent parallel subtask is likely to improve quality or reduce total context. ")
b.WriteString("Do not delegate small direct work, duplicate exploration, or split sequential steps merely because a specialist is available. ")
b.WriteString("When delegating, give a bounded assignment and consume its concise evidence instead of importing a long transcript.\n")
b.WriteString("Available specialists (call Task with the matching name when the task fits its purpose):\n")
for _, info := range options.Specialists {
name := strings.TrimSpace(info.Name)
+7 -11
View File
@@ -31,15 +31,10 @@ work.
tests, or config before you modify behavior. Never edit a file you have not
read.
2. **Plan.** For multi-step work, call update_plan with an ordered checklist and
keep it live. The plan bar is the user's progress signal — call update_plan
after EACH concrete unit of work (every file written, every command run), not
just at coarse milestones: mark the finished item completed and the next one
in_progress before you start it. A plan stuck at 0/N while files are landing is
a bug, not economy — these calls are cheap, expected, and the update_plan cards
are hidden from the transcript, so frequent updates cost the user nothing and
never clutter the conversation. Keep at most one item in_progress, and never
batch the updates to the end of the turn. Skip the plan for trivial one-step
tasks.
keep it current at meaningful milestones. Mark completed work and the next
in-progress step together when practical; do not spend separate model turns
updating the plan after every file or command. Keep at most one item
in_progress. Skip the plan for trivial one-step tasks.
3. **Implement.** Make focused changes that match the surrounding code's style,
naming, and conventions. Prefer the smallest change that fully solves the
problem. Avoid broad refactors, unrelated rewrites, dependency churn, and
@@ -57,8 +52,9 @@ work.
file tools - read_file, list_directory, glob, grep, write_file, edit_file,
apply_patch - over shelling out to cat/sed/awk/python for file operations.
They are safer, reviewable, and produce clean diffs.
- Make one tool call per file. Do not batch multi-file writes into a single
shell or script invocation.
- 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
+7 -3
View File
@@ -27,6 +27,8 @@ Usage:
Not meant to be run interactively — point your editor's ACP / external-agent
setting at "zero acp".`
var acpSignalContext = signalContext
// runACP serves ACP over stdio so an editor can drive ZERO's agent core. It
// speaks JSON-RPC 2.0 (newline-delimited JSON) on stdin/stdout; stderr stays free
// for human-readable diagnostics. The session lifecycle maps onto ZERO's own
@@ -42,7 +44,9 @@ func runACP(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int
}
}
conn := acp.NewConn(deps.stdin, stdout)
// The process's own stdin belongs exclusively to this Conn for the life of
// the command, so cancellation may close it to interrupt an idle read.
conn := acp.NewOwnedConn(deps.stdin, stdout)
acp.NewAgent(conn, acp.Deps{
ResolveConfig: deps.resolveConfig,
DiscoverModels: func(ctx context.Context, profile config.ProviderProfile) ([]providermodeldiscovery.Model, error) {
@@ -73,9 +77,9 @@ func runACP(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int
AgentInfo: acp.Implementation{Name: "zero", Version: version},
})
ctx, stop := signalContext()
ctx, stop := acpSignalContext()
defer stop()
if err := conn.Serve(ctx); err != nil && ctx.Err() == nil {
if err := conn.Serve(ctx); err != nil {
return writeAppError(stderr, "acp: "+err.Error(), exitCrash)
}
return exitSuccess
+120
View File
@@ -0,0 +1,120 @@
package cli
import (
"bytes"
"context"
"errors"
"io"
"strings"
"sync"
"testing"
"time"
)
type acpTestReader func([]byte) (int, error)
func (r acpTestReader) Read(p []byte) (int, error) { return r(p) }
type acpTestWriter func([]byte) (int, error)
func (w acpTestWriter) Write(p []byte) (int, error) { return w(p) }
func installACPTestContext(t *testing.T) context.CancelFunc {
t.Helper()
ctx, cancel := context.WithCancel(context.Background())
previous := acpSignalContext
acpSignalContext = func() (context.Context, context.CancelFunc) {
return ctx, cancel
}
t.Cleanup(func() {
cancel()
acpSignalContext = previous
})
return cancel
}
func TestRunACPCancellationPreservesTerminalReadError(t *testing.T) {
cancel := installACPTestContext(t)
wantErr := errors.New("transport read failed")
reader := acpTestReader(func(p []byte) (int, error) {
return copy(p, `{"jsonrpc":"1.0","id":1}`+"\n"), wantErr
})
writeStarted := make(chan struct{}, 1)
releaseWrite := make(chan struct{})
var stdout bytes.Buffer
writer := acpTestWriter(func(p []byte) (int, error) {
writeStarted <- struct{}{}
<-releaseWrite
return stdout.Write(p)
})
var stderr bytes.Buffer
done := make(chan int, 1)
go func() {
done <- runACP(nil, writer, &stderr, fillAppDeps(appDeps{stdin: reader}))
}()
select {
case <-writeStarted:
case <-time.After(time.Second):
t.Fatal("ACP did not reach the response write")
}
// Keep Serve in the synchronous write until cancellation is observable. The
// regression requires the genuine read error to win over that cancellation.
cancel()
close(releaseWrite)
select {
case code := <-done:
if code != exitCrash {
t.Fatalf("exit code = %d, want crash %d", code, exitCrash)
}
case <-time.After(time.Second):
t.Fatal("ACP did not exit after terminal read error")
}
if got := stderr.String(); !strings.Contains(got, "acp: "+wantErr.Error()) {
t.Fatalf("stderr = %q, want terminal read error", got)
}
}
func TestRunACPIdleCancellationExitsCleanly(t *testing.T) {
cancel := installACPTestContext(t)
pipeReader, pipeWriter := io.Pipe()
t.Cleanup(func() { _ = pipeWriter.Close() })
readStarted := make(chan struct{}, 1)
reader := &acpNotifyingReadCloser{ReadCloser: pipeReader, readStarted: readStarted}
var stdout, stderr bytes.Buffer
done := make(chan int, 1)
go func() {
done <- runACP(nil, &stdout, &stderr, fillAppDeps(appDeps{stdin: reader}))
}()
select {
case <-readStarted:
case <-time.After(time.Second):
t.Fatal("ACP did not begin reading")
}
cancel()
select {
case code := <-done:
if code != exitSuccess {
t.Fatalf("exit code = %d, want success %d; stderr: %s", code, exitSuccess, stderr.String())
}
case <-time.After(time.Second):
t.Fatal("ACP did not exit after cancellation")
}
if stderr.Len() != 0 {
t.Fatalf("stderr = %q, want empty", stderr.String())
}
}
type acpNotifyingReadCloser struct {
io.ReadCloser
readStarted chan<- struct{}
once sync.Once
}
func (r *acpNotifyingReadCloser) Read(p []byte) (int, error) {
r.once.Do(func() { r.readStarted <- struct{}{} })
return r.ReadCloser.Read(p)
}
+3
View File
@@ -585,6 +585,9 @@ func fillAppDeps(deps appDeps) appDeps {
baseProbeProviderHealth := deps.probeProviderHealth
deps.probeProviderHealth = func(ctx context.Context, options providerhealth.Options) providerhealth.Result {
options.Profile = applyStoredProviderKeyAt(options.Profile, userConfigPath)
if options.OAuthResolver == nil {
options.OAuthResolver, _ = oauthLoginForProfile(options.Profile)
}
return baseProbeProviderHealth(ctx, options)
}
return deps
+5 -4
View File
@@ -40,10 +40,11 @@ func runContext(args []string, stdout io.Writer, stderr io.Writer, deps appDeps)
return writeAppError(stderr, err.Error(), exitCrash)
}
report, err := contextreport.Build(contextreport.Options{
WorkspaceRoot: workspaceRoot,
Provider: resolved.Provider,
Registry: newCoreRegistry(workspaceRoot),
ContextWindow: modelContextWindow(modelRegistry, resolved.Provider.Model),
WorkspaceRoot: workspaceRoot,
Provider: resolved.Provider,
Registry: newCoreRegistry(workspaceRoot),
ContextWindow: modelContextWindow(modelRegistry, resolved.Provider.Model),
DeferThreshold: resolved.Tools.DeferThreshold,
})
if err != nil {
return writeAppError(stderr, err.Error(), exitCrash)
+3 -4
View File
@@ -77,11 +77,10 @@ func TestRegisterToolSearchIfEligibleSkipsWhenThresholdZero(t *testing.T) {
}
}
func TestDeferredEligibleCountIgnoresCoreTools(t *testing.T) {
func TestDeferredEligibleCountIncludesOptionalBuiltins(t *testing.T) {
registry := newCoreRegistry(t.TempDir())
// newCoreRegistry holds only built-ins; none implement Deferred().
if got := deferredEligibleCount(registry, agent.PermissionModeAuto, nil, nil); got != 0 {
t.Fatalf("deferredEligibleCount(core) = %d, want 0", got)
if got := deferredEligibleCount(registry, agent.PermissionModeAuto, nil, nil); got < 3 {
t.Fatalf("deferredEligibleCount(core) = %d, want at least 3 optional built-ins", got)
}
}
+62
View File
@@ -5,6 +5,8 @@ import (
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
@@ -12,6 +14,7 @@ import (
"time"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/oauth"
"github.com/Gitlawb/zero/internal/providerhealth"
"github.com/Gitlawb/zero/internal/sessions"
)
@@ -109,6 +112,65 @@ func TestRunDoctorConnectivityProbesProvider(t *testing.T) {
}
}
func TestRunDoctorConnectivityUsesOAuthLogin(t *testing.T) {
tokenPath := filepath.Join(t.TempDir(), "oauth-tokens.json")
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
t.Setenv("ZERO_OAUTH_STORAGE", "file")
t.Setenv("ZERO_OAUTH_TOKENS_PATH", tokenPath)
store, err := oauth.NewStore(oauth.StoreOptions{FilePath: tokenPath})
if err != nil {
t.Fatalf("NewStore: %v", err)
}
if err := store.Save(oauth.ProviderKey("xai"), oauth.Token{
AccessToken: "oauth-doctor-token",
ExpiresAt: time.Now().Add(time.Hour),
}); err != nil {
t.Fatalf("seed OAuth token: %v", err)
}
var gotAuth string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotAuth = r.Header.Get("Authorization")
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"data":[]}`))
}))
defer server.Close()
profile := config.ProviderProfile{
Name: "xai",
CatalogID: "xai",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: server.URL + "/v1",
Model: "grok-4",
}
var stdout bytes.Buffer
var stderr bytes.Buffer
exitCode := runWithDeps([]string{"doctor", "--connectivity"}, &stdout, &stderr, appDeps{
getwd: func() (string, error) {
return t.TempDir(), nil
},
resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) {
return config.ResolvedConfig{
Provider: profile,
Providers: []config.ProviderProfile{profile},
ActiveProvider: profile.Name,
}, nil
},
probeProviderHealth: providerhealth.Probe,
now: fixedCLITime("2026-07-29T20:00:00Z"),
})
if exitCode != exitSuccess {
t.Fatalf("doctor exit = %d, want success; stderr=%q\n%s", exitCode, stderr.String(), stdout.String())
}
if gotAuth != "Bearer oauth-doctor-token" {
t.Fatalf("Authorization = %q, want OAuth bearer", gotAuth)
}
if !strings.Contains(stdout.String(), "[pass] provider.connectivity") {
t.Fatalf("doctor did not report passing connectivity:\n%s", stdout.String())
}
}
func TestRunDoctorConnectivityReportsProviderHealthFailure(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
@@ -19,6 +19,7 @@ func TestProviderCatalogRuntimeProvidersShowOnboardingSetupHints(t *testing.T) {
{id: "openai", name: "OpenAI", transport: "openai", defaultModel: "gpt-4.1"},
{id: "groq", name: "Groq", transport: "openai-compatible", defaultModel: "llama-3.3-70b-versatile"},
{id: "longcat", name: "LongCat", transport: "openai-compatible", defaultModel: "LongCat-2.0"},
{id: "fireworks", name: "Fireworks AI", transport: "openai-compatible", defaultModel: "accounts/fireworks/models/kimi-k2p7-code"},
{id: "ollama-cloud", name: "Ollama Cloud", transport: "openai-compatible", defaultModel: "qwen3-coder:480b"},
{id: "ollama", name: "Ollama Local", transport: "openai-compatible", defaultModel: "llama3.1"},
}
+115 -16
View File
@@ -8,6 +8,7 @@ import (
"os/exec"
"runtime"
"strings"
"sync"
"time"
)
@@ -17,7 +18,11 @@ func LoadProviderCommand(command string) (FileConfig, error) {
stdout, stderr, err := runProviderCommand(command, providerCommandTimeout)
if err != nil {
if errors.Is(err, errProviderCommandTimeout) {
return FileConfig{}, fmt.Errorf("provider command timed out after 5s")
// Report which phase ran out of time, and anything the command
// managed to say. This used to be a bare string with the underlying
// error and stderr both dropped, which left a user whose command was
// slow on a cold machine with nothing to act on.
return FileConfig{}, fmt.Errorf("provider command timed out after %s: %w%s", providerCommandTimeout, err, commandOutput(stderr))
}
return FileConfig{}, fmt.Errorf("provider command failed: %w%s", err, commandOutput(stderr))
}
@@ -40,20 +45,111 @@ func LoadProviderCommand(command string) (FileConfig, error) {
var errProviderCommandTimeout = errors.New("provider command timeout")
// providerCommandDrainGrace bounds how long the timeout path waits for cmd.Wait
// to return after the process tree has been terminated. Terminate normally makes
// that immediate; this exists so a descendant that refuses to die cannot extend
// a call whose deadline has already expired.
const providerCommandDrainGrace = 2 * time.Second
// syncBuffer serialises the command's I/O pump against reads from this package.
// On the timeout path runProviderCommand can return while cmd.Wait is still
// draining, so the buffers outlive the call and a plain bytes.Buffer would be
// read and written concurrently.
type syncBuffer struct {
mu sync.Mutex
buf bytes.Buffer
}
func (b *syncBuffer) Write(p []byte) (int, error) {
b.mu.Lock()
defer b.mu.Unlock()
return b.buf.Write(p)
}
// snapshot copies out what has been written so far. A copy rather than the
// buffer's own slice, which the pump may still append to.
func (b *syncBuffer) snapshot() []byte {
b.mu.Lock()
defer b.mu.Unlock()
return append([]byte(nil), b.buf.Bytes()...)
}
// awaitWithGrace waits for done and gives up after grace, reporting whether the
// wait completed rather than expired.
//
// Split out from the timeout path so the bound can be asserted on its own. In
// place it is effectively untestable: Terminate kills the tree, so Wait returns
// promptly and an unbounded receive would behave identically in every test that
// is not pathological, which is exactly how the old unbounded receive survived.
func awaitWithGrace(done <-chan error, grace time.Duration) bool {
timer := time.NewTimer(grace)
defer timer.Stop()
select {
case <-done:
return true
case <-timer.C:
return false
}
}
// runProviderCommand runs a provider command under a real deadline.
//
// "Real" is the point: timeout used to bound only the wait, so the call could
// take arbitrarily longer than it. Two phases sat outside it. Starting the
// process ran before the clock began, and on Windows that is CREATE_SUSPENDED,
// job-object creation and assignment, and a system-wide thread snapshot, none of
// which is fast on a cold or contended machine. Draining after termination was
// unbounded as well. Both are inside the budget now, so the deadline is an upper
// bound rather than a floor.
func runProviderCommand(command string, timeout time.Duration) ([]byte, []byte, error) {
cmd := shellCommand(command)
// Bound the time Wait spends draining I/O pipes after the process exits,
// in case an orphaned descendant still holds the write ends open.
cmd.WaitDelay = time.Second
var stdout bytes.Buffer
var stderr bytes.Buffer
cmd.Stdout = &stdout
cmd.Stderr = &stderr
stdout := &syncBuffer{}
stderr := &syncBuffer{}
cmd.Stdout = stdout
cmd.Stderr = stderr
proc, err := startCommandProcess(cmd)
if err != nil {
return nil, nil, err
// Armed before the process is started, so process creation is inside the
// budget rather than ahead of it.
deadline := time.NewTimer(timeout)
defer deadline.Stop()
type startResult struct {
proc *commandProcess
err error
}
// Buffered, so the goroutine can finish and exit even when nobody is left
// waiting on it because the deadline already fired.
started := make(chan startResult, 1)
go func() {
proc, err := startCommandProcess(cmd)
started <- startResult{proc: proc, err: err}
}()
var proc *commandProcess
select {
case result := <-started:
if result.err != nil {
return stdout.snapshot(), stderr.snapshot(), result.err
}
proc = result.proc
case <-deadline.C:
// The start is still in flight and will eventually hand back whatever it
// created, so cleanup is handed off rather than abandoned. Wait reaps the
// process; without it a started-but-orphaned command becomes a zombie.
go func() {
result := <-started
if result.err != nil {
return
}
result.proc.Terminate()
_ = cmd.Wait()
result.proc.Close()
}()
return stdout.snapshot(), stderr.snapshot(), fmt.Errorf("%w: starting the command exceeded %s", errProviderCommandTimeout, timeout)
}
defer proc.Close()
@@ -62,9 +158,6 @@ func runProviderCommand(command string, timeout time.Duration) ([]byte, []byte,
done <- cmd.Wait()
}()
timer := time.NewTimer(timeout)
defer timer.Stop()
select {
case err := <-done:
if err != nil {
@@ -83,13 +176,19 @@ func runProviderCommand(command string, timeout time.Duration) ([]byte, []byte,
// and tests unable to tell which one happened. Callers that only ask
// errors.Is(err, errProviderCommandTimeout) are unaffected, so the
// message LoadProviderCommand reports is unchanged.
return stdout.Bytes(), stderr.Bytes(), fmt.Errorf("%w: %w", errProviderCommandTimeout, err)
return stdout.snapshot(), stderr.snapshot(), fmt.Errorf("%w: %w", errProviderCommandTimeout, err)
}
return stdout.Bytes(), stderr.Bytes(), err
case <-timer.C:
return stdout.snapshot(), stderr.snapshot(), err
case <-deadline.C:
proc.Terminate()
<-done
return stdout.Bytes(), stderr.Bytes(), errProviderCommandTimeout
// Bounded, unlike the bare receive this replaces. Terminate kills the
// tree so Wait almost always returns at once, but "almost always" was
// doing real work here: a tree that resists termination used to hold the
// call open indefinitely, long past the deadline that had just expired.
// Returning early leaves cmd.Wait running, which is why the buffers are
// synchronised.
awaitWithGrace(done, providerCommandDrainGrace)
return stdout.snapshot(), stderr.snapshot(), fmt.Errorf("%w: the command did not finish within %s", errProviderCommandTimeout, timeout)
}
}
+130
View File
@@ -414,3 +414,133 @@ func itoa(value int) string {
}
return string(digits[index:])
}
// The whole point of #811: providerCommandTimeout has to be an upper bound on
// how long the call takes, not a floor. Process creation used to run before the
// clock started and the post-termination drain was unbounded, so a slow or
// contended machine could block startup for far longer than the error text
// promised.
//
// This one guards the end-to-end contract against gross regressions. It is NOT
// the assertion that pins the drain bound: Terminate kills the tree, so Wait
// returns promptly and this passes even with the drain unbounded. That property
// is pinned by TestAwaitWithGraceGivesUp, and the start-phase bound by
// TestRunProviderCommandCountsProcessStartAgainstTheDeadline.
func TestRunProviderCommandBoundsTotalDuration(t *testing.T) {
pidFile := filepath.Join(t.TempDir(), "bound.pid")
command := writeCommand(t, commandScript{SleepSeconds: 60, PidFile: pidFile})
const budget = 2 * time.Second
start := time.Now()
_, _, err := runProviderCommand(command, budget)
elapsed := time.Since(start)
if !errors.Is(err, errProviderCommandTimeout) {
t.Fatalf("error = %v, want the timeout sentinel", err)
}
// Budget, plus the bounded drain, plus room for a loaded runner. A failure
// here means some phase is outside the deadline again, which is the bug.
ceiling := budget + providerCommandDrainGrace + 8*time.Second
if elapsed > ceiling {
t.Fatalf("returned after %s with a %s budget; the timeout is not bounding the call (ceiling %s)", elapsed, budget, ceiling)
}
assertProcessTerminatedIfStarted(t, pidFile)
}
// The timeout error has to carry something an operator can act on. It used to be
// a bare string with the underlying error and the command's stderr both dropped,
// so a user whose provider command was slow got no clue which phase ran out or
// what the command had said.
func TestLoadProviderCommandTimeoutReportsDiagnostics(t *testing.T) {
command := writeCommand(t, commandScript{
Stderr: "provider is warming up",
SleepSeconds: 30,
PidFile: filepath.Join(t.TempDir(), "diag.pid"),
})
_, err := LoadProviderCommand(command)
if err == nil {
t.Fatal("LoadProviderCommand() error = nil, want timeout")
}
if !strings.Contains(err.Error(), "timed out") {
t.Fatalf("error = %q, want a timeout message", err.Error())
}
// Wrapped, so callers can still classify it.
if !errors.Is(err, errProviderCommandTimeout) {
t.Fatalf("error = %v, want errors.Is against the timeout sentinel", err)
}
// And it must say which phase, rather than only that time ran out.
if !strings.Contains(err.Error(), "did not finish") && !strings.Contains(err.Error(), "starting the command") {
t.Fatalf("error = %q, want it to name the phase that exceeded the deadline", err.Error())
}
}
// The output buffers are written by cmd.Wait's I/O pump, which keeps running
// after the timeout path has returned them to the caller, so a plain
// bytes.Buffer would be read and written concurrently. Asserted on the buffer
// directly: driving it end to end does not reliably produce the overlap, because
// a terminated command usually stops writing before the snapshot is taken, so
// such a test passes with the locking removed and proves nothing.
//
// Only meaningful under -race, which CI runs.
func TestSyncBufferAllowsConcurrentWriteAndSnapshot(t *testing.T) {
buffer := &syncBuffer{}
writing := make(chan struct{})
go func() {
defer close(writing)
for i := 0; i < 2000; i++ {
_, _ = buffer.Write([]byte("x"))
}
}()
for i := 0; i < 2000; i++ {
_ = buffer.snapshot()
}
<-writing
if len(buffer.snapshot()) != 2000 {
t.Fatalf("buffer holds %d bytes, want 2000", len(buffer.snapshot()))
}
}
// The drain after termination must be bounded. Asserted directly on the helper,
// because in place it is untestable without a tree that refuses to die: Wait
// normally returns at once, which is precisely why an unbounded receive sat
// there unnoticed. A channel that never fires is the pathological case made
// cheap.
func TestAwaitWithGraceGivesUp(t *testing.T) {
never := make(chan error)
start := time.Now()
if awaitWithGrace(never, 50*time.Millisecond) {
t.Fatal("reported completion for a wait that never finished")
}
if elapsed := time.Since(start); elapsed > 5*time.Second {
t.Fatalf("gave up after %s, want roughly the grace period", elapsed)
}
// Still reports completion when the wait does finish, or the timeout path
// would always claim the tree hung.
done := make(chan error, 1)
done <- nil
if !awaitWithGrace(done, time.Second) {
t.Fatal("reported expiry for a wait that had already finished")
}
}
// Process creation has to be inside the budget. It used to run before the clock
// started, and on Windows that phase is CREATE_SUSPENDED plus job-object setup
// plus a system-wide thread snapshot, none of it fast on a contended machine.
//
// A budget that is already spent when the call begins can only be exceeded
// during the start, so the phase named in the error is what distinguishes the
// two orderings: with the clock started afterwards, the start would always
// complete first and the failure would come from the wait instead.
func TestRunProviderCommandCountsProcessStartAgainstTheDeadline(t *testing.T) {
command := writeCommand(t, commandScript{Stdout: `{"name":"cmd","provider":"openai","apiKey":"sk-x","model":"m"}`})
_, _, err := runProviderCommand(command, time.Nanosecond)
if !errors.Is(err, errProviderCommandTimeout) {
t.Fatalf("error = %v, want the timeout sentinel", err)
}
if !strings.Contains(err.Error(), "starting the command") {
t.Fatalf("error = %q, want the start phase to be the one that exceeded the deadline; "+
"anything else means process creation is outside the budget again", err.Error())
}
}
+2 -2
View File
@@ -113,7 +113,7 @@ type PreferencesConfig struct {
// ZERO_THEME, so a /theme choice survives restart. Empty = unset (defaults auto).
Theme string `json:"theme,omitempty"`
// Recaps is a tri-state: nil (unset) defaults to ON; an explicit false means
// the user turned post-turn recaps off. A *bool is its own tri-state, so no
// 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"`
}
@@ -130,7 +130,7 @@ type RecentModelEntry struct {
// MaxRecentModels caps the persisted/displayed recent-selection history.
const MaxRecentModels = 5
// RecapsEnabled reports whether post-turn recaps are on. Unset defaults to ON.
// RecapsEnabled reports whether idle recaps are on. Unset defaults to ON.
func (p PreferencesConfig) RecapsEnabled() bool {
return p.Recaps == nil || *p.Recaps
}
+1 -1
View File
@@ -546,7 +546,7 @@ func SetRecentModels(path string, entries []RecentModelEntry) (FileConfig, error
return cfg, nil
}
// SetRecapsEnabled persists the post-turn recap preference, mirroring
// SetRecapsEnabled persists the idle recap preference, mirroring
// SetFavoriteModels (read-modify-atomic-write).
func SetRecapsEnabled(path string, enabled bool) (FileConfig, error) {
path = strings.TrimSpace(path)
+40 -9
View File
@@ -40,6 +40,7 @@ type Options struct {
Provider config.ProviderProfile
Registry *tools.Registry
ContextWindow int
DeferThreshold int
ProjectContextFiles []string
}
@@ -56,6 +57,8 @@ type Report struct {
FreeTokens int `json:"freeTokens"`
UsedFraction float64 `json:"usedFraction"`
ToolCount int `json:"toolCount"`
AvailableToolCount int `json:"availableToolCount,omitempty"`
DeferredToolCount int `json:"deferredToolCount,omitempty"`
ProjectGuidelineFile string `json:"projectGuidelineFile,omitempty"`
Categories []Category `json:"categories"`
}
@@ -113,8 +116,10 @@ func Build(options Options) (Report, error) {
categories = append(categories, category(CategoryWorkspaceMap, "Workspace map", estimateTextTokens(workspaceMap), report.ContextWindow))
}
toolCount, toolTokens := estimateRegistryTools(options.Registry)
toolCount, availableTools, deferredTools, toolTokens := estimateRegistryTools(options.Registry, options.DeferThreshold)
report.ToolCount = toolCount
report.AvailableToolCount = availableTools
report.DeferredToolCount = deferredTools
if toolTokens > 0 {
categories = append(categories, category(CategoryTools, "Tools", toolTokens, report.ContextWindow))
}
@@ -175,7 +180,11 @@ func Format(report Report) string {
lines = append(lines, fmt.Sprintf("usage: %s tokens (context window unknown)", compactNumber(report.UsedTokens)))
}
if report.ToolCount > 0 {
lines = append(lines, fmt.Sprintf("tools: %d", report.ToolCount))
toolLine := fmt.Sprintf("tools: %d exposed", report.ToolCount)
if report.AvailableToolCount > report.ToolCount {
toolLine += fmt.Sprintf(", %d available, %d deferred", report.AvailableToolCount, report.DeferredToolCount)
}
lines = append(lines, toolLine)
}
if report.ProjectGuidelineFile != "" {
lines = append(lines, "project_guidelines: "+report.ProjectGuidelineFile)
@@ -212,29 +221,51 @@ func category(key string, name string, tokens int, contextWindow int) Category {
return cat
}
func estimateRegistryTools(registry *tools.Registry) (int, int) {
func estimateRegistryTools(registry *tools.Registry, deferThreshold int) (int, int, int, int) {
if registry == nil {
return 0, 0
return 0, 0, 0, 0
}
all := registry.All()
sort.Slice(all, func(left int, right int) bool {
return all[left].Name() < all[right].Name()
})
total := 0
count := 0
var visible []tools.Tool
var deferred []tools.Tool
for _, tool := range all {
if tool.Safety().Permission == tools.PermissionDeny {
continue
}
count++
visible = append(visible, tool)
if tools.IsDeferred(tool) {
deferred = append(deferred, tool)
}
}
active := deferThreshold > 0 && len(deferred) >= deferThreshold
exposed := visible
if active {
exposed = make([]tools.Tool, 0, len(visible)-len(deferred)+1)
for _, tool := range visible {
if !tools.IsDeferred(tool) && tool.Name() != tools.ToolSearchToolName {
exposed = append(exposed, tool)
}
}
exposed = append(exposed, tools.NewToolSearchTool(registry))
}
total := 0
for _, tool := range exposed {
total += estimateTextTokens(tool.Name())
total += estimateTextTokens(tool.Description())
description := tool.Description()
if active && tool.Name() == tools.ToolSearchToolName {
description = tools.BuildToolSearchDescription(deferred)
}
total += estimateTextTokens(description)
if encoded, err := json.Marshal(tool.Parameters()); err == nil {
total += estimateTextTokens(string(encoded))
}
total += toolDefinitionOverheadTokens
}
return count, total
return len(exposed), len(visible), len(deferred), total
}
func readProjectGuidelines(root string, names []string) (string, string) {
@@ -130,6 +130,35 @@ func TestBuildHasStableJSONContractAndCategoryMath(t *testing.T) {
}
}
func TestBuildReportsInitialDeferredToolSurface(t *testing.T) {
root := t.TempDir()
registry := tools.NewRegistry()
for _, tool := range tools.CoreToolsScoped(root, nil) {
registry.Register(tool)
}
eager, err := Build(Options{WorkspaceRoot: root, Registry: registry})
if err != nil {
t.Fatal(err)
}
deferred, err := Build(Options{WorkspaceRoot: root, Registry: registry, DeferThreshold: 3})
if err != nil {
t.Fatal(err)
}
if deferred.DeferredToolCount < 3 {
t.Fatalf("DeferredToolCount = %d, want at least 3", deferred.DeferredToolCount)
}
if deferred.AvailableToolCount != eager.ToolCount {
t.Fatalf("available tools = %d, want eager total %d", deferred.AvailableToolCount, eager.ToolCount)
}
if deferred.ToolCount >= deferred.AvailableToolCount {
t.Fatalf("initial exposed tools = %d, available = %d; expected a smaller initial surface", deferred.ToolCount, deferred.AvailableToolCount)
}
if categoryByKey(deferred, CategoryTools).Tokens >= categoryByKey(eager, CategoryTools).Tokens {
t.Fatalf("deferred tool schemas did not reduce the initial context")
}
}
func TestBuildClampsFreeBudgetWhenOverContextWindow(t *testing.T) {
root := t.TempDir()
writeTestFile(t, root, "AGENTS.md", strings.Repeat("very large project rule\n", 1000))
+62 -18
View File
@@ -2,6 +2,7 @@ package imageinput
import (
"bytes"
"errors"
"fmt"
"net/http"
"os"
@@ -16,6 +17,13 @@ import (
// clipboard, or (nil, "", nil) when the clipboard has no image. Called when
// text clipboard is empty (the user pasted a screenshot). The media type is
// sniffed from the bytes, not trusted from the clipboard.
// ErrClipboardImageUnreadable reports that the clipboard holds an image the host
// has no way to extract. It is deliberately distinct from "no image on the
// clipboard", which is an ordinary no-op: this one is actionable, and silence is
// the wrong response to it.
var ErrClipboardImageUnreadable = errors.New(
"the clipboard holds an image but no helper on this host can read it; on macOS install pngpaste (brew install pngpaste) or a Python with pyobjc")
func ReadClipboardImage() ([]byte, string, error) {
data, err := readClipboardImageBytes()
if err != nil {
@@ -95,35 +103,71 @@ func readClipboardImageWindows() ([]byte, error) {
// readClipboardImageDarwin uses osascript to check for and read a clipboard
// image. Returns (nil, nil) when no image is present.
func readClipboardImageDarwin() ([]byte, error) {
// Check clipboard info for image classes.
check := `osascript -e 'clipboard info'`
out, err := exec.Command("sh", "-c", check).Output()
if err != nil {
return nil, nil
// darwinClipboardInfo reports what the pasteboard is currently offering, and
// darwinClipboardExtract pulls the PNG bytes out of it.
//
// Both are vars so the two-step logic below can be driven in a test. The step
// that matters, "the clipboard holds an image and nothing here can extract it",
// only happens when the first succeeds and the second fails, which cannot be
// arranged by stubbing the caller: a test that substitutes ReadClipboardImage
// asserts the plumbing around this function rather than the decision inside it,
// and stays green with that decision reverted.
var (
darwinClipboardInfo = func() (string, error) {
out, err := exec.Command("sh", "-c", `osascript -e 'clipboard info'`).Output()
return string(out), err
}
info := string(out)
if !strings.Contains(info, "PNG") && !strings.Contains(info, "JPEG") && !strings.Contains(info, "TIFF") && !strings.Contains(info, "GIF") {
return nil, nil
}
// Write clipboard image to a temp file via AppleScript, then read it.
// Using pngpaste if available, falling back to a Python one-liner.
cmd := exec.Command("sh", "-c", `pngpaste - 2>/dev/null || python3 -c "
darwinClipboardExtract = func() ([]byte, error) {
// pngpaste if available, falling back to a Python one-liner. Neither
// ships with macOS: pngpaste is Homebrew-only and Apple dropped the
// bundled PyObjC, so a stock Mac fails here every time.
cmd := exec.Command("sh", "-c", `pngpaste - 2>/dev/null || python3 -c "
import AppKit, sys
pb = AppKit.NSPasteboard.generalPasteboard()
data = pb.dataForType_(AppKit.NSPasteboardTypePNG)
if data:
sys.stdout.buffer.write(data.bytes())
"`)
var stdout bytes.Buffer
cmd.Stdout = &stdout
if err := cmd.Run(); err != nil {
var stdout bytes.Buffer
cmd.Stdout = &stdout
if err := cmd.Run(); err != nil {
return nil, err
}
return stdout.Bytes(), nil
}
)
// darwinClipboardOffersImage reports whether the pasteboard info line names an
// image class.
func darwinClipboardOffersImage(info string) bool {
for _, class := range []string{"PNG", "JPEG", "TIFF", "GIF"} {
if strings.Contains(info, class) {
return true
}
}
return false
}
func readClipboardImageDarwin() ([]byte, error) {
info, err := darwinClipboardInfo()
if err != nil {
return nil, nil
}
if stdout.Len() == 0 {
if !darwinClipboardOffersImage(info) {
return nil, nil
}
return stdout.Bytes(), nil
// Past this point the clipboard is KNOWN to hold an image, because the probe
// above said so. A failure here is therefore not "no image", it is "there is
// an image and nothing on this host can extract it". Reporting that as
// absence is what left users staring at a paste that silently did nothing.
data, err := darwinClipboardExtract()
if err != nil {
return nil, ErrClipboardImageUnreadable
}
if len(data) == 0 {
return nil, ErrClipboardImageUnreadable
}
return data, nil
}
// readClipboardImageLinux tries wl-paste (Wayland) then xclip (X11) to read
@@ -0,0 +1,84 @@
package imageinput
import (
"errors"
"testing"
)
// Drives readClipboardImageDarwin's own decision rather than the plumbing around
// it. Both shell steps are substituted, so this runs on any platform and, more
// importantly, actually reaches the branch under test: a test that stubs
// ReadClipboardImage instead stays green with the production change reverted,
// which is how the first version of this coverage slipped through.
func TestReadClipboardImageDarwinDistinguishesUnreadableFromAbsent(t *testing.T) {
restore := func() func() {
info, extract := darwinClipboardInfo, darwinClipboardExtract
return func() { darwinClipboardInfo, darwinClipboardExtract = info, extract }
}()
t.Cleanup(restore)
const pngInfo = `«class PNGf», «class 8BPS»`
t.Run("image present but no extractor is actionable", func(t *testing.T) {
darwinClipboardInfo = func() (string, error) { return pngInfo, nil }
darwinClipboardExtract = func() ([]byte, error) { return nil, errors.New("exit status 127") }
data, err := readClipboardImageDarwin()
if !errors.Is(err, ErrClipboardImageUnreadable) {
t.Fatalf("err = %v, want ErrClipboardImageUnreadable; a stock Mac gets silence instead of a remedy", err)
}
if data != nil {
t.Fatalf("data = %q, want none", data)
}
})
t.Run("extractor producing nothing is also actionable", func(t *testing.T) {
darwinClipboardInfo = func() (string, error) { return pngInfo, nil }
darwinClipboardExtract = func() ([]byte, error) { return nil, nil }
if _, err := readClipboardImageDarwin(); !errors.Is(err, ErrClipboardImageUnreadable) {
t.Fatalf("err = %v, want ErrClipboardImageUnreadable for an extractor that exits 0 with no bytes", err)
}
})
// The other half of the distinction, and the reason this cannot simply always
// report an error: an empty clipboard is an ordinary no-op, and turning that
// into a notice would make every stray Ctrl+V nag the user.
t.Run("no image on the clipboard stays silent", func(t *testing.T) {
darwinClipboardInfo = func() (string, error) { return `«class utf8»`, nil }
darwinClipboardExtract = func() ([]byte, error) {
t.Fatal("extractor ran despite the clipboard offering no image")
return nil, nil
}
data, err := readClipboardImageDarwin()
if err != nil || data != nil {
t.Fatalf("data=%q err=%v, want a silent no-op", data, err)
}
})
t.Run("probe failure stays silent", func(t *testing.T) {
darwinClipboardInfo = func() (string, error) { return "", errors.New("osascript missing") }
darwinClipboardExtract = func() ([]byte, error) {
t.Fatal("extractor ran despite the info probe failing")
return nil, nil
}
if data, err := readClipboardImageDarwin(); err != nil || data != nil {
t.Fatalf("data=%q err=%v, want a silent no-op", data, err)
}
})
t.Run("working extractor returns the bytes", func(t *testing.T) {
darwinClipboardInfo = func() (string, error) { return pngInfo, nil }
darwinClipboardExtract = func() ([]byte, error) { return []byte("png-bytes"), nil }
data, err := readClipboardImageDarwin()
if err != nil {
t.Fatalf("err = %v, want none", err)
}
if string(data) != "png-bytes" {
t.Fatalf("data = %q, want the extractor's bytes", data)
}
})
}
+1
View File
@@ -138,6 +138,7 @@ var descriptors = []Descriptor{
openAICompat("groq", "Groq", "https://api.groq.com/openai/v1", "llama-3.3-70b-versatile", []string{"GROQ_API_KEY"}),
openAICompat("deepseek", "DeepSeek", "https://api.deepseek.com/v1", "deepseek-chat", []string{"DEEPSEEK_API_KEY"}),
openAICompat("together", "Together AI", "https://api.together.xyz/v1", "meta-llama/Llama-3.3-70B-Instruct-Turbo", []string{"TOGETHER_API_KEY"}),
openAICompat("fireworks", "Fireworks AI", "https://api.fireworks.ai/inference/v1", "accounts/fireworks/models/kimi-k2p7-code", []string{"FIREWORKS_API_KEY"}, "fireworks-ai"),
openAICompat("dashscope", "DashScope", "https://dashscope-intl.aliyuncs.com/compatible-mode/v1", "qwen-plus", []string{"DASHSCOPE_API_KEY", "QWEN_API_KEY"}, "qwen"),
openAICompat("moonshot", "Moonshot AI", "https://api.moonshot.ai/v1", "kimi-k2-0905-preview", []string{"MOONSHOT_API_KEY"}, "kimi"),
openAICompat("atlascloud", "Atlas Cloud", "https://api.atlascloud.ai/v1", "qwen/qwen3.5-flash", []string{"ATLASCLOUD_API_KEY"}, "atlas cloud", "atlas"),
+24 -1
View File
@@ -22,6 +22,7 @@ var expectedCatalogIDs = []string{
"groq",
"deepseek",
"together",
"fireworks",
"dashscope",
"moonshot",
"atlascloud",
@@ -155,6 +156,28 @@ func TestLongCatDescriptor(t *testing.T) {
}
}
func TestFireworksDescriptor(t *testing.T) {
descriptor, err := Require("fireworks")
if err != nil {
t.Fatalf("Require(fireworks) error = %v", err)
}
if descriptor.Name != "Fireworks AI" {
t.Fatalf("Name = %q, want Fireworks AI", descriptor.Name)
}
if descriptor.DefaultBaseURL != "https://api.fireworks.ai/inference/v1" {
t.Fatalf("DefaultBaseURL = %q, want Fireworks OpenAI-compatible endpoint", descriptor.DefaultBaseURL)
}
if descriptor.DefaultModel != "accounts/fireworks/models/kimi-k2p7-code" {
t.Fatalf("DefaultModel = %q, want accounts/fireworks/models/kimi-k2p7-code", descriptor.DefaultModel)
}
if descriptor.Transport != TransportOpenAICompatible {
t.Fatalf("Transport = %q, want %q", descriptor.Transport, TransportOpenAICompatible)
}
if !reflect.DeepEqual(descriptor.AuthEnvVars, []string{"FIREWORKS_API_KEY"}) {
t.Fatalf("AuthEnvVars = %#v, want FIREWORKS_API_KEY", descriptor.AuthEnvVars)
}
}
func TestAtlasCloudDescriptor(t *testing.T) {
descriptor, err := Require("atlascloud")
if err != nil {
@@ -346,7 +369,7 @@ func TestListByTransportPreservesCatalogOrder(t *testing.T) {
TransportBedrock: {"bedrock"},
TransportVertex: {"vertex"},
TransportAnthropicCompat: {"minimax", "minimaxi-cn", "opencode-go-anthropic-compatible", "custom-anthropic-compatible"},
TransportOpenAICompat: {"gitlawb-opengateway", "aimlapi", "ollama-cloud", "ollama", "lmstudio", "openrouter", "huggingface", "chatgpt", "groq", "deepseek", "together", "dashscope", "moonshot", "atlascloud", "longcat", "nvidia-nim", "mistral", "github", "xai", "venice", "xiaomi-mimo", "bankr", "zai", "zai-cn", "kilocode", "opencode", "opencode-go", "atomic-chat", "chatgpt-proxy", "custom-openai-compatible"},
TransportOpenAICompat: {"gitlawb-opengateway", "aimlapi", "ollama-cloud", "ollama", "lmstudio", "openrouter", "huggingface", "chatgpt", "groq", "deepseek", "together", "fireworks", "dashscope", "moonshot", "atlascloud", "longcat", "nvidia-nim", "mistral", "github", "xai", "venice", "xiaomi-mimo", "bankr", "zai", "zai-cn", "kilocode", "opencode", "opencode-go", "atomic-chat", "chatgpt-proxy", "custom-openai-compatible"},
}
for transport, wantIDs := range cases {
+73 -11
View File
@@ -46,12 +46,13 @@ const (
const defaultTimeout = 5 * time.Second
type Options struct {
Profile config.ProviderProfile
Connectivity bool
HTTPClient *http.Client
Resolver Resolver
Timeout time.Duration
UserAgent string
Profile config.ProviderProfile
Connectivity bool
HTTPClient *http.Client
Resolver Resolver
OAuthResolver providerio.TokenResolver
Timeout time.Duration
UserAgent string
}
type Resolver interface {
@@ -207,11 +208,26 @@ func Probe(ctx context.Context, options Options) Result {
"providerKind": metadata.ProviderKind,
}, profile))
if credentialRequired(profile, metadata.ProviderKind) && !hasCredential(profile) {
oauthConfigured := options.OAuthResolver != nil
if oauthConfigured {
var resolveErr error
profile, oauthConfigured, resolveErr = profileWithOAuthCredential(ctx, profile, options.OAuthResolver, false)
if resolveErr != nil {
result.add(check("provider.auth", "Provider auth", StatusFail, CategoryAuth, "Provider OAuth credential could not be resolved: "+resolveErr.Error(), map[string]any{
"oauthLogin": true,
}, profile))
return result.finalize()
}
}
if credentialRequired(profile, metadata.ProviderKind) && !hasCredential(profile) && !oauthConfigured {
result.add(check("provider.auth", "Provider auth", StatusFail, CategoryAuth, fmt.Sprintf("Provider %s requires API credentials.", providerName(profile)), credentialDetails(profile), profile))
return result.finalize()
}
if hasCredential(profile) {
if oauthConfigured {
result.add(check("provider.auth", "Provider auth", StatusPass, CategoryAuth, fmt.Sprintf("Provider %s has OAuth credentials configured.", providerName(profile)), map[string]any{
"oauthLogin": true,
}, profile))
} else if hasCredential(profile) {
result.add(check("provider.auth", "Provider auth", StatusPass, CategoryAuth, fmt.Sprintf("Provider %s has credentials configured.", providerName(profile)), credentialDetails(profile), profile))
} else {
result.add(check("provider.auth", "Provider auth", StatusPass, CategoryAuth, fmt.Sprintf("Provider %s does not require API credentials.", providerName(profile)), credentialDetails(profile), profile))
@@ -286,9 +302,35 @@ func connectivityCheck(ctx context.Context, profile config.ProviderProfile, kind
if err != nil {
return classifyTransportError(err, profile)
}
defer func() {
_ = response.Body.Close()
}()
defer func(body io.ReadCloser) {
_ = body.Close()
}(response.Body)
if response.StatusCode == http.StatusUnauthorized && options.OAuthResolver != nil {
refreshed, ok, refreshErr := profileWithOAuthCredential(requestCtx, options.Profile, options.OAuthResolver, true)
if refreshErr != nil {
return check("provider.connectivity", "Provider connectivity", StatusFail, CategoryAuth, "Provider OAuth credential refresh failed: "+refreshErr.Error(), map[string]any{
"statusCode": response.StatusCode,
}, profile)
}
if ok {
_ = response.Body.Close()
profile = refreshed
request, allowLoopbackOrPrivate, err = healthRequest(requestCtx, profile, kind, options)
if err != nil {
return check("provider.connectivity", "Provider connectivity", StatusFail, CategoryConfig, err.Error(), nil, profile)
}
if options.HTTPClient == nil {
client = newConnectivityClient(timeout, options.Resolver, sensitiveAuthHeaderNames(profile, kind), allowLoopbackOrPrivate)
}
response, err = client.Do(request)
if err != nil {
return classifyTransportError(err, profile)
}
defer func(body io.ReadCloser) {
_ = body.Close()
}(response.Body)
}
}
body, err := io.ReadAll(io.LimitReader(response.Body, 64*1024))
if err != nil {
@@ -320,6 +362,23 @@ func connectivityCheck(ctx context.Context, profile config.ProviderProfile, kind
}, profile)
}
func profileWithOAuthCredential(ctx context.Context, profile config.ProviderProfile, resolver providerio.TokenResolver, forceRefresh bool) (config.ProviderProfile, bool, error) {
if resolver == nil {
return profile, false, nil
}
header, value, ok, err := resolver(ctx, forceRefresh)
if err != nil || !ok {
return profile, false, err
}
if strings.TrimSpace(header) == "" {
header = "Authorization"
}
profile.AuthHeader = header
profile.AuthScheme = ""
profile.AuthHeaderValue = value
return profile, true, nil
}
func healthRequest(ctx context.Context, profile config.ProviderProfile, kind config.ProviderKind, options Options) (*http.Request, bool, error) {
baseURL, err := resolvedBaseURL(profile, kind)
if err != nil {
@@ -748,6 +807,9 @@ func redact(message string, profile config.ProviderProfile) string {
func providerSecrets(profile config.ProviderProfile) []string {
secrets := []string{profile.APIKey, profile.AuthHeaderValue}
if fields := strings.Fields(profile.AuthHeaderValue); len(fields) == 2 {
secrets = append(secrets, fields[1])
}
for _, value := range profile.CustomHeaders {
if strings.TrimSpace(value) != "" {
secrets = append(secrets, value)
@@ -180,6 +180,182 @@ func TestProbeResolvedKindConnectivityAuthErrorRedactsSecret(t *testing.T) {
}
}
func TestProbeConnectivityUsesOAuthCredential(t *testing.T) {
var gotAuth string
resolverCalled := false
client := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
gotAuth = r.Header.Get("Authorization")
return &http.Response{
StatusCode: http.StatusOK,
Status: "200 OK",
Body: io.NopCloser(strings.NewReader(`{"data":[]}`)),
Header: make(http.Header),
}, nil
})}
result := Probe(context.Background(), Options{
Profile: config.ProviderProfile{
Name: "xai",
CatalogID: "xai",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "https://api.example.com/v1",
Model: "grok-4",
},
Connectivity: true,
HTTPClient: client,
Resolver: staticResolver{addr: netip.MustParseAddr("93.184.216.34")},
OAuthResolver: func(context.Context, bool) (string, string, bool, error) {
resolverCalled = true
return "Authorization", "Bearer oauth-test-token", true, nil
},
})
if result.Status != StatusPass {
t.Fatalf("Status = %q, want pass: %#v", result.Status, result.Checks)
}
if !resolverCalled {
t.Fatal("OAuth resolver was not called")
}
if gotAuth != "Bearer oauth-test-token" {
t.Fatalf("Authorization = %q, want OAuth bearer", gotAuth)
}
}
func TestProbeConnectivityOAuthCredentialsAreRedacted(t *testing.T) {
const (
staticKey = "STATICKEYVALUE1234567890abcdef"
oauthToken = "hf_VASANTHSENTINEL0011223344"
)
var authHeaders []string
client := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
authHeaders = append(authHeaders, r.Header.Get("Authorization"))
return &http.Response{
StatusCode: http.StatusUnauthorized,
Status: "401 Unauthorized",
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"invalid credentials ` + staticKey + ` and ` + oauthToken + `"}}`)),
Header: make(http.Header),
}, nil
})}
result := Probe(context.Background(), Options{
Profile: config.ProviderProfile{
Name: "xai",
CatalogID: "xai",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "https://api.example.com/v1",
APIKey: staticKey,
Model: "grok-4",
},
Connectivity: true,
HTTPClient: client,
Resolver: staticResolver{addr: netip.MustParseAddr("93.184.216.34")},
OAuthResolver: func(context.Context, bool) (string, string, bool, error) {
return "Authorization", "Bearer " + oauthToken, true, nil
},
})
if len(authHeaders) != 2 {
t.Fatalf("Authorization headers = %#v, want initial request and one refresh retry", authHeaders)
}
for _, header := range authHeaders {
if header != "Bearer "+oauthToken {
t.Fatalf("Authorization = %q, want OAuth bearer", header)
}
}
check := result.Check("provider.connectivity")
if check == nil || check.Category != CategoryAuth {
t.Fatalf("provider.connectivity = %#v, want auth failure", check)
}
if strings.Contains(check.Message, staticKey) {
t.Fatalf("static API key leaked in message: %q", check.Message)
}
if strings.Contains(check.Message, oauthToken) {
t.Fatalf("OAuth token leaked in message: %q", check.Message)
}
if strings.Count(check.Message, "[REDACTED]") != 2 {
t.Fatalf("expected both credentials to be redacted, got %q", check.Message)
}
}
func TestProbeWithoutConnectivityRejectsUnavailableOAuthCredential(t *testing.T) {
resolverCalled := false
result := Probe(context.Background(), Options{
Profile: config.ProviderProfile{
Name: "xai",
CatalogID: "xai",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "https://api.example.com/v1",
Model: "grok-4",
},
OAuthResolver: func(context.Context, bool) (string, string, bool, error) {
resolverCalled = true
return "", "", false, nil
},
})
if !resolverCalled {
t.Fatal("OAuth resolver was not called")
}
if result.Status != StatusFail {
t.Fatalf("Status = %q, want fail: %#v", result.Status, result.Checks)
}
check := result.Check("provider.auth")
if check == nil || check.Status != StatusFail || check.Category != CategoryAuth {
t.Fatalf("provider.auth = %#v, want auth failure", check)
}
}
func TestProbeConnectivityRefreshesOAuthCredentialAfterUnauthorized(t *testing.T) {
var authHeaders []string
client := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
authHeaders = append(authHeaders, r.Header.Get("Authorization"))
status := http.StatusOK
body := `{"data":[]}`
if len(authHeaders) == 1 {
status = http.StatusUnauthorized
body = `{"error":{"message":"expired"}}`
}
return &http.Response{
StatusCode: status,
Status: http.StatusText(status),
Body: io.NopCloser(strings.NewReader(body)),
Header: make(http.Header),
}, nil
})}
var refreshFlags []bool
result := Probe(context.Background(), Options{
Profile: config.ProviderProfile{
Name: "xai",
CatalogID: "xai",
ProviderKind: config.ProviderKindOpenAICompatible,
BaseURL: "https://api.example.com/v1",
Model: "grok-4",
},
Connectivity: true,
HTTPClient: client,
Resolver: staticResolver{addr: netip.MustParseAddr("93.184.216.34")},
OAuthResolver: func(_ context.Context, forceRefresh bool) (string, string, bool, error) {
refreshFlags = append(refreshFlags, forceRefresh)
token := "stale-token"
if forceRefresh {
token = "fresh-token"
}
return "Authorization", "Bearer " + token, true, nil
},
})
if result.Status != StatusPass {
t.Fatalf("Status = %q, want pass after refresh: %#v", result.Status, result.Checks)
}
if len(refreshFlags) != 2 || refreshFlags[0] || !refreshFlags[1] {
t.Fatalf("OAuth refresh flags = %#v, want [false true]", refreshFlags)
}
if len(authHeaders) != 2 || authHeaders[0] != "Bearer stale-token" || authHeaders[1] != "Bearer fresh-token" {
t.Fatalf("Authorization headers = %#v, want stale then fresh bearer", authHeaders)
}
}
func TestProbeConnectivityClassifiesTimeout(t *testing.T) {
client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
return nil, context.DeadlineExceeded
+8
View File
@@ -122,6 +122,14 @@ var curatedModels = map[string][]Model{
{ID: "deepseek-ai/DeepSeek-R1", Description: "reasoning model"},
{ID: "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", Description: "multimodal model"},
},
"fireworks": {
// Serverless-only curated set: ordinary API-key setups hit /inference/v1
// and cannot use on-demand deployment IDs.
{ID: "accounts/fireworks/models/kimi-k2p7-code", Description: "catalog default"},
{ID: "accounts/fireworks/models/deepseek-v4-flash", Description: "fast coding model"},
{ID: "accounts/fireworks/models/gpt-oss-120b", Description: "general model"},
{ID: "accounts/fireworks/models/deepseek-v4-pro", Description: "reasoning model"},
},
"dashscope": {
{ID: "qwen-plus", Description: "catalog default"},
{ID: "qwen-max", Description: "strong general model"},
@@ -28,6 +28,11 @@ func TestModelsAreProviderScoped(t *testing.T) {
want: []string{"llama-3.3-70b-versatile", "openai/gpt-oss-120b"},
notWant: []string{"gpt-4.1", "claude-sonnet-4.5"},
},
{
provider: "fireworks",
want: []string{"accounts/fireworks/models/kimi-k2p7-code", "accounts/fireworks/models/deepseek-v4-flash"},
notWant: []string{"gpt-4.1", "llama-3.3-70b-versatile"},
},
{
provider: "chatgpt",
want: []string{"gpt-5.5", "gpt-5.4", "gpt-5.4-mini", "gpt-5.3-codex-spark"},
+2
View File
@@ -132,6 +132,8 @@ func ModelsDevProviderID(provider providercatalog.Descriptor) string {
return "zai"
case "minimaxi-cn":
return "minimax"
case "fireworks":
return "fireworks-ai"
default:
return strings.TrimSpace(provider.ID)
}
@@ -131,6 +131,7 @@ func TestModelsDevProviderIDMapsZeroAliases(t *testing.T) {
"ollama-cloud": "ollama-cloud",
"zai-cn": "zai",
"minimaxi-cn": "minimax",
"fireworks": "fireworks-ai",
}
for zeroID, want := range tests {
provider, ok := providercatalog.Get(zeroID)
+3 -1
View File
@@ -25,7 +25,9 @@ type TokenResolver func(ctx context.Context, forceRefresh bool) (header string,
//
// base carries the API-key auth config (key + default header/scheme). setExtra
// (optional) sets the provider's non-auth headers on every attempt. Transient
// retries (429/503/529) are handled by the inner SendWithRetry.
// retries are handled by the inner SendWithRetry: 429/503/529 on their own
// budget, and provably pre-send transport failures (refused or timed-out
// connect, DNS, TLS handshake) on a separate, shorter one.
func SendWithAuthRetry(
ctx context.Context,
client *http.Client,
+292 -28
View File
@@ -3,6 +3,9 @@ package providerio
import (
"bytes"
"context"
"errors"
"io"
"net"
"net/http"
"strconv"
"strings"
@@ -17,13 +20,22 @@ import (
// (previously only the OpenAI provider retried; Anthropic and Gemini surfaced the
// first failure).
//
// What is retried: ONLY 429 (rate limit) and 503 (service unavailable) — the
// statuses where the server explicitly did NOT accept the request, so replaying
// it cannot duplicate work. Other 5xx (500/502/504) and transport/network errors
// are NOT retried: a completion POST is non-idempotent and may have reached and
// been processed by the server, so replaying it could duplicate
// (billable) work. Only the INITIAL request is ever in scope; once the response
// body starts streaming it is never re-issued.
// What is retried: 429 (rate limit) and 503 (service unavailable) — the statuses
// where the server explicitly did NOT accept the request — plus PROVABLY
// pre-send transport failures (DNS resolution, a refused/unreachable dial, a TLS
// handshake timeout) where no request bytes ever left this host. In every one of
// those cases we KNOW the server did not receive and process the request, so a
// replay cannot duplicate work. Other 5xx (500/502/504) and every ambiguous or
// post-send transport error (connection reset, broken pipe, EOF, a generic i/o
// timeout, context deadline) are NOT retried: a completion POST is
// non-idempotent and may already have reached and been processed by the server,
// so replaying it could duplicate (billable) work. Only the INITIAL request is
// ever in scope; once the response body starts streaming it is never re-issued.
// Pre-send classification is by syscall errno (portable across POSIX and
// Windows), scoped to an Op=="dial" *net.OpError so the same errno raised on a
// post-send read/write can't be misread as pre-send. NXDOMAIN is treated as
// permanent (not retried), and the pre-send path fails fast on its own
// sub-second schedule rather than the seconds-long 429 backoff.
const defaultMaxRetryAttempts = 6
@@ -37,12 +49,32 @@ const maxBackoff = 30 * time.Second
// schedule is 2s, 4s, 8s, 16s, then maxBackoff. A var so tests can shrink it.
var retryBackoffBase = 2 * time.Second
// preSendMaxAttempts bounds retries of a provably pre-send transport failure.
// Unlike a 429 (whose window is measured in seconds), a refused/unreachable dial
// either recovers within a few hundred ms or is a deterministic misconfiguration
// (wrong host/port, dead local daemon) that will never recover — so a short,
// low-count schedule fails fast instead of stalling the agent ~60s on the
// rate-limit schedule. This bounds the pre-send path independently of the caller's
// maxAttempts, which is tuned for 429/503.
const preSendMaxAttempts = 3
// preSendBackoffBase is the first wait before replaying a pre-send transport
// failure. Modeled on the mid-stream reconnect base (internal/agent/reconnect.go),
// it is sub-second so a transient dial blip recovers quickly and a permanent one
// fails fast. A var so tests can shrink it.
var preSendBackoffBase = 500 * time.Millisecond
// SendWithRetry issues an HTTP request, retrying ONLY the safe-to-replay server
// responses (429 and 503, see ShouldRetryStatus) up to maxAttempts — backing off
// between tries and honoring a server Retry-After header and context
// cancellation. Other 5xx and transport/network errors are returned immediately,
// never replayed (see the package note). The request is rebuilt from body each
// attempt; setHeader (if non-nil) sets headers on every attempt.
// cancellation. The one transport-error exception is a PROVABLY pre-send failure
// (see isPreSendTransportError): a refused/unreachable dial or DNS failure where
// no request bytes left this host, retried on its own short schedule
// (preSendMaxAttempts / preSendBackoff). Every other 5xx and every ambiguous or
// post-send transport error is returned immediately, never replayed, because a
// non-idempotent completion POST may already have been received. The request is
// rebuilt from body each attempt; setHeader (if non-nil) sets headers on every
// attempt.
//
// It returns the final *http.Response (which the caller inspects for a non-2xx
// status, exactly as before) or an error for a network failure / context
@@ -60,7 +92,47 @@ func SendWithRetry(
if maxAttempts <= 0 {
maxAttempts = defaultMaxRetryAttempts
}
for attempt := 1; ; attempt++ {
if client == nil {
client = http.DefaultClient
}
// Redirects are still FOLLOWED, so a redirecting gateway keeps working, but
// entering one makes that attempt non-retryable. The hazard is that client.Do
// can transmit the original POST to the first host, follow a 307/308, and then
// fail the dial to the redirect target: that error is an Op=="dial" error, so
// isPreSendTransportError would read it as "nothing was sent" and replay a
// completion the first host already received. Observing the redirect rather
// than blocking it keeps both properties, since a request that was redirected
// is never replayed while redirecting endpoints stay usable. The caller's own
// CheckRedirect still runs and still decides whether to proceed.
redirected := false
callerCheckRedirect := client.CheckRedirect
redirectAwareClient := *client
redirectAwareClient.CheckRedirect = func(request *http.Request, via []*http.Request) error {
redirected = true
if callerCheckRedirect != nil {
return callerCheckRedirect(request, via)
}
// Mirror net/http's default policy, which a nil CheckRedirect would have
// applied: stop after 10 hops.
if len(via) >= 10 {
return errors.New("stopped after 10 redirects")
}
return nil
}
// The two retry budgets are fully independent, each with its own counter and
// its own backoff ordinal, because they answer different questions: a
// dial/DNS/TLS failure is a fast local blip (preSendMaxAttempts, sub-second
// schedule) while a 429/503 is a server-side window measured in seconds
// (maxAttempts, retryBackoffBase). A single shared loop counter coupled them
// in both directions: status retries starved the pre-send budget, and
// pre-send retries shortened the status window and started its backoff at a
// later rung. Neither kind of failure now consumes or advances the other.
preSendAttempts := 0
statusAttempts := 0
for {
// Per attempt: only a redirect entered by THIS attempt can make its own
// transport error unsafe to replay.
redirected = false
request, err := http.NewRequestWithContext(ctx, method, url, bytes.NewReader(body))
if err != nil {
return nil, err
@@ -70,32 +142,78 @@ func SendWithRetry(
}
connectSpan := trace.FromContext(ctx).Span(trace.SpanProviderConnect)
response, err := client.Do(request)
response, err := redirectAwareClient.Do(request)
connectSpan.End()
if err != nil {
// A transport failure on a POST does NOT mean the server didn't receive
// it — the request may have arrived and be generating a (billable,
// non-idempotent) completion while only the response/connection failed.
// Replaying it could duplicate that work, so surface the error instead
// of retrying. Only 429/503 responses (below), where we KNOW the request
// was not accepted, are safe to retry.
// Context cancellation always surfaces as cancellation, never a retry.
if ctx.Err() != nil {
return nil, ctx.Err()
}
// A transport failure on a POST does NOT usually mean the server didn't
// receive it — the request may have arrived and be generating a
// (billable, non-idempotent) completion while only the response or
// connection failed. Replaying that could duplicate work, so it is
// surfaced immediately. The one safe exception is a PROVABLY pre-send
// failure (see isPreSendTransportError): the request never left this
// host, so a replay cannot duplicate anything — exactly like a 429/503
// where we KNOW the request was not accepted. The pre-send path uses
// its own short bound and sub-second schedule (preSendMaxAttempts /
// preSendBackoff), not the seconds-long 429 schedule, so a permanent
// dial failure fails fast instead of stalling the agent for ~60s.
// A redirect means the original POST already left this host, so the
// dial that failed belongs to a later hop and replaying would duplicate
// a completion the first host may already have processed.
if isPreSendTransportError(err) && !redirected {
preSendAttempts++
if preSendAttempts < preSendMaxAttempts {
if r := trace.FromContext(ctx); r != nil {
r.Counter(trace.CounterRetryCount, 1)
}
if preSendBackoff(ctx, preSendAttempts) {
continue
}
return nil, ctx.Err()
}
}
return nil, err
}
if ShouldRetryStatus(response.StatusCode) && attempt < maxAttempts {
if r := trace.FromContext(ctx); r != nil {
r.Counter(trace.CounterRetryCount, 1)
// Deliberately NOT guarded on `redirected`, unlike the transport-error
// branch above. The two differ in what is known about the request. A
// transport failure after a redirect leaves it unknown whether the
// redirect target received and began processing the POST, so replaying
// could duplicate billable work. Here every hop answered: a 307/308 is the
// origin declining and delegating, and a retryable status is the target
// declining as well, so nothing accepted the request and a replay
// duplicates nothing.
//
// The exposure is also identical to a retryable status with no redirect in
// play, which this package has always retried. Guarding on redirects here
// would not close a gap, it would make a rate-limited completion behind a
// redirecting gateway fail outright instead of backing off.
// The transport hop succeeded, so the connect phase is healthy again and
// the pre-send budget starts over. It bounds CONSECUTIVE connect failures,
// not the lifetime of the call: without this reset, two early dial retries
// followed by a 429 and then one more dial blip would surface a transport
// error while the status budget still had room, which is the coupling
// these two counters were split apart to remove. Total work stays bounded,
// since every reset costs a status attempt and those are capped.
preSendAttempts = 0
if ShouldRetryStatus(response.StatusCode) {
statusAttempts++
if statusAttempts < maxAttempts {
if r := trace.FromContext(ctx); r != nil {
r.Counter(trace.CounterRetryCount, 1)
}
wait := RetryAfter(response)
_ = response.Body.Close()
if Backoff(ctx, statusAttempts, wait) {
continue
}
// Backoff aborted: the only reason it returns false is ctx cancellation.
return nil, ctx.Err()
}
wait := RetryAfter(response)
_ = response.Body.Close()
if Backoff(ctx, attempt, wait) {
continue
}
// Backoff aborted: the only reason it returns false is ctx cancellation.
return nil, ctx.Err()
}
// Success, a non-retryable status, or retries exhausted on a retryable
@@ -109,6 +227,119 @@ func SendWithRetry(
}
}
// isPreSendTransportError reports whether a transport error PROVES no request
// bytes reached the server, so replaying the request cannot duplicate a
// billable completion. Only narrow, unambiguous pre-connect failures qualify: a
// recoverable DNS lookup failure, a refused/unreachable dial, and a TLS
// handshake timeout (the handshake completes before any HTTP request bytes are
// written).
//
// The refused/unreachable errnos (ECONNREFUSED/ENETUNREACH/EHOSTUNREACH) are
// matched by syscall errno via errors.Is — portable across platforms, since Go
// maps the WSA* codes on Windows too — but ONLY when carried by an Op=="dial"
// *net.OpError. errors.Is matches an errno anywhere in the chain regardless of
// which socket operation raised it, and the kernel raises these same errnos on
// an established connection when a route drops mid-generation; scoping to the
// dial phase is what makes "no request bytes were sent" structural rather than
// argued, so a post-send read/write can never replay a non-idempotent POST.
// NXDOMAIN is treated as permanent (deterministic) and is NOT retried.
//
// Ambiguous or post-send failures are deliberately excluded and checked FIRST,
// so an error that is both is never treated as pre-send: a reset or broken pipe
// can follow a request that was already sent, and a bare "i/o timeout" covers
// read timeouts after send as well as dial timeouts, so none of these prove the
// request was not received and a non-idempotent POST must not be replayed on
// them.
func isPreSendTransportError(err error) bool {
if err == nil {
return false
}
// DNS resolution precedes any connection, so a lookup failure proves nothing
// was sent — but only retry a lookup that could plausibly succeed on a replay.
// NXDOMAIN (IsNotFound) is authoritative and deterministic (a mistyped or
// non-existent host), so retrying it just stalls the agent through the whole
// backoff schedule before failing anyway; a timeout or temporary SERVFAIL is
// worth another try. errors.As unwraps url.Error/OpError to reach the DNSError.
var dnsErr *net.DNSError
if errors.As(err, &dnsErr) {
return !dnsErr.IsNotFound
}
// Post-send / ambiguous failures, excluded first. EOF and reset are matched
// by identity so wording (or a hostname containing "eof") can't fool them.
if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) || isConnResetErrno(err) {
return false
}
// A CONNECT-phase timeout is provably pre-send for the same structural reason
// a refused dial is: the connection was never established, so no request
// bytes left this host. It has to be admitted BEFORE the wording exclusions
// below, because net surfaces it with "i/o timeout" in the message, which the
// blanket exclusion would otherwise reject. The gate is what keeps that safe:
// a post-send read timeout is also Timeout(), but carries Op "read", so it
// never reaches here.
if opErr := connectOpError(err); opErr != nil && opErr.Timeout() {
return true
}
msg := strings.ToLower(err.Error())
if strings.Contains(msg, "connection reset") ||
strings.Contains(msg, "broken pipe") ||
strings.Contains(msg, "i/o timeout") {
return false
}
// A refused/unreachable connection is provably pre-send ONLY when it was
// raised during the DIAL (connection-establishment) phase. The same errnos
// are also raised by the kernel on an ESTABLISHED connection — when a route
// drops or an ICMP unreachable arrives mid-generation — which is post-send.
// errors.Is matches an errno anywhere in the chain regardless of the
// operation, so we first require an Op=="dial" *net.OpError: that makes "no
// request bytes left this host" structural, so a post-send read/write carrying
// one of these can never replay a non-idempotent POST. dialPreSendErrnos is
// platform-specific: on Windows a refused dial carries the raw WSA* errno
// (WSAECONNREFUSED etc.), which does NOT satisfy errors.Is against the POSIX
// syscall.ECONNREFUSED, so the Windows list adds those codes. The string
// markers only catch a POSIX dial error already flattened past its errno.
if connectOpError(err) != nil {
for _, errno := range dialPreSendErrnos {
if errors.Is(err, errno) {
return true
}
}
switch {
case strings.Contains(msg, "connection refused"),
strings.Contains(msg, "network is unreachable"),
strings.Contains(msg, "no route to host"):
return true
}
}
// A TLS handshake times out after the TCP dial connects but before any HTTP
// request bytes are written, so it is structurally pre-send. net/http surfaces
// it as its own error string, not a *net.OpError, so it is matched by wording
// independent of the dial gate.
if strings.Contains(msg, "tls handshake timeout") {
return true
}
return false
}
// connectOpError returns the outermost *net.OpError in err's chain when it comes
// from the connection-establishment phase, else nil. Go tags a direct dial with
// Op "dial" and a failed CONNECT through an HTTP proxy with Op "proxyconnect";
// a failure on an already-established connection carries Op "read"/"write".
// This is how isPreSendTransportError scopes the pre-send match to the connect
// phase so a post-send failure can't be misclassified as pre-send.
//
// "proxyconnect" belongs here for the same reason "dial" does: when HTTPS_PROXY
// is set and the CONNECT fails, the tunnel to the model host was never built, so
// no request bytes reached it. Omitting it meant every pre-send failure behind a
// corporate proxy fell through to no-retry, which is the environment where
// transient connect failures are most common.
func connectOpError(err error) *net.OpError {
var opErr *net.OpError
if errors.As(err, &opErr) && (opErr.Op == "dial" || opErr.Op == "proxyconnect") {
return opErr
}
return nil
}
// ShouldRetryStatus reports whether an HTTP status is safe to retry for a
// non-idempotent completion POST: 429 (Too Many Requests), 503 (Service
// Unavailable), and 529 (Anthropic's "overloaded"). All mean the server
@@ -137,6 +368,39 @@ func Backoff(ctx context.Context, attempt int, retryAfter time.Duration) bool {
}
}
// preSendBackoff waits before pre-send retry attempt N (1-based) on a short
// sub-second schedule (500ms, 1s, ...), returning false if the context is
// cancelled during the wait. Separate from Backoff's seconds-long rate-limit
// schedule because a pre-send dial failure recovers in milliseconds or never.
func preSendBackoff(ctx context.Context, attempt int) bool {
timer := time.NewTimer(preSendBackoffWait(attempt))
defer timer.Stop()
select {
case <-ctx.Done():
return false
case <-timer.C:
return true
}
}
// preSendBackoffWait computes the pre-send wait before retry attempt N (1-based):
// exponential from preSendBackoffBase (500ms, 1s, 2s, ...), capped at maxBackoff.
// The exponent is clamped so a large attempt count cannot overflow.
func preSendBackoffWait(attempt int) time.Duration {
exponent := attempt - 1
if exponent > 5 {
exponent = 5
}
if exponent < 0 {
exponent = 0
}
wait := preSendBackoffBase * time.Duration(1<<exponent)
if wait > maxBackoff {
wait = maxBackoff
}
return wait
}
// backoffWait computes the wait before retry attempt N (1-based): Retry-After
// when supplied, else exponential from retryBackoffBase, both capped at
// maxBackoff. The exponent is clamped so a large attempt count cannot overflow.
@@ -0,0 +1,32 @@
//go:build !windows && !plan9
package providerio
import (
"errors"
"syscall"
)
// dialPreSendErrnos are the syscall errnos a refused or unreachable DIAL raises
// on this platform. On POSIX the standard syscall constants are exactly what the
// net package surfaces, so matching them via errors.Is is sufficient. Plan 9 is
// excluded because it does not define these BSD-style errno constants (it uses
// string errors), so the constants below would not compile there; Zero targets
// linux/darwin/windows, so no Plan 9 variant is needed.
var dialPreSendErrnos = []error{
syscall.ECONNREFUSED,
syscall.ENETUNREACH,
syscall.EHOSTUNREACH,
// A connect timeout is normally caught by the Timeout() check on the connect
// gate; this covers the case where the errno survives but the timeout flag
// does not, e.g. an error rebuilt from the errno alone.
syscall.ETIMEDOUT,
}
// isConnResetErrno reports whether err carries a connection-reset errno, the
// post-send exclusion in isPreSendTransportError. It lives here (per platform)
// so retry.go references no syscall constants and stays buildable on Plan 9,
// which defines none of them.
func isConnResetErrno(err error) bool {
return errors.Is(err, syscall.ECONNRESET)
}
@@ -0,0 +1,15 @@
//go:build plan9
package providerio
// Plan 9 has no BSD-style errno constants (it uses string errors), so the
// errno-based classification is unavailable here: isConnResetErrno returns false
// and dialPreSendErrnos is empty. retry.go still classifies via its string
// markers on this platform, the pre-send inclusions (connection refused /
// network unreachable / no route to host) and the separate post-send
// "connection reset" exclusion each keep working through the wording checks.
// Zero targets linux/darwin/windows, so this variant only keeps the package
// building on Plan 9 rather than adding real support.
var dialPreSendErrnos = []error{}
func isConnResetErrno(error) bool { return false }
@@ -0,0 +1,40 @@
//go:build windows
package providerio
import (
"errors"
"syscall"
"golang.org/x/sys/windows"
)
// dialPreSendErrnos are the errnos a refused or unreachable DIAL raises on
// Windows. A real Windows dial failure carries the raw Winsock error
// (WSAECONNREFUSED = 10061, WSAENETUNREACH = 10051, WSAEHOSTUNREACH = 10065), and
// errors.Is against the POSIX syscall.ECONNREFUSED returns FALSE there because
// Go's Windows syscall.ECONNREFUSED is a distinct value the net package never
// produces. So the WSA codes must be matched explicitly, or a refused dial would
// silently never be retried on Windows. The POSIX constants stay in the list too
// so an error built with them (e.g. a test fixture) still classifies.
var dialPreSendErrnos = []error{
syscall.ECONNREFUSED,
syscall.ENETUNREACH,
syscall.EHOSTUNREACH,
syscall.ETIMEDOUT,
windows.WSAECONNREFUSED,
windows.WSAENETUNREACH,
windows.WSAEHOSTUNREACH,
// Same reasoning as the POSIX ETIMEDOUT entry, with the Winsock code a real
// Windows connect timeout actually carries.
windows.WSAETIMEDOUT,
}
// isConnResetErrno reports whether err carries a connection-reset errno, the
// post-send exclusion in isPreSendTransportError. It matches the same
// syscall.ECONNRESET the shared path used before the errno classification was
// split out per platform, so behavior on Windows is unchanged; it lives here
// so retry.go references no syscall constants and stays buildable on Plan 9.
func isConnResetErrno(err error) bool {
return errors.Is(err, syscall.ECONNRESET)
}
+510 -3
View File
@@ -1,15 +1,34 @@
//go:build !plan9
package providerio
import (
"context"
"errors"
"io"
"net"
"net/http"
"net/http/httptest"
"os"
"sync/atomic"
"syscall"
"testing"
"time"
)
// wrapDialErrno builds the error shape a dial failure has (a *net.OpError
// wrapping an *os.SyscallError wrapping the errno) so errors.Is reaches the
// errno exactly as in production. These fixtures use the POSIX syscall.Exxx
// constants, which do NOT reproduce a real Windows dial: that carries distinct
// windows.WSA* errnos (Go's syscall.ECONNREFUSED is a value the net package
// never produces on Windows). They still classify on Windows only because
// dialPreSendErrnos there also lists the POSIX constants for exactly this
// fixture reason; the real Windows dial branch is covered separately by
// TestIsPreSendTransportErrorRealRefusedDial, which makes an actual refused dial.
func wrapDialErrno(op string, errno syscall.Errno) error {
return &net.OpError{Op: op, Net: "tcp", Err: os.NewSyscallError("connectex", errno)}
}
type roundTripperFunc func(*http.Request) (*http.Response, error)
func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
@@ -35,6 +54,152 @@ func TestSendWithRetryDoesNotReplayTransportErrors(t *testing.T) {
}
}
// A PROVABLY pre-send transport failure (no request bytes left this host) is
// safe to replay and must be retried, bounded by preSendMaxAttempts (its own
// short schedule), unlike an ambiguous post-send failure. Real dial failures
// arrive as an Op=="dial" *net.OpError, so the errno cases are the production
// shape on every platform (Windows included); the string case exercises the
// wording fallback for a dial error already flattened past its errno.
func TestSendWithRetryReplaysProvablyPreSendErrors(t *testing.T) {
shrinkBackoff(t)
cases := map[string]error{
"errno refused (dial)": wrapDialErrno("dial", syscall.ECONNREFUSED),
"errno network unreachable": wrapDialErrno("dial", syscall.ENETUNREACH),
"errno host unreachable": wrapDialErrno("dial", syscall.EHOSTUNREACH),
"dial string fallback": &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connect: connection refused")},
"tls handshake timeout": errors.New("net/http: TLS handshake timeout"),
"dns timeout": &net.DNSError{Err: "server misbehaving", Name: "nope.invalid", IsTimeout: true},
}
for name, transportErr := range cases {
t.Run(name, func(t *testing.T) {
var calls int32
client := &http.Client{Transport: roundTripperFunc(func(*http.Request) (*http.Response, error) {
atomic.AddInt32(&calls, 1)
return nil, transportErr
})}
// maxAttempts=6 (the default) proves the pre-send path is bounded by
// preSendMaxAttempts, not the caller's 429-tuned maxAttempts.
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "http://example.invalid", []byte("{}"), nil, 6)
if resp != nil {
_ = resp.Body.Close()
}
if err == nil {
t.Fatal("expected the transport error to surface after retries are exhausted")
}
if got := atomic.LoadInt32(&calls); got != preSendMaxAttempts {
t.Fatalf("pre-send error tried %d times, want %d (preSendMaxAttempts)", got, preSendMaxAttempts)
}
})
}
}
// An ambiguous transport failure that could have followed a sent request must
// NOT be replayed: a non-idempotent completion POST could duplicate billable
// work. This is the safety line the fix must not cross.
func TestSendWithRetryDoesNotReplayAmbiguousTransportErrors(t *testing.T) {
for name, transportErr := range map[string]error{
"generic i/o timeout": errors.New("dial tcp 1.2.3.4:443: i/o timeout"),
"broken pipe": errors.New("write tcp: broken pipe"),
"unexpected eof": io.ErrUnexpectedEOF,
// The pre-send errnos are also raised on an ESTABLISHED connection: a route
// dropping mid-generation surfaces on the pending read/write, which is
// post-send. Scoping to Op=="dial" must keep these from replaying the POST.
"host unreachable on read is post-send": wrapDialErrno("read", syscall.EHOSTUNREACH),
"net unreachable on write is post-send": wrapDialErrno("write", syscall.ENETUNREACH),
// NXDOMAIN is authoritative and deterministic, so retrying it would only
// stall the agent before failing anyway.
"dns nxdomain": &net.DNSError{Err: "no such host", Name: "nope.invalid", IsNotFound: true},
} {
t.Run(name, func(t *testing.T) {
var calls int32
client := &http.Client{Transport: roundTripperFunc(func(*http.Request) (*http.Response, error) {
atomic.AddInt32(&calls, 1)
return nil, transportErr
})}
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "http://example.invalid", []byte("{}"), nil, 3)
if resp != nil {
_ = resp.Body.Close()
}
if err == nil {
t.Fatal("expected the transport error to surface")
}
if got := atomic.LoadInt32(&calls); got != 1 {
t.Fatalf("ambiguous transport error replayed %d times, want 1 (no retry)", got)
}
})
}
}
// A REAL refused dial (nothing listening on the target) must classify as
// pre-send on every platform. This is the regression the errno-constant fixtures
// miss: on Windows the kernel raises WSAECONNREFUSED, which errors.Is does NOT
// match against syscall.ECONNREFUSED, so before dialPreSendErrnos carried the WSA
// codes a refused dial returned false here and Windows never retried it. Using a
// real Dial exercises the platform's true error shape.
func TestIsPreSendTransportErrorRealRefusedDial(t *testing.T) {
// Reserve an ephemeral port, then close it, so a dial there is refused.
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
addr := ln.Addr().String()
if cerr := ln.Close(); cerr != nil {
t.Fatalf("close listener: %v", cerr)
}
conn, dialErr := (&net.Dialer{Timeout: 2 * time.Second}).Dial("tcp", addr)
if dialErr == nil {
_ = conn.Close()
t.Skip("expected a refused dial but the port accepted a connection")
}
if !isPreSendTransportError(dialErr) {
t.Fatalf("real refused dial not classified pre-send: %v", dialErr)
}
}
func TestIsPreSendTransportError(t *testing.T) {
cases := []struct {
name string
err error
want bool
}{
{"nil", nil, false},
// DNS: a lookup that could recover is pre-send; NXDOMAIN is permanent.
{"dns timeout", &net.DNSError{Err: "server misbehaving", Name: "x.invalid", IsTimeout: true}, true},
{"dns temporary", &net.DNSError{Err: "server misbehaving", Name: "x.invalid", IsTemporary: true}, true},
{"dns nxdomain is permanent", &net.DNSError{Err: "no such host", Name: "x.invalid", IsNotFound: true}, false},
{"tls handshake timeout", errors.New("net/http: TLS handshake timeout"), true},
// Errno-wrapped DIAL failures: how a real refused/unreachable dial arrives
// on EVERY platform, including Windows where the wording differs entirely.
{"errno refused (dial)", wrapDialErrno("dial", syscall.ECONNREFUSED), true},
{"errno network unreachable (dial)", wrapDialErrno("dial", syscall.ENETUNREACH), true},
{"errno host unreachable (dial)", wrapDialErrno("dial", syscall.EHOSTUNREACH), true},
{"dial operror string fallback", &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connect: connection refused")}, true},
// The SAME errnos raised on a post-send read/write are NOT pre-send — the
// kernel reports them on an established connection when a route drops.
{"host unreachable on read is post-send", wrapDialErrno("read", syscall.EHOSTUNREACH), false},
{"net unreachable on write is post-send", wrapDialErrno("write", syscall.ENETUNREACH), false},
{"errno reset is post-send", wrapDialErrno("read", syscall.ECONNRESET), false},
// A refused/unreachable error already flattened past its *net.OpError can't
// be proven pre-send, so it is NOT retried (conservative direction).
{"flattened refused, no opError", errors.New("dial tcp 127.0.0.1:1: connect: connection refused"), false},
{"connection reset", errors.New("read tcp: connection reset by peer"), false},
{"broken pipe", errors.New("write tcp: broken pipe"), false},
{"unexpected eof", io.ErrUnexpectedEOF, false},
{"eof", io.EOF, false},
{"generic io timeout", errors.New("dial tcp: i/o timeout"), false},
{"context deadline", context.DeadlineExceeded, false},
// Exclusion is checked before inclusion, even for a dial OpError: an "i/o
// timeout" in the message wins over the refused wording.
{"exclusion wins over inclusion", &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused / i/o timeout")}, false},
{"unrelated", errors.New("some other error"), false},
}
for _, c := range cases {
if got := isPreSendTransportError(c.err); got != c.want {
t.Errorf("%s: isPreSendTransportError(%v) = %v, want %v", c.name, c.err, got, c.want)
}
}
}
func TestShouldRetryStatus(t *testing.T) {
cases := map[int]bool{
http.StatusOK: false,
@@ -94,12 +259,17 @@ func TestBackoffWaitsThenReturnsTrue(t *testing.T) {
}
}
// shrinkBackoff makes retry waits negligible for the duration of a test.
// shrinkBackoff makes both retry schedules (429 and pre-send) negligible for the
// duration of a test.
func shrinkBackoff(t *testing.T) {
t.Helper()
saved := retryBackoffBase
savedRetry, savedPreSend := retryBackoffBase, preSendBackoffBase
retryBackoffBase = time.Millisecond
t.Cleanup(func() { retryBackoffBase = saved })
preSendBackoffBase = time.Millisecond
t.Cleanup(func() {
retryBackoffBase = savedRetry
preSendBackoffBase = savedPreSend
})
}
func TestBackoffWaitSchedule(t *testing.T) {
@@ -126,6 +296,30 @@ func TestBackoffWaitSchedule(t *testing.T) {
}
}
// The pre-send schedule is sub-second and doubles, far shorter than the 429
// schedule: a permanent dial failure fails in ~1.5s across preSendMaxAttempts
// (500ms + 1s) instead of stalling the agent ~60s on 2/4/8/16/30s. This is the
// second half of the fix for a mistyped host / dead local daemon hanging a turn.
func TestPreSendBackoffWaitSchedule(t *testing.T) {
cases := []struct {
attempt int
want time.Duration
}{
{1, 500 * time.Millisecond},
{2, 1 * time.Second},
{3, 2 * time.Second},
// The exponent clamps at 5 (500ms<<5 = 16s) so a large attempt count can't
// overflow; in practice preSendMaxAttempts caps retries at attempt 2, so
// only the 500ms/1s rungs are ever reached.
{50, 16 * time.Second},
}
for _, c := range cases {
if got := preSendBackoffWait(c.attempt); got != c.want {
t.Errorf("preSendBackoffWait(%d) = %v, want %v", c.attempt, got, c.want)
}
}
}
func TestSendWithRetryRetriesThenSucceeds(t *testing.T) {
shrinkBackoff(t)
var hits int32
@@ -193,3 +387,316 @@ func TestSendWithRetryReturnsLastResponseAfterMaxAttempts(t *testing.T) {
t.Fatalf("server hit %d times, want 2 (maxAttempts)", got)
}
}
// Redirects are still followed, so a gateway that answers the completion endpoint
// with a 307/308 keeps working: the caller gets the final response, not the 3xx.
func TestSendWithRetryFollowsRedirects(t *testing.T) {
var calls int32
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
if atomic.AddInt32(&calls, 1) == 1 {
return &http.Response{
StatusCode: http.StatusTemporaryRedirect,
Header: http.Header{"Location": {"https://redirect-target.invalid/v1"}},
Body: http.NoBody,
Request: r,
}, nil
}
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Request: r}, nil
})}
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://origin.invalid/v1", []byte("{}"), nil, 3)
if err != nil {
t.Fatalf("a redirected completion must succeed: %v", err)
}
_ = resp.Body.Close()
if resp.StatusCode != http.StatusOK {
t.Fatalf("status = %d, want 200 (the redirect should have been followed)", resp.StatusCode)
}
if got := atomic.LoadInt32(&calls); got != 2 {
t.Fatalf("transport called %d times, want 2 (origin then redirect target)", got)
}
}
// Once an attempt has entered a redirect, the original POST has already left this
// host, so a dial failure on the redirect hop must NOT be replayed even though it
// arrives as an Op=="dial" error that isPreSendTransportError would otherwise
// treat as provably pre-send. Replaying it would re-bill a completion the first
// host may already have processed.
func TestSendWithRetryDoesNotReplayAfterRedirect(t *testing.T) {
shrinkBackoff(t)
var calls int32
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
if atomic.AddInt32(&calls, 1) == 1 {
return &http.Response{
StatusCode: http.StatusTemporaryRedirect,
Header: http.Header{"Location": {"https://redirect-target.invalid/v1"}},
Body: http.NoBody,
Request: r,
}, nil
}
// The dial to the redirect target fails: pre-send in shape, but post-send
// in fact, because the origin already received the POST.
return nil, wrapDialErrno("dial", syscall.ECONNREFUSED)
})}
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://origin.invalid/v1", []byte("{}"), nil, 6)
if resp != nil {
_ = resp.Body.Close()
}
if err == nil {
t.Fatal("expected the redirect-hop dial failure to surface as an error")
}
// Exactly two: the origin request and the single redirect hop. A third would
// mean the redirected POST was replayed.
if got := atomic.LoadInt32(&calls); got != 2 {
t.Fatalf("redirected POST was replayed: transport called %d times, want 2", got)
}
}
// A caller that supplies its own CheckRedirect keeps it: the retry path observes
// redirects, it does not take the decision away from the caller.
func TestSendWithRetryPreservesCallerRedirectPolicy(t *testing.T) {
var policyCalls int32
client := &http.Client{
CheckRedirect: func(*http.Request, []*http.Request) error {
atomic.AddInt32(&policyCalls, 1)
return http.ErrUseLastResponse
},
Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
return &http.Response{
StatusCode: http.StatusTemporaryRedirect,
Header: http.Header{"Location": {"https://redirect-target.invalid/v1"}},
Body: http.NoBody,
Request: r,
}, nil
}),
}
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://origin.invalid/v1", []byte("{}"), nil, 3)
if err != nil {
t.Fatalf("caller policy asked to stop at the 3xx, so it must surface: %v", err)
}
_ = resp.Body.Close()
if resp.StatusCode != http.StatusTemporaryRedirect {
t.Fatalf("status = %d, want the 307 surfaced per the caller's policy", resp.StatusCode)
}
if got := atomic.LoadInt32(&policyCalls); got != 1 {
t.Fatalf("caller CheckRedirect called %d times, want 1 (it must not be discarded)", got)
}
}
// The pre-send retry budget is independent of the 429/503 status retries: two
// rate-limit responses must not consume the pre-send allowance, so the first
// refused dial after them still gets its own preSendMaxAttempts tries.
func TestSendWithRetryPreSendBudgetSurvivesStatusRetries(t *testing.T) {
shrinkBackoff(t)
var calls int32
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
if atomic.AddInt32(&calls, 1) <= 2 {
return &http.Response{StatusCode: http.StatusTooManyRequests, Body: http.NoBody, Request: r}, nil
}
return nil, wrapDialErrno("dial", syscall.ECONNREFUSED)
})}
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://x.invalid/v1", []byte("{}"), nil, 6)
if resp != nil {
_ = resp.Body.Close()
}
if err == nil {
t.Fatal("expected the exhausted pre-send failure to surface as an error")
}
// Two status responses (calls 1-2), then the pre-send failure gets its own
// budget: preSendMaxAttempts=3 means it retries twice more before returning.
if got := atomic.LoadInt32(&calls); got != 2+preSendMaxAttempts {
t.Fatalf("pre-send budget eaten by status retries: transport called %d times, want %d (2 status + %d pre-send)", got, 2+preSendMaxAttempts, preSendMaxAttempts)
}
}
// The inverse sequence: pre-send retries must not consume the 429/503 budget or
// advance its backoff rung either. Two safely replayed refused dials followed by
// rate limiting must still leave the status path its full maxAttempts window,
// otherwise a couple of dial blips silently shorten every later rate-limit wait.
func TestSendWithRetryStatusBudgetSurvivesPreSendRetries(t *testing.T) {
shrinkBackoff(t)
const maxAttempts = 3
var calls int32
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
if atomic.AddInt32(&calls, 1) <= 2 {
return nil, wrapDialErrno("dial", syscall.ECONNREFUSED)
}
return &http.Response{StatusCode: http.StatusTooManyRequests, Body: http.NoBody, Request: r}, nil
})}
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://x.invalid/v1", []byte("{}"), nil, maxAttempts)
if err != nil {
t.Fatalf("exhausted status retries surface the response, not an error: %v", err)
}
_ = resp.Body.Close()
if resp.StatusCode != http.StatusTooManyRequests {
t.Fatalf("status = %d, want 429", resp.StatusCode)
}
// Calls 1-2 are the retried dial failures; the status path then gets its own
// full window (maxAttempts requests) before the 429 is surfaced.
if got := atomic.LoadInt32(&calls); got != 2+maxAttempts {
t.Fatalf("status budget eaten by pre-send retries: transport called %d times, want %d (2 pre-send + %d status)", got, 2+maxAttempts, maxAttempts)
}
}
// A redirect followed by a RETRYABLE STATUS is still retried, and deliberately
// so. This is the case a reviewer flagged as a possible replay of billable work,
// and the asymmetry with the transport-error guard above is intentional rather
// than an oversight, so it is pinned here.
//
// The two situations differ in what is known about the request:
//
// - Transport error after a redirect: the original POST reached the origin and
// the dial to the redirect target then failed, so whether the target received
// and began processing it is UNKNOWN. Replaying could duplicate work, which
// is why that path refuses to retry.
// - Retryable status after a redirect: every hop answered. A 307/308 is the
// origin declining to process and delegating, and a 429/503/529 is the target
// declining as well. Nothing accepted the request, so replaying duplicates
// nothing.
//
// The risk here is also exactly the risk of a 429 without any redirect, which
// this package has always retried; the redirect adds no new ambiguity. Refusing
// to retry would instead make a rate-limited completion behind a redirecting
// gateway fail outright rather than back off.
func TestSendWithRetryRetriesRedirectedRetryableStatus(t *testing.T) {
var calls int32
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
switch atomic.AddInt32(&calls, 1) {
case 1, 3:
// The origin declines and delegates, both times.
return &http.Response{
StatusCode: http.StatusTemporaryRedirect,
Header: http.Header{"Location": {"https://redirect-target.invalid/v1"}},
Body: http.NoBody,
Request: r,
}, nil
case 2:
// The target declines too: rate limited, so nothing was accepted.
return &http.Response{StatusCode: http.StatusTooManyRequests, Body: http.NoBody, Request: r}, nil
default:
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Request: r}, nil
}
})}
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://origin.invalid/v1", []byte("{}"), nil, 3)
if err != nil {
t.Fatalf("a rate-limited completion behind a redirect must still retry: %v", err)
}
_ = resp.Body.Close()
if resp.StatusCode != http.StatusOK {
t.Fatalf("status = %d, want 200 after the retry succeeded", resp.StatusCode)
}
// origin, target(429), origin, target(200).
if got := atomic.LoadInt32(&calls); got != 4 {
t.Fatalf("transport called %d times, want 4 (redirect, 429, redirect, 200)", got)
}
}
// A CONNECT-phase timeout is provably pre-send: the connection was never
// established, so no request bytes left this host. It reaches the classifier
// carrying "i/o timeout", which the blanket wording exclusion rejects, so
// admitting it depends on the connect gate running first.
//
// The post-send counterpart is the point of the table: a read timeout is also
// Timeout(), and must stay excluded, or a reply that timed out mid-generation
// would replay a completion the server already billed.
func TestPreSendClassifiesConnectTimeouts(t *testing.T) {
timeoutErr := &timeoutError{}
for name, testCase := range map[string]struct {
err error
want bool
}{
"dial timeout": {
err: &net.OpError{Op: "dial", Net: "tcp", Err: timeoutErr},
want: true,
},
"proxy connect timeout": {
err: &net.OpError{Op: "proxyconnect", Net: "tcp", Err: timeoutErr},
want: true,
},
"read timeout stays post-send": {
err: &net.OpError{Op: "read", Net: "tcp", Err: timeoutErr},
want: false,
},
"write timeout stays post-send": {
err: &net.OpError{Op: "write", Net: "tcp", Err: timeoutErr},
want: false,
},
} {
t.Run(name, func(t *testing.T) {
if got := isPreSendTransportError(testCase.err); got != testCase.want {
t.Fatalf("isPreSendTransportError(%v) = %v, want %v", testCase.err, got, testCase.want)
}
})
}
}
// A failed CONNECT through an HTTP proxy is pre-send for the same structural
// reason a refused dial is, but net tags it Op "proxyconnect", so a gate
// matching only "dial" left every pre-send failure behind a corporate proxy
// unretried.
func TestPreSendClassifiesProxyConnectRefusal(t *testing.T) {
err := &net.OpError{Op: "proxyconnect", Net: "tcp", Err: syscall.ECONNREFUSED}
if !isPreSendTransportError(err) {
t.Fatal("a refused proxy CONNECT was not classified pre-send; no request bytes reached the model host")
}
// The op still has to be a connect phase. The same errno on an established
// connection is post-send and must not replay a POST.
post := &net.OpError{Op: "read", Net: "tcp", Err: syscall.ECONNREFUSED}
if isPreSendTransportError(post) {
t.Fatal("a read-phase error was classified pre-send")
}
}
// timeoutError is a net.Error whose Timeout() is true, matching the shape
// net/http surfaces for a connect deadline.
type timeoutError struct{}
func (timeoutError) Error() string { return "i/o timeout" }
func (timeoutError) Timeout() bool { return true }
func (timeoutError) Temporary() bool { return true }
// The pre-send budget bounds CONSECUTIVE connect failures, not the lifetime of
// the call. Without a reset on a successful hop, an early run of dial blips
// permanently spends the budget, so a later blip surfaces as a transport error
// while the status budget still has room. That is the coupling the two counters
// were split apart to remove, in the one ordering the earlier tests did not
// cover: pre-send, then status, then pre-send again.
func TestPreSendBudgetResetsAfterSuccessfulHop(t *testing.T) {
refused := &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED}
var calls int32
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
switch atomic.AddInt32(&calls, 1) {
case 1, 2:
// Two connect failures, spending the pre-send budget under the old
// lifetime accounting.
return nil, refused
case 3:
// A hop that reaches the server. It declines, so the request was not
// accepted and retrying stays safe.
return &http.Response{StatusCode: http.StatusTooManyRequests, Body: http.NoBody, Request: r}, nil
case 4:
// One more connect blip. With a lifetime cap this surfaces as an
// error; the call should still recover.
return nil, refused
default:
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Request: r}, nil
}
})}
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://origin.invalid/v1", []byte("{}"), nil, 3)
if err != nil {
t.Fatalf("interleaved connect and status failures returned %v; the pre-send budget did not reset after the hop that reached the server", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
t.Fatalf("status = %d, want 200", resp.StatusCode)
}
if got := atomic.LoadInt32(&calls); got != 5 {
t.Fatalf("transport called %d times, want 5 (dial, dial, 429, dial, 200)", got)
}
}
+14
View File
@@ -212,6 +212,14 @@ func credentialDenyReadPathsForEnvironment(env credentialPathEnvironment, allowR
candidates = append(candidates,
filepath.Join(home, ".aws"),
filepath.Join(home, ".azure"),
// git's credential store backend, which holds host passwords and
// personal access tokens in cleartext. Denied rather than the whole
// of ~/.ssh, because these cost nothing functionally: git reads them
// through a credential helper for authentication, not for identity,
// so a sandboxed git still works and simply cannot authenticate as
// the user. SSH key material is a harder trade and is tracked
// separately (#815).
filepath.Join(home, ".git-credentials"),
)
}
npmUserConfig := strings.TrimSpace(env.NPMUserConfig)
@@ -233,6 +241,12 @@ func credentialDenyReadPathsForEnvironment(env credentialPathEnvironment, allowR
if configHome != "" {
candidates = append(candidates,
filepath.Join(configHome, "zero"),
// The XDG location of the same git credential store. Denying the file
// rather than the directory is deliberate: ~/.config/git also holds
// the global git config, which userGitConfigReadPaths grants on
// purpose so a sandboxed git can read user.name and aliases instead
// of failing outright.
filepath.Join(configHome, "git", "credentials"),
)
}
cloudSDKConfig := strings.TrimSpace(env.CloudSDKConfig)
+93
View File
@@ -0,0 +1,93 @@
package sandbox
import (
"os"
"path/filepath"
"testing"
)
// git's credential store holds host passwords and personal access tokens in
// cleartext, in one of two locations depending on whether the user is on the
// XDG layout. Neither was denied, so a sandboxed command could read them
// (#815).
//
// Scoped to the credential files on purpose. Denying ~/.ssh as well would stop
// a sandboxed git push over SSH from working, which is a functional trade that
// issue tracks separately; these two cost nothing, because git reads them for
// authentication rather than identity.
func TestCredentialDenyReadPathsCoversGitCredentialStores(t *testing.T) {
home := t.TempDir()
configHome := filepath.Join(home, ".config")
if err := os.MkdirAll(filepath.Join(configHome, "git"), 0o700); err != nil {
t.Fatal(err)
}
// The builder only emits paths that exist, so these have to be real.
gitCredentials := filepath.Join(home, ".git-credentials")
xdgCredentials := filepath.Join(configHome, "git", "credentials")
gitConfig := filepath.Join(configHome, "git", "config")
for _, path := range []string{gitCredentials, xdgCredentials, gitConfig} {
if err := os.WriteFile(path, []byte("x"), 0o600); err != nil {
t.Fatal(err)
}
}
denied := credentialDenyReadPathsForEnvironment(credentialPathEnvironment{
Home: home,
ConfigHome: configHome,
}, nil)
covered := func(target string) bool {
norm := normalizeProfilePath(target)
for _, entry := range denied {
if entry == norm || pathWithinRoot(entry, norm) {
return true
}
}
return false
}
if !covered(gitCredentials) {
t.Fatalf("~/.git-credentials is readable; deny list = %v", denied)
}
if !covered(xdgCredentials) {
t.Fatalf("~/.config/git/credentials is readable; deny list = %v", denied)
}
// The global git config must stay readable. userGitConfigReadPaths grants it
// deliberately so a sandboxed git can read user.name and aliases rather than
// failing with "unable to access", so denying the directory instead of the
// credential file would break that on purpose-built setups.
if covered(gitConfig) {
t.Fatalf("~/.config/git/config was denied; a sandboxed git loses its identity config")
}
}
// An explicit read grant still wins, matching how every other entry behaves: a
// user who deliberately scopes a workspace over one of these paths is not
// overridden by the default deny.
//
// Asserted from both sides on purpose. Checking only that the path is absent
// under an allowRead is satisfied just as well by a deny rule that was never
// added, so the run without allowRead has to establish that there is something
// there for the grant to override.
func TestGitCredentialDenyYieldsToExplicitAllowRead(t *testing.T) {
home := t.TempDir()
gitCredentials := filepath.Join(home, ".git-credentials")
if err := os.WriteFile(gitCredentials, []byte("x"), 0o600); err != nil {
t.Fatal(err)
}
target := normalizeProfilePath(gitCredentials)
listed := func(entries []string) bool {
for _, entry := range entries {
if entry == target {
return true
}
}
return false
}
if !listed(credentialDenyReadPathsForEnvironment(credentialPathEnvironment{Home: home}, nil)) {
t.Fatalf("~/.git-credentials is not denied without an allowRead; nothing for the grant to override")
}
if listed(credentialDenyReadPathsForEnvironment(credentialPathEnvironment{Home: home}, []string{home})) {
t.Fatalf("explicit allowRead %q did not re-include the credential store", home)
}
}
+7
View File
@@ -1027,6 +1027,13 @@ func scrubSensitiveEnv(env []string, additionalKeys ...string) []string {
"GH_TOKEN",
"ZERO_WEBSEARCH_API_KEY",
"ZERO_DAEMON_REMOTE_TOKEN",
// The file form of the same bridge token. TokenFromEnv accepts either,
// so scrubbing only the inline variable left the pointer readable, and
// the default sandbox posture is read-all. There is no fallback
// location to guess at, the path comes from this variable alone, so
// removing it closes the leak rather than half of it. Same reasoning as
// GOOGLE_APPLICATION_CREDENTIALS below.
"ZERO_DAEMON_REMOTE_TOKEN_FILE",
}
for _, descriptor := range providercatalog.All() {
for _, key := range descriptor.AuthEnvVars {
+5
View File
@@ -699,6 +699,11 @@ func TestScrubSensitiveEnv(t *testing.T) {
"XAI_API_KEY=xai-12345",
"HUGGINGFACE_API_KEY=hf_12345",
"GOOGLE_APPLICATION_CREDENTIALS=/home/user/sa-key.json",
// Both forms of the daemon bridge token. The inline one was already
// scrubbed; the file form was not, which left a sandboxed command the
// path to read it from (#677).
"ZERO_DAEMON_REMOTE_TOKEN=bridge-token-inline",
"ZERO_DAEMON_REMOTE_TOKEN_FILE=/home/user/.zero/remote-token",
"COMPANY_LLM_SECRET=custom-secret",
"ZERO_OAUTH_MY_SVC_CLIENT_SECRET=oauth-secret",
"zero_oauth_second_client_secret=case-insensitive-secret",
@@ -9,6 +9,7 @@ import (
"os"
"os/exec"
"path/filepath"
"strconv"
"strings"
"testing"
"time"
@@ -386,6 +387,127 @@ func runWindowsRealSmokeCommand(t *testing.T, runnerExe string, base WindowsSand
}
}
// The write jail must hold on a path whose DACL grants Everyone write access.
//
// A WRITE_RESTRICTED token runs TWO checks for a write and needs both to pass:
// the normal one against its enabled SIDs, and a second against its RESTRICTED
// SID list. The jail is built on the second check only succeeding where Zero has
// explicitly ACL'd a capability SID. Putting the World SID (S-1-1-0) in the
// restricted list breaks that globally — every principal is a member of
// Everyone, so on a DACL that grants Everyone write, the restricted half passes
// for free and confinement collapses back to the ordinary user's own
// permissions, which is precisely the boundary the sandbox exists to be
// stricter than.
//
// The runner already states this rule for the SIDs it refuses to add ("None of
// the granted SIDs can be added to the restricted list without collapsing the
// write jail"); Everyone was in the list anyway, since the original sandbox
// baseline.
//
// Realistic rather than theoretical: administrators open share roots with
// Everyone:F, and third-party installers ship loose ACLs. It needs no privilege,
// no symlink and no race. Checked and ruled out: C:\Users\Public\Documents
// grants BATCH, not Everyone, so the common "public folder" case is not this.
func TestWindowsRestrictedTokenDeniesWritesToEveryoneWritablePaths(t *testing.T) {
if os.Getenv("ZERO_SANDBOX_REAL_SMOKE") != "1" {
t.Skip("set ZERO_SANDBOX_REAL_SMOKE=1 to run real Windows sandbox smoke tests")
}
runnerExe := realSmokeExecutable(t, "ZERO_WINDOWS_COMMAND_RUNNER_EXE", WindowsSandboxCommandRunnerName)
root := t.TempDir()
outside := t.TempDir()
everyoneDir := filepath.Join(outside, "everyone-writable")
if err := os.MkdirAll(everyoneDir, 0o700); err != nil {
t.Fatalf("MkdirAll everyone-writable: %v", err)
}
// Granted through the production ACL applier, so the hostile DACL is built
// the same way a real one is rather than by a test-only shortcut.
snapshot, applied, err := applyWindowsACLPathGroup(windowsACLPathGroup{
Path: everyoneDir,
Entries: []WindowsACLEntry{{
Action: WindowsACLAllowWrite,
Path: everyoneDir,
Capability: "S-1-1-0",
}},
})
if err != nil {
t.Fatalf("grant Everyone write: %v", err)
}
if !applied {
t.Fatal("precondition: the Everyone grant did not apply, so there is no bypass to test for")
}
t.Cleanup(func() { _ = rollbackWindowsACLSnapshots([]windowsACLSnapshot{snapshot}) })
sandboxHome := filepath.Join(root, ".zero-sandbox")
// No DenyRead on purpose: that is what makes the runner choose the
// WRITE_RESTRICTED token, which is the default posture and the one whose
// restricted-SID list this test is about.
profile := PermissionProfile{
FileSystem: FileSystemPolicy{
Kind: FileSystemRestricted,
ReadRoots: []string{root},
WriteRoots: []WritableRoot{{Root: root, ProtectedMetadataNames: []string{".git", ".zero", ".agents"}}},
IncludePlatformRoots: true,
AllowTemp: true,
},
Network: NetworkPolicy{Mode: NetworkDeny},
}
config := WindowsSandboxCommandArgsOptions{
SandboxHome: sandboxHome,
CommandCWD: root,
WorkspaceRoots: []string{root},
PermissionProfile: profile,
SandboxLevel: WindowsSandboxLevelUnelevated,
}
// Granted root first. A jail that denies everything would satisfy the two
// assertions below while being useless, and this separates the two outcomes.
insideMarker := filepath.Join(root, "write-ok.txt")
runWindowsRealSmokeCommand(t, runnerExe, config, []string{
"cmd.exe", "/d", "/s", "/c", "echo ok>" + insideMarker,
}, 0)
if contents, err := os.ReadFile(insideMarker); err != nil || strings.TrimSpace(string(contents)) != "ok" {
t.Fatalf("sandboxed write to a granted root = %q, %v; want ok", contents, err)
}
// Control: an ordinary directory outside every granted root. This is the
// behaviour the Everyone case must match.
controlMarker := filepath.Join(outside, "control-denied.txt")
runWindowsRealSmokeCommand(t, runnerExe, config, deniedWriteCommand(controlMarker), deniedWriteExitCode)
if _, err := os.Stat(controlMarker); err == nil {
t.Fatal("precondition: the sandbox allowed a write to an ordinary path outside every granted root, so this test cannot measure the Everyone case")
}
// The assertion. Same as the control in every respect except the DACL.
everyoneMarker := filepath.Join(everyoneDir, "everyone-denied.txt")
runWindowsRealSmokeCommand(t, runnerExe, config, deniedWriteCommand(everyoneMarker), deniedWriteExitCode)
if _, err := os.Stat(everyoneMarker); err == nil {
t.Error("the sandbox wrote outside every granted root because the path grants Everyone write; the restricted-SID list is satisfied by a SID every principal carries")
} else if !os.IsNotExist(err) {
t.Errorf("stat the Everyone-writable marker: %v", err)
}
}
// deniedWriteExitCode is chosen so a denial cannot be confused with the runner
// failing to start the command at all.
//
// The obvious spelling of these assertions is `echo leaked>path` expecting exit
// 1, but the runner itself exits 1 for its own errors — a failed marker
// validation, a bad argument, a token it could not build. A test written that
// way passes both when the sandbox denies the write and when nothing ever ran,
// and the second case proves nothing while looking identical to success. On a
// test whose whole job is catching a confinement regression, that is a failure
// mode to design out rather than hope about.
//
// 77 is arbitrary beyond being outside the range the runner produces for itself.
const deniedWriteExitCode = 77
// deniedWriteCommand attempts a write and reports deniedWriteExitCode when the
// redirect is refused, so the exit code also proves cmd.exe actually ran.
func deniedWriteCommand(marker string) []string {
return []string{"cmd.exe", "/d", "/s", "/c", "echo leaked>" + marker + " || exit " + strconv.Itoa(deniedWriteExitCode)}
}
func powershellSingleQuote(value string) string {
out := "'"
for _, r := range value {
+45 -9
View File
@@ -88,19 +88,55 @@ func createWindowsRestrictedTokenFromBase(base windows.Token, capabilitySIDs []w
if err != nil {
return 0, err
}
worldSID, err := windows.CreateWellKnownSid(windows.WinWorldSid)
if err != nil {
return 0, fmt.Errorf("create world SID: %w", err)
}
entries := make([]windows.SIDAndAttributes, 0, len(capabilitySIDs)+2)
for _, sid := range capabilitySIDs {
entries = append(entries, windows.SIDAndAttributes{Sid: sid.sid})
}
entries = append(entries,
windows.SIDAndAttributes{Sid: sidFromBytes(logonSID)},
windows.SIDAndAttributes{Sid: worldSID},
)
entries = append(entries, windows.SIDAndAttributes{Sid: sidFromBytes(logonSID)})
// The World SID (S-1-1-0, Everyone) is added ONLY to the token that does not
// carry WRITE_RESTRICTED, and putting it back unconditionally would reopen a
// write-jail bypass.
//
// A restricted SID is a key to every object whose DACL names it. That is why
// the runner refuses to add the user SID, Administrators or SYSTEM — "each
// has write access nearly everywhere". Everyone is the broadest of the lot:
// every principal carries it, so under WRITE_RESTRICTED the second
// (restricted-SID) check passes for free on any path whose DACL grants
// Everyone write, and confinement falls back to the ordinary user's own
// permissions — the exact boundary this token exists to be stricter than. It
// needs no privilege, no symlink and no race: an Everyone-writable directory
// is enough, and share roots opened with Everyone:F and loose installer ACLs
// supply them. It was present from the original sandbox baseline,
// uncommented, and under WRITE_RESTRICTED nothing depends on it, because that
// flag already exempts reads from the restricted-SID check.
//
// Without the flag it IS load-bearing and cannot simply be dropped. The
// restricted-SID check then applies to READS too, and default Windows DACLs
// grant BUILTIN\Users rather than anything in this list, so a token without
// Everyone cannot open cmd.exe — the process dies at launch with
// STATUS_ACCESS_DENIED (0xC0000022) before it runs anything. That path is
// only taken when the profile configures DenyRead, which is already the
// posture that trades capability for read-deny enforcement (#612).
//
// So the bypass survives for DenyRead profiles, deliberately and narrowly,
// rather than being traded for a sandbox that cannot start a command. Closing
// it there needs a different mechanism (a read-side grant that is not a
// universal group), tracked separately.
//
// The logon SID above stays in both modes: it is this token's own rather than
// a broad group, and broadenWindowsRestrictedTokenDefaultDacl depends on it so
// the process can use pipes and events it creates for itself.
//
// Anything added here needs the same scrutiny — Authenticated Users, Users,
// INTERACTIVE and BATCH would each produce this bypass on a DACL naming them.
if !writeRestricted {
worldSID, err := windows.CreateWellKnownSid(windows.WinWorldSid)
if err != nil {
return 0, fmt.Errorf("create world SID: %w", err)
}
entries = append(entries, windows.SIDAndAttributes{Sid: worldSID})
}
// WRITE_RESTRICTED scopes the restricted-SID check to write-type accesses:
// reads use only the normal token identity, so the sandboxed process can
+26 -2
View File
@@ -43,7 +43,7 @@ func TestUpdateTitle(t *testing.T) {
t.Fatalf("title not trimmed/stored: %q", updated.Title)
}
if updated.UpdatedAt != before.UpdatedAt {
t.Fatalf("UpdatedAt must not change on retitle: before=%q after=%q", before.UpdatedAt, updated.UpdatedAt)
t.Fatalf("UpdatedAt must not change on rename: before=%q after=%q", before.UpdatedAt, updated.UpdatedAt)
}
if updated.EventCount != before.EventCount {
t.Fatalf("EventCount changed: before=%d after=%d", before.EventCount, updated.EventCount)
@@ -67,7 +67,7 @@ func TestUpdateTitle(t *testing.T) {
// An unchanged title is a no-op (still succeeds).
if _, err := store.UpdateTitle(session.SessionID, "Clean Generated Title"); err != nil {
t.Fatalf("no-op retitle should succeed: %v", err)
t.Fatalf("no-op rename should succeed: %v", err)
}
// An invalid session id is rejected.
@@ -76,6 +76,30 @@ func TestUpdateTitle(t *testing.T) {
}
}
func TestUpdateTitleIfCurrent(t *testing.T) {
store := newTitleTestStore(t)
session, err := store.Create(CreateInput{Title: "Automatic title"})
if err != nil {
t.Fatalf("create: %v", err)
}
updated, applied, err := store.UpdateTitleIfCurrent(session.SessionID, "Automatic title", "Generated title")
if err != nil || !applied || updated.Title != "Generated title" {
t.Fatalf("matching update = (%#v, %v, %v)", updated, applied, err)
}
if _, err := store.UpdateTitle(session.SessionID, "Manual title"); err != nil {
t.Fatalf("manual rename: %v", err)
}
updated, applied, err = store.UpdateTitleIfCurrent(session.SessionID, "Generated title", "Late generated title")
if err != nil {
t.Fatalf("stale update: %v", err)
}
if applied || updated.Title != "Manual title" {
t.Fatalf("stale update should preserve manual title: (%#v, %v)", updated, applied)
}
}
func TestUpdateModel(t *testing.T) {
store := newTitleTestStore(t)
session, err := store.Create(CreateInput{ModelID: "model-a"})
+35 -1
View File
@@ -719,7 +719,7 @@ func (store *Store) appendPreparedEventsLocked(sessionID string, inputs []prepar
// serialized under the same per-session lock as AppendEvent and re-reads the
// latest metadata under that lock before rewriting, so a concurrent append can't
// clobber the new title (nor the title clobber a concurrent append's event
// count/timestamp). UpdatedAt is deliberately left untouched: a retitle is not
// count/timestamp). UpdatedAt is deliberately left untouched: a rename is not
// activity, so it must not reorder the session in the resumable list. A blank
// title is rejected so a failed model generation can never erase a useful
// first-message title, and an unchanged title is a no-op (no rewrite/fsync).
@@ -751,6 +751,40 @@ func (store *Store) UpdateTitle(sessionID string, title string) (Metadata, error
return session, nil
}
// UpdateTitleIfCurrent replaces a session title only when it still matches
// expected. It lets background automatic naming avoid overwriting a newer manual
// rename while keeping the check and write under the same per-session lock.
func (store *Store) UpdateTitleIfCurrent(sessionID string, expected string, title string) (Metadata, bool, error) {
if !ValidSessionID(sessionID) {
return Metadata{}, false, fmt.Errorf("invalid zero session id %q", sessionID)
}
trimmed := strings.TrimSpace(title)
if trimmed == "" {
return Metadata{}, false, fmt.Errorf("zero session title is required")
}
unlock, err := store.lockSession(sessionID)
if err != nil {
return Metadata{}, false, err
}
defer unlock()
session, err := store.readMetadata(sessionID)
if err != nil {
return Metadata{}, false, err
}
if session.Title != strings.TrimSpace(expected) {
return session, false, nil
}
if session.Title == trimmed {
return session, true, nil
}
session.Title = trimmed
if err := store.writeMetadata(session); err != nil {
return Metadata{}, false, err
}
return session, true, nil
}
// UpdateModel replaces a session's selected model without changing its activity
// timestamp or event counters. An empty model clears the session override.
func (store *Store) UpdateModel(sessionID string, modelID string) (Metadata, error) {
+180 -27
View File
@@ -12,6 +12,12 @@ import (
// team is at its slot cap the member is queued and launches when a slot frees
// (it stays pending in the coordinator until then).
func (s *Swarm) Spawn(pol Policy, teamName, agentType, task, cwd string) (string, error) {
release, err := s.beginLifecycleAdmission()
if err != nil {
return "", err
}
defer release()
def, err := s.registry.Lookup(agentType)
if err != nil {
return "", err
@@ -23,15 +29,44 @@ func (s *Swarm) Spawn(pol Policy, teamName, agentType, task, cwd string) (string
}
s.rememberCwd(id, cwd)
spec := s.buildSpec(pol, id, id, team, def, task, cwd)
s.dispatch(spec)
s.dispatchAdmitted(spec)
return id, nil
}
// beginLifecycleAdmission admits one unit of lifecycle work if shutdown has not
// begun. On success it returns a release func the caller MUST call exactly once
// (a deferred call is the norm) — Close blocks until every admitted unit has
// released. On shutdown it returns a nil func and ErrSwarmClosed, and there is
// nothing to release.
//
// It deliberately does NOT keep lifecycleMu held for the caller's duration, which
// an earlier revision did. Close needs the write lock before it can cancel, so a
// launcher that blocks in Launch until its context is cancelled would deadlock
// against every Close caller: the launch holds the read lock while Close waits
// for the write lock to issue the cancellation neither side can reach. Holding the
// lock only across the flag check and the counter increment keeps the ordering
// guarantee (no admission once closed is set) without letting external launcher
// code sit inside the critical section.
//
// The read lock is enough to make the WaitGroup safe: Close flips closed under
// the write lock, so an Add here either completes before Close acquires it (and is
// therefore visible to the later Wait) or observes closed and is refused.
func (s *Swarm) beginLifecycleAdmission() (func(), error) {
s.lifecycleMu.RLock()
defer s.lifecycleMu.RUnlock()
if s.closed {
return nil, ErrSwarmClosed
}
s.lifecycleWork.Add(1)
return s.lifecycleWork.Done, nil
}
// dispatch admits a spec to its team (launching now or queuing for a slot).
func (s *Swarm) dispatch(spec MemberSpec) {
// The caller must hold an admission ticket from beginLifecycleAdmission.
func (s *Swarm) dispatchAdmitted(spec MemberSpec) {
t := s.team(spec.Team)
if t.admit(spec) {
s.launch(t, spec)
s.launchAdmitted(t, spec)
}
// Otherwise the spec is queued; the coordinator task stays pending until a
// slot frees and afterExit launches it.
@@ -39,17 +74,73 @@ func (s *Swarm) dispatch(spec MemberSpec) {
// launch starts a member for spec and supervises it. A synchronous launch
// failure fails the task and frees the slot.
func (s *Swarm) launch(t *Team, spec MemberSpec) {
handle, err := s.launcher.Launch(s.baseCtx, spec)
// The caller must hold an admission ticket from beginLifecycleAdmission.
func (s *Swarm) launchAdmitted(t *Team, spec MemberSpec) {
handle, err := s.launchMemberAdmitted(spec)
if err != nil {
_ = s.coord.Fail(spec.TaskID, "launch: "+err.Error())
s.afterExit(t)
if err == ErrSwarmClosed {
t.releaseSlot()
return
}
s.afterExitAdmitted(t)
return
}
m := &Member{ID: spec.ID, AgentType: spec.AgentType, TaskID: spec.TaskID, handle: handle}
t.addMember(m)
_ = s.coord.SetStatus(spec.TaskID, StatusRunning)
go s.watch(t, m, spec)
committed := s.commitLaunch(handle, func() {
m := &Member{ID: spec.ID, AgentType: spec.AgentType, TaskID: spec.TaskID, handle: handle}
t.addMember(m)
_ = s.coord.SetStatus(spec.TaskID, StatusRunning)
s.watchers.Add(1)
go func() {
defer s.watchers.Done()
s.watch(t, m, spec)
}()
})
if !committed {
_ = s.coord.Fail(spec.TaskID, "launch: "+ErrSwarmClosed.Error())
t.releaseSlot()
}
}
// commitLaunch closes the window between MemberLauncher.Launch returning and the
// swarm taking ownership of the handle it produced. Launch is external code and
// therefore runs without lifecycleMu held, so Close may have won after
// launchMemberAdmitted's pre-check but before Launch returned. adopt runs under
// the read lock so the shutdown check and everything the caller does to take
// ownership (registration, status transition, watcher Add, handle swap) are
// atomic with respect to Close setting closed and beginning its waits.
//
// It reports whether the handle was adopted. When shutdown won, the handle is
// reaped before returning false: its launch context was already cancelled by
// Close, and MemberHandle has no separate cancellation operation — Launch's
// context is its cancellation contract — so waiting is how a
// successfully-created process/goroutine avoids being abandoned. The caller
// records the terminal outcome for the task it was launching.
func (s *Swarm) commitLaunch(handle MemberHandle, adopt func()) bool {
s.lifecycleMu.RLock()
if s.closed {
s.lifecycleMu.RUnlock()
_, _ = handle.Wait()
return false
}
defer s.lifecycleMu.RUnlock()
adopt()
return true
}
// launchMemberAdmitted re-checks shutdown at the last internal boundary before
// invoking the external launcher. A lifecycle ticket proves the operation won
// admission before Close, but the caller may not reach its launch until after
// Close has set closed; in that case no new member should be started.
func (s *Swarm) launchMemberAdmitted(spec MemberSpec) (MemberHandle, error) {
s.lifecycleMu.RLock()
closed := s.closed
s.lifecycleMu.RUnlock()
if closed {
return nil, ErrSwarmClosed
}
return s.launcher.Launch(s.baseCtx, spec)
}
// watch awaits a member, applies bounded relaunch on temporary failures, records
@@ -60,42 +151,98 @@ func (s *Swarm) launch(t *Team, spec MemberSpec) {
// life (Member.ID == MemberSpec.ID); a retry never reuses the struct for a
// different spec.
func (s *Swarm) watch(t *Team, m *Member, spec MemberSpec) {
res, err := m.handle.Wait()
if err != nil {
if isRetryable(err) && m.restarts < maxMemberRestarts {
if nh, relErr := s.launcher.Launch(s.baseCtx, spec); relErr == nil {
m.restarts++
m.handle = nh
go s.watch(t, m, spec)
return
for {
res, err := m.handle.Wait()
if err != nil {
if isRetryable(err) && m.restarts < maxMemberRestarts && s.relaunchAdmitted(m, spec) {
continue
}
// fall through to record the original error if relaunch fails
// Fall through if shutdown started or the relaunch failed.
// res.SessionID is preserved by the handle even on error, so a member that
// ran then failed stays drillable; a pure launch error carries an empty id.
_ = s.coord.FailWithSession(m.TaskID, memberError(err), res.SessionID)
} else {
_ = s.coord.CompleteWithSession(m.TaskID, res.Result, res.SessionID)
}
// res.SessionID is preserved by the handle even on error, so a member that
// ran then failed stays drillable; a pure launch error carries an empty id.
_ = s.coord.FailWithSession(m.TaskID, memberError(err), res.SessionID)
} else {
_ = s.coord.CompleteWithSession(m.TaskID, res.Result, res.SessionID)
break
}
t.removeMember(m.ID)
s.afterExit(t)
}
// relaunchAdmitted performs one bounded relaunch of an existing member's spec,
// reporting whether the member was relaunched and its new handle adopted. A
// relaunch is lifecycle work like any other, so it runs the same three gates an
// initial launch does: admission (refused outright once shutdown began), the
// pre-Launch re-check, and the post-Launch commit — the last of which matters
// because the production FuncLauncher returns a live handle immediately without
// consulting its context, so a Close landing between the pre-check and the return
// would otherwise resume supervising a member started after shutdown. When any
// gate refuses, the caller records the member's terminal failure instead.
func (s *Swarm) relaunchAdmitted(m *Member, spec MemberSpec) bool {
release, err := s.beginLifecycleAdmission()
if err != nil {
return false
}
defer release()
handle, err := s.launchMemberAdmitted(spec)
if err != nil {
return false
}
return s.commitLaunch(handle, func() {
m.restarts++
m.handle = handle
})
}
// afterExit releases the just-vacated slot and launches the next queued spec, if
// any. Each exit drains at most one queued member; that member's own exit drains
// the next, so the queue empties one-per-slot without unbounded recursion.
func (s *Swarm) afterExit(t *Team) {
release, err := s.beginLifecycleAdmission()
if err != nil {
t.releaseSlot()
return
}
defer release()
s.afterExitAdmitted(t)
}
// afterExitAdmitted drains at most one queued spec while lifecycle admission is
// held open. The caller must hold an admission ticket from beginLifecycleAdmission.
func (s *Swarm) afterExitAdmitted(t *Team) {
next, ok := t.onExit()
if !ok {
return
}
s.launch(t, next)
// This call's own ticket only proves shutdown hadn't begun when it was
// granted; Close may have set closed in the meantime, and t.onExit above
// already removed next from t.queue, so Close's clearQueue sweep can no
// longer see it (nor will it wait for one, since this ticket is what Close's
// lifecycleWork.Wait() is blocked on). Re-check before launching so a dequeue
// that straddles the shutdown boundary fails the task in place of a launch a
// shutting-down swarm never wanted started.
s.lifecycleMu.RLock()
closed := s.closed
s.lifecycleMu.RUnlock()
if closed {
t.releaseSlot()
_ = s.coord.Fail(next.TaskID, ErrSwarmClosed.Error())
return
}
s.launchAdmitted(t, next)
}
// Handoff transfers a task to a fresh member of toAgentType, delivering a note to
// the new member's inbox and marking the original task handed-off. It returns the
// new task id. A handoff of an already-terminal task is rejected (fail closed).
func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string) (string, error) {
release, err := s.beginLifecycleAdmission()
if err != nil {
return "", err
}
defer release()
task, ok := s.coord.Get(taskID)
if !ok {
return "", fmt.Errorf("%w: %s", ErrUnknownTask, taskID)
@@ -131,7 +278,7 @@ func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string)
// Retire the original task (it has been re-delegated).
_ = s.coord.SetStatus(taskID, StatusHandedOff)
spec := s.buildSpec(pol, newID, newID, team, def, handoffTask, cwd)
s.dispatch(spec)
s.dispatchAdmitted(spec)
return newID, nil
}
@@ -139,6 +286,12 @@ func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string)
// (e.g. a crashed worker) onto fresh members of toAgentType, returning the
// adopted task ids. Terminal tasks and tasks with a live owner are left alone.
func (s *Swarm) AdoptOrphans(pol Policy, teamName, toAgentType string) ([]string, error) {
release, err := s.beginLifecycleAdmission()
if err != nil {
return nil, err
}
defer release()
def, err := s.registry.Lookup(toAgentType)
if err != nil {
return nil, err
@@ -161,7 +314,7 @@ func (s *Swarm) AdoptOrphans(pol Policy, teamName, toAgentType string) ([]string
cwd := s.cwdFor(task.ID)
s.rememberCwd(task.ID, cwd)
spec := s.buildSpec(pol, newAgent, task.ID, team, def, task.Description, cwd)
s.dispatch(spec)
s.dispatchAdmitted(spec)
adopted = append(adopted, task.ID)
}
return adopted, nil
+720
View File
@@ -2,6 +2,7 @@ package swarm
import (
"context"
"errors"
"strings"
"sync"
"testing"
@@ -217,6 +218,141 @@ func TestPermanentErrorNoRetry(t *testing.T) {
}
}
func TestLifecycleAdmissionRejectedAfterClose(t *testing.T) {
l := newLauncher(okFor)
sw := newSwarmFor(t, l)
sw.Close()
if _, err := sw.Spawn(Policy{}, "team", "teammate", "task", ""); !errors.Is(err, ErrSwarmClosed) {
t.Fatalf("Spawn after Close error = %v, want ErrSwarmClosed", err)
}
if _, err := sw.Handoff(Policy{}, "team", "task", "teammate", "note"); !errors.Is(err, ErrSwarmClosed) {
t.Fatalf("Handoff after Close error = %v, want ErrSwarmClosed", err)
}
if _, err := sw.AdoptOrphans(Policy{}, "team", "teammate"); !errors.Is(err, ErrSwarmClosed) {
t.Fatalf("AdoptOrphans after Close error = %v, want ErrSwarmClosed", err)
}
if got := len(l.recorded()); got != 0 {
t.Fatalf("launches after Close = %d, want 0", got)
}
}
func TestCloseDoesNotLaunchQueuedMembers(t *testing.T) {
gate := make(chan struct{})
l := newLauncher(okFor)
l.gate = gate
sw := newSwarmFor(t, l)
var ids []string
for i := 0; i < 3; i++ {
id, err := sw.Spawn(Policy{}, "team", "teammate", "task", "")
if err != nil {
t.Fatalf("Spawn %d: %v", i, err)
}
ids = append(ids, id)
}
if got := len(l.recorded()); got != 2 {
t.Fatalf("initial launches = %d, want 2 with one queued", got)
}
sw.Close()
if got := len(l.recorded()); got != 2 {
t.Fatalf("launches after Close = %d, want queued member not launched", got)
}
for _, id := range ids {
task, ok := sw.Coordinator().Get(id)
if !ok || !task.Status.terminal() {
t.Fatalf("task %s after Close = %+v, want terminal", id, task)
}
}
}
func TestClosePreventsMemberRetry(t *testing.T) {
started := make(chan struct{}, maxMemberRestarts+1)
release := make(chan struct{})
l := FuncLauncher{Run: func(context.Context, MemberSpec) (MemberResult, error) {
started <- struct{}{}
<-release
return MemberResult{}, ErrMemberTemporary
}}
sw := newSwarmFor(t, l)
_, err := sw.Spawn(Policy{}, "team", "teammate", "task", "")
if err != nil {
t.Fatalf("Spawn: %v", err)
}
select {
case <-started:
case <-time.After(3 * time.Second):
t.Fatal("initial member did not start")
}
closed := make(chan struct{})
go func() {
sw.Close()
close(closed)
}()
waitFor(t, "swarm closed state", func() bool {
sw.lifecycleMu.RLock()
defer sw.lifecycleMu.RUnlock()
return sw.closed
})
close(release)
select {
case <-closed:
case <-time.After(3 * time.Second):
t.Fatal("Close did not return after retryable member exit")
}
select {
case <-started:
t.Fatal("member retried after Close")
default:
}
}
func TestCloseWaitsForMemberWatchers(t *testing.T) {
release := make(chan struct{})
l := FuncLauncher{Run: func(context.Context, MemberSpec) (MemberResult, error) {
<-release
return MemberResult{Result: "done"}, nil
}}
sw := newSwarmFor(t, l)
id, err := sw.Spawn(Policy{}, "team", "teammate", "task", "")
if err != nil {
t.Fatalf("Spawn: %v", err)
}
waitFor(t, "task running", func() bool {
task, ok := sw.Coordinator().Get(id)
return ok && task.Status == StatusRunning
})
const callers = 3
closing := make(chan struct{}, callers)
closed := make(chan struct{}, callers)
for i := 0; i < callers; i++ {
go func() {
closing <- struct{}{}
sw.Close()
closed <- struct{}{}
}()
}
for i := 0; i < callers; i++ {
<-closing
}
select {
case <-closed:
t.Fatal("a Close caller returned before the member watcher exited")
case <-time.After(50 * time.Millisecond):
}
close(release)
for i := 0; i < callers; i++ {
select {
case <-closed:
case <-time.After(3 * time.Second):
t.Fatal("all Close callers did not return after the member watcher exited")
}
}
}
func TestHandoffDeliversNoteAndRetiresOriginal(t *testing.T) {
gate := make(chan struct{})
l := newLauncher(okFor)
@@ -364,3 +500,587 @@ func TestSpawnUnknownAgentType(t *testing.T) {
t.Fatal("Spawn with unknown agent type must error")
}
}
// blockingLauncher blocks inside Launch until its context is cancelled — the
// shape of a real launcher that waits on a slot, a daemon connection, or a
// sandbox handshake before it can return a handle.
type blockingLauncher struct {
entered chan struct{}
once sync.Once
}
func (l *blockingLauncher) Launch(ctx context.Context, spec MemberSpec) (MemberHandle, error) {
l.once.Do(func() { close(l.entered) })
<-ctx.Done()
return nil, ctx.Err()
}
// TestCloseCancelsLauncherBlockedInLaunch is the regression test for holding
// lifecycle admission across MemberLauncher.Launch. When admission was a read lock
// held for the caller's whole duration, this deadlocked: the spawn sat inside
// Launch waiting for its context, while Close waited for the write lock it needed
// before it could cancel that very context. Neither side could progress, so both
// the spawn and every Close caller hung.
//
// Admission is now a counted ticket, so Close can flip the flag, cancel, and then
// wait the ticket out.
func TestCloseCancelsLauncherBlockedInLaunch(t *testing.T) {
launcher := &blockingLauncher{entered: make(chan struct{})}
sw := newSwarmFor(t, launcher)
spawned := make(chan error, 1)
go func() {
_, err := sw.Spawn(Policy{}, "team", "teammate", "task", "")
spawned <- err
}()
select {
case <-launcher.entered:
case <-time.After(3 * time.Second):
t.Fatal("launcher never entered Launch")
}
closed := make(chan struct{})
go func() {
sw.Close()
close(closed)
}()
select {
case <-closed:
case <-time.After(5 * time.Second):
t.Fatal("Close deadlocked against a launcher blocked in Launch")
}
select {
case <-spawned:
case <-time.After(3 * time.Second):
t.Fatal("Spawn never returned after Close cancelled the launch context")
}
}
// TestCloseWaitsForAdmittedSpawn pins the other half of the contract: Close is a
// barrier, so it must not return while an admitted spawn is still running. Without
// the lifecycleWork wait, Close could finish while this launch was mid-flight.
func TestCloseWaitsForAdmittedSpawn(t *testing.T) {
entered := make(chan struct{})
finish := make(chan struct{})
launcher := &gatedLaunchLauncher{entered: entered, finish: finish}
sw := newSwarmFor(t, launcher)
go func() {
_, _ = sw.Spawn(Policy{}, "team", "teammate", "task", "")
}()
select {
case <-entered:
case <-time.After(3 * time.Second):
t.Fatal("launcher never entered Launch")
}
closed := make(chan struct{})
go func() {
sw.Close()
close(closed)
}()
select {
case <-closed:
t.Fatal("Close returned while an admitted spawn was still in flight")
case <-time.After(200 * time.Millisecond):
}
close(finish)
select {
case <-closed:
case <-time.After(5 * time.Second):
t.Fatal("Close did not return after the admitted spawn finished")
}
}
// TestAdmittedDispatchDoesNotLaunchAfterClose covers an operation that won a
// lifecycle ticket before shutdown but did not reach dispatch until after Close
// set closed. The ticket keeps Close from returning early, but it must not permit
// a new external launch with the already-cancelled swarm context.
func TestAdmittedDispatchDoesNotLaunchAfterClose(t *testing.T) {
l := newLauncher(okFor)
sw := newSwarmFor(t, l)
sw.maxTeamSize = 1
if _, err := sw.coord.Register("late-task", "late-agent", "team", "late work"); err != nil {
t.Fatalf("Register late task: %v", err)
}
release, err := sw.beginLifecycleAdmission()
if err != nil {
t.Fatalf("beginLifecycleAdmission: %v", err)
}
closed := make(chan struct{})
go func() {
sw.Close()
close(closed)
}()
waitFor(t, "swarm closed state", func() bool {
sw.lifecycleMu.RLock()
defer sw.lifecycleMu.RUnlock()
return sw.closed
})
sw.dispatchAdmitted(MemberSpec{
ID: "late-agent", TaskID: "late-task", AgentType: "teammate", Team: "team",
})
release()
select {
case <-closed:
case <-time.After(3 * time.Second):
t.Fatal("Close did not return after the admitted dispatch finished")
}
if got := len(l.recorded()); got != 0 {
t.Fatalf("launches after closed was set = %d, want 0", got)
}
team := sw.team("team")
if got := team.Running(); got != 0 {
t.Fatalf("team running after rejected launch = %d, want 0", got)
}
task, ok := sw.Coordinator().Get("late-task")
if !ok {
t.Fatal("late task vanished")
}
if task.Status != StatusFailed || !strings.Contains(task.Err, ErrSwarmClosed.Error()) {
t.Fatalf("late task after Close = %+v, want failed with ErrSwarmClosed", task)
}
}
// TestCloseBetweenLaunchPrecheckAndReturn rejects and reaps a handle created by
// a Launch that straddles shutdown. The launch has passed the pre-check when it
// enters the gate; Close then sets closed and cancels its context before the
// launcher is allowed to return successfully.
func TestCloseBetweenLaunchPrecheckAndReturn(t *testing.T) {
launcher := &closeRaceLauncher{
entered: make(chan struct{}),
finish: make(chan struct{}),
waited: make(chan struct{}),
}
sw := newSwarmFor(t, launcher)
spawned := make(chan string, 1)
go func() {
id, _ := sw.Spawn(Policy{}, "team", "teammate", "task", "")
spawned <- id
}()
select {
case <-launcher.entered:
case <-time.After(3 * time.Second):
t.Fatal("launcher never entered Launch")
}
closed := make(chan struct{})
go func() {
sw.Close()
close(closed)
}()
waitFor(t, "swarm closed state", func() bool {
sw.lifecycleMu.RLock()
defer sw.lifecycleMu.RUnlock()
return sw.closed
})
close(launcher.finish)
var id string
select {
case id = <-spawned:
case <-time.After(3 * time.Second):
t.Fatal("Spawn did not return after Launch")
}
select {
case <-closed:
case <-time.After(3 * time.Second):
t.Fatal("Close did not return after late handle was reaped")
}
select {
case <-launcher.waited:
default:
t.Fatal("late handle was not reaped")
}
team := sw.team("team")
if got := team.Running(); got != 0 {
t.Fatalf("team running after close-raced Launch = %d, want 0", got)
}
if got := len(team.liveAgents()); got != 0 {
t.Fatalf("registered members after close-raced Launch = %d, want 0", got)
}
task, ok := sw.Coordinator().Get(id)
if !ok {
t.Fatal("task vanished")
}
if task.Status != StatusFailed || !strings.Contains(task.Err, ErrSwarmClosed.Error()) {
t.Fatalf("task after close-raced Launch = %+v, want failed with ErrSwarmClosed", task)
}
}
type closeRaceLauncher struct {
entered chan struct{}
finish chan struct{}
waited chan struct{}
once sync.Once
}
func (l *closeRaceLauncher) Launch(ctx context.Context, spec MemberSpec) (MemberHandle, error) {
l.once.Do(func() { close(l.entered) })
<-l.finish
return &closeRaceHandle{id: spec.ID, ctx: ctx, waited: l.waited}, nil
}
type closeRaceHandle struct {
id string
ctx context.Context
waited chan struct{}
once sync.Once
}
func (h *closeRaceHandle) ID() string { return h.id }
func (h *closeRaceHandle) Wait() (MemberResult, error) {
<-h.ctx.Done()
h.once.Do(func() { close(h.waited) })
return MemberResult{}, h.ctx.Err()
}
// TestCloseBetweenRetryLaunchPrecheckAndReturn is the retry counterpart of
// TestCloseBetweenLaunchPrecheckAndReturn: a bounded relaunch wins admission and
// passes the pre-check, then Close lands while the launcher is inside Launch. The
// relaunched handle must be reaped rather than adopted, so the watcher never
// resumes supervising a member started after shutdown began — a relaunch whose
// member goes on to succeed must not be recorded as the task's outcome.
func TestCloseBetweenRetryLaunchPrecheckAndReturn(t *testing.T) {
launcher := &retryCloseRaceLauncher{
entered: make(chan struct{}),
finish: make(chan struct{}),
late: &retryLateHandle{id: "teammate-1"},
}
sw := newSwarmFor(t, launcher)
id, err := sw.Spawn(Policy{}, "team", "teammate", "task", "")
if err != nil {
t.Fatalf("Spawn: %v", err)
}
// The first member exits retryable, so the watcher is now inside the relaunch.
select {
case <-launcher.entered:
case <-time.After(3 * time.Second):
t.Fatal("launcher never entered the relaunch Launch")
}
closed := make(chan struct{})
go func() {
sw.Close()
close(closed)
}()
waitFor(t, "swarm closed state", func() bool {
sw.lifecycleMu.RLock()
defer sw.lifecycleMu.RUnlock()
return sw.closed
})
close(launcher.finish)
select {
case <-closed:
case <-time.After(3 * time.Second):
t.Fatal("Close did not return after the late relaunch handle was reaped")
}
if got := launcher.launchCount(); got != 2 {
t.Fatalf("launch attempts = %d, want 2 (initial + one relaunch)", got)
}
if got := launcher.late.waitCount(); got != 1 {
t.Fatalf("late relaunch handle waits = %d, want 1 (reaped, not supervised)", got)
}
team := sw.team("team")
if got := team.Running(); got != 0 {
t.Fatalf("team running after close-raced relaunch = %d, want 0", got)
}
if got := len(team.liveAgents()); got != 0 {
t.Fatalf("registered members after close-raced relaunch = %d, want 0", got)
}
task, ok := sw.Coordinator().Get(id)
if !ok {
t.Fatal("task vanished")
}
if task.Status != StatusFailed {
t.Fatalf("task after close-raced relaunch = %+v, want failed, not the late member's outcome", task)
}
if task.Result != "" {
t.Fatalf("task result after close-raced relaunch = %q, want the late member's result discarded", task.Result)
}
}
// retryCloseRaceLauncher fails its first member with a retryable error, then
// blocks inside the relaunch's Launch until the test releases it.
type retryCloseRaceLauncher struct {
mu sync.Mutex
calls int
entered chan struct{}
finish chan struct{}
late *retryLateHandle
once sync.Once
}
func (l *retryCloseRaceLauncher) Launch(_ context.Context, spec MemberSpec) (MemberHandle, error) {
l.mu.Lock()
l.calls++
first := l.calls == 1
l.mu.Unlock()
if first {
return &temporaryHandle{id: spec.ID}, nil
}
l.once.Do(func() { close(l.entered) })
<-l.finish
return l.late, nil
}
func (l *retryCloseRaceLauncher) launchCount() int {
l.mu.Lock()
defer l.mu.Unlock()
return l.calls
}
// temporaryHandle exits immediately with a retryable error, driving the watcher
// into its bounded relaunch path.
type temporaryHandle struct{ id string }
func (h *temporaryHandle) ID() string { return h.id }
func (h *temporaryHandle) Wait() (MemberResult, error) {
return MemberResult{}, ErrMemberTemporary
}
// retryLateHandle models a member whose work does not consult the cancelled
// launch context and reports success. Adopting it after shutdown would record a
// post-close result on the task, so a correct shutdown reaps it exactly once and
// discards what it returns.
type retryLateHandle struct {
id string
mu sync.Mutex
waits int
}
func (h *retryLateHandle) ID() string { return h.id }
func (h *retryLateHandle) Wait() (MemberResult, error) {
h.mu.Lock()
h.waits++
h.mu.Unlock()
return MemberResult{Result: "late result", SessionID: "late-session"}, nil
}
func (h *retryLateHandle) waitCount() int {
h.mu.Lock()
defer h.mu.Unlock()
return h.waits
}
func TestCloseRacesLifecycleAdmission(t *testing.T) {
gate := make(chan struct{})
l := newLauncher(okFor)
l.gate = gate
sw := newSwarmFor(t, l)
const spawns = 64
start := make(chan struct{})
results := make(chan struct {
id string
err error
}, spawns)
var callers sync.WaitGroup
callers.Add(spawns)
for i := 0; i < spawns; i++ {
go func() {
defer callers.Done()
<-start
id, err := sw.Spawn(Policy{}, "team", "teammate", "racing task", "")
results <- struct {
id string
err error
}{id: id, err: err}
}()
}
closed := make(chan struct{})
go func() {
<-start
sw.Close()
close(closed)
}()
close(start)
callers.Wait()
<-closed
close(results)
for result := range results {
if result.err != nil {
if !errors.Is(result.err, ErrSwarmClosed) {
t.Fatalf("Spawn racing Close error = %v, want ErrSwarmClosed", result.err)
}
continue
}
task, ok := sw.Coordinator().Get(result.id)
if !ok || !task.Status.terminal() {
t.Fatalf("admitted task %q after Close = %+v, want terminal", result.id, task)
}
}
for _, task := range sw.Coordinator().List() {
if !task.Status.terminal() {
t.Fatalf("task %q remained %s after Close", task.ID, task.Status)
}
}
}
// TestCloseFailsQueuedSpecOnLateCreatedTeam is the regression test for queue
// sweeping after shutdown begins. Spawn/Handoff/AdoptOrphans calls that obtained
// tickets before closed flipped can still create a team and append to its queue.
// Close must wait for those calls before both snapshotting teams and sweeping
// their queues, or a late-created team's queued task remains pending forever.
//
// This drives that interleaving directly: two operations obtain tickets before
// Close, then the first creates and fills a previously unseen one-slot team and
// the second queues behind it after Close has flipped closed.
func TestCloseFailsQueuedSpecOnLateCreatedTeam(t *testing.T) {
sw := newSwarmFor(t, newLauncher(okFor))
sw.maxTeamSize = 1
if _, err := sw.coord.Register("late-task", "late-agent", "team", "late work"); err != nil {
t.Fatalf("Register late task: %v", err)
}
firstRelease, err := sw.beginLifecycleAdmission()
if err != nil {
t.Fatalf("first beginLifecycleAdmission: %v", err)
}
secondRelease, err := sw.beginLifecycleAdmission()
if err != nil {
firstRelease()
t.Fatalf("second beginLifecycleAdmission: %v", err)
}
runningSpec := MemberSpec{ID: "running-agent", TaskID: "running-task", AgentType: "teammate", Team: "team"}
queuedSpec := MemberSpec{ID: "late-agent", TaskID: "late-task", AgentType: "teammate", Team: "team"}
closed := make(chan struct{})
go func() {
sw.Close()
close(closed)
}()
waitFor(t, "swarm closed state", func() bool {
sw.lifecycleMu.RLock()
defer sw.lifecycleMu.RUnlock()
return sw.closed
})
select {
case <-closed:
t.Fatal("Close returned while an admitted ticket was still open")
default:
}
// Neither operation touched a team before Close took its old, buggy snapshot.
// Finish the team-admission portion of both operations while their tickets are
// still held: the first creates the team and fills its only slot; the second
// appends to that late-created team's queue.
team := sw.team("team")
if !team.admit(runningSpec) {
t.Fatal("first spec should take the late-created team's only slot")
}
firstRelease()
if team.admit(queuedSpec) {
t.Fatal("second spec should queue on the late-created team")
}
secondRelease()
select {
case <-closed:
case <-time.After(3 * time.Second):
t.Fatal("Close did not return after the admitted ticket released")
}
if got := team.QueueDepth(); got != 0 {
t.Fatalf("queue depth after Close = %d, want 0 (late spec swept)", got)
}
task, ok := sw.Coordinator().Get("late-task")
if !ok {
t.Fatal("late-queued task vanished")
}
if task.Status != StatusFailed {
t.Fatalf("late-queued task status = %v, want %v (failed, not stranded)", task.Status, StatusFailed)
}
if !strings.Contains(task.Err, ErrSwarmClosed.Error()) {
t.Fatalf("late-queued task err = %q, want it to mention %q", task.Err, ErrSwarmClosed.Error())
}
}
// TestAfterExitAdmittedFailsDequeuedSpecAfterClose is the regression test for
// jatmn's second #776 finding: afterExitAdmitted used to dequeue a queued spec
// (t.onExit) and launch it unconditionally once its own ticket was granted,
// even though Close can flip closed in the gap between that dequeue — which
// makes the spec invisible to Close's clearQueue sweep — and the launch. A
// launcher that performs work before honoring its cancelled context could then
// start a brand new worker after shutdown had already begun.
func TestAfterExitAdmittedFailsDequeuedSpecAfterClose(t *testing.T) {
l := newLauncher(okFor)
sw := newSwarmFor(t, l)
if _, err := sw.coord.Register("queued-task", "queued-agent", "team", "queued work"); err != nil {
t.Fatalf("Register: %v", err)
}
spec := MemberSpec{ID: "queued-agent", TaskID: "queued-task", AgentType: "teammate", Team: "team"}
// Occupy the only slot, then queue spec behind it.
team := &Team{Name: "team", members: map[string]*Member{}, maxSize: 1}
if !team.admit(MemberSpec{ID: "running-agent", TaskID: "running-task"}) {
t.Fatal("first spec should take the only slot")
}
if team.admit(spec) {
t.Fatal("second spec should queue, not launch, over the cap")
}
// Simulate shutdown racing in between this call's ticket admission (already
// granted, hence no beginLifecycleAdmission call here) and its dequeue+launch,
// exactly like the finding describes: closed is already set by the time
// afterExitAdmitted's post-dequeue check runs.
sw.lifecycleMu.Lock()
sw.closed = true
sw.lifecycleMu.Unlock()
sw.afterExitAdmitted(team) // running-agent "exits" -> dequeues spec
if got := team.Running(); got != 0 {
t.Fatalf("team running after close-raced drain = %d, want 0 (slot released, not held for a launch)", got)
}
task, ok := sw.Coordinator().Get("queued-task")
if !ok {
t.Fatal("dequeued task vanished")
}
if task.Status != StatusFailed {
t.Fatalf("dequeued task status = %v, want %v (failed, not launched)", task.Status, StatusFailed)
}
if !strings.Contains(task.Err, ErrSwarmClosed.Error()) {
t.Fatalf("dequeued task err = %q, want it to mention %q", task.Err, ErrSwarmClosed.Error())
}
for _, s := range l.recorded() {
if s.ID == "queued-agent" {
t.Fatal("a spec dequeued after shutdown began must never launch")
}
}
}
// gatedLaunchLauncher blocks in Launch until the test releases it, ignoring
// context cancellation so the test controls exactly when the admitted work ends.
type gatedLaunchLauncher struct {
entered chan struct{}
finish chan struct{}
once sync.Once
}
func (l *gatedLaunchLauncher) Launch(_ context.Context, spec MemberSpec) (MemberHandle, error) {
l.once.Do(func() { close(l.entered) })
<-l.finish
return &funcHandle{id: spec.ID, done: closedChan()}, nil
}
func closedChan() chan struct{} {
ch := make(chan struct{})
close(ch)
return ch
}
+12 -13
View File
@@ -121,11 +121,12 @@ type Scheduler struct {
ctx context.Context
cancel context.CancelFunc
mu sync.Mutex
jobs map[string]*scheduledJob
closed bool
seq atomic.Uint64
wg sync.WaitGroup
mu sync.Mutex
jobs map[string]*scheduledJob
closed bool
seq atomic.Uint64
wg sync.WaitGroup
closeOnce sync.Once
}
// newScheduler builds a Scheduler bound to sw. Its context derives from the
@@ -229,15 +230,13 @@ func (s *Scheduler) List() []JobStatus {
// Close cancels every job and waits for their loops to exit. Safe to call more
// than once.
func (s *Scheduler) Close() {
s.mu.Lock()
if s.closed {
s.closeOnce.Do(func() {
s.mu.Lock()
s.closed = true
s.cancel()
s.mu.Unlock()
return
}
s.closed = true
s.cancel()
s.mu.Unlock()
s.wg.Wait()
s.wg.Wait()
})
}
// run is one job's timing loop. It requests a fresh one-shot ticker per interval
+19
View File
@@ -164,6 +164,25 @@ func TestSchedulerCloseStopsJobs(t *testing.T) {
}
}
func TestSchedulerCreatedAfterSwarmCloseIsClosed(t *testing.T) {
sw := newSwarmFor(t, newLauncher(okFor))
sw.Close()
sched := sw.Scheduler()
sched.mu.Lock()
closed := sched.closed
sched.mu.Unlock()
if !closed {
t.Fatal("scheduler created after Swarm.Close was not closed")
}
if _, err := sched.Add(Policy{}, "team", "teammate", "t", "", Schedule{Every: time.Hour}); err == nil {
t.Fatal("Add on a scheduler created after Swarm.Close should fail")
}
if got := len(sched.List()); got != 0 {
t.Fatalf("scheduler created after Swarm.Close has %d jobs, want 0", got)
}
}
func TestSchedulerAddRejectedAfterContextCancel(t *testing.T) {
sw := newSwarmFor(t, newLauncher(okFor))
sched := sw.Scheduler()
+103 -13
View File
@@ -64,6 +64,18 @@ type Swarm struct {
baseCtx context.Context
cancel context.CancelFunc
// lifecycleMu guards closed and orders admission against shutdown. It is held
// only for the flag check plus the lifecycleWork increment — never across a
// launcher call, which would deadlock Close (see beginLifecycleAdmission).
lifecycleMu sync.RWMutex
closed bool
// lifecycleWork counts admitted-but-unfinished lifecycle operations (spawn,
// handoff, adoption, retry, queue drain). Close waits it out so no admitted
// operation is still touching the swarm when Close returns.
lifecycleWork sync.WaitGroup
watchers sync.WaitGroup
closeOnce sync.Once
mu sync.Mutex
teams map[string]*Team
taskCwd map[string]string // taskID -> cwd, for handoff/adoption relaunch
@@ -137,31 +149,93 @@ func New(opts Options) (*Swarm, error) {
}
// Close cancels every running member's context and releases resources. It is
// safe to call more than once. The scheduler is closed first so no new spawn
// fires after shutdown begins.
// safe to call more than once. It closes lifecycle admission before canceling
// members, fails queued tasks, and waits for every member watcher to exit.
func (s *Swarm) Close() {
s.mu.Lock()
sched := s.scheduler
s.mu.Unlock()
if sched != nil {
sched.Close()
}
if s.cancel != nil {
s.cancel()
}
s.closeOnce.Do(func() {
// Closing admission and cancelling are separate steps on purpose. The flag
// flips under lifecycleMu so no further work is admitted, and the lock is
// released BEFORE cancelling: a member already being launched holds no lock
// here, but it may be blocked inside MemberLauncher.Launch waiting for the
// context this cancel provides, and it can only release its admission ticket
// once that cancellation lands.
s.lifecycleMu.Lock()
s.closed = true
s.mu.Lock()
sched := s.scheduler
s.mu.Unlock()
s.lifecycleMu.Unlock()
if s.cancel != nil {
s.cancel()
}
if sched != nil {
sched.Close()
}
// Wait for every already-admitted lifecycle operation (spawn, handoff,
// adoption, retry, queue drain) to finish BEFORE sweeping queues below.
// Such work can create teams as well as append to their queues, so snapshot
// the teams only after every ticket is back. At that point no team or queue
// can gain a new entry, and the sweep cannot miss late-admitted work.
s.lifecycleWork.Wait()
s.mu.Lock()
teams := make([]*Team, 0, len(s.teams))
for _, team := range s.teams {
teams = append(teams, team)
}
s.mu.Unlock()
for _, team := range teams {
for _, spec := range team.clearQueue() {
_ = s.coord.Fail(spec.TaskID, ErrSwarmClosed.Error())
}
}
// Every watcher is added while its admitting lifecycle ticket is still held,
// so lifecycleWork.Wait above also orders every Add before this Wait. The
// watcher itself can outlive that ticket, which is why both waits are needed.
s.watchers.Wait()
})
}
// Scheduler returns the swarm's recurring-spawn scheduler, creating it on first
// use. Scheduling is opt-in: until a job is added the scheduler does nothing.
//
// The post-close handoff between this method and Close relies on the ORDER in
// which each reads/writes s.scheduler and s.closed, not on holding one lock across
// both: Close takes lifecycleMu.Lock() (excluding this RLock), sets closed, reads
// whatever s.scheduler currently is, and only THEN releases lifecycleMu and calls
// sched.Close() on what it captured. So a call here either:
// - runs entirely before Close's write lock — sees closed == false, and the
// scheduler it creates/returns is exactly the one Close's later snapshot will
// capture and close; or
// - runs entirely after Close's write lock — sees closed == true (memory model:
// the RLock/Lock pair orders the read after Close's write) and closes the
// scheduler itself here, because Close already captured its own snapshot
// (possibly nil, if this call is what first creates one) and will not look
// again.
//
// Either way exactly one of the two call sites closes the scheduler that ends up
// installed, and neither can observe a "created after close but never closed"
// scheduler — that would require reading closed == false here while also reading
// a scheduler Close had already passed over, which the lock ordering rules out.
func (s *Swarm) Scheduler() *Scheduler {
s.lifecycleMu.RLock()
closed := s.closed
s.mu.Lock()
defer s.mu.Unlock()
if s.scheduler == nil {
s.scheduler = newScheduler(s)
}
return s.scheduler
sched := s.scheduler
s.mu.Unlock()
s.lifecycleMu.RUnlock()
if closed {
sched.Close()
}
return sched
}
// ErrSwarmClosed reports lifecycle work submitted after shutdown begins.
var ErrSwarmClosed = errors.New("swarm: closed")
// rememberCwd records a task's working dir so a handoff/adoption relaunch keeps it.
func (s *Swarm) rememberCwd(taskID, cwd string) {
s.mu.Lock()
@@ -288,6 +362,22 @@ func (t *Team) onExit() (MemberSpec, bool) {
return MemberSpec{}, false
}
func (t *Team) releaseSlot() {
t.mu.Lock()
if t.running > 0 {
t.running--
}
t.mu.Unlock()
}
func (t *Team) clearQueue() []MemberSpec {
t.mu.Lock()
defer t.mu.Unlock()
queued := t.queue
t.queue = nil
return queued
}
func (t *Team) addMember(m *Member) {
t.mu.Lock()
t.members[m.ID] = m
+97 -14
View File
@@ -80,8 +80,19 @@ func (tool applyPatchTool) RunWithOptions(ctx context.Context, args map[string]a
return errorResult("Error applying patch: " + err.Error())
}
var createdTargets []string
var fullySuppliedTargets []string
wholeBefore := map[string]bool{}
if options.FileTracker != nil {
createdTargets = missingPatchTargets(applyRoot, patch)
fullySuppliedTargets = completeCreatedPatchTargets(applyRoot, patch)
for _, path := range patchHeaderPaths(patch) {
if path == "" || path == "/dev/null" {
continue
}
if absolute, _, rerr := resolveWorkspaceTargetPath(applyRoot, path); rerr == nil {
wholeBefore[absolute] = options.FileTracker.SeenWhole(absolute)
}
}
}
command := exec.CommandContext(ctx, "git", "apply", "--whitespace=nowarn", patchPath)
@@ -102,13 +113,28 @@ func (tool applyPatchTool) RunWithOptions(ctx context.Context, args map[string]a
result := okResult(summary)
result.ChangedFiles = changedFilesFromPatch(relativeRoot, patch)
result.Display = Display{Summary: summary, Kind: "diff", Preview: capPreviewDiff(patch)}
// git apply already rejects a patch whose context drifted, so it has its own
// staleness guard. Drop any tracked baseline for the files it rewrote so a
// subsequent edit_file/write_file re-reads instead of false-flagging the
// patch's own change as an external modification.
fullySupplied := make(map[string]bool, len(fullySuppliedTargets))
for _, absolute := range fullySuppliedTargets {
fullySupplied[absolute] = true
}
// Re-baseline files changed by this tool. When the model had already seen the
// whole input (or supplied a complete new file), the exact patch plus that
// input determines the whole output, so a follow-up edit needs no wasted read.
// Partial observations remain conservative and are cleared by Record.
for _, changed := range result.ChangedFiles {
if absolute, _, rerr := resolveScopedPath(tool.workspaceRoot, tool.scope, changed); rerr == nil {
options.FileTracker.Forget(absolute)
content, readErr := os.ReadFile(absolute)
if readErr != nil {
options.FileTracker.Forget(absolute)
continue
}
info, _ := os.Stat(absolute)
wasWhole := wholeBefore[absolute] || fullySupplied[absolute]
options.FileTracker.Record(absolute, content, info)
if wasWhole {
lines := lineCount(string(content))
options.FileTracker.RecordSeenRange(absolute, 1, lines, lines)
}
}
}
recordCreatedPatchTargets(options.FileTracker, createdTargets)
@@ -134,6 +160,60 @@ func missingPatchTargets(root string, patch string) []string {
return missing
}
// completeCreatedPatchTargets returns only files whose full contents are
// supplied by a /dev/null creation patch. A missing rename/copy destination is
// created by git too, but its bytes come from an unread source and must not gain
// whole-file observation credit.
func completeCreatedPatchTargets(root string, patch string) []string {
seen := map[string]bool{}
var created []string
oldRemaining, newRemaining := 0, 0
inHunk := false
fromDevNull := false
for _, line := range strings.Split(strings.ReplaceAll(patch, "\r\n", "\n"), "\n") {
if inHunk && (oldRemaining > 0 || newRemaining > 0) {
switch {
case strings.HasPrefix(line, "-"):
oldRemaining--
case strings.HasPrefix(line, "+"):
newRemaining--
case strings.HasPrefix(line, "\\"):
default:
oldRemaining--
newRemaining--
}
continue
}
inHunk = false
switch {
case strings.HasPrefix(line, "diff --git "):
fromDevNull = false
case strings.HasPrefix(line, "@@"):
oldRemaining, newRemaining = parseHunkCounts(line)
inHunk = oldRemaining > 0 || newRemaining > 0
case strings.HasPrefix(line, "--- "):
fromDevNull = patchFileHeaderPath(line) == "/dev/null"
case strings.HasPrefix(line, "+++ "):
path := patchFileHeaderPath(line)
if !fromDevNull || path == "" || path == "/dev/null" {
fromDevNull = false
continue
}
fromDevNull = false
absolute, _, err := resolveWorkspaceTargetPath(root, stripPatchPrefix(path))
if err != nil || seen[absolute] {
continue
}
if _, err := os.Stat(absolute); !os.IsNotExist(err) {
continue
}
seen[absolute] = true
created = append(created, absolute)
}
}
return created
}
func recordCreatedPatchTargets(tracker *FileTracker, missingBefore []string) {
if tracker == nil {
return
@@ -240,15 +320,7 @@ func patchHeaderPaths(patch string) []string {
oldRemaining, newRemaining = parseHunkCounts(line)
inHunk = oldRemaining > 0 || newRemaining > 0
case strings.HasPrefix(line, "--- "), strings.HasPrefix(line, "+++ "):
// Take everything after the fixed 4-char prefix instead of splitting on
// spaces, so a path that contains spaces survives. Drop a trailing
// tab-delimited timestamp (unified-diff convention) and unquote a
// C-quoted git path (git quotes names with spaces/specials) (L18).
rest := line[len("--- "):] // "--- " and "+++ " are both 4 bytes
if tab := strings.IndexByte(rest, '\t'); tab >= 0 {
rest = rest[:tab]
}
if p := strings.TrimSpace(unquoteGitPath(rest)); p != "" && p != "/dev/null" {
if p := patchFileHeaderPath(line); p != "" && p != "/dev/null" {
paths = append(paths, stripPatchPrefix(p))
}
}
@@ -256,6 +328,17 @@ func patchHeaderPaths(patch string) []string {
return paths
}
func patchFileHeaderPath(line string) string {
if len(line) < len("--- ") {
return ""
}
rest := line[len("--- "):] // "--- " and "+++ " are both 4 bytes
if tab := strings.IndexByte(rest, '\t'); tab >= 0 {
rest = rest[:tab]
}
return strings.TrimSpace(unquoteGitPath(rest))
}
// parseHunkCounts reads the old/new line counts from a "@@ -a,b +c,d @@" header.
// A missing count (e.g. "@@ -a +c @@") means 1 per unified-diff convention.
//
+1
View File
@@ -43,6 +43,7 @@ func NewScopedBashTool(workspaceRoot string, scope PathScope) Tool {
baseTool: baseTool{
name: "bash",
description: "Execute a shell command inside the workspace (or an explicitly granted extra directory) after permission is granted. " + shellGuidance,
deferred: true,
parameters: Schema{
Type: "object",
Properties: map[string]PropertySchema{
+8 -1
View File
@@ -625,7 +625,14 @@ func TestBashToolRequireEscalatedMsysGuard(t *testing.T) {
t.Run("approved require_escalated still blocks an unrelated syntax issue", func(t *testing.T) {
result := registry.RunWithOptions(context.Background(), "bash", map[string]any{
"command": `cd /d/tmp/zero-pr-158 && dir`,
// No `&&` here on purpose. The POSIX-style path is what this case is
// about, and on Windows PowerShell 5.1 a chained command trips the
// separate windows_powershell_version preflight first, which shadows
// the block being asserted. That only shows up on 5.1: the runners
// have PowerShell 7, where && is valid and the version check does not
// fire, so the substitution kept this green in CI while failing on the
// shell a stock Windows box actually gets.
"command": `cd /d/tmp/zero-pr-158`,
"sandbox_permissions": string(SandboxPermissionsRequireEscalated),
}, RunOptions{
PermissionGranted: true,
+71
View File
@@ -74,6 +74,11 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
if cerr := options.FileTracker.CheckConflict(absolutePath, contentBytes); cerr != nil {
return errorResult(fileConflictMessage(relativePath))
}
if options.FileTracker != nil {
if _, tracked := options.FileTracker.Version(absolutePath); !tracked {
return errorResult(fileUnseenMessage(relativePath))
}
}
content := string(contentBytes)
occurrences := strings.Count(content, oldString)
@@ -131,7 +136,11 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
if !replaceAll && occurrences > 1 {
return errorResult(fmt.Sprintf("Error: old_string matches %d locations in %s. Either make old_string more specific, or pass replace_all: true to replace every occurrence.", occurrences, relativePath))
}
if options.FileTracker != nil && !allOccurrencesSeen(options.FileTracker, absolutePath, content, oldString, replaceAll) {
return errorResult(fileUnseenMessage(relativePath))
}
previouslySeenWhole := options.FileTracker.SeenWhole(absolutePath)
updated := strings.Replace(content, oldString, newString, 1)
replacedCount := 1
if replaceAll {
@@ -141,12 +150,14 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
if updated == content {
return okResult("No changes: new_string is identical to old_string.")
}
editedSpans := replacementByteSpans(content, oldString, newString, replaceAll)
if err := recheckScopedWriteTarget(tool.workspaceRoot, tool.scope, requestedPath); err != nil {
return errorResult("Error writing " + relativePath + ": " + err.Error())
}
if err := os.WriteFile(absolutePath, []byte(updated), 0o644); err != nil {
return errorResult("Error writing " + relativePath + ": " + err.Error())
}
modelKnownContent := updated
// Optional format-on-write (ZERO_FORMAT_ON_WRITE). Must run BEFORE the
// FileTracker re-baseline: recording pre-format content would make the very
// next edit look like an external modification and trip the conflict guard.
@@ -155,6 +166,15 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
// compare against the current on-disk state, not the pre-edit version.
newInfo, _ := os.Stat(absolutePath)
options.FileTracker.Record(absolutePath, []byte(updated), newInfo)
if updated == modelKnownContent {
if previouslySeenWhole {
options.FileTracker.RecordSeenRange(absolutePath, 1, lineCount(updated), lineCount(updated))
} else {
for _, span := range editedSpans {
options.FileTracker.RecordSeenBytes(absolutePath, span.start, span.end, len(updated))
}
}
}
suffix := ""
if replacedCount != 1 {
@@ -169,3 +189,54 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
result.Display = Display{Summary: fmt.Sprintf("Edited %s", relativePath), Kind: "diff", Preview: boundedUnifiedDiff(relativePath, content, updated)}
return result
}
// replacementByteSpans returns every half-open post-edit byte range whose exact
// contents came from newString. Offsets are translated from the original with a
// cumulative delta so replace_all remains correct after earlier replacements.
func replacementByteSpans(content, oldString, newString string, replaceAll bool) []lineRange {
spans := make([]lineRange, 0, 1)
offset := 0
delta := 0
for {
index := strings.Index(content[offset:], oldString)
if index < 0 {
return spans
}
index += offset
start := index + delta
spans = append(spans, lineRange{start: start, end: start + len(newString)})
if !replaceAll {
return spans
}
offset = index + len(oldString)
delta += len(newString) - len(oldString)
}
}
func allOccurrencesSeen(tracker *FileTracker, path, content, oldString string, replaceAll bool) bool {
offset := 0
for {
index := strings.Index(content[offset:], oldString)
if index < 0 {
return true
}
index += offset
start, end := lineSpanForOffset(content, index, len(oldString))
if !tracker.SeenBytes(path, index, index+len(oldString)) && !tracker.SeenRange(path, start, end) {
return false
}
if !replaceAll {
return true
}
offset = index + len(oldString)
}
}
func lineSpanForOffset(content string, offset, length int) (int, int) {
if offset < 0 {
return 1, 1
}
start := strings.Count(content[:offset], "\n") + 1
end := start + strings.Count(content[offset:min(len(content), offset+length)], "\n")
return start, end
}
+97 -22
View File
@@ -18,7 +18,7 @@ import (
)
func TestIndependentExecCommandConstructorsShareDefaultManager(t *testing.T) {
root := t.TempDir()
root := execTestRoot(t)
execTool := NewScopedExecCommandTool(root, nil, nil)
writeTool := NewWriteStdinTool(nil)
@@ -90,7 +90,7 @@ func TestExecCommandToolDescribesHostShellSyntax(t *testing.T) {
}
func TestExecCommandReturnsSessionAndWriteStdinPollsCompletion(t *testing.T) {
root := t.TempDir()
root := execTestRoot(t)
manager := newExecSessionManager()
execTool := NewScopedExecCommandTool(root, nil, manager)
writeTool := NewWriteStdinTool(manager)
@@ -138,7 +138,7 @@ func TestExecCommandReturnsSessionAndWriteStdinPollsCompletion(t *testing.T) {
}
func TestExecCommandRequireEscalatedBypassesNativeSandboxAfterApproval(t *testing.T) {
root := t.TempDir()
root := execTestRoot(t)
manager := newExecSessionManager()
registry := NewRegistry()
registry.Register(NewScopedExecCommandTool(root, nil, manager))
@@ -187,7 +187,7 @@ func TestExecCommandRequireEscalatedBypassesMsysGuardAfterApproval(t *testing.T)
if runtime.GOOS != "windows" {
t.Skip("windows-only MSYS sandbox guard")
}
root := t.TempDir()
root := execTestRoot(t)
manager := newExecSessionManager()
registry := NewRegistry()
registry.Register(NewScopedExecCommandTool(root, nil, manager))
@@ -224,7 +224,7 @@ func TestExecCommandRequireEscalatedBypassesMsysGuardAfterApproval(t *testing.T)
}
func TestExecCommandReturnsExitCodeWhenCommandCompletesDuringInitialYield(t *testing.T) {
root := t.TempDir()
root := execTestRoot(t)
manager := newExecSessionManager()
execTool := NewScopedExecCommandTool(root, nil, manager)
@@ -263,7 +263,7 @@ func TestExecCommandReportsWorkspaceChanges(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("uses POSIX shell syntax")
}
root := t.TempDir()
root := execTestRoot(t)
result := NewScopedExecCommandTool(root, nil, newExecSessionManager()).Run(context.Background(), map[string]any{
"cmd": "mkdir -p src node_modules/pkg && printf 'export {}' > src/main.ts && printf generated > node_modules/pkg/index.js",
"yield_time_ms": 30000,
@@ -357,7 +357,7 @@ func executionRequestHasCapability(request execution.Request, kind execution.Cap
}
func TestExecCommandForegroundServerReturnsSessionAndServesHTTP(t *testing.T) {
root := t.TempDir()
root := execTestRoot(t)
manager := newExecSessionManager()
execTool := NewScopedExecCommandTool(root, nil, manager)
writeTool := NewWriteStdinTool(manager)
@@ -411,7 +411,7 @@ func parseListeningAddress(output string) string {
}
func TestExecCommandReapsFinishedUnpolledSession(t *testing.T) {
root := t.TempDir()
root := execTestRoot(t)
manager := execution.NewProcessManager(execution.ProcessManagerOptions{CompletedRetention: 10 * time.Millisecond})
execTool := NewScopedExecCommandTool(root, nil, manager)
@@ -537,7 +537,7 @@ func TestWriteStdinInterruptTerminatesSession(t *testing.T) {
}
func TestWriteStdinRejectsInputForNonTTYSession(t *testing.T) {
root := t.TempDir()
root := execTestRoot(t)
manager := newExecSessionManager()
execTool := NewScopedExecCommandTool(root, nil, manager)
writeTool := NewWriteStdinTool(manager)
@@ -568,9 +568,43 @@ func TestWriteStdinRejectsInputForNonTTYSession(t *testing.T) {
manager.Stop(sessionID)
}
// execTestRoot is t.TempDir with a cleanup that tolerates Windows still holding
// the directory.
//
// A terminated process does not release its handles the instant the session
// reports it exited, and Windows refuses to remove a directory with an open
// handle where POSIX does not. t.TempDir's own cleanup removes once and fails
// the test on a sharing violation, which turns ordinary teardown timing into a
// test failure with nothing actually wrong underneath it.
func execTestRoot(t *testing.T) string {
t.Helper()
root, err := os.MkdirTemp("", "zero-exec-test-")
if err != nil {
t.Fatalf("create test root: %v", err)
}
t.Cleanup(func() {
deadline := time.Now().Add(10 * time.Second)
for {
if err := os.RemoveAll(root); err == nil {
return
}
if time.Now().After(deadline) {
// Deliberately not a failure. The directory is under the OS temp
// root and the process holding it has already been terminated, so
// the worst case is a stale directory, not a broken test.
t.Logf("test root %s still held after the cleanup deadline; leaving it", root)
return
}
time.Sleep(50 * time.Millisecond)
}
})
return root
}
func TestWriteStdinStopIntentTerminatesNonTTYSession(t *testing.T) {
for _, chars := range []string{`\u0003`, "exit\n"} {
root := t.TempDir()
root := execTestRoot(t)
manager := newExecSessionManager()
execTool := NewScopedExecCommandTool(root, nil, manager)
writeTool := NewWriteStdinTool(manager)
@@ -587,16 +621,46 @@ func TestWriteStdinStopIntentTerminatesNonTTYSession(t *testing.T) {
t.Fatalf("session_id is not numeric: %v", err)
}
result := writeTool.Run(context.Background(), map[string]any{
"session_id": sessionID,
"chars": chars,
"yield_time_ms": 1000,
})
if result.Status != StatusOK {
t.Fatalf("stop input %q status = %s: %s", chars, result.Status, result.Output)
}
if result.Meta["session_id"] != "" {
t.Fatalf("stop input %q should not leave session running, meta=%#v output=%q", chars, result.Meta, result.Output)
// Re-send the stop until the session reports it has exited, rather than
// waiting a fixed slice and asserting on whatever is true by then.
//
// A single 1s yield used to be enough because the session ran under
// cmd.exe. Under PowerShell the tree is slower to stand up and tear down,
// and the FIRST stop in a fresh process pays the interpreter's cold start,
// so it routinely returns with termination already under way but not
// finished. The old assertion read that as "the stop did not work" and
// failed, intermittently, on whichever of the two inputs happened to run
// first. CI never saw it because the runners have PowerShell 7 while a
// stock Windows box only has 5.1.
//
// Repeating the same input rather than polling with empty chars is
// deliberate: Terminate is a no-op against an already-dead tree, so it is
// safe to repeat, and it keeps the interrupt flag that an empty-chars poll
// would drop, which is what the interrupted assertion below needs.
var result Result
// Every poll's output is kept, not just the last one. Each Continue
// returns only what the session produced since the previous collection,
// so a marker printed just before the poll that observes the exit lands
// in the earlier result and would be invisible to a check against the
// final one alone.
var collected strings.Builder
deadline := time.Now().Add(30 * time.Second)
for {
result = writeTool.Run(context.Background(), map[string]any{
"session_id": sessionID,
"chars": chars,
"yield_time_ms": 500,
})
collected.WriteString(result.Output)
if result.Status != StatusOK {
t.Fatalf("stop input %q status = %s: %s", chars, result.Status, result.Output)
}
if result.Meta["session_id"] == "" {
break
}
if time.Now().After(deadline) {
t.Fatalf("stop input %q left the session running past the deadline, meta=%#v output=%q", chars, result.Meta, result.Output)
}
}
if result.Meta["exit_code"] == "" {
t.Fatalf("stop input %q should report exit_code, meta=%#v output=%q", chars, result.Meta, result.Output)
@@ -604,6 +668,17 @@ func TestWriteStdinStopIntentTerminatesNonTTYSession(t *testing.T) {
if result.Meta["interrupted"] != "true" {
t.Fatalf("stop input %q should report interrupted metadata, meta=%#v output=%q", chars, result.Meta, result.Output)
}
// The session has to have been killed rather than allowed to finish.
// Without this the retry loop makes the test unfalsifiable: the helper
// sleeps 5s, well inside the deadline, so a terminate that did nothing at
// all would let the sleep end on its own and every assertion above would
// still hold. The interrupted flag does not close the hole either, since
// it echoes the request rather than the outcome. The helper prints this
// line only on the way out of a completed sleep, so its absence is what
// separates the two, and it says nothing about how long the loop took.
if strings.Contains(collected.String(), "long sleep finished") {
t.Fatalf("stop input %q let the session run to completion instead of terminating it, output=%q", chars, collected.String())
}
}
}
@@ -643,7 +718,7 @@ func TestExecCommandTTYSessionAcceptsInputOnLinux(t *testing.T) {
if runtime.GOOS != "linux" {
t.Skip("pty transport is currently implemented for linux")
}
root := t.TempDir()
root := execTestRoot(t)
manager := newExecSessionManager()
execTool := NewScopedExecCommandTool(root, nil, manager)
writeTool := NewWriteStdinTool(manager)
@@ -681,7 +756,7 @@ func TestExecCommandTTYSessionAcceptsInputOnLinux(t *testing.T) {
}
func TestExecSessionSnapshotsAndStopAll(t *testing.T) {
root := t.TempDir()
root := execTestRoot(t)
manager := newExecSessionManager()
execTool := NewScopedExecCommandTool(root, nil, manager).(execCommandTool)
+259 -6
View File
@@ -79,6 +79,147 @@ func TestEditFileSucceedsWhenUnchangedAndRebaselines(t *testing.T) {
}
}
func TestApplyPatchPreservesWholeFileObservationForFollowUpEdit(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "f.txt")
if err := os.WriteFile(path, []byte("alpha\nbeta\n"), 0o644); err != nil {
t.Fatal(err)
}
registry := safetyRegistry(t, dir)
registry.Register(NewScopedApplyPatchTool(dir, nil))
tracker := NewFileTracker()
opts := grantedOpts(tracker)
if result := registry.RunWithOptions(context.Background(), "read_file", map[string]any{"path": "f.txt"}, opts); result.Status != StatusOK {
t.Fatalf("read_file failed: %s", result.Output)
}
patch := strings.Join([]string{
"diff --git a/f.txt b/f.txt",
"--- a/f.txt",
"+++ b/f.txt",
"@@ -1,2 +1,2 @@",
"-alpha",
"+gamma",
" beta",
"",
}, "\n")
if result := registry.RunWithOptions(context.Background(), "apply_patch", map[string]any{"patch": patch}, opts); result.Status != StatusOK {
if gitApplyUnavailable(result.Output) {
t.Skipf("git binary unavailable: %s", result.Output)
}
t.Fatalf("apply_patch failed: %s", result.Output)
}
result := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
"path": "f.txt", "old_string": "beta", "new_string": "delta",
}, opts)
if result.Status != StatusOK {
t.Fatalf("follow-up edit should not require a wasted re-read: %s", result.Output)
}
}
func TestApplyPatchRenameOrCopyDoesNotCreditUnreadDestination(t *testing.T) {
for _, tc := range []struct {
name string
patch string
}{
{
name: "rename",
patch: strings.Join([]string{
"diff --git a/source.txt b/destination.txt",
"similarity index 100%",
"rename from source.txt",
"rename to destination.txt",
"",
}, "\n"),
},
{
name: "copy",
patch: strings.Join([]string{
"diff --git a/source.txt b/destination.txt",
"similarity index 100%",
"copy from source.txt",
"copy to destination.txt",
"",
}, "\n"),
},
} {
t.Run(tc.name, func(t *testing.T) {
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "source.txt"), []byte("unread source\n"), 0o644); err != nil {
t.Fatal(err)
}
registry := safetyRegistry(t, dir)
registry.Register(NewScopedApplyPatchTool(dir, nil))
tracker := NewFileTracker()
opts := grantedOpts(tracker)
if result := registry.RunWithOptions(context.Background(), "apply_patch", map[string]any{"patch": tc.patch}, opts); result.Status != StatusOK {
if gitApplyUnavailable(result.Output) {
t.Skipf("git binary unavailable: %s", result.Output)
}
t.Fatalf("apply_patch failed: %s", result.Output)
}
destination, err := filepath.EvalSymlinks(filepath.Join(dir, "destination.txt"))
if err != nil {
t.Fatal(err)
}
if tracker.SeenWhole(destination) {
t.Fatal("rename/copy destination was credited as fully seen without a read")
}
result := registry.RunWithOptions(context.Background(), "write_file", map[string]any{
"path": "destination.txt", "content": "blind overwrite\n", "overwrite": true,
}, opts)
if result.Status != StatusError || !strings.Contains(result.Output, "has not been read exactly") {
t.Fatalf("blind destination overwrite was not refused: %#v", result)
}
})
}
}
func TestApplyPatchCompleteCreationCreditsSuppliedFile(t *testing.T) {
dir := t.TempDir()
registry := safetyRegistry(t, dir)
registry.Register(NewScopedApplyPatchTool(dir, nil))
tracker := NewFileTracker()
opts := grantedOpts(tracker)
patch := strings.Join([]string{
"diff --git a/new.txt b/new.txt",
"new file mode 100644",
"--- /dev/null",
"+++ b/new.txt",
"@@ -0,0 +1 @@",
"+supplied content",
"",
}, "\n")
if result := registry.RunWithOptions(context.Background(), "apply_patch", map[string]any{"patch": patch}, opts); result.Status != StatusOK {
if gitApplyUnavailable(result.Output) {
t.Skipf("git binary unavailable: %s", result.Output)
}
t.Fatalf("apply_patch failed: %s", result.Output)
}
created, err := filepath.EvalSymlinks(filepath.Join(dir, "new.txt"))
if err != nil {
t.Fatal(err)
}
if !tracker.SeenWhole(created) {
t.Fatal("complete /dev/null creation was not credited as supplied content")
}
result := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
"path": "new.txt", "old_string": "supplied", "new_string": "updated",
}, opts)
if result.Status != StatusOK {
t.Fatalf("follow-up edit of supplied file failed: %s", result.Output)
}
updated, err := os.ReadFile(filepath.Join(dir, "new.txt"))
if err != nil {
t.Fatal(err)
}
if got, want := strings.ReplaceAll(string(updated), "\r\n", "\n"), "updated content\n"; got != want {
t.Fatalf("follow-up edit content = %q, want %q", got, want)
}
}
func TestWriteFileOverwriteRefusedWhenChangedOnDisk(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "f.txt")
@@ -103,7 +244,7 @@ func TestWriteFileOverwriteRefusedWhenChangedOnDisk(t *testing.T) {
}
}
func TestWriteFileOverwriteAllowedForUntrackedFile(t *testing.T) {
func TestWriteFileOverwriteRequiresExactRead(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "f.txt")
if err := os.WriteFile(path, []byte("v1\n"), 0o644); err != nil {
@@ -112,15 +253,127 @@ func TestWriteFileOverwriteAllowedForUntrackedFile(t *testing.T) {
registry := safetyRegistry(t, dir)
tracker := NewFileTracker()
// Never read through a tool → no baseline → overwrite is the model's call.
// Never read through a tool: an overwrite would discard unseen content.
res := registry.RunWithOptions(context.Background(), "write_file", map[string]any{
"path": "f.txt", "content": "v3\n", "overwrite": true,
}, grantedOpts(tracker))
if res.Status != StatusOK {
t.Fatalf("untracked overwrite should be allowed, got %q", res.Output)
if res.Status != StatusError || !strings.Contains(res.Output, "has not been read exactly") {
t.Fatalf("unseen overwrite should be refused, got %q", res.Output)
}
if got, _ := os.ReadFile(path); string(got) != "v3\n" {
t.Fatalf("file = %q, want v3", got)
if got, _ := os.ReadFile(path); string(got) != "v1\n" {
t.Fatalf("file = %q, want original content", got)
}
}
func TestEditFileAllowsSeenRangeAndRejectsUnseenRange(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "f.txt")
if err := os.WriteFile(path, []byte("alpha\nbeta\ngamma\n"), 0o644); err != nil {
t.Fatal(err)
}
registry := safetyRegistry(t, dir)
tracker := NewFileTracker()
opts := grantedOpts(tracker)
if res := registry.RunWithOptions(context.Background(), "read_file", map[string]any{
"path": "f.txt", "start_line": 2, "end_line": 2,
}, opts); res.Status != StatusOK {
t.Fatalf("partial read failed: %s", res.Output)
}
if res := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
"path": "f.txt", "old_string": "gamma", "new_string": "delta",
}, opts); res.Status != StatusError || !strings.Contains(res.Output, "has not been read exactly") {
t.Fatalf("unseen edit should be refused, got %q", res.Output)
}
if res := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
"path": "f.txt", "old_string": "beta", "new_string": "bravo",
}, opts); res.Status != StatusOK {
t.Fatalf("seen-range edit should succeed, got %q", res.Output)
}
if content, err := os.ReadFile(path); err != nil || !strings.Contains(string(content), "bravo") {
t.Fatalf("seen-range edit did not update the file: content=%q err=%v", content, err)
}
}
func TestEditFileReplaceAllRecordsEveryShiftedReplacementSpan(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "f.txt")
if err := os.WriteFile(path, []byte("before\nneedle\nmiddle\nneedle\nafter\n"), 0o644); err != nil {
t.Fatal(err)
}
registry := safetyRegistry(t, dir)
tracker := NewFileTracker()
opts := grantedOpts(tracker)
if res := registry.RunWithOptions(context.Background(), "read_file", map[string]any{
"path": "f.txt", "start_line": 2, "end_line": 4,
}, opts); res.Status != StatusOK {
t.Fatalf("partial read failed: %q", res.Output)
}
if res := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
"path": "f.txt",
"old_string": "needle",
"new_string": "left\nright",
"replace_all": true,
}, opts); res.Status != StatusOK {
t.Fatalf("replace_all failed: %q", res.Output)
}
// This follow-up targets both generated "right" lines. It succeeds only if
// replace_all recorded the shifted second replacement as seen.
if res := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
"path": "f.txt",
"old_string": "right",
"new_string": "done",
"replace_all": true,
}, opts); res.Status != StatusOK {
t.Fatalf("follow-up edit of every replacement should succeed: %q", res.Output)
}
if content, err := os.ReadFile(path); err != nil || strings.Count(string(content), "done") != 2 {
t.Fatalf("unexpected follow-up content: content=%q err=%v", content, err)
}
}
func TestWriteFileCanOverwriteNewlyCreatedEmptyFile(t *testing.T) {
dir := t.TempDir()
registry := safetyRegistry(t, dir)
tracker := NewFileTracker()
opts := grantedOpts(tracker)
if res := registry.RunWithOptions(context.Background(), "write_file", map[string]any{
"path": "empty.txt", "content": "",
}, opts); res.Status != StatusOK {
t.Fatalf("empty create failed: %q", res.Output)
}
if res := registry.RunWithOptions(context.Background(), "write_file", map[string]any{
"path": "empty.txt", "content": "now populated\n", "overwrite": true,
}, opts); res.Status != StatusOK {
t.Fatalf("empty overwrite should succeed: %q", res.Output)
}
if content, err := os.ReadFile(filepath.Join(dir, "empty.txt")); err != nil || string(content) != "now populated\n" {
t.Fatalf("unexpected overwritten content: content=%q err=%v", content, err)
}
}
func TestMinifiedReadDoesNotAuthorizeExactEdit(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "f.go")
if err := os.WriteFile(path, []byte("package p\n\nvar value = 1\n"), 0o644); err != nil {
t.Fatal(err)
}
registry := NewRegistry()
registry.Register(NewScopedReadMinifiedFileTool(dir, nil))
registry.Register(NewScopedEditFileTool(dir, nil))
tracker := NewFileTracker()
opts := grantedOpts(tracker)
if res := registry.RunWithOptions(context.Background(), "read_minified_file", map[string]any{"path": "f.go"}, opts); res.Status != StatusOK {
t.Fatalf("minified read failed: %s", res.Output)
}
res := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
"path": "f.go", "old_string": "var value = 1", "new_string": "var value = 2",
}, opts)
if res.Status != StatusError || !strings.Contains(res.Output, "has not been read exactly") {
t.Fatalf("minified view must not authorize exact edit, got %q", res.Output)
}
}
+227 -1
View File
@@ -5,6 +5,7 @@ import (
"encoding/hex"
"errors"
"os"
"sort"
"sync"
"time"
)
@@ -34,6 +35,7 @@ type FileVersion struct {
type FileTracker struct {
mu sync.Mutex
versions map[string]FileVersion
seen map[string]fileObservation
// created preserves the order in which brand-new files (paths that did not
// exist before a write tool created them) were first written this session.
// A map would lose that order and complicates de-duplication, so a slice
@@ -43,7 +45,36 @@ type FileTracker struct {
}
func NewFileTracker() *FileTracker {
return &FileTracker{versions: make(map[string]FileVersion)}
return &FileTracker{
versions: make(map[string]FileVersion),
seen: make(map[string]fileObservation),
}
}
type lineRange struct {
start int
end int
}
type fileObservation struct {
whole bool
// total is the file's line count as of the last recorded read, so coverage
// can be judged from the ranges alone. Without it a file read in pieces
// could never be known to have been seen in full.
total int
ranges []lineRange
totalBytes int
byteRanges []lineRange // zero-based, half-open byte intervals
}
type pendingFileObservation struct {
path string
output string
hash string
start int
end int
total int
byteMode bool
}
// Record stores the version of absPath given its content and optional stat info.
@@ -63,10 +94,200 @@ func (tracker *FileTracker) RecordHash(absPath string, hash string, info os.File
version.MTime = info.ModTime()
}
tracker.mu.Lock()
if previous, ok := tracker.versions[absPath]; !ok || previous.Hash != hash {
delete(tracker.seen, absPath)
}
tracker.versions[absPath] = version
tracker.mu.Unlock()
}
// RecordSeenRange records exact, non-elided source lines returned to the model.
// Ranges accumulate while the content hash remains unchanged.
func (tracker *FileTracker) RecordSeenRange(absPath string, start, end, total int) {
if tracker == nil || start < 1 {
return
}
tracker.mu.Lock()
defer tracker.mu.Unlock()
observation := tracker.seen[absPath]
if total == 0 {
observation.whole = true
observation.ranges = nil
tracker.seen[absPath] = observation
return
}
if end < start {
return
}
// A new line count means the file changed shape since the last read, so
// previously recorded ranges no longer describe it.
if observation.total != 0 && observation.total != total {
observation = fileObservation{}
}
observation.total = total
if start == 1 && end >= total {
observation.whole = true
observation.ranges = nil
tracker.seen[absPath] = observation
return
}
observation.ranges = append(observation.ranges, lineRange{start: start, end: end})
tracker.seen[absPath] = observation
}
// RecordSeenBytes records an exact, non-elided byte interval returned to the
// model. Byte intervals make oversized single-line files recoverable without
// crediting the model for the omitted middle of a truncated line read.
func (tracker *FileTracker) RecordSeenBytes(absPath string, start, end, total int) {
if tracker == nil || start < 0 || end < start || total < 0 || end > total {
return
}
tracker.mu.Lock()
defer tracker.mu.Unlock()
observation := tracker.seen[absPath]
if total == 0 {
observation.whole = true
observation.byteRanges = nil
tracker.seen[absPath] = observation
return
}
if observation.totalBytes != 0 && observation.totalBytes != total {
observation = fileObservation{}
}
observation.totalBytes = total
if start == 0 && end >= total {
observation.whole = true
observation.byteRanges = nil
tracker.seen[absPath] = observation
return
}
observation.byteRanges = append(observation.byteRanges, lineRange{start: start, end: end})
tracker.seen[absPath] = observation
}
// coversFully reports whether ranges together cover every line in [start, end].
//
// Ranges are merged rather than scanned line by line: a caller asking about a
// large file would otherwise walk every line against every recorded range, and
// the answer is the same.
func coversFully(ranges []lineRange, start int, end int) bool {
if start > end {
return true
}
if len(ranges) == 0 {
return false
}
sorted := make([]lineRange, len(ranges))
copy(sorted, ranges)
sort.Slice(sorted, func(left int, right int) bool {
if sorted[left].start != sorted[right].start {
return sorted[left].start < sorted[right].start
}
return sorted[left].end < sorted[right].end
})
next := start
for _, seen := range sorted {
if seen.start > next {
return false
}
if seen.end >= next {
next = seen.end + 1
}
if next > end {
return true
}
}
return next > end
}
// coversBytes reports whether half-open byte ranges cover [start, end).
func coversBytes(ranges []lineRange, start, end int) bool {
if start >= end {
return true
}
if len(ranges) == 0 {
return false
}
sorted := make([]lineRange, len(ranges))
copy(sorted, ranges)
sort.Slice(sorted, func(left, right int) bool {
if sorted[left].start != sorted[right].start {
return sorted[left].start < sorted[right].start
}
return sorted[left].end < sorted[right].end
})
next := start
for _, seen := range sorted {
if seen.start > next {
return false
}
if seen.end > next {
next = seen.end
}
if next >= end {
return true
}
}
return false
}
// SeenRange reports whether every line in the requested range was returned
// exactly to the model for the currently tracked version.
func (tracker *FileTracker) SeenRange(absPath string, start, end int) bool {
if tracker == nil {
return true
}
tracker.mu.Lock()
defer tracker.mu.Unlock()
observation, ok := tracker.seen[absPath]
if !ok {
return false
}
if observation.whole {
return true
}
return coversFully(observation.ranges, start, end)
}
// SeenBytes reports whether every byte in [start, end) was returned exactly to
// the model for the currently tracked version.
func (tracker *FileTracker) SeenBytes(absPath string, start, end int) bool {
if tracker == nil {
return true
}
tracker.mu.Lock()
defer tracker.mu.Unlock()
observation, ok := tracker.seen[absPath]
if !ok {
return false
}
return observation.whole || coversBytes(observation.byteRanges, start, end)
}
func (tracker *FileTracker) SeenWhole(absPath string) bool {
if tracker == nil {
return true
}
tracker.mu.Lock()
defer tracker.mu.Unlock()
observation, ok := tracker.seen[absPath]
if !ok {
return false
}
if observation.whole {
return true
}
// Derived from the ranges rather than tracked as its own flag. The flag was
// only ever set by a SINGLE read covering the file, so a file read in two
// halves stayed "not seen whole" forever even though SeenRange agreed every
// line had been seen — and write_file, which gates on this, then refused the
// overwrite with advice to read the file that could not change the answer.
if observation.total > 0 && coversFully(observation.ranges, 1, observation.total) {
return true
}
return observation.totalBytes > 0 && coversBytes(observation.byteRanges, 0, observation.totalBytes)
}
// Version returns the recorded version for absPath and whether one exists.
func (tracker *FileTracker) Version(absPath string) (FileVersion, bool) {
if tracker == nil {
@@ -119,6 +340,7 @@ func (tracker *FileTracker) Forget(absPath string) {
}
tracker.mu.Lock()
delete(tracker.versions, absPath)
delete(tracker.seen, absPath)
tracker.mu.Unlock()
}
@@ -153,3 +375,7 @@ func fileConflictMessage(relativePath string) string {
return "Error writing " + relativePath + ": " + ErrFileChangedOnDisk.Error() +
" (it may have been edited outside Zero). Re-read it with read_file, then re-apply your change so you do not overwrite the newer content."
}
func fileUnseenMessage(relativePath string) string {
return "Error writing " + relativePath + ": the intended change depends on content that has not been read exactly in this session. Read the affected lines with read_file, or use byte_offset/byte_limit for an oversized single line, then retry the edit."
}
@@ -0,0 +1,96 @@
package tools
import (
"path/filepath"
"testing"
)
// SeenWhole used to be a flag set only by a SINGLE read that covered the file,
// while SeenRange judged coverage from the accumulated ranges. The two then
// disagreed about the same reads, and write_file — which gates on SeenWhole —
// refused overwrites it should have allowed.
func TestSeenWholeIsDerivedFromAccumulatedRanges(t *testing.T) {
path := filepath.Join(t.TempDir(), "f.go")
for name, testCase := range map[string]struct {
reads [][3]int // start, end, total
want bool
}{
"one read covering the file": {
reads: [][3]int{{1, 100, 100}},
want: true,
},
"two adjacent reads covering it together": {
reads: [][3]int{{1, 50, 100}, {51, 100, 100}},
want: true,
},
"overlapping reads covering it together": {
reads: [][3]int{{1, 60, 100}, {40, 100, 100}},
want: true,
},
"out-of-order reads covering it together": {
reads: [][3]int{{51, 100, 100}, {1, 50, 100}},
want: true,
},
"a gap in the middle": {
reads: [][3]int{{1, 40, 100}, {60, 100, 100}},
want: false,
},
"missing the first line": {
reads: [][3]int{{2, 100, 100}},
want: false,
},
"missing the last line": {
reads: [][3]int{{1, 99, 100}},
want: false,
},
"nothing read at all": {
reads: nil,
want: false,
},
} {
t.Run(name, func(t *testing.T) {
tracker := NewFileTracker()
for _, read := range testCase.reads {
tracker.RecordSeenRange(path, read[0], read[1], read[2])
}
if got := tracker.SeenWhole(path); got != testCase.want {
t.Errorf("SeenWhole = %v, want %v after reads %v", got, testCase.want, testCase.reads)
}
})
}
}
// SeenWhole and SeenRange must not disagree about the same reads.
func TestSeenWholeAgreesWithSeenRange(t *testing.T) {
path := filepath.Join(t.TempDir(), "f.go")
tracker := NewFileTracker()
tracker.RecordSeenRange(path, 1, 50, 100)
tracker.RecordSeenRange(path, 51, 100, 100)
if !tracker.SeenRange(path, 1, 100) {
t.Fatal("SeenRange says the whole file was not seen")
}
if !tracker.SeenWhole(path) {
t.Error("SeenRange covers every line but SeenWhole disagrees")
}
}
// A file that changed length since the recorded reads must not stay "seen":
// the old ranges describe a shape the file no longer has.
func TestSeenWholeResetsWhenTheFileChangesLength(t *testing.T) {
path := filepath.Join(t.TempDir(), "f.go")
tracker := NewFileTracker()
tracker.RecordSeenRange(path, 1, 100, 100)
if !tracker.SeenWhole(path) {
t.Fatal("precondition: the whole file was read")
}
// The file grew; only the first 100 lines of 200 have been seen.
tracker.RecordSeenRange(path, 1, 100, 200)
if tracker.SeenWhole(path) {
t.Error("file grew to 200 lines but is still reported as seen whole")
}
if !tracker.SeenRange(path, 1, 100) {
t.Error("the lines that were read should still count as seen")
}
}
@@ -0,0 +1,263 @@
package tools
import (
"context"
"fmt"
"os"
"path/filepath"
"strconv"
"strings"
"testing"
"unicode/utf8"
)
// A file past the read byte budget cannot be read whole in one call: the read is
// truncated and RecordSeenRange is skipped, which is correct — the model did not
// see those lines. But SeenWhole was a flag only a single whole-file read could
// set, so no sequence of scoped reads could ever clear it either, and write_file
// refused the overwrite forever while telling the model to read the file again.
func TestLargeFileBecomesWritableAfterScopedReadsCoverIt(t *testing.T) {
dir := t.TempDir()
// read_file resolves the workspace root through EvalSymlinks before it
// records anything, so the tracker is keyed on the resolved path. On macOS
// t.TempDir() sits under /var, a symlink to /private/var, and asserting
// against the raw path would look up a key that never existed — passing on
// Windows and failing there.
resolvedDir, err := filepath.EvalSymlinks(dir)
if err != nil {
t.Fatalf("resolve temp dir: %v", err)
}
path := filepath.Join(resolvedDir, "big.go")
const lines = 2400
var builder strings.Builder
for index := 0; index < lines; index++ {
fmt.Fprintf(&builder, "line %06d %s\n", index, strings.Repeat("x", 60))
}
if writeErr := os.WriteFile(path, []byte(builder.String()), 0o600); writeErr != nil {
t.Fatalf("seed file: %v", writeErr)
}
tracker := NewFileTracker()
options := RunOptions{FileTracker: tracker}
ctx := context.Background()
read := NewReadFileTool(dir).(interface {
RunWithOptions(context.Context, map[string]any, RunOptions) Result
})
write := NewScopedWriteFileTool(dir, nil).(interface {
RunWithOptions(context.Context, map[string]any, RunOptions) Result
})
// One unbounded read is byte-truncated, so nothing is recorded.
full := read.RunWithOptions(ctx, map[string]any{"path": "big.go"}, options)
if full.Meta["truncation_reason"] != "byte_budget" {
t.Fatalf("precondition: expected the file to exceed the read byte budget, got reason %q", full.Meta["truncation_reason"])
}
if tracker.SeenWhole(path) {
t.Fatal("a truncated read must not credit the model with the whole file")
}
// Scoped reads that together cover it must.
for _, span := range [][2]int{{1, 1200}, {1201, 2400}} {
result := read.RunWithOptions(ctx, map[string]any{
"path": "big.go", "start_line": span[0], "end_line": span[1],
}, options)
if result.Status != StatusOK {
t.Fatalf("scoped read %v: %s", span, result.Output)
}
if result.Meta["truncation_reason"] == "byte_budget" {
t.Fatalf("scoped read %v was itself truncated; pick smaller spans", span)
}
}
if !tracker.SeenWhole(path) {
t.Fatal("scoped reads covered every line but the file is still not seen whole")
}
overwrite := write.RunWithOptions(ctx, map[string]any{
"path": "big.go", "content": "replaced\n", "overwrite": true,
}, options)
if overwrite.Status != StatusOK {
t.Fatalf("overwrite refused after the file was fully read in scoped reads: %s", overwrite.Output)
}
}
func TestSingleLongLineBecomesWritableAfterExactByteReads(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "bundle.min.js")
content := strings.Repeat("x", readOutputBudgetBytes*3)
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
t.Fatalf("seed file: %v", err)
}
resolvedPath, err := filepath.EvalSymlinks(path)
if err != nil {
t.Fatalf("resolve file: %v", err)
}
tracker := NewFileTracker()
options := RunOptions{FileTracker: tracker}
ctx := context.Background()
read := NewReadFileTool(dir).(interface {
RunWithOptions(context.Context, map[string]any, RunOptions) Result
})
write := NewScopedWriteFileTool(dir, nil).(interface {
RunWithOptions(context.Context, map[string]any, RunOptions) Result
})
full := read.RunWithOptions(ctx, map[string]any{"path": "bundle.min.js"}, options)
if full.Meta["truncation_reason"] != "byte_budget" {
t.Fatalf("precondition: expected byte truncation, got %#v", full.Meta)
}
if tracker.SeenWhole(resolvedPath) {
t.Fatal("a truncated line read must not credit the whole file")
}
const chunkSize = 64 * 1024
for offset := 0; offset < len(content); offset += chunkSize {
result := read.RunWithOptions(ctx, map[string]any{
"path": "bundle.min.js",
"byte_offset": offset,
"byte_limit": chunkSize,
}, options)
if result.Status != StatusOK || result.Truncated {
t.Fatalf("byte read at %d was not exact: status=%s truncated=%v output=%q", offset, result.Status, result.Truncated, result.Output)
}
}
if !tracker.SeenWhole(resolvedPath) {
t.Fatal("exact byte reads covered the file but it is not seen whole")
}
overwrite := write.RunWithOptions(ctx, map[string]any{
"path": "bundle.min.js", "content": "replaced\n", "overwrite": true,
}, options)
if overwrite.Status != StatusOK {
t.Fatalf("overwrite refused after exact byte reads: %s", overwrite.Output)
}
}
func TestExactByteReadsAdvanceAlongUTF8Boundaries(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "unicode.min.js")
content := strings.Repeat("😀", 40_000)
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
t.Fatal(err)
}
resolvedPath, err := filepath.EvalSymlinks(path)
if err != nil {
t.Fatal(err)
}
tracker := NewFileTracker()
options := RunOptions{FileTracker: tracker}
read := NewReadFileTool(dir).(interface {
RunWithOptions(context.Context, map[string]any, RunOptions) Result
})
offset := 0
for offset < len(content) {
result := read.RunWithOptions(context.Background(), map[string]any{
"path": "unicode.min.js", "byte_offset": offset, "byte_limit": 65_535,
}, options)
if result.Status != StatusOK || result.Truncated || !utf8.ValidString(result.Output) {
t.Fatalf("invalid exact UTF-8 read at %d: %#v", offset, result)
}
next, parseErr := strconv.Atoi(result.Meta["next_byte_offset"])
if parseErr != nil || next <= offset {
t.Fatalf("invalid continuation offset %q after %d", result.Meta["next_byte_offset"], offset)
}
offset = next
}
if !tracker.SeenWhole(resolvedPath) {
t.Fatal("UTF-8-safe byte reads covered the file but it is not seen whole")
}
}
func TestRegistryTruncationDoesNotCreditUndeliveredByteRanges(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "unicode.min.js")
content := strings.Repeat("界", 84_000)
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
t.Fatal(err)
}
resolvedPath, err := filepath.EvalSymlinks(path)
if err != nil {
t.Fatal(err)
}
registry := safetyRegistry(t, dir)
tracker := NewFileTracker()
result := registry.RunWithOptions(context.Background(), "read_file", map[string]any{
"path": "unicode.min.js", "byte_offset": 0, "byte_limit": readFileByteChunkMax,
}, grantedOpts(tracker))
if !result.Truncated {
t.Fatalf("precondition: multibyte chunk was not truncated at the registry boundary: %#v", result.Meta)
}
next, err := strconv.Atoi(result.Meta["next_byte_offset"])
if err != nil || next <= 0 {
t.Fatalf("invalid continuation offset %q", result.Meta["next_byte_offset"])
}
if tracker.SeenBytes(resolvedPath, 0, next) || tracker.SeenWhole(resolvedPath) {
t.Fatal("registry-truncated bytes were credited as delivered to the model")
}
}
func TestRegistryTruncationDoesNotCreditWholeLineRead(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "unicode.go")
content := strings.Repeat("界界界界界界界界\n", 4_000)
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
t.Fatal(err)
}
resolvedPath, err := filepath.EvalSymlinks(path)
if err != nil {
t.Fatal(err)
}
registry := safetyRegistry(t, dir)
tracker := NewFileTracker()
result := registry.RunWithOptions(context.Background(), "read_file", map[string]any{
"path": "unicode.go",
}, grantedOpts(tracker))
if !result.Truncated {
t.Fatalf("precondition: multibyte line read was not truncated at the registry boundary: %#v", result.Meta)
}
if tracker.SeenWhole(resolvedPath) {
t.Fatal("registry-truncated line output credited the whole file")
}
}
func TestPartialByteEditDoesNotCreditAnEntireLongLine(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "bundle.min.js")
content := "UNIQUE_TARGET" + strings.Repeat("x", readOutputBudgetBytes*2)
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
t.Fatal(err)
}
resolvedPath, err := filepath.EvalSymlinks(path)
if err != nil {
t.Fatal(err)
}
registry := safetyRegistry(t, dir)
tracker := NewFileTracker()
opts := grantedOpts(tracker)
read := registry.RunWithOptions(context.Background(), "read_file", map[string]any{
"path": "bundle.min.js", "byte_offset": 0, "byte_limit": 64,
}, opts)
if read.Status != StatusOK || read.Truncated {
t.Fatalf("exact byte read failed: %#v", read)
}
edit := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
"path": "bundle.min.js", "old_string": "UNIQUE_TARGET", "new_string": "UPDATED_TARGET",
}, opts)
if edit.Status != StatusOK {
t.Fatalf("edit of exactly observed bytes failed: %s", edit.Output)
}
if tracker.SeenWhole(resolvedPath) {
t.Fatal("editing observed bytes credited the unread remainder of the long line")
}
overwrite := registry.RunWithOptions(context.Background(), "write_file", map[string]any{
"path": "bundle.min.js", "content": "replacement", "overwrite": true,
}, opts)
if overwrite.Status != StatusError {
t.Fatalf("whole overwrite should remain blocked after a partial byte read, got %#v", overwrite)
}
}
+25 -3
View File
@@ -61,15 +61,37 @@ func TestFormatOnWriteFormatsAndKeepsTrackerConsistent(t *testing.T) {
t.Fatalf("expected gofmt-formatted content, got %q", onDisk)
}
// The tracker must have been re-baselined to the POST-format content: a
// follow-up edit must not trip the external-modification conflict guard.
// The tracker is re-baselined to the post-format bytes, but those bytes were
// not returned exactly to the model, so an edit must require an exact read.
edit := NewScopedEditFileTool(dir, nil).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
"path": "a.go",
"old_string": "func A() {",
"new_string": "func B() {",
}, RunOptions{FileTracker: tracker})
if edit.Status != StatusError || !strings.Contains(edit.Output, "has not been read exactly") {
t.Fatalf("formatter-modified content must require an exact read: %q", edit.Output)
}
read := NewScopedReadFileTool(dir, nil).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
"path": "a.go",
}, RunOptions{FileTracker: tracker})
if read.Status != StatusOK {
t.Fatalf("exact read failed: %q", read.Output)
}
edit = NewScopedEditFileTool(dir, nil).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
"path": "a.go",
"old_string": "func A() {",
"new_string": "func B( ) {",
}, RunOptions{FileTracker: tracker})
if edit.Status != StatusOK {
t.Fatalf("follow-up edit must not conflict after formatting: %q", edit.Output)
t.Fatalf("follow-up edit after exact read failed: %q", edit.Output)
}
edit = NewScopedEditFileTool(dir, nil).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
"path": "a.go",
"old_string": "func B() {",
"new_string": "func C() {",
}, RunOptions{FileTracker: tracker})
if edit.Status != StatusError || !strings.Contains(edit.Output, "has not been read exactly") {
t.Fatalf("formatter-modified edit must require an exact read: %q", edit.Output)
}
}
+9
View File
@@ -45,6 +45,7 @@ func newBrowserInstallTool(options localcontrol.BrowserOptions) Tool {
return browserInstallTool{
baseTool: baseTool{
name: "browser_install",
deferred: true,
description: "Install the local browser runtime used by browser automation.",
parameters: Schema{
Type: "object",
@@ -86,6 +87,7 @@ func newBrowserLaunchToolWithLauncher(options localcontrol.BrowserOptions, launc
return browserLaunchTool{
baseTool: baseTool{
name: "browser_launch",
deferred: true,
description: "Launch a supported local Chromium/Electron app with Chrome DevTools enabled and attach browser automation to it. Use this for installed Electron apps such as Discord instead of launching GUI apps through shell commands.",
parameters: Schema{
Type: "object",
@@ -165,6 +167,7 @@ func newBrowserConnectTool(options localcontrol.BrowserOptions) Tool {
return browserConnectTool{
baseTool: baseTool{
name: "browser_connect",
deferred: true,
description: "Attach local browser automation to an existing Chrome/Electron DevTools endpoint. For Electron apps, first launch the app with --remote-debugging-port=<port>, then connect to that port instead of using desktop control.",
parameters: Schema{
Type: "object",
@@ -205,6 +208,7 @@ func newBrowserOpenTool(options localcontrol.BrowserOptions) Tool {
return browserOpenTool{
baseTool: baseTool{
name: "browser_open",
deferred: true,
description: "Open a URL in the local browser automation session.",
parameters: Schema{
Type: "object",
@@ -245,6 +249,7 @@ func newBrowserSnapshotTool(options localcontrol.BrowserOptions) Tool {
return browserSnapshotTool{
baseTool: baseTool{
name: "browser_snapshot",
deferred: true,
description: "Return an accessibility snapshot from the current local browser automation session.",
parameters: Schema{
Type: "object",
@@ -288,6 +293,7 @@ func newBrowserClickTool(options localcontrol.BrowserOptions) Tool {
return browserClickTool{
baseTool: baseTool{
name: "browser_click",
deferred: true,
description: "Click a ref from browser_snapshot in the current local browser automation session.",
parameters: Schema{
Type: "object",
@@ -328,6 +334,7 @@ func newBrowserTypeTool(options localcontrol.BrowserOptions) Tool {
return browserTypeTool{
baseTool: baseTool{
name: "browser_type",
deferred: true,
description: "Type text into a ref from browser_snapshot in the current local browser automation session.",
parameters: Schema{
Type: "object",
@@ -369,6 +376,7 @@ func newBrowserPressTool(options localcontrol.BrowserOptions) Tool {
return browserPressTool{
baseTool: baseTool{
name: "browser_press",
deferred: true,
description: "Press a keyboard key in the current local browser automation session.",
parameters: Schema{
Type: "object",
@@ -454,6 +462,7 @@ func newBrowserActionTool(options localcontrol.BrowserOptions) Tool {
return browserActionTool{
baseTool: baseTool{
name: "browser_action",
deferred: true,
description: "Run an allowed action against the current local browser automation session.",
parameters: Schema{
Type: "object",
+1
View File
@@ -40,6 +40,7 @@ func newCaptureArtifactTool(options LocalControlArtifactOptions) Tool {
return captureArtifactTool{
baseTool: baseTool{
name: "capture_artifact",
deferred: true,
description: "Capture local browser, desktop, or terminal state into the configured artifact directory.",
parameters: Schema{
Type: "object",
+3
View File
@@ -27,6 +27,7 @@ func newDesktopWindowsTool(options localcontrol.DesktopOptions) Tool {
return desktopWindowsTool{
baseTool: baseTool{
name: "desktop_windows",
deferred: true,
description: "List native desktop windows through the local desktop automation helper.",
parameters: Schema{
Type: "object",
@@ -70,6 +71,7 @@ func newDesktopSnapshotTool(options localcontrol.DesktopOptions) Tool {
return desktopSnapshotTool{
baseTool: baseTool{
name: "desktop_snapshot",
deferred: true,
description: "Read a native desktop window accessibility snapshot through the local desktop automation helper.",
parameters: Schema{
Type: "object",
@@ -117,6 +119,7 @@ func newDesktopActionTool(options localcontrol.DesktopOptions) Tool {
return desktopActionTool{
baseTool: baseTool{
name: "desktop_action",
deferred: true,
description: "Run an allowed native desktop action through the local desktop automation helper.",
parameters: Schema{
Type: "object",
+1
View File
@@ -23,6 +23,7 @@ func newTerminalSessionTool(options localcontrol.TerminalOptions) Tool {
return terminalSessionTool{
baseTool: baseTool{
name: "terminal_session",
deferred: true,
description: "Launch or control a local virtual terminal session through the local terminal automation helper.",
parameters: Schema{
Type: "object",
+2 -1
View File
@@ -32,7 +32,8 @@ func NewScopedLSPNavigateTool(workspaceRoot string, scope PathScope) Tool {
root := normalizeWorkspaceRoot(workspaceRoot)
return lspNavigateTool{
baseTool: baseTool{
name: "lsp_navigate",
name: "lsp_navigate",
deferred: true,
description: "Navigate code semantically via the language server: jump to a symbol's " +
"definition, find all references, find implementations of an interface/method, or " +
"search workspace symbols by name. Use this instead of grep when you need precise " +
+4 -1
View File
@@ -57,6 +57,9 @@ func applySelfManagedOutputBudget(tool Tool, toolName string, args map[string]an
category := resolveOutputCategory(tool, toolName, args)
budgeted := budgetSemanticOutput(result.Output, category, budget)
if result.Truncated && budgeted.spillPath == "" {
budgeted.spillPath = result.Meta["spill_path"]
}
if result.Truncated && !budgeted.truncated {
budgeted.truncated = true
budgeted.reason = result.Meta["truncation_reason"]
@@ -64,7 +67,7 @@ func applySelfManagedOutputBudget(tool Tool, toolName string, args map[string]an
budgeted.reason = "upstream_tool_budget"
}
}
if budgeted.truncated {
if budgeted.truncated && budgeted.spillPath == "" {
budgeted = attachExistingSpill(toolName, result.Output, budget, budgeted)
}
result.Output = budgeted.text
+121 -9
View File
@@ -10,8 +10,11 @@ import (
"os"
"strconv"
"strings"
"unicode/utf8"
)
const readFileByteChunkMax = 64 * 1024
type readFileTool struct {
baseTool
workspaceRoot string
@@ -28,14 +31,16 @@ func NewScopedReadFileTool(workspaceRoot string, scope PathScope) Tool {
return readFileTool{
baseTool: baseTool{
name: "read_file",
description: "Read a file with optional 1-based inclusive line range and max line cap.",
description: "Read a file by 1-based inclusive line range, or by exact byte chunks for oversized single-line files.",
parameters: Schema{
Type: "object",
Properties: map[string]PropertySchema{
"path": {Type: "string", Description: "Path of the file to read."},
"start_line": {Type: "integer", Description: "1-based inclusive line number to start reading from.", Minimum: intPtr(1)},
"end_line": {Type: "integer", Description: "1-based inclusive line number to stop reading at.", Minimum: intPtr(1)},
"max_lines": {Type: "integer", Description: "Maximum number of lines to return.", Minimum: intPtr(1)},
"path": {Type: "string", Description: "Path of the file to read."},
"start_line": {Type: "integer", Description: "1-based inclusive line number to start reading from.", Minimum: intPtr(1)},
"end_line": {Type: "integer", Description: "1-based inclusive line number to stop reading at.", Minimum: intPtr(1)},
"max_lines": {Type: "integer", Description: "Maximum number of lines to return.", Minimum: intPtr(1)},
"byte_offset": {Type: "integer", Description: "Zero-based byte offset for an exact chunk read. Use with byte_limit for oversized single-line files; cannot be combined with line ranges.", Minimum: intPtr(0)},
"byte_limit": {Type: "integer", Description: "Maximum source bytes returned in exact byte mode. Defaults to 65536.", Default: readFileByteChunkMax, Minimum: intPtr(1), Maximum: intPtr(readFileByteChunkMax)},
},
Required: []string{"path"},
AdditionalProperties: false,
@@ -76,6 +81,20 @@ func (tool readFileTool) run(args map[string]any, options RunOptions, directBudg
if err != nil {
return errorResult("Error: Invalid arguments for read_file: " + err.Error())
}
_, hasByteOffset := args["byte_offset"]
_, hasByteLimit := args["byte_limit"]
byteMode := hasByteOffset || hasByteLimit
byteOffset, err := intArg(args, "byte_offset", 0, 0, 0)
if err != nil {
return errorResult("Error: Invalid arguments for read_file: " + err.Error())
}
byteLimit, err := intArg(args, "byte_limit", readFileByteChunkMax, 1, readFileByteChunkMax)
if err != nil {
return errorResult("Error: Invalid arguments for read_file: " + err.Error())
}
if byteMode && (args["start_line"] != nil || args["end_line"] != nil || args["max_lines"] != nil) {
return errorResult("Error: Invalid arguments for read_file: byte_offset/byte_limit cannot be combined with line ranges")
}
absolutePath, relativePath, err := resolveScopedReadPath(tool.workspaceRoot, tool.scope, requestedPath)
if err != nil {
@@ -91,12 +110,61 @@ func (tool readFileTool) run(args map[string]any, options RunOptions, directBudg
// Stat is best-effort: a missing FileInfo only drops the diagnostic size/mtime,
// not the authoritative content hash.
options.FileTracker.RecordHash(absolutePath, stats.hash, stats.info)
if byteMode {
result, seenStart, seenEnd := renderReadFileBytes(absolutePath, relativePath, stats.bytes, byteOffset, byteLimit)
if result.Status == StatusOK && !result.Truncated && seenEnd > seenStart {
if options.deferFileObservation {
result.pendingFileObservation = &pendingFileObservation{
path: absolutePath, output: result.Output, hash: stats.hash,
start: seenStart, end: seenEnd, total: stats.bytes, byteMode: true,
}
} else {
options.FileTracker.RecordSeenBytes(absolutePath, seenStart, seenEnd, stats.bytes)
result.Meta["file_version"] = stats.hash
result.Meta["seen_bytes"] = fmt.Sprintf("%d-%d", seenStart, seenEnd)
}
}
return result
}
maxBytes := 0
if directBudget {
maxBytes = readOutputBudgetBytes
}
return renderReadFileRange(absolutePath, relativePath, stats.lines, startLine, endLine, maxLines, maxBytes)
result := renderReadFileRange(absolutePath, relativePath, stats.lines, startLine, endLine, maxLines, maxBytes)
if result.Status == StatusOK && result.Meta["truncation_reason"] != "byte_budget" {
seenStart, seenEnd := renderedReadRange(stats.lines, startLine, endLine, maxLines)
if options.deferFileObservation {
result.pendingFileObservation = &pendingFileObservation{
path: absolutePath, output: result.Output, hash: stats.hash,
start: seenStart, end: seenEnd, total: stats.lines,
}
} else {
options.FileTracker.RecordSeenRange(absolutePath, seenStart, seenEnd, stats.lines)
if result.Meta == nil {
result.Meta = map[string]string{}
}
result.Meta["file_version"] = stats.hash
result.Meta["seen_lines"] = fmt.Sprintf("%d-%d", seenStart, seenEnd)
}
}
return result
}
func renderedReadRange(total, start, end, maxLines int) (int, int) {
if start > total {
return start, start
}
if end == 0 || end > total {
end = total
}
if end < start {
end = start
}
if maxLines > 0 && end-start+1 > maxLines {
end = start + maxLines - 1
}
return start, end
}
func renderReadFileRange(absolutePath string, relativePath string, total int, startLine int, endLine int, maxLines int, maxBytes int) Result {
@@ -140,9 +208,9 @@ func renderReadFileRange(absolutePath string, relativePath string, total int, st
// Tool.Run calls preserve the legacy prefix-only byte budget.
var budgetedOutput *outputBudgetBuilder
if maxBytes > 0 {
budgetedOutput = newOutputBudgetBuilder(maxBytes, "use start_line/end_line or max_lines to continue with a smaller range")
budgetedOutput = newOutputBudgetBuilder(maxBytes, "use start_line/end_line or max_lines for normal files; use byte_offset/byte_limit for an oversized single line")
} else {
budgetedOutput = newHeadTailOutputBudgetBuilder(readOutputBudgetBytes, "use start_line/end_line or max_lines to continue with a smaller range")
budgetedOutput = newHeadTailOutputBudgetBuilder(readOutputBudgetBytes, "use start_line/end_line or max_lines for normal files; use byte_offset/byte_limit for an oversized single line")
}
budgetedOutput.WriteString(header)
budgetedOutput.WriteString("\n")
@@ -186,6 +254,7 @@ func renderReadFileRange(absolutePath string, relativePath string, total int, st
type readFileStats struct {
lines int
bytes int
hash string
info os.FileInfo
}
@@ -200,6 +269,7 @@ func scanReadFileStats(path string) (readFileStats, error) {
hasher := sha256.New()
reader := bufio.NewReader(file)
lines := 0
bytes := 0
for {
raw, _, err := readRawLine(reader)
if err == io.EOF {
@@ -211,13 +281,55 @@ func scanReadFileStats(path string) (readFileStats, error) {
if _, err := hasher.Write(raw); err != nil {
return readFileStats{}, err
}
bytes += len(raw)
lines++
}
if lines == 0 {
lines = 1
}
info, _ := file.Stat()
return readFileStats{lines: lines, hash: hex.EncodeToString(hasher.Sum(nil)), info: info}, nil
return readFileStats{lines: lines, bytes: bytes, hash: hex.EncodeToString(hasher.Sum(nil)), info: info}, nil
}
func renderReadFileBytes(path, relativePath string, total, requestedStart, limit int) (Result, int, int) {
if requestedStart >= total {
return okResult(fmt.Sprintf("File: %s\n(byte_offset %d is past the end of the file, which has %d bytes)", relativePath, requestedStart, total)), 0, 0
}
file, err := os.Open(path)
if err != nil {
return errorResult("Error reading file " + relativePath + ": " + err.Error()), 0, 0
}
defer file.Close()
want := min(limit, total-requestedStart)
data := make([]byte, want)
read, err := file.ReadAt(data, int64(requestedStart))
if err != nil && err != io.EOF {
return errorResult("Error reading file " + relativePath + ": " + err.Error()), 0, 0
}
data = data[:read]
start := requestedStart
for len(data) > 0 && data[0]&0xc0 == 0x80 {
data = data[1:]
start++
}
for trim := 0; !utf8.Valid(data) && trim < utf8.UTFMax && len(data) > 0; trim++ {
data = data[:len(data)-1]
}
if !utf8.Valid(data) {
return errorResult("Error reading file " + relativePath + ": exact byte mode requires UTF-8 text"), 0, 0
}
end := start + len(data)
if end <= start {
return errorResult("Error reading file " + relativePath + ": byte_limit is too small to return the next UTF-8 character"), 0, 0
}
header := fmt.Sprintf("File: %s (bytes %d-%d of %d)", relativePath, start, end-1, total)
output := header + "\n\n" + string(data)
if end < total {
output += fmt.Sprintf("\n\n[continue with byte_offset=%d]", end)
}
return Result{Status: StatusOK, Output: output, Meta: map[string]string{"next_byte_offset": strconv.Itoa(end)}}, start, end
}
func appendReadFileRange(output *outputBudgetBuilder, path string, startLine int, selectedLines int, width int) error {
+2 -1
View File
@@ -29,6 +29,7 @@ func NewScopedReadMinifiedFileTool(workspaceRoot string, scope PathScope) Tool {
baseTool: baseTool{
name: "read_minified_file",
description: "Read a source file in a dense, token-cheap form: comments and redundant whitespace removed, no line numbers. Use it to scan or understand code for far fewer tokens than read_file. For exact text, comments, line numbers, or before editing, use read_file instead.",
deferred: true,
parameters: Schema{
Type: "object",
Properties: map[string]PropertySchema{
@@ -112,7 +113,7 @@ func (tool readMinifiedFileTool) run(args map[string]any, options RunOptions, di
meta["estimated_tokens_saved"] = strconv.Itoa(savedTokens)
toolResult := Result{Status: StatusOK, Output: output, Meta: meta}
if directBudget {
return applyLegacyByteBudgetToResult(toolResult, readOutputBudgetBytes, "use read_file with start_line/end_line or max_lines for a smaller exact range")
return applyLegacyByteBudgetToResult(toolResult, readOutputBudgetBytes, "use read_file with start_line/end_line or max_lines for normal files, or byte_offset/byte_limit for an oversized single line")
}
return toolResult
}
+38
View File
@@ -2,7 +2,9 @@ package tools
import (
"context"
"fmt"
"sort"
"strings"
"github.com/Gitlawb/zero/internal/redaction"
"github.com/Gitlawb/zero/internal/sandbox"
@@ -29,6 +31,10 @@ type RunOptions struct {
// disk outside Zero since it was last read. nil disables the feature entirely
// (the read/write tools behave exactly as before).
FileTracker *FileTracker
// DeferFileObservationCommit lets a caller with an additional output layer
// commit exact-read credit after that final layer has run.
DeferFileObservationCommit bool
deferFileObservation bool
// EnabledTools / DisabledTools carry the run's operator tool filters so a
// filter-aware tool (tool_search) never discloses or loads an operator-hidden
// tool. They use the same allow/deny semantics as the agent's filter gate:
@@ -130,17 +136,23 @@ func (registry *Registry) RunWithOptions(ctx context.Context, name string, args
// args/paths) are redacted at the boundary just like tool output. The output
// ceiling runs after the scrub so the transcript and the spill file agree on
// what was hidden.
commitFileObservation := !options.DeferFileObservationCommit
options.deferFileObservation = true
selfManagedOutput := false
var tool Tool
var ok bool
defer func() {
result = scrubResultSecrets(result)
result = reduceCommandOutput(name, args, result)
if selfManagedOutput {
result = applySelfManagedOutputBudget(tool, name, args, result)
} else {
result = applyRegistryOutputBudget(tool, name, args, result)
result = enforceOutputCeiling(name, result)
}
if commitFileObservation {
result = registry.CommitFileObservation(result, options.FileTracker)
}
}()
tool, ok = registry.Get(name)
@@ -220,6 +232,32 @@ func (registry *Registry) RunWithOptions(ctx context.Context, name string, args
return res
}
// CommitFileObservation grants exact-read credit only when the tool's complete
// output remains intact at the final boundary that will be sent to the model.
func (registry *Registry) CommitFileObservation(result Result, tracker *FileTracker) Result {
observation := result.pendingFileObservation
result.pendingFileObservation = nil
if observation == nil || tracker == nil || result.Status != StatusOK || result.Truncated ||
!strings.HasPrefix(result.Output, observation.output) {
return result
}
if observation.byteMode {
tracker.RecordSeenBytes(observation.path, observation.start, observation.end, observation.total)
} else {
tracker.RecordSeenRange(observation.path, observation.start, observation.end, observation.total)
}
if result.Meta == nil {
result.Meta = map[string]string{}
}
result.Meta["file_version"] = observation.hash
if observation.byteMode {
result.Meta["seen_bytes"] = fmt.Sprintf("%d-%d", observation.start, observation.end)
} else {
result.Meta["seen_lines"] = fmt.Sprintf("%d-%d", observation.start, observation.end)
}
return result
}
func effectiveToolPermission(tool Tool, args map[string]any) Permission {
if permissioner, ok := tool.(ArgsPermissioner); ok {
return permissioner.PermissionForArgs(args)
+41
View File
@@ -515,3 +515,44 @@ func TestIsDeferralEligibleDecouplesFromDeferred(t *testing.T) {
t.Fatal("IsDeferralEligible(optedOut) = true, want false")
}
}
func TestOptionalBuiltinsUseDeferredDiscovery(t *testing.T) {
wantDeferred := map[string]bool{
"bash": true,
"browser_action": true,
"browser_click": true,
"browser_connect": true,
"browser_install": true,
"browser_launch": true,
"browser_open": true,
"browser_press": true,
"browser_snapshot": true,
"browser_type": true,
"capture_artifact": true,
"desktop_action": true,
"desktop_snapshot": true,
"desktop_windows": true,
"lsp_navigate": true,
"read_minified_file": true,
"terminal_session": true,
"web_fetch": true,
"web_search": true,
}
for _, tool := range BuiltinCatalog(t.TempDir()) {
if _, listed := wantDeferred[tool.Name()]; listed {
if !IsDeferred(tool) {
t.Errorf("%s should be deferred", tool.Name())
}
delete(wantDeferred, tool.Name())
} else if tool.Name() == "read_file" || tool.Name() == "grep" || tool.Name() == ExecCommandToolName || tool.Name() == "edit_file" {
if IsDeferred(tool) {
t.Errorf("essential tool %s must stay eager", tool.Name())
}
} else if IsDeferred(tool) {
t.Errorf("unexpected deferred builtin %s", tool.Name())
}
}
for missing := range wantDeferred {
t.Errorf("deferred builtin %s missing from catalog", missing)
}
}
+95
View File
@@ -0,0 +1,95 @@
package tools
import (
"fmt"
"strconv"
"strings"
)
const commandReducerMinPassingLines = 12
// reduceCommandOutput removes repetitive success-only lines from recognized
// simple commands. Unsupported or compound shell input passes through exactly,
// and every changed result retains the pre-reduction output as an artifact.
func reduceCommandOutput(toolName string, args map[string]any, result Result) Result {
if toolName != ExecCommandToolName && toolName != "bash" {
return result
}
command, err := commandTextForReducer(toolName, args)
if err != nil || !isSimpleGoTestCommand(command) {
return result
}
reduced, omitted := reduceGoTestPassingPackages(result.Output)
if omitted < commandReducerMinPassingLines || len(reduced) >= len(result.Output) {
return result
}
originalBytes := len(result.Output)
originalTokens := estimateOutputTokens(result.Output)
spillPath := spillTruncatedOutput(toolName, result.Output)
if spillPath == "" {
return result
}
notice := fmt.Sprintf("[zero] compacted %d passing package lines; exact output saved to %s", omitted, spillPath)
result.Output = strings.TrimRight(reduced, "\n") + "\n" + notice
result.Truncated = true
if result.Meta == nil {
result.Meta = map[string]string{}
}
result.Meta["command_output_reduced"] = "true"
result.Meta["command_output_omitted_lines"] = strconv.Itoa(omitted)
result.Meta["command_output_original_bytes"] = strconv.Itoa(originalBytes)
result.Meta["command_output_original_tokens"] = strconv.Itoa(originalTokens)
result.Meta["command_output_retained_tokens"] = strconv.Itoa(estimateOutputTokens(result.Output))
result.Meta["spill_path"] = spillPath
result.Meta["truncation_reason"] = "command_aware_reduction"
return result
}
func commandTextForReducer(toolName string, args map[string]any) (string, error) {
if toolName == "bash" {
return bashCommandArg(args)
}
return execCommandArg(args)
}
func isSimpleGoTestCommand(command string) bool {
command = strings.TrimSpace(command)
if command == "" || strings.ContainsAny(command, "|;&<>\n\r`") || strings.Contains(command, "$(") {
return false
}
fields := strings.Fields(command)
return len(fields) >= 2 && fields[0] == "go" && fields[1] == "test"
}
func reduceGoTestPassingPackages(output string) (string, int) {
lines := strings.Split(output, "\n")
kept := make([]string, 0, len(lines))
omitted := 0
for _, line := range lines {
packageSuccess := strings.HasPrefix(line, "ok \t") || strings.HasPrefix(line, "? \t") ||
strings.HasPrefix(line, "ok ") || strings.HasPrefix(line, "? ")
if packageSuccess && !strings.Contains(line, "coverage:") && !strings.Contains(line, "[no test files]") {
omitted++
continue
}
kept = append(kept, line)
}
if omitted == 0 {
return output, 0
}
insertAt := len(kept)
for index, line := range kept {
if strings.HasPrefix(line, "exit_code:") || strings.HasPrefix(line, "session_id:") {
insertAt = index
break
}
}
summary := fmt.Sprintf("[zero] %d passing package lines omitted", omitted)
kept = append(kept, "")
copy(kept[insertAt+1:], kept[insertAt:])
kept[insertAt] = summary
return strings.Join(kept, "\n"), omitted
}
+133
View File
@@ -0,0 +1,133 @@
package tools
import (
"fmt"
"os"
"strings"
"testing"
)
func TestReduceCommandOutputCompactsPassingGoPackagesAndKeepsRawArtifact(t *testing.T) {
var lines []string
lines = append(lines, "output:")
for index := 0; index < commandReducerMinPassingLines+3; index++ {
lines = append(lines, fmt.Sprintf("ok \texample.test/pkg%d\t0.01s", index))
}
lines = append(lines, "exit_code: 0")
original := strings.Join(lines, "\n")
result := reduceCommandOutput(ExecCommandToolName, map[string]any{"cmd": "go test ./..."}, Result{
Status: StatusOK,
Output: original,
})
if result.Output == original || !strings.Contains(result.Output, "passing package lines omitted") {
t.Fatalf("expected compact result, got %q", result.Output)
}
if !strings.Contains(result.Output, "exit_code: 0") {
t.Fatalf("exit status was lost: %q", result.Output)
}
spillPath := result.Meta["spill_path"]
spilled, err := os.ReadFile(spillPath)
if err != nil {
t.Fatalf("read raw artifact: %v", err)
}
if string(spilled) != original {
t.Fatalf("raw artifact differs from original output")
}
budgeted := applySelfManagedOutputBudget(
NewExecCommandTool(t.TempDir(), newExecSessionManager()),
ExecCommandToolName,
map[string]any{"cmd": "go test ./..."},
result,
)
if budgeted.Meta["spill_path"] != spillPath {
t.Fatalf("budget boundary replaced the exact raw artifact: got %q want %q", budgeted.Meta["spill_path"], spillPath)
}
spilled, err = os.ReadFile(budgeted.Meta["spill_path"])
if err != nil || string(spilled) != original {
t.Fatalf("budgeted raw artifact was not preserved: err=%v", err)
}
}
func TestReduceCommandOutputFailsOpenForCompoundAndFailedOutput(t *testing.T) {
failure := "output:\n--- FAIL: TestThing\npanic: broken\nFAIL\texample.test/pkg\nexit_code: 1"
for _, command := range []string{"go test ./... | tee test.log", "go test ./... && echo done"} {
result := reduceCommandOutput(ExecCommandToolName, map[string]any{"cmd": command}, Result{Status: StatusError, Output: failure})
if result.Output != failure || result.Meta != nil {
t.Fatalf("compound command %q should pass through, got %#v", command, result)
}
}
result := reduceCommandOutput(ExecCommandToolName, map[string]any{"cmd": "go test ./..."}, Result{Status: StatusError, Output: failure})
if result.Output != failure {
t.Fatalf("failure evidence should pass through when there is no repetitive success bulk: %q", result.Output)
}
}
func TestReduceGoTestPassingPackagesPreservesCoverageResults(t *testing.T) {
output := strings.Join([]string{
"ok \texample.test/a\t0.01s\tcoverage: 82.4% of statements",
"? \texample.test/b\t[no test files]",
"ok \texample.test/c\t0.01s",
}, "\n")
reduced, omitted := reduceGoTestPassingPackages(output)
if !strings.Contains(reduced, "coverage: 82.4% of statements") {
t.Fatalf("coverage result was removed: %q", reduced)
}
if !strings.Contains(reduced, "[no test files]") {
t.Fatalf("no-test-files result was removed: %q", reduced)
}
if omitted != 1 {
t.Fatalf("omitted=%d, want only the plain passing line omitted", omitted)
}
}
func TestReduceGoTestPassingPackagesPreservesIndentedTestLogs(t *testing.T) {
output := strings.Join([]string{
" ok user-provided test log",
" ? another test log",
"ok \texample.test/pkg\t0.01s",
}, "\n")
reduced, omitted := reduceGoTestPassingPackages(output)
if !strings.Contains(reduced, " ok user-provided test log") ||
!strings.Contains(reduced, " ? another test log") {
t.Fatalf("indented test log was removed: %q", reduced)
}
if omitted != 1 {
t.Fatalf("omitted=%d, want only the package result omitted", omitted)
}
}
func TestSelfManagedBudgetPreservesRawSpillWhenReducedOutputAlsoTruncates(t *testing.T) {
raw := "raw output\n" + strings.Repeat("ok \texample.test/raw\t0.01s\n", 1000)
rawSpill := spillTruncatedOutput(ExecCommandToolName, raw)
if rawSpill == "" {
t.Fatal("failed to create raw spill")
}
reduced := strings.Repeat("reduced but still large\n", 1000)
result := applySelfManagedOutputBudget(
NewExecCommandTool(t.TempDir(), newExecSessionManager()),
ExecCommandToolName,
map[string]any{"cmd": "go test ./...", "max_output_tokens": 100},
Result{
Status: StatusOK,
Output: reduced,
Truncated: true,
Meta: map[string]string{
"spill_path": rawSpill,
"truncation_reason": "command_aware_reduction",
},
},
)
if result.Meta["spill_path"] != rawSpill {
t.Fatalf("budget boundary replaced raw spill: got %q want %q", result.Meta["spill_path"], rawSpill)
}
spilled, err := os.ReadFile(result.Meta["spill_path"])
if err != nil || string(spilled) != raw {
t.Fatalf("raw spill was not preserved: err=%v", err)
}
}
+9
View File
@@ -115,6 +115,9 @@ type Result struct {
ChangeSummaries []execution.Change
// Display carries a short, structured summary for the TUI / stream.
Display Display
// pendingFileObservation is proposed by read_file and committed only after
// the final model-visible output boundary confirms the exact content survived.
pendingFileObservation *pendingFileObservation
}
// Display carries a short, structured summary of a tool result for the TUI/stream.
@@ -159,6 +162,7 @@ type baseTool struct {
parameters Schema
safety Safety
capabilities ToolCapabilities // zero value = EffectUnknown, not thread-safe
deferred bool
}
func (tool baseTool) Name() string {
@@ -177,6 +181,11 @@ func (tool baseTool) Safety() Safety {
return tool.safety
}
// Deferred reports whether this built-in is discoverable on demand.
func (tool baseTool) Deferred() bool {
return tool.deferred
}
func okResult(output string) Result {
return Result{Status: StatusOK, Output: output}
}
+1
View File
@@ -107,6 +107,7 @@ func newWebFetchToolWithClientAndResolver(client *http.Client, resolver webFetch
baseTool: baseTool{
name: "web_fetch",
description: "Fetch text from a public remote HTTP or HTTPS URL after network permission is granted. Do not use for localhost, private network URLs, or local dev servers; use bash with curl for those.",
deferred: true,
parameters: Schema{
Type: "object",
Properties: map[string]PropertySchema{
+1
View File
@@ -58,6 +58,7 @@ func newWebSearchToolWithBackend(backend searchBackend) Tool {
baseTool: baseTool{
name: "web_search",
description: "Search the web for a query and return ranked results (title, URL, snippet). Complements web_fetch, which retrieves a single known URL.",
deferred: true,
parameters: Schema{
Type: "object",
Properties: map[string]PropertySchema{
+7
View File
@@ -75,6 +75,9 @@ func (tool writeFileTool) RunWithOptions(ctx context.Context, args map[string]an
// stale view. Only read current bytes when there is a baseline to compare,
// so a first-touch create/overwrite stays a single write with no extra read.
if existed {
if options.FileTracker != nil && !options.FileTracker.SeenWhole(absolutePath) {
return errorResult(fileUnseenMessage(relativePath))
}
if _, tracked := options.FileTracker.Version(absolutePath); tracked {
// Fail CLOSED: if the tracked file can't be re-read to verify it, refuse
// the overwrite rather than clobbering a file whose current state is
@@ -107,6 +110,7 @@ func (tool writeFileTool) RunWithOptions(ctx context.Context, args map[string]an
if err := os.WriteFile(absolutePath, []byte(content), 0o644); err != nil {
return errorResult("Error writing file " + relativePath + ": " + err.Error())
}
modelKnownContent := content
// Optional format-on-write (ZERO_FORMAT_ON_WRITE). Must run BEFORE the
// FileTracker baseline: recording pre-format content would make the very
// next edit look like an external modification and trip the conflict guard.
@@ -115,6 +119,9 @@ func (tool writeFileTool) RunWithOptions(ctx context.Context, args map[string]an
// session compares against what is now on disk.
newInfo, _ := os.Stat(absolutePath)
options.FileTracker.Record(absolutePath, []byte(content), newInfo)
if content == modelKnownContent {
options.FileTracker.RecordSeenRange(absolutePath, 1, lineCount(content), lineCount(content))
}
if !existed {
options.FileTracker.RecordCreated(absolutePath)
}
+1 -1
View File
@@ -206,7 +206,7 @@ func (m model) leaveBTW() (model, tea.Cmd) {
func btwCommandUnavailable(command parsedCommand) bool {
arg := strings.ToLower(strings.TrimSpace(command.text))
switch command.kind {
case commandNew, commandResume, commandRetitle, commandSpec, commandLoop, commandGoal,
case commandNew, commandResume, commandRename, commandSpec, commandLoop, commandGoal,
commandRewind, commandCompact, commandSTTModel, commandMCP:
return true
case commandModel:
+15 -2
View File
@@ -23,12 +23,16 @@ type clipboardImageMsg struct {
err error
}
// readClipboardImage is the OS clipboard image reader. A var so tests can
// substitute a stub instead of depending on the developer's real clipboard.
var readClipboardImage = imageinput.ReadClipboardImage
// readClipboardImageCmd reads the OS clipboard for image content off the
// Update goroutine. Returns a clipboardImageMsg with the bytes, or nil (no
// command) if there is no image — the caller treats nil as a silent no-op.
func readClipboardImageCmd() tea.Cmd {
return func() tea.Msg {
data, mediaType, err := imageinput.ReadClipboardImage()
data, mediaType, err := readClipboardImage()
if err != nil {
return clipboardImageMsg{err: err}
}
@@ -54,7 +58,8 @@ func pasteFromClipboardCmd() tea.Cmd {
// right-click paste (clipboardReadMsg) so a bracketed paste and a right-click
// paste behave identically. Surfaces with no editable text field (a permission/
// spec prompt, the MCP manager, an open picker, the detailed transcript) swallow
// the paste; empty content is a no-op.
// the paste. The session rename editor accepts text only; empty content is a
// no-op there rather than an image probe.
func (m model) routePaste(content string) (tea.Model, tea.Cmd) {
// A paste is a deliberate action, same as a keypress or click — it means
// the user moved on to something else, so it disarms a stale Esc
@@ -71,6 +76,14 @@ func (m model) routePaste(content string) (tea.Model, tea.Cmd) {
// toggle instead of going solid.
m.lastCharTime = m.now()
m.composerCursorVisible = true
if m.renamePrompt != nil {
if content == "" {
return m, nil
}
var cmd tea.Cmd
m.input, cmd = m.input.Update(tea.PasteMsg{Content: sanitizeComposerInput(content)})
return m, cmd
}
if content == "" {
// Empty text clipboard — the user may have pasted a screenshot.
// Probe the OS clipboard for image content asynchronously.
+5 -5
View File
@@ -27,7 +27,7 @@ const (
commandPlan
commandSearch
commandResume
commandRetitle
commandRename
commandSpec
commandInit
commandCompact
@@ -226,11 +226,11 @@ var commandDefinitions = []commandDefinition{
kind: commandResume,
},
{
name: "/retitle",
usage: "/retitle",
name: "/rename",
usage: "/rename [title]",
group: commandGroupSession,
description: "Generate concise titles for resumable sessions.",
kind: commandRetitle,
description: "Rename the current session (no arg opens an editor).",
kind: commandRename,
},
{
name: "/spec",
+1 -1
View File
@@ -281,7 +281,7 @@ func (m model) composerPositionAtMouse(msg tea.MouseMsg) (int, bool) {
func (m model) composerMouseSelectionBlocked() bool {
return m.transcriptDetailed || m.setup.visible || m.providerWizard != nil || m.mcpAddWizard != nil ||
m.mcpManager != nil || m.picker != nil || m.suggestionsActive()
m.mcpManager != nil || m.picker != nil || m.renamePrompt != nil || m.suggestionsActive()
}
func (m model) composerPositionAtVisualCell(x int, y int, width int) (int, bool) {
+112 -5
View File
@@ -6,6 +6,8 @@ import (
"testing"
tea "charm.land/bubbletea/v2"
"github.com/Gitlawb/zero/internal/imageinput"
)
func TestComposerInsertNewlineAtCursor(t *testing.T) {
@@ -129,22 +131,77 @@ func TestSanitizeComposerPastePreservesNewlines(t *testing.T) {
}
}
func TestCtrlVDoesNotPasteIntoComposer(t *testing.T) {
// Ctrl+V must never insert TEXT: bracketed paste is the only text path, so
// letting Bubble's own Ctrl+V binding also read the clipboard would double-paste.
// It may return a command, but that command is the image probe below, never a
// text paste, so the composer contents are unchanged either way.
func TestCtrlVDoesNotPasteTextIntoComposer(t *testing.T) {
m := newModel(context.Background(), Options{})
m.input.SetValue("hello")
m.input.CursorEnd()
updated, cmd := m.Update(testKeyCtrl('v'))
updated, _ := m.Update(testKeyCtrl('v'))
next := updated.(model)
if cmd != nil {
t.Fatal("ctrl+v should not run the textinput clipboard paste command")
}
if got := next.composerValue(); got != "hello" {
t.Fatalf("composer value after ctrl+v = %q, want unchanged", got)
}
}
// Ctrl+V probes the clipboard for an image (#534). A clipboard holding a
// screenshot produces no bracketed paste at all, so without this the empty-paste
// image probe in routePaste never runs and Ctrl+V does nothing, which is exactly
// what users reported: right-click paste attached the image, Ctrl+V did not.
func TestCtrlVProbesClipboardForImage(t *testing.T) {
// Stubbed BEFORE Update, so the command Update returns is the one that runs.
// Asserting only that some command came back, then running a freshly built
// probe, would pass even if the key never reached the image route at all.
original := readClipboardImage
t.Cleanup(func() { readClipboardImage = original })
readClipboardImage = func() ([]byte, string, error) {
return []byte("png-bytes"), "image/png", nil
}
// Both modifiers: macOS reports Command as ModSuper, so a handler matching
// only ModCtrl leaves Cmd+V doing nothing on the platform where screenshots
// are most often on the clipboard.
for name, key := range map[string]tea.KeyPressMsg{
"ctrl+v": testKeyCtrl('v'),
"cmd+v": testKeyPressMod('v', tea.ModSuper),
} {
t.Run(name, func(t *testing.T) {
m := newModel(context.Background(), Options{})
_, cmd := m.Update(key)
if cmd == nil {
t.Fatal("no command issued; the clipboard image probe never ran")
}
image, ok := execCmd(cmd).(clipboardImageMsg)
if !ok {
t.Fatal("the command issued was not the clipboard image probe")
}
if string(image.data) != "png-bytes" || image.mediaType != "image/png" {
t.Fatalf("unexpected image message: %+v", image)
}
})
}
}
// With no image on the clipboard the probe stays silent: Ctrl+V while copying
// text must not emit a notice or disturb the composer. Driven through Update
// rather than by calling the probe directly, so this covers the real key route
// and not just the command in isolation.
func TestClipboardImageProbeSilentWithoutImage(t *testing.T) {
original := readClipboardImage
t.Cleanup(func() { readClipboardImage = original })
readClipboardImage = func() ([]byte, string, error) { return nil, "", nil }
m := newModel(context.Background(), Options{})
_, cmd := m.Update(testKeyCtrl('v'))
if msg := execCmd(cmd); msg != nil {
t.Fatalf("probe emitted %T with no image on the clipboard, want no message", msg)
}
}
func TestPastedMultilineComposerContentRendersAsPreview(t *testing.T) {
paste := strings.Join([]string{
"Create a book library dashboard page with the Bootstrap theme.",
@@ -465,3 +522,53 @@ func TestComposerTerminalWordKeybindings(t *testing.T) {
})
}
}
// A clipboard that holds an image the host cannot extract must say so. This is
// the macOS case: neither pngpaste nor PyObjC ships with the OS, so the reader
// fails on a stock Mac, and reporting that as "no image" meant Ctrl+V did
// nothing at all with no explanation. Silence is the wrong answer to a
// condition the user can act on.
func TestClipboardImageUnreadableSurfacesToUser(t *testing.T) {
original := readClipboardImage
t.Cleanup(func() { readClipboardImage = original })
readClipboardImage = func() ([]byte, string, error) {
return nil, "", imageinput.ErrClipboardImageUnreadable
}
m := newModel(context.Background(), Options{})
_, cmd := m.Update(testKeyCtrl('v'))
if cmd == nil {
t.Fatal("no command issued; the clipboard image probe never ran")
}
msg := execCmd(cmd)
image, ok := msg.(clipboardImageMsg)
if !ok {
t.Fatalf("probe returned %T, want clipboardImageMsg carrying the failure", msg)
}
if image.err == nil {
t.Fatal("the unreadable-clipboard failure was swallowed; the user is told nothing")
}
// The message has to name the remedy, since the whole point is that the user
// can fix this by installing something.
if !strings.Contains(image.err.Error(), "pngpaste") {
t.Fatalf("error = %q, want it to name what to install", image.err.Error())
}
// And it must reach the transcript as an error row rather than stopping at
// the message, which is where the old silent no-op ended.
updated, _ := m.Update(image)
next, ok := updated.(model)
if !ok {
t.Fatalf("Update returned %T, want model", updated)
}
found := false
for _, row := range next.transcript {
if row.kind == rowError && strings.Contains(row.text, "Clipboard image read failed") {
found = true
break
}
}
if !found {
t.Fatal("the failure never reached the transcript; the user still sees nothing")
}
}
+8
View File
@@ -34,6 +34,14 @@ func keyCtrl(msg tea.KeyMsg, code rune) bool {
return keyCode(msg) == code && keyHasMod(msg, tea.ModCtrl)
}
// keySuper matches the Command key on macOS, which the rest of the TUI models as
// ModSuper (see keybindings.go). Kept separate from keyCtrl because a binding
// that means Ctrl on Linux and Windows usually means Command on macOS, and the
// two arrive as different modifiers.
func keySuper(msg tea.KeyMsg, code rune) bool {
return keyCode(msg) == code && keyHasMod(msg, tea.ModSuper)
}
func keyCtrlArrow(msg tea.KeyMsg, code rune) bool {
return keyIs(msg, code) && keyHasMod(msg, tea.ModCtrl)
}
+113 -51
View File
@@ -102,6 +102,7 @@ type model struct {
doctorInFlight bool
doctorFrame int
activeSession sessions.Metadata
pendingSessionTitle string
sessionEvents []sessions.Event
btw btwState
// btwRunIDSeq is the highest run ID issued by any completed or abandoned BTW
@@ -112,14 +113,8 @@ type model struct {
// already been attempted this process, so a finished turn re-fires the title
// generator at most once per session (even before its async result lands).
// Lazily initialized.
titledSessions map[string]bool
// retitle* drive the sequential /retitle backfill: queued session ids still
// awaiting a title, whether a backfill is running, and its progress counters.
retitleQueue []string
retitleActive bool
retitleTotal int
retitleDone int
retitleOK int
titledSessions map[string]bool
renamePrompt *sessionRenamePrompt
usageTracker *usage.Tracker
sessionCompactor SessionCompactor
prService *PrService
@@ -221,6 +216,9 @@ type model struct {
// renders the live elapsed time from it so a long or stalled turn never looks
// like a frozen terminal (for ANY provider, not just slow ones). Zero = idle.
turnStartedAt time.Time
// turnTimer is shared with the agent command so both the live status and the
// settled "worked for" duration exclude time blocked on a user permission.
turnTimer *activeTurnTimer
// lastCharTime tracks when the last non-Enter key was received, for paste detection.
lastCharTime time.Time
// lastKeyTime tracks every keypress timestamp for burst calculation.
@@ -474,8 +472,14 @@ type model struct {
// pins), this is maintained by recordRecentModel on every successful
// switch and persisted via config.SetRecentModels.
recentModels []config.RecentModelEntry
recapsEnabled bool // post-turn "※ recap:" line (config: recaps on|off)
recappedRuns map[int]bool // per-run guard so a recap fires at most once per turn
recapsEnabled bool // idle orientation note (config: recaps on|off)
recapSeq int
recapRunning bool
recapCancel context.CancelFunc
recapTimerCancel context.CancelFunc
recapIdleArmed bool
recapIdleRunID int
idleRecap string
modelPickerLoading bool
modelPickerLoadingProviderID string
modelPickerLoadError string
@@ -1069,7 +1073,7 @@ func (m *model) stopPRWatcher() {
func (m model) noBlockingModal() bool {
return m.pendingPermission == nil && m.pendingAskUser == nil && m.pendingSpecReview == nil &&
m.providerWizard == nil && m.mcpAddWizard == nil && m.mcpManager == nil && m.picker == nil &&
m.sttKeyPrompt == nil
m.sttKeyPrompt == nil && m.renamePrompt == nil
}
func (m model) quit() (tea.Model, tea.Cmd) {
@@ -1157,6 +1161,11 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.printInFlight = false
return m.drainFlushQueue()
}
var recapActivityCmd tea.Cmd
switch msg.(type) {
case tea.KeyPressMsg, tea.PasteMsg, tea.MouseMsg:
m, recapActivityCmd = m.resetIdleRecapAfterActivity()
}
next, cmd := m.updateModel(msg)
nm, ok := next.(model)
if !ok {
@@ -1165,7 +1174,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
nm = nm.syncChatScroll()
nm, mouseCmd := nm.syncMouseCapture()
nm, flushCmd := nm.settleTranscript()
return nm, batchCommands(cmd, mouseCmd, flushCmd)
return nm, batchCommands(cmd, mouseCmd, flushCmd, recapActivityCmd)
}
func batchCommands(cmds ...tea.Cmd) tea.Cmd {
@@ -1352,6 +1361,9 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
if m.sttKeyPrompt != nil {
return m.handleSTTKeyPromptKey(msg)
}
if m.renamePrompt != nil {
return m.handleSessionRenameKey(msg)
}
m.transcriptSelection = transcriptSelectionState{}
m.composerSelection = composerSelectionState{}
m.clearMouseSelection()
@@ -1703,6 +1715,27 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
m.input.SetWidth(maxInt(20, m.chatColumnWidth()-14))
return m, nil
}
case keyCtrl(msg, 'v'), keySuper(msg, 'v'):
// Ctrl+V probes the clipboard for an IMAGE only. Text pasting stays
// exclusively on the terminal's bracketed-paste path (Bubble's own
// Ctrl+V binding is disabled in newModel for exactly that reason), so
// this cannot double-insert text. It is needed because a clipboard
// holding a screenshot produces no bracketed paste at all: the terminal
// has no text to send, so routePaste never runs and its empty-content
// image probe never fires. Right-click paste reached that probe only
// because it always delivers a clipboardReadMsg, empty or not.
// readClipboardImageCmd yields no message when the clipboard holds no
// image, so Ctrl+V with text on the clipboard stays a no-op here and is
// handled by the bracketed paste exactly as before.
//
// Cmd+V is matched too, since macOS reports Command as ModSuper rather
// than ModCtrl. That only helps on terminals that deliver the key to the
// application: one that handles Cmd+V itself pastes the clipboard TEXT
// and sends no key event, so an image-only clipboard still produces
// nothing for this to react to.
if m.noBlockingModal() {
return m, readClipboardImageCmd()
}
case keyCtrl(msg, 'f'):
if m.picker != nil && m.picker.kind == pickerModel {
if m.modelPickerIsLoading() {
@@ -2444,15 +2477,15 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
var titleCmd, recapCmd tea.Cmd
if msg.err == nil {
m, titleCmd = m.maybeAutoTitleActiveSession()
// Post-turn recap (gated on the recaps preference): one short sentence
// summarizing the turn's final answer, shown as a "※ recap:" footnote.
// Arm an idle recap only after a successful answer. The timer performs
// no provider work unless the session stays untouched long enough.
var finalAnswer string
for _, row := range msg.rows {
if row.kind == rowAssistant && row.final {
finalAnswer = row.text
}
}
m, recapCmd = m.maybeRecapTurn(msg.runID, finalAnswer)
m, recapCmd = m.maybeScheduleIdleRecap(msg.runID, finalAnswer)
}
// End-of-turn git sweep: catch file mutations the tool stream couldn't
// report (bash scaffolding, subagent edits) so the FILES sidebar is
@@ -2485,6 +2518,8 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
return next, tea.Batch(pendingClearCmd, titleCmd, recapCmd, sweepCmd, queuedCmd, loopTickCmd, goalCmd)
case sessionTitleGeneratedMsg:
return m.handleSessionTitleGenerated(msg)
case recapIdleMsg:
return m.handleRecapIdle(msg)
case recapGeneratedMsg:
return m.handleRecapGenerated(msg)
case compactResultMsg:
@@ -2927,6 +2962,12 @@ func (m model) pinnedTitleBar(width int) string {
func (m model) footerView(width int) string {
var footer strings.Builder
if m.renamePrompt != nil {
footer.WriteString(m.sessionRenamePromptView(width))
footer.WriteString("\n")
footer.WriteString(m.statusLine(width))
return footer.String()
}
// While an ask-user questionnaire is active it REPLACES the composer box (the
// text box becomes the questionnaire): render the tabbed prompt + status line and
// skip the plan panel / idle hints / composer for a focused modal.
@@ -2965,6 +3006,8 @@ func (m model) footerView(width int) string {
// so the footer height is unchanged.
if copyStatus := strings.TrimSpace(m.copyStatus); copyStatus != "" {
footer.WriteString(rightAlignedLine(zeroTheme.ink.Render(copyStatus), width))
} else if recap := strings.TrimSpace(m.idleRecap); recap != "" {
footer.WriteString(fitStyledLine(" "+zeroTheme.faint.Render("※ "+recap), width))
} else if left, right := m.composerIdleHint(), m.jumpToBottomHint(); left != "" || right != "" {
footer.WriteString(fitStyledLine(joinHeaderLine(" "+left, right, width), width))
}
@@ -3015,9 +3058,18 @@ func (m model) composerIdleHint() string {
case tierNarrow:
hint = "? shortcuts"
case tierMedium:
hint = fmt.Sprintf("? shortcuts · Ctrl+X cmds · %s sidebar", sidebarKey)
parts := []string{"? shortcuts", "Ctrl+X cmds"}
if m.sidebarAvailable() {
parts = append(parts, sidebarKey+" sidebar")
}
hint = strings.Join(parts, " · ")
default:
hint = fmt.Sprintf("? shortcuts · Ctrl+X cmds · %s sidebar · %s detail · %s copy · Shift+Tab mode", sidebarKey, detailKey, mouseKey)
parts := []string{"? shortcuts", "Ctrl+X cmds"}
if m.sidebarAvailable() {
parts = append(parts, sidebarKey+" sidebar")
}
parts = append(parts, detailKey+" detail", mouseKey+" copy", "Shift+Tab mode")
hint = strings.Join(parts, " · ")
}
return zeroTheme.faint.Render(hint)
}
@@ -3481,7 +3533,7 @@ func (m model) workingStatusLine() string {
// (reasoning, waiting on the model, or running a tool).
line += zeroTheme.faint.Render(" · " + m.workingActivity())
if !m.turnStartedAt.IsZero() {
line += zeroTheme.faint.Render(" · " + formatWorkingElapsed(m.now().Sub(m.turnStartedAt)))
line += zeroTheme.faint.Render(" · " + formatWorkingElapsed(m.activeTurnElapsed(m.turnStartedAt)))
}
// Live token estimate so the working line visibly climbs as the model reasons
// and writes, instead of a static figure. Shown from the start of the turn (at
@@ -3566,7 +3618,7 @@ const quietWorkingHint = 8 * time.Second
// the ticking number was the only signal, and it looks identical whether real
// (if slow) content is coming or nothing ever will.
func (m model) quietGenerationHint() string {
if m.activeRunID == 0 {
if m.activeRunID == 0 || m.pendingPermission != nil {
return ""
}
last := m.lastStreamActivity
@@ -4115,6 +4167,10 @@ func (m model) resolvePermissionWithReason(decision permissionDecision, reason s
})
}
m.pendingPermission = nil
// Time spent at the prompt is user wait, not provider silence. Restart the
// quiet-generation clock so resuming a long-blocked run does not immediately
// claim the model has been inactive for the entire approval interval.
m.lastStreamActivity = m.now()
return m, nil
}
@@ -4499,21 +4555,11 @@ func (m model) dispatchCommand(command parsedCommand) (tea.Model, tea.Cmd) {
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: text})
}
return m, nil
case commandRetitle:
if m.pending {
m.transcript = reduceTranscript(m.transcript, transcriptAction{
kind: actionAppendError,
text: "Cannot retitle sessions while a run is active.",
})
return m, nil
case commandRename:
if title := strings.TrimSpace(command.text); title != "" {
return m.renameActiveSession(title), nil
}
text := ""
var retitleCmd tea.Cmd
m, retitleCmd, text = m.startSessionRetitle()
if text != "" {
m.transcript = reduceTranscript(m.transcript, transcriptAction{kind: actionAppendSystem, text: text})
}
return m, retitleCmd
return m.openSessionRenamePrompt(), nil
case commandSpec:
return m.handleSpecCommand(command.text)
case commandInit:
@@ -4817,6 +4863,7 @@ func (m model) launchPrompt(prompt string) (model, tea.Cmd) {
// (normal prompt + spec draft/impl) keeps these in sync — a missing
// turnStartedAt previously dropped the elapsed timer on spec-mode runs.
func (m model) beginRun(cancel context.CancelFunc) model {
m = m.cancelIdleRecap()
if m.prepareRunCompletionWarning != nil {
m.prepareRunCompletionWarning()
}
@@ -4838,6 +4885,7 @@ func (m model) beginRun(cancel context.CancelFunc) model {
// suppressed by a stale preference.
m.sidebarHidden = false
m.turnStartedAt = m.now()
m.turnTimer = newActiveTurnTimer(m.turnStartedAt)
m.lastStreamActivity = m.turnStartedAt
m.turnStreamedRunes = 0
m.spinnerTicking = true
@@ -5018,9 +5066,21 @@ func selfCorrectAutonomyForMode(mode agent.PermissionMode) string {
func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt string, images []zeroruntime.ImageBlock, runOptions tuiAgentRunOptions) tea.Cmd {
return func() tea.Msg {
started := m.now()
// firstTokenAt is stamped when the first token (reasoning or text) streams,
// so the turn can report time-to-first-token alongside total wall time.
var firstTokenAt time.Time
if m.turnTimer != nil {
m.turnTimer.start(started)
}
// firstTokenElapsed is stamped from the pause-aware turn timer when the
// first reasoning or text token streams, so TTFT and total elapsed use
// the same clock.
var firstTokenElapsed time.Duration
firstTokenSeen := false
stampFirstToken := func(at time.Time) {
if firstTokenSeen {
return
}
firstTokenSeen = true
firstTokenElapsed = m.activeTurnElapsedAt(started, at)
}
toolCalls := 0
rows := []transcriptRow{}
usageEvents := []zeroruntime.Usage{}
@@ -5136,11 +5196,10 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
onText := options.OnText
options.OnText = func(delta string) {
if firstTokenAt.IsZero() {
firstTokenAt = m.now()
}
now := m.now()
stampFirstToken(now)
if strings.TrimSpace(reasoningText) != "" {
flushReasoning(m.now())
flushReasoning(now)
}
m.sendAgentText(runID, delta)
if onText != nil {
@@ -5157,6 +5216,12 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
}
onPermissionRequest := options.OnPermissionRequest
options.OnPermissionRequest = func(ctx context.Context, request agent.PermissionRequest) (agent.PermissionDecision, error) {
if m.turnTimer != nil {
m.turnTimer.pause(m.now())
defer func() {
m.turnTimer.resume(m.now())
}()
}
if onPermissionRequest != nil {
return onPermissionRequest(ctx, request)
}
@@ -5234,8 +5299,8 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
onReasoning := options.OnReasoning
options.OnReasoning = func(delta string) {
now := m.now()
if firstTokenAt.IsZero() && strings.TrimSpace(delta) != "" {
firstTokenAt = now
if strings.TrimSpace(delta) != "" {
stampFirstToken(now)
}
if strings.TrimSpace(reasoningText) == "" && strings.TrimSpace(delta) != "" {
reasoningStarted = now
@@ -5346,6 +5411,7 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
tool: result.Name,
status: result.Status,
detail: toolResultDetail(result),
meta: result.Meta,
runID: runID,
changedFiles: result.ChangedFiles,
changeSummaries: result.ChangeSummaries,
@@ -5460,7 +5526,7 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
Type: sessions.EventError,
Payload: map[string]any{"message": err.Error()},
})
return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, err: err, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.now().Sub(started)}
return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, err: err, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.activeTurnElapsed(started)}
}
if runOptions.specDraft {
if result.StopReason != agent.StopReasonSpecReviewRequired || specReview == nil || specReview.SpecID == "" || specReview.SpecFilePath == "" {
@@ -5470,17 +5536,13 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
Type: sessions.EventError,
Payload: map[string]any{"message": err.Error()},
})
return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, err: err, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.now().Sub(started)}
return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, err: err, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.activeTurnElapsed(started)}
}
flushReasoning(m.now())
return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, specReview: specReview, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.now().Sub(started)}
return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, specReview: specReview, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: m.activeTurnElapsed(started)}
}
flushReasoning(m.now())
elapsed := m.now().Sub(started)
ttft := time.Duration(0)
if !firstTokenAt.IsZero() {
ttft = firstTokenAt.Sub(started)
}
elapsed := m.activeTurnElapsed(started)
rows = append(rows, transcriptRow{
kind: rowAssistant,
text: result.FinalAnswer,
@@ -5498,7 +5560,7 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
"content": result.FinalAnswer,
},
})
return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: elapsed, ttft: ttft}
return agentResponseMsg{runID: runID, rows: rows, usageEvents: usageEvents, usageModelID: usageModelID, sessionEvents: sessionEvents, goalAware: goalAwareRun, turnTools: toolCalls, turnElapsed: elapsed, ttft: firstTokenElapsed}
}
}
+72 -9
View File
@@ -886,8 +886,9 @@ func TestResumeCommandListsRecentSessions(t *testing.T) {
if next.picker == nil || next.picker.kind != pickerSession {
t.Fatalf("expected /resume to open the session picker, got picker=%#v", next.picker)
}
// Every row carries the session title (in the Label, after the timestamp) and
// resolves to / shows the session id (Value + Meta), for both sessions.
// Every row carries the session title after the timestamp and resolves by its
// hidden session id. Raw ids are implementation detail and must not consume
// the visible title width.
findByID := func(id string) (pickerItem, bool) {
for _, item := range next.picker.items {
if item.Value == id {
@@ -904,17 +905,52 @@ func TestResumeCommandListsRecentSessions(t *testing.T) {
if !strings.Contains(item.Label, want.title) {
t.Fatalf("picker Label %q should contain the title %q", item.Label, want.title)
}
if !strings.Contains(item.Meta, want.id) {
t.Fatalf("picker %q Meta should show the id %q, got %q", want.title, want.id, item.Meta)
if item.Meta != "" {
t.Fatalf("picker %q should not expose raw session id metadata, got %q", want.title, item.Meta)
}
}
// The picker overlay renders the titles and ids.
// The picker overlay renders clean title rows plus a position indicator.
view := viewString(next.View())
for _, want := range []string{"Resume a session", "Newer", "Older", first.SessionID, second.SessionID} {
for _, want := range []string{"Resume a session", "Newer", "Older", "1 / 2"} {
if !strings.Contains(view, want) {
t.Fatalf("session picker view missing %q:\n%s", want, view)
}
}
for _, id := range []string{first.SessionID, second.SessionID} {
if strings.Contains(view, id) {
t.Fatalf("session picker view should not expose raw id %q:\n%s", id, view)
}
}
// IDs remain searchable even though they are hidden from normal rows.
searchByID := *next.picker
searchByID.query = first.SessionID
searchByID.applyQuery()
if len(searchByID.items) != 1 || searchByID.items[0].Value != first.SessionID {
t.Fatalf("hidden session id should remain searchable, got %#v", searchByID.items)
}
// An empty result set reports an unambiguous zero position.
noResults := next
noResults.picker.query = "__missing_session__"
noResults.picker.applyQuery()
if view := viewString(noResults.View()); !strings.Contains(view, "0 / 0") {
t.Fatalf("empty session search should show 0 / 0:\n%s", view)
}
}
func TestSessionPickerLabelAlignsTitles(t *testing.T) {
today := sessionPickerLabel("20:47:50", "Today title")
older := sessionPickerLabel("Jul 24 10:47", "Older title")
todayColumn := strings.Index(today, "Today title")
olderColumn := strings.Index(older, "Older title")
if todayColumn < 0 || olderColumn < 0 {
t.Fatalf("sessionPickerLabel omitted a title: today=%q older=%q", today, older)
}
if todayColumn != olderColumn {
t.Fatalf("title columns differ: today=%d (%q), older=%d (%q)", todayColumn, today, olderColumn, older)
}
}
func TestResumePickerSelectionHydratesSession(t *testing.T) {
@@ -2501,9 +2537,36 @@ func TestComposerIdleHintAndJumpCue(t *testing.T) {
m.width, m.height = 100, 30
m.transcript = append(m.transcript, transcriptRow{kind: rowAssistant, text: "hi", final: true})
// Idle, empty composer, managed mode -> the discoverability hint shows.
if h := plainRender(t, m.composerIdleHint()); !strings.Contains(h, "shortcuts") {
t.Fatalf("expected idle hint, got %q", h)
for _, test := range []struct {
name string
width int
}{
{name: "medium", width: 80},
{name: "full", width: 100},
} {
t.Run(test.name, func(t *testing.T) {
idle := m
idle.width = test.width
// Idle, empty composer, managed mode -> the discoverability hint shows.
hint := plainRender(t, idle.composerIdleHint())
if !strings.Contains(hint, "shortcuts") {
t.Fatalf("expected idle hint, got %q", hint)
}
if strings.Contains(hint, "sidebar") {
t.Fatalf("empty sidebar should not be advertised, got %q", hint)
}
withSidebar := idle
withSidebar.plan.steps = []planStep{{content: "inspect footer", status: "in_progress"}}
if hint := plainRender(t, withSidebar.composerIdleHint()); !strings.Contains(hint, "Ctrl+B sidebar") {
t.Fatalf("available sidebar should be advertised, got %q", hint)
}
withSidebar.sidebarHidden = true
if hint := plainRender(t, withSidebar.composerIdleHint()); !strings.Contains(hint, "Ctrl+B sidebar") {
t.Fatalf("collapsed sidebar should keep its restore shortcut, got %q", hint)
}
})
}
// Hidden during a run.
busy := m
+162 -41
View File
@@ -13,21 +13,33 @@ import (
)
const (
// recapIdleDelay keeps recap generation off the critical path of ordinary
// turns. A recap is useful when returning to a quiet session, not immediately
// after reading the answer that it would otherwise repeat.
recapIdleDelay = 4 * time.Minute
// recapTimeout bounds a single recap generation so a hung provider can never
// wedge the background command.
recapTimeout = 30 * time.Second
// recapMaxAnswerChars bounds how much of the final answer feeds the recap
// prompt — the gist is enough, and it keeps the call cheap.
recapMaxAnswerChars = 1600
// recapMaxContextChars bounds the recent conversation excerpt sent to the
// provider. Goal and plan hints are added separately.
recapMaxContextChars = 4800
recapMaxRows = 8
)
const recapSystemPrompt = "You write ONE short, plain-English sentence summarizing what the assistant just did this turn, for a transcript footnote. " +
"Reply with ONLY that sentence: no preamble, no markdown, no bullet list, no trailing question."
const recapSystemPrompt = "The user is returning to an idle coding session. In under 40 words and 1-2 plain sentences, state the overall goal or current task, then the next actionable step. Reply with only the recap: no markdown, preamble, or trailing question."
// recapIdleMsg marks the end of an inactivity window. seq makes abandoned
// timers cheap: any subsequent user activity or run invalidates the message.
type recapIdleMsg struct {
runID int
seq int
}
// recapGeneratedMsg carries the outcome of a background recap generation back to
// the Update loop.
type recapGeneratedMsg struct {
runID int
seq int
recap string
err error
}
@@ -43,21 +55,20 @@ func cleanGeneratedRecap(raw string) string {
return ""
}
// generateRecap asks the provider for a one-sentence recap of the turn's final
// answer. Provider-shaped exactly like generateSessionTitle: system + a single
// user turn, no tools.
func generateRecap(ctx context.Context, provider zeroruntime.Provider, answer string) (string, error) {
// generateRecap asks the provider for a compact orientation note using bounded
// session context. It is an isolated side request with no tools.
func generateRecap(ctx context.Context, provider zeroruntime.Provider, sessionContext string) (string, error) {
if provider == nil {
return "", errors.New("no provider configured")
}
answer = strings.TrimSpace(answer)
if answer == "" {
return "", errors.New("no answer to recap")
sessionContext = strings.TrimSpace(sessionContext)
if sessionContext == "" {
return "", errors.New("no session context to recap")
}
request := zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{
{Role: zeroruntime.MessageRoleSystem, Content: recapSystemPrompt},
{Role: zeroruntime.MessageRoleUser, Content: "Assistant's final answer:\n\n" + cutRunes(answer, recapMaxAnswerChars) + "\n\nOne-sentence recap:"},
{Role: zeroruntime.MessageRoleUser, Content: sessionContext},
},
}
stream, err := provider.StreamCompletion(ctx, request)
@@ -75,51 +86,158 @@ func generateRecap(ctx context.Context, provider zeroruntime.Provider, answer st
return recap, nil
}
// generateRecapCmd builds the background command that generates a recap for the
// turn, off the Update goroutine (mirrors generateSessionTitleCmd).
func (m model) generateRecapCmd(runID int, answer string) tea.Cmd {
// recapContext builds the bounded orientation material used by the idle recap.
// Goal and plan state are authoritative; recent user/final-assistant rows fill
// in enough conversational context when neither exists.
func (m model) recapContext() string {
var sections []string
if goal := m.activeSession.Goal; goal != nil && strings.TrimSpace(goal.Objective) != "" {
sections = append(sections, "Overall goal: "+strings.TrimSpace(goal.Objective)+"\nGoal status: "+string(goal.Status))
}
if current := strings.TrimSpace(currentStepContent(m.plan.steps)); current != "" {
sections = append(sections, "Current task: "+current)
}
rows := make([]string, 0, recapMaxRows)
used := 0
for i := len(m.transcript) - 1; i >= 0 && len(rows) < recapMaxRows && used < recapMaxContextChars; i-- {
row := m.transcript[i]
role := ""
switch {
case row.kind == rowUser:
role = "User"
case row.kind == rowAssistant && row.final:
role = "Assistant"
default:
continue
}
content := strings.TrimSpace(row.text)
if content == "" {
continue
}
content = cutRunes(content, recapMaxContextChars-used)
rows = append(rows, role+": "+content)
used += len([]rune(content))
}
for left, right := 0, len(rows)-1; left < right; left, right = left+1, right-1 {
rows[left], rows[right] = rows[right], rows[left]
}
if len(rows) > 0 {
sections = append(sections, "Recent conversation:\n"+strings.Join(rows, "\n"))
}
return strings.Join(sections, "\n\n")
}
// generateRecapCmd builds the cancellable background command that generates an
// idle recap off the Update goroutine.
func (m model) generateRecapCmd(ctx context.Context, cancel context.CancelFunc, runID, seq int) tea.Cmd {
provider := m.provider
sessionContext := m.recapContext()
return func() tea.Msg {
ctx, cancel := context.WithTimeout(context.Background(), recapTimeout)
defer cancel()
recap, err := generateRecap(ctx, provider, answer)
return recapGeneratedMsg{runID: runID, recap: recap, err: err}
recap, err := generateRecap(ctx, provider, sessionContext)
return recapGeneratedMsg{runID: runID, seq: seq, recap: recap, err: err}
}
}
// maybeRecapTurn fires a one-shot recap generation after a successful turn, when
// recaps are enabled, a provider exists, and there is a final answer. It is a
// no-op otherwise, and the per-run gate prevents a double-fire for the same turn.
func (m model) maybeRecapTurn(runID int, answer string) (model, tea.Cmd) {
// maybeScheduleIdleRecap arms a one-shot inactivity window after a successful
// turn. It performs no provider work until that window actually elapses.
func (m model) maybeScheduleIdleRecap(runID int, answer string) (model, tea.Cmd) {
if !m.recapsEnabled || m.provider == nil || strings.TrimSpace(answer) == "" {
return m, nil
}
if m.recappedRuns[runID] {
return m, nil
}
if m.recappedRuns == nil {
m.recappedRuns = map[int]bool{}
}
m.recappedRuns[runID] = true
return m, m.generateRecapCmd(runID, answer)
m.recapIdleArmed = true
m.recapIdleRunID = runID
return m.scheduleIdleRecap(runID)
}
// handleRecapGenerated appends the finished recap as a footnote row. A failed or
// empty generation is silent (and releases the per-run gate so a future turn can
// recap normally).
func (m model) scheduleIdleRecap(runID int) (model, tea.Cmd) {
m.recapSeq++
seq := m.recapSeq
ctx, cancel := context.WithCancel(context.Background())
m.recapTimerCancel = cancel
return m, func() tea.Msg {
timer := time.NewTimer(recapIdleDelay)
defer timer.Stop()
select {
case <-timer.C:
return recapIdleMsg{runID: runID, seq: seq}
case <-ctx.Done():
return nil
}
}
}
// cancelIdleRecap invalidates pending timers, aborts provider work, and clears
// the ephemeral orientation note as soon as the user resumes activity.
func (m model) cancelIdleRecap() model {
m.recapIdleArmed = false
m.recapIdleRunID = 0
return m.clearIdleRecapWork()
}
func (m model) clearIdleRecapWork() model {
m.recapSeq++
if m.recapTimerCancel != nil {
m.recapTimerCancel()
m.recapTimerCancel = nil
}
if m.recapCancel != nil {
m.recapCancel()
m.recapCancel = nil
}
m.recapRunning = false
m.idleRecap = ""
return m
}
// resetIdleRecapAfterActivity restarts the inactivity window for a completed
// turn. Activity still cancels in-flight generation and clears the ephemeral
// note, but scrolling or reading no longer disables recaps until another turn.
func (m model) resetIdleRecapAfterActivity() (model, tea.Cmd) {
armed := m.recapIdleArmed
runID := m.recapIdleRunID
m = m.clearIdleRecapWork()
if !armed || !m.recapsEnabled || m.provider == nil || runID != m.runID {
return m, nil
}
return m.scheduleIdleRecap(runID)
}
// handleRecapIdle starts generation only if the same run is still genuinely
// idle. Busy UI, a draft, a modal, or newer activity silently abandons it.
func (m model) handleRecapIdle(msg recapIdleMsg) (model, tea.Cmd) {
if msg.seq != m.recapSeq || msg.runID != m.runID || !m.recapsEnabled ||
m.provider == nil || m.pending || m.compactInFlight || m.recapRunning ||
strings.TrimSpace(m.composerValue()) != "" || !m.noBlockingModal() {
return m, nil
}
if strings.TrimSpace(m.recapContext()) == "" {
return m, nil
}
m.recapTimerCancel = nil
ctx, cancel := context.WithTimeout(context.Background(), recapTimeout)
m.recapCancel = cancel
m.recapRunning = true
return m, m.generateRecapCmd(ctx, cancel, msg.runID, msg.seq)
}
// handleRecapGenerated publishes an ephemeral orientation note above the
// composer. Stale, cancelled, failed, or now-busy results are silent.
func (m model) handleRecapGenerated(msg recapGeneratedMsg) (model, tea.Cmd) {
// Drop a recap whose run is no longer the latest: a newer turn has started, so
// appending now would land the recap on the wrong (current) conversation.
if msg.runID != m.runID {
delete(m.recappedRuns, msg.runID)
if msg.seq != m.recapSeq || msg.runID != m.runID {
return m, nil
}
m.recapCancel = nil
m.recapRunning = false
m.recapIdleArmed = false
m.recapIdleRunID = 0
recap := strings.TrimSpace(msg.recap)
if msg.err != nil || recap == "" {
delete(m.recappedRuns, msg.runID)
if msg.err != nil || recap == "" || m.pending || m.compactInFlight ||
strings.TrimSpace(m.composerValue()) != "" || !m.noBlockingModal() {
return m, nil
}
m.transcript = appendTranscriptRow(m.transcript, transcriptRow{kind: rowRecap, text: recap})
m.idleRecap = recap
return m, nil
}
@@ -130,6 +248,9 @@ func (m model) handleConfigCommand(arg string) (model, string) {
switch arg {
case "recaps on", "recaps off":
m.recapsEnabled = arg == "recaps on"
if !m.recapsEnabled {
m = m.cancelIdleRecap()
}
if err := m.persistRecapsEnabled(); err != nil {
return m, "Config\nFailed to save recaps preference: " + err.Error()
}
+186 -69
View File
@@ -1,16 +1,20 @@
package tui
import (
"context"
"os"
"path/filepath"
"strings"
"testing"
"time"
tea "charm.land/bubbletea/v2"
"github.com/Gitlawb/zero/internal/sessions"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
// TestHandleConfigCommandReportsPersistError: a failed persist surfaces an error
// instead of falsely reporting "recaps: on/off".
func TestHandleConfigCommandReportsPersistError(t *testing.T) {
// Point the config path under a regular file, so the write can't create it.
dir := t.TempDir()
parent := filepath.Join(dir, "afile")
if err := os.WriteFile(parent, []byte("x"), 0o600); err != nil {
@@ -22,85 +26,179 @@ func TestHandleConfigCommandReportsPersistError(t *testing.T) {
}
}
// TestMaybeRecapTurnGating: a recap fires only when enabled, a provider exists,
// and there's a final answer — and at most once per run.
func TestMaybeRecapTurnGating(t *testing.T) {
func TestMaybeScheduleIdleRecapDefersProviderWork(t *testing.T) {
provider := &fakeProvider{}
withProvider := func() model {
return model{recapsEnabled: true, provider: &fakeProvider{}}
return model{recapsEnabled: true, provider: provider}
}
// Disabled -> no cmd.
if _, cmd := (model{recapsEnabled: false, provider: &fakeProvider{}}).maybeRecapTurn(1, "did X"); cmd != nil {
t.Error("recaps disabled: want no cmd")
if _, cmd := (model{provider: provider}).maybeScheduleIdleRecap(1, "did X"); cmd != nil {
t.Error("recaps disabled: want no timer")
}
// No provider -> no cmd.
if _, cmd := (model{recapsEnabled: true}).maybeRecapTurn(1, "did X"); cmd != nil {
t.Error("no provider: want no cmd")
if _, cmd := (model{recapsEnabled: true}).maybeScheduleIdleRecap(1, "did X"); cmd != nil {
t.Error("no provider: want no timer")
}
// Empty answer -> no cmd.
if _, cmd := withProvider().maybeRecapTurn(1, " "); cmd != nil {
t.Error("empty answer: want no cmd")
if _, cmd := withProvider().maybeScheduleIdleRecap(1, " "); cmd != nil {
t.Error("empty answer: want no timer")
}
// Enabled + provider + answer -> a cmd, and the per-run gate is set.
m, cmd := withProvider().maybeRecapTurn(7, "Built the website")
m, cmd := withProvider().maybeScheduleIdleRecap(7, "Built the website")
if cmd == nil || m.recapSeq != 1 {
t.Fatalf("successful turn should arm one idle timer: seq=%d cmd=%v", m.recapSeq, cmd != nil)
}
if len(provider.requests) != 0 {
t.Fatal("scheduling an idle recap must not call the provider")
}
}
func TestHandleRecapIdleUsesGoalPlanAndRecentConversation(t *testing.T) {
provider := &fakeProvider{events: []zeroruntime.StreamEvent{
{Type: zeroruntime.StreamEventText, Content: "The goal is to ship safer edits. Next, run the cross-platform checks."},
{Type: zeroruntime.StreamEventDone},
}}
m := newModel(context.Background(), Options{Provider: provider})
m.recapsEnabled = true
m.runID = 4
m.recapSeq = 9
m.activeSession.Goal = &sessions.Goal{Objective: "Ship safer edits", Status: sessions.GoalStatusActive}
m.plan.steps = []planStep{{content: "Run cross-platform checks", status: "in_progress"}}
m.transcript = []transcriptRow{
{kind: rowUser, text: "Make edit safety reliable"},
{kind: rowAssistant, text: "The implementation is ready.", final: true},
}
next, cmd := m.handleRecapIdle(recapIdleMsg{runID: 4, seq: 9})
if cmd == nil || !next.recapRunning || next.recapCancel == nil {
t.Fatal("an eligible idle session should start cancellable recap generation")
}
if len(provider.requests) != 0 {
t.Fatal("provider work must remain inside the returned command")
}
msg, ok := cmd().(recapGeneratedMsg)
if !ok {
t.Fatalf("command returned %T, want recapGeneratedMsg", msg)
}
if len(provider.requests) != 1 || len(provider.requests[0].Messages) != 2 {
t.Fatalf("recap request = %#v", provider.requests)
}
contextMessage := provider.requests[0].Messages[1].Content
for _, want := range []string{"Overall goal: Ship safer edits", "Current task: Run cross-platform checks", "User: Make edit safety reliable", "Assistant: The implementation is ready."} {
if !strings.Contains(contextMessage, want) {
t.Errorf("recap context missing %q:\n%s", want, contextMessage)
}
}
next, _ = next.handleRecapGenerated(msg)
if !strings.Contains(next.idleRecap, "cross-platform checks") || next.recapRunning {
t.Fatalf("generated idle recap not published correctly: recap=%q running=%v", next.idleRecap, next.recapRunning)
}
}
func TestRecapIdleAndGeneratedMessagesAreDroppedWhenStaleOrBusy(t *testing.T) {
provider := &fakeProvider{}
base := newModel(context.Background(), Options{Provider: provider})
base.recapsEnabled = true
base.runID = 3
base.recapSeq = 5
base.transcript = []transcriptRow{{kind: rowAssistant, text: "done", final: true}}
cases := []struct {
name string
edit func(*model)
msg recapIdleMsg
}{
{name: "stale activity", msg: recapIdleMsg{runID: 3, seq: 4}},
{name: "newer run", msg: recapIdleMsg{runID: 2, seq: 5}},
{name: "pending run", msg: recapIdleMsg{runID: 3, seq: 5}, edit: func(m *model) { m.pending = true }},
{name: "draft", msg: recapIdleMsg{runID: 3, seq: 5}, edit: func(m *model) { m.input.SetValue("still typing") }},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
m := base
if tc.edit != nil {
tc.edit(&m)
}
if _, cmd := m.handleRecapIdle(tc.msg); cmd != nil {
t.Fatal("ineligible idle message started provider work")
}
})
}
base.idleRecap = ""
base, _ = base.handleRecapGenerated(recapGeneratedMsg{runID: 3, seq: 4, recap: "stale"})
if base.idleRecap != "" {
t.Fatal("stale generated recap must be dropped")
}
}
func TestUserActivityCancelsAndClearsIdleRecap(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
m := newModel(context.Background(), Options{})
m.recapSeq = 2
m.recapRunning = true
m.recapCancel = cancel
m.idleRecap = "Continue with the release checks."
updated, _ := m.Update(tea.KeyPressMsg{Code: 'x', Text: "x"})
next := updated.(model)
if next.recapSeq != 3 || next.recapRunning || next.recapCancel != nil || next.idleRecap != "" {
t.Fatalf("activity did not clear recap state: seq=%d running=%v cancel=%v recap=%q", next.recapSeq, next.recapRunning, next.recapCancel != nil, next.idleRecap)
}
select {
case <-ctx.Done():
default:
t.Fatal("activity did not cancel in-flight recap context")
}
}
func TestUserActivityRearmsPendingIdleRecap(t *testing.T) {
provider := &fakeProvider{}
m := newModel(context.Background(), Options{Provider: provider})
m.recapsEnabled = true
m.runID = 7
m.transcript = []transcriptRow{{kind: rowAssistant, text: "done", final: true}}
m, initial := m.maybeScheduleIdleRecap(7, "done")
if initial == nil {
t.Fatal("precondition: successful turn did not arm recap")
}
oldTimerDone := make(chan tea.Msg, 1)
go func() { oldTimerDone <- initial() }()
updated, cmd := m.Update(tea.MouseMotionMsg{})
next := updated.(model)
if cmd == nil {
t.Fatal("enabled run with an answer should dispatch a recap cmd")
t.Fatal("activity canceled the pending recap without rearming the idle timer")
}
if !m.recappedRuns[7] {
t.Error("the per-run gate should be set")
if next.recapSeq <= m.recapSeq {
t.Fatalf("activity did not invalidate the old timer: before=%d after=%d", m.recapSeq, next.recapSeq)
}
// Second call for the same run -> no cmd (already recapped).
if _, cmd2 := m.maybeRecapTurn(7, "Built the website"); cmd2 != nil {
t.Error("a second recap for the same run must not fire")
select {
case msg := <-oldTimerDone:
if msg != nil {
t.Fatalf("canceled timer returned %#v, want nil", msg)
}
case <-time.After(time.Second):
t.Fatal("activity left the old four-minute timer running")
}
}
// TestHandleRecapGenerated: a successful recap appends a "※ recap:" row; a
// failed/empty one appends nothing and releases the gate.
func TestHandleRecapGenerated(t *testing.T) {
// runID matches the latest run -> the recap appends.
m := model{runID: 5, recappedRuns: map[int]bool{5: true}}
m, _ = m.handleRecapGenerated(recapGeneratedMsg{runID: 5, recap: "Built a fashion site with cart and blog"})
if len(m.transcript) != 1 {
t.Fatalf("a recap for the current run should append one row, got %d", len(m.transcript))
}
last := m.transcript[len(m.transcript)-1]
if last.kind != rowRecap || !strings.Contains(last.text, "fashion site") {
t.Errorf("recap row wrong: %+v", last)
}
if got := plainRender(t, renderRecapRow(last, 100)); !strings.Contains(got, "※ recap:") {
t.Errorf("recap render should show the ※ marker, got %q", got)
}
// A recap from a STALE run (a newer turn started) is dropped, not appended to
// the wrong conversation; the gate is released.
stale := model{runID: 8, recappedRuns: map[int]bool{6: true}}
stale, _ = stale.handleRecapGenerated(recapGeneratedMsg{runID: 6, recap: "old run recap"})
if len(stale.transcript) != 0 {
t.Error("a stale-run recap must not be appended")
}
if stale.recappedRuns[6] {
t.Error("a stale-run recap must release the per-run gate")
}
// Failed recap: no row, gate released.
m2 := model{runID: 9, recappedRuns: map[int]bool{9: true}}
m2, _ = m2.handleRecapGenerated(recapGeneratedMsg{runID: 9, err: errExampleRecap})
if len(m2.transcript) != 0 {
t.Error("a failed recap must append nothing")
}
if m2.recappedRuns[9] {
t.Error("a failed recap must release the per-run gate")
func TestIdleRecapRendersAboveComposer(t *testing.T) {
m := newModel(context.Background(), Options{})
m.width = 100
m.height = 30
m.idleRecap = "Continue by running the release checks."
got := plainRender(t, m.footerView(100))
if !strings.Contains(got, "※ Continue by running the release checks.") {
t.Fatalf("footer missing idle recap:\n%s", got)
}
}
// TestHandleConfigCommandTogglesRecaps: /config recaps on|off flips the flag and
// reports it; unknown args are rejected.
func TestHandleConfigCommandTogglesRecaps(t *testing.T) {
m := model{recapsEnabled: true}
m := model{recapsEnabled: true, idleRecap: "old recap"}
m, text := m.handleConfigCommand("recaps off")
if m.recapsEnabled || !strings.Contains(text, "recaps: off") {
t.Errorf("recaps off: enabled=%v text=%q", m.recapsEnabled, text)
if m.recapsEnabled || m.idleRecap != "" || !strings.Contains(text, "recaps: off") {
t.Errorf("recaps off: enabled=%v recap=%q text=%q", m.recapsEnabled, m.idleRecap, text)
}
m, text = m.handleConfigCommand("recaps on")
if !m.recapsEnabled || !strings.Contains(text, "recaps: on") {
@@ -111,8 +209,27 @@ func TestHandleConfigCommandTogglesRecaps(t *testing.T) {
}
}
var errExampleRecap = errExample("recap failed")
func TestHandleConfigCommandCancelsInFlightRecap(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
m := model{
recapsEnabled: true,
recapRunning: true,
recapCancel: cancel,
idleRecap: "old recap",
}
type errExample string
func (e errExample) Error() string { return string(e) }
m, _ = m.handleConfigCommand("recaps off")
if m.recapRunning || m.recapCancel != nil || m.idleRecap != "" {
t.Fatalf(
"disabling recaps did not clear in-flight state: running=%v cancel=%v recap=%q",
m.recapRunning,
m.recapCancel != nil,
m.idleRecap,
)
}
select {
case <-ctx.Done():
default:
t.Fatal("disabling recaps did not cancel in-flight generation context")
}
}
+45 -2
View File
@@ -1477,13 +1477,13 @@ func renderToolResultCard(row transcriptRow, width int, rc rowContext, opts card
collapsedFooter = collapsedToolFooter(row.detail)
}
if collapsedFooter != "" && !row.expanded {
head := toolCardHead(name, headTarget, headArg, "", row.detail, row.text, false, nameStyle, rc.auto[key], width, opts)
head := toolCardHead(name, headTarget, headArg, toolResultBudgetTag(row.meta), row.detail, row.text, false, nameStyle, rc.auto[key], width, opts)
return toolCard(head, glyph, nil, collapsedFooter, borderStyle, width)
}
bodyOpts := opts
bodyOpts.expanded = row.expanded
body := toolCardBody(name, rc.hints[key], rc.args[key], row.detail, width, bodyOpts, failed)
head := toolCardHead(name, headTarget, headArg, body.headTag, row.detail, row.text, false, nameStyle, rc.auto[key], width, opts)
head := toolCardHead(name, headTarget, headArg, joinToolHeadTags(body.headTag, toolResultBudgetTag(row.meta)), row.detail, row.text, false, nameStyle, rc.auto[key], width, opts)
footer := body.footer
if collapsedFooter != "" && row.expanded && footer == "" {
footer = "▾ collapse"
@@ -1491,6 +1491,49 @@ func renderToolResultCard(row transcriptRow, width int, rc rowContext, opts card
return toolCard(head, glyph, body.lines, footer, borderStyle, width)
}
func joinToolHeadTags(tags ...string) string {
var present []string
for _, tag := range tags {
if tag = strings.TrimSpace(tag); tag != "" {
present = append(present, tag)
}
}
return strings.Join(present, " · ")
}
func toolResultBudgetTag(meta map[string]string) string {
if len(meta) == 0 {
return ""
}
reduced := meta["command_output_reduced"] == "true" || meta["truncated"] == "true"
if !reduced {
return ""
}
original, originalErr := strconv.Atoi(meta["output_budget_estimated_original_tokens"])
retained, retainedErr := strconv.Atoi(meta["output_budget_estimated_retained_tokens"])
if commandOriginal, err := strconv.Atoi(meta["command_output_original_tokens"]); err == nil {
original, originalErr = commandOriginal, nil
}
if commandRetained, err := strconv.Atoi(meta["command_output_retained_tokens"]); err == nil {
retained, retainedErr = commandRetained, nil
}
tag := "compact output"
if originalErr == nil && retainedErr == nil && original > retained {
tag = fmt.Sprintf("compact %s→%s tok", compactCount(original), compactCount(retained))
}
if strings.TrimSpace(meta["spill_path"]) != "" {
tag += " · raw saved"
}
return tag
}
func compactCount(value int) string {
if value < 1000 {
return strconv.Itoa(value)
}
return fmt.Sprintf("%.1fk", float64(value)/1000)
}
// confirmationVerbPattern matches a single-line success confirmation that only
// restates the action + target: a leading verb ("Created", "Overwrote",
// "Successfully created directory", "Wrote", "Deleted", …) optionally followed
+44 -7
View File
@@ -29,8 +29,13 @@ func (m model) ensureActiveSession(prompt string) (model, error) {
return m, nil
}
title := strings.TrimSpace(m.pendingSessionTitle)
manuallyNamed := title != ""
if !manuallyNamed {
title = tuiSessionTitle(prompt)
}
session, err := m.sessionStore.Create(sessions.CreateInput{
Title: tuiSessionTitle(prompt),
Title: title,
Cwd: m.cwd,
ModelID: m.modelName,
Provider: m.providerName,
@@ -39,7 +44,14 @@ func (m model) ensureActiveSession(prompt string) (model, error) {
return m, err
}
m.activeSession = session
m.pendingSessionTitle = ""
m.sessionEvents = []sessions.Event{}
if manuallyNamed {
if m.titledSessions == nil {
m.titledSessions = map[string]bool{}
}
m.titledSessions[session.SessionID] = true
}
return m, nil
}
@@ -55,6 +67,7 @@ func (m model) startNewSession() model {
previousID := m.activeSession.SessionID
m.activeSession = sessions.Metadata{}
m.pendingSessionTitle = ""
m.sessionEvents = nil
// Reset the per-session usage + compaction display so the new session starts
@@ -217,6 +230,7 @@ func (m model) handleResumeCommand(args string) (model, string) {
// the already-active session, whose loops belong to it, not a "previous" one.
previousID := m.activeSession.SessionID
m.activeSession = *session
m.pendingSessionTitle = ""
m.sessionEvents = append([]sessions.Event{}, events...)
if m.providerName == "" {
m.providerName = session.Provider
@@ -392,17 +406,16 @@ func (m model) newSessionPicker() *commandPicker {
if !m.sessionHasResumableContent(meta.SessionID) {
continue
}
// Lead with the timestamp so same-titled sessions (e.g. the same first
// prompt run several times) are visually distinct; the id (right, faint)
// stays for reference and is what selection actually resolves.
// Lead with a fixed-width timestamp so titles form one scannable column.
// The raw id remains the selection/search value but stays out of the row:
// rendering it consumed half the picker and truncated the useful title.
label := displayValue(meta.Title, "untitled")
if when := sessionWhen(meta.UpdatedAt, now); when != "" {
label = when + " " + label
label = sessionPickerLabel(when, label)
}
items = append(items, pickerItem{
Label: label,
Value: meta.SessionID,
Meta: meta.SessionID,
})
}
if len(items) == 0 {
@@ -417,6 +430,12 @@ func (m model) newSessionPicker() *commandPicker {
}
}
const sessionPickerTimeWidth = len("Jan 02 15:04")
func sessionPickerLabel(when, title string) string {
return fmt.Sprintf("%-*s %s", sessionPickerTimeWidth, when, title)
}
// sessionHasResumableContent reports whether a session has anything worth
// resuming: a tool call/result, or a non-user message with real content (not the
// no-output guardrail stop). Empty/failed runs return false and are hidden from
@@ -479,7 +498,7 @@ func (m model) sessionHasResumableContent(sessionID string) bool {
// anything worth resuming: a tool call/result, or a non-user message with real
// content (not the no-output guardrail stop). It is the pure core of
// sessionHasResumableContent so callers that already hold the events (e.g. the
// /retitle scan) don't re-read them.
// session picker refresh) don't re-read them.
func eventsHaveResumableContent(events []sessions.Event) bool {
for _, event := range events {
switch event.Type {
@@ -616,6 +635,7 @@ func transcriptRowsFromSessionEvents(events []sessions.Event) []transcriptRow {
tool: name,
status: status,
detail: output,
meta: payloadStringMap(payload, "meta"),
changedFiles: payloadStringSlice(payload, "changedFiles"),
changeSummaries: payloadExecutionChanges(payload, "changeSummaries"),
})
@@ -809,6 +829,23 @@ func payloadStringSlice(payload map[string]any, key string) []string {
}
}
func payloadStringMap(payload map[string]any, key string) map[string]string {
value, ok := payloadMap(payload, key)
if !ok {
return nil
}
out := make(map[string]string, len(value))
for name, raw := range value {
if text, ok := raw.(string); ok {
out[name] = text
}
}
if len(out) == 0 {
return nil
}
return out
}
func payloadExecutionChanges(payload map[string]any, key string) []execution.Change {
value, ok := payload[key]
if !ok {
+106
View File
@@ -0,0 +1,106 @@
package tui
import (
"fmt"
"strings"
tea "charm.land/bubbletea/v2"
"charm.land/lipgloss/v2"
)
type sessionRenamePrompt struct{}
func (m model) openSessionRenamePrompt() model {
if m.sessionStore == nil {
return m.appendSessionRenameError("session storage is unavailable")
}
m.clearComposer()
title := m.activeSession.Title
if m.activeSession.SessionID == "" {
title = m.pendingSessionTitle
}
m.input.SetValue(title)
m.input.CursorEnd()
m.renamePrompt = &sessionRenamePrompt{}
return m
}
func (m model) handleSessionRenameKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
switch {
case keyIs(msg, tea.KeyEsc), keyCtrl(msg, 'c'):
m.renamePrompt = nil
m.clearComposer()
return m, nil
case keyIs(msg, tea.KeyEnter):
title := m.input.Value()
m.renamePrompt = nil
m.clearComposer()
return m.renameActiveSession(title), nil
default:
var cmd tea.Cmd
m.input, cmd = m.input.Update(msg)
return m, cmd
}
}
func (m model) renameActiveSession(title string) model {
title = cutRunes(strings.TrimSpace(title), tuiSessionTitleLimit)
if title == "" {
return m.appendSessionRenameError("session name cannot be empty")
}
if m.sessionStore == nil {
return m.appendSessionRenameError("session storage is unavailable")
}
if m.activeSession.SessionID == "" {
m.pendingSessionTitle = title
m.transcript = reduceTranscript(m.transcript, transcriptAction{
kind: actionAppendSystem,
text: "Session renamed to " + title + ".",
})
return m
}
updated, err := m.sessionStore.UpdateTitle(m.activeSession.SessionID, title)
if err != nil {
return m.appendSessionRenameError(err.Error())
}
m.activeSession.Title = updated.Title
if m.titledSessions == nil {
m.titledSessions = map[string]bool{}
}
m.titledSessions[updated.SessionID] = true
m.transcript = reduceTranscript(m.transcript, transcriptAction{
kind: actionAppendSystem,
text: "Session renamed to " + updated.Title + ".",
})
return m
}
func (m model) appendSessionRenameError(detail string) model {
m.transcript = reduceTranscript(m.transcript, transcriptAction{
kind: actionAppendError,
text: fmt.Sprintf("Could not rename session: %s.", strings.TrimSuffix(strings.TrimSpace(detail), ".")),
})
return m
}
func (m model) sessionRenamePromptView(width int) string {
if width <= 0 {
width = defaultStartupWidth
}
input := m.input
input.Prompt = "> "
input.Placeholder = "Type a name"
innerWidth := maxInt(1, width-4)
line := renderComposerInput(
input,
composerState{text: input.Value(), cursor: input.Position()},
innerWidth,
m.composerCursorVisible,
composerSelectionState{},
)
lines := append(strings.Split(line, "\n"),
zeroTheme.line.Render(strings.Repeat("─", innerWidth)),
zeroTheme.faint.Render("Enter save Esc cancel"),
)
return styledBlockFillTitle(width, "Rename session", lines, zeroTheme.lineStrong, lipgloss.NewStyle())
}
+190
View File
@@ -0,0 +1,190 @@
package tui
import (
"context"
"strings"
"testing"
tea "charm.land/bubbletea/v2"
"github.com/Gitlawb/zero/internal/sessions"
)
func renameTestModel(t *testing.T) (model, *sessions.Store, sessions.Metadata) {
t.Helper()
store := testSessionStore(t)
session, err := store.Create(sessions.CreateInput{Title: "Current session name"})
if err != nil {
t.Fatalf("create session: %v", err)
}
m := newModel(context.Background(), Options{SessionStore: store})
m.activeSession = session
return m, store, session
}
func TestRenameCommandRenamesCurrentSessionWithoutAgentRun(t *testing.T) {
m, store, session := renameTestModel(t)
m.input.SetValue("/rename Better session name ")
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if cmd != nil {
t.Fatal("/rename must not start an agent or background command")
}
if next.activeSession.Title != "Better session name" {
t.Fatalf("active title = %q", next.activeSession.Title)
}
stored, err := store.Get(session.SessionID)
if err != nil || stored == nil {
t.Fatalf("get renamed session: %v", err)
}
if stored.Title != "Better session name" {
t.Fatalf("stored title = %q", stored.Title)
}
if !next.titledSessions[session.SessionID] {
t.Fatal("manual rename should suppress later automatic naming")
}
if !transcriptContains(next.transcript, "Session renamed to Better session name.") {
t.Fatalf("missing rename confirmation: %#v", next.transcript)
}
}
func TestRenameCommandCapsSessionTitle(t *testing.T) {
m, store, session := renameTestModel(t)
longTitle := strings.Repeat("界", tuiSessionTitleLimit+20)
m.input.SetValue("/rename " + longTitle)
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if cmd != nil {
t.Fatal("/rename must not start an agent or background command")
}
want := cutRunes(longTitle, tuiSessionTitleLimit)
if next.activeSession.Title != want {
t.Fatalf("active title has %d runes, want %d", len([]rune(next.activeSession.Title)), tuiSessionTitleLimit)
}
stored, err := store.Get(session.SessionID)
if err != nil || stored == nil || stored.Title != want {
t.Fatalf("stored capped title = %#v, err=%v", stored, err)
}
}
func TestRenameEditorBordersWrappedTitleLines(t *testing.T) {
m, _, _ := renameTestModel(t)
m.input.SetValue(strings.Repeat("x", tuiSessionTitleLimit))
m.input.CursorEnd()
m.renamePrompt = &sessionRenamePrompt{}
view := plainRender(t, m.sessionRenamePromptView(80))
lines := strings.Split(view, "\n")
if len(lines) < 6 {
t.Fatalf("expected wrapped rename editor, got:\n%s", view)
}
for _, line := range lines[1 : len(lines)-1] {
if !strings.HasPrefix(line, "│ ") || !strings.HasSuffix(line, " │") {
t.Fatalf("wrapped editor row lost its border: %q\n%s", line, view)
}
}
}
func TestBareRenameOpensPrefilledEditorAndSaves(t *testing.T) {
m, store, session := renameTestModel(t)
m.input.SetValue("/rename")
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if cmd != nil || next.renamePrompt == nil {
t.Fatalf("bare /rename should open editor without a command: prompt=%#v cmd=%v", next.renamePrompt, cmd)
}
if next.input.Value() != "Current session name" {
t.Fatalf("editor value = %q", next.input.Value())
}
view := plainRender(t, next.View())
for _, want := range []string{"Rename session", "Current session name", "Enter save", "Esc cancel"} {
if !strings.Contains(view, want) {
t.Fatalf("rename editor missing %q:\n%s", want, view)
}
}
next.input.SetValue("Edited in prompt")
updated, cmd = next.Update(testKey(tea.KeyEnter))
next = updated.(model)
if cmd != nil || next.renamePrompt != nil {
t.Fatalf("saving should close editor without a command: prompt=%#v cmd=%v", next.renamePrompt, cmd)
}
stored, err := store.Get(session.SessionID)
if err != nil || stored == nil || stored.Title != "Edited in prompt" {
t.Fatalf("stored session after editor save = %#v, err=%v", stored, err)
}
}
func TestRenameNamesFreshSessionBeforeFirstPrompt(t *testing.T) {
store := testSessionStore(t)
m := newModel(context.Background(), Options{SessionStore: store})
m.input.SetValue("/rename Fresh session name")
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if cmd != nil {
t.Fatal("fresh /rename must stay local")
}
if next.activeSession.SessionID != "" || next.pendingSessionTitle != "Fresh session name" {
t.Fatalf("fresh rename state = active:%q pending:%q", next.activeSession.SessionID, next.pendingSessionTitle)
}
next, err := next.ensureActiveSession("first real prompt")
if err != nil {
t.Fatalf("ensure active session: %v", err)
}
if next.activeSession.Title != "Fresh session name" || next.pendingSessionTitle != "" {
t.Fatalf("created session title = %q, pending=%q", next.activeSession.Title, next.pendingSessionTitle)
}
if !next.titledSessions[next.activeSession.SessionID] {
t.Fatal("a manually named fresh session must skip automatic naming")
}
stored, err := store.Get(next.activeSession.SessionID)
if err != nil || stored == nil || stored.Title != "Fresh session name" {
t.Fatalf("stored fresh session = %#v, err=%v", stored, err)
}
}
func TestRenameEditorRejectsBlankAndEscCancels(t *testing.T) {
m, store, session := renameTestModel(t)
m = m.openSessionRenamePrompt()
m.input.SetValue(" ")
updated, _ := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if next.renamePrompt != nil {
t.Fatal("blank submit should close the editor")
}
if !transcriptContains(next.transcript, "session name cannot be empty") {
t.Fatalf("missing blank-name error: %#v", next.transcript)
}
stored, _ := store.Get(session.SessionID)
if stored == nil || stored.Title != "Current session name" {
t.Fatalf("blank rename changed title: %#v", stored)
}
next = next.openSessionRenamePrompt()
next.input.SetValue("Do not save")
updated, _ = next.Update(testKey(tea.KeyEsc))
next = updated.(model)
if next.renamePrompt != nil || next.input.Value() != "" {
t.Fatalf("Esc should close and clear editor: prompt=%#v input=%q", next.renamePrompt, next.input.Value())
}
stored, _ = store.Get(session.SessionID)
if stored == nil || stored.Title != "Current session name" {
t.Fatalf("cancelled rename changed title: %#v", stored)
}
}
func TestRetitleCommandIsRemoved(t *testing.T) {
if command, ok := resolveCommand("/retitle"); ok {
t.Fatalf("/retitle should not resolve, got %#v", command)
}
command, ok := resolveCommand("/rename")
if !ok || command.kind != commandRename {
t.Fatalf("/rename should resolve, got ok=%v command=%#v", ok, command)
}
}
+94 -3
View File
@@ -6,6 +6,7 @@ import (
"os"
"path/filepath"
"strings"
"sync/atomic"
"testing"
"time"
@@ -19,9 +20,10 @@ import (
)
type scriptedProvider struct {
scripts [][]zeroruntime.StreamEvent
requests []zeroruntime.CompletionRequest
calls int
scripts [][]zeroruntime.StreamEvent
requests []zeroruntime.CompletionRequest
beforeCall func(int)
calls int
}
func (provider *scriptedProvider) StreamCompletion(
@@ -39,6 +41,9 @@ func (provider *scriptedProvider) StreamCompletion(
if index >= len(provider.scripts) {
index = len(provider.scripts) - 1
}
if provider.beforeCall != nil {
provider.beforeCall(index)
}
ch := make(chan zeroruntime.StreamEvent, len(provider.scripts[index]))
for _, event := range provider.scripts[index] {
ch <- event
@@ -358,6 +363,92 @@ func TestPromptSubmitPersistsPermissionSessionEvents(t *testing.T) {
}
}
func TestPermissionWaitDoesNotCountTowardTurnElapsed(t *testing.T) {
store := testSessionStore(t)
root := t.TempDir()
base := time.Date(2026, 7, 29, 12, 0, 0, 0, time.UTC)
var clockNanos atomic.Int64
clockNanos.Store(base.UnixNano())
provider := &scriptedProvider{scripts: [][]zeroruntime.StreamEvent{
writeFileToolScript("call_write", "notes.txt", "hello"),
textScript("write blocked"),
}}
provider.beforeCall = func(index int) {
switch index {
case 0:
clockNanos.Store(base.Add(2 * time.Second).UnixNano())
case 1:
clockNanos.Store(base.Add(95 * time.Second).UnixNano())
}
}
registry := tools.NewRegistry()
registry.Register(tools.NewScopedWriteFileTool(root, nil))
runtimeMessageCh := make(chan tea.Msg, 8)
m := newPermissionTestModel(root, provider, registry, store, nil, runtimeMessageCh)
m.now = func() time.Time {
return time.Unix(0, clockNanos.Load()).UTC()
}
m.input.SetValue("write notes")
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if cmd == nil {
t.Fatal("expected prompt submit to start an agent run")
}
finalCh := make(chan tea.Msg, 1)
go func() {
finalCh <- execCmd(cmd)
}()
for received := 0; received < 4; {
runtimeMsg := receiveRuntimeMessage(t, runtimeMessageCh)
updated, _ = next.Update(runtimeMsg)
next = updated.(model)
switch runtimeMsg.(type) {
case toolCallStreamStartMsg, toolCallStreamDeltaMsg:
continue
case permissionRequestMsg:
clockNanos.Store(base.Add(92 * time.Second).UnixNano())
if status := plainRender(t, next.workingStatusLine()); !strings.Contains(status, "· 2s ·") || strings.Contains(status, "still generating") {
t.Fatalf("live elapsed advanced during permission wait: %q", status)
}
updated, _ = next.Update(testKeyText("d"))
next = updated.(model)
if hint := next.quietGenerationHint(); hint != "" {
t.Fatalf("permission wait leaked into quiet-generation age: %q", hint)
}
}
received++
}
finalMsg := receiveFinalMessage(t, finalCh)
response, ok := finalMsg.(agentResponseMsg)
if !ok {
t.Fatalf("expected agent response, got %T", finalMsg)
}
if response.turnElapsed != 5*time.Second {
t.Fatalf("turn elapsed = %s, want 5s of active work", response.turnElapsed)
}
if response.ttft != 5*time.Second {
t.Fatalf("ttft = %s, want 5s with permission wait excluded", response.ttft)
}
if response.ttft > response.turnElapsed {
t.Fatalf("ttft = %s exceeds turn elapsed = %s", response.ttft, response.turnElapsed)
}
updated, _ = next.Update(finalMsg)
next = updated.(model)
answer, ok := findTranscriptRow(next.transcript, rowAssistant)
if !ok {
t.Fatalf("expected final assistant row, got %#v", next.transcript)
}
if answer.turnElapsed != 5*time.Second {
t.Fatalf("assistant turn elapsed = %s, want 5s", answer.turnElapsed)
}
}
func TestPermissionPromptAllowWritesFileAndRecordsDecision(t *testing.T) {
store := testSessionStore(t)
root := t.TempDir()
Loaded 100 of 110 files, more files were not shown because too many files have changed in this diff. Show more