Compare commits

..
Author SHA1 Message Date
anandh8x 72440a85bf fix(tui): align active voice status with capture key 2026-08-26 16:25:15 +05:30
anandh8x ed38f57fab fix(tui): clarify voice capture controls 2026-08-26 15:38:35 +05:30
anandh8x 7e6f5e2a4a fix(tui): keep voice mode from swallowing spaces 2026-08-26 14:54:23 +05:30
PierrunoYTandClaude Opus 5 27b319ca88 test(execution): fail loudly when ps cannot answer instead of skipping (#961)
* test(execution): fail loudly when ps cannot answer instead of skipping

Every test in process_unix_test.go proves something about zombie detection, and
every one of those proofs is read out of ps. When ps could not answer, the
tests skipped — which is indistinguishable from success in CI output — and one
assertion passed vacuously for the same reason. A regression in ps parsing or
availability would have shown up as green.

requireProcessState probes BOTH ps pipelines against this test process, which
is by definition alive and not a zombie, so a working ps must report a live
state for it. When it cannot:

- Under CI (or ZERO_REQUIRE_PS=1) this is a hard failure. ps is present on every
  image used today, so its absence means the pipeline regressed and that has to
  surface immediately.
- Otherwise it stays a skip, so a contributor on an unsupported box is not
  blocked.

Both pipelines are probed because they are deliberately separate: the tests
assert what the production process-table parser concludes, so this file keeps
its own per-PID ps oracle rather than borrowing that parser to establish its own
preconditions. One regression must not be able to silence both sides at once.

Once the probe passes, failing to observe an EXPECTED state is a defect rather
than an environment problem, so the remaining escapes become failures:

- The zombie waits (awaitZombie) fail instead of skipping. The caller never
  waits on the child, so its zombie row persists until reaped — reaching the
  deadline means the child never exited or ps stopped answering, and the two
  are reported distinctly.
- A zombie group that is not signalable fails instead of skipping: a zombie
  still occupies the process table, so if kill(-pid, 0) fails the case under
  test is not present and the assertions after it would prove nothing.
- processSignalTarget returning anything but the individual PID fails instead
  of skipping: the child inherited this process's group, so its PGID is this
  process's PID and can never equal its own PID.

The vacuous assertion in TestTerminateProcessTreeStopsRunningGroup now requires
ok. That is sound rather than merely stricter: the test never waits on the
child, so the group's zombie row must still be present after termination, and
an unknown state means ps stopped answering rather than that the group is
clean.

Verified on Linux (WSL, go1.26.6) rather than by cross-compilation alone:

- all three tests pass unmodified
- with a stub ps that exits 1: CI unset skips with the reason named; CI=true
  fails, naming which pipeline could not answer
- the same broken-ps run against the ORIGINAL file produced two silent skips
  plus one unrelated failure after 12s, versus an immediate, explicit failure now

Windows is unaffected (the file is !windows); go vet passes for linux, darwin
and windows.

Fixes #863

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01U389qmQUhoZB3YtXkFmiTq

* test(execution): treat an unsignalable zombie group as a kernel difference

The macOS Smoke job failed:

    the zombie group 5931 is not signalable (operation not permitted)

That was my error. I converted this skip along with the ps-availability ones,
but it never had anything to do with ps: whether kill(2) may signal a group
whose only member is a zombie is a KERNEL difference. Linux permits it, so
signalTargetRunning goes on to consult ps — the branch this test was written
for. Darwin returns EPERM, so signalTargetRunning short-circuits on the kill and
never reaches ps.

Rather than restore the blanket skip, the test now records which branch answers:

- Signal permitted: assert processGroupHasRunningMember directly, so a ps-side
  regression cannot hide behind the kill result. This is new coverage; the old
  test only checked signalTargetRunning, which conflates the two paths.
- Signal refused: log it and continue.

The outcome assertion — a zombie-only group must not count as running — now runs
on EVERY platform. Previously Darwin skipped before reaching it, so this is
strictly more coverage there, not less.

Any errno other than EPERM still fails: it means the group is not in the state
the test needs, and nothing after it would be meaningful.

TerminateProcessTree is asserted only where the signal is permitted, because
terminateTarget's own kill(2) would hit the same EPERM and return it as a
cleanup error. Worth a maintainer's eye: that may be a real gap in
TerminateProcessTree on such a platform — a zombie-only tree is already exited
and arguably should terminate cleanly rather than report failure, which is the
same spurious-cleanup-failure class #774 set out to remove. Confirming that
needs a Darwin host and is a production change, so it is flagged rather than
guessed at here.

Verified on Linux (WSL, go1.26.6): all three tests pass, including the new
direct ps assertion. The Darwin branch was exercised by forcing the EPERM path
on Linux — it logs both explanations and passes with the outcome assertion
still enforced. go vet passes for linux, darwin and windows.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01U389qmQUhoZB3YtXkFmiTq

---------

Co-authored-by: Claude Opus 5 (1M context) <[email protected]>
2026-08-26 13:22:43 +08:00
Anandan 1fde941803 perf: reduce startup and turn overhead (#932)
* perf(tui): disable duplicate composer cursor blink

* perf(tui): stop composer wakeups after idle

* refactor(tools): publish immutable registry snapshots

* perf(startup): defer optional tool discovery

* perf(provider): reuse compatible response sessions

* feat(provider): add native patch tool input

* fix(provider): recognize saved chat sessions

* perf(tui): render only on screen changes

* fix(runtime): address performance review findings

* fix(usage): expose cache efficiency details

* perf(agent): reduce redundant bounded-task turns

* fix: harden runtime review edge cases

* fix: address runtime review findings

* build: restore official Bubble Tea dependency

* fix(cli): redact MCP registration errors

* fix: harden turn session regressions

* fix: keep run tools and patch roots consistent

* fix: bind freeform patch writes to scoped roots

* test: keep patch fixtures out of worktree

* fix(tools): reject mixed-root freeform patches
2026-08-26 13:22:05 +08:00
KRATOS 7acac7acee fix(sandbox): refcount temporary roots so one holder's cleanup cannot revoke another's (#911)
* fix(sandbox): a temporary root is refcounted, so one holder's cleanup cannot revoke another's

Split out of #829 as an independent fix, per @Vasanthdev2004's review asking for
the self-contained pieces to arrive separately.

A temporary grant was add-then-remove with no notion of who still needed the
root. The second caller to ask for a root already present got a NO-OP undo, and
the first caller's cleanup then removed the root out from under it — so a live
grant lost its access while its holder still believed it had it.

Two read-only tools in the same parallel batch, both blocked on the same
directory, is exactly that shape, and read-only tools are precisely the ones the
batch runs concurrently.

tempReads/tempWrites now count live holders, and the root is only stripped when
the last one releases.

The release also performs both mutations under ONE hold of the lock. Dropping it
between them left a window where the root was still in readRoots but no longer in
tempReads: a concurrent AddTemporaryRead landing there reads it as a PERMANENT
root, hands its caller a no-op undo, and this call then strips the root — so that
caller believes it holds access it has already silently lost.

Independent of the remaining #829 work; builds and tests against current main on
its own.

Origin-Session: local-13d543 | Claude Code | 5 prompts
Origin-Snapshot: 6b5eb8ba4e5b
Origin-Session: local-c962d7 | Claude Code | 3 prompts
Origin-Snapshot: d2b8f44a9abc

* fix(sandbox): a read covered by a TEMPORARY write root takes a reference too

Found by a self-review of this branch: the refcount closed the defect within
reads and within writes, but not across them.

AddTemporaryRead returned a no-op undo whenever writeRootCoversLocked said a
write root already covered the path, and the comment there claimed that cover was
permanent. It is not: writeRootCoversLocked scans extraRoots, and
AddTemporaryWrite appends TEMPORARY write roots into that same slice. So a read
covered by another holder's temporary write took no reference, and that holder's
cleanup silently revoked the reader's access — the very defect the refcount
exists to close, reached across the read/write boundary instead of within one
side of it.

Demonstrated before fixing, and the regression fails without it:

  ReadRoots before release = [ws, …/outer]
  ReadRoots after  release = [ws]
  the write holder's cleanup revoked a live read grant on …/outer/inner

The read now takes a reference on the covering temporary write root, and the
genuinely permanent case — the workspace root, or a session-scoped grant — still
returns a no-op undo.

The regression also asserts the other direction: once the reader releases too,
the root is gone, so the reference does not leak the grant past its last holder.

Origin-Session: local-0484ba | Claude Code | 6 prompts
Origin-Snapshot: e6d90ff3ce46
Origin-Session: local-13d543 | Claude Code | 5 prompts
Origin-Snapshot: 6b5eb8ba4e5b
Origin-Session: local-c962d7 | Claude Code | 3 prompts
Origin-Snapshot: d2b8f44a9abc

* test(sandbox): make the concurrency test concurrent, and portable off macOS

Both raised by CodeRabbit. The third finding on this PR — a read covered by a
TEMPORARY write root taking no reference — was already fixed in 87a722cb with the
regression it asked for; this covers the two that were still live, and both are
cases of a test reporting more than it checked.

THE CONCURRENCY TEST WAS NOT CONCURRENT. The gate spun until hasReadRoot saw
somebody holding the root, which is satisfied by the FIRST holder, then opened
the release. The spin wins that race against goroutines still being scheduled:
instrumented over 200 runs, 194 peaked at a single simultaneous holder and not
one ever reached eight. A test named for concurrent holders was exercising one
holder at a time. Against a deliberately broken refcount, where the first release
strips the root, it passed 33 times in 40 — it was a flaky test that mostly
agreed with broken code. Counting all eight acquisitions makes the overlap real:
200 of 200 runs now peak at eight, and the same broken refcount is caught 40
times in 40. Holders signal the barrier on every path including failure, so a
holder that returns early fails the test rather than hanging it.

THE HELPER SKIPPED ON WINDOWS. scopeOutsideRoots reached for /Users/Shared and
fell back to /var/empty, neither of which exists there, so every test built on it
skipped and reported green having asserted nothing. It now uses the home
directory, which no platform lists as a default write root, matching what
scope_extra_read_test.go already does — and then PROVES the choice against
defaultTempWriteRoots rather than assuming it, so a future default root turns
these tests into an honest skip instead of a silent no-op.

Origin-Session: local-abff1c | Claude Code | 2 prompts
Origin-Snapshot: d2f269b81f33
Origin-Session: local-13d543 | Claude Code | 5 prompts
Origin-Snapshot: 6b5eb8ba4e5b
Origin-Session: local-c962d7 | Claude Code | 3 prompts
Origin-Snapshot: d2b8f44a9abc

* fix(sandbox): a permanent grant outlives the temporary one, and an undo is per holder

Both P1s from @anandh8x, both reproduced before changing anything, and both are
the same mistake in different places: treating "something covers this path" as
though it answered "and will it still be there afterwards".

A PERMANENT GRANT WAS REVOKED BY A TEMPORARY HOLDER. extraRoots holds temporary
write roots alongside permanent ones, so Add asked only whether the path was
already covered. A session-scoped grant made while a temporary holder happened to
cover it recorded nothing at all, and disappeared the moment that holder
released. Outliving the request that prompted it is the entire difference between
Add and AddTemporaryWrite, and it was not happening. AddRead had it too,
including across the boundary: a permanent read covered only by a temporary WRITE
root died with that root.

Permanent coverage is now looked for first and across every root, so a temporary
cover cannot decide the answer when a permanent one is also present. The same
root held temporarily is promoted in place — its holders' releases become no-ops,
which is what permanence means here. A BROADER temporary root over a narrower
permanent grant cannot be fixed by promotion, so the narrower root is recorded in
its own right; that is the case he specifically called out, and it is the one a
promotion-only fix would have missed.

ONE HOLDER'S UNDO IS ONE HOLDER'S REFERENCE. The comment on releaseTemporaryRead
claimed each undo is called exactly once and that flooring the count at zero
covered the rest. Flooring stops the count going negative; it does nothing about
a holder's second call consuming a DIFFERENT holder's reference. With two
readers, calling the first's undo twice took the count 2 -> 1 -> 0 and removed a
root the second was still using — the exact revocation this refcount exists to
prevent, reached through a duplicate call rather than a sibling's cleanup, and a
deferred cleanup in a retry path is how a caller does that by accident. Every
undo handed out is now wrapped so it drops its own reference at most once, and
the comment says what is actually true.

Both mutation-checked: unwrapping the undo, and ignoring permanence in Add, each
fail the test that covers them. Race detector clean.

Origin-Session: local-abff1c | Claude Code | 3 prompts
Origin-Snapshot: 07900397c62e
Origin-Session: local-13d543 | Claude Code | 5 prompts
Origin-Snapshot: 6b5eb8ba4e5b
Origin-Session: local-c962d7 | Claude Code | 3 prompts
Origin-Snapshot: d2b8f44a9abc

* chore(sandbox): re-request an automated review of this head

No code change. CodeRabbit's verdict on this PR is pinned to ff85dbcc, a commit
that no longer exists on the branch, and a stale changes-requested keeps the PR
at mergeStateStatus=BLOCKED exactly as a human one would. It re-reviews on push
rather than on request, and repeated @coderabbitai triggers over five days
produced nothing, so this is the only lever available without maintainer rights.

The cost is that this dismisses @Vasanthdev2004's and @anandh8x's approvals of
4a7063aa. The tree they approved is byte-identical to this one.

Origin-Session: local-8cd239 | Claude Code | 6 prompts
Origin-Snapshot: 149ab98453ad
Origin-Session: local-13d543 | Claude Code | 5 prompts
Origin-Snapshot: 6b5eb8ba4e5b
Origin-Session: local-c962d7 | Claude Code | 3 prompts
Origin-Snapshot: d2b8f44a9abc

* test(sandbox): a fixture failure fails, a report cannot deadlock, and AddRead's
own branches are covered

All raised by CodeRabbit on the full review. Every one verified before changing
anything.

## The concurrent test could hang instead of failing

failures was buffered at one slot per holder, but each goroutine can report at
BOTH hasReadRoot checks — its own grant going missing while held, and again
after the release gate. Nothing drains the channel until after wg.Wait(), so the
ninth report blocks forever, wg.Done() is never reached and the wait hangs. The
test would deadlock in exactly the case it exists to catch: a refcount bug
severe enough to fail more than eight of eight holders. Buffered at 2*holders.

## A broken fixture reported itself as a pass

TestANarrowerRequestHoldsTheCoveringGrant called t.Skip when os.MkdirAll failed.
The directory is the fixture, not a precondition the environment may reasonably
withhold, and the sibling test above it already uses t.Fatal for the same call.

## Two tests assumed something they never checked

Both tests in scope_extra_read_test.go place their grant under $HOME precisely
because NewScope seeds /tmp and $TMPDIR as write roots — a grant that collapsed
into one of those would be writable for a reason unrelated to the property under
test. $HOME is normally clear of them, but a sandboxed or CI environment that
points HOME at a temp directory turns both tests into something that proves
nothing while still reporting PASS.

homeGrantOutsideTempRoots now builds the fixture for both and skips with the
reason when the assumption does not hold. Verified by running with HOME set to a
temp directory: both tests SKIP naming the covering root, where before they
passed. The path is resolved before comparing, because the temp roots are
resolved and /var is a symlink to /private/var on macOS.

## AddRead's branches were only covered from the write side

Two subtests added: a permanent read under a temporary WRITE root, and a narrow
permanent read under a BROAD temporary read root. Both mutation-checked by
removing the corresponding temporary-coverage guard in AddRead — each fails.

The first mutation I ran did not fail its test, and the finding looked wrong: the
pattern being edited occurs in Add, AddRead and AddTemporaryWrite alike, and a
first-match edit had changed Add. Re-run against the line inside AddRead, it
fails as claimed. Recorded because a mutation that misses its target is
indistinguishable from a test that holds.

## Also asserted: the refcount table is empty, not merely invisible

hasReadRoot answers what a caller can see. An entry left at zero in tempReads is
not visible that way and would still leak one map entry per acquire/release
cycle, so the contention test now checks the terminal state directly.

go test -race ./internal/sandbox/ -count=5: clean.
Unrelated and pre-existing here: TestRunDoctorFormatsRedactedProviderDiagnostics
and TestRunDoctorConnectivityProbesProvider exit 3 in this environment.

Origin-Session: local-8cd239 | Claude Code | 12 prompts
Origin-Snapshot: dd397730a138
Origin-Session: local-13d543 | Claude Code | 5 prompts
Origin-Snapshot: 6b5eb8ba4e5b
Origin-Session: local-c962d7 | Claude Code | 3 prompts
Origin-Snapshot: d2b8f44a9abc

* fix(sandbox): a reader outliving its writer keeps read authority, not write

Reported by @jatmn, and it is a privilege escalation rather than a lifetime bug.

AddTemporaryRead, given a path already covered by a temporary WRITE root, took a
reference on that write root and returned an undo that released it. The lifetime
was then correct — the reader no longer lost its access when the writer cleaned
up — but the capability was wrong. extraRoots feeds validate(), which authorises
WRITES, so once the write holder released first the path stayed in extraRoots for
as long as the read holder lived. A caller that asked only to READ could write
anywhere below the former write root after the write grant had ended.

His diagnosis names the root cause exactly: one write-capability reference was
doing two jobs, tracking a lifetime and standing in for a read-only dependent.

The lifetime is now kept as what the caller actually asked for — a temporary READ
root on its own path. AddTemporaryRead consults permanentWriteRootCoversLocked
rather than writeRootCoversLocked, so a merely temporary write cover no longer
short-circuits it into borrowing write authority.

Mutation-checked: restoring writeRootCoversLocked in that branch compiles and
fails the regression on three separate assertions — the reader keeps no read root
of its own, ReadRoots() loses the path, and the writer's cleanup revokes a live
read grant.

Two things a verification pass turned up that are worth recording:

The comment justifying this design said readRoots "feeds ReadRoots() and nothing
else". That was false — it also feeds ExtraReadRoots(). The security argument
survives, because neither confers write authority and validate() reads
workspaceRoot plus extraRoots, which this path never enters. But that sentence IS
the written justification for the whole choice, so it now lists both and says why
neither grants write.

The "the reader releases first" subtest passes on the unfixed tree: the
escalation needs the WRITER to go first. It is a real no-regression control for
the opposite ordering, not evidence the finding is closed, and it is renamed to
say so rather than padding the count.

Rebased onto ad34dc8d as requested. go test -race ./internal/sandbox/ -count=5:
clean. Pre-existing here and on main, unrelated:
TestRunDoctorFormatsRedactedProviderDiagnostics and
TestRunDoctorConnectivityProbesProvider exit 3 in this environment.

Origin-Session: local-c962d7 | Claude Code | 3 prompts
Origin-Snapshot: d2b8f44a9abc

* fix(sandbox): a nested write holder keeps only the subtree it asked for

Reported by CodeRabbit, and it is the write-side twin of the read-side
escalation @jatmn found — I had disclosed it on this PR as still open, so this
closes it rather than deferring it again.

AddTemporaryWrite, given a path already covered by a broader TEMPORARY write
root, took a reference on that covering root. The lifetime was right and the
capability was wrong, exactly as on the read side: extraRoots is what validate()
authorises writes from, so keeping the covering root alive handed the nested
holder everything its neighbour had asked for.

Reproduced:

  both held           tempWrites={outer:2}                extraRoots=[... outer]
  after outer release tempWrites={outer:1}                extraRoots=[... outer]
  -> the inner-only holder can WRITE outer/sibling/not-mine.txt

A narrower request now records its own root. Each holder gets exactly the
authority it asked for and a lifetime of its own:

  both held           tempWrites={outer:1, inner:1}       extraRoots=[... outer inner]
  after outer release tempWrites={inner:1}                extraRoots=[... inner]
  -> sibling denied, inner retained

Only an EXACT match shares a refcount entry; a narrower one never does, which is
the distinction the escalation turned on. The nesting costs a redundant
extraRoots entry while both holders are live, which is not wrong — the broader
root already permits everything the narrower one does.

Three regressions: the escalation itself (broad holder releases first), the
opposite order as a control, and two holders of the same root sharing one entry.

Mutation: restoring the covering-root borrow compiles and fails the escalation
test on both the sibling and the covering root itself.

Honest note on a second mutation. Removing the permanent-cover early return
compiles and fails NOTHING. That is correct rather than a coverage gap: a
permanently covered request would then record its own temporary root, which is
redundant but harmless, since the permanent root outlives it and already permits
the writes. The early return is a cleanliness guard, not a correctness one, and
I would rather say so than write a test that pretends otherwise.

go test -race ./internal/sandbox/ -count=5: clean. Pre-existing here and on main:
TestRunDoctorFormatsRedactedProviderDiagnostics and
TestRunDoctorConnectivityProbesProvider exit 3 in this environment.

Origin-Session: local-c962d7 | Claude Code | 5 prompts
Origin-Snapshot: 3bb420a0a97b

* test(sandbox): the outside-roots fixture fails loudly instead of skipping

Raised by CodeRabbit, and the helper had already written the argument against
itself: its own comment says "A skip is not a pass" while three of its four exits
were skips.

Every caller of scopeOutsideRoots asserts an AUTHORIZATION boundary — who may
write where, and whose release revokes it. A skipped fixture means that assertion
never ran and the build reported green anyway, which is precisely the failure the
helper was written to close. It closed it for one case (a home directory that
resolves nowhere on Windows) and left the others open.

Three exits become fatal:

  os.UserHomeDir failing        — resolves on every platform Zero supports
  os.MkdirTemp(home) failing    — a home that will not take a temp directory is a
                                  broken machine, not a limited one
  os.MkdirAll under base        — base was just created successfully, so a
                                  subdirectory failing under it is the same

If any fires, the right outcome is a red build somebody fixes, not silent
non-coverage of the sandbox's write rules.

ONE SKIP STAYS, and it is a different kind. When the grant path is already inside
a default write root, the fixture built correctly and the machine is fine — the
grant simply cannot demonstrate anything, because that path is writable before
any temporary grant exists. Running on would assert a permission that was never
in question. The comment now says which kind it is, so the distinction survives
the next reader.

Verified the fatal path fires: forcing the MkdirTemp error compiles and fails
TestANestedWriteHolderKeepsOnlyItsOwnSubtree at the fixture with the reason,
where it previously reported a pass.

go test -race ./internal/sandbox/ -count=3: clean. Pre-existing here and on main:
TestRunDoctorFormatsRedactedProviderDiagnostics and
TestRunDoctorConnectivityProbesProvider exit 3 in this environment.

Origin-Session: local-c962d7 | Claude Code | 13 prompts
Origin-Snapshot: a599377c09e0

* fix(sandbox): least authority for nested readers, and a hermetic grant fixture

@jatmn's three findings, taken as one pass over the grant construct rather than
as three local patches — his point that the same invariant kept not being
reapplied to the symmetric branch was fair, and this is the symmetric branch.

## The nested reader kept its neighbour's tree

AddTemporaryRead had exactly the shape AddTemporaryWrite was fixed for one
commit earlier: reader B asking for outer/inner took a reference on reader A's
covering root. Lifetime correct, capability wrong. readRoots feeds validateRead
and the native sandbox profile, so once A released, B could read outer/sibling.

Reproduced: tempReads={outer:2}; after A releases, the inner-only reader reads a
sibling it never asked for.

The rule is now the same on both sides: permanent coverage decides first, an
EXACT temporary root shares one refcount, and a narrower request records its own
root. Mutation: restoring the borrow fires both escalation assertions.

## Slice order no longer decides permanent coverage

permanentReadRootCoversLocked asks the whole question over both lists before any
insertion order can answer it, and AddRead now uses the same predicate instead of
its own scan, so the two cannot drift.

Honest note: with the borrow gone, this mutation is no longer catchable. Making
the scan return at the first covering entry compiles and fails nothing, because
AddTemporaryRead no longer borrows anything — it either no-ops or records its own
root, and neither escalates. The guard now prevents a redundant refcount entry
rather than an escalation. Verified by measurement, not assumed.

## A test that asserted the defect

TestANarrowerRequestHoldsTheCoveringGrant proved a lifetime property by
asserting the BROAD root was still present after the broad holder released —
which is the escalation written down as a requirement. It is replaced by tests
that assert lifetime and authority separately: the requested path stays readable,
the sibling and the covering root do not, plus the reverse release order,
exact-root sharing, and a permanent narrow root not borrowing a broad temporary
one.

## The obsolete predicate that broke CI

writeRootCoversLocked lost its last caller when the three coverage decisions
moved to permanentWriteRootCoversLocked and exact-root logic. make lint-static
reproduced the Windows `unused` failure locally. Removed rather than suppressed.
This is on me: I did not run lint-static before the previous push, which is a
required step, and CI caught what I should have.

## Fixtures out of the real home

scopeOutsideRoots and homeGrantOutsideTempRoots both built fixtures under
os.UserHomeDir(). Fixed at the cause jatmn identified rather than by relaxing the
setup checks: NewScope seeds the system temp directory as a permanent write root,
which is what made a t.TempDir() grant vacuous and sent the fixture to HOME in the
first place. Both helpers now build the Scope DIRECTLY — the pattern
scopeOutsideDefaults and temporaryWriteFixture already used — so no defaults are
seeded, t.TempDir() is outside by construction, and nothing touches HOME.

Each helper now asserts the NEGATIVE against its own scope before returning: the
target must be unwritable, and unreadable where relevant, before any grant
exists. That is stronger than the old check against defaultTempWriteRoots(),
which asked about a scope these tests no longer build.

Verified with HOME readable but not writable (caches redirected): the sandbox
suite passes. The one remaining failure there is
TestSeatbeltEnforcesExtraWriteRoots in seatbelt_integration_darwin_test.go, which
is NOT in this PR's diff — pre-existing, same shape, and worth its own issue
rather than riding along here.

gofmt, go vet, lint-static (0 issues), go test -race ./internal/sandbox/ -count=5:
all clean. Pre-existing here and on main:
TestRunDoctorFormatsRedactedProviderDiagnostics and
TestRunDoctorConnectivityProbesProvider exit 3 in this environment.

Origin-Session: local-c962d7 | Claude Code | 14 prompts
Origin-Snapshot: a599377c09e0

* refactor(sandbox): drop the child-grant accessors this PR does not consume

Reported by @jatmn.

ExtraRoots and ExtraReadRoots have no production caller anywhere in the tree —
only their own definitions and this PR's tests. They arrived with the initial
split because the extraction was made at the file level: the temporary-root
implementation came across while the CLI and runner wiring that consumes them
stayed behind in #829's child read-grant propagation area. That leaves unrelated
production surface in a branch scoped to refcounting.

Both accessors are removed, along with the getter-only propagation assertion that
existed to exercise them. The refcount regressions that used ExtraRoots only to
inspect beyond-workspace write state now use a test-only helper over Roots(),
because that inspection is a test concern rather than an API this branch needs.

Nothing about the refcount, permanence, least-authority, atomic-release or
hermetic-fixture work changes, and no child propagation is wired here — that
would widen the scope rather than fix the extraction boundary.

One comment needed correcting rather than deleting. The justification for keeping
a reader's lifetime as a READ root named the lists readRoots feeds, and that
sentence has now been wrong in both directions: an earlier draft said "ReadRoots
and nothing else" while ExtraReadRoots also read it, and the accessor it was
corrected to mention is the one being removed here. It now describes what ships
and says why the list has to stay exhaustive as accessors change.

Rebased onto 6fe0d1ed, 0 behind. lint-static: 0 issues.
go test -race ./internal/sandbox/ -count=3: clean. Pre-existing here and on main:
TestRunDoctorFormatsRedactedProviderDiagnostics and
TestRunDoctorConnectivityProbesProvider exit 3 in this environment.
2026-08-26 13:21:25 +08:00
KRATOS 1649008506 fix(tools): an edit keeps the reads it did not disturb (#908)
* fix(tools): an edit keeps the reads it did not disturb

Split out of #829 as an independent fix, per @Vasanthdev2004's review asking for
the small self-contained pieces to arrive separately.

A successful edit re-baselined the file through RecordHash, which drops EVERY
recorded read range. A file read in three pieces therefore lost all three to a
single two-line edit, and the next six edits into regions that had already been
read were refused as unseen — the model had to re-read what it had just read in
order to keep working.

RecordEdit replaces that on the edit path. Because the edit is ours, the changed
span is known exactly, so ranges the edit did not touch are carried across,
ranges after it are shifted by the line delta, and a range the edit spans is
split around it.

An EXTERNAL change still drops everything, deliberately: once a formatter has
rewritten the file we no longer know which line holds what was read, and the
conservative drop is the only honest answer there.

"Seen whole" is now decided in ONE place. SeenWhole derives the answer from the
ranges because the raw flag is only ever set by a single covering read, but
RecordEdit branched on the flag — so a file read whole in two chunks reported
seen-whole true, then false after one edit, putting back the write_file refusal
the derived answer exists to prevent.

countLines now reports the count a READER would give. strings.Split leaves an
empty final element for content ending in a newline, so it counted 4 lines for a
3-line file; since SeenWhole asks whether the ranges cover 1..total, that
inflated total made full coverage unreachable for a file just read in full.
Nearly every text file ends in a newline, so this was the common case.

Origin-Session: local-de382f | Claude Code | 9 prompts
Origin-Snapshot: 92ae33c95cd1
Origin-Session: local-c962d7 | Claude Code | 3 prompts
Origin-Snapshot: ae653221fa15

* chore(tools): re-request an automated review of this head

No code change. CodeRabbit's verdict on this PR is pinned to 6f0cd6c4, a commit
that was force-pushed away five days ago, and a stale changes-requested keeps the
PR at mergeStateStatus=BLOCKED exactly as a human one would. It re-reviews on
push rather than on request, and repeated @coderabbitai triggers over five days
produced nothing, so this is the only lever available without maintainer rights.

The cost is that this dismisses @Vasanthdev2004's and @anandh8x's approvals of
0b6fcf54, which is an unhappy trade for an empty commit — the tree they approved
is byte-identical to this one.

Origin-Session: local-8cd239 | Claude Code | 6 prompts
Origin-Snapshot: 149ab98453ad
Origin-Session: local-c962d7 | Claude Code | 3 prompts
Origin-Snapshot: ae653221fa15

* perf(tools): derive the changed-line span from bytes, not from per-line slices

Reported by @jatmn.

changedLineSpan split both file versions into []string, so it allocated one
string header per line for EACH version — work that scales with the size of the
file rather than the size of the edit. countLines then split a third time just to
take a length. All of it ran before RecordEdit had checked whether there was a
tracked observation to update at all.

He measured 32,036,640 bytes for two 2 MB versions of a million short lines, and
extrapolated roughly 1.6 GB of headers for a 100 MB file before counting the
content, the updated bytes and the hashes already live beside them.

The memory was not the sharp end. edit_file writes the updated bytes BEFORE
calling RecordEdit, so an out-of-memory kill lands after the user's file has
changed and before the tracker baseline catches up: the file and the record of it
disagree, and nothing says so. His instruction was to fix the allocation source
rather than add a size guard, and that a guard — if used at all — would have to
reject before os.WriteFile rather than during post-write re-baselining. Taken as
written: there is no new limit, and the same edits are accepted as before.

Two identical byte prefixes share every line that ends inside them, so counting
newlines in the common prefix counts the unchanged leading lines directly. Same
for the suffix, with one rule that the line-array version got for free: the line
a shared tail STARTS in is only shared when the tail begins on a line boundary in
both versions, or its opening bytes differ and it is a changed line that merely
ends the same way.

Measured on his own scenario, two 2 MB versions:

  changedLineSpan   32,036,640 bytes  ->  0 allocations
  countLines        per-line slice    ->  0 allocations

## The equivalence is tested, not asserted

A rewrite of a diff primitive is only safe if it is behaviour-preserving, so the
replaced implementation is kept in the test file as an oracle and both are asked
the same questions. That caught 206 mismatches in my first attempt at the byte
scan — empty versions, and trailing lines that are only partially shared. Neither
appears in the hand-written cases, and I would have shipped it.

Two regressions: the oracle comparison, and an absolute allocation ceiling rather
than a ratio, so a future rewrite that reintroduces per-line slices fails rather
than merely getting slower. Mutation: restoring strings.Split in either function
allocates 3.2 MB and 3.6 MB against a 4 KB bound.

splitLinesForTracking lost its last caller and is removed — the same obsolete-
helper lint failure that broke Windows CI on #911, caught here by running
make lint-static before pushing rather than after. 0 issues.

Rebased onto ad34dc8d, 0 behind. go test -race ./internal/tools/: clean.
Pre-existing here and on main: TestRunDoctorFormatsRedactedProviderDiagnostics
and TestRunDoctorConnectivityProbesProvider exit 3 in this environment.

Origin-Session: local-c962d7 | Claude Code | 17 prompts
Origin-Snapshot: a599377c09e0
2026-08-26 13:20:48 +08:00
KRATOS 6fe0d1edcd fix(lockutil): use kernel-held advisory locks (#950)
* fix(lockutil): use kernel-held advisory locks

* fix(lockutil): harden cross-platform lock files

* test(lockutil): require one concurrent holder
2026-08-25 21:20:18 +08:00
e59bf78cab fix(sandbox): Make WSL backend tests deterministic. (#948)
Co-authored-by: euxaristia <[email protected]>
Co-authored-by: cairn-code <[email protected]>
2026-08-25 21:19:31 +08:00
Kevin CodexandClaude Fable 5 6645f9744b perf(agent): cheap compaction — dedicated summarizer model, input cap… (#532)
* perf(agent): cheap compaction — dedicated summarizer model, input caps, free reactive prune, window re-derive

Compaction summarization was the single most expensive recurring event
in a long run: the full elided middle — verbatim tool bodies included —
went to the session's main model at full price (a ~100k-token call near
a 200k window, recursively multiplied on overflow).

Four changes, each independently safe:

- Dedicated summarizer model. Compaction summaries route to a cheap
  model resolved as ZERO_COMPACTION_MODEL > preferences.compactionModel
  > a curated default on official endpoints (Haiku 4.5 / Gemini
  Flash-Lite / gpt-4.1-mini; custom-endpoint catalogs are unknowable so
  they get none). Built lazily on the first paid compaction; ANY
  failure permanently falls back to the main provider for the run, so a
  misconfigured summarizer can never break compaction. "main" at any
  level forces the old behavior.

- Summarizer input caps. renderTranscript clamps tool-result bodies to
  2,000 chars and tool-call arguments to 500 (head+tail, rune-safe) —
  the summary needs what was done and decided, not raw payloads. Cuts
  summarizer input 5-10x on tool-heavy sessions.

- Free reactive prune. recover() now runs the same prune-first stage as
  the proactive path: a context-limit error is first answered by
  pruning stale tool bodies at zero cost, without consuming the
  one-shot reactive budget; the paid summarizer only runs when a later
  error finds nothing left to prune.

- Window re-derivation on model switch. Options.ContextWindowFor
  (wired to the model registry in exec and the TUI) lets the loop
  re-derive the compaction threshold after escalate_model, closing the
  documented KNOWN LIMITATION where a switch to a smaller-window model
  could overflow and a larger one over-compacted.

## Rebase notes (2026-08-25, onto 2749447c)

Rebased across 214 commits. main now summarizes through the context planner
(#879) and switches models via turn sessions, so the agent side was reworked
rather than merged textually:

- Summarizer selection lives on compactionState and is resolved lazily per
  main model through summarizeWithPlanner (planner + trace preserved); a
  failing dedicated summarizer is marked broken and the same slice is
  retried once on the main provider.
- Options.Summarizer now receives the main model in force:
  `func(ctx, mainModelID string) (Provider, error)`, and may return
  (nil, nil) when no dedicated model applies. After a successful
  escalate_model switch the loop calls compactor.switchModel, which
  re-derives the window and re-resolves the summarizer — a run that started
  on the cheap model and escalated regains cheap summaries (review P2).
- recover()'s free prune stage runs before the one-shot reactive guard
  (pruning is free and cannot loop) and mirrors maybeCompact's two prune
  passes; the free retry is taken only when the pruned history fits the
  window the provider enforces (with a 5% estimate margin), otherwise the
  same call falls through to the paid summarizer. Gating on the 70%
  proactive threshold would have sent histories in the 70-95% band to the
  summarizer unnecessarily (review P1 and the follow-up question).
- ContextWindowFor is wired through the turn-session switch path; the
  planner's window is updated alongside the compactor's.

Tests: the PR's suite is ported to the new seams, plus insufficient-prune
fall-through, free prune after the budget is spent, window-vs-threshold
gating, and summarizer re-resolution across a model switch.

* fix(agent): branch on a dedicated-summarizer flag instead of comparing providers

Review follow-ups for #532:

- summarizeProvider now reports whether the dedicated summarizer was
  selected and summarizeClosure branches on that flag; the previous
  `chosen == provider` compared Provider interface values, which panics
  when the dynamic type is not comparable. Regression test uses a provider
  with a slice field for both the summarizer and the main provider.
- CompactionModelID's curated-default test covers the official OpenAI
  profile alongside Anthropic and Google.

Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_012JqvAf7zsP1qYFVZayKXHq

---------

Co-authored-by: Claude Fable 5 <[email protected]>
2026-08-25 21:18:04 +08:00
108 changed files with 6636 additions and 1458 deletions

No files matched your search

+7 -3
View File
@@ -371,9 +371,13 @@ go run ./cmd/zero-perf-bench
Experimental: `ZERO_OPENAI_TURN_SESSION=1` enables the optimized OpenAI turn
session (background connection prewarm + request-prefix telemetry) for headless
`zero exec` runs against official OpenAI profiles. Off by default; `0`/`false`
disable. A/B-benchmark it by running the same `zero-perf-bench` suite with the
variable unset and set.
`zero exec` runs against official OpenAI profiles. Off by default; `0`, `false`,
or `off` disable it. A/B-benchmark it by running the same `zero-perf-bench` suite
with the variable unset and set.
Native ChatGPT Responses sessions are enabled by default. Set
`ZERO_CHATGPT_TURN_SESSION=0`, `false`, or `off` to restore stateless HTTP/SSE
transport.
### Code Quality and Security Checks
+1 -1
View File
@@ -4,7 +4,7 @@ go 1.26.6
require (
charm.land/bubbles/v2 v2.1.1
charm.land/bubbletea/v2 v2.0.8
charm.land/bubbletea/v2 v2.0.9
charm.land/lipgloss/v2 v2.0.5
github.com/Microsoft/go-winio v0.6.2
github.com/alecthomas/chroma/v2 v2.27.0
+2 -4
View File
@@ -1,7 +1,7 @@
charm.land/bubbles/v2 v2.1.1 h1:7r55WzBxpo/R3z98hGmY7KKPd3ET6vsf0Fb9sDHOV60=
charm.land/bubbles/v2 v2.1.1/go.mod h1:GE6M31gaWZVXzGw73OeuTTgy4lX+OtkH0E5ymnNsHxo=
charm.land/bubbletea/v2 v2.0.8 h1:SxTJMhCAI3lbPmy4SgX5LWZ24AdINr4I6UEqzZvYJuY=
charm.land/bubbletea/v2 v2.0.8/go.mod h1:2SkdgoTXluXJHOUwAoRlRXF/28vklb1rFl6GcgV1/ss=
charm.land/bubbletea/v2 v2.0.9 h1:DpJCMWKgzQK8SJv4zbKKFHAI10ymWy/evClPFk0k0f8=
charm.land/bubbletea/v2 v2.0.9/go.mod h1:2SkdgoTXluXJHOUwAoRlRXF/28vklb1rFl6GcgV1/ss=
charm.land/lipgloss/v2 v2.0.5 h1:kbNxgeeUOYv5J0YdpxFjfvf3dFvqH8Aci4zB6xqFtrY=
charm.land/lipgloss/v2 v2.0.5/go.mod h1:9oqhxt4yxIMe6q5A4kHr44DremZk7J9UNh74GlWa5nc=
github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
@@ -64,8 +64,6 @@ github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavM
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
golang.org/x/exp v0.0.0-20260611194520-c48552f49976 h1:X8Hz2ImujgbmetVuW+w2YkyZChE3cBpZi2P158rTG9M=
golang.org/x/exp v0.0.0-20260611194520-c48552f49976/go.mod h1:vnf4pv9iKZXY58sQE1L86zmNWJ4159e1RkcWiLCkeEY=
golang.org/x/image v0.44.0 h1:+tDekMZED9+LrtB3G5xzRggpVh9CARjZqROla3R3R+I=
golang.org/x/image v0.44.0/go.mod h1:V8K3KE9KKKE+pLpQDOeN18w9oacNSvy1tDOirTu4xtY=
golang.org/x/image v0.45.0 h1:FMb1nTbH5H9vF55SriQHgFw5GnNL9Jg6L25BwXKzhB0=
golang.org/x/image v0.45.0/go.mod h1:n62x/7RqlwXDvGsSU4u6IUTUf6KghUZ9Bt7cG/T9Fx4=
golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek=
+6
View File
@@ -146,11 +146,17 @@ func estimateToolDefTokens(tools []zeroruntime.ToolDefinition) int {
for _, tool := range tools {
total += ApproxTextTokens(tool.Name)
total += ApproxTextTokens(tool.Description)
total += ApproxTextTokens(string(tool.Type))
if len(tool.Parameters) > 0 {
if encoded, err := json.Marshal(tool.Parameters); err == nil {
total += ApproxTextTokens(string(encoded))
}
}
if tool.Format != nil {
total += ApproxTextTokens(tool.Format.Type)
total += ApproxTextTokens(tool.Format.Syntax)
total += ApproxTextTokens(tool.Format.Definition)
}
total += 4 // per-tool overhead
}
return total
+4 -1
View File
@@ -126,8 +126,11 @@ func capCompactionBrief(brief string) (string, bool) {
func compactionToolCallLine(call zeroruntime.ToolCall) string {
detail := ""
if call.Freeform && call.Name == "apply_patch" {
detail = strings.Join(patchPaths(call.Arguments), ", ")
}
var arguments map[string]any
if json.Unmarshal([]byte(strings.TrimSpace(call.Arguments)), &arguments) == nil {
if detail == "" && json.Unmarshal([]byte(strings.TrimSpace(call.Arguments)), &arguments) == nil {
for _, key := range []string{"file_path", "path", "url", "query", "pattern", "prompt", "description", "question"} {
if value, ok := arguments[key].(string); ok && strings.TrimSpace(value) != "" {
detail = value
+2 -1
View File
@@ -73,6 +73,7 @@ func TestProjectCompactionInputRetainsAskUserExchangeAndSemanticArguments(t *tes
{Role: zeroruntime.MessageRoleTool, ToolCallID: "ask", Content: "Which database should remain?: Postgres only; do not touch MySQL."},
{Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{
{ID: "patch", Name: "apply_patch", Arguments: `{"patch":"*** Begin Patch\n*** Update File: internal/db.go\n*** End Patch"}`},
{ID: "raw-patch", Name: "apply_patch", Freeform: true, Arguments: "*** Begin Patch\n*** Update File: internal/api.go\n*** Add File: internal/config.go\n*** End Patch\n"},
{ID: "fetch", Name: "web_fetch", Arguments: `{"url":"https://example.com/spec"}`},
{ID: "task", Name: "task", Arguments: `{"prompt":"Inspect the Windows implementation"}`},
}},
@@ -80,7 +81,7 @@ func TestProjectCompactionInputRetainsAskUserExchangeAndSemanticArguments(t *tes
}
content := projectCompactionInput(messages).messages[0].Content
for _, want := range []string{"Which database should remain?", "Postgres only; do not touch MySQL", "internal/db.go", "https://example.com/spec", "Inspect the Windows implementation"} {
for _, want := range []string{"Which database should remain?", "Postgres only; do not touch MySQL", "internal/db.go", "internal/api.go", "internal/config.go", "https://example.com/spec", "Inspect the Windows implementation"} {
if !strings.Contains(content, want) {
t.Fatalf("projection missing %q: %q", want, content)
}
+21
View File
@@ -275,6 +275,27 @@ func TestEstimateToolDefTokensCountsDefinitions(t *testing.T) {
if estimateToolDefTokens(one) <= 0 {
t.Fatal("a tool definition (name + description + schema) should estimate > 0 tokens")
}
typed := append([]zeroruntime.ToolDefinition(nil), one...)
typed[0].Type = zeroruntime.ToolDefinitionFreeform
if estimateToolDefTokens(typed) <= estimateToolDefTokens(one) {
t.Fatal("a non-empty tool type must increase the estimated request cost")
}
for _, test := range []struct {
name string
format zeroruntime.ToolDefinitionFormat
}{
{name: "type", format: zeroruntime.ToolDefinitionFormat{Type: "grammar"}},
{name: "syntax", format: zeroruntime.ToolDefinitionFormat{Syntax: "lark"}},
{name: "definition", format: zeroruntime.ToolDefinitionFormat{Definition: "start: WORD"}},
} {
t.Run("format "+test.name, func(t *testing.T) {
formatted := append([]zeroruntime.ToolDefinition(nil), one...)
formatted[0].Format = &test.format
if estimateToolDefTokens(formatted) <= estimateToolDefTokens(one) {
t.Fatalf("a non-empty format %s must increase the estimated request cost", test.name)
}
})
}
two := append(append([]zeroruntime.ToolDefinition{}, one...), zeroruntime.ToolDefinition{
Name: "write_file",
Description: "Write contents to a file in the workspace.",
+4
View File
@@ -102,6 +102,10 @@ func copyToolDefinitions(toolDefs []zeroruntime.ToolDefinition) []zeroruntime.To
for index, definition := range toolDefs {
copied[index] = definition
copied[index].Parameters = copySchemaMap(definition.Parameters)
if definition.Format != nil {
format := *definition.Format
copied[index].Format = &format
}
}
return copied
}
+3
View File
@@ -15,6 +15,7 @@ func TestContextPlannerPreservesProviderRequest(t *testing.T) {
}
toolDefs := []zeroruntime.ToolDefinition{{
Name: "read_file", Description: "Read a file",
Format: &zeroruntime.ToolDefinitionFormat{Type: "grammar", Syntax: "lark", Definition: "start: PATH"},
Parameters: map[string]any{
"type": "object",
"required": []string{"path"},
@@ -53,10 +54,12 @@ func TestContextPlannerPreservesProviderRequest(t *testing.T) {
toolDefs[0].Parameters["required"].([]string)[0] = "other"
toolDefs[0].Parameters["properties"].(map[string]any)["path"].(map[string]any)["type"] = "integer"
toolDefs[0].Parameters["anyOf"].([]any)[0].(map[string]any)["type"] = "array"
toolDefs[0].Format.Definition = "changed"
if plan.Request.Messages[1].Content != "inspect this" || plan.Request.Messages[1].Images[0].Data[0] != 1 {
t.Fatalf("planned messages alias caller state: %#v", plan.Request.Messages[1])
}
if plan.Request.Tools[0].Name != "read_file" ||
plan.Request.Tools[0].Format == nil || plan.Request.Tools[0].Format.Definition != "start: PATH" ||
plan.Request.Tools[0].Parameters["required"].([]string)[0] != "path" ||
plan.Request.Tools[0].Parameters["properties"].(map[string]any)["path"].(map[string]any)["type"] != "string" ||
plan.Request.Tools[0].Parameters["anyOf"].([]any)[0].(map[string]any)["type"] != "object" {
+378
View File
@@ -0,0 +1,378 @@
package agent
import (
"context"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"github.com/Gitlawb/zero/internal/sandbox"
"github.com/Gitlawb/zero/internal/tools"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
type customApplyPatchTool struct {
calls int
}
type fixedPathScope struct {
roots []string
}
func (scope fixedPathScope) Roots() []string {
return append([]string(nil), scope.roots...)
}
func newFixedPathScope(t *testing.T, roots ...string) fixedPathScope {
t.Helper()
resolved := make([]string, 0, len(roots))
for _, root := range roots {
absolute, err := filepath.Abs(root)
if err != nil {
t.Fatalf("Abs(%q): %v", root, err)
}
if physical, err := filepath.EvalSymlinks(absolute); err == nil {
absolute = physical
}
resolved = append(resolved, absolute)
}
return fixedPathScope{roots: resolved}
}
func (*customApplyPatchTool) Name() string { return "apply_patch" }
func (*customApplyPatchTool) Description() string { return "custom JSON tool" }
func (*customApplyPatchTool) Parameters() tools.Schema {
return tools.Schema{Type: "object", Properties: map[string]tools.PropertySchema{"value": {Type: "string"}}}
}
func (*customApplyPatchTool) Safety() tools.Safety {
return tools.Safety{Permission: tools.PermissionAllow, SideEffect: tools.SideEffectRead}
}
func (tool *customApplyPatchTool) Run(context.Context, map[string]any) tools.Result {
tool.calls++
return tools.Result{Status: tools.StatusOK}
}
func TestFreeformApplyPatchUsesExistingToolExecutionPath(t *testing.T) {
root := t.TempDir()
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(root, nil))
patch := "*** Begin Patch\n*** Add File: hello.txt\n+hello\n*** End Patch\n"
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1",
Name: "apply_patch",
Arguments: patch,
Freeform: true,
}, PermissionModeUnsafe, Options{Cwd: root, FileTracker: tools.NewFileTracker()})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusOK {
t.Fatalf("freeform apply_patch status = %s: %s", result.Status, result.Output)
}
content, err := os.ReadFile(filepath.Join(root, "hello.txt"))
if err != nil {
t.Fatalf("read patched file: %v", err)
}
if string(content) != "hello\n" {
t.Fatalf("patched content = %q, want hello newline", content)
}
}
func TestFreeformApplyPatchSupportsGrantedExtraRoot(t *testing.T) {
root := t.TempDir()
extra := t.TempDir()
scope, err := sandbox.NewScope(root, []string{extra})
if err != nil {
t.Fatalf("NewScope: %v", err)
}
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(root, scope))
target := filepath.Join(extra, "hello.txt")
patch := "*** Begin Patch\n*** Add File: " + target + "\n+hello\n*** End Patch\n"
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1", Name: "apply_patch", Arguments: patch, Freeform: true,
}, PermissionModeUnsafe, Options{Cwd: root, FileTracker: tools.NewFileTracker()})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusOK {
t.Fatalf("freeform extra-root apply_patch status = %s: %s", result.Status, result.Output)
}
content, err := os.ReadFile(target)
if err != nil {
t.Fatalf("read extra-root patched file: %v", err)
}
if string(content) != "hello\n" {
t.Fatalf("patched content = %q, want hello newline", content)
}
}
func TestFreeformApplyPatchRejectsMixedWorkspaceAndExtraRoots(t *testing.T) {
for _, extraHasMatchingPath := range []bool{false, true} {
name := "extra target missing"
if extraHasMatchingPath {
name = "extra target exists"
}
t.Run(name, func(t *testing.T) {
workspace := t.TempDir()
extra := t.TempDir()
if err := os.MkdirAll(filepath.Join(workspace, "src"), 0o755); err != nil {
t.Fatal(err)
}
workspaceTarget := filepath.Join(workspace, "src", "a.go")
if err := os.WriteFile(workspaceTarget, []byte("old\n"), 0o644); err != nil {
t.Fatal(err)
}
extraTarget := filepath.Join(extra, "src", "a.go")
if extraHasMatchingPath {
if err := os.MkdirAll(filepath.Dir(extraTarget), 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(extraTarget, []byte("old\n"), 0o644); err != nil {
t.Fatal(err)
}
}
scope, err := sandbox.NewScope(workspace, []string{extra})
if err != nil {
t.Fatal(err)
}
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(workspace, scope))
extraCreated := filepath.Join(extra, "b.go")
patch := strings.Join([]string{
"*** Begin Patch",
"*** Update File: src/a.go",
"@@",
"-old",
"+new",
"*** Add File: " + extraCreated,
"+extra",
"*** End Patch",
}, "\n")
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1", Name: "apply_patch", Arguments: patch, Freeform: true,
}, PermissionModeUnsafe, Options{Cwd: workspace, FileTracker: tools.NewFileTracker()})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusError || !strings.Contains(result.Output, "mixes workspace-relative paths with an extra write root") {
t.Fatalf("mixed-root patch result = %s: %s", result.Status, result.Output)
}
if content, err := os.ReadFile(workspaceTarget); err != nil || string(content) != "old\n" {
t.Fatalf("workspace target changed: content=%q err=%v", content, err)
}
if extraHasMatchingPath {
if content, err := os.ReadFile(extraTarget); err != nil || string(content) != "old\n" {
t.Fatalf("extra-root target changed: content=%q err=%v", content, err)
}
}
if _, err := os.Stat(extraCreated); !os.IsNotExist(err) {
t.Fatalf("mixed-root patch created extra file: %v", err)
}
})
}
}
func TestFreeformApplyPatchAcceptsOneSemanticRoot(t *testing.T) {
for _, test := range []struct {
name string
useExtraRoot bool
}{
{name: "workspace"},
{name: "extra root", useExtraRoot: true},
} {
t.Run(test.name, func(t *testing.T) {
workspace := t.TempDir()
extra := t.TempDir()
scope, err := sandbox.NewScope(workspace, []string{extra})
if err != nil {
t.Fatal(err)
}
root := workspace
if test.useExtraRoot {
root = extra
}
if err := os.MkdirAll(filepath.Join(root, "src"), 0o755); err != nil {
t.Fatal(err)
}
target := filepath.Join(root, "src", "a.go")
if err := os.WriteFile(target, []byte("old\n"), 0o644); err != nil {
t.Fatal(err)
}
updatePath := target
if !test.useExtraRoot {
updatePath = filepath.Join("src", "a.go")
}
created := filepath.Join(root, "b.go")
patch := strings.Join([]string{
"*** Begin Patch",
"*** Update File: " + updatePath,
"@@",
"-old",
"+new",
"*** Add File: " + created,
"+created",
"*** End Patch",
}, "\n")
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(workspace, scope))
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1", Name: "apply_patch", Arguments: patch, Freeform: true,
}, PermissionModeUnsafe, Options{Cwd: workspace, FileTracker: tools.NewFileTracker()})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusOK {
t.Fatalf("single-root patch status = %s: %s", result.Status, result.Output)
}
if content, err := os.ReadFile(target); err != nil || string(content) != "new\n" {
t.Fatalf("updated content=%q err=%v", content, err)
}
if content, err := os.ReadFile(created); err != nil || string(content) != "created\n" {
t.Fatalf("created content=%q err=%v", content, err)
}
})
}
}
func TestFreeformApplyPatchRejectsAbsolutePathOutsideGrantedRoots(t *testing.T) {
root := t.TempDir()
granted := t.TempDir()
outside := t.TempDir()
scope := newFixedPathScope(t, root, granted)
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(root, scope))
target := filepath.Join(outside, "blocked.txt")
patch := "*** Begin Patch\n*** Add File: " + target + "\n+blocked\n*** End Patch\n"
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1", Name: "apply_patch", Arguments: patch, Freeform: true,
}, PermissionModeUnsafe, Options{Cwd: root})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusError {
t.Fatalf("ungranted absolute patch status = %s, want error: %s", result.Status, result.Output)
}
if _, err := os.Stat(target); !os.IsNotExist(err) {
t.Fatalf("ungranted absolute patch created target: %v", err)
}
}
func TestPreparedFreeformPatchRejectsIntermediateSymlinkSwap(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("creating directory symlinks requires privileges on Windows")
}
root := t.TempDir()
extra := t.TempDir()
outside := t.TempDir()
intermediate := filepath.Join(extra, "nested")
if err := os.Mkdir(intermediate, 0o755); err != nil {
t.Fatalf("Mkdir: %v", err)
}
scope := newFixedPathScope(t, root, extra)
tool := tools.NewScopedApplyPatchTool(root, scope)
target := filepath.Join(intermediate, "blocked.txt")
patch := "*** Begin Patch\n*** Add File: " + target + "\n+blocked\n*** End Patch\n"
args, err := tools.PrepareFreeformApplyPatchArguments(tool, patch)
if err != nil {
t.Fatalf("PrepareFreeformApplyPatchArguments: %v", err)
}
if err := os.Remove(intermediate); err != nil {
t.Fatalf("Remove intermediate: %v", err)
}
if err := os.Symlink(outside, intermediate); err != nil {
t.Fatalf("Symlink: %v", err)
}
result := tool.Run(context.Background(), args)
if result.Status != tools.StatusError {
t.Fatalf("symlink-swapped patch status = %s, want error: %s", result.Status, result.Output)
}
if _, err := os.Stat(filepath.Join(outside, "blocked.txt")); !os.IsNotExist(err) {
t.Fatalf("symlink-swapped patch wrote outside root: %v", err)
}
}
func TestUnknownFreeformToolFailsClosed(t *testing.T) {
result, err := executeToolCall(context.Background(), tools.NewRegistry(), zeroruntime.ToolCall{
ID: "call-1", Name: "unknown", Arguments: "raw", Freeform: true,
}, PermissionModeUnsafe, Options{})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusError {
t.Fatalf("unknown freeform status = %s, want error", result.Status)
}
}
func TestCustomApplyPatchCannotUseFreeformContract(t *testing.T) {
custom := &customApplyPatchTool{}
registry := tools.NewRegistry()
registry.Register(custom)
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1", Name: "apply_patch", Arguments: "raw patch", Freeform: true,
}, PermissionModeUnsafe, Options{})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusError || !strings.Contains(result.Output, "Unsupported freeform tool call") {
t.Fatalf("custom freeform apply_patch result = %#v, want unsupported error", result)
}
if custom.calls != 0 {
t.Fatalf("custom apply_patch executed %d times", custom.calls)
}
}
func TestDeniedFreeformApplyPatchDoesNotWrite(t *testing.T) {
root := t.TempDir()
registry := tools.NewRegistry()
registry.Register(tools.NewScopedApplyPatchTool(root, nil))
patch := "*** Begin Patch\n*** Add File: denied.txt\n+must not exist\n*** End Patch\n"
result, err := executeToolCall(context.Background(), registry, zeroruntime.ToolCall{
ID: "call-1",
Name: "apply_patch",
Arguments: patch,
Freeform: true,
}, PermissionModeAsk, Options{
Cwd: root,
OnPermissionRequest: func(context.Context, PermissionRequest) (PermissionDecision, error) {
return PermissionDecision{Action: PermissionDecisionDeny, Reason: "test denial"}, nil
},
})
if err != nil {
t.Fatalf("executeToolCall: %v", err)
}
if result.Status != tools.StatusError {
t.Fatalf("denied freeform status = %s, want error", result.Status)
}
if _, err := os.Stat(filepath.Join(root, "denied.txt")); !os.IsNotExist(err) {
t.Fatalf("denied freeform patch changed the workspace: %v", err)
}
}
func TestHistorySafeToolCallsPreservesRawFreeformArguments(t *testing.T) {
raw := "*** Begin Patch\n*** Update File: internal/main.go\n@@\n-old\n+new\n*** End Patch\n"
safe := historySafeToolCalls([]ToolCall{{
ID: "call-1", ProviderCallID: "provider-1", Name: "apply_patch", Arguments: raw, Freeform: true,
}})
if len(safe) != 1 || safe[0].Arguments != raw || safe[0].ProviderCallID != "provider-1" || !safe[0].Freeform {
t.Fatalf("freeform history changed raw call: %#v", safe)
}
}
func TestAbortedToolResultPreservesProviderCallID(t *testing.T) {
messages := appendAbortedToolResults(nil, []ToolCall{{
ID: "call-1", ProviderCallID: "provider-1", Name: "read_file",
}})
if len(messages) != 1 || messages[0].ToolCallID != "call-1" || messages[0].ToolCallProviderID != "provider-1" || !messages[0].IsError {
t.Fatalf("aborted tool result lost provider identity: %#v", messages)
}
}
+8 -9
View File
@@ -33,11 +33,11 @@ const (
// knows before spending more tool turns.
toolOnlyProgressReminderAt = 6
// planReminderTurn is the turn (1-based) by the end of which a multi-step
// task should have called update_plan; if it hasn't, a one-time reminder is
// injected. Set to 3 (not 2) so short, legitimate two-step tasks finish
// without a spurious planning nag.
planReminderTurn = 3
// planReminderTurn and planReminderToolThreshold keep bounded implementation
// work plan-free while still nudging a genuinely sustained run that has not
// established any structured progress. Both thresholds must be met.
planReminderTurn = 7
planReminderToolThreshold = 7
// planToolName is the planning tool the loop watches for by name.
planToolName = "update_plan"
@@ -591,13 +591,12 @@ func (state *guardState) planReminder(turn int) string {
return planStaleReminder(state.toolCallsSincePlanUpdate)
}
// NOT-CALLED reminder: by the end of planReminderTurn the model should have
// called update_plan if it's doing a multi-step task (>=1 other tool call).
// One-shot for the whole run.
// NOT-CALLED reminder: once both the turn and tool-call thresholds are met,
// sustained work should have called update_plan. One-shot for the whole run.
if !state.notCalledReminderSent &&
!state.planEverCalled &&
turn >= planReminderTurn &&
state.totalToolCalls >= 1 {
state.totalToolCalls >= planReminderToolThreshold {
state.notCalledReminderSent = true
return planNotCalledReminder()
}
+43 -8
View File
@@ -257,19 +257,16 @@ func TestRunDoesNotCountDroppedToolCallTurnsAsEmpty(t *testing.T) {
}
func TestRunInjectsPlanNotCalledReminderForMultiStepTask(t *testing.T) {
const expectedThreshold = 7
if planReminderToolThreshold != expectedThreshold {
t.Fatalf("plan reminder tool threshold = %d, want %d", planReminderToolThreshold, expectedThreshold)
}
root := t.TempDir()
writeAgentTestFile(t, root+"/notes.txt", "alpha")
registry := tools.NewRegistry()
registry.Register(tools.NewScopedReadFileTool(root, nil))
provider := &mockProvider{
turns: [][]zeroruntime.StreamEvent{
toolTurn("call-1", "read_file", `{"path":"notes.txt"}`), // turn 1: other tool call
toolTurn("call-2", "read_file", `{"path":"notes.txt"}`), // turn 2: still no update_plan
toolTurn("call-3", "read_file", `{"path":"notes.txt"}`), // turn 3: reminder fires here
textTurn("done"),
},
}
provider := &mockProvider{turns: append(repeatedReadTurns(expectedThreshold), textTurn("done"))}
result, err := Run(context.Background(), "go", provider, Options{
Registry: registry,
@@ -287,6 +284,44 @@ func TestRunInjectsPlanNotCalledReminderForMultiStepTask(t *testing.T) {
}
}
func TestRunDoesNotInjectPlanReminderForBoundedToolSequence(t *testing.T) {
const expectedThreshold = 7
root := t.TempDir()
writeAgentTestFile(t, root+"/notes.txt", "alpha")
registry := tools.NewRegistry()
registry.Register(tools.NewScopedReadFileTool(root, nil))
provider := &mockProvider{turns: append(repeatedReadTurns(expectedThreshold-1), textTurn("done"))}
result, err := Run(context.Background(), "go", provider, Options{Registry: registry, MaxTurns: 12})
if err != nil {
t.Fatal(err)
}
if result.FinalAnswer != "done" {
t.Fatalf("expected final answer, got %q", result.FinalAnswer)
}
if count := countUserMessagesContaining(result.Messages, planNotCalledReminderMarker); count != 0 {
t.Fatalf("bounded tool sequence must not draw a plan reminder, got %d", count)
}
}
func TestPlanReminderWaitsForTurnThresholdAfterEnoughToolCalls(t *testing.T) {
state := guardState{totalToolCalls: 7}
if got := state.planReminder(6); got != "" {
t.Fatalf("planReminder(6) = %q, want no reminder before turn threshold", got)
}
if state.notCalledReminderSent {
t.Fatal("early turn marked the not-called reminder as sent")
}
}
func repeatedReadTurns(count int) [][]zeroruntime.StreamEvent {
turns := make([][]zeroruntime.StreamEvent, 0, count)
for range count {
turns = append(turns, toolTurn("call", "read_file", `{"path":"notes.txt"}`))
}
return turns
}
func TestRunDoesNotInjectPlanReminderForTrivialTask(t *testing.T) {
root := t.TempDir()
writeAgentTestFile(t, root+"/notes.txt", "alpha")
+41 -12
View File
@@ -141,6 +141,7 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
}
options.Trace.Counter(trace.CounterInputTokens, int64(usage.InputTokens))
options.Trace.Counter(trace.CounterCachedInputTokens, int64(usage.CachedInputTokens))
options.Trace.Counter(trace.CounterCacheWriteTokens, int64(usage.CacheWriteTokens))
options.Trace.Counter(trace.CounterOutputTokens, int64(usage.OutputTokens))
}
}
@@ -697,11 +698,12 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options)
turnRequestedModel = toolResult.RequestedModel
}
messages = append(messages, zeroruntime.Message{
Role: zeroruntime.MessageRoleTool,
Content: toolResult.ModelOutput(),
ToolCallID: toolResult.ToolCallID,
IsError: toolResult.Status == tools.StatusError,
ChangedFiles: append([]string(nil), toolResult.ChangedFiles...),
Role: zeroruntime.MessageRoleTool,
Content: toolResult.ModelOutput(),
ToolCallID: toolResult.ToolCallID,
ToolCallProviderID: call.ProviderCallID,
IsError: toolResult.Status == tools.StatusError,
ChangedFiles: append([]string(nil), toolResult.ChangedFiles...),
})
// Images ride a following USER message rather than the tool result
// above. Every provider drops images on a tool-role message —
@@ -1045,6 +1047,9 @@ func historySafeToolCalls(calls []ToolCall) []ToolCall {
safe := make([]ToolCall, len(calls))
for i, call := range calls {
safe[i] = call
if call.Freeform {
continue
}
args := strings.TrimSpace(call.Arguments)
switch {
case args == "":
@@ -1128,8 +1133,28 @@ func decodeToolArguments(arguments string, v any) error {
}
func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCall, permissionMode PermissionMode, options Options) (ToolResult, error) {
tool, toolFound := registry.Get(call.Name)
args := map[string]any{}
if call.Arguments != "" {
if call.Freeform {
if !toolFound || !tools.IsBuiltInApplyPatch(tool) {
return ToolResult{
ToolCallID: call.ID,
Name: call.Name,
Status: tools.StatusError,
Output: "Error: Unsupported freeform tool call for " + call.Name + ".",
}, nil
}
prepared, err := tools.PrepareFreeformApplyPatchArguments(tool, call.Arguments)
if err != nil {
return ToolResult{
ToolCallID: call.ID,
Name: call.Name,
Status: tools.StatusError,
Output: "Error: Invalid freeform apply_patch input: " + err.Error(),
}, nil
}
args = prepared
} else if call.Arguments != "" {
if err := decodeToolArguments(call.Arguments, &args); err != nil {
return ToolResult{
ToolCallID: call.ID,
@@ -1157,7 +1182,6 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
DenialReason: DenialFiltered,
}, nil
}
tool, toolFound := registry.Get(call.Name)
if (permissionMode == PermissionModeSpecDraft || permissionMode == PermissionModePlan) && toolFound && !ToolAdvertised(tool, permissionMode) {
modeName := string(permissionMode)
return ToolResult{
@@ -3220,11 +3244,15 @@ func cachedRuntimeToolDefinition(defCache map[string]zeroruntime.ToolDefinition,
// runtimeToolDefinition renders a tool's advertised definition (name, description,
// JSON-schema parameters) as sent to the provider.
func runtimeToolDefinition(tool tools.Tool) zeroruntime.ToolDefinition {
return zeroruntime.ToolDefinition{
definition := zeroruntime.ToolDefinition{
Name: tool.Name(),
Description: tool.Description(),
Parameters: schemaToRuntimeMap(tool.Parameters()),
}
if tools.IsBuiltInApplyPatch(tool) {
definition.Type = zeroruntime.ToolDefinitionFreeform
}
return definition
}
func ToolVisible(tool tools.Tool, permissionMode PermissionMode, enabledTools []string, disabledTools []string) bool {
@@ -3379,10 +3407,11 @@ func loadedToolsFromResult(meta map[string]string) []string {
func appendAbortedToolResults(messages []Message, remaining []ToolCall) []Message {
for _, call := range remaining {
messages = append(messages, zeroruntime.Message{
Role: zeroruntime.MessageRoleTool,
Content: abortedToolResultNotice,
ToolCallID: call.ID,
IsError: true,
Role: zeroruntime.MessageRoleTool,
Content: abortedToolResultNotice,
ToolCallID: call.ID,
ToolCallProviderID: call.ProviderCallID,
IsError: true,
})
}
return messages
+4 -1
View File
@@ -4096,7 +4096,7 @@ func TestRunDoesNotFlagCleanToolOutput(t *testing.T) {
// caller's callback AND stamps token counters, and the run completes.
func TestRunTracingWrapperStampsUsage(t *testing.T) {
provider := &mockProvider{turns: [][]zeroruntime.StreamEvent{{
{Type: zeroruntime.StreamEventUsage, Usage: zeroruntime.Usage{InputTokens: 100, CachedInputTokens: 20, OutputTokens: 40}},
{Type: zeroruntime.StreamEventUsage, Usage: zeroruntime.Usage{PromptTokens: 100, CachedInputTokens: 20, CacheWriteTokens: 10, CompletionTokens: 40}},
{Type: zeroruntime.StreamEventText, Content: "done"},
{Type: zeroruntime.StreamEventDone},
}}}
@@ -4135,6 +4135,9 @@ func TestRunTracingWrapperStampsUsage(t *testing.T) {
if got := tr.Counter(trace.CounterCachedInputTokens); got != 20 {
t.Fatalf("cached input token counter = %d, want 20", got)
}
if got := tr.Counter(trace.CounterCacheWriteTokens); got != 10 {
t.Fatalf("cache-write token counter = %d, want 10", got)
}
if got := tr.Counter(trace.CounterOutputTokens); got != 40 {
t.Fatalf("output token counter = %d, want 40", got)
}
+13
View File
@@ -31,6 +31,19 @@ func (t countingSchemaTool) Run(_ context.Context, _ map[string]any) tools.Resul
return tools.Result{Status: tools.StatusOK}
}
func TestRuntimeToolDefinitionMarksOnlyBuiltInApplyPatchFreeform(t *testing.T) {
builtIn := runtimeToolDefinition(tools.NewScopedApplyPatchTool(t.TempDir(), nil))
if builtIn.Type != zeroruntime.ToolDefinitionFreeform {
t.Fatalf("built-in apply_patch type = %q, want freeform", builtIn.Type)
}
calls := 0
custom := runtimeToolDefinition(countingSchemaTool{name: "apply_patch", calls: &calls})
if custom.Type != "" {
t.Fatalf("custom apply_patch type = %q, want legacy function shape", custom.Type)
}
}
// The cache renders a tool's schema once and reuses it across calls, and its
// output is identical to the uncached path.
func TestPartitionToolsCacheRendersOnceAndMatchesUncached(t *testing.T) {
+6
View File
@@ -58,12 +58,18 @@ func toolSubstrings(exposed []zeroruntime.ToolDefinition) (toolsSubstr, schemaSu
for _, def := range exposed {
writeLengthPrefixed(&toolsSB, def.Name)
writeLengthPrefixed(&toolsSB, def.Description)
writeLengthPrefixed(&toolsSB, string(def.Type))
// Parameters is rendered to a canonical JSON string per tool. The
// schema-render cache in internal/agent (used by partitionToolsCached)
// guarantees the same render is byte-identical across turns for the
// same tool, so this substring is a stable hash input.
writeLengthPrefixed(&schemaSB, def.Name)
writeLengthPrefixed(&schemaSB, canonicalSchemaString(def.Parameters))
if def.Format != nil {
writeLengthPrefixed(&schemaSB, def.Format.Type)
writeLengthPrefixed(&schemaSB, def.Format.Syntax)
writeLengthPrefixed(&schemaSB, def.Format.Definition)
}
}
return toolsSB.String(), schemaSB.String()
}
+32
View File
@@ -73,6 +73,38 @@ func TestComputePrefixFingerprintSchemaChangeChangesSchemaHash(t *testing.T) {
}
}
func TestComputePrefixFingerprintIncludesToolTypeAndFormat(t *testing.T) {
opts := Options{Cwd: t.TempDir(), SystemPrompt: "core"}
base := zeroruntime.ToolDefinition{Name: "apply_patch", Description: "Patch files"}
baseFingerprint := ComputePrefixFingerprint(opts, []zeroruntime.ToolDefinition{base})
typed := base
typed.Type = zeroruntime.ToolDefinitionFreeform
if got := ComputePrefixFingerprint(opts, []zeroruntime.ToolDefinition{typed}); got.ToolsHash == baseFingerprint.ToolsHash {
t.Fatal("ToolsHash must change when tool Type changes")
}
formatCases := []struct {
name string
format zeroruntime.ToolDefinitionFormat
}{
{name: "type", format: zeroruntime.ToolDefinitionFormat{Type: "grammar"}},
{name: "syntax", format: zeroruntime.ToolDefinitionFormat{Syntax: "lark"}},
{name: "definition", format: zeroruntime.ToolDefinitionFormat{Definition: "start: PATCH"}},
{name: "empty", format: zeroruntime.ToolDefinitionFormat{}},
}
for _, test := range formatCases {
t.Run(test.name, func(t *testing.T) {
withFormat := base
withFormat.Format = &test.format
got := ComputePrefixFingerprint(opts, []zeroruntime.ToolDefinition{withFormat})
if got.SchemaHash == baseFingerprint.SchemaHash {
t.Fatalf("SchemaHash must change for %s format input", test.name)
}
})
}
}
// TestComputePrefixFingerprintSystemPromptChangeChangesBaseHash asserts the
// most common drift path: a system-prompt edit (e.g. an agent.md update, a
// model addendum change) must move the base-instructions hash, which moves
+19 -18
View File
@@ -30,13 +30,10 @@ work.
read-before-edit discipline: inspect the target file and nearby callers,
tests, or config before you modify behavior. Never edit a file you have not
read.
2. **Plan.** Use update_plan for implementation or investigation that has at
least three meaningful dependent steps, then keep it current at milestones.
Mark completed work and the next in-progress step together when practical;
do not spend separate turns updating it after every file or command. Keep at
most one item in_progress. Skip plans for simple lookups, explanations, code
navigation, and other short tasks. Never create a plan after the work is
already complete merely to describe what you did.
2. **Plan.** Use update_plan for multi-component, sequencing-uncertain, or long
work. Skip it for bounded changes in one component. Update only when the
overall phase changes, keep at most one item in_progress, and never create a
plan after the work is complete.
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
@@ -55,9 +52,9 @@ work.
edit_file, apply_patch, write_file - over shelling out to
cat/sed/awk/python for file operations.
They are safer, reviewable, and produce clean diffs.
- Prefer read_minified_file when initially exploring source code; it preserves
code while removing comments and redundant whitespace. Use read_file when
exact text, comments, or line numbers are needed.
- Use read_minified_file only to explore large or unfamiliar source. Use
read_file for small or likely edit targets and exact text, comments, or line
numbers. Do not read both unless a new need for exact content appears.
- 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.
@@ -76,21 +73,24 @@ work.
- Solve the problem as posed, not a more general version of it. Add no
speculative abstraction, configurability, or handling for cases that cannot
occur, and nothing the user did not ask for. A small diff can still be
over-built; if a 200-line solution could be 50, rewrite it.
over-built; if a 200-line solution could be 50, rewrite it. Once the required
behavior passes validation, stop instead of adding optional cleanup.
- Preserve behavior you were not asked to change. Do not delete or rewrite code
you did not author unless the task requires it; if you must, say so.
## Testing gate (mandatory)
- After any change to code, verify after edits by running the project's
validators before you summarize or commit: tests, type-checks, linters, and/or
the build, as appropriate. Scope them to the change while iterating; reserve
full-suite runs for milestones.
- If you are unsure which validators apply, search the repo (Makefile, package
manifests, CI config) to find them.
- After code changes, run the project's documented appropriate validators.
Scope iteration checks, combine compatible commands when diagnostics stay
useful, and run the final set once after the last edit. Rerun only after a
change or failure.
- If unsure which entry points apply, inspect the Makefile, manifests, or CI;
do not invent a parallel flow.
- Never claim a task is done, and never commit, while validators are failing. If
they fail, fix the cause and rerun; do not paper over it. If you could not run
a validator, say so explicitly rather than implying success.
- A successful native patch confirms the edit. Reread only when the next change
needs exact content or the result was ambiguous.
## Tool use
@@ -148,6 +148,7 @@ work.
- Default to concise, skimmable output. Lead with the answer or the result.
- Use GitHub-flavored Markdown: headings to structure longer replies, fenced
code blocks for code, and `inline code` for file paths, commands, symbols, and
short snippets. Reference code as `file:line` so it is clickable.
short snippets. Reference code with file paths and include line numbers when
already known; do not look them up solely for the final summary.
- Report outcomes faithfully: if tests failed, show it; if a step was skipped,
say so; when something is done and verified, state it plainly without hedging.
+5 -6
View File
@@ -60,12 +60,11 @@ const openAIPromptAddendum = `<model_guidance>
longer answers, fenced code blocks for code, and ` + "`inline code`" + ` for paths,
commands, and symbols.
- Strongly prefer the native file tools (read_file, list_directory, grep, glob,
edit_file, apply_patch, write_file) over shelling out to cat/sed/awk/python
for file work. read_file, edit_file and write_file take one file per call;
a single apply_patch may span several files when they form one coherent
change. Do not batch file writes into a script. Independent tool calls
(several reads, several edit_file calls to different files) belong in the
same turn.
edit_file, apply_patch, write_file) over shelling out to cat/sed/awk/python for
file work. Use them directly; never run apply_patch through exec_command.
read_file, edit_file and write_file take one file per call, while one coherent
apply_patch may span several files. Keep independent native calls in the same
turn, and do not batch file writes into a script.
- Persist until the task is fully handled this turn: gather context, implement,
run the validators, and report — do not stop at a partial result.
</model_guidance>`
@@ -29,6 +29,9 @@ func TestBuildSystemPromptAppendsModelAddendum(t *testing.T) {
if got := buildSystemPrompt(Options{Model: "gpt-5"}); !strings.Contains(got, openAIPromptAddendum) {
t.Fatalf("expected the OpenAI addendum in the gpt-5 prompt")
}
if got := strings.Join(strings.Fields(strings.ToLower(buildSystemPrompt(Options{Model: "gpt-5"}))), " "); !strings.Contains(got, "never run apply_patch through exec_command") {
t.Fatalf("OpenAI prompt must route edits through the native patch tool:\n%s", got)
}
// Claude is aligned with the core prompt and gets no family addendum now that
// comment discipline is universal; it must not pick up another family's block.
claude := buildSystemPrompt(Options{Model: "claude-opus-4-6"})
+11 -1
View File
@@ -29,7 +29,10 @@ func TestMain(m *testing.M) {
}
func TestCoreSystemPromptIncludesCodingQualityRules(t *testing.T) {
prompt := strings.ToLower(buildSystemPrompt(Options{}))
prompt := strings.Join(strings.Fields(strings.ToLower(buildSystemPrompt(Options{}))), " ")
if strings.Contains(prompt, "clickable file paths") {
t.Fatalf("system prompt must not claim inline-code file paths are clickable: %s", prompt)
}
for _, want := range []string{
"read-before-edit",
@@ -44,6 +47,13 @@ func TestCoreSystemPromptIncludesCodingQualityRules(t *testing.T) {
"do not recognize",
"scaled to the work",
"comment density",
"skip it for bounded changes in one component",
"do not read both unless a new need for exact content appears",
"combine compatible commands when diagnostics stay useful",
"do not invent a parallel flow",
"a successful native patch confirms the edit",
"once the required behavior passes validation, stop",
"do not look them up solely for the final summary",
} {
if !strings.Contains(prompt, want) {
t.Fatalf("expected core system prompt to include %q, got:\n%s", want, buildSystemPrompt(Options{}))
+48 -12
View File
@@ -839,18 +839,20 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
mcpTokenStore = nil
err = nil
}
criticalMCPConfig, optionalMCPConfig := splitMCPStartupConfig(mcpConfig)
mcpRuntime := mcpToolRuntime(noopMCPRuntime{})
if len(mcpConfig.Servers) > 0 {
mcpRuntime, err = deps.registerMCPTools(context.Background(), registry, mcpConfig, mcp.RegisterOptions{
if len(criticalMCPConfig.Servers) > 0 {
runtime, registerErr := deps.registerMCPTools(context.Background(), registry, criticalMCPConfig, mcp.RegisterOptions{
PermissionStore: mcpPermissionStore,
Autonomy: mcp.AutonomyLow,
Execution: executionRunner,
WorkspaceRoot: workspaceRoot,
})
}
if err != nil {
closeMCPRuntime(stderr, mcpRuntime)
return writeAppError(stderr, err.Error(), 1)
if registerErr != nil {
closeMCPRuntime(stderr, runtime)
return writeAppError(stderr, redaction.ErrorMessage(registerErr, redaction.Options{}), 1)
}
mcpRuntime = runtime
}
defer closeMCPRuntime(stderr, mcpRuntime)
// A server that could not be reached or validated is skipped, not fatal (one
@@ -909,11 +911,35 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
}
// Activate deferred MCP-tool loading for the interactive run only when the
// VISIBLE deferred-eligible count meets the resolved threshold, matching exec.
// The registry is complete (core + specialist + MCP + plugins) here, so the
// count is accurate; below threshold this is a no-op and the surface is
// unchanged. The interactive surface applies no operator tool filters, so
// enabled/disabled are nil — matching the AgentOptions below.
// This first pass covers every prompt-critical tool. Optional built-in MCP
// defaults re-run the same gate after publishing their atomic batch, so a
// later catalog generation cannot miss its loader. The interactive surface
// applies no operator tool filters, so enabled/disabled are nil — matching the
// AgentOptions below.
registerToolSearchIfEligible(registry, resolved.Tools.DeferThreshold, permissionMode, nil, nil)
var optionalMCPRuntime *optionalMCPStartup
var awaitToolReadiness func(context.Context)
if len(optionalMCPConfig.Servers) > 0 {
optionalMCPRuntime = startOptionalMCP(
context.Background(),
registry,
optionalMCPConfig,
mcp.RegisterOptions{
PermissionStore: mcpPermissionStore,
Autonomy: mcp.AutonomyLow,
Execution: executionRunner,
WorkspaceRoot: workspaceRoot,
},
deps.registerMCPTools,
func() {
registerToolSearchIfEligible(registry, resolved.Tools.DeferThreshold, permissionMode, nil, nil)
},
)
defer closeMCPRuntime(stderr, optionalMCPRuntime)
awaitToolReadiness = func(ctx context.Context) {
optionalMCPRuntime.Await(ctx, optionalMCPPromptGrace)
}
}
lastKnownMCPConfig := mcpConfig
fileTracker := tools.NewFileTracker()
var scratchBaseline scratchFileBaseline
@@ -951,8 +977,17 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
CompactionModel: resolved.Preferences.CompactionModel,
Provider: provider,
NewProvider: deps.newProvider,
ProbeProviderHealth: deps.probeProviderHealth,
UserAgent: userAgent(),
NewTurnSessionProvider: func(profile config.ProviderProfile, provider zeroruntime.Provider) zeroruntime.TurnSessionProvider {
if provider == nil {
return nil
}
if optimized, ok := providers.OptimizedTurnSessions(profile, provider, providers.Options{}); ok {
return optimized
}
return providers.DefaultTurnSessions(profile, provider, providers.Options{})
},
ProbeProviderHealth: deps.probeProviderHealth,
UserAgent: userAgent(),
PrepareRunCompletionWarning: func() {
scratchBaseline = scratchFileSnapshot(workspaceRoot)
},
@@ -960,6 +995,7 @@ func runInteractiveTUIWithSetup(stderr io.Writer, deps appDeps, permissionMode a
return scratchFileWarning(workspaceRoot, scratchBaseline)
},
Registry: registry,
AwaitToolReadiness: awaitToolReadiness,
SessionStore: sessionStore,
PeerService: peerService,
SandboxStore: sandboxStore,
+77 -1
View File
@@ -17,6 +17,7 @@ import (
"github.com/Gitlawb/zero/internal/agent"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/mcp"
"github.com/Gitlawb/zero/internal/redaction"
"github.com/Gitlawb/zero/internal/tools"
"github.com/Gitlawb/zero/internal/tui"
"github.com/Gitlawb/zero/internal/update"
@@ -486,6 +487,74 @@ func TestRunNoArgsLaunchesTUIWithMCPState(t *testing.T) {
}
}
func TestRunNoArgsPaintsBeforeOptionalDefaultMCPIsReady(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
cwd := t.TempDir()
userConfigPath := filepath.Join(t.TempDir(), "zero", "config.json")
permissionStore, err := mcp.NewPermissionStore(mcp.StoreOptions{FilePath: filepath.Join(t.TempDir(), "mcp-permissions.json")})
if err != nil {
t.Fatalf("NewPermissionStore() error = %v", err)
}
tokenStore, err := mcp.NewTokenStore(mcp.TokenStoreOptions{FilePath: filepath.Join(t.TempDir(), "mcp-oauth.json")})
if err != nil {
t.Fatalf("NewTokenStore() error = %v", err)
}
started := make(chan struct{})
release := make(chan struct{})
configResult := make(chan error, 1)
exitCode := runWithDeps([]string{}, &stdout, &stderr, appDeps{
getwd: func() (string, error) { return cwd, nil },
userConfigPath: func() (string, error) { return userConfigPath, nil },
resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) {
return config.ResolvedConfig{MaxTurns: 8}, nil
},
resolveMCPConfig: func(string, bool) (config.MCPConfig, error) {
return config.MCPConfig{Servers: config.DefaultMCPServers()}, nil
},
newMCPStore: func() (*mcp.PermissionStore, error) { return permissionStore, nil },
newMCPTokenStore: func() (*mcp.TokenStore, error) { return tokenStore, nil },
registerMCPTools: func(_ context.Context, registry *tools.Registry, cfg config.MCPConfig, _ mcp.RegisterOptions) (mcpToolRuntime, error) {
var configErr error
if _, ok := cfg.Servers["exa"]; !ok || len(cfg.Servers) != 1 {
configErr = fmt.Errorf("optional MCP config = %#v, want only exa", cfg.Servers)
}
configResult <- configErr
close(started)
<-release
registry.Register(cliFakeMCPRegistryTool{})
return noopMCPRuntime{}, nil
},
runTUI: func(ctx context.Context, options tui.Options) int {
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("optional MCP initialization did not start")
}
if options.AwaitToolReadiness == nil {
t.Fatal("TUI did not receive optional tool readiness barrier")
}
if _, ok := options.Registry.Get("mcp_docs_lookup"); ok {
t.Fatal("optional MCP tool was published before registration completed")
}
close(release)
options.AwaitToolReadiness(ctx)
if _, ok := options.Registry.Get("mcp_docs_lookup"); !ok {
t.Fatal("optional MCP tool was not visible after readiness completed")
}
return 0
},
})
if exitCode != 0 {
t.Fatalf("exitCode = %d stderr=%s", exitCode, stderr.String())
}
if configErr := <-configResult; configErr != nil {
t.Fatal(configErr)
}
}
func TestTUIMCPCommandUsesLastGoodConfigOnRefreshError(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
@@ -565,6 +634,7 @@ func TestRunNoArgsClosesPartialMCPRuntimeWhenRegistrationFails(t *testing.T) {
t.Fatalf("NewTokenStore() error = %v", err)
}
runtimeClosed := false
secret := "sk-proj-" + strings.Repeat("a", 23) + "0"
exitCode := runWithDeps([]string{}, &stdout, &stderr, appDeps{
getwd: func() (string, error) {
@@ -591,7 +661,7 @@ func TestRunNoArgsClosesPartialMCPRuntimeWhenRegistrationFails(t *testing.T) {
return closeFunc(func() error {
runtimeClosed = true
return nil
}), errors.New("register mcp tools failed")
}), fmt.Errorf("register mcp tools failed: %s", secret)
},
runTUI: func(ctx context.Context, options tui.Options) int {
t.Fatal("TUI should not launch after MCP registration fails")
@@ -605,6 +675,12 @@ func TestRunNoArgsClosesPartialMCPRuntimeWhenRegistrationFails(t *testing.T) {
if !strings.Contains(stderr.String(), "register mcp tools failed") {
t.Fatalf("stderr missing registration error: %s", stderr.String())
}
if strings.Contains(stderr.String(), secret) {
t.Fatalf("stderr leaked MCP registration secret: %s", stderr.String())
}
if !strings.Contains(stderr.String(), redaction.RedactedSecret) {
t.Fatalf("stderr missing redaction marker: %s", stderr.String())
}
if !runtimeClosed {
t.Fatal("partial MCP runtime was not closed after registration error")
}
+3
View File
@@ -943,6 +943,9 @@ func registerToolSearchIfEligible(registry *tools.Registry, deferThreshold int,
if deferThreshold <= 0 {
return
}
if _, registered := registry.Get(tools.ToolSearchToolName); registered {
return
}
if deferredEligibleCount(registry, permissionMode, enabledTools, disabledTools) < deferThreshold {
return
}
+15 -9
View File
@@ -267,23 +267,29 @@ func streamJSONPermissionEventType(event agent.PermissionEvent) streamjson.Event
func (writer *execEventWriter) usage(usage agent.Usage) {
if writer.format == execOutputJSON {
writer.writeJSON(map[string]any{
"type": "usage",
"prompt_tokens": usage.PromptTokens,
"completion_tokens": usage.CompletionTokens,
"total_tokens": usage.TotalTokens(),
"type": "usage",
"prompt_tokens": usage.PromptTokens,
"completion_tokens": usage.CompletionTokens,
"cached_input_tokens": usage.CachedInputTokens,
"cache_write_tokens": usage.CacheWriteTokens,
"total_tokens": usage.TotalTokens(),
})
return
}
if writer.format == execOutputStreamJSON {
promptTokens := usage.EffectiveInputTokens()
completionTokens := usage.EffectiveOutputTokens()
cachedInputTokens := usage.CachedInputTokens
cacheWriteTokens := usage.CacheWriteTokens
totalTokens := usage.TotalTokens()
writer.writeStreamJSON(streamjson.Event{
Type: streamjson.EventUsage,
RunID: writer.runID,
PromptTokens: &promptTokens,
CompletionTokens: &completionTokens,
TotalTokens: &totalTokens,
Type: streamjson.EventUsage,
RunID: writer.runID,
PromptTokens: &promptTokens,
CompletionTokens: &completionTokens,
CachedInputTokens: &cachedInputTokens,
CacheWriteTokens: &cacheWriteTokens,
TotalTokens: &totalTokens,
})
}
}
+36
View File
@@ -52,3 +52,39 @@ func TestExecWriterPropagatesToolResultTruncation(t *testing.T) {
})
}
}
func TestExecWriterPropagatesCacheUsage(t *testing.T) {
for _, format := range []execOutputFormat{execOutputJSON, execOutputStreamJSON} {
t.Run(string(format), func(t *testing.T) {
var stdout bytes.Buffer
writer := execEventWriter{
stdout: &stdout,
format: format,
runID: "run_usage",
streamedText: &strings.Builder{},
}
writer.usage(agent.Usage{
InputTokens: 100,
OutputTokens: 20,
PromptTokens: 100,
CompletionTokens: 20,
CachedInputTokens: 60,
CacheWriteTokens: 10,
})
if writer.err != nil {
t.Fatalf("usage: %v", writer.err)
}
var payload map[string]any
if err := json.Unmarshal(bytes.TrimSpace(stdout.Bytes()), &payload); err != nil {
t.Fatalf("decode output %q: %v", stdout.String(), err)
}
cachedKey, writeKey := "cached_input_tokens", "cache_write_tokens"
if format == execOutputStreamJSON {
cachedKey, writeKey = "cachedInputTokens", "cacheWriteTokens"
}
if payload[cachedKey] != float64(60) || payload[writeKey] != float64(10) {
t.Fatalf("cache usage missing from payload: %#v", payload)
}
})
}
}
+147
View File
@@ -0,0 +1,147 @@
package cli
import (
"context"
"sync"
"time"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/mcp"
"github.com/Gitlawb/zero/internal/tools"
)
const (
optionalMCPPromptGrace = time.Second
optionalMCPCloseGrace = 250 * time.Millisecond
)
// splitMCPStartupConfig keeps explicitly configured servers on the
// prompt-critical path. Only unchanged built-in defaults may initialize after
// first paint; their failure was already intentionally silent.
func splitMCPStartupConfig(cfg config.MCPConfig) (critical config.MCPConfig, optional config.MCPConfig) {
critical.Servers = map[string]config.MCPServerConfig{}
optional.Servers = map[string]config.MCPServerConfig{}
for name, server := range cfg.Servers {
if server.Disabled {
continue
}
if config.IsUnconfiguredDefault(name, server) {
optional.Servers[name] = server
continue
}
critical.Servers[name] = server
}
return critical, optional
}
type optionalMCPStartup struct {
cancel context.CancelFunc
done chan struct{}
mu sync.Mutex
runtime mcpToolRuntime
closeOnce sync.Once
runtimeCloseOnce sync.Once
closeErr error
}
func startOptionalMCP(
parent context.Context,
registry *tools.Registry,
cfg config.MCPConfig,
options mcp.RegisterOptions,
register func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error),
onReady func(),
) *optionalMCPStartup {
ctx, cancel := context.WithCancel(parent)
startup := &optionalMCPStartup{cancel: cancel, done: make(chan struct{})}
go func() {
runtime, err := register(ctx, registry, cfg, options)
if err == nil && onReady != nil {
onReady()
}
startup.mu.Lock()
startup.runtime = runtime
startup.mu.Unlock()
close(startup.done)
}()
return startup
}
// Await gives optional tools one bounded opportunity to join the next turn's
// immutable registry snapshot. A timeout leaves this turn on the existing core
// tool set; a later turn observes the completed generation.
func (startup *optionalMCPStartup) Await(ctx context.Context, timeout time.Duration) bool {
if startup == nil {
return true
}
select {
case <-startup.done:
return true
default:
}
if timeout <= 0 {
select {
case <-startup.done:
return true
case <-ctx.Done():
return false
}
}
timer := time.NewTimer(timeout)
defer timer.Stop()
select {
case <-startup.done:
return true
case <-timer.C:
return false
case <-ctx.Done():
return false
}
}
func (startup *optionalMCPStartup) Close() error {
if startup == nil {
return nil
}
startup.closeOnce.Do(func() {
startup.cancel()
reaped := make(chan struct{})
go func() {
defer close(reaped)
<-startup.done
startup.closeRuntime()
}()
timer := time.NewTimer(optionalMCPCloseGrace)
defer timer.Stop()
select {
case <-reaped:
case <-timer.C:
}
})
startup.mu.Lock()
defer startup.mu.Unlock()
return startup.closeErr
}
func (startup *optionalMCPStartup) closeRuntime() {
startup.runtimeCloseOnce.Do(func() {
startup.mu.Lock()
runtime := startup.runtime
startup.mu.Unlock()
if runtime == nil {
return
}
err := runtime.Close()
startup.mu.Lock()
if err != nil && startup.closeErr == nil {
startup.closeErr = err
}
startup.mu.Unlock()
})
}
func (startup *optionalMCPStartup) Skipped() []mcp.SkippedServer {
return nil
}
+187
View File
@@ -0,0 +1,187 @@
package cli
import (
"context"
"errors"
"testing"
"time"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/mcp"
"github.com/Gitlawb/zero/internal/tools"
)
type closeTrackingMCPRuntime struct {
closed chan struct{}
err error
}
func (runtime *closeTrackingMCPRuntime) Close() error {
close(runtime.closed)
return runtime.err
}
func (*closeTrackingMCPRuntime) Skipped() []mcp.SkippedServer { return nil }
func TestSplitMCPStartupConfigDefersOnlyUnconfiguredDefaults(t *testing.T) {
defaults := config.DefaultMCPServers()
cfg := config.MCPConfig{Servers: map[string]config.MCPServerConfig{
"exa": defaults["exa"],
"docs": {Type: "http", URL: "https://docs.example.test/mcp"},
"disabled": {Type: "http", URL: "https://disabled.example.test/mcp", Disabled: true},
}}
critical, optional := splitMCPStartupConfig(cfg)
if _, ok := critical.Servers["docs"]; !ok || len(critical.Servers) != 1 {
t.Fatalf("critical servers = %#v, want only docs", critical.Servers)
}
if _, ok := optional.Servers["exa"]; !ok || len(optional.Servers) != 1 {
t.Fatalf("optional servers = %#v, want only exa", optional.Servers)
}
}
func TestOptionalMCPAwaitTimesOutThenObservesPublishedTools(t *testing.T) {
registry := tools.NewRegistry()
release := make(chan struct{})
started := make(chan struct{})
startup := startOptionalMCP(
context.Background(),
registry,
config.MCPConfig{Servers: config.DefaultMCPServers()},
mcp.RegisterOptions{},
func(ctx context.Context, registry *tools.Registry, _ config.MCPConfig, _ mcp.RegisterOptions) (mcpToolRuntime, error) {
close(started)
select {
case <-release:
registry.RegisterBatch([]tools.Tool{tools.NewScopedReadFileTool(t.TempDir(), nil)})
return noopMCPRuntime{}, nil
case <-ctx.Done():
return noopMCPRuntime{}, ctx.Err()
}
},
nil,
)
defer func() { _ = startup.Close() }()
<-started
if startup.Await(t.Context(), time.Millisecond) {
t.Fatal("optional startup unexpectedly completed before release")
}
if got := len(registry.All()); got != 0 {
t.Fatalf("registry published tools before readiness: %d", got)
}
close(release)
if !startup.Await(t.Context(), time.Second) {
t.Fatal("optional startup did not complete after release")
}
if _, ok := registry.Get("read_file"); !ok {
t.Fatal("completed optional startup did not publish its tool batch")
}
}
func TestOptionalMCPFailedRegistrationSkipsReadinessCallback(t *testing.T) {
readyCalls := 0
startup := startOptionalMCP(
context.Background(),
tools.NewRegistry(),
config.MCPConfig{Servers: config.DefaultMCPServers()},
mcp.RegisterOptions{},
func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
return noopMCPRuntime{}, errors.New("connect failed")
},
func() { readyCalls++ },
)
defer func() { _ = startup.Close() }()
if !startup.Await(t.Context(), time.Second) {
t.Fatal("failed optional startup did not settle")
}
if readyCalls != 0 {
t.Fatalf("readiness callback ran after failed registration: %d", readyCalls)
}
}
func TestOptionalMCPCloseCancelsBlockedStartup(t *testing.T) {
started := make(chan struct{})
startup := startOptionalMCP(
context.Background(),
tools.NewRegistry(),
config.MCPConfig{Servers: config.DefaultMCPServers()},
mcp.RegisterOptions{},
func(ctx context.Context, _ *tools.Registry, _ config.MCPConfig, _ mcp.RegisterOptions) (mcpToolRuntime, error) {
close(started)
<-ctx.Done()
return noopMCPRuntime{}, nil
},
nil,
)
<-started
if err := startup.Close(); err != nil {
t.Fatalf("Close() error = %v", err)
}
}
func TestOptionalMCPCloseReturnsRuntimeErrorWithinGrace(t *testing.T) {
closeErr := errors.New("close failed")
runtime := &closeTrackingMCPRuntime{closed: make(chan struct{}), err: closeErr}
startup := startOptionalMCP(
context.Background(),
tools.NewRegistry(),
config.MCPConfig{Servers: config.DefaultMCPServers()},
mcp.RegisterOptions{},
func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
return runtime, nil
},
nil,
)
if !startup.Await(t.Context(), time.Second) {
t.Fatal("optional startup did not settle")
}
if err := startup.Close(); !errors.Is(err, closeErr) {
t.Fatalf("Close() error = %v, want %v", err, closeErr)
}
}
func TestOptionalMCPCloseBoundsRegistrationThatIgnoresCancellation(t *testing.T) {
started := make(chan struct{})
release := make(chan struct{})
released := false
defer func() {
if !released {
close(release)
}
}()
runtime := &closeTrackingMCPRuntime{closed: make(chan struct{})}
startup := startOptionalMCP(
context.Background(),
tools.NewRegistry(),
config.MCPConfig{Servers: config.DefaultMCPServers()},
mcp.RegisterOptions{},
func(context.Context, *tools.Registry, config.MCPConfig, mcp.RegisterOptions) (mcpToolRuntime, error) {
close(started)
<-release
return runtime, nil
},
nil,
)
<-started
closeResult := make(chan error, 1)
go func() { closeResult <- startup.Close() }()
select {
case err := <-closeResult:
if err != nil {
t.Fatalf("Close() error = %v", err)
}
case <-time.After(time.Second):
t.Fatal("Close blocked on registration after cancellation")
}
close(release)
released = true
select {
case <-runtime.closed:
case <-time.After(time.Second):
t.Fatal("late optional MCP runtime was not closed")
}
}
+9 -66
View File
@@ -5,31 +5,23 @@ import (
"fmt"
"os"
"path/filepath"
"sync/atomic"
"time"
"github.com/Gitlawb/zero/internal/lockutil"
)
// Cross-process lock tuning. The lock is held only for a single metadata
// read-modify-write (milliseconds), never across a job's exec, so the timeout is
// generous and the stale threshold sits far above any real hold.
// read-modify-write (milliseconds), never across a job's exec, so the timeout
// is generous relative to any real hold.
const (
cronLockTimeout = 10 * time.Second
cronLockStaleAfter = 60 * time.Second
cronLockRetryDelay = 20 * time.Millisecond
)
var cronLockSeq atomic.Uint64
// lockJob takes a cross-process exclusive lock for one job by O_EXCL-creating a
// sibling "<id>.lock" file next to the job directory (so Remove's RemoveAll of
// the job dir never deletes a live lock). It serializes the read-modify-write of
// a job's metadata across concurrent schedulers and commands. The lock file is
// removed on release; a stale lock from a crashed holder (older than
// cronLockStaleAfter) is reclaimed. Wall-clock time.Now is used deliberately
// (not the injectable Store.now) so lock timing never depends on a frozen test
// clock and the stale check compares against real file mtimes.
// lockJob takes a cross-process advisory lock for one job on a stable sibling
// "<id>.lock" file next to the job directory. It serializes the read-modify-write
// of a job's metadata across concurrent schedulers and commands. The kernel
// releases the lock when a holder exits; the file is never renamed or removed.
func (s *Store) lockJob(id string) (func(), error) {
if !validID(id) {
return nil, fmt.Errorf("invalid cron job id %q", id)
@@ -38,64 +30,15 @@ func (s *Store) lockJob(id string) (func(), error) {
if err := os.MkdirAll(filepath.Dir(lockPath), 0o700); err != nil {
return nil, err
}
token := fmt.Sprintf("%d-%d-%d", os.Getpid(), time.Now().UnixNano(), cronLockSeq.Add(1))
deadline := time.Now().Add(cronLockTimeout)
for {
f, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
lock, err := lockutil.TryAcquireFileLock(lockPath)
if err == nil {
// A partial write would leave a lock file without our token, so the
// releaser could never delete it — stranding the lock. Fail closed.
if _, werr := f.WriteString(token); werr != nil {
_ = f.Close()
_ = lockutil.RemoveLockFile(lockPath)
return nil, fmt.Errorf("cron: write job lock: %w", werr)
}
if cerr := f.Close(); cerr != nil {
_ = lockutil.RemoveLockFile(lockPath)
return nil, fmt.Errorf("cron: close job lock: %w", cerr)
}
var released bool
return func() {
if released {
return
}
released = true
if data, rerr := os.ReadFile(lockPath); rerr == nil && string(data) == token {
_ = lockutil.RemoveLockFile(lockPath)
}
}, nil
return func() { _ = lock.Release() }, nil
}
// On Windows a concurrent holder's os.Remove leaves the lock file in a
// "delete pending" state, so an O_EXCL create races it with
// ERROR_ACCESS_DENIED (os.ErrPermission) rather than ErrExist. Treat that
// as contention and retry, exactly like ErrExist — otherwise the lock
// spuriously fails under concurrency on Windows.
if !errors.Is(err, os.ErrExist) && !errors.Is(err, os.ErrPermission) {
if !errors.Is(err, lockutil.ErrLockHeld) {
return nil, fmt.Errorf("cron: acquire job lock: %w", err)
}
// Reclaim a stale lock left by a crashed holder — atomically (H3). A blind
// Remove lets two racers both "reclaim" and recreate, so both hold the lock;
// reclaimStaleLock renames the file aside (only one rename of a given file
// wins) and restores it if it turns out fresh, so a live lock is never deleted
// out from under its holder.
if info, statErr := os.Stat(lockPath); statErr == nil && time.Since(info.ModTime()) > cronLockStaleAfter {
cleared, rerr := lockutil.ReclaimStaleLock(lockPath, token, func(reclaimedPath string) bool {
info, err := os.Stat(reclaimedPath)
return err == nil && time.Since(info.ModTime()) <= cronLockStaleAfter
})
if rerr != nil {
// Reclaim hit a hard failure: the rename aside failed outright, or a
// live holder's lock could not be put back (the lock path may be
// missing, so re-acquiring would break mutual exclusion). Fail closed
// instead of spinning to the deadline.
return nil, fmt.Errorf("cron: reclaim stale job lock: %w", rerr)
}
if cleared {
continue // cleared a genuinely stale lock; retry the O_EXCL create now
}
// Lost the reclaim race (or it was actually fresh) — fall through to the
// bounded wait instead of hot-spinning on a reclaim that never wins (L13).
}
if !time.Now().Before(deadline) {
return nil, fmt.Errorf("cron: timed out acquiring job lock for %q", id)
}
+1 -1
View File
@@ -64,7 +64,7 @@ func TestMutateRemovedJobReturnsNotFound(t *testing.T) {
if !errors.Is(err, ErrJobNotFound) {
t.Fatalf("Mutate on removed job = %v, want ErrJobNotFound", err)
}
// And the lock file left no litter / did not resurrect the job dir.
// The stable advisory-lock file does not resurrect the removed job dir.
if _, gerr := store.Get(job.ID); !errors.Is(gerr, ErrJobNotFound) {
t.Fatalf("job should remain removed, Get = %v", gerr)
}
+35 -67
View File
@@ -3,107 +3,60 @@ package daemon
import (
"errors"
"fmt"
"io/fs"
"os"
"strconv"
"strings"
"sync/atomic"
"github.com/Gitlawb/zero/internal/lockutil"
)
// Single-instance lock. Mirrors reference-daemon-code-agent-js/supervisor.js's
// lock file: a PID file created with O_EXCL. A second start fails; a STALE lock
// left by a dead daemon (the recorded PID is no longer alive) is reclaimed so the
// daemon recovers from an unclean shutdown without manual cleanup.
// Single-instance lock. The stable PID file carries a kernel advisory lock for
// the daemon's full lifetime. Process exit releases the lock automatically, so
// stale metadata needs no rename/remove recovery protocol.
// ErrAlreadyRunning is returned when a live daemon already holds the lock.
var ErrAlreadyRunning = errors.New("daemon: another instance is already running")
// fileLock is an acquired single-instance lock.
type fileLock struct {
path string
lock *lockutil.FileLock
}
// processAlive reports whether pid is a live process. Implemented per-platform
// (lock_posix.go / lock_windows.go). It is a package var so tests can stub it.
var processAlive = osProcessAlive
// acquireLock takes the single-instance lock at path, reclaiming a stale lock
// whose recorded PID is dead. isAlive overrides the liveness check (tests pass a
// stub); nil uses the real processAlive.
// acquireLock takes the single-instance advisory lock at path. isAlive is used
// only to enrich a contention error with the current holder's PID; the kernel
// lock, not PID metadata, is authoritative.
func acquireLock(path string, isAlive func(pid int) bool) (*fileLock, error) {
if isAlive == nil {
isAlive = processAlive
}
// At most two passes: create, or detect-stale-then-retry once.
for attempt := 0; attempt < 2; attempt++ {
f, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
if err == nil {
// A failed PID write would leave a malformed lock file that another
// process reads as stale (unparsable PID) and wrongly reclaims, breaking
// the single-instance guarantee — so on write failure, remove it and fail.
if _, werr := fmt.Fprintf(f, "%d\n", os.Getpid()); werr != nil {
_ = f.Close()
_ = lockutil.RemoveLockFile(path)
return nil, werr
}
if cerr := f.Close(); cerr != nil {
_ = lockutil.RemoveLockFile(path)
return nil, cerr
}
return &fileLock{path: path}, nil
}
if !errors.Is(err, fs.ErrExist) && !errors.Is(err, os.ErrPermission) {
lock, err := lockutil.TryAcquireFileLock(path)
if err != nil {
if !errors.Is(err, lockutil.ErrLockHeld) {
return nil, err
}
pid, perr := readPidFile(path)
if perr == nil && pid > 0 && isAlive(pid) {
return nil, fmt.Errorf("%w (pid %d)", ErrAlreadyRunning, pid)
}
// Stale lock (dead PID or unreadable) — reclaim it atomically, then retry the
// O_EXCL create. A blind Remove here races: two daemons starting at once could
// both read the stale PID, both Remove, and then one Removes the OTHER's
// freshly-created lock — leaving both "holding" the single-instance lock.
// reclaimStaleLock renames the file aside so only one racer wins the rename,
// and restores it if a live holder reacquired in the gap (D6).
if _, rerr := reclaimStaleLock(path, isAlive); rerr != nil {
// Reclaim hit a hard failure: the rename aside failed outright, or a
// live holder's lock could not be put back (the lock path may be
// missing, so re-acquiring would break the single-instance guarantee).
// Fail closed instead of spinning to the deadline.
return nil, fmt.Errorf("daemon: reclaim stale lock: %w", rerr)
}
return nil, ErrAlreadyRunning
}
return nil, ErrAlreadyRunning
if err := lock.WriteMetadata([]byte(fmt.Sprintf("%d\n", os.Getpid()))); err != nil {
return nil, errors.Join(err, lock.Release())
}
return &fileLock{lock: lock}, nil
}
// daemonLockSeq makes each reclaim attempt's sidelined filename unique per process.
var daemonLockSeq atomic.Uint64
// reclaimStaleLock atomically reclaims a single-instance lock whose recorded
// PID is dead, via lockutil.ReclaimStaleLock: the lock file is renamed aside
// (only one racer can win the rename) and stolen only if the moved file's PID
// is still dead; a holder that reacquired the lock in the gap between the
// stale check and the rename carries a LIVE pid and is restored rather than
// stolen. Returns true only when a genuinely stale lock was removed. A
// non-nil error means the caller must fail closed instead of re-acquiring
// (see lockutil.ReclaimStaleLock).
func reclaimStaleLock(path string, isAlive func(pid int) bool) (bool, error) {
suffix := fmt.Sprintf("%d-%d", os.Getpid(), daemonLockSeq.Add(1))
return lockutil.ReclaimStaleLock(path, suffix, func(reclaimedPath string) bool {
pid, err := readPidFile(reclaimedPath)
return err == nil && pid > 0 && isAlive(pid)
})
}
// release removes the lock file. Safe to call once. An already-missing lock
// file is not an error (lockutil.RemoveLockFile swallows it on every platform).
// release drops the kernel lock and closes its handle. The stable PID file is
// deliberately retained so contenders always address the same inode.
func (l *fileLock) release() error {
if l == nil || l.path == "" {
if l == nil || l.lock == nil {
return nil
}
return lockutil.RemoveLockFile(l.path)
return l.lock.Release()
}
// readPidFile reads and parses the PID recorded in a lock file.
@@ -112,5 +65,20 @@ func readPidFile(path string) (int, error) {
if err != nil {
return 0, err
}
return strconv.Atoi(strings.TrimSpace(string(data)))
metadata := strings.TrimSpace(string(data))
if pid, err := strconv.Atoi(metadata); err == nil {
return pid, nil
}
pidText, sequenceText, ok := strings.Cut(metadata, "-")
if !ok || pidText == "" || sequenceText == "" || strings.Contains(sequenceText, "-") {
return 0, fmt.Errorf("daemon: parse lock PID metadata %q", metadata)
}
if _, err := strconv.ParseUint(sequenceText, 10, 64); err != nil {
return 0, fmt.Errorf("daemon: parse lock holder sequence: %w", err)
}
pid, err := strconv.Atoi(pidText)
if err != nil {
return 0, fmt.Errorf("daemon: parse lock PID: %w", err)
}
return pid, nil
}
+54 -42
View File
@@ -32,27 +32,48 @@ func TestLockSingleInstance(t *testing.T) {
_ = l2.release()
}
func TestLockStaleRecovery(t *testing.T) {
// A PID liveness probe is diagnostic only. Even if it reports the recorded PID
// dead, a contended kernel lock proves that the holder is active and must win.
// The old rename-aside protocol trusted this stale observation and admitted a
// second daemon while the first one was still inside its critical section.
func TestLockDoesNotOverrideKernelHolderWithStalePIDProbe(t *testing.T) {
path := filepath.Join(t.TempDir(), "daemon.lock")
// Simulate a stale lock from a crashed daemon: a PID file whose process is
// dead.
first, err := acquireLock(path, func(int) bool { return true })
if err != nil {
t.Fatalf("first acquireLock: %v", err)
}
defer first.release()
second, err := acquireLock(path, func(int) bool { return false })
if second != nil {
_ = second.release()
}
if !errors.Is(err, ErrAlreadyRunning) {
t.Fatalf("contended acquire with stale PID probe = %v, want ErrAlreadyRunning", err)
}
}
func TestLockIgnoresStaleMetadataWhenKernelLockIsFree(t *testing.T) {
path := filepath.Join(t.TempDir(), "daemon.lock")
// A crashed daemon may leave PID metadata behind, but its kernel lock is
// released automatically, so the next daemon can acquire without moving it.
if err := os.WriteFile(path, []byte("4242\n"), 0o600); err != nil {
t.Fatalf("write stale lock: %v", err)
}
dead := func(int) bool { return false }
l, err := acquireLock(path, dead)
if err != nil {
t.Fatalf("stale-lock recovery failed: %v", err)
t.Fatalf("acquire over stale metadata: %v", err)
}
// The lock now records OUR pid, not the stale one.
data, _ := os.ReadFile(path)
if strings.TrimSpace(string(data)) != strconv.Itoa(os.Getpid()) {
t.Fatalf("reclaimed lock pid = %q, want %d", strings.TrimSpace(string(data)), os.Getpid())
t.Fatalf("lock pid = %q, want %d", strings.TrimSpace(string(data)), os.Getpid())
}
_ = l.release()
}
func TestLockReleaseRemovesFile(t *testing.T) {
func TestLockReleaseKeepsStableFile(t *testing.T) {
path := filepath.Join(t.TempDir(), "daemon.lock")
l, err := acquireLock(path, func(int) bool { return true })
if err != nil {
@@ -61,46 +82,37 @@ func TestLockReleaseRemovesFile(t *testing.T) {
if err := l.release(); err != nil {
t.Fatalf("release: %v", err)
}
if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("lock file still present after release: %v", err)
if _, err := os.Stat(path); err != nil {
t.Fatalf("release removed stable lock file: %v", err)
}
}
func TestReclaimStaleLockRemovesDeadHolder(t *testing.T) {
func TestReadPidFileMetadataFormats(t *testing.T) {
path := filepath.Join(t.TempDir(), "daemon.lock")
if err := os.WriteFile(path, []byte("4242\n"), 0o600); err != nil {
t.Fatalf("write stale lock: %v", err)
}
if ok, err := reclaimStaleLock(path, func(int) bool { return false }); err != nil || !ok {
t.Fatalf("reclaimStaleLock must report a genuinely stale (dead-PID) lock reclaimed (ok=%v err=%v)", ok, err)
}
if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("a reclaimed stale lock must be removed, stat = %v", err)
}
}
func TestReclaimStaleLockRestoresLiveHolder(t *testing.T) {
// If a holder reacquires the lock in the gap between the stale check and the
// rename, the moved file carries a LIVE pid; reclaim must restore it untouched
// rather than steal it — otherwise two daemons both "hold" the lock (D6).
path := filepath.Join(t.TempDir(), "daemon.lock")
if err := os.WriteFile(path, []byte("4242\n"), 0o600); err != nil {
t.Fatalf("write lock: %v", err)
}
if ok, err := reclaimStaleLock(path, func(int) bool { return true }); err != nil || ok {
t.Fatalf("reclaimStaleLock must NOT report a live-PID lock reclaimed (ok=%v err=%v)", ok, err)
}
data, err := os.ReadFile(path)
if err != nil {
t.Fatalf("live lock must be restored in place, read = %v", err)
}
if strings.TrimSpace(string(data)) != "4242" {
t.Fatalf("restored lock content = %q, want %q (unchanged)", strings.TrimSpace(string(data)), "4242")
}
// No sidelined ".stale" leftovers.
matches, _ := filepath.Glob(path + ".stale.*")
if len(matches) != 0 {
t.Fatalf("reclaim left sidelined files: %v", matches)
for _, test := range []struct {
name string
metadata string
wantPID int
wantErr bool
}{
{name: "legacy PID", metadata: "4242\n", wantPID: 4242},
{name: "PID and holder sequence", metadata: "4242-7\n", wantPID: 4242},
{name: "missing sequence", metadata: "4242-", wantErr: true},
{name: "non-numeric sequence", metadata: "4242-next", wantErr: true},
{name: "extra component", metadata: "4242-7-8", wantErr: true},
} {
t.Run(test.name, func(t *testing.T) {
if err := os.WriteFile(path, []byte(test.metadata), 0o600); err != nil {
t.Fatal(err)
}
got, err := readPidFile(path)
if (err != nil) != test.wantErr {
t.Fatalf("readPidFile() error = %v, wantErr %v", err, test.wantErr)
}
if !test.wantErr && got != test.wantPID {
t.Fatalf("readPidFile() = %d, want %d", got, test.wantPID)
}
})
}
}
+6 -2
View File
@@ -124,12 +124,16 @@ func TestServerEndToEnd(t *testing.T) {
case <-time.After(3 * time.Second):
t.Fatal("Serve did not return after shutdown")
}
// Socket, status, and lock files are removed on exit.
for _, p := range []string{paths.Socket, paths.Status, paths.Lock} {
// Runtime endpoints are removed on exit. The advisory lock file remains as
// a stable inode, but the kernel lock on it has been released.
for _, p := range []string{paths.Socket, paths.Status} {
if _, err := os.Stat(p); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("file %s not cleaned up after shutdown: %v", p, err)
}
}
if _, err := os.Stat(paths.Lock); err != nil {
t.Fatalf("stable lock file missing after shutdown: %v", err)
}
}
func TestServerSecondInstanceFails(t *testing.T) {
+157 -29
View File
@@ -3,13 +3,93 @@
package execution
import (
"errors"
"os"
"os/exec"
"strconv"
"strings"
"syscall"
"testing"
"time"
)
// requireProcessState refuses to let a broken ps pipeline masquerade as passing
// coverage. Every test in this file proves something about zombie detection,
// and every one of those proofs is read out of ps — so if ps cannot report
// state, the tests verify nothing whatever they do next.
//
// The old code answered that with t.Skip, which is indistinguishable from
// success in CI output. This probes BOTH ps pipelines against this test
// process, which is by definition alive and not a zombie, so a working ps must
// report a live state for it. Failing that:
//
// - In CI (or with ZERO_REQUIRE_PS=1) it is a hard failure. ps is present on
// every image used today, so its absence means the pipeline regressed and
// that has to surface immediately rather than as skipped-green.
// - Elsewhere it stays a skip, so a contributor on an exotic box is not
// blocked by an environment the project does not claim to support.
//
// This only covers "ps cannot report state at all". Once it passes, a later
// failure to observe an EXPECTED state is a real defect, and the tests below
// treat it as one.
func requireProcessState(t *testing.T) {
t.Helper()
self := os.Getpid()
var reasons []string
// The production path: ps -A -o pid=,pgid=,stat=, parsed by
// processTableStateMatches. This is what signalTargetRunning consults.
if running, ok := processIsRunning(self); !ok || !running {
reasons = append(reasons, "the process table read used by signalTargetRunning reported no live state for this test process")
}
// This file's independent oracle: ps -o stat= -p <pid>. It is deliberately
// NOT the production parser, so a regression in that parser cannot hide by
// breaking the test's own precondition check in the same way.
if zombie, ok := processIsZombie(self); !ok || zombie {
reasons = append(reasons, "the per-PID state read used by this test file reported no live state for this test process")
}
if len(reasons) == 0 {
return
}
reason := "cannot read process state, so ps-based zombie detection is unverifiable here: " + strings.Join(reasons, "; ")
if requireWorkingPS() {
t.Fatal(reason + " — failing rather than skipping because CI (or ZERO_REQUIRE_PS) is set, where ps is expected to work")
}
t.Skip(reason)
}
// requireWorkingPS reports whether an unusable ps must fail the run instead of
// skipping it. ZERO_REQUIRE_PS lets a developer opt into the strict behavior
// locally to reproduce what CI does.
func requireWorkingPS() bool {
if strings.TrimSpace(os.Getenv("ZERO_REQUIRE_PS")) != "" {
return true
}
return strings.TrimSpace(os.Getenv("CI")) != ""
}
// awaitZombie blocks until pid is an unreaped zombie. The caller must not have
// waited on the child, so the zombie is a stable state rather than a race: it
// persists until something reaps it. Failing to reach it therefore means the
// child never exited or ps stopped reporting its state, and both are defects
// rather than reasons to skip.
func awaitZombie(t *testing.T, pid int) {
t.Helper()
deadline := time.Now().Add(5 * time.Second)
for {
zombie, ok := processIsZombie(pid)
if ok && zombie {
return
}
if time.Now().After(deadline) {
if !ok {
t.Fatalf("pid %d left the process table without being reaped; this test never waited on it, so its zombie row must persist", pid)
}
t.Fatalf("pid %d did not reach the zombie state within 5s", pid)
}
time.Sleep(10 * time.Millisecond)
}
}
// TestTerminateProcessTreeTreatsZombieGroupAsExited is the regression test for
// counting unreaped processes as alive: kill(2) with signal 0 succeeds for a
// zombie, so a group whose only remaining member is waiting to be reaped used to
@@ -17,6 +97,8 @@ import (
// as "did not exit after SIGKILL" — a spurious cleanup failure whose timing
// depended on when the reaper ran.
func TestTerminateProcessTreeTreatsZombieGroupAsExited(t *testing.T) {
requireProcessState(t)
cmd := exec.Command("sh", "-c", "exit 0")
ConfigureProcessGroup(cmd)
if err := cmd.Start(); err != nil {
@@ -27,23 +109,59 @@ func TestTerminateProcessTreeTreatsZombieGroupAsExited(t *testing.T) {
// is exactly the state a terminated tree passes through.
t.Cleanup(func() { _ = cmd.Wait() })
deadline := time.Now().Add(5 * time.Second)
for {
if zombie, ok := processIsZombie(pid); ok && zombie {
break
awaitZombie(t, pid)
// Whether a zombie-only group can be signalled is a KERNEL difference, not a
// ps one. On Linux the zombie still occupies the process table and
// kill(-pid, 0) succeeds, so signalTargetRunning goes on to consult ps —
// which is the branch this test was written for. Darwin returns EPERM for
// the same call, so signalTargetRunning short-circuits on the kill and never
// reaches the ps check.
//
// The OUTCOME is identical and is asserted on every platform below. Only the
// branch that produces it differs, so record that instead of skipping the
// whole test (which would drop Darwin's coverage of the outcome) or failing
// it (which would report a supported platform as broken). Any errno other
// than EPERM means the group is not in the state this test needs.
psPathReachable := true
if err := syscall.Kill(-pid, syscall.Signal(0)); err != nil {
if !errors.Is(err, syscall.EPERM) {
t.Fatalf("zombie group %d is unsignalable with an unexpected error (%v); the zombie case is not present, so nothing below would be meaningful", pid, err)
}
if time.Now().After(deadline) {
t.Skip("child did not reach the zombie state; ps output is unavailable in this environment")
}
time.Sleep(10 * time.Millisecond)
}
if syscall.Kill(-pid, syscall.Signal(0)) != nil {
t.Skip("the zombie group is no longer signalable; cannot exercise the case")
psPathReachable = false
t.Logf("kernel refuses to signal zombie-only group %d (EPERM); signalTargetRunning answers from kill(2) rather than ps here, so the ps-based zombie check is asserted directly on platforms that allow the signal", pid)
}
// Where the signal IS permitted, pin the ps-based check itself. This is the
// regression #774 was about, and asserting only signalTargetRunning would
// let a ps-side regression hide behind the kill result.
if psPathReachable {
running, ok := processGroupHasRunningMember(pid)
if !ok {
t.Fatalf("process group %d has no rows while its zombie is unreaped; ps stopped reporting the group rather than the group being gone", pid)
}
if running {
t.Fatal("the ps-based group check must not report a zombie-only group as running")
}
}
// Valid on every platform: whichever branch answers, a zombie-only group is
// not running. On a kernel that refuses the signal this is the whole of the
// coverage, and it is more than the previous skip provided.
if signalTargetRunning(-pid) {
t.Fatal("a group holding only a zombie must not count as running")
}
if !psPathReachable {
// terminateTarget's own kill(2) would hit the same EPERM and be returned
// as a cleanup error, so this assertion cannot be made where the kernel
// refuses the signal. That may be a real gap in TerminateProcessTree on
// such a platform rather than a test limitation — a zombie-only tree is
// already exited and should terminate cleanly — but confirming and
// changing that is a production fix, not this test's business.
t.Log("skipping the TerminateProcessTree assertion: its kill(2) would hit the same EPERM as the probe above")
return
}
if err := TerminateProcessTree(pid, 50*time.Millisecond, 10*time.Millisecond); err != nil {
t.Fatalf("TerminateProcessTree on an already-exited tree: %v", err)
}
@@ -53,6 +171,8 @@ func TestTerminateProcessTreeTreatsZombieGroupAsExited(t *testing.T) {
// a live member must still be seen as running and then actually stopped, so the
// zombie tolerance above cannot degrade into ignoring real processes.
func TestTerminateProcessTreeStopsRunningGroup(t *testing.T) {
requireProcessState(t)
cmd := exec.Command("sh", "-c", "sleep 30")
ConfigureProcessGroup(cmd)
if err := cmd.Start(); err != nil {
@@ -67,7 +187,15 @@ func TestTerminateProcessTreeStopsRunningGroup(t *testing.T) {
if err := TerminateProcessTree(pid, time.Second, 10*time.Millisecond); err != nil {
t.Fatalf("TerminateProcessTree: %v", err)
}
if running, ok := processGroupHasRunningMember(pid); ok && running {
// ok reports whether ANY row for this group was found. This test never waits
// on the child, so its zombie row must still be there — an unknown state
// means ps stopped answering, not that the group is gone. Treating that as a
// pass is what let this assertion succeed without checking anything.
running, ok := processGroupHasRunningMember(pid)
if !ok {
t.Fatalf("process group %d has no rows after termination; the unreaped child's zombie row must persist, so the group's state is undeterminable rather than clean", pid)
}
if running {
t.Fatal("the group still has a running member after termination")
}
}
@@ -82,6 +210,8 @@ func TestTerminateProcessTreeStopsRunningGroup(t *testing.T) {
// "unknown" resulted every time, and signalTargetRunning conservatively (and
// incorrectly) treated a zombie individual-PID target as still running.
func TestSignalTargetRunningTreatsZombieIndividualPIDAsExited(t *testing.T) {
requireProcessState(t)
cmd := exec.Command("sh", "-c", "exit 0")
// Deliberately do NOT call ConfigureProcessGroup: the child inherits this
// test process's group, so it is not its own leader and processSignalTarget
@@ -92,40 +222,38 @@ func TestSignalTargetRunningTreatsZombieIndividualPIDAsExited(t *testing.T) {
pid := cmd.Process.Pid
t.Cleanup(func() { _ = cmd.Wait() })
deadline := time.Now().Add(5 * time.Second)
for {
if zombie, ok := processIsZombie(pid); ok && zombie {
break
}
if time.Now().After(deadline) {
t.Skip("child did not reach the zombie state; ps output is unavailable in this environment")
}
time.Sleep(10 * time.Millisecond)
}
awaitZombie(t, pid)
target, err := processSignalTarget(pid)
if err != nil {
t.Fatalf("processSignalTarget: %v", err)
}
// The child inherited this process's group, so its PGID is this process's
// PID and can never equal its own PID. A negative target here would mean
// processSignalTarget misread the group, which is the very thing the
// individual-PID fallback exists to handle.
if target != pid {
t.Skip("child unexpectedly became its own group leader; cannot exercise the individual-PID fallback")
t.Fatalf("processSignalTarget(%d) = %d, want the individual PID: the child inherited this process's group, so it is not its own leader", pid, target)
}
if signalTargetRunning(target) {
t.Fatal("a zombie individual-PID target must not count as running")
}
}
// processIsZombie reports a process's zombie state via ps. ok is false when the
// state could not be read.
// processIsZombie reports a process's zombie state via ps. It is deliberately a
// separate, per-PID ps invocation rather than a reuse of the production
// process-table parser: these tests assert what that parser concludes, so
// borrowing it to establish their own preconditions would let one regression
// silence both sides at once.
//
// ok is false when no state could be read, which after requireProcessState has
// passed means the process left the table rather than that ps is broken.
func processIsZombie(pid int) (zombie bool, ok bool) {
output, err := exec.Command("ps", "-o", "stat=", "-p", strconv.Itoa(pid)).Output()
if err != nil {
return false, false
}
state := string(output)
for len(state) > 0 && (state[0] == ' ' || state[0] == '\n' || state[0] == '\t') {
state = state[1:]
}
state := strings.TrimLeft(string(output), " \n\t")
if state == "" {
return false, false
}
+9 -63
View File
@@ -5,93 +5,39 @@ import (
"fmt"
"os"
"path/filepath"
"sync/atomic"
"time"
"github.com/Gitlawb/zero/internal/lockutil"
)
// Cross-process lock tuning for the audit log. The lock is held only across a
// single read-then-append (milliseconds), so the timeout is generous and the
// stale threshold sits far above any real hold.
// single read-then-append (milliseconds), so the timeout is generous relative
// to any real hold.
const (
auditLockTimeout = 10 * time.Second
auditLockStaleAfter = 60 * time.Second
auditLockRetryDelay = 20 * time.Millisecond
)
var auditLockSeq atomic.Uint64
// lockAudit takes a cross-process exclusive lock on the audit log by
// O_EXCL-creating a sibling "<auditPath>.lock" file (removed on release). It makes
// lockAudit takes a cross-process advisory lock on the audit log through a
// stable sibling "<auditPath>.lock" file. It makes
// the lastSequence()+append in append() atomic across processes; the audit log is
// a shared global file, so without it two processes can read the same last
// sequence and write duplicate numbers. A stale lock from a crashed holder (older
// than auditLockStaleAfter) is reclaimed. Wall-clock time is used deliberately so
// lock timing never depends on the store's injectable clock and the stale check
// compares against real file mtimes. This mirrors internal/cron/lock.go and
// internal/oauth/lock.go. Uses internal/lockutil for lock restore and reclaim.
// sequence and write duplicate numbers. Kernel lock lifetime provides crash
// recovery without moving or deleting the lock path.
func (store *AuditStore) lockAudit() (func(), error) {
lockPath := store.auditPath + ".lock"
if err := os.MkdirAll(filepath.Dir(lockPath), 0o700); err != nil {
return nil, err
}
token := fmt.Sprintf("%d-%d-%d", os.Getpid(), time.Now().UnixNano(), auditLockSeq.Add(1))
deadline := time.Now().Add(auditLockTimeout)
for {
f, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
lock, err := lockutil.TryAcquireFileLock(lockPath)
if err == nil {
// A partial write would leave a lock file without our token, so the
// releaser could never delete it — stranding the lock. Fail closed.
if _, werr := f.WriteString(token); werr != nil {
_ = f.Close()
_ = lockutil.RemoveLockFile(lockPath)
return nil, fmt.Errorf("hooks: write audit lock: %w", werr)
}
if cerr := f.Close(); cerr != nil {
_ = lockutil.RemoveLockFile(lockPath)
return nil, fmt.Errorf("hooks: close audit lock: %w", cerr)
}
var released bool
return func() {
if released {
return
}
released = true
if data, rerr := os.ReadFile(lockPath); rerr == nil && string(data) == token {
_ = lockutil.RemoveLockFile(lockPath)
}
}, nil
return func() { _ = lock.Release() }, nil
}
// On Windows a concurrent holder's os.Remove leaves the lock file in a
// "delete pending" state, so an O_EXCL create races it with
// ERROR_ACCESS_DENIED (os.ErrPermission) rather than ErrExist. Treat both
// as contention and retry.
if !errors.Is(err, os.ErrExist) && !errors.Is(err, os.ErrPermission) {
if !errors.Is(err, lockutil.ErrLockHeld) {
return nil, fmt.Errorf("hooks: acquire audit lock: %w", err)
}
// Reclaim a stale lock left by a crashed holder — atomically (H3). A blind
// Remove lets two racers both reclaim + recreate and so both hold the lock;
// reclaimStaleLock renames the file aside (only one rename wins) and restores
// it if it turns out fresh, so a live lock is never deleted out from under it.
if info, statErr := os.Stat(lockPath); statErr == nil && time.Since(info.ModTime()) > auditLockStaleAfter {
cleared, rerr := lockutil.ReclaimStaleLock(lockPath, token, func(reclaimedPath string) bool {
info, err := os.Stat(reclaimedPath)
return err == nil && time.Since(info.ModTime()) <= auditLockStaleAfter
})
if rerr != nil {
// Reclaim hit a hard failure: the rename aside failed outright, or a
// live holder's lock could not be put back (the lock path may be
// missing, so re-acquiring would break mutual exclusion). Fail closed
// instead of spinning to the deadline.
return nil, fmt.Errorf("hooks: reclaim stale audit lock: %w", rerr)
}
if cleared {
continue
}
// Lost the reclaim race (or it was actually fresh) — fall through to the
// bounded wait rather than hot-spinning on a reclaim that never wins.
}
if !time.Now().Before(deadline) {
return nil, fmt.Errorf("hooks: timed out acquiring audit lock")
}
+111 -48
View File
@@ -1,64 +1,127 @@
// Package lockutil provides the platform-specific file primitives behind the
// O_EXCL lock files in cron, daemon, hooks, oauth, and swarm: a no-overwrite
// restore for locks that were sidelined during a stale-reclaim attempt, and a
// lock file remover with one cross-platform contract (missing files are a
// no-op; Windows retries transient sharing violations).
// Package lockutil provides cross-process advisory file locks used by cron,
// daemon, hooks, oauth, and swarm. Locks are held by the kernel for the lifetime
// of an open file handle, so process exit releases them without a stale-file
// reclaim protocol. The lock file itself is stable: holders never rename or
// remove it while another process may still have its inode locked.
package lockutil
import (
"io"
"errors"
"fmt"
"os"
"path/filepath"
"strings"
"sync"
"sync/atomic"
)
// restoreByCopy restores reclaimed to path without overwriting an existing
// path, as a fallback for when the platform's primary no-replace primitive
// (hard link on POSIX, MoveFileEx on Windows) fails for a reason other than
// the destination existing. It stages a full copy under a private name next
// to reclaimed and publishes it to path with publish (the same no-replace
// primitive the caller's platform uses for the primary restore), so path
// never appears in a partially-copied state: a crash between staging and
// publish leaves path exactly as it was (missing or, if a competing holder
// created it, untouched), never a truncated file that a PID/content-based
// liveness check could mistake for dead. Leaving path missing here would let
// the next O_EXCL create succeed while the sidelined holder is still live,
// breaking mutual exclusion. publish keeps the no-overwrite guarantee: a new
// holder that appeared in the meantime wins and this returns os.ErrExist. The
// copy resets the lock's mtime to now, which only makes the restored lock
// look fresher; that is safe, since it is being handed back to a live holder.
func restoreByCopy(reclaimed, path string, publish func(from, to string) error) error {
staged := reclaimed + ".copy"
src, err := os.Open(reclaimed)
var (
// ErrLockHeld reports that another process or goroutine currently holds the
// advisory lock. Callers may wait and retry until their own deadline.
ErrLockHeld = errors.New("lockutil: lock is held")
// ErrLockReleased reports an attempt to update metadata after Release.
ErrLockReleased = errors.New("lockutil: lock is released")
)
var holderSequence atomic.Uint64
// FileLock is an exclusive advisory lock held through an open file handle.
// Release is idempotent. The lock path remains on disk after release; deleting
// or replacing it would let different processes lock different inodes.
type FileLock struct {
mu sync.Mutex
file *os.File
state platformLockState
released bool
releaseErr error
}
// TryAcquireFileLock attempts to take an exclusive advisory lock without
// waiting. A contended lock returns ErrLockHeld. The kernel releases a lock if
// its process exits, so an old unlocked file never needs stale reclamation.
func TryAcquireFileLock(path string) (*FileLock, error) {
return TryAcquireFileLockAt(filepath.Dir(path), path)
}
// TryAcquireFileLockAt attempts to lock path while confining every component
// below root to a handle-relative traversal. root is the caller's trusted
// boundary and may itself be reached through a legitimate link; links and
// reparse points below it are rejected before metadata can be written.
func TryAcquireFileLockAt(root, path string) (*FileLock, error) {
if root == "" {
return nil, errors.New("lockutil: lock root is empty")
}
relative, err := filepath.Rel(root, path)
if err != nil {
return err
return nil, fmt.Errorf("lockutil: resolve lock path relative to root: %w", err)
}
dst, err := os.OpenFile(staged, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
relative = filepath.Clean(relative)
if relative == "." || relative == ".." || filepath.IsAbs(relative) || strings.HasPrefix(relative, ".."+string(filepath.Separator)) {
return nil, fmt.Errorf("lockutil: lock path %q escapes root %q", path, root)
}
file, err := openLockFileAt(root, relative, path)
if err != nil {
_ = src.Close()
return err
return nil, fmt.Errorf("lockutil: open lock file: %w", err)
}
_, err = io.Copy(dst, src)
// Close the source before removing anything below: Go opens files without
// FILE_SHARE_DELETE on Windows, so deleting reclaimed or staged while src
// is open would fail with a sharing violation.
_ = src.Close()
state, contended, err := tryLockFile(file)
if err != nil {
_ = dst.Close()
_ = os.Remove(staged)
return err
_ = file.Close()
return nil, fmt.Errorf("lockutil: acquire advisory lock: %w", err)
}
if err := dst.Close(); err != nil {
_ = os.Remove(staged)
return err
if contended {
_ = file.Close()
return nil, ErrLockHeld
}
if err := publish(staged, path); err != nil {
_ = os.Remove(staged)
return err
lock := &FileLock{file: file, state: state}
metadata := []byte(fmt.Sprintf("%d-%d\n", os.Getpid(), holderSequence.Add(1)))
if err := lock.WriteMetadata(metadata); err != nil {
return nil, errors.Join(err, lock.Release())
}
return lock, nil
}
// WriteMetadata replaces the diagnostic contents of the lock file while the
// caller holds it. Metadata does not establish ownership; the kernel lock does.
func (lock *FileLock) WriteMetadata(data []byte) error {
if lock == nil {
return ErrLockReleased
}
lock.mu.Lock()
defer lock.mu.Unlock()
if lock.released || lock.file == nil {
return ErrLockReleased
}
if err := lock.file.Truncate(0); err != nil {
return fmt.Errorf("lockutil: truncate lock metadata: %w", err)
}
if _, err := lock.file.Seek(0, 0); err != nil {
return fmt.Errorf("lockutil: seek lock metadata: %w", err)
}
if _, err := lock.file.Write(data); err != nil {
return fmt.Errorf("lockutil: write lock metadata: %w", err)
}
// publish may leave staged behind (a hard-link publish keeps both names),
// so remove it explicitly; a move-based publish already made this a no-op.
_ = os.Remove(staged)
// The lock is back at path, so the restore has succeeded; failing to clean
// up the sidelined name must not be reported as a failed restore.
_ = RemoveLockFile(reclaimed)
return nil
}
// Release unlocks and closes the held file handle. It deliberately leaves the
// stable lock file in place. Repeated calls return the first release result.
func (lock *FileLock) Release() error {
if lock == nil {
return nil
}
lock.mu.Lock()
defer lock.mu.Unlock()
if lock.released {
return lock.releaseErr
}
lock.released = true
unlockErr := unlockFile(lock.file, &lock.state)
closeErr := lock.file.Close()
lock.file = nil
if err := errors.Join(unlockErr, closeErr); err != nil {
lock.releaseErr = fmt.Errorf("lockutil: release advisory lock: %w", err)
}
return lock.releaseErr
}
+82 -33
View File
@@ -4,43 +4,92 @@ package lockutil
import (
"errors"
"io/fs"
"fmt"
"os"
"path/filepath"
"strings"
"golang.org/x/sys/unix"
)
// RestoreLockFile restores a sidelined lock file on non-Windows platforms. It
// uses os.Link so a competing lock created at path in the meantime wins: the
// link fails with os.ErrExist instead of overwriting it. If the link fails for
// any other reason (hardlink-incapable filesystems such as FAT or some
// FUSE/network mounts, ENOSPC, EPERM), it falls back to an O_EXCL copy rather
// than leaving path missing and the live holder's lock stranded in reclaimed.
// Once the lock is back at path the restore has succeeded, so a failed cleanup
// of the sidelined name is not reported as an error; the leftover file is
// invisible to the lock protocol.
func RestoreLockFile(reclaimed, path string) error {
err := os.Link(reclaimed, path)
type platformLockState struct{}
func openLockFileAt(root, relative, displayPath string) (_ *os.File, resultErr error) {
parentFD, err := unix.Open(root, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_CLOEXEC, 0)
if err != nil {
return nil, fmt.Errorf("open lock root: %w", err)
}
defer func() {
if parentFD >= 0 {
resultErr = errors.Join(resultErr, unix.Close(parentFD))
}
}()
components := strings.Split(filepath.Clean(relative), string(filepath.Separator))
for _, component := range components[:len(components)-1] {
nextFD, openErr := unix.Openat(parentFD, component, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0)
if openErr != nil {
return nil, fmt.Errorf("open lock path component %q: %w", component, openErr)
}
if closeErr := unix.Close(parentFD); closeErr != nil {
_ = unix.Close(nextFD)
parentFD = -1
return nil, fmt.Errorf("close lock path component: %w", closeErr)
}
parentFD = nextFD
}
name := components[len(components)-1]
fd := -1
// Darwin can transiently return ENOENT when many callers race to create the
// same O_NOFOLLOW entry. The bound parent handle cannot be redirected, so a
// bounded retry preserves first-use contention without weakening no-follow.
for attempt := 0; attempt < 10; attempt++ {
fd, err = unix.Openat(parentFD, name, unix.O_CREAT|unix.O_RDWR|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0o600)
if !errors.Is(err, unix.ENOENT) {
break
}
}
if err != nil {
return nil, err
}
file := os.NewFile(uintptr(fd), displayPath)
if file == nil {
_ = unix.Close(fd)
return nil, errors.New("convert lock handle to file")
}
var stat unix.Stat_t
if err := unix.Fstat(fd, &stat); err != nil {
return nil, errors.Join(fmt.Errorf("inspect lock file: %w", err), file.Close())
}
if stat.Mode&unix.S_IFMT != unix.S_IFREG {
return nil, errors.Join(errors.New("refusing non-regular lock file"), file.Close())
}
if stat.Nlink != 1 {
return nil, errors.Join(errors.New("refusing multiply-linked lock file"), file.Close())
}
if err := unix.Close(parentFD); err != nil {
parentFD = -1
return nil, errors.Join(fmt.Errorf("close lock parent: %w", err), file.Close())
}
parentFD = -1
return file, nil
}
func tryLockFile(file *os.File) (platformLockState, bool, error) {
err := unix.Flock(int(file.Fd()), unix.LOCK_EX|unix.LOCK_NB)
if err == nil {
_ = RemoveLockFile(reclaimed)
return platformLockState{}, false, nil
}
if errors.Is(err, unix.EWOULDBLOCK) || errors.Is(err, unix.EAGAIN) {
return platformLockState{}, true, nil
}
return platformLockState{}, false, err
}
func unlockFile(file *os.File, _ *platformLockState) error {
if file == nil {
return nil
}
if errors.Is(err, os.ErrExist) {
return err
}
return restoreByCopy(reclaimed, path, os.Link)
}
// isReclaimContended reports whether a failed rename-aside of a suspected
// stale lock means it was lost to a racer rather than a hard failure. POSIX
// rename has no contention errno (a lost race surfaces only as ENOENT, which
// ReclaimStaleLock already treats as benign), so nothing extra is benign here.
func isReclaimContended(error) bool { return false }
// RemoveLockFile removes a lock file on non-Windows platforms. Removing an
// already-missing file reports nil, matching the Windows implementation, so
// callers see one cross-platform contract.
func RemoveLockFile(path string) error {
if err := os.Remove(path); err != nil && !errors.Is(err, fs.ErrNotExist) {
return err
}
return nil
return unix.Flock(int(file.Fd()), unix.LOCK_UN)
}
+205 -101
View File
@@ -1,116 +1,220 @@
package lockutil
import (
"bufio"
"errors"
"fmt"
"io"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
"sync"
"testing"
)
func TestRestoreLockFile(t *testing.T) {
tempDir := t.TempDir()
path := filepath.Join(tempDir, "lock")
reclaimed := filepath.Join(tempDir, "lock.stale.token")
// 1. Successful restore when target does not exist
if err := os.WriteFile(reclaimed, []byte("token"), 0o600); err != nil {
t.Fatal(err)
}
if err := RestoreLockFile(reclaimed, path); err != nil {
t.Fatalf("RestoreLockFile failed: %v", err)
}
if data, err := os.ReadFile(path); err != nil || string(data) != "token" {
t.Fatalf("restored lock content = %q, err %v; want %q intact", data, err, "token")
}
if _, err := os.Stat(reclaimed); !os.IsNotExist(err) {
t.Fatalf("expected sidelined lock to be cleaned up/removed: %v", err)
}
// 2. Fail when target already exists (prevent overwrite)
if err := os.WriteFile(reclaimed, []byte("token"), 0o600); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("competing"), 0o600); err != nil {
t.Fatal(err)
}
err := RestoreLockFile(reclaimed, path)
if err == nil {
t.Fatal("expected RestoreLockFile to fail when target exists")
}
if !errors.Is(err, os.ErrExist) {
t.Fatalf("expected os.ErrExist, got: %v", err)
}
// The sidelined lock must still exist on failure, and the competing lock
// must be untouched.
if _, err := os.Stat(reclaimed); err != nil {
t.Fatalf("expected sidelined lock to still exist: %v", err)
}
if data, _ := os.ReadFile(path); string(data) != "competing" {
t.Fatalf("competing lock was clobbered: %q", data)
}
}
func TestRestoreByCopy(t *testing.T) {
tempDir := t.TempDir()
path := filepath.Join(tempDir, "lock")
reclaimed := filepath.Join(tempDir, "lock.stale.token")
// 1. Restores content and cleans up the sidelined file when the target is
// free (the fallback taken when the platform's no-replace primitive fails
// for a reason other than the destination existing).
if err := os.WriteFile(reclaimed, []byte("token"), 0o600); err != nil {
t.Fatal(err)
}
if err := restoreByCopy(reclaimed, path, os.Link); err != nil {
t.Fatalf("restoreByCopy failed: %v", err)
}
if data, err := os.ReadFile(path); err != nil || string(data) != "token" {
t.Fatalf("restored lock content = %q, err %v; want %q intact", data, err, "token")
}
if _, err := os.Stat(reclaimed); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expected sidelined lock to be cleaned up after copy: %v", err)
}
if _, err := os.Stat(reclaimed + ".copy"); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expected staged copy to be cleaned up: %v", err)
}
// 2. Keeps the no-overwrite guarantee: a competing lock at the target wins
// and the sidelined lock survives for the caller's ErrExist cleanup.
if err := os.WriteFile(reclaimed, []byte("token"), 0o600); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("competing"), 0o600); err != nil {
t.Fatal(err)
}
if err := restoreByCopy(reclaimed, path, os.Link); !errors.Is(err, os.ErrExist) {
t.Fatalf("expected os.ErrExist when target exists, got: %v", err)
}
if data, _ := os.ReadFile(path); string(data) != "competing" {
t.Fatalf("competing lock was clobbered: %q", data)
}
if _, err := os.Stat(reclaimed); err != nil {
t.Fatalf("expected sidelined lock to still exist: %v", err)
}
if _, err := os.Stat(reclaimed + ".copy"); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expected staged copy to be cleaned up after a failed publish: %v", err)
}
}
func TestRemoveLockFile(t *testing.T) {
func TestFileLockSerializesAndKeepsStablePath(t *testing.T) {
path := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(path, []byte("token"), 0o600); err != nil {
t.Fatal(err)
first, err := TryAcquireFileLock(path)
if err != nil {
t.Fatalf("first acquire: %v", err)
}
if err := RemoveLockFile(path); err != nil {
t.Fatalf("RemoveLockFile failed: %v", err)
firstInfo, err := os.Stat(path)
if err != nil {
t.Fatalf("stat first lock path: %v", err)
}
if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expected lock file to be removed: %v", err)
if _, err := TryAcquireFileLock(path); !errors.Is(err, ErrLockHeld) {
t.Fatalf("contended acquire = %v, want ErrLockHeld", err)
}
// Removing an already-missing lock file is a no-op on every platform.
if err := RemoveLockFile(path); err != nil {
t.Fatalf("RemoveLockFile on a missing file = %v, want nil", err)
if err := first.Release(); err != nil {
t.Fatalf("release first: %v", err)
}
releasedInfo, err := os.Stat(path)
if err != nil {
t.Fatalf("release removed stable lock path: %v", err)
}
if !os.SameFile(firstInfo, releasedInfo) {
t.Fatal("release replaced the stable lock file")
}
second, err := TryAcquireFileLock(path)
if err != nil {
t.Fatalf("acquire after release: %v", err)
}
secondInfo, err := os.Stat(path)
if err != nil {
t.Fatalf("stat second lock path: %v", err)
}
if !os.SameFile(firstInfo, secondInfo) {
t.Fatal("reacquire used a different lock file")
}
if err := second.Release(); err != nil {
t.Fatalf("release second: %v", err)
}
if err := second.Release(); err != nil {
t.Fatalf("idempotent release: %v", err)
}
}
func TestFileLockMetadata(t *testing.T) {
path := filepath.Join(t.TempDir(), "lock")
lock, err := TryAcquireFileLock(path)
if err != nil {
t.Fatal(err)
}
if err := lock.WriteMetadata([]byte("holder\n")); err != nil {
t.Fatalf("write metadata: %v", err)
}
data, err := os.ReadFile(path)
if err != nil || string(data) != "holder\n" {
t.Fatalf("metadata = %q, err %v", data, err)
}
if err := lock.Release(); err != nil {
t.Fatal(err)
}
if err := lock.WriteMetadata(nil); !errors.Is(err, ErrLockReleased) {
t.Fatalf("write after release = %v, want ErrLockReleased", err)
}
}
func TestFileLockConcurrentFirstAcquisition(t *testing.T) {
path := filepath.Join(t.TempDir(), "lock")
const contenders = 200
start := make(chan struct{})
results := make(chan error, contenders)
locks := make(chan *FileLock, contenders)
var ready sync.WaitGroup
ready.Add(contenders)
for range contenders {
go func() {
ready.Done()
<-start
lock, err := TryAcquireFileLock(path)
if lock != nil {
locks <- lock
}
results <- err
}()
}
ready.Wait()
close(start)
acquired := 0
for range contenders {
err := <-results
switch {
case err == nil:
acquired++
case !errors.Is(err, ErrLockHeld):
t.Errorf("concurrent first acquire: %v", err)
}
}
close(locks)
for lock := range locks {
if err := lock.Release(); err != nil {
t.Error(err)
}
}
if acquired != 1 {
t.Fatalf("concurrent first acquire: got %d holders, want 1", acquired)
}
}
func TestFileLockRefusesMultiplyLinkedEntry(t *testing.T) {
dir := t.TempDir()
target := filepath.Join(dir, "target")
lockPath := filepath.Join(dir, "lock")
const sentinel = "do not overwrite"
if err := os.WriteFile(target, []byte(sentinel), 0o600); err != nil {
t.Fatal(err)
}
if err := os.Link(target, lockPath); err != nil {
t.Fatalf("create redirected hard link: %v", err)
}
if lock, err := TryAcquireFileLock(lockPath); err == nil {
_ = lock.Release()
t.Fatal("TryAcquireFileLock accepted a multiply-linked entry")
}
data, err := os.ReadFile(target)
if err != nil {
t.Fatal(err)
}
if string(data) != sentinel {
t.Fatalf("redirected target was modified: %q", data)
}
}
func TestFileLockAtRejectsEscapingPath(t *testing.T) {
base := t.TempDir()
outside := filepath.Join(filepath.Dir(base), "outside.lock")
if lock, err := TryAcquireFileLockAt(base, outside); err == nil {
_ = lock.Release()
t.Fatal("TryAcquireFileLockAt accepted a path outside its root")
}
if _, err := os.Stat(outside); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("escaping path was created: %v", err)
}
}
// TestFileLockReleasedByProcessExit proves crash recovery is provided by the
// kernel lock lifetime. The child exits without calling Release; the next
// process can then lock the same persistent file without moving or deleting it.
func TestFileLockReleasedByProcessExit(t *testing.T) {
path := filepath.Join(t.TempDir(), "lock")
cmd := exec.Command(os.Args[0], "-test.run=^TestFileLockHelperProcess$")
cmd.Env = append(os.Environ(), "ZERO_LOCKUTIL_HELPER=1", "ZERO_LOCKUTIL_PATH="+path)
stdout, err := cmd.StdoutPipe()
if err != nil {
t.Fatal(err)
}
stdin, err := cmd.StdinPipe()
if err != nil {
t.Fatal(err)
}
cmd.Stderr = os.Stderr
if err := cmd.Start(); err != nil {
t.Fatal(err)
}
ready, err := bufio.NewReader(stdout).ReadString('\n')
if err != nil || strings.TrimSpace(ready) != "ready" {
_ = cmd.Process.Kill()
_ = cmd.Wait()
t.Fatalf("helper readiness = %q, err %v", ready, err)
}
if _, err := TryAcquireFileLock(path); !errors.Is(err, ErrLockHeld) {
_ = cmd.Process.Kill()
_ = cmd.Wait()
t.Fatalf("acquire while helper holds lock = %v, want ErrLockHeld", err)
}
if err := stdin.Close(); err != nil {
t.Fatal(err)
}
if err := cmd.Wait(); err != nil {
t.Fatalf("helper exit: %v", err)
}
lock, err := TryAcquireFileLock(path)
if err != nil {
t.Fatalf("acquire after helper exit: %v", err)
}
if err := lock.Release(); err != nil {
t.Fatal(err)
}
}
func TestFileLockHelperProcess(t *testing.T) {
if os.Getenv("ZERO_LOCKUTIL_HELPER") != "1" {
return
}
lock, err := TryAcquireFileLock(os.Getenv("ZERO_LOCKUTIL_PATH"))
if err != nil {
fmt.Fprintln(os.Stderr, err)
os.Exit(2)
}
if _, err := fmt.Fprintln(os.Stdout, "ready"); err != nil {
os.Exit(3)
}
_, _ = io.Copy(io.Discard, os.Stdin)
runtime.KeepAlive(lock) // deliberately exit without Release
}
+183 -58
View File
@@ -4,77 +4,202 @@ package lockutil
import (
"errors"
"fmt"
"os"
"syscall"
"time"
"path/filepath"
"strings"
"unsafe"
"golang.org/x/sys/windows"
)
// RestoreLockFile restores a sidelined lock file on Windows. It moves the file
// with a no-replace MoveFileEx so a competing lock created at path in the
// meantime wins: the move fails with os.ErrExist instead of overwriting it. If
// the move fails for any other reason (ERROR_SHARING_VIOLATION or
// ERROR_ACCESS_DENIED under concurrent access to the file), it falls back to
// an O_EXCL copy rather than leaving path missing and the live holder's lock
// stranded in reclaimed.
func RestoreLockFile(reclaimed, path string) error {
err := moveFileNoReplace(reclaimed, path)
if err == nil || errors.Is(err, os.ErrExist) {
return err
}
return restoreByCopy(reclaimed, path, moveFileNoReplace)
type platformLockState struct {
overlapped windows.Overlapped
}
// moveFileNoReplace renames from to to, failing if to already exists. It calls
// MoveFileExW directly via golang.org/x/sys/windows (the standard library's
// syscall package does not export MoveFileEx on this platform) with no flags,
// so an existing destination fails with ERROR_ALREADY_EXISTS, which satisfies
// errors.Is(err, os.ErrExist) (pinned by
// TestMoveFileNoReplaceMapsAlreadyExistsToErrExist), instead of overwriting it
// the way os.Rename does on Windows (it passes MOVEFILE_REPLACE_EXISTING to
// match POSIX rename semantics cross-platform).
func moveFileNoReplace(from, to string) error {
fromPtr, err := windows.UTF16PtrFromString(from)
const lockByteOffsetHigh = 1
func openLockFileAt(root, relative, displayPath string) (_ *os.File, resultErr error) {
rootPtr, err := windows.UTF16PtrFromString(root)
if err != nil {
return err
return nil, err
}
toPtr, err := windows.UTF16PtrFromString(to)
parent, err := windows.CreateFile(
rootPtr,
windows.FILE_READ_ATTRIBUTES|windows.FILE_TRAVERSE|windows.SYNCHRONIZE,
windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE|windows.FILE_SHARE_DELETE,
nil,
windows.OPEN_EXISTING,
windows.FILE_FLAG_BACKUP_SEMANTICS,
0,
)
if err != nil {
return err
return nil, fmt.Errorf("open lock root: %w", err)
}
return windows.MoveFileEx(fromPtr, toPtr, 0)
}
// isReclaimContended reports whether a failed rename-aside of a suspected
// stale lock means the file was concurrently open or pending deletion (which
// surfaces as ERROR_SHARING_VIOLATION or ERROR_ACCESS_DENIED on Windows)
// rather than a hard failure, so ReclaimStaleLock reports a lost race instead
// of an error. Neither errno satisfies errors.Is(err, os.ErrExist), and
// ERROR_SHARING_VIOLATION does not satisfy errors.Is(err, os.ErrPermission)
// either, so the classification compares raw errnos like RemoveLockFile below.
func isReclaimContended(err error) bool {
var errno syscall.Errno
return errors.As(err, &errno) && (errno == windows.ERROR_SHARING_VIOLATION || errno == windows.ERROR_ACCESS_DENIED)
}
// RemoveLockFile removes a lock file on Windows, retrying on the sharing
// violation or access denied errors that are common under heavy concurrent
// contention. Removing an already-missing file reports nil, matching the
// non-Windows implementation, so callers see one cross-platform contract.
func RemoveLockFile(path string) error {
var err error
for i := 0; i < 15; i++ {
err = os.Remove(path)
if err == nil || os.IsNotExist(err) {
return nil
defer func() {
if parent != 0 {
resultErr = errors.Join(resultErr, windows.CloseHandle(parent))
}
var errno syscall.Errno
if errors.As(err, &errno) && (errno == windows.ERROR_SHARING_VIOLATION || errno == windows.ERROR_ACCESS_DENIED) {
time.Sleep(5 * time.Millisecond)
}()
components := strings.Split(filepath.Clean(relative), string(filepath.Separator))
for _, component := range components[:len(components)-1] {
next, openErr := openLockDirectoryAt(parent, component)
if openErr != nil {
return nil, fmt.Errorf("open lock path component %q: %w", component, openErr)
}
if closeErr := windows.CloseHandle(parent); closeErr != nil {
_ = windows.CloseHandle(next)
parent = 0
return nil, fmt.Errorf("close lock path component: %w", closeErr)
}
parent = next
}
handle, err := openLockFileHandleAt(parent, components[len(components)-1])
if err != nil {
return nil, err
}
file := os.NewFile(uintptr(handle), displayPath)
if file == nil {
_ = windows.CloseHandle(handle)
return nil, errors.New("convert lock handle to file")
}
if err := windows.CloseHandle(parent); err != nil {
parent = 0
return nil, errors.Join(fmt.Errorf("close lock parent: %w", err), file.Close())
}
parent = 0
return file, nil
}
func openLockDirectoryAt(parent windows.Handle, name string) (windows.Handle, error) {
handle, err := openLockPathAt(
parent,
name,
windows.FILE_READ_ATTRIBUTES|windows.FILE_TRAVERSE|windows.SYNCHRONIZE,
windows.FILE_OPEN,
windows.FILE_ATTRIBUTE_DIRECTORY,
windows.FILE_DIRECTORY_FILE|windows.FILE_SYNCHRONOUS_IO_NONALERT|windows.FILE_OPEN_REPARSE_POINT,
)
if err != nil {
return 0, err
}
if err := validateWindowsLockHandle(handle, true); err != nil {
return 0, errors.Join(err, windows.CloseHandle(handle))
}
return handle, nil
}
func openLockFileHandleAt(parent windows.Handle, name string) (windows.Handle, error) {
for attempt := 0; attempt < 10; attempt++ {
handle, err := openLockPathAt(
parent,
name,
windows.GENERIC_READ|windows.GENERIC_WRITE|windows.SYNCHRONIZE,
windows.FILE_OPEN,
windows.FILE_ATTRIBUTE_NORMAL,
windows.FILE_NON_DIRECTORY_FILE|windows.FILE_SYNCHRONOUS_IO_NONALERT|windows.FILE_OPEN_REPARSE_POINT,
)
if err == nil {
if err := validateWindowsLockHandle(handle, false); err != nil {
return 0, errors.Join(err, windows.CloseHandle(handle))
}
return handle, nil
}
if !windowsLockPathMissing(err) {
return 0, err
}
handle, err = openLockPathAt(
parent,
name,
windows.GENERIC_READ|windows.GENERIC_WRITE|windows.SYNCHRONIZE,
windows.FILE_CREATE,
windows.FILE_ATTRIBUTE_NORMAL,
windows.FILE_NON_DIRECTORY_FILE|windows.FILE_SYNCHRONOUS_IO_NONALERT|windows.FILE_OPEN_REPARSE_POINT,
)
if err == windows.STATUS_OBJECT_NAME_COLLISION || errors.Is(err, windows.ERROR_ALREADY_EXISTS) {
continue
}
return err
if err != nil {
return 0, err
}
if err := validateWindowsLockHandle(handle, false); err != nil {
return 0, errors.Join(err, windows.CloseHandle(handle))
}
return handle, nil
}
return err
return 0, errors.New("lock file creation did not stabilize")
}
func windowsLockPathMissing(err error) bool {
return err == windows.STATUS_NO_SUCH_FILE ||
err == windows.STATUS_OBJECT_NAME_NOT_FOUND ||
err == windows.STATUS_OBJECT_PATH_NOT_FOUND ||
errors.Is(err, windows.ERROR_FILE_NOT_FOUND) ||
errors.Is(err, windows.ERROR_PATH_NOT_FOUND)
}
func openLockPathAt(parent windows.Handle, name string, access, disposition, attributes, options uint32) (windows.Handle, error) {
objectName, err := windows.NewNTUnicodeString(name)
if err != nil {
return 0, err
}
objectAttributes := &windows.OBJECT_ATTRIBUTES{
Length: uint32(unsafe.Sizeof(windows.OBJECT_ATTRIBUTES{})),
RootDirectory: parent,
ObjectName: objectName,
Attributes: windows.OBJ_CASE_INSENSITIVE | windows.OBJ_DONT_REPARSE,
}
var handle windows.Handle
err = windows.NtCreateFile(
&handle,
access,
objectAttributes,
&windows.IO_STATUS_BLOCK{},
nil,
attributes,
windows.FILE_SHARE_READ|windows.FILE_SHARE_WRITE|windows.FILE_SHARE_DELETE,
disposition,
options,
0,
0,
)
return handle, err
}
func validateWindowsLockHandle(handle windows.Handle, wantDirectory bool) error {
var info windows.ByHandleFileInformation
if err := windows.GetFileInformationByHandle(handle, &info); err != nil {
return fmt.Errorf("inspect lock path handle: %w", err)
}
isDirectory := info.FileAttributes&windows.FILE_ATTRIBUTE_DIRECTORY != 0
if info.FileAttributes&windows.FILE_ATTRIBUTE_REPARSE_POINT != 0 || isDirectory != wantDirectory {
return errors.New("refusing non-file or reparse-point lock path")
}
if !wantDirectory && info.NumberOfLinks != 1 {
return errors.New("refusing multiply-linked lock file")
}
return nil
}
func tryLockFile(file *os.File) (platformLockState, bool, error) {
state := platformLockState{overlapped: windows.Overlapped{OffsetHigh: lockByteOffsetHigh}}
flags := uint32(windows.LOCKFILE_EXCLUSIVE_LOCK | windows.LOCKFILE_FAIL_IMMEDIATELY)
err := windows.LockFileEx(windows.Handle(file.Fd()), flags, 0, 1, 0, &state.overlapped)
if err == nil {
return state, false, nil
}
if errors.Is(err, windows.ERROR_LOCK_VIOLATION) || errors.Is(err, windows.ERROR_SHARING_VIOLATION) {
return platformLockState{}, true, nil
}
return platformLockState{}, false, err
}
func unlockFile(file *os.File, state *platformLockState) error {
if file == nil {
return nil
}
return windows.UnlockFileEx(windows.Handle(file.Fd()), 0, 1, 0, &state.overlapped)
}
@@ -1,83 +0,0 @@
//go:build windows
package lockutil
import (
"errors"
"os"
"path/filepath"
"testing"
"time"
"golang.org/x/sys/windows"
)
// TestMoveFileNoReplaceMapsAlreadyExistsToErrExist pins the errno coupling
// every caller keys its cleanup on: with no flags, MoveFileEx fails with
// ERROR_ALREADY_EXISTS when the destination exists, and that errno must
// satisfy errors.Is(err, os.ErrExist).
func TestMoveFileNoReplaceMapsAlreadyExistsToErrExist(t *testing.T) {
dir := t.TempDir()
from := filepath.Join(dir, "from")
to := filepath.Join(dir, "to")
if err := os.WriteFile(from, []byte("sidelined"), 0o600); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(to, []byte("competing"), 0o600); err != nil {
t.Fatal(err)
}
err := moveFileNoReplace(from, to)
if !errors.Is(err, windows.ERROR_ALREADY_EXISTS) {
t.Fatalf("MoveFileEx with an existing destination = %v, want ERROR_ALREADY_EXISTS", err)
}
if !errors.Is(err, os.ErrExist) {
t.Fatalf("ERROR_ALREADY_EXISTS must satisfy errors.Is(err, os.ErrExist), got: %v", err)
}
if data, _ := os.ReadFile(to); string(data) != "competing" {
t.Fatalf("destination was overwritten: %q", data)
}
}
// TestRemoveLockFileRetriesSharingViolation holds the file open briefly (Go
// opens files without FILE_SHARE_DELETE on Windows, so a concurrent DeleteFile
// fails with ERROR_SHARING_VIOLATION while the handle is open); RemoveLockFile
// must ride its retry loop (15 x 5ms) past the window and succeed once the
// handle closes.
func TestRemoveLockFileRetriesSharingViolation(t *testing.T) {
path := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(path, []byte("token"), 0o600); err != nil {
t.Fatal(err)
}
f, err := os.Open(path)
if err != nil {
t.Fatal(err)
}
go func() {
time.Sleep(10 * time.Millisecond)
_ = f.Close()
}()
if err := RemoveLockFile(path); err != nil {
t.Fatalf("RemoveLockFile should succeed once the contending handle closes: %v", err)
}
if _, err := os.Stat(path); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expected lock file to be removed: %v", err)
}
}
// TestRemoveLockFileSharingViolationExhausts verifies the retry loop is
// bounded: with the contending handle held for the whole call, RemoveLockFile
// surfaces the sharing violation instead of retrying forever.
func TestRemoveLockFileSharingViolationExhausts(t *testing.T) {
path := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(path, []byte("token"), 0o600); err != nil {
t.Fatal(err)
}
f, err := os.Open(path)
if err != nil {
t.Fatal(err)
}
defer f.Close()
if err := RemoveLockFile(path); !errors.Is(err, windows.ERROR_SHARING_VIOLATION) {
t.Fatalf("RemoveLockFile with a held handle = %v, want ERROR_SHARING_VIOLATION", err)
}
}
-82
View File
@@ -1,82 +0,0 @@
package lockutil
import (
"errors"
"os"
)
// restoreLockFile is swappable so tests can force the fail-closed path of
// ReclaimStaleLock, which requires both the fast restore and its no-replace
// fallback (with its own copy fallback) to fail; that cannot be provoked
// portably on a healthy filesystem.
var restoreLockFile = restoreLiveLock
// restoreLiveLock puts a lock that turned out to be live back at path after
// ReclaimStaleLock moved it aside to inspect it. It first tries a fast,
// replacing rename straight from reclaimed to path: a single syscall, which
// keeps the window during which path does not exist (and so is open to a
// fourth process's unrelated O_EXCL create landing on it) as short as
// possible. RestoreLockFile's no-replace restore (with its slower copy
// fallback on some failures) is a correctness-preserving fallback for when
// the fast path itself fails (e.g. a cross-device sidelined name): it is a
// much longer version of the identical race, but still detects rather than
// silently clobbers a competing lock, which the fast path's replacing rename
// cannot do. Neither path makes the race impossible, only unlikely; see
// ReclaimStaleLock's doc comment.
func restoreLiveLock(reclaimed, path string) error {
if err := os.Rename(reclaimed, path); err == nil {
return nil
}
return RestoreLockFile(reclaimed, path)
}
// ReclaimStaleLock atomically reclaims a suspected-stale lock file. It renames
// lockPath aside to "<lockPath>.stale.<suffix>" (only one racer can win the
// rename of a given file, so two racers can never both reclaim the same lock),
// then consults isLive on the moved file; if the lock turns out to be live (a
// holder reacquired it in the gap between the caller's stale check and the
// rename) it is restored rather than stolen. The suffix must be unique per
// acquirer attempt. Returns true only when a genuinely stale lock was removed,
// so the caller knows it is safe to retry its exclusive create immediately; on
// a lost race it returns false. A non-nil error means either the rename aside
// failed for a reason that is not contention (so retrying cannot help and the
// caller should fail fast instead of spinning to its deadline), or a live
// holder's lock could not be restored (both restoreLiveLock's fast path and
// its no-replace fallback failed), so lockPath may be missing; callers must
// fail closed instead of re-acquiring. The sidelined file is removed on every
// restore failure: once the restore has failed it has no protocol function
// (release only consults the lock path), so keeping it would only leak files.
//
// The live-restore path has an inherent, unclosed race: between the rename
// aside above and restoreLiveLock putting the lock back, lockPath does not
// exist, so an unrelated caller's O_EXCL create can legitimately succeed
// there. restoreLiveLock's fast path then silently overwrites that new
// claimant's lock file, which does not corrupt release (it is
// ownership-aware, so the new claimant's later release safely no-ops against
// content it no longer owns) but does mean the new claimant can still run
// its critical section concurrently with the original live holder. Making
// this race actually impossible would need an OS-level advisory lock (flock
// / LockFileEx) held for a holder's whole critical section, checked
// non-destructively instead of by moving the file; restoreLiveLock only
// shrinks the window to roughly one syscall, it does not close it.
func ReclaimStaleLock(lockPath, suffix string, isLive func(reclaimedPath string) bool) (bool, error) {
reclaimed := lockPath + ".stale." + suffix
if err := os.Rename(lockPath, reclaimed); err != nil {
if errors.Is(err, os.ErrNotExist) || isReclaimContended(err) {
return false, nil // another racer already moved/removed it, or it vanished
}
return false, err
}
if isLive(reclaimed) {
// Put the live lock back instead of stealing it, and let the caller wait.
if rerr := restoreLockFile(reclaimed, lockPath); rerr != nil {
_ = RemoveLockFile(reclaimed)
if !errors.Is(rerr, os.ErrExist) {
return false, rerr
}
}
return false, nil
}
_ = RemoveLockFile(reclaimed)
return true, nil
}
-31
View File
@@ -1,31 +0,0 @@
//go:build !windows
package lockutil
import (
"os"
"path/filepath"
"testing"
)
// TestReclaimStaleLockPropagatesRenameFailure pins the fail-fast contract of
// the rename aside: a failure that is not a lost race must reach the caller.
// Mapping it to the benign lost-race result would leave acquirers spinning
// until their deadline and reporting a timeout instead of the real cause. A
// directory squatting on the sidelined name makes os.Rename fail with EISDIR,
// which matches neither os.ErrNotExist nor a contention errno.
func TestReclaimStaleLockPropagatesRenameFailure(t *testing.T) {
lockPath := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(lockPath, []byte("stale-holder"), 0o600); err != nil {
t.Fatal(err)
}
if err := os.Mkdir(lockPath+".stale.tok", 0o755); err != nil {
t.Fatal(err)
}
if ok, err := ReclaimStaleLock(lockPath, "tok", func(string) bool { return false }); err == nil || ok {
t.Fatalf("a hard rename failure must surface an error (ok=%v err=%v)", ok, err)
}
if _, err := os.Stat(lockPath); err != nil {
t.Fatalf("the lock file must be left in place on a failed rename: %v", err)
}
}
-112
View File
@@ -1,112 +0,0 @@
package lockutil
import (
"errors"
"os"
"path/filepath"
"testing"
)
func TestReclaimStaleLock(t *testing.T) {
lockPath := filepath.Join(t.TempDir(), "lock")
dead := func(string) bool { return false }
live := func(string) bool { return true }
// A lock the predicate reports dead is reclaimed and removed.
if err := os.WriteFile(lockPath, []byte("crashed-holder"), 0o600); err != nil {
t.Fatal(err)
}
if ok, err := ReclaimStaleLock(lockPath, "tok-a", dead); err != nil || !ok {
t.Fatalf("a dead lock should be reclaimed (ok=%v err=%v)", ok, err)
}
if _, err := os.Stat(lockPath); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("reclaimed dead lock should be gone, stat err=%v", err)
}
// A LIVE lock (a holder reacquired in the gap) must be RESTORED intact.
if err := os.WriteFile(lockPath, []byte("live-holder"), 0o600); err != nil {
t.Fatal(err)
}
if ok, err := ReclaimStaleLock(lockPath, "tok-b", live); err != nil || ok {
t.Fatalf("a live lock must not be reclaimed (ok=%v err=%v)", ok, err)
}
if data, err := os.ReadFile(lockPath); err != nil || string(data) != "live-holder" {
t.Fatalf("live lock must be left intact, got %q err %v", data, err)
}
// A missing lock reports no reclaim (nothing to steal).
_ = os.Remove(lockPath)
if ok, err := ReclaimStaleLock(lockPath, "tok-c", live); err != nil || ok {
t.Fatalf("a missing lock should not report a reclaim (ok=%v err=%v)", ok, err)
}
}
// restoreLiveLock's fast path is a single replacing rename, chosen to
// minimize (not eliminate) the window during which lockPath is absent and
// open to an unrelated O_EXCL create; see ReclaimStaleLock's doc comment for
// why that residual race exists. This test pins the resulting behavior: the
// fast path overwrites a competing file at path rather than detecting it,
// unlike RestoreLockFile's own no-replace contract.
func TestRestoreLiveLockFastPathOverwritesCompetingFile(t *testing.T) {
dir := t.TempDir()
reclaimed := filepath.Join(dir, "lock.stale.tok")
path := filepath.Join(dir, "lock")
if err := os.WriteFile(reclaimed, []byte("original-holder"), 0o600); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(path, []byte("new-claimant"), 0o600); err != nil {
t.Fatal(err)
}
if err := restoreLiveLock(reclaimed, path); err != nil {
t.Fatalf("restoreLiveLock failed: %v", err)
}
data, err := os.ReadFile(path)
if err != nil || string(data) != "original-holder" {
t.Fatalf("path = %q, err %v; want the original holder's content restored", data, err)
}
if _, err := os.Stat(reclaimed); !errors.Is(err, os.ErrNotExist) {
t.Fatalf("expected reclaimed to be consumed by the rename: %v", err)
}
}
func TestReclaimStaleLockFailsClosedOnRestoreError(t *testing.T) {
// When both the no-replace restore and its copy fallback fail (only
// provokable via the seam; a healthy filesystem cannot produce it), the
// caller must receive an error so it fails closed instead of re-acquiring a
// missing lock path, and the sidelined file must not leak.
restoreLockFile = func(reclaimed, path string) error { return errors.New("restore failed") }
defer func() { restoreLockFile = restoreLiveLock }()
lockPath := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(lockPath, []byte("live-holder"), 0o600); err != nil {
t.Fatal(err)
}
ok, err := ReclaimStaleLock(lockPath, "tok", func(string) bool { return true })
if err == nil || ok {
t.Fatalf("a failed restore must surface an error (ok=%v err=%v)", ok, err)
}
if matches, _ := filepath.Glob(lockPath + ".stale.*"); len(matches) != 0 {
t.Fatalf("a failed restore must not leak sidelined files: %v", matches)
}
}
func TestReclaimStaleLockDropsSidelinedWhenNewHolderWins(t *testing.T) {
// An os.ErrExist restore failure means a new holder recreated the lock
// path; that is not an error for the caller, and the sidelined file is
// dropped rather than leaked.
restoreLockFile = func(reclaimed, path string) error { return os.ErrExist }
defer func() { restoreLockFile = restoreLiveLock }()
lockPath := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(lockPath, []byte("live-holder"), 0o600); err != nil {
t.Fatal(err)
}
ok, err := ReclaimStaleLock(lockPath, "tok", func(string) bool { return true })
if err != nil || ok {
t.Fatalf("losing to a new holder is not an error (ok=%v err=%v)", ok, err)
}
if matches, _ := filepath.Glob(lockPath + ".stale.*"); len(matches) != 0 {
t.Fatalf("the sidelined file must be dropped when a new holder wins: %v", matches)
}
}
-50
View File
@@ -1,50 +0,0 @@
//go:build windows
package lockutil
import (
"os"
"path/filepath"
"testing"
)
// TestReclaimStaleLockTreatsSharingViolationAsLostRace pins the errno
// classification the Windows acquire loops depend on: a rename aside that
// fails because the lock file is concurrently open (Go opens files without
// FILE_SHARE_DELETE, so the rename fails with ERROR_SHARING_VIOLATION or
// ERROR_ACCESS_DENIED) is contention, not a hard failure, and must report a
// benign lost race so the caller keeps waiting.
func TestReclaimStaleLockTreatsSharingViolationAsLostRace(t *testing.T) {
lockPath := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(lockPath, []byte("held"), 0o600); err != nil {
t.Fatal(err)
}
f, err := os.Open(lockPath)
if err != nil {
t.Fatal(err)
}
defer f.Close()
if ok, rerr := ReclaimStaleLock(lockPath, "tok", func(string) bool { return true }); rerr != nil || ok {
t.Fatalf("a sharing-violation rename must report a benign lost race (ok=%v err=%v)", ok, rerr)
}
if _, err := os.Stat(lockPath); err != nil {
t.Fatalf("the lock file must be left in place: %v", err)
}
}
// TestReclaimStaleLockPropagatesRenameFailure pins the fail-fast contract of
// the rename aside on Windows: a failure that is not contention must reach
// the caller instead of masquerading as a lost race. A quote is an invalid
// filename character, so the sidelined name fails with ERROR_INVALID_NAME.
func TestReclaimStaleLockPropagatesRenameFailure(t *testing.T) {
lockPath := filepath.Join(t.TempDir(), "lock")
if err := os.WriteFile(lockPath, []byte("stale-holder"), 0o600); err != nil {
t.Fatal(err)
}
if ok, err := ReclaimStaleLock(lockPath, `bad"name`, func(string) bool { return false }); err == nil || ok {
t.Fatalf("a hard rename failure must surface an error (ok=%v err=%v)", ok, err)
}
if _, err := os.Stat(lockPath); err != nil {
t.Fatalf("the lock file must be left in place on a failed rename: %v", err)
}
}
+4 -2
View File
@@ -176,9 +176,11 @@ func RegisterTools(ctx context.Context, registry *tools.Registry, cfg config.MCP
staged = append(staged, tool)
}
}
for _, tool := range staged {
registry.Register(tool)
registered := make([]tools.Tool, 0, len(staged))
for index := range staged {
registered = append(registered, staged[index])
}
registry.RegisterBatch(registered)
return runtime, nil
}
+56
View File
@@ -70,6 +70,62 @@ func TestRegisterToolsAddsPromptGatedMCPTools(t *testing.T) {
}
}
func TestRegisterToolsPublishesServerToolsInOneGeneration(t *testing.T) {
registry := tools.NewRegistry()
before := registry.Snapshot().Generation
client := &fakeToolClient{listed: []RemoteTool{
{Name: "lookup", Description: "Lookup documentation"},
{Name: "search", Description: "Search documentation"},
}}
runtime, err := RegisterTools(context.Background(), registry, config.MCPConfig{Servers: map[string]config.MCPServerConfig{
"docs": {Type: "stdio", Command: "docs-mcp"},
}}, RegisterOptions{ClientFactory: func(context.Context, Server) (ToolClient, error) {
return client, nil
}})
if err != nil {
t.Fatalf("RegisterTools() error = %v", err)
}
defer runtime.Close()
after := registry.Snapshot()
if after.Generation != before+1 {
t.Fatalf("registry generation = %d, want %d", after.Generation, before+1)
}
for _, name := range []string{"mcp_docs_lookup", "mcp_docs_search"} {
if _, ok := registry.Get(name); !ok {
t.Fatalf("expected %s to be published in the batch", name)
}
}
}
func TestRegisterToolsPublishesNoneWhenServerToolValidationFails(t *testing.T) {
registry := tools.NewRegistry()
before := registry.Snapshot().Generation
client := &fakeToolClient{listed: []RemoteTool{
{Name: "lookup", Description: "Lookup documentation"},
{Name: "", Description: "Invalid nameless tool"},
}}
runtime, err := RegisterTools(context.Background(), registry, config.MCPConfig{Servers: map[string]config.MCPServerConfig{
"docs": {Type: "stdio", Command: "docs-mcp"},
}}, RegisterOptions{ClientFactory: func(context.Context, Server) (ToolClient, error) {
return client, nil
}})
if err != nil {
t.Fatalf("RegisterTools() error = %v", err)
}
defer runtime.Close()
if got := registry.Snapshot().Generation; got != before {
t.Fatalf("registry generation = %d, want unchanged %d", got, before)
}
if _, ok := registry.Get("mcp_docs_lookup"); ok {
t.Fatal("valid prefix of rejected server batch was published")
}
if skipped := runtime.Skipped(); len(skipped) != 1 || skipped[0].Name != "docs" || skipped[0].Err == nil {
t.Fatalf("Skipped() = %#v, want one validation failure for docs", skipped)
}
}
func TestRegisterToolsMarksPersistentlyApprovedToolsAllow(t *testing.T) {
store, err := NewPermissionStore(StoreOptions{
FilePath: filepath.Join(t.TempDir(), "permissions.json"),
+7 -64
View File
@@ -5,24 +5,16 @@ import (
"fmt"
"os"
"path/filepath"
"sync/atomic"
"time"
"github.com/Gitlawb/zero/internal/lockutil"
)
const (
fileLockTimeout = 5 * time.Second
fileLockStaleAfter = 30 * time.Second
)
const fileLockTimeout = 5 * time.Second
var lockSeq atomic.Uint64
// acquireFileLock takes a cross-process exclusive lock by creating lockPath with
// O_EXCL. It retries with a short backoff until a timeout, breaking a lock whose
// file is older than fileLockStaleAfter (so a crashed holder cannot deadlock the
// store). Release is ownership-aware: it removes the lock only if it still holds
// our token, so a stale-broken holder cannot delete a newer holder's lock.
// acquireFileLock takes a cross-process advisory lock through a stable lock
// file. It retries with a short backoff until a timeout. Kernel lock lifetime
// provides crash recovery, so the path is never renamed or removed.
func acquireFileLock(lockPath string, now func() time.Time) (func(), error) {
if now == nil {
now = time.Now
@@ -30,64 +22,15 @@ func acquireFileLock(lockPath string, now func() time.Time) (func(), error) {
if err := os.MkdirAll(filepath.Dir(lockPath), 0o700); err != nil {
return nil, err
}
token := fmt.Sprintf("%d-%d-%d", os.Getpid(), now().UnixNano(), lockSeq.Add(1))
deadline := now().Add(fileLockTimeout)
for {
f, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600)
lock, err := lockutil.TryAcquireFileLock(lockPath)
if err == nil {
// A partial/failed write would leave a lock file without our token, so
// the releaser below could never delete it — stranding the lock for
// other processes. Fail closed: remove the file and surface the error.
if _, werr := f.WriteString(token); werr != nil {
_ = f.Close()
_ = lockutil.RemoveLockFile(lockPath)
return nil, fmt.Errorf("oauth: write token lock: %w", werr)
}
if cerr := f.Close(); cerr != nil {
_ = lockutil.RemoveLockFile(lockPath)
return nil, fmt.Errorf("oauth: close token lock: %w", cerr)
}
var released bool
return func() {
if released {
return
}
released = true
if data, rerr := os.ReadFile(lockPath); rerr == nil && string(data) == token {
_ = lockutil.RemoveLockFile(lockPath)
}
}, nil
return func() { _ = lock.Release() }, nil
}
// On Windows a concurrent holder's os.Remove leaves the lock file in a
// "delete pending" state, so an O_EXCL create races it with
// ERROR_ACCESS_DENIED (os.ErrPermission) rather than ErrExist. Treat that
// as contention and retry, exactly like ErrExist — otherwise the lock
// spuriously fails under concurrency on Windows.
if !errors.Is(err, os.ErrExist) && !errors.Is(err, os.ErrPermission) {
if !errors.Is(err, lockutil.ErrLockHeld) {
return nil, fmt.Errorf("oauth: acquire token lock: %w", err)
}
// Reclaim a stale lock left by a crashed holder — atomically (H3). A blind
// Remove lets two racers both reclaim + recreate and so both hold the lock;
// reclaimStaleLock renames the file aside (only one rename wins) and restores
// it if it turns out fresh, so a live lock is never deleted out from under it.
if info, statErr := os.Stat(lockPath); statErr == nil && time.Since(info.ModTime()) > fileLockStaleAfter {
cleared, rerr := lockutil.ReclaimStaleLock(lockPath, token, func(reclaimedPath string) bool {
info, err := os.Stat(reclaimedPath)
return err == nil && time.Since(info.ModTime()) <= fileLockStaleAfter
})
if rerr != nil {
// Reclaim hit a hard failure: the rename aside failed outright, or a
// live holder's lock could not be put back (the lock path may be
// missing, so re-acquiring would break mutual exclusion). Fail closed
// instead of spinning to the deadline.
return nil, fmt.Errorf("oauth: reclaim stale token lock: %w", rerr)
}
if cleared {
continue
}
// Lost the reclaim race (or it was actually fresh) — fall through to the
// bounded wait rather than hot-spinning on a reclaim that never wins.
}
if now().After(deadline) {
return nil, fmt.Errorf("oauth: timed out acquiring token lock %s", filepath.Base(lockPath))
}
+22 -2
View File
@@ -191,6 +191,7 @@ type TurnBenchResult struct {
type TurnBenchTotals struct {
InputTokens int64 `json:"inputTokens"`
CachedInputTokens int64 `json:"cachedInputTokens"`
CacheWriteTokens int64 `json:"cacheWriteTokens"`
OutputTokens int64 `json:"outputTokens"`
ModelRequests int64 `json:"modelRequests"`
ToolCalls int64 `json:"toolCalls"`
@@ -479,6 +480,7 @@ func topLatencySources(perSpan map[string]SpanStats, top int) []LatencySource {
func aggregateTotals(totals *TurnBenchTotals, tr *trace.TurnTrace) {
totals.InputTokens += tr.Counter(trace.CounterInputTokens)
totals.CachedInputTokens += tr.Counter(trace.CounterCachedInputTokens)
totals.CacheWriteTokens += tr.Counter(trace.CounterCacheWriteTokens)
totals.OutputTokens += tr.Counter(trace.CounterOutputTokens)
totals.ModelRequests += tr.Counter(trace.CounterModelRequests)
totals.ToolCalls += tr.Counter(trace.CounterToolCalls)
@@ -531,8 +533,9 @@ func FormatTurnBenchSummary(result TurnBenchResult) string {
lines = append(lines, fmt.Sprintf(" %-18s %10s %5.1f%%", src.Span, FormatMetric(src.TotalMs, "ms"), src.Share*100))
}
}
lines = append(lines, fmt.Sprintf("totals: in=%d (cached %d) out=%d | requests=%d tools=%d retries=%d reconnects=%d compactions=%d",
result.Totals.InputTokens, result.Totals.CachedInputTokens, result.Totals.OutputTokens,
lines = append(lines, fmt.Sprintf("totals: in=%d (cache-read %d, cache-write %d, uncached %d) out=%d | requests=%d tools=%d retries=%d reconnects=%d compactions=%d",
result.Totals.InputTokens, result.Totals.CachedInputTokens, result.Totals.CacheWriteTokens,
uncachedInputTokens(result.Totals), result.Totals.OutputTokens,
result.Totals.ModelRequests, result.Totals.ToolCalls, result.Totals.Retries,
result.Totals.Reconnects, result.Totals.Compactions))
for _, class := range sortedClasses(result.PerClass) {
@@ -545,6 +548,23 @@ func FormatTurnBenchSummary(result TurnBenchResult) string {
return strings.Join(lines, "\n")
}
func uncachedInputTokens(totals TurnBenchTotals) int64 {
remaining := totals.InputTokens
if remaining <= 0 {
return 0
}
for _, cached := range []int64{totals.CachedInputTokens, totals.CacheWriteTokens} {
if cached <= 0 {
continue
}
if cached >= remaining {
return 0
}
remaining -= cached
}
return remaining
}
// classTier returns the oracle tier label a class was classified into, so the
// per-class roll-up states explicitly which oracle (if any) its "passed" count
// is measured against.
+33
View File
@@ -74,6 +74,7 @@ func TestRunTurnBenchAggregation(t *testing.T) {
"t3": cannedTrace(150, 20, 1500),
"t4": cannedTrace(300, 10, 1000),
}
canned["t1"].Counters = append(canned["t1"].Counters, trace.Counter{Name: trace.CounterCacheWriteTokens, Value: 125})
cfg := TurnBenchConfig{
Model: "fake-model",
Iterations: 1,
@@ -155,6 +156,9 @@ func TestRunTurnBenchAggregation(t *testing.T) {
if result.Totals.InputTokens != 5500 {
t.Fatalf("inputTokens = %d, want 5500", result.Totals.InputTokens)
}
if result.Totals.CacheWriteTokens != 125 {
t.Fatalf("cacheWriteTokens = %d, want 125", result.Totals.CacheWriteTokens)
}
if result.Totals.OutputTokens != 2750 {
t.Fatalf("outputTokens = %d, want 2750", result.Totals.OutputTokens)
}
@@ -334,6 +338,35 @@ func TestFormatTurnBenchSummaryNamesTopSources(t *testing.T) {
}
}
func TestUncachedInputTokensClampsInconsistentTotals(t *testing.T) {
tests := []struct {
name string
totals TurnBenchTotals
want int64
}{
{name: "normal", totals: TurnBenchTotals{InputTokens: 100, CachedInputTokens: 30, CacheWriteTokens: 20}, want: 50},
{name: "cached exceeds input", totals: TurnBenchTotals{InputTokens: 100, CachedInputTokens: 120}, want: 0},
{name: "write exceeds remainder", totals: TurnBenchTotals{InputTokens: 100, CachedInputTokens: 20, CacheWriteTokens: 90}, want: 0},
{name: "overflowing inconsistent counters", totals: TurnBenchTotals{CachedInputTokens: 1<<63 - 1, CacheWriteTokens: 1<<63 - 1}, want: 0},
{name: "overflowing counters with input", totals: TurnBenchTotals{InputTokens: 100, CachedInputTokens: 1<<63 - 1, CacheWriteTokens: 1<<63 - 1}, want: 0},
{name: "negative cache counters ignored", totals: TurnBenchTotals{InputTokens: 100, CachedInputTokens: -20, CacheWriteTokens: -10}, want: 100},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
if got := uncachedInputTokens(test.totals); got != test.want {
t.Fatalf("uncachedInputTokens(%+v) = %d, want %d", test.totals, got, test.want)
}
})
}
summary := FormatTurnBenchSummary(TurnBenchResult{Totals: TurnBenchTotals{
CachedInputTokens: 1<<63 - 1,
CacheWriteTokens: 1<<63 - 1,
}})
if !strings.Contains(summary, "uncached 0") {
t.Fatalf("summary did not clamp inconsistent cache totals:\n%s", summary)
}
}
func TestLoadBaselineManifest(t *testing.T) {
path := filepath.Join("manifests", "baseline.json")
set, err := LoadTaskSet(path)
+238 -60
View File
@@ -62,6 +62,7 @@ const (
responsesEventReasoningDelta = "response.reasoning_summary_text.delta"
responsesEventOutputTextDone = "response.output_text.done"
responsesEventFunctionArgsDelta = "response.function_call_arguments.delta"
responsesEventCustomInputDelta = "response.custom_tool_call_input.delta"
responsesEventContentPartDone = "response.content_part.done"
responsesEventOutputItemDone = "response.output_item.done"
responsesEventCompleted = "response.completed"
@@ -70,17 +71,37 @@ const (
responsesEventIncomplete = "response.incomplete"
)
const applyPatchLarkGrammar = `start: begin_patch hunk+ end_patch
begin_patch: "*** Begin Patch" LF
end_patch: "*** End Patch" LF?
hunk: add_hunk | delete_hunk | update_hunk
add_hunk: "*** Add File: " filename LF add_line+
delete_hunk: "*** Delete File: " filename LF
update_hunk: "*** Update File: " filename LF change_move? change?
filename: /(.+)/
add_line: "+" /(.*)/ LF -> line
change_move: "*** Move to: " filename LF
change: (change_context | change_line)+ eof_line?
change_context: ("@@" | "@@ " /(.+)/) LF
change_line: ("+" | "-" | " ") /(.*)/ LF
eof_line: "*** End of File" LF
%import common.LF`
// responsesRequest is the wire shape POSTed to {baseURL}/responses.
type responsesRequest struct {
Model string `json:"model"`
Instructions string `json:"instructions"`
Input []inputItem `json:"input"`
Stream bool `json:"stream"`
Store bool `json:"store"`
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
Tools []responsesTool `json:"tools,omitempty"`
Reasoning *responsesReasoning `json:"reasoning,omitempty"`
ServiceTier string `json:"service_tier,omitempty"`
Type string `json:"type,omitempty"`
Model string `json:"model"`
Instructions string `json:"instructions"`
PreviousResponseID string `json:"previous_response_id,omitempty"`
Input []inputItem `json:"input"`
Stream bool `json:"stream"`
Store bool `json:"store"`
ParallelToolCalls bool `json:"parallel_tool_calls"`
MaxOutputTokens int `json:"max_output_tokens,omitempty"`
Tools []responsesTool `json:"tools,omitempty"`
Reasoning *responsesReasoning `json:"reasoning,omitempty"`
ServiceTier string `json:"service_tier,omitempty"`
PromptCacheKey string `json:"prompt_cache_key,omitempty"`
}
// responsesReasoning carries the reasoning controls for the Responses API. The
@@ -105,6 +126,7 @@ type inputItem struct {
CallID string `json:"call_id,omitempty"`
Name string `json:"name,omitempty"`
Arguments string `json:"arguments,omitempty"`
Input string `json:"input,omitempty"`
// function_call_output fields
Output any `json:"output,omitempty"`
}
@@ -118,10 +140,11 @@ type contentItem struct {
}
type responsesTool struct {
Type string `json:"type"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
Parameters map[string]any `json:"parameters,omitempty"`
Type string `json:"type"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
Parameters map[string]any `json:"parameters,omitempty"`
Format *zeroruntime.ToolDefinitionFormat `json:"format,omitempty"`
}
// responsesEvent is the decoded form of one Responses SSE data payload.
@@ -132,6 +155,7 @@ type responsesEvent struct {
Delta string `json:"delta,omitempty"`
// item payloads (response.output_item.added / done)
ItemID string `json:"item_id,omitempty"`
CallID string `json:"call_id,omitempty"`
// OutputIndex is a *int so a real 0 (the first output) is distinguishable from
// "absent" — a plain int defaulting to 0 dropped a no-item_id call's args (M1).
OutputIndex *int `json:"output_index,omitempty"`
@@ -150,6 +174,7 @@ type itemPayload struct {
CallID string `json:"call_id,omitempty"`
Name string `json:"name,omitempty"`
Arguments string `json:"arguments,omitempty"`
Input string `json:"input,omitempty"`
Status string `json:"status,omitempty"`
}
@@ -173,7 +198,8 @@ type outputDetails struct {
}
type inputDetails struct {
CachedTokens int `json:"cached_tokens"`
CachedTokens int `json:"cached_tokens"`
CacheWriteTokens int `json:"cache_write_tokens"`
}
type errorPayload struct {
@@ -187,8 +213,10 @@ type errorPayload struct {
// response.output_item.done.
type toolCallBuilder struct {
ID string
CallID string
Name string
Arguments strings.Builder
Freeform bool
started bool
ended bool
}
@@ -197,15 +225,57 @@ type toolCallBuilder struct {
// event has been emitted. It is the Responses-API equivalent of the
// chat-completions toolState in provider.go.
type responsesState struct {
toolCalls map[string]*toolCallBuilder
usage *usagePayload
done bool
toolCalls map[string]*toolCallBuilder
toolOrder []string
text strings.Builder
responseID string
usage *usagePayload
emitted bool
done bool
}
func newResponsesState() *responsesState {
return &responsesState{toolCalls: map[string]*toolCallBuilder{}}
}
func (state *responsesState) outputInputItems() []inputItem {
items := make([]inputItem, 0, len(state.toolOrder)+1)
if state.text.Len() > 0 {
items = append(items, inputItem{
Type: "message",
Role: "assistant",
Content: []contentItem{{Type: "output_text", Text: state.text.String()}},
})
}
for _, key := range state.toolOrder {
builder := state.toolCalls[key]
if builder == nil || builder.ID == "" || builder.Name == "" {
continue
}
callID := builder.CallID
if callID == "" {
callID = builder.ID
}
itemType := "function_call"
if builder.Freeform {
itemType = "custom_tool_call"
}
item := inputItem{
Type: itemType,
ID: builder.ID,
CallID: callID,
Name: builder.Name,
}
if builder.Freeform {
item.Input = builder.Arguments.String()
} else {
item.Arguments = builder.Arguments.String()
}
items = append(items, item)
}
return items
}
// buildResponsesRequest converts a runtime CompletionRequest into the
// Responses API wire format. System messages fold into a top-level
// `instructions` field; everything else (user/assistant turns, tool
@@ -215,11 +285,14 @@ func (p *CodexProvider) buildResponsesRequest(request zeroruntime.CompletionRequ
return nil, errors.New("codex provider: model is required")
}
req := &responsesRequest{
Model: p.inner.model,
Stream: true,
MaxOutputTokens: p.inner.maxTokens,
Model: p.inner.model,
Stream: true,
ParallelToolCalls: true,
MaxOutputTokens: p.inner.maxTokens,
PromptCacheKey: request.PromptCacheKey,
}
instructions := []string{}
customCallIDs := map[string]string{}
for _, msg := range request.Messages {
switch msg.Role {
case zeroruntime.MessageRoleSystem:
@@ -234,13 +307,31 @@ func (p *CodexProvider) buildResponsesRequest(request zeroruntime.CompletionRequ
return nil, err
}
req.Input = append(req.Input, items...)
for _, call := range msg.ToolCalls {
if call.Freeform && call.ID != "" {
callID := call.ProviderCallID
if callID == "" {
callID = call.ID
}
customCallIDs[call.ID] = callID
}
}
case zeroruntime.MessageRoleTool:
if msg.ToolCallID == "" {
return nil, fmt.Errorf("codex provider: tool message missing tool call id")
}
outputType := "function_call_output"
callID := msg.ToolCallProviderID
if callID == "" {
callID = msg.ToolCallID
}
if customID, custom := customCallIDs[msg.ToolCallID]; custom {
outputType = "custom_tool_call_output"
callID = customID
}
req.Input = append(req.Input, inputItem{
Type: "function_call_output",
CallID: msg.ToolCallID,
Type: outputType,
CallID: callID,
Output: msg.Content,
})
default:
@@ -249,12 +340,25 @@ func (p *CodexProvider) buildResponsesRequest(request zeroruntime.CompletionRequ
}
req.Instructions = strings.Join(instructions, "\n\n")
for _, tool := range request.Tools {
req.Tools = append(req.Tools, responsesTool{
Type: "function",
Name: tool.Name,
Description: tool.Description,
Parameters: tool.Parameters,
})
definition := tool
if definition.Name == "apply_patch" && definition.Type == zeroruntime.ToolDefinitionFreeform {
definition.Description = "The apply_patch tool edits files from a raw structured patch. Use absolute file paths in patch headers to target a granted extra write root. Do not wrap the patch in JSON."
definition.Format = &zeroruntime.ToolDefinitionFormat{Type: "grammar", Syntax: "lark", Definition: applyPatchLarkGrammar}
}
wireTool := responsesTool{
Type: string(definition.Type),
Name: definition.Name,
Description: definition.Description,
Parameters: definition.Parameters,
Format: definition.Format,
}
if wireTool.Type == "" {
wireTool.Type = string(zeroruntime.ToolDefinitionFunction)
}
if wireTool.Type == string(zeroruntime.ToolDefinitionFreeform) {
wireTool.Parameters = nil
}
req.Tools = append(req.Tools, wireTool)
}
// Codex / o-series reasoning models take the effort tier nested under
// `reasoning` on the Responses API (the chat-completions `reasoning_effort`
@@ -312,13 +416,26 @@ func (p *CodexProvider) assistantInputItems(msg zeroruntime.Message) ([]inputIte
if tc.Name == "" {
return nil, errors.New("codex provider: assistant tool call missing name")
}
items = append(items, inputItem{
Type: "function_call",
ID: tc.ID,
CallID: tc.ID,
Name: tc.Name,
Arguments: tc.Arguments,
})
callID := tc.ProviderCallID
if callID == "" {
callID = tc.ID
}
itemType := "function_call"
if tc.Freeform {
itemType = "custom_tool_call"
}
item := inputItem{
Type: itemType,
ID: tc.ID,
CallID: callID,
Name: tc.Name,
}
if tc.Freeform {
item.Input = tc.Arguments
} else {
item.Arguments = tc.Arguments
}
items = append(items, item)
}
return items, nil
}
@@ -479,6 +596,18 @@ func (p *CodexProvider) emitResponsesEvent(
state.done = true
return false
}
return p.emitParsedResponsesEvent(ctx, &event, state, events)
}
// emitParsedResponsesEvent converts one already-decoded Responses event into
// zero or more runtime events. WebSocket sessions use this path to avoid
// decoding each frame twice; the SSE path retains emitResponsesEvent above.
func (p *CodexProvider) emitParsedResponsesEvent(
ctx context.Context,
event *responsesEvent,
state *responsesState,
events chan<- zeroruntime.StreamEvent,
) bool {
if event.Type == "" {
// A payload that parses as JSON but carries no `type` field is a
// malformed Responses event — the type is the discriminator that
@@ -500,12 +629,17 @@ func (p *CodexProvider) emitResponsesEvent(
// the runtime. response.incomplete is folded into a length finish at
// scan-end (state.done stays false so the wrapper emits the
// StreamEventDone with FinishReasonLength).
if event.Response != nil && event.Response.ID != "" {
state.responseID = event.Response.ID
}
return true
case responsesEventOutputItemAdded:
p.handleOutputItemAdded(ctx, &event, state, events)
p.handleOutputItemAdded(ctx, event, state, events)
return true
case responsesEventOutputTextDelta:
if event.Delta != "" {
state.emitted = true
state.text.WriteString(event.Delta)
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventText,
Content: event.Delta,
@@ -517,6 +651,7 @@ func (p *CodexProvider) emitResponsesEvent(
// phase shows progress (and keeps the activity clock fresh) instead of
// looking like a hang. Requested via reasoning.summary="auto".
if event.Delta != "" {
state.emitted = true
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventReasoning,
Content: event.Delta,
@@ -524,13 +659,16 @@ func (p *CodexProvider) emitResponsesEvent(
}
return true
case responsesEventFunctionArgsDelta:
p.handleFunctionArgsDelta(ctx, &event, state, events)
p.handleToolInputDelta(ctx, event, state, events, false)
return true
case responsesEventCustomInputDelta:
p.handleToolInputDelta(ctx, event, state, events, true)
return true
case responsesEventOutputItemDone:
p.handleOutputItemDone(ctx, &event, state, events)
p.handleOutputItemDone(ctx, event, state, events)
return true
case responsesEventCompleted, responsesEventFailed:
return p.handleTerminalResponse(ctx, &event, state, events)
return p.handleTerminalResponse(ctx, event, state, events)
case responsesEventError:
message := event.Message
if event.Code != "" {
@@ -570,22 +708,38 @@ func (p *CodexProvider) handleOutputItemAdded(
return
}
switch event.Item.Type {
case "function_call":
case "function_call", "custom_tool_call":
key := p.toolCallKey(event)
if key == "" {
return
}
builder := &toolCallBuilder{
ID: key,
Name: event.Item.Name,
builder, exists := state.toolCalls[key]
if !exists {
builder = &toolCallBuilder{ID: key, Freeform: event.Item.Type == "custom_tool_call"}
state.toolCalls[key] = builder
state.toolOrder = append(state.toolOrder, key)
}
if event.Item.CallID != "" {
builder.CallID = event.Item.CallID
}
if event.Item.Name != "" {
builder.Name = event.Item.Name
}
builder.Freeform = builder.Freeform || event.Item.Type == "custom_tool_call"
if builder.Freeform && event.Item.Input != "" && builder.Arguments.Len() == 0 {
builder.Arguments.WriteString(event.Item.Input)
}
state.emitted = true
if !builder.started {
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventToolCallStart,
ToolCallID: key,
ToolCallProviderID: builder.CallID,
ToolName: builder.Name,
ToolCallFreeform: builder.Freeform,
})
builder.started = true
}
state.toolCalls[key] = builder
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventToolCallStart,
ToolCallID: key,
ToolName: event.Item.Name,
})
builder.started = true
}
}
@@ -593,11 +747,12 @@ func (p *CodexProvider) handleOutputItemAdded(
// to the in-flight builder and emits StreamEventToolCallDelta. The
// Codex backend uses either `item_id` or `output_index` to attribute
// the delta; we honor both.
func (p *CodexProvider) handleFunctionArgsDelta(
func (p *CodexProvider) handleToolInputDelta(
ctx context.Context,
event *responsesEvent,
state *responsesState,
events chan<- zeroruntime.StreamEvent,
freeform bool,
) {
key := p.toolCallKey(event)
if key == "" {
@@ -610,10 +765,16 @@ func (p *CodexProvider) handleFunctionArgsDelta(
// an intermediate proxy). Treat the delta as the start of a new
// tool call with no known name yet — the eventual
// response.output_item.done carries the final name.
builder = &toolCallBuilder{ID: key}
builder = &toolCallBuilder{ID: key, Freeform: freeform}
state.toolCalls[key] = builder
state.toolOrder = append(state.toolOrder, key)
}
builder.Freeform = builder.Freeform || freeform
if event.CallID != "" && builder.CallID == "" {
builder.CallID = event.CallID
}
builder.Arguments.WriteString(event.Delta)
state.emitted = true
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventToolCallDelta,
ToolCallID: key,
@@ -632,7 +793,7 @@ func (p *CodexProvider) handleOutputItemDone(
state *responsesState,
events chan<- zeroruntime.StreamEvent,
) {
if event.Item == nil || event.Item.Type != "function_call" {
if event.Item == nil || (event.Item.Type != "function_call" && event.Item.Type != "custom_tool_call") {
return
}
key := p.toolCallKey(event)
@@ -641,20 +802,31 @@ func (p *CodexProvider) handleOutputItemDone(
}
builder, ok := state.toolCalls[key]
if !ok {
builder = &toolCallBuilder{ID: key}
builder = &toolCallBuilder{ID: key, Freeform: event.Item.Type == "custom_tool_call"}
state.toolCalls[key] = builder
state.toolOrder = append(state.toolOrder, key)
}
if event.Item.Name != "" {
builder.Name = event.Item.Name
}
if event.Item.Arguments != "" && builder.Arguments.Len() == 0 {
builder.Arguments.WriteString(event.Item.Arguments)
if event.Item.CallID != "" {
builder.CallID = event.Item.CallID
}
builder.Freeform = builder.Freeform || event.Item.Type == "custom_tool_call"
finalInput := event.Item.Arguments
if builder.Freeform {
finalInput = event.Item.Input
}
if finalInput != "" && builder.Arguments.Len() == 0 {
builder.Arguments.WriteString(finalInput)
}
if !builder.started {
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventToolCallStart,
ToolCallID: key,
ToolName: builder.Name,
Type: zeroruntime.StreamEventToolCallStart,
ToolCallID: key,
ToolCallProviderID: builder.CallID,
ToolName: builder.Name,
ToolCallFreeform: builder.Freeform,
})
builder.started = true
}
@@ -664,6 +836,7 @@ func (p *CodexProvider) handleOutputItemDone(
ToolName: builder.Name,
})
builder.ended = true
state.emitted = true
}
// handleTerminalResponse emits the final usage + done event for a
@@ -693,6 +866,7 @@ func (p *CodexProvider) handleTerminalResponse(
return false
}
state.usage = event.Response.Usage
state.responseID = event.Response.ID
if event.Response.Usage != nil {
usage := zeroruntime.Usage{
InputTokens: event.Response.Usage.InputTokens,
@@ -700,6 +874,7 @@ func (p *CodexProvider) handleTerminalResponse(
}
if event.Response.Usage.InputTokensDetails != nil {
usage.CachedInputTokens = event.Response.Usage.InputTokensDetails.CachedTokens
usage.CacheWriteTokens = event.Response.Usage.InputTokensDetails.CacheWriteTokens
}
if event.Response.Usage.OutputTokensDetails != nil {
usage.ReasoningTokens = event.Response.Usage.OutputTokensDetails.ReasoningTokens
@@ -735,7 +910,7 @@ func (p *CodexProvider) handleTerminalResponse(
state.done = true
return false
}
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone})
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{Type: zeroruntime.StreamEventDone, ResponseID: event.Response.ID})
state.done = true
return false
}
@@ -754,6 +929,9 @@ func (p *CodexProvider) toolCallKey(event *responsesEvent) string {
if event.Item != nil && event.Item.CallID != "" {
return event.Item.CallID
}
if event.CallID != "" {
return event.CallID
}
if event.OutputIndex != nil {
// output_index is present (a *int distinguishes a real 0 — the first output
// in the response — from "absent"), so key on it even when 0 instead of
+421
View File
@@ -0,0 +1,421 @@
package openai
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/url"
"reflect"
"strings"
"sync"
"time"
"github.com/coder/websocket"
"github.com/Gitlawb/zero/internal/providers/providerio"
"github.com/Gitlawb/zero/internal/trace"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
const (
responsesWebSocketBetaHeader = "OpenAI-Beta"
responsesWebSocketBetaValue = "responses_websockets=2026-02-06"
responsesWebSocketRequest = "response.create"
previousResponseNotFoundCode = "previous_response_not_found"
webSocketConnectionLimitCode = "websocket_connection_limit_reached"
responsesWebSocketReadLimit = 16 << 20
)
var errCodexTurnSessionBusy = errors.New("codex turn session already has an active stream")
// NewCodexTurnSessionProvider enables one Responses WebSocket connection for
// each agent run. The session falls back to the existing HTTP/SSE provider when
// the endpoint cannot upgrade or loses its response chain.
func NewCodexTurnSessionProvider(provider *CodexProvider, caps zeroruntime.ProviderCapabilities) zeroruntime.TurnSessionProvider {
return &codexTurnSessionProvider{provider: provider, caps: caps}
}
type codexTurnSessionProvider struct {
provider *CodexProvider
caps zeroruntime.ProviderCapabilities
}
func (provider *codexTurnSessionProvider) OpenTurnSession(ctx context.Context) (zeroruntime.TurnSession, error) {
if provider == nil || provider.provider == nil || provider.provider.inner == nil {
return nil, errors.New("codex turn session requires a provider")
}
sessionCtx, cancel := context.WithCancel(ctx)
return &codexTurnSession{
ctx: sessionCtx,
cancel: cancel,
provider: provider.provider,
prewarmDone: make(chan struct{}),
}, nil
}
func (provider *codexTurnSessionProvider) Capabilities() zeroruntime.ProviderCapabilities {
return provider.caps
}
type codexTurnSession struct {
ctx context.Context
cancel context.CancelFunc
provider *CodexProvider
mu sync.Mutex
connection *websocket.Conn
webSocketOff bool
lastRequest *responsesRequest
lastResponseID string
lastOutput []inputItem
streamActive bool
prewarmOnce sync.Once
prewarmDone chan struct{}
closeOnce sync.Once
}
// Prewarm starts the WebSocket handshake without blocking prompt assembly. The
// first Stream waits for this bounded attempt; failure simply selects HTTP/SSE.
func (session *codexTurnSession) Prewarm(ctx context.Context) error {
session.prewarmOnce.Do(func() {
go func() {
defer close(session.prewarmDone)
prewarmCtx, cancel := context.WithTimeout(session.ctx, prewarmTimeout)
defer cancel()
recorder := trace.FromContext(ctx)
span := recorder.Span(trace.SpanProviderPrewarm)
defer span.End()
recorder.Counter(trace.CounterPrewarmAttempts, 1)
connection, err := session.provider.dialResponsesWebSocket(prewarmCtx)
session.mu.Lock()
defer session.mu.Unlock()
if err != nil {
session.webSocketOff = true
return
}
session.connection = connection
}()
})
return nil
}
func (session *codexTurnSession) Stream(ctx context.Context, request zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) {
_ = session.Prewarm(ctx)
select {
case <-session.prewarmDone:
case <-ctx.Done():
return nil, ctx.Err()
}
fullRequest, err := session.provider.buildResponsesRequest(request)
if err != nil {
return nil, err
}
session.mu.Lock()
if session.streamActive {
session.mu.Unlock()
return nil, errCodexTurnSessionBusy
}
if session.webSocketOff || session.connection == nil {
session.mu.Unlock()
trace.FromContext(ctx).Counter(trace.CounterResponsesHTTPFallback, 1)
return session.provider.StreamCompletion(ctx, request)
}
session.streamActive = true
connection := session.connection
wireRequest := *fullRequest
wireRequest.Type = responsesWebSocketRequest
if delta, ok := responseInputDelta(session.lastRequest, session.lastOutput, fullRequest); ok && session.lastResponseID != "" {
wireRequest.PreviousResponseID = session.lastResponseID
wireRequest.Input = delta
trace.FromContext(ctx).Counter(trace.CounterResponseChainReused, 1)
} else if session.lastRequest != nil {
trace.FromContext(ctx).Counter(trace.CounterResponseChainReset, 1)
session.lastRequest = nil
session.lastResponseID = ""
session.lastOutput = nil
}
session.mu.Unlock()
body, err := json.Marshal(wireRequest)
if err != nil {
session.releaseStream()
return nil, fmt.Errorf("encode codex websocket request: %w", err)
}
events := make(chan zeroruntime.StreamEvent, 16)
go func() {
defer close(events)
defer session.releaseStream()
session.streamWebSocket(ctx, connection, body, fullRequest, request, events)
}()
return events, nil
}
func (session *codexTurnSession) releaseStream() {
session.mu.Lock()
session.streamActive = false
session.mu.Unlock()
}
func (session *codexTurnSession) Compact(context.Context, zeroruntime.CompletionRequest) ([]zeroruntime.Message, error) {
return nil, zeroruntime.ErrCompactionUnsupported
}
func (session *codexTurnSession) Close() error {
session.closeOnce.Do(func() {
session.cancel()
session.prewarmOnce.Do(func() { close(session.prewarmDone) })
select {
case <-session.prewarmDone:
case <-time.After(prewarmTimeout):
}
session.mu.Lock()
connection := session.connection
session.connection = nil
session.webSocketOff = true
session.mu.Unlock()
if connection != nil {
_ = connection.CloseNow()
}
})
return nil
}
func responseInputDelta(previous *responsesRequest, output []inputItem, current *responsesRequest) ([]inputItem, bool) {
if previous == nil || current == nil || !responsesRequestPropertiesEqual(previous, current) {
return nil, false
}
prefix := make([]inputItem, 0, len(previous.Input)+len(output))
prefix = append(prefix, previous.Input...)
prefix = append(prefix, output...)
if len(current.Input) <= len(prefix) || !reflect.DeepEqual(current.Input[:len(prefix)], prefix) {
return nil, false
}
return append([]inputItem(nil), current.Input[len(prefix):]...), true
}
func responsesRequestPropertiesEqual(left, right *responsesRequest) bool {
return left.Model == right.Model &&
left.Instructions == right.Instructions &&
left.Stream == right.Stream &&
left.Store == right.Store &&
left.ParallelToolCalls == right.ParallelToolCalls &&
left.MaxOutputTokens == right.MaxOutputTokens &&
reflect.DeepEqual(left.Tools, right.Tools) &&
reflect.DeepEqual(left.Reasoning, right.Reasoning) &&
left.ServiceTier == right.ServiceTier &&
left.PromptCacheKey == right.PromptCacheKey
}
func (session *codexTurnSession) streamWebSocket(
ctx context.Context,
connection *websocket.Conn,
body []byte,
fullRequest *responsesRequest,
runtimeRequest zeroruntime.CompletionRequest,
events chan<- zeroruntime.StreamEvent,
) {
if err := connection.Write(ctx, websocket.MessageText, body); err != nil {
session.disableWebSocket(connection)
session.forwardHTTP(ctx, runtimeRequest, events)
return
}
state := newResponsesState()
for {
readCtx := ctx
cancel := func() {}
idleTimeout := session.provider.inner.streamIdleTimeout
if idleTimeout > 0 {
readCtx, cancel = context.WithTimeout(ctx, idleTimeout)
}
messageType, data, err := connection.Read(readCtx)
readTimedOut := errors.Is(readCtx.Err(), context.DeadlineExceeded) && ctx.Err() == nil
cancel()
if err != nil {
session.disableWebSocket(connection)
if ctx.Err() != nil {
return
}
if !state.emitted {
session.forwardHTTP(ctx, runtimeRequest, events)
return
}
errorMessage := err.Error()
if readTimedOut {
errorMessage = providerio.StreamTimeoutMessage(providerio.ErrStreamIdle, idleTimeout)
}
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventError,
Error: session.provider.redact("provider stream error: " + errorMessage),
})
return
}
if messageType != websocket.MessageText {
continue
}
var responseEvent responsesEvent
decodeErr := json.Unmarshal(data, &responseEvent)
if decodeErr == nil {
if (responseEvent.Code == previousResponseNotFoundCode || responseEvent.Code == webSocketConnectionLimitCode) &&
!state.emitted {
session.disableWebSocket(connection)
session.forwardHTTP(ctx, runtimeRequest, events)
return
}
if responseEvent.Type == responsesEventIncomplete {
session.resetResponseChain(connection)
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventDone,
FinishReason: zeroruntime.FinishReasonLength,
})
return
}
}
keepReading := false
if decodeErr == nil {
keepReading = session.provider.emitParsedResponsesEvent(ctx, &responseEvent, state, events)
} else {
keepReading = session.provider.emitResponsesEvent(ctx, string(data), state, events)
}
if keepReading {
continue
}
if responseEvent.Type == responsesEventCompleted && responseEvent.Response != nil &&
responseEvent.Response.Error == nil && responseEvent.Response.Status != "failed" && state.responseID != "" {
session.mu.Lock()
if session.connection == connection && !session.webSocketOff {
requestCopy := *fullRequest
requestCopy.Input = append([]inputItem(nil), fullRequest.Input...)
session.lastRequest = &requestCopy
session.lastResponseID = state.responseID
session.lastOutput = state.outputInputItems()
}
session.mu.Unlock()
}
return
}
}
func (session *codexTurnSession) resetResponseChain(connection *websocket.Conn) {
session.mu.Lock()
defer session.mu.Unlock()
if session.connection != connection {
return
}
session.lastRequest = nil
session.lastResponseID = ""
session.lastOutput = nil
}
func (session *codexTurnSession) forwardHTTP(ctx context.Context, request zeroruntime.CompletionRequest, events chan<- zeroruntime.StreamEvent) {
trace.FromContext(ctx).Counter(trace.CounterResponsesHTTPFallback, 1)
stream, err := session.provider.StreamCompletion(ctx, request)
if err != nil {
providerio.SendEvent(ctx, events, zeroruntime.StreamEvent{
Type: zeroruntime.StreamEventError,
Error: session.provider.redact(err.Error()),
})
return
}
for event := range stream {
providerio.SendEvent(ctx, events, event)
}
}
func (session *codexTurnSession) disableWebSocket(connection *websocket.Conn) {
session.mu.Lock()
if session.connection == connection {
session.connection = nil
session.webSocketOff = true
session.lastRequest = nil
session.lastResponseID = ""
session.lastOutput = nil
}
session.mu.Unlock()
_ = connection.CloseNow()
}
func (provider *CodexProvider) dialResponsesWebSocket(ctx context.Context) (*websocket.Conn, error) {
endpoint, err := responsesWebSocketURL(provider.inner.endpoint)
if err != nil {
return nil, err
}
for authTry := 0; authTry < 2; authTry++ {
headers, err := provider.responsesWebSocketHeaders(ctx, authTry == 1)
if err != nil {
return nil, err
}
connection, response, err := websocket.Dial(ctx, endpoint, &websocket.DialOptions{
HTTPClient: provider.inner.httpClient,
HTTPHeader: headers,
CompressionMode: websocket.CompressionNoContextTakeover,
})
if err == nil {
connection.SetReadLimit(responsesWebSocketReadLimit)
return connection, nil
}
if response == nil || response.StatusCode != http.StatusUnauthorized || provider.inner.oauthResolver == nil || authTry == 1 {
return nil, err
}
}
return nil, errors.New("codex websocket authentication failed")
}
func (provider *CodexProvider) responsesWebSocketHeaders(ctx context.Context, forceRefresh bool) (http.Header, error) {
request, err := http.NewRequestWithContext(ctx, http.MethodGet, provider.inner.endpoint, nil)
if err != nil {
return nil, err
}
base := providerio.AuthHeaders{
APIKey: provider.inner.apiKey,
DefaultAuthHeader: "Authorization",
DefaultAuthScheme: "Bearer",
AuthHeader: provider.inner.authHeader,
AuthScheme: provider.inner.authScheme,
AuthHeaderValue: provider.inner.authHeaderValue,
CustomHeaders: provider.inner.customHeaders,
}
if resolver := provider.inner.oauthResolver; resolver != nil {
header, value, ok, err := resolver(ctx, forceRefresh)
if err != nil {
return nil, err
}
if ok {
providerio.ApplyAuthHeaders(request, providerio.AuthHeaders{CustomHeaders: provider.inner.customHeaders})
if strings.TrimSpace(header) == "" {
header = "Authorization"
}
request.Header.Set(header, value)
} else {
providerio.ApplyAuthHeaders(request, base)
}
} else {
providerio.ApplyAuthHeaders(request, base)
}
provider.injectCodexHeaders(request)
request.Header.Set(responsesWebSocketBetaHeader, responsesWebSocketBetaValue)
return request.Header.Clone(), nil
}
func responsesWebSocketURL(endpoint string) (string, error) {
parsed, err := url.Parse(strings.TrimSpace(endpoint))
if err != nil {
return "", err
}
switch parsed.Scheme {
case "http":
parsed.Scheme = "ws"
case "https":
parsed.Scheme = "wss"
default:
return "", fmt.Errorf("unsupported Responses websocket URL scheme %q", parsed.Scheme)
}
return parsed.String(), nil
}
@@ -0,0 +1,590 @@
package openai
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
"github.com/coder/websocket"
"github.com/Gitlawb/zero/internal/trace"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
func TestCodexTurnSessionChainsOnlyAppendOnlyRequests(t *testing.T) {
var mu sync.Mutex
requests := []map[string]any{}
betaHeader := ""
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if !strings.EqualFold(request.Header.Get("Upgrade"), "websocket") {
http.Error(writer, "websocket required", http.StatusUpgradeRequired)
return
}
mu.Lock()
betaHeader = request.Header.Get(responsesWebSocketBetaHeader)
mu.Unlock()
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
defer connection.CloseNow()
for index := 0; index < 3; index++ {
_, body, err := connection.Read(request.Context())
if err != nil {
return
}
var payload map[string]any
if err := json.Unmarshal(body, &payload); err != nil {
return
}
mu.Lock()
requests = append(requests, payload)
mu.Unlock()
switch index {
case 0:
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.created","response":{"id":"resp-1","status":"in_progress"}}`,
`{"type":"response.output_item.added","item_id":"call-1","item":{"type":"function_call","id":"call-1","call_id":"call-1","name":"read_file"}}`,
`{"type":"response.function_call_arguments.delta","item_id":"call-1","delta":"{}"}`,
`{"type":"response.output_item.done","item_id":"call-1","item":{"type":"function_call","id":"call-1","call_id":"call-1","name":"read_file","arguments":"{}"}}`,
`{"type":"response.completed","response":{"id":"resp-1","status":"completed"}}`,
)
case 1:
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.output_text.delta","delta":"done"}`,
`{"type":"response.completed","response":{"id":"resp-2","status":"completed"}}`,
)
case 2:
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.completed","response":{"id":"resp-3","status":"completed"}}`,
)
}
}
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
// A disabled stream watchdog must use the parent context rather than creating
// an immediately-expired read deadline.
provider.inner.streamIdleTimeout = 0
session := openCodexSession(t, provider)
defer session.Close()
recorder := trace.NewRecorder("session-1", "run-1", "test")
recorder.Start()
streamContext := trace.WithContext(t.Context(), recorder)
base := zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{
{Role: zeroruntime.MessageRoleSystem, Content: "Use tools carefully."},
{Role: zeroruntime.MessageRoleUser, Content: "Read the file."},
},
Tools: []zeroruntime.ToolDefinition{{Name: "read_file", Parameters: map[string]any{"type": "object"}}},
PromptCacheKey: "session-1",
}
first := collectCodexSessionEvents(t, session, base)
if !hasToolCallEnd(first, "call-1") {
t.Fatalf("first stream did not produce the tool call: %#v", first)
}
second := base
second.Messages = append(append([]zeroruntime.Message(nil), base.Messages...),
zeroruntime.Message{Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{{ID: "call-1", Name: "read_file", Arguments: "{}"}}},
zeroruntime.Message{Role: zeroruntime.MessageRoleTool, ToolCallID: "call-1", Content: "contents"},
)
secondEvents := collectCodexSessionEventsContext(t, streamContext, session, second)
if got := joinedText(secondEvents); got != "done" {
t.Fatalf("second stream text = %q, want done", got)
}
third := second
third.ReasoningEffort = "high"
third.Messages = append(append([]zeroruntime.Message(nil), second.Messages...),
zeroruntime.Message{Role: zeroruntime.MessageRoleAssistant, Content: "done"},
zeroruntime.Message{Role: zeroruntime.MessageRoleUser, Content: "Continue."},
)
collectCodexSessionEventsContext(t, streamContext, session, third)
finishedTrace := recorder.Finish()
if got := finishedTrace.Counter(trace.CounterResponseChainReused); got != 1 {
t.Fatalf("response-chain reuse count = %d, want 1", got)
}
if got := finishedTrace.Counter(trace.CounterResponseChainReset); got != 1 {
t.Fatalf("response-chain reset count = %d, want 1", got)
}
mu.Lock()
defer mu.Unlock()
if len(requests) != 3 {
t.Fatalf("websocket requests = %d, want 3", len(requests))
}
if betaHeader != responsesWebSocketBetaValue {
t.Fatalf("OpenAI-Beta = %q, want %q", betaHeader, responsesWebSocketBetaValue)
}
if requests[0]["type"] != responsesWebSocketRequest || requests[0]["previous_response_id"] != nil {
t.Fatalf("first request chaining fields = %#v", requests[0])
}
if requests[0]["parallel_tool_calls"] != true || requests[0]["prompt_cache_key"] != "session-1" {
t.Fatalf("first request stable fields = %#v", requests[0])
}
if requests[1]["previous_response_id"] != "resp-1" {
t.Fatalf("second previous_response_id = %#v", requests[1]["previous_response_id"])
}
secondInput, _ := requests[1]["input"].([]any)
if len(secondInput) != 1 || secondInput[0].(map[string]any)["type"] != "function_call_output" {
t.Fatalf("second incremental input = %#v", requests[1]["input"])
}
if requests[2]["previous_response_id"] != nil {
t.Fatalf("incompatible request reused response chain: %#v", requests[2])
}
thirdInput, _ := requests[2]["input"].([]any)
if len(thirdInput) <= 1 {
t.Fatalf("incompatible request did not send full input: %#v", thirdInput)
}
}
func TestCodexTurnSessionFallsBackWhenPreviousResponseIsMissing(t *testing.T) {
var mu sync.Mutex
webSocketRequests := 0
httpRequests := 0
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if strings.EqualFold(request.Header.Get("Upgrade"), "websocket") {
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
defer connection.CloseNow()
for index := 0; index < 2; index++ {
if _, _, err := connection.Read(request.Context()); err != nil {
return
}
mu.Lock()
webSocketRequests++
mu.Unlock()
if index == 0 {
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.output_item.added","item_id":"call-1","item":{"type":"function_call","id":"call-1","call_id":"call-1","name":"read_file"}}`,
`{"type":"response.output_item.done","item_id":"call-1","item":{"type":"function_call","id":"call-1","call_id":"call-1","name":"read_file","arguments":"{}"}}`,
`{"type":"response.completed","response":{"id":"resp-1","status":"completed"}}`,
)
continue
}
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.error","code":"previous_response_not_found","message":"missing"}`,
)
}
return
}
mu.Lock()
httpRequests++
mu.Unlock()
writer.Header().Set("Content-Type", "text/event-stream")
_, _ = fmt.Fprint(writer, "data: {\"type\":\"response.output_text.delta\",\"delta\":\"fallback\"}\n\n")
_, _ = fmt.Fprint(writer, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-http\",\"status\":\"completed\"}}\n\n")
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
session := openCodexSession(t, provider)
defer session.Close()
base := zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Read it."}},
Tools: []zeroruntime.ToolDefinition{{Name: "read_file", Parameters: map[string]any{"type": "object"}}},
}
collectCodexSessionEvents(t, session, base)
second := base
second.Messages = append(append([]zeroruntime.Message(nil), base.Messages...),
zeroruntime.Message{Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{{ID: "call-1", Name: "read_file", Arguments: "{}"}}},
zeroruntime.Message{Role: zeroruntime.MessageRoleTool, ToolCallID: "call-1", Content: "contents"},
)
events := collectCodexSessionEvents(t, session, second)
if got := joinedText(events); got != "fallback" {
t.Fatalf("fallback stream text = %q, want fallback; events=%#v", got, events)
}
for _, event := range events {
if event.Type == zeroruntime.StreamEventError {
t.Fatalf("response-chain recovery leaked an error: %#v", event)
}
}
mu.Lock()
defer mu.Unlock()
if webSocketRequests != 2 || httpRequests != 1 {
t.Fatalf("requests = websocket:%d http:%d, want 2/1", webSocketRequests, httpRequests)
}
}
func TestCodexTurnSessionFallsBackWhenSocketClosesBeforeOutput(t *testing.T) {
var mu sync.Mutex
httpRequests := 0
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if strings.EqualFold(request.Header.Get("Upgrade"), "websocket") {
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
if _, _, err := connection.Read(request.Context()); err == nil {
_ = connection.CloseNow()
}
return
}
mu.Lock()
httpRequests++
mu.Unlock()
writer.Header().Set("Content-Type", "text/event-stream")
_, _ = fmt.Fprint(writer, "data: {\"type\":\"response.output_text.delta\",\"delta\":\"recovered\"}\n\n")
_, _ = fmt.Fprint(writer, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-http\",\"status\":\"completed\"}}\n\n")
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
session := openCodexSession(t, provider)
defer session.Close()
recorder := trace.NewRecorder("session-fallback", "run-fallback", "test")
recorder.Start()
events := collectCodexSessionEventsContext(t, trace.WithContext(t.Context(), recorder), session, zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Hello"}},
})
if got := joinedText(events); got != "recovered" {
t.Fatalf("fallback text = %q, want recovered", got)
}
mu.Lock()
defer mu.Unlock()
if httpRequests != 1 {
t.Fatalf("HTTP fallback requests = %d, want 1", httpRequests)
}
if got := recorder.Finish().Counter(trace.CounterResponsesHTTPFallback); got != 1 {
t.Fatalf("responses HTTP fallback counter = %d, want 1", got)
}
}
func TestCodexTurnSessionKeepsWebSocketAfterIncompleteResponse(t *testing.T) {
var mu sync.Mutex
webSocketRequests := 0
httpRequests := 0
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if strings.EqualFold(request.Header.Get("Upgrade"), "websocket") {
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
defer connection.CloseNow()
for index := 0; index < 2; index++ {
if _, _, err := connection.Read(request.Context()); err != nil {
return
}
mu.Lock()
webSocketRequests++
mu.Unlock()
if index == 0 {
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.incomplete","response":{"id":"resp-1","status":"incomplete"}}`,
)
continue
}
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.output_text.delta","delta":"websocket"}`,
`{"type":"response.completed","response":{"id":"resp-2","status":"completed"}}`,
)
}
return
}
mu.Lock()
httpRequests++
mu.Unlock()
writer.Header().Set("Content-Type", "text/event-stream")
_, _ = fmt.Fprint(writer, "data: {\"type\":\"response.output_text.delta\",\"delta\":\"fallback\"}\n\n")
_, _ = fmt.Fprint(writer, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-http\",\"status\":\"completed\"}}\n\n")
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
session := openCodexSession(t, provider)
defer session.Close()
request := zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Continue."}},
}
first := collectCodexSessionEvents(t, session, request)
if len(first) != 1 || first[0].Type != zeroruntime.StreamEventDone || first[0].FinishReason != zeroruntime.FinishReasonLength {
t.Fatalf("incomplete response events = %#v, want length completion", first)
}
second := collectCodexSessionEvents(t, session, request)
if got := joinedText(second); got != "websocket" {
t.Fatalf("second turn text = %q, want websocket; events=%#v", got, second)
}
mu.Lock()
defer mu.Unlock()
if webSocketRequests != 2 || httpRequests != 0 {
t.Fatalf("requests = websocket:%d http:%d, want 2/0", webSocketRequests, httpRequests)
}
}
func TestCodexTurnSessionDoesNotReplayAfterVisibleOutput(t *testing.T) {
var mu sync.Mutex
httpRequests := 0
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
if strings.EqualFold(request.Header.Get("Upgrade"), "websocket") {
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
if _, _, err := connection.Read(request.Context()); err == nil {
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.output_text.delta","delta":"visible"}`,
)
}
_ = connection.Close(websocket.StatusInternalError, "interrupted")
return
}
mu.Lock()
httpRequests++
mu.Unlock()
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
session := openCodexSession(t, provider)
defer session.Close()
stream, err := session.Stream(t.Context(), zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Hello"}},
})
if err != nil {
t.Fatalf("Stream: %v", err)
}
sawText, sawError := false, false
for event := range stream {
sawText = sawText || event.Type == zeroruntime.StreamEventText
sawError = sawError || event.Type == zeroruntime.StreamEventError
}
if !sawText || !sawError {
t.Fatalf("visible interrupted stream = text:%v error:%v", sawText, sawError)
}
mu.Lock()
defer mu.Unlock()
if httpRequests != 0 {
t.Fatalf("visible interrupted stream was replayed over HTTP %d time(s)", httpRequests)
}
}
func TestCodexTurnSessionReportsWebSocketIdleTimeout(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
defer connection.CloseNow()
if _, _, err := connection.Read(request.Context()); err != nil {
return
}
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.output_text.delta","delta":"visible"}`,
)
<-request.Context().Done()
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
provider.inner.streamIdleTimeout = 200 * time.Millisecond
session := openCodexSession(t, provider)
defer session.Close()
stream, err := session.Stream(t.Context(), zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Hello"}},
})
if err != nil {
t.Fatalf("Stream: %v", err)
}
var streamError string
for event := range stream {
if event.Type == zeroruntime.StreamEventError {
streamError = event.Error
break
}
}
if !strings.Contains(streamError, "idle timeout after 200ms") {
t.Fatalf("stream error = %q, want idle-timeout detail", streamError)
}
}
func TestCodexTurnSessionRejectsOverlappingStreams(t *testing.T) {
started := make(chan struct{})
release := make(chan struct{})
var releaseOnce sync.Once
releaseStream := func() { releaseOnce.Do(func() { close(release) }) }
t.Cleanup(releaseStream)
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
connection, err := websocket.Accept(writer, request, nil)
if err != nil {
return
}
defer connection.CloseNow()
if _, _, err := connection.Read(request.Context()); err != nil {
return
}
close(started)
<-release
writeWebSocketEvents(request.Context(), connection,
`{"type":"response.completed","response":{"id":"resp-1","status":"completed"}}`,
)
}))
defer server.Close()
provider := newCodexSessionTestProvider(t, server)
session := openCodexSession(t, provider)
defer session.Close()
request := zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Hello"}},
}
first, err := session.Stream(t.Context(), request)
if err != nil {
t.Fatalf("first Stream: %v", err)
}
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("first WebSocket stream did not start")
}
if _, err := session.Stream(t.Context(), request); !errors.Is(err, errCodexTurnSessionBusy) {
t.Fatalf("overlapping Stream error = %v, want %v", err, errCodexTurnSessionBusy)
}
releaseStream()
for range first {
}
}
func TestCodexTurnSessionHTTPFallbackRedactsSetupError(t *testing.T) {
const secret = "sk-secret-fallback"
provider, err := NewCodexProvider(CodexOptions{Options: Options{
APIKey: secret,
BaseURL: "https://chatgpt.example/backend-api/codex",
Model: "gpt-test",
}})
if err != nil {
t.Fatalf("NewCodexProvider: %v", err)
}
session := &codexTurnSession{provider: provider}
events := make(chan zeroruntime.StreamEvent, 1)
session.forwardHTTP(t.Context(), zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRole(secret), Content: "invalid"}},
}, events)
select {
case event := <-events:
if event.Type != zeroruntime.StreamEventError {
t.Fatalf("fallback event type = %s, want error", event.Type)
}
if strings.Contains(event.Error, secret) {
t.Fatalf("fallback error leaked credential: %q", event.Error)
}
case <-time.After(time.Second):
t.Fatal("fallback did not emit setup error")
}
}
func TestCodexTurnSessionCloseDoesNotStartPrewarm(t *testing.T) {
provider, err := NewCodexProvider(CodexOptions{Options: Options{
APIKey: "test-token",
BaseURL: "https://chatgpt.example/backend-api/codex",
Model: "gpt-test",
}})
if err != nil {
t.Fatalf("NewCodexProvider: %v", err)
}
session, err := NewCodexTurnSessionProvider(provider, zeroruntime.ProviderCapabilities{}).OpenTurnSession(t.Context())
if err != nil {
t.Fatalf("OpenTurnSession: %v", err)
}
done := make(chan error, 1)
go func() { done <- session.Close() }()
select {
case err := <-done:
if err != nil {
t.Fatalf("Close() error = %v", err)
}
case <-time.After(200 * time.Millisecond):
t.Fatal("Close blocked without a prewarm attempt")
}
}
func newCodexSessionTestProvider(t *testing.T, server *httptest.Server) *CodexProvider {
t.Helper()
provider, err := NewCodexProvider(CodexOptions{
Options: Options{
APIKey: "test-token",
BaseURL: server.URL,
Model: "gpt-test",
HTTPClient: server.Client(),
},
AccountID: "account-test",
})
if err != nil {
t.Fatalf("NewCodexProvider: %v", err)
}
return provider
}
func openCodexSession(t *testing.T, provider *CodexProvider) zeroruntime.TurnSession {
t.Helper()
session, err := NewCodexTurnSessionProvider(provider, zeroruntime.ProviderCapabilities{}).OpenTurnSession(t.Context())
if err != nil {
t.Fatalf("OpenTurnSession: %v", err)
}
if err := session.Prewarm(t.Context()); err != nil {
t.Fatalf("Prewarm: %v", err)
}
return session
}
func collectCodexSessionEvents(t *testing.T, session zeroruntime.TurnSession, request zeroruntime.CompletionRequest) []zeroruntime.StreamEvent {
t.Helper()
return collectCodexSessionEventsContext(t, t.Context(), session, request)
}
func collectCodexSessionEventsContext(t *testing.T, ctx context.Context, session zeroruntime.TurnSession, request zeroruntime.CompletionRequest) []zeroruntime.StreamEvent {
t.Helper()
stream, err := session.Stream(ctx, request)
if err != nil {
t.Fatalf("Stream: %v", err)
}
events := []zeroruntime.StreamEvent{}
for event := range stream {
events = append(events, event)
if event.Type == zeroruntime.StreamEventError {
t.Fatalf("stream error: %s", event.Error)
}
}
return events
}
func writeWebSocketEvents(ctx context.Context, connection *websocket.Conn, events ...string) {
for _, event := range events {
if err := connection.Write(ctx, websocket.MessageText, []byte(event)); err != nil {
return
}
}
}
func hasToolCallEnd(events []zeroruntime.StreamEvent, id string) bool {
for _, event := range events {
if event.Type == zeroruntime.StreamEventToolCallEnd && event.ToolCallID == id {
return true
}
}
return false
}
func joinedText(events []zeroruntime.StreamEvent) string {
var builder strings.Builder
for _, event := range events {
if event.Type == zeroruntime.StreamEventText {
builder.WriteString(event.Content)
}
}
return builder.String()
}
@@ -2,6 +2,8 @@ package openai
import (
"context"
"reflect"
"strings"
"testing"
"github.com/Gitlawb/zero/internal/zeroruntime"
@@ -25,6 +27,126 @@ func TestToolCallKeyOutputIndexZero(t *testing.T) {
}
}
func TestResponsesStatePreservesDistinctCallID(t *testing.T) {
provider := &CodexProvider{}
state := newResponsesState()
events := make(chan zeroruntime.StreamEvent, 4)
provider.handleOutputItemAdded(context.Background(), &responsesEvent{
ItemID: "item-1",
Item: &itemPayload{
Type: "function_call",
ID: "item-1",
CallID: "call-1",
Name: "read_file",
},
}, state, events)
items := state.outputInputItems()
if len(items) != 1 || items[0].ID != "item-1" || items[0].CallID != "call-1" {
t.Fatalf("captured function call = %#v, want item id item-1 and call id call-1", items)
}
}
func TestCodexCustomApplyPatchStreamProducesFreeformCall(t *testing.T) {
provider := &CodexProvider{}
state := newResponsesState()
events := make(chan zeroruntime.StreamEvent, 16)
patch := "*** Begin Patch\n*** Add File: hello.txt\n+hello\n*** End Patch\n"
for _, payload := range []string{
`{"type":"response.output_item.added","item_id":"item-1","item":{"type":"custom_tool_call","call_id":"call-1","name":"apply_patch"}}`,
`{"type":"response.custom_tool_call_input.delta","item_id":"item-1","call_id":"call-1","delta":"*** Begin Patch\n*** Add File: hello.txt\n"}`,
`{"type":"response.custom_tool_call_input.delta","item_id":"item-1","call_id":"call-1","delta":"+hello\n*** End Patch\n"}`,
`{"type":"response.output_item.done","item_id":"item-1","item":{"type":"custom_tool_call","call_id":"call-1","name":"apply_patch","input":"*** Begin Patch\n*** Add File: hello.txt\n+hello\n*** End Patch\n"}}`,
`{"type":"response.completed","response":{"id":"resp-1","status":"completed"}}`,
} {
if keepReading := provider.emitResponsesEvent(context.Background(), payload, state, events); !keepReading && !state.done {
t.Fatalf("stream stopped before terminal event for %s", payload)
}
}
close(events)
collected := zeroruntime.CollectStream(context.Background(), events)
if len(collected.ToolCalls) != 1 {
t.Fatalf("tool calls = %#v, want one", collected.ToolCalls)
}
call := collected.ToolCalls[0]
if call.ID != "item-1" || call.ProviderCallID != "call-1" || call.Name != "apply_patch" || !call.Freeform || call.Arguments != patch {
t.Fatalf("custom apply_patch call = %#v", call)
}
}
func TestCodexApplyPatchUsesFreeformWireContract(t *testing.T) {
provider, err := NewCodexProvider(CodexOptions{Options: Options{
APIKey: "test-token",
BaseURL: "https://chatgpt.example/backend-api/codex",
Model: "gpt-test",
}})
if err != nil {
t.Fatalf("NewCodexProvider: %v", err)
}
patch := "*** Begin Patch\n*** Add File: hello.txt\n+hello\n*** End Patch\n"
request, err := provider.buildResponsesRequest(zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{
{Role: zeroruntime.MessageRoleUser, Content: "Add the file."},
{Role: zeroruntime.MessageRoleAssistant, ToolCalls: []zeroruntime.ToolCall{{
ID: "item-1", ProviderCallID: "call-1", Name: "apply_patch", Arguments: patch, Freeform: true,
}}},
{Role: zeroruntime.MessageRoleTool, ToolCallID: "item-1", ToolCallProviderID: "call-1", Content: "Done!"},
},
Tools: []zeroruntime.ToolDefinition{{
Name: "apply_patch", Description: "Apply a patch.", Type: zeroruntime.ToolDefinitionFreeform,
}},
})
if err != nil {
t.Fatalf("buildResponsesRequest: %v", err)
}
if len(request.Tools) != 1 || request.Tools[0].Type != string(zeroruntime.ToolDefinitionFreeform) ||
request.Tools[0].Format == nil || request.Tools[0].Format.Syntax != "lark" || request.Tools[0].Parameters != nil {
t.Fatalf("apply_patch tool wire shape = %#v", request.Tools)
}
if !strings.Contains(request.Tools[0].Format.Definition, "*** Begin Patch") {
t.Fatalf("apply_patch grammar = %q", request.Tools[0].Format.Definition)
}
if !strings.Contains(request.Tools[0].Description, "absolute file paths") || !strings.Contains(request.Tools[0].Description, "granted extra write root") {
t.Fatalf("apply_patch description does not explain extra-root paths: %q", request.Tools[0].Description)
}
if len(request.Input) != 3 || request.Input[1].Type != "custom_tool_call" || request.Input[1].ID != "item-1" ||
request.Input[1].CallID != "call-1" || request.Input[1].Input != patch ||
request.Input[2].Type != "custom_tool_call_output" || request.Input[2].CallID != "call-1" {
t.Fatalf("custom apply_patch replay input = %#v", request.Input)
}
}
func TestCodexPreservesCustomApplyPatchFunction(t *testing.T) {
provider, err := NewCodexProvider(CodexOptions{Options: Options{
APIKey: "test-token",
BaseURL: "https://chatgpt.example/backend-api/codex",
Model: "gpt-test",
}})
if err != nil {
t.Fatalf("NewCodexProvider: %v", err)
}
parameters := map[string]any{
"type": "object",
"properties": map[string]any{"value": map[string]any{"type": "string"}},
}
request, err := provider.buildResponsesRequest(zeroruntime.CompletionRequest{
Messages: []zeroruntime.Message{{Role: zeroruntime.MessageRoleUser, Content: "Call the custom tool."}},
Tools: []zeroruntime.ToolDefinition{{
Name: "apply_patch", Description: "Custom integration tool.", Parameters: parameters,
}},
})
if err != nil {
t.Fatalf("buildResponsesRequest: %v", err)
}
if len(request.Tools) != 1 {
t.Fatalf("tools = %#v, want one", request.Tools)
}
tool := request.Tools[0]
if tool.Type != string(zeroruntime.ToolDefinitionFunction) || tool.Format != nil ||
tool.Description != "Custom integration tool." || !reflect.DeepEqual(tool.Parameters, parameters) {
t.Fatalf("custom apply_patch function was rewritten: %#v", tool)
}
}
func TestHandleTerminalResponseNilPayload(t *testing.T) {
p := &CodexProvider{}
@@ -90,11 +212,14 @@ func TestHandleTerminalResponseFailedPayloadWithoutError(t *testing.T) {
// A response.completed with a payload and no error remains a clean done.
ok := make(chan zeroruntime.StreamEvent, 8)
p.handleTerminalResponse(context.Background(),
&responsesEvent{Type: responsesEventCompleted, Response: &responsePayload{Status: "completed"}}, &responsesState{}, ok)
&responsesEvent{Type: responsesEventCompleted, Response: &responsePayload{ID: "resp-1", Status: "completed"}}, &responsesState{}, ok)
close(ok)
for ev := range ok {
if ev.Type == zeroruntime.StreamEventError {
t.Error("response.completed with a non-error payload must NOT emit an error")
}
if ev.Type == zeroruntime.StreamEventDone && ev.ResponseID != "resp-1" {
t.Errorf("done response id = %q, want resp-1", ev.ResponseID)
}
}
}
+3 -3
View File
@@ -771,7 +771,7 @@ func TestCodexProviderParsesResponsesTextDeltas(t *testing.T) {
`{"type":"response.output_text.delta","delta":"hello "}`,
`{"type":"response.output_text.delta","delta":"world"}`,
`{"type":"response.output_text.delta","delta":"!"}`,
`{"type":"response.completed","response":{"id":"resp-1","status":"completed","usage":{"input_tokens":7,"output_tokens":3,"output_tokens_details":{"reasoning_tokens":1},"input_tokens_details":{"cached_tokens":2}}}}`,
`{"type":"response.completed","response":{"id":"resp-1","status":"completed","usage":{"input_tokens":7,"output_tokens":3,"output_tokens_details":{"reasoning_tokens":1},"input_tokens_details":{"cached_tokens":2,"cache_write_tokens":1}}}}`,
)
defer srv.Close()
@@ -809,8 +809,8 @@ func TestCodexProviderParsesResponsesTextDeltas(t *testing.T) {
if usage == nil {
t.Fatal("expected a usage event, got none")
}
if usage.InputTokens != 7 || usage.OutputTokens != 3 || usage.CachedInputTokens != 2 || usage.ReasoningTokens != 1 {
t.Fatalf("usage = %+v, want input=7 output=3 cached=2 reasoning=1", *usage)
if usage.InputTokens != 7 || usage.OutputTokens != 3 || usage.CachedInputTokens != 2 || usage.CacheWriteTokens != 1 || usage.ReasoningTokens != 1 {
t.Fatalf("usage = %+v, want input=7 output=3 cached=2 cache-write=1 reasoning=1", *usage)
}
if !gotDone {
t.Fatal("expected a done event, got none")
+1
View File
@@ -357,6 +357,7 @@ func (provider *Provider) emitChunk(
PromptTokens: chunk.Usage.PromptTokens,
CompletionTokens: chunk.Usage.CompletionTokens,
CachedInputTokens: chunk.Usage.PromptTokensDetails.CachedTokens,
CacheWriteTokens: chunk.Usage.PromptTokensDetails.CacheWriteTokens,
ReasoningTokens: chunk.Usage.CompletionTokensDetails.ReasoningTokens,
},
})
+2 -2
View File
@@ -260,14 +260,14 @@ func TestStreamCompletionAppliesCustomAuthAndHeaders(t *testing.T) {
func TestStreamCompletionEmitsTextUsageAndDone(t *testing.T) {
provider := newTestProvider(t, func(w http.ResponseWriter, r *http.Request) {
writeSSE(w, `{"choices":[{"delta":{"content":"hello "}}]}`)
writeSSE(w, `{"choices":[{"delta":{"content":"zero"}}],"usage":{"prompt_tokens":12,"completion_tokens":5,"prompt_tokens_details":{"cached_tokens":3}}}`)
writeSSE(w, `{"choices":[{"delta":{"content":"zero"}}],"usage":{"prompt_tokens":12,"completion_tokens":5,"prompt_tokens_details":{"cached_tokens":3,"cache_write_tokens":2}}}`)
writeSSE(w, `[DONE]`)
})
events := collectProviderEvents(t, provider)
assertEvent(t, events[0], zeroruntime.StreamEventText, "hello ")
assertEvent(t, events[1], zeroruntime.StreamEventText, "zero")
if events[2].Type != zeroruntime.StreamEventUsage || events[2].Usage.PromptTokens != 12 || events[2].Usage.CompletionTokens != 5 || events[2].Usage.CachedInputTokens != 3 {
if events[2].Type != zeroruntime.StreamEventUsage || events[2].Usage.PromptTokens != 12 || events[2].Usage.CompletionTokens != 5 || events[2].Usage.CachedInputTokens != 3 || events[2].Usage.CacheWriteTokens != 2 {
t.Fatalf("unexpected usage event: %#v", events[2])
}
if events[3].Type != zeroruntime.StreamEventDone {
+6 -1
View File
@@ -195,7 +195,12 @@ func (s *turnSession) computeFingerprint(request zeroruntime.CompletionRequest)
func wireToolsDigest(tools []zeroruntime.ToolDefinition) string {
rendered := make([]string, 0, len(tools))
for _, tool := range tools {
schema, err := json.Marshal(tool.Parameters)
wireShape := struct {
Type zeroruntime.ToolDefinitionType `json:"type,omitempty"`
Parameters map[string]any `json:"parameters,omitempty"`
Format *zeroruntime.ToolDefinitionFormat `json:"format,omitempty"`
}{Type: tool.Type, Parameters: tool.Parameters, Format: tool.Format}
schema, err := json.Marshal(wireShape)
if err != nil {
rendered = append(rendered, tool.Name+"\n"+tool.Description+"\n__non_json:"+tool.Name)
continue
+2 -1
View File
@@ -109,7 +109,8 @@ type usage struct {
}
type promptTokenDetails struct {
CachedTokens int `json:"cached_tokens"`
CachedTokens int `json:"cached_tokens"`
CacheWriteTokens int `json:"cache_write_tokens"`
}
type completionTokenDetails struct {
+43 -12
View File
@@ -15,9 +15,27 @@ import (
// value (e.g. "1") to enable — the same boolean idiom as ZERO_FORMAT_ON_WRITE.
const openaiTurnSessionEnv = "ZERO_OPENAI_TURN_SESSION"
// ChatGPT Responses sessions are enabled by default for their native provider.
// Setting this to 0/false/off restores the stateless HTTP/SSE transport.
const chatGPTTurnSessionEnv = "ZERO_CHATGPT_TURN_SESSION"
func openaiTurnSessionEnabled() bool {
value := strings.TrimSpace(os.Getenv(openaiTurnSessionEnv))
return value != "" && value != "0" && !strings.EqualFold(value, "false")
return value != "" && !explicitEnvFalse(value)
}
func chatGPTTurnSessionEnabled() bool {
value := strings.TrimSpace(os.Getenv(chatGPTTurnSessionEnv))
return value == "" || !explicitEnvFalse(value)
}
func explicitEnvFalse(value string) bool {
switch strings.ToLower(strings.TrimSpace(value)) {
case "0", "false", "off":
return true
default:
return false
}
}
// OptimizedTurnSessions returns the gated optimized TurnSessionProvider for an
@@ -25,20 +43,33 @@ func openaiTurnSessionEnabled() bool {
// is ineligible. Callers leave agent Options.TurnSessionProvider nil on false,
// so the loop's default wrap keeps today's behavior byte-identical.
//
// Eligibility is deliberately narrow: the resolved provider kind must be
// official OpenAI (NOT openai-compatible gateways — the constructor merges the
// two, so this branches on the resolved kind explicitly), not the ChatGPT
// Codex catalog, and the provider value must be the concrete *openai.Provider
// (a fake, a Codex provider, or nil falls back to the default path safely).
// No base-URL check on top of the kind: prewarm and fingerprint telemetry are
// harmless against any host, and kind==openai mirrors the existing
// official-OpenAI precedent used for prompt_cache_key.
// Eligibility is deliberately narrow: native ChatGPT Responses uses its
// concrete Codex provider and defaults on with a kill switch; official OpenAI
// chat-completions keeps its existing opt-in prewarm session. Compatible
// gateways and foreign provider values always retain the default adapter.
func OptimizedTurnSessions(profile config.ProviderProfile, provider zeroruntime.Provider, options Options) (zeroruntime.TurnSessionProvider, bool) {
if !openaiTurnSessionEnabled() {
resolved, err := resolveProfile(profile, options)
if err != nil {
return nil, false
}
resolved, err := resolveProfile(profile, options)
if err != nil || resolved.providerKind != config.ProviderKindOpenAI || isCodexCatalog(profile, resolved) {
if isCodexCatalog(profile, resolved) {
if !chatGPTTurnSessionEnabled() {
return nil, false
}
concrete, ok := provider.(*openai.CodexProvider)
if !ok {
return nil, false
}
caps, err := resolveCapabilities(profile, options)
if err != nil {
return nil, false
}
return openai.NewCodexTurnSessionProvider(concrete, caps), true
}
if resolved.providerKind != config.ProviderKindOpenAI {
return nil, false
}
if !openaiTurnSessionEnabled() {
return nil, false
}
concrete, ok := provider.(*openai.Provider)
+72 -5
View File
@@ -5,6 +5,7 @@ import (
"testing"
"github.com/Gitlawb/zero/internal/config"
"github.com/Gitlawb/zero/internal/providers/openai"
"github.com/Gitlawb/zero/internal/zeroruntime"
)
@@ -45,7 +46,7 @@ func TestOptimizedTurnSessionsDefaultOff(t *testing.T) {
func TestOptimizedTurnSessionsFalseyValues(t *testing.T) {
profile := openaiEligibleProfile()
provider := buildGateProvider(t, profile)
for _, value := range []string{"0", "false", "FALSE", "False", " ", ""} {
for _, value := range []string{"0", "false", "FALSE", "False", "off", "OFF", " ", ""} {
t.Setenv(openaiTurnSessionEnv, value)
if _, ok := OptimizedTurnSessions(profile, provider, Options{}); ok {
t.Fatalf("gate enabled for falsey value %q", value)
@@ -59,6 +60,21 @@ func TestOptimizedTurnSessionsFalseyValues(t *testing.T) {
}
}
func TestChatGPTTurnSessionGateValues(t *testing.T) {
for _, value := range []string{"0", "false", "FALSE", "off", "OFF"} {
t.Setenv(chatGPTTurnSessionEnv, value)
if chatGPTTurnSessionEnabled() {
t.Fatalf("ChatGPT gate enabled for falsey value %q", value)
}
}
for _, value := range []string{"", " ", "1", "true", "on"} {
t.Setenv(chatGPTTurnSessionEnv, value)
if !chatGPTTurnSessionEnabled() {
t.Fatalf("ChatGPT gate disabled for default/truthy value %q", value)
}
}
}
func TestOptimizedTurnSessionsOpenAIEligible(t *testing.T) {
t.Setenv(openaiTurnSessionEnv, "1")
profile := openaiEligibleProfile()
@@ -88,13 +104,64 @@ func TestOptimizedTurnSessionsRejectsCompatible(t *testing.T) {
}
func TestOptimizedTurnSessionsRejectsCodexCatalog(t *testing.T) {
t.Setenv(openaiTurnSessionEnv, "1")
t.Setenv(chatGPTTurnSessionEnv, "0")
profile := openaiEligibleProfile()
profile.Name = "chatgpt"
profile.CatalogID = "chatgpt"
if _, ok := OptimizedTurnSessions(profile, buildChatGPTGateProvider(t, profile), Options{}); ok {
t.Fatal("the ChatGPT Codex catalog must honor the turn-session kill switch")
}
}
func TestOptimizedTurnSessionsEnablesChatGPTResponsesByDefault(t *testing.T) {
t.Setenv(chatGPTTurnSessionEnv, "")
profile := openaiEligibleProfile()
profile.Name = "chatgpt"
profile.CatalogID = "chatgpt"
provider := buildChatGPTGateProvider(t, profile)
turnSessions, ok := OptimizedTurnSessions(profile, provider, Options{})
if !ok || turnSessions == nil {
t.Fatal("ChatGPT Responses provider did not receive its default turn session")
}
if got := turnSessions.Capabilities().Model; got != profile.Model {
t.Fatalf("Capabilities().Model = %q, want %q", got, profile.Model)
}
}
func TestOptimizedTurnSessionsRecognizesCompatibleLabeledChatGPTProfile(t *testing.T) {
t.Setenv(chatGPTTurnSessionEnv, "")
profile := openaiEligibleProfile()
profile.Name = "chatgpt"
profile.CatalogID = "chatgpt"
profile.ProviderKind = config.ProviderKindOpenAICompatible
if turnSessions, ok := OptimizedTurnSessions(profile, buildChatGPTGateProvider(t, profile), Options{}); !ok || turnSessions == nil {
t.Fatal("ChatGPT catalog profile must use its native session regardless of the stored compatibility label")
}
}
func buildChatGPTGateProvider(t *testing.T, profile config.ProviderProfile) *openai.CodexProvider {
t.Helper()
provider, err := openai.NewCodexProvider(openai.CodexOptions{
Options: openai.Options{
APIKey: "test-token",
BaseURL: "https://chatgpt.example/backend-api/codex",
Model: profile.Model,
},
AccountID: "account-test",
})
if err != nil {
t.Fatalf("NewCodexProvider: %v", err)
}
return provider
}
func TestOptimizedTurnSessionsRejectsNonCodexProviderForChatGPT(t *testing.T) {
t.Setenv(chatGPTTurnSessionEnv, "1")
profile := openaiEligibleProfile()
profile.CatalogID = "chatgpt"
// Build the provider from the base profile: the point under test is the
// catalog check, not Codex construction (which needs OAuth wiring).
if _, ok := OptimizedTurnSessions(profile, buildGateProvider(t, openaiEligibleProfile()), Options{}); ok {
t.Fatal("the ChatGPT Codex catalog must not get the optimized session")
t.Fatal("a ChatGPT catalog profile with a non-Codex provider must use the default session")
}
}
+2
View File
@@ -7,6 +7,7 @@ import (
type BackendOptions struct {
GOOS string
LookupExecutable func(string) (string, error)
DetectWSL func() WSLInfo
}
type Backend struct {
@@ -53,6 +54,7 @@ func SelectBackend(options BackendOptions) Backend {
return NewSandboxManager(SandboxManagerOptions{
GOOS: options.GOOS,
LookupExecutable: options.LookupExecutable,
DetectWSL: options.DetectWSL,
}).Backend()
}
+15
View File
@@ -45,6 +45,7 @@ func TestSelectBackendChoosesPlatformAdapterWithFallback(t *testing.T) {
}
return "", errors.New("missing")
},
DetectWSL: func() WSLInfo { return WSLInfo{} },
})
if backend.Name != BackendUnavailable || backend.Available {
t.Fatalf("linux backend = %#v, want native sandbox unavailable without Linux helper", backend)
@@ -54,6 +55,20 @@ func TestSelectBackendChoosesPlatformAdapterWithFallback(t *testing.T) {
}
})
t.Run("linux helper missing under WSL reports WSL fallback", func(t *testing.T) {
backend := SelectBackend(BackendOptions{
GOOS: "linux",
LookupExecutable: func(string) (string, error) { return "", errors.New("missing") },
DetectWSL: func() WSLInfo { return WSLInfo{IsWSL: true, IsWSL2: true} },
})
if backend.Name != BackendWSL || backend.Available || !backend.Fallback {
t.Fatalf("linux backend = %#v, want explicit WSL fallback", backend)
}
if !strings.Contains(backend.Message, "WSL2") {
t.Fatalf("linux fallback message = %q, want WSL2 diagnostic", backend.Message)
}
})
t.Run("darwin sandbox exec available", func(t *testing.T) {
backend := SelectBackend(BackendOptions{
GOOS: "darwin",
+7 -3
View File
@@ -51,6 +51,7 @@ const (
type SandboxManagerOptions struct {
GOOS string
LookupExecutable func(string) (string, error)
DetectWSL func() WSLInfo
Backend Backend
}
@@ -96,7 +97,7 @@ func NewSandboxManager(options SandboxManagerOptions) SandboxManager {
goos = runtime.GOOS
}
if backend.Name == "" {
backend = selectPlatformBackend(goos, options.LookupExecutable)
backend = selectPlatformBackend(goos, options.LookupExecutable, options.DetectWSL)
}
if backend.Platform == "" {
backend.Platform = goos
@@ -168,10 +169,13 @@ func lookupExecutable(name string) (string, error) {
return "", errors.New("executable file not found")
}
func selectPlatformBackend(goos string, lookup func(string) (string, error)) Backend {
func selectPlatformBackend(goos string, lookup func(string) (string, error), detect func() WSLInfo) Backend {
if lookup == nil {
lookup = lookupExecutable
}
if detect == nil {
detect = detectWSL
}
switch goos {
case "linux":
if helper, err := lookup(LinuxSandboxHelperName); err == nil && helper != "" {
@@ -180,7 +184,7 @@ func selectPlatformBackend(goos string, lookup func(string) (string, error)) Bac
}
return nativeBackend(goos, BackendLinuxBwrap, helper, "Linux sandbox helper available")
}
if info := detectWSL(); info.IsWSL {
if info := detect(); info.IsWSL {
return wslBackend(goos, info)
}
return unavailableBackend(goos, "Linux sandbox helper is not available")
+239 -25
View File
@@ -19,6 +19,17 @@ type Scope struct {
workspaceRoot string
readRoots []string
extraRoots []string
// tempReads / tempWrites count how many LIVE temporary grants depend on a
// root, so one holder's cleanup cannot revoke another's access.
//
// Without them a temporary grant was add-then-remove with no notion of who
// still needed it: the second caller to ask for a root already present got a
// NO-OP undo, and the first caller's cleanup removed the root out from under
// it. Two read-only tools in the same parallel batch, both blocked on the
// same directory, is exactly that shape — and read-only tools are precisely
// the ones the batch runs concurrently.
tempReads map[string]int
tempWrites map[string]int
}
// NewScope builds a scope for workspaceRoot plus the given extra roots. The
@@ -81,15 +92,88 @@ func (s *Scope) Add(path string) (string, error) {
}
s.mu.Lock()
defer s.mu.Unlock()
for _, existing := range append([]string{s.workspaceRoot}, s.extraRoots...) {
if pathWithinRoot(existing, root) {
// COVERED BY WHAT. extraRoots holds temporary write roots alongside
// permanent ones, so "something already covers this" was not the question
// worth asking: a permanent grant made over a path a temporary holder
// happened to cover recorded nothing, and vanished the moment that holder
// released. A session-scoped grant outliving the request that prompted it is
// the entire difference between Add and AddTemporaryWrite.
//
// Permanent coverage is looked for FIRST and across every root, because a
// path can be covered twice and the temporary cover must not decide the
// answer when a permanent one is also present.
if s.permanentWriteRootCoversLocked(root) {
return root, nil
}
for _, existing := range s.extraRoots {
// The same root, held temporarily: promote it in place. Its holders'
// releases become no-ops, which is what permanence means here.
if existing == root && s.temporaryWriteLocked(existing) {
delete(s.tempWrites, root)
return root, nil
}
}
// Either uncovered, or covered only by a BROADER temporary root that is
// going to be released. Recording the narrower root in its own right is what
// survives that release.
s.extraRoots = append(s.extraRoots, root)
return root, nil
}
// temporaryWriteLocked reports whether root is held by a temporary write grant
// rather than a session-scoped one. Callers must hold the lock.
func (s *Scope) temporaryWriteLocked(root string) bool {
_, temporary := s.tempWrites[root]
return temporary
}
// permanentWriteRootCoversLocked reports whether root sits under write authority
// that outlives every temporary holder: the workspace root, or a session-scoped
// grant. Coverage by a TEMPORARY write root deliberately does not count —
// it is going to be released, so nothing recorded on the strength of it
// survives. Callers must hold the lock.
//
// One helper rather than a copy per caller because Add, AddRead and
// AddTemporaryRead all turn on the same question, and three copies of "covered,
// but by what" is how one of them ends up answering it differently.
// permanentReadRootCoversLocked reports whether root is already readable by a
// grant that outlives every temporary holder.
//
// SLICE ORDER MUST NOT DECIDE THIS. The old check walked readRoots and stopped at
// the first entry covering the path, so a BROAD TEMPORARY root sitting earlier in
// the slice shadowed a narrower PERMANENT one later — and the caller was handed a
// reference on somebody else's grant when it needed no reference at all. Asking
// the whole question first, over both lists, makes the answer independent of
// insertion order.
//
// Write roots count: a grant that permits writing permits reading the same tree.
func (s *Scope) permanentReadRootCoversLocked(root string) bool {
if s.permanentWriteRootCoversLocked(root) {
return true
}
for _, existing := range s.readRoots {
if pathWithinRoot(existing, root) && !s.temporaryReadLocked(existing) {
return true
}
}
return false
}
func (s *Scope) permanentWriteRootCoversLocked(root string) bool {
for _, existing := range append([]string{s.workspaceRoot}, s.extraRoots...) {
if pathWithinRoot(existing, root) && !s.temporaryWriteLocked(existing) {
return true
}
}
return false
}
// temporaryReadLocked is temporaryWriteLocked for read roots.
func (s *Scope) temporaryReadLocked(root string) bool {
_, temporary := s.tempReads[root]
return temporary
}
// AddRead grants read-only access under path. If the path is already covered by
// a writable root, no separate read root is stored.
func (s *Scope) AddRead(path string) (string, error) {
@@ -99,11 +183,19 @@ func (s *Scope) AddRead(path string) (string, error) {
}
s.mu.Lock()
defer s.mu.Unlock()
if s.writeRootCoversLocked(root) {
// Same rule as Add: coverage that is going to be released is not coverage a
// permanent grant can rely on, whichever side of the read/write boundary it
// sits on. A permanent read covered only by a temporary WRITE root died with
// that root just as surely.
if s.permanentWriteRootCoversLocked(root) {
return root, nil
}
if s.permanentReadRootCoversLocked(root) {
return root, nil
}
for _, existing := range s.readRoots {
if pathWithinRoot(existing, root) {
if existing == root && s.temporaryReadLocked(existing) {
delete(s.tempReads, root)
return root, nil
}
}
@@ -118,16 +210,63 @@ func (s *Scope) AddTemporaryRead(path string) (string, func(), error) {
}
s.mu.Lock()
defer s.mu.Unlock()
if s.writeRootCoversLocked(root) {
// Genuinely permanent — the workspace root, or a session-scoped grant. It
// outlives this reader, so there is nothing to release.
if s.permanentWriteRootCoversLocked(root) {
return root, func() {}, nil
}
for _, existing := range s.readRoots {
if pathWithinRoot(existing, root) {
return root, func() {}, nil
}
// A TEMPORARY write root covering this path is NOT borrowed, and this is the
// one place the distinction bites hardest.
//
// The dependency on that root is real — the write holder's cleanup must not
// revoke this reader, which is the defect the refcount exists to close,
// reached across the read/write boundary. But taking a reference on the
// WRITE root pays for the lifetime with the authority: the reference keeps
// the root in extraRoots, extraRoots is what Roots() feeds validate(), and
// validate() is WRITE authorization. So a reader outliving its writer went on
// writing anywhere under a write grant that had already ended — not merely
// under the path it asked to read.
//
// One reference cannot be both the lifetime and the capability. The lifetime
// is kept below, as a temporary READ root for the path this caller actually
// asked for. readRoots feeds ReadRoots(), which confers no write authority —
// validate() authorises writes from workspaceRoot plus extraRoots, and this
// path never enters that list. So the reader survives its writer with read
// authority and only read authority.
//
// That claim is the written justification for the whole design, so it has to
// stay exhaustive as the accessors change. It has already been wrong once, in
// both directions: an earlier draft said readRoots fed ReadRoots() and nothing
// else while ExtraReadRoots() also read it, and that accessor has since been
// removed as belonging to a different change.
if s.permanentReadRootCoversLocked(root) {
return root, func() {}, nil
}
if s.tempReads == nil {
s.tempReads = map[string]int{}
}
// THE SAME ROOT, ANOTHER HOLDER: one entry, two references. Only an EXACT
// match shares — a narrower request must not, which is the whole point below.
if _, held := s.tempReads[root]; held {
s.tempReads[root]++
return root, oncePerHolder(func() { s.releaseTemporaryRead(root) }), nil
}
// A NARROWER REQUEST RECORDS ITS OWN ROOT, even under a broader temporary
// read that currently covers it.
//
// Borrowing the covering root's reference kept this reader alive by keeping
// the BROAD ROOT alive, and readRoots feeds validateRead and the native
// sandbox profile. So once the broad reader released, a caller that had asked
// to read one subdirectory was left able to read the whole tree its neighbour
// had asked for, siblings included. Reported by @jatmn, and it is the read
// twin of the write escalation fixed above: the refcount key decides not only
// WHEN an entry disappears but WHICH path stays authorised while it lives.
s.readRoots = append(s.readRoots, root)
return root, func() { s.removeReadRoot(root) }, nil
if s.tempReads == nil {
s.tempReads = map[string]int{}
}
s.tempReads[root] = 1
return root, oncePerHolder(func() { s.releaseTemporaryRead(root) }), nil
}
func (s *Scope) AddTemporaryWrite(path string) (string, func(), error) {
@@ -137,32 +276,107 @@ func (s *Scope) AddTemporaryWrite(path string) (string, func(), error) {
}
s.mu.Lock()
defer s.mu.Unlock()
if s.writeRootCoversLocked(root) {
// A PERMANENT grant is the only cover that needs no bookkeeping: it outlives
// every temporary holder, so there is nothing to keep alive and nothing to
// release.
if s.permanentWriteRootCoversLocked(root) {
return root, func() {}, nil
}
s.extraRoots = append(s.extraRoots, root)
return root, func() { s.removeWriteRoot(root) }, nil
}
func (s *Scope) writeRootCoversLocked(root string) bool {
for _, existing := range append([]string{s.workspaceRoot}, s.extraRoots...) {
if pathWithinRoot(existing, root) {
return true
}
if s.tempWrites == nil {
s.tempWrites = map[string]int{}
}
return false
// THE SAME ROOT, ANOTHER HOLDER: one entry, two references. Only an exact
// match shares an entry — a NARROWER request must not, which is the whole
// point below.
if _, held := s.tempWrites[root]; held {
s.tempWrites[root]++
return root, oncePerHolder(func() { s.releaseTemporaryWrite(root) }), nil
}
// A NARROWER REQUEST RECORDS ITS OWN ROOT, even when a broader temporary
// grant currently covers it.
//
// Borrowing a reference on the covering root kept this holder alive, which
// was the point, but it kept the COVERING ROOT alive to do it — and
// extraRoots is what validate() authorises writes from. So once the broad
// holder released, a caller that had asked to write one subdirectory was
// left holding write authority over the whole tree its neighbour had asked
// for, including siblings it never named. Reported by CodeRabbit, and it is
// the write-side twin of the read-side escalation @jatmn found: one
// reference cannot be both a lifetime and a capability, whichever side of
// the read/write boundary it sits on.
//
// Recording the requested root separately gives each holder exactly the
// authority it asked for and a lifetime of its own. The nesting costs an
// extra entry in extraRoots while both are live, which is redundant but not
// wrong — the broader root already permits everything the narrower one does.
s.extraRoots = append(s.extraRoots, root)
s.tempWrites[root] = 1
return root, oncePerHolder(func() { s.releaseTemporaryWrite(root) }), nil
}
func (s *Scope) removeReadRoot(root string) {
// releaseTemporaryRead drops ONE holder's reference and removes the root only
// when the last one is gone.
//
// The count alone cannot make this safe, and an earlier comment here claimed it
// could — that each undo is called exactly once, and flooring at zero handled
// the rest. Flooring stops the count going negative; it does not stop one
// holder's second call consuming a DIFFERENT holder's reference. With two
// readers, calling the first's undo twice took the count 2 -> 1 -> 0 and removed
// a root the second was still using. Idempotency belongs to the closure, which
// is the thing that knows whose reference it is, so every undo handed out is
// wrapped in oncePerHolder.
func (s *Scope) releaseTemporaryRead(root string) {
s.mu.Lock()
defer s.mu.Unlock()
remaining, tracked := s.tempReads[root]
if !tracked {
s.mu.Unlock()
return
}
remaining--
if remaining > 0 {
s.tempReads[root] = remaining
s.mu.Unlock()
return
}
// Both mutations under ONE hold. Dropping the lock between them opened a
// window where the root was still in readRoots but no longer in tempReads:
// a concurrent AddTemporaryRead landing there reads it as a PERMANENT root,
// hands its caller a no-op undo, and then this call strips the root — so
// that caller believes it holds access it has already silently lost.
delete(s.tempReads, root)
s.readRoots = removeScopeRoot(s.readRoots, root)
s.mu.Unlock()
}
func (s *Scope) removeWriteRoot(root string) {
// releaseTemporaryWrite is releaseTemporaryRead for write roots. Two functions
// rather than one generic helper because they guard different slices and the
// generic version would take the slice by name — which is how the wrong one
// gets passed.
func (s *Scope) releaseTemporaryWrite(root string) {
s.mu.Lock()
defer s.mu.Unlock()
remaining, tracked := s.tempWrites[root]
if !tracked {
s.mu.Unlock()
return
}
remaining--
if remaining > 0 {
s.tempWrites[root] = remaining
s.mu.Unlock()
return
}
// Same single-hold rule as releaseTemporaryRead above.
delete(s.tempWrites, root)
s.extraRoots = removeScopeRoot(s.extraRoots, root)
s.mu.Unlock()
}
// oncePerHolder makes one holder's undo safe to call more than once. The
// reference it drops is that holder's own, so a second call must do nothing
// rather than reach into the shared count and take somebody else's.
func oncePerHolder(release func()) func() {
var once sync.Once
return func() { once.Do(release) }
}
func removeScopeRoot(roots []string, root string) []string {
+56
View File
@@ -0,0 +1,56 @@
package sandbox
import (
"path/filepath"
"testing"
)
// THE ANTI-ESCALATION PROPERTY. A read grant (AddRead, where --add-read-dir lands)
// is READABLE but NOT WRITABLE. This is what makes propagating a parent's
// request_permissions read grant to a child safe: the child can audit the granted
// path but cannot modify it. Emitting the grant as a WRITE root (--add-dir) would
// break exactly this.
func TestAReadGrantIsReadableButNotWritable(t *testing.T) {
scope, granted := grantOutsideDefaults(t)
root, err := scope.AddRead(granted)
if err != nil {
t.Fatalf("AddRead: %v", err)
}
target := filepath.Join(root, "audit-me.go")
if block := scope.validateRead(target); block != nil {
t.Fatalf("a read grant did not allow READING %q: %v — the audit would fail 'outside the workspace'", target, block)
}
if block := scope.validate(target); block == nil {
t.Fatalf("a read grant ALLOWED WRITING %q — a read grant escalated to write", target)
}
}
// grantOutsideDefaults returns a scope and a directory that scope does not
// already cover, both under test-owned storage.
//
// THE CAUSE, NOT THE SYMPTOM. NewScope seeds the system temporary directory as a
// permanent write root, so a plain t.TempDir() grant is writable before any
// grant exists and the anti-escalation assertions below prove nothing. The first
// fix placed the fixture under os.UserHomeDir() to find ground the defaults do
// not own. That worked and bought two problems: a hermetic runner may expose
// HOME read-only, so these tests failed at MkdirTemp before asserting anything,
// and on an ordinary machine they wrote into the developer's real home and left
// zero-scope-* debris whenever a run was interrupted. Reported by @jatmn.
//
// Building the Scope directly removes the seeded defaults, which is what made
// HOME necessary in the first place. t.TempDir() is then outside by
// construction, the fixture is hermetic, and nothing touches HOME.
func grantOutsideDefaults(t *testing.T) (*Scope, string) {
t.Helper()
workspace := resolvedFixturePath(t, t.TempDir())
granted := resolvedFixturePath(t, t.TempDir())
scope := &Scope{workspaceRoot: workspace}
// PROVED AGAINST THIS SCOPE. The point of the helper is that the grant is
// what changes access, so the target has to begin unauthorised — asserted
// here once rather than trusted in every caller.
if scope.validate(filepath.Join(granted, "probe.txt")) == nil {
t.Fatalf("%s is writable before any grant exists; the assertions below would prove nothing", granted)
}
return scope, granted
}
@@ -0,0 +1,595 @@
package sandbox
import (
"os"
"path/filepath"
"slices"
"sync"
"testing"
)
// scopeOutsideDefaults builds a Scope directly instead of through NewScope.
//
// NewScope adds the system temp directory as a write root, so a target under
// t.TempDir() is already covered and AddTemporaryRead returns before it ever
// touches the refcount. A test built that way exercises none of this and passes
// against the bug — which is how the first version of this test passed.
func scopeOutsideDefaults(t *testing.T) (*Scope, string) {
t.Helper()
workspace := t.TempDir()
target := filepath.Join(t.TempDir(), "shared")
if err := os.MkdirAll(target, 0o700); err != nil {
t.Fatalf("mkdir target: %v", err)
}
// Built directly, NOT via NewScope: real directories are needed because
// normalizeScopeRoot requires them to exist, but NewScope's default temp
// write root would then cover the target and short-circuit the path.
return &Scope{workspaceRoot: workspace}, target
}
// The refcount and the root list have to move together.
//
// Release used to delete the refcount, drop the lock, then strip the root in a
// second acquisition. In that window the root was still in readRoots but no
// longer in tempReads, so a concurrent AddTemporaryRead read it as a PERMANENT
// root and handed its caller a no-op undo — then the release stripped it and
// that caller silently lost access it believed it held.
func TestReleasingATemporaryReadCannotStripALiveGrant(t *testing.T) {
scope, target := scopeOutsideDefaults(t)
// First holder establishes the root.
first, releaseFirst, err := scope.AddTemporaryRead(target)
if err != nil {
t.Fatalf("AddTemporaryRead: %v", err)
}
if len(scope.tempReads) == 0 {
t.Fatal("precondition: the grant should be refcounted, not covered by a default root")
}
// Second holder takes a reference on the same root.
second, releaseSecond, err := scope.AddTemporaryRead(target)
if err != nil {
t.Fatalf("AddTemporaryRead (second): %v", err)
}
// The first holder leaving must not revoke the second's access.
releaseFirst()
if block := scope.validateRead(second); block != nil {
t.Fatalf("the second holder lost its grant when the first released: %v", block)
}
// Only the last release retires the root.
releaseSecond()
if block := scope.validateRead(first); block == nil {
t.Error("the root outlived its last holder")
}
}
// Under contention the same invariant has to hold: while a grant is live, its
// root is readable.
func TestTemporaryReadGrantsSurviveConcurrentReleases(t *testing.T) {
scope, target := scopeOutsideDefaults(t)
var wg sync.WaitGroup
var mu sync.Mutex
var failures []string
for i := 0; i < 24; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := 0; j < 200; j++ {
granted, undo, err := scope.AddTemporaryRead(target)
if err != nil {
mu.Lock()
failures = append(failures, "AddTemporaryRead: "+err.Error())
mu.Unlock()
return
}
if block := scope.validateRead(granted); block != nil {
mu.Lock()
failures = append(failures, "root not readable while a grant was live")
mu.Unlock()
}
undo()
}
}()
}
wg.Wait()
if len(failures) > 0 {
t.Fatalf("%d failures, first: %s", len(failures), failures[0])
}
}
// A read covered by a TEMPORARY WRITE root must take a reference too.
//
// permanentWriteRootCoversLocked scans extraRoots, and AddTemporaryWrite appends
// temporary write roots there — so "covered by a write root" does not mean
// "covered permanently". Without a reference the write holder's cleanup silently
// revoked the reader's access: the same defect the refcount closes, reached
// across the read/write boundary instead of within one side of it.
func TestAReadCoveredByATemporaryWriteSurvivesItsRelease(t *testing.T) {
workspace := t.TempDir()
outer := filepath.Join(t.TempDir(), "outer")
inner := filepath.Join(outer, "inner")
if err := os.MkdirAll(inner, 0o700); err != nil {
t.Fatal(err)
}
// Built directly for the same reason scopeOutsideDefaults is: NewScope's
// default temp write root would cover the fixture and short-circuit this.
scope := &Scope{workspaceRoot: workspace}
_, releaseWrite, err := scope.AddTemporaryWrite(outer)
if err != nil {
t.Fatalf("temporary write: %v", err)
}
readRoot, releaseRead, err := scope.AddTemporaryRead(inner)
if err != nil {
t.Fatalf("temporary read: %v", err)
}
covered := func() bool {
for _, root := range scope.ReadRoots() {
if pathWithinRoot(root, readRoot) {
return true
}
}
return false
}
if !covered() {
t.Fatal("the reader did not hold access before any release")
}
// The WRITE holder finishes first; the reader has not released.
releaseWrite()
if !covered() {
t.Errorf("the write holder's cleanup revoked a live read grant on %s", readRoot)
}
// And once the reader releases too, the root is genuinely gone.
releaseRead()
if covered() {
t.Errorf("the root outlived its last holder: %v", scope.ReadRoots())
}
}
func hasExtraRoot(scope *Scope, root string) bool {
for _, existing := range writeRootsBeyondWorkspace(scope) {
if existing == root {
return true
}
}
return false
}
// A PERMANENT GRANT MUST OUTLIVE THE TEMPORARY ONE IT LANDED ON TOP OF.
//
// extraRoots holds temporary write roots alongside permanent ones, so Add asked
// only "is this already covered" — and a session-scoped grant made while a
// temporary holder happened to cover the path recorded nothing, then vanished
// when that holder released. Outliving the request that prompted it is the
// entire difference between Add and AddTemporaryWrite.
func TestAPermanentGrantSurvivesTheTemporaryOneItCovered(t *testing.T) {
t.Run("write, same root", func(t *testing.T) {
scope, _, outside := scopeOutsideRoots(t)
_, releaseTemp, err := scope.AddTemporaryWrite(outside)
if err != nil {
t.Fatal(err)
}
if _, err := scope.Add(outside); err != nil {
t.Fatal(err)
}
releaseTemp()
if !hasExtraRoot(scope, outside) {
t.Error("the temporary holder's release revoked a permanent write grant")
}
})
t.Run("read, same root", func(t *testing.T) {
scope, _, outside := scopeOutsideRoots(t)
_, releaseTemp, err := scope.AddTemporaryRead(outside)
if err != nil {
t.Fatal(err)
}
if _, err := scope.AddRead(outside); err != nil {
t.Fatal(err)
}
releaseTemp()
if !hasReadRoot(scope, outside) {
t.Error("the temporary holder's release revoked a permanent read grant")
}
})
// The nested shape: a BROAD temporary root covering a NARROW permanent one.
// Promoting in place cannot help here — the narrow root has to be recorded
// in its own right, or it goes when the broad one does.
t.Run("broad temporary over narrow permanent", func(t *testing.T) {
scope, _, outside := scopeOutsideRoots(t)
inner := filepath.Join(outside, "inner")
if err := os.MkdirAll(inner, 0o700); err != nil {
t.Fatal(err)
}
_, releaseBroad, err := scope.AddTemporaryWrite(outside)
if err != nil {
t.Fatal(err)
}
if _, err := scope.Add(inner); err != nil {
t.Fatal(err)
}
releaseBroad()
covered := false
for _, existing := range writeRootsBeyondWorkspace(scope) {
if pathWithinRoot(existing, inner) {
covered = true
}
}
if !covered {
t.Error("a narrower permanent grant died with the broader temporary root it sat under")
}
})
// THE READ MIRROR OF BOTH SHAPES ABOVE. AddRead makes two claims that the
// subtests so far only made for Add: the loop over extraRoots skips coverage
// that is a temporary WRITE, and the loop over readRoots skips coverage that
// is a temporary READ. Each claim is a separate branch, and neither was
// exercised from the read side.
t.Run("permanent read under a temporary write", func(t *testing.T) {
scope, _, outside := scopeOutsideRoots(t)
_, releaseWrite, err := scope.AddTemporaryWrite(outside)
if err != nil {
t.Fatal(err)
}
if _, err := scope.AddRead(outside); err != nil {
t.Fatal(err)
}
releaseWrite()
if !hasReadRoot(scope, outside) {
t.Error("a temporary write holder's release revoked a permanent read grant")
}
})
t.Run("broad temporary read over narrow permanent read", func(t *testing.T) {
scope, _, outside := scopeOutsideRoots(t)
inner := filepath.Join(outside, "inner")
if err := os.MkdirAll(inner, 0o700); err != nil {
t.Fatal(err)
}
_, releaseBroad, err := scope.AddTemporaryRead(outside)
if err != nil {
t.Fatal(err)
}
if _, err := scope.AddRead(inner); err != nil {
t.Fatal(err)
}
releaseBroad()
if block := scope.validateRead(filepath.Join(inner, "audit-me.go")); block != nil {
t.Errorf("a narrower permanent read died with the broader temporary read it sat under: %v", block)
}
})
}
// ONE HOLDER'S UNDO DROPS ONE HOLDER'S REFERENCE, HOWEVER OFTEN IT IS CALLED.
//
// The count flooring at zero stops it going negative; it does not stop a
// holder's second call consuming somebody else's reference. With two readers,
// calling the first's undo twice took the count 2 -> 1 -> 0 and removed a root
// the second was still using — the exact revocation this refcount exists to
// prevent, reached through a duplicate call rather than a sibling's cleanup.
// Deferred cleanups in a retry path are how a real caller does this by accident.
func TestOneHoldersUndoIsIdempotent(t *testing.T) {
scope, workspace, outside := scopeOutsideRoots(t)
_, undoFirst, err := scope.AddTemporaryRead(outside)
if err != nil {
t.Fatal(err)
}
_, undoSecond, err := scope.AddTemporaryRead(outside)
if err != nil {
t.Fatal(err)
}
undoFirst()
undoFirst()
undoFirst()
if !hasReadRoot(scope, outside) {
t.Fatal("repeated calls to one holder's undo revoked a live holder's access")
}
undoSecond()
if hasReadRoot(scope, outside) {
t.Error("the root outlived its last holder")
}
// The write side takes the same rule. Built directly for the same reason the
// helper is: NewScope would seed the system temp directory as a permanent
// write root, which already covers this fixture and makes the grant vacuous.
writeScope := &Scope{workspaceRoot: workspace}
_, undoWriteFirst, err := writeScope.AddTemporaryWrite(outside)
if err != nil {
t.Fatal(err)
}
_, undoWriteSecond, err := writeScope.AddTemporaryWrite(outside)
if err != nil {
t.Fatal(err)
}
undoWriteFirst()
undoWriteFirst()
if !hasExtraRoot(writeScope, outside) {
t.Fatal("repeated calls to one write holder's undo revoked a live holder's access")
}
undoWriteSecond()
if hasExtraRoot(writeScope, outside) {
t.Error("the write root outlived its last holder")
}
}
// temporaryWriteFixture builds a nested pair of real directories and a scope
// that covers neither: an OUTER directory a temporary write grant will take,
// and an INNER one below it that a temporary read grant will ask for.
//
// Built directly rather than through NewScope, for the reason
// scopeOutsideDefaults documents: NewScope seeds the system temp directory as a
// PERMANENT write root, which would cover the whole fixture and send every
// branch under test down the "genuinely permanent" path instead. Both paths are
// returned symlink-resolved because normalizeScopeRoot resolves what it stores,
// and on macOS the fixture's /var spelling is a symlink to /private/var.
func temporaryWriteFixture(t *testing.T) (scope *Scope, workspace, outer, inner string) {
t.Helper()
outer = filepath.Join(t.TempDir(), "outer")
inner = filepath.Join(outer, "inner")
if err := os.MkdirAll(inner, 0o700); err != nil {
t.Fatal(err)
}
workspace = resolvedFixturePath(t, t.TempDir())
return &Scope{workspaceRoot: workspace}, workspace, resolvedFixturePath(t, outer), resolvedFixturePath(t, inner)
}
func resolvedFixturePath(t *testing.T, path string) string {
t.Helper()
resolved, err := filepath.EvalSymlinks(path)
if err != nil {
t.Fatalf("resolve fixture %s: %v", path, err)
}
return resolved
}
// A READER MUST NOT INHERIT ITS WRITER'S AUTHORITY.
//
// The reader's dependency on a covering TEMPORARY WRITE root was recorded as a
// reference on that root — one reference standing for both the lifetime and the
// capability. Release the WRITER first and the reference kept the path in
// extraRoots, which is the list Roots() feeds validate(), so the surviving
// read-only holder could WRITE anywhere below a write grant that had already
// ended — including paths it never asked to read.
//
// This is TestAReadGrantIsReadableButNotWritable's property for temporary
// grants: read authority stays read authority, whichever holder leaves first.
func TestAReaderOutlivingItsWriterKeepsReadAuthorityOnly(t *testing.T) {
t.Run("the writer releases first", func(t *testing.T) {
scope, workspace, outer, inner := temporaryWriteFixture(t)
writeRoot, releaseWrite, err := scope.AddTemporaryWrite(outer)
if err != nil {
t.Fatalf("temporary write: %v", err)
}
if writeRoot != outer {
t.Fatalf("temporary write granted %q, want %q", writeRoot, outer)
}
readRoot, releaseRead, err := scope.AddTemporaryRead(inner)
if err != nil {
t.Fatalf("temporary read: %v", err)
}
if readRoot != inner {
t.Fatalf("temporary read granted %q, want %q", readRoot, inner)
}
// The reader holds a read root of its own and NO reference on the write
// root. Both halves matter: the first is what survives the writer, the
// second is what stops the write root surviving with it.
//
// Errorf, not Fatalf: these say WHY the behavior below breaks, and a
// bookkeeping check that aborts the run hides the behavior it explains.
scope.mu.RLock()
writeHolders := scope.tempWrites[outer]
readHolders := scope.tempReads[inner]
scope.mu.RUnlock()
if writeHolders != 1 {
t.Errorf("tempWrites[%s] = %d, want 1: the reader took a reference on the WRITE root", outer, writeHolders)
}
if readHolders != 1 {
t.Errorf("tempReads[%s] = %d, want 1: the reader kept no read root of its own", inner, readHolders)
}
// The WRITE holder finishes; only the read-only holder is left.
releaseWrite()
if got := writeRootsBeyondWorkspace(scope); len(got) != 0 {
t.Errorf("ExtraRoots() = %v, want none: the reader held the write root open", got)
}
if got, want := scope.ReadRoots(), []string{workspace, inner}; !slices.Equal(got, want) {
t.Errorf("ReadRoots() = %v, want %v", got, want)
}
target := filepath.Join(inner, "audit-me.go")
if block := scope.validateRead(target); block != nil {
t.Errorf("the writer's cleanup revoked a live read grant on %q: %v", target, block)
}
block := scope.validate(target)
if block == nil {
t.Fatalf("a read-only holder may still WRITE %q after the write grant ended", target)
}
if block.Code != BlockOutsideWorkspace {
t.Errorf("write block code = %q, want %q", block.Code, BlockOutsideWorkspace)
}
if block.Path != target {
t.Errorf("write block path = %q, want %q", block.Path, target)
}
// The escalation was never confined to the path the reader asked for —
// a borrowed reference carries the whole covering root — so a sibling
// the reader never named has to be denied too.
sibling := filepath.Join(outer, "sibling.txt")
if block := scope.validate(sibling); block == nil {
t.Errorf("a read-only holder of %q may still WRITE %q, which it never asked to read", inner, sibling)
}
releaseRead()
if got, want := scope.ReadRoots(), []string{workspace}; !slices.Equal(got, want) {
t.Errorf("ReadRoots() = %v, want %v: the read root outlived its last holder", got, want)
}
})
// A CONTROL, NOT COVERAGE FOR THE FINDING. This subtest passes on the
// unfixed tree too — the escalation needs the WRITER to go first, which the
// subtest above exercises. It earns its place by pinning that the opposite
// release order was not broken while fixing the dangerous one, and it is
// labelled so nobody counts it as evidence the finding is closed.
t.Run("the reader releases first (control: passes unfixed)", func(t *testing.T) {
scope, _, outer, inner := temporaryWriteFixture(t)
_, releaseWrite, err := scope.AddTemporaryWrite(outer)
if err != nil {
t.Fatalf("temporary write: %v", err)
}
_, releaseRead, err := scope.AddTemporaryRead(inner)
if err != nil {
t.Fatalf("temporary read: %v", err)
}
// The other order, which the fix must not break: the writer keeps FULL
// authority for as long as it holds the grant.
releaseRead()
buildLog := filepath.Join(outer, "build.log")
if block := scope.validate(buildLog); block != nil {
t.Fatalf("the reader's cleanup revoked the live write grant on %q: %v", buildLog, block)
}
if got, want := writeRootsBeyondWorkspace(scope), []string{outer}; !slices.Equal(got, want) {
t.Errorf("ExtraRoots() = %v, want %v", got, want)
}
releaseWrite()
if block := scope.validate(buildLog); block == nil {
t.Errorf("the write root %q outlived its only holder", outer)
}
if block := scope.validateRead(filepath.Join(inner, "audit-me.go")); block == nil {
t.Errorf("the read root %q outlived its only holder", inner)
}
})
}
// A NARROWER WRITE HOLDER MUST NOT INHERIT ITS NEIGHBOUR'S TREE.
//
// The write-side twin of the read-side escalation above, and the same root
// cause: a nested holder used to take a reference on the BROADER temporary root
// that happened to cover it. That kept the nested holder alive, which was the
// point, but it kept the covering root alive to do it — and extraRoots is what
// validate() authorises writes from. Once the broad holder released, a caller
// that had asked for one subdirectory held write authority over the whole tree,
// siblings included. Reported by CodeRabbit.
func TestANestedWriteHolderKeepsOnlyItsOwnSubtree(t *testing.T) {
scope, _, outsideBase := scopeOutsideRoots(t)
outer := filepath.Join(outsideBase, "outer")
inner := filepath.Join(outer, "inner")
sibling := filepath.Join(outer, "sibling")
for _, dir := range []string{outer, inner, sibling} {
if err := os.MkdirAll(dir, 0o700); err != nil {
t.Fatal(err)
}
}
_, releaseOuter, err := scope.AddTemporaryWrite(outer)
if err != nil {
t.Fatalf("temporary write on the broad root: %v", err)
}
_, releaseInner, err := scope.AddTemporaryWrite(inner)
if err != nil {
t.Fatalf("temporary write on the nested root: %v", err)
}
// THE BROAD HOLDER GOES FIRST — the nested holder outliving it is what
// exposes whose authority it is actually holding.
releaseOuter()
if block := scope.validate(filepath.Join(inner, "mine.txt")); block != nil {
t.Errorf("the nested holder lost the subtree it asked for: %v", block.Reason)
}
if block := scope.validate(filepath.Join(sibling, "not-mine.txt")); block == nil {
t.Error("the nested holder can write a SIBLING it never asked for: the broad holder's authority outlived the broad holder")
}
if block := scope.validate(filepath.Join(outer, "not-mine.txt")); block == nil {
t.Error("the nested holder can write the covering root itself after its holder released")
}
releaseInner()
if block := scope.validate(filepath.Join(inner, "mine.txt")); block == nil {
t.Error("the nested root outlived its only holder")
}
}
// The opposite order, as a control: the nested holder releasing first must not
// disturb the broad holder that is still live.
func TestANestedWriteHolderReleasingFirstLeavesTheBroadGrant(t *testing.T) {
scope, _, outsideBase := scopeOutsideRoots(t)
outer := filepath.Join(outsideBase, "outer")
inner := filepath.Join(outer, "inner")
for _, dir := range []string{outer, inner} {
if err := os.MkdirAll(dir, 0o700); err != nil {
t.Fatal(err)
}
}
_, releaseOuter, err := scope.AddTemporaryWrite(outer)
if err != nil {
t.Fatal(err)
}
_, releaseInner, err := scope.AddTemporaryWrite(inner)
if err != nil {
t.Fatal(err)
}
releaseInner()
if block := scope.validate(filepath.Join(inner, "still-ok.txt")); block != nil {
t.Errorf("the nested holder's release revoked the broad holder's coverage: %v", block.Reason)
}
releaseOuter()
if block := scope.validate(filepath.Join(inner, "gone.txt")); block == nil {
t.Error("the tree outlived every holder")
}
}
// TWO HOLDERS OF THE SAME ROOT SHARE ONE ENTRY. Only an exact match shares —
// a narrower request must not, which is what the escalation test above pins.
func TestTwoHoldersOfTheSameWriteRootShareOneEntry(t *testing.T) {
scope, _, outsideBase := scopeOutsideRoots(t)
target := filepath.Join(outsideBase, "shared")
if err := os.MkdirAll(target, 0o700); err != nil {
t.Fatal(err)
}
_, releaseFirst, err := scope.AddTemporaryWrite(target)
if err != nil {
t.Fatal(err)
}
_, releaseSecond, err := scope.AddTemporaryWrite(target)
if err != nil {
t.Fatal(err)
}
releaseFirst()
if block := scope.validate(filepath.Join(target, "still-held.txt")); block != nil {
t.Errorf("one holder's release revoked the other's grant on the same root: %v", block.Reason)
}
releaseSecond()
if block := scope.validate(filepath.Join(target, "gone.txt")); block == nil {
t.Error("the shared root outlived both holders")
}
}
// writeRootsBeyondWorkspace is Roots() without the workspace entry.
//
// These tests need to see which BEYOND-WORKSPACE write roots a scope currently
// holds. There was an ExtraRoots() accessor that answered exactly this, but it
// had no production caller — it belongs to #829's child-grant propagation, whose
// consumer side is not on this branch, and @jatmn was right that carrying it here
// leaves unrelated production surface in a PR scoped to refcounting. The
// inspection is a test concern, so it lives in the tests.
func writeRootsBeyondWorkspace(scope *Scope) []string {
all := scope.Roots()
if len(all) == 0 {
return nil
}
return append([]string(nil), all[1:]...)
}
+366
View File
@@ -0,0 +1,366 @@
package sandbox
import (
"os"
"path/filepath"
"sync"
"testing"
)
// scopeOutsideRoots builds a workspace and an unrelated directory that no
// default write root already covers.
//
// NOT t.TempDir(). /tmp and $TMPDIR are default write roots, so a temporary
// grant for a path under one is a no-op — every assertion here would pass
// against code that grants nothing at all. That is RULES §2.3 in its exact
// form, and it caught the first version of this test.
// scopeOutsideRoots builds a workspace and a directory outside it, both under
// test-owned storage, and returns a Scope that treats the second as genuinely
// outside.
//
// THE FIXTURE PROBLEM AND ITS ACTUAL CAUSE. NewScope seeds the system temporary
// directory as a permanent write root, so an ordinary t.TempDir() grant is
// covered before any temporary grant exists and every assertion built on it is
// vacuous. The first fix for that reached into os.UserHomeDir() to find ground
// the defaults do not already own — which worked, and bought two new problems:
// a hermetic runner may expose HOME read-only, so the tests failed at MkdirTemp
// before asserting anything; and on an ordinary machine they wrote fixtures into
// the developer's real home and left zeromax-scope-* debris behind whenever a
// run was interrupted. Reported by @jatmn.
//
// The cause is the seeded defaults, not the location, so the Scope is built
// DIRECTLY here — the same pattern scopeOutsideDefaults and temporaryWriteFixture
// already use. With no defaults seeded, t.TempDir() is outside by construction,
// the fixture is hermetic, and nothing touches HOME.
//
// Paths are symlink-resolved because the scope stores resolved roots and macOS
// hands out /var/... temp directories that resolve to /private/var/...; an
// unresolved fixture path compares unequal to the root the scope just recorded.
func scopeOutsideRoots(t *testing.T) (scope *Scope, workspace string, outside string) {
t.Helper()
workspace = resolvedFixturePath(t, t.TempDir())
outside = filepath.Join(resolvedFixturePath(t, t.TempDir()), "outside")
// FATAL, NOT SKIP. Every caller of this helper asserts an AUTHORIZATION
// boundary, so a skipped fixture means that assertion never ran and the build
// reported green anyway. t.TempDir() already failed the test if it could not
// provide storage, so a subdirectory failing under it is a broken machine.
if err := os.MkdirAll(outside, 0o700); err != nil {
t.Fatalf("cannot prepare %s: %v", outside, err)
}
scope = &Scope{workspaceRoot: workspace}
// PROVED AGAINST THIS SCOPE, not against production's default list. The point
// of the helper is that a grant here changes something, and the only thing
// that can answer that is the scope the test will actually use. Checking
// defaultTempWriteRoots() instead asks about a scope nobody builds here — and
// says "prove nothing" for a t.TempDir() path that this scope, seeded with no
// defaults, genuinely does not cover.
//
// Asserting the NEGATIVE up front is also what keeps every caller honest: it
// establishes that the target begins unauthorised, so a later success is the
// grant's doing and not the fixture's placement.
if scope.validate(filepath.Join(outside, "probe.txt")) == nil {
t.Fatalf("%s is writable before any grant exists, so a temporary grant here would prove nothing", outside)
}
if scope.validateRead(filepath.Join(outside, "probe.txt")) == nil {
t.Fatalf("%s is readable before any grant exists, so a temporary read grant here would prove nothing", outside)
}
return scope, workspace, outside
}
func hasReadRoot(scope *Scope, root string) bool {
for _, existing := range scope.ReadRoots() {
if existing == root {
return true
}
}
return false
}
// THE DEFECT: one holder's cleanup revoked another's access. Two read-only
// tools in the same parallel batch, both blocked on the same directory, is
// exactly this — and read-only tools are the ones the batch runs concurrently.
func TestATemporaryReadSurvivesASiblingsCleanup(t *testing.T) {
scope, _, outside := scopeOutsideRoots(t)
// The probe that the fix must make unnecessary: without it this reads
// true -> false.
if hasReadRoot(scope, outside) {
t.Fatal("the outside root is already covered; this test proves nothing")
}
_, undoA, err := scope.AddTemporaryRead(outside)
if err != nil {
t.Fatal(err)
}
_, undoB, err := scope.AddTemporaryRead(outside)
if err != nil {
t.Fatal(err)
}
if !hasReadRoot(scope, outside) {
t.Fatal("the grant did not take effect")
}
undoA()
if !hasReadRoot(scope, outside) {
t.Fatal("A's cleanup revoked the root while B still held a grant")
}
undoB()
if hasReadRoot(scope, outside) {
t.Fatal("the root outlived its last holder")
}
}
// A NARROWER REQUEST KEEPS ITS OWN SUBTREE AND NOTHING WIDER.
//
// THIS TEST USED TO ASSERT THE OPPOSITE, and the earlier version is worth
// recording because it is how the escalation survived a review round. It was
// written to prove a LIFETIME property — the broad holder's cleanup must not
// revoke a live narrow holder — and it proved that by asserting the BROAD root
// was still present afterwards. That is the defect stated as a requirement: the
// narrow holder was keeping its neighbour's whole tree alive in order to keep
// itself alive, so the surviving reader could read siblings it never asked for.
//
// Lifetime and capability are two properties and one reference cannot carry
// both. The assertions below name them separately: the requested path stays
// readable while its holder lives, AND a sibling under the released broad root
// is denied as soon as the holder that authorised it exits. Reported by @jatmn.
func TestANarrowerReaderKeepsOnlyItsOwnSubtree(t *testing.T) {
scope, _, outside := scopeOutsideRoots(t)
nested := filepath.Join(outside, "nested")
sibling := filepath.Join(outside, "sibling")
for _, dir := range []string{nested, sibling} {
// FATAL, NOT SKIP. These directories are the fixture, not a precondition
// the environment might reasonably withhold — skipping on them reports a
// broken setup as a passing run.
if err := os.MkdirAll(dir, 0o700); err != nil {
t.Fatal(err)
}
}
_, undoBroad, err := scope.AddTemporaryRead(outside)
if err != nil {
t.Fatal(err)
}
_, undoNarrow, err := scope.AddTemporaryRead(nested)
if err != nil {
t.Fatal(err)
}
// THE BROAD HOLDER GOES FIRST. The narrow holder outliving it is what exposes
// whose authority it is actually holding.
undoBroad()
if block := scope.validateRead(filepath.Join(nested, "mine.txt")); block != nil {
t.Errorf("the broad holder's cleanup revoked the subtree the narrow holder asked for: %v", block.Reason)
}
if block := scope.validateRead(filepath.Join(sibling, "not-mine.txt")); block == nil {
t.Error("the narrow reader can read a SIBLING it never asked for: the broad holder's authority outlived the broad holder")
}
if block := scope.validateRead(filepath.Join(outside, "not-mine.txt")); block == nil {
t.Error("the narrow reader can read the covering root itself after its holder released")
}
undoNarrow()
if block := scope.validateRead(filepath.Join(nested, "mine.txt")); block == nil {
t.Error("the nested root outlived its only holder")
}
}
// The opposite order, as a control: the narrow reader releasing first must not
// disturb the broad reader that is still live.
func TestANarrowerReaderReleasingFirstLeavesTheBroadGrant(t *testing.T) {
scope, _, outside := scopeOutsideRoots(t)
nested := filepath.Join(outside, "nested")
if err := os.MkdirAll(nested, 0o700); err != nil {
t.Fatal(err)
}
_, undoBroad, err := scope.AddTemporaryRead(outside)
if err != nil {
t.Fatal(err)
}
_, undoNarrow, err := scope.AddTemporaryRead(nested)
if err != nil {
t.Fatal(err)
}
undoNarrow()
if block := scope.validateRead(filepath.Join(nested, "still-ok.txt")); block != nil {
t.Errorf("the narrow holder's release revoked the broad holder's coverage: %v", block.Reason)
}
undoBroad()
if block := scope.validateRead(filepath.Join(nested, "gone.txt")); block == nil {
t.Error("the tree outlived every holder")
}
}
// TWO READERS OF THE SAME ROOT SHARE ONE ENTRY. Only an exact match shares; a
// narrower one never does, which is what the escalation test above pins.
func TestTwoReadersOfTheSameRootShareOneEntry(t *testing.T) {
scope, _, outside := scopeOutsideRoots(t)
_, undoFirst, err := scope.AddTemporaryRead(outside)
if err != nil {
t.Fatal(err)
}
_, undoSecond, err := scope.AddTemporaryRead(outside)
if err != nil {
t.Fatal(err)
}
undoFirst()
if block := scope.validateRead(filepath.Join(outside, "still-held.txt")); block != nil {
t.Errorf("one reader's release revoked the other's grant on the same root: %v", block.Reason)
}
undoSecond()
if block := scope.validateRead(filepath.Join(outside, "gone.txt")); block == nil {
t.Error("the shared root outlived both holders")
}
}
// A NARROW PERMANENT READ IS NOT BORROWED FROM A BROAD TEMPORARY ONE. The old
// scan walked readRoots in order and stopped at the first entry covering the
// path, so a broad TEMPORARY root earlier in the slice shadowed a narrower
// PERMANENT one later — and the caller took a reference it did not need on a
// grant that was about to end.
func TestAPermanentReadIsNotShadowedByABroadTemporaryOne(t *testing.T) {
scope, _, outside := scopeOutsideRoots(t)
nested := filepath.Join(outside, "nested")
if err := os.MkdirAll(nested, 0o700); err != nil {
t.Fatal(err)
}
// The broad TEMPORARY grant is taken first, so it sits earlier in readRoots.
_, undoBroad, err := scope.AddTemporaryRead(outside)
if err != nil {
t.Fatal(err)
}
if _, err := scope.AddRead(nested); err != nil {
t.Fatal(err)
}
undoBroad()
if block := scope.validateRead(filepath.Join(nested, "mine.txt")); block != nil {
t.Errorf("a PERMANENT read grant died with a broad temporary one that merely preceded it: %v", block.Reason)
}
}
// A PERMANENT root is not refcounted and must never be removed by a temporary
// holder's cleanup — the undo for a request it already covers is genuinely
// nothing.
func TestATemporaryGrantNeverRevokesAPermanentRoot(t *testing.T) {
scope, _, outside := scopeOutsideRoots(t)
if _, err := scope.AddRead(outside); err != nil {
t.Fatal(err)
}
_, undo, err := scope.AddTemporaryRead(outside)
if err != nil {
t.Fatal(err)
}
undo()
if !hasReadRoot(scope, outside) {
t.Fatal("a temporary holder's cleanup removed a permanent read root")
}
}
// The same rule for WRITE roots, or a write grant keeps the bug reads no longer
// have.
func TestATemporaryWriteSurvivesASiblingsCleanup(t *testing.T) {
scope, _, outside := scopeOutsideRoots(t)
covered := func() bool {
for _, existing := range scope.Roots() {
if existing == outside {
return true
}
}
return false
}
_, undoA, err := scope.AddTemporaryWrite(outside)
if err != nil {
t.Fatal(err)
}
_, undoB, err := scope.AddTemporaryWrite(outside)
if err != nil {
t.Fatal(err)
}
if !covered() {
t.Fatal("the write grant did not take effect")
}
undoA()
if !covered() {
t.Fatal("A's cleanup revoked the write root while B still held a grant")
}
undoB()
if covered() {
t.Fatal("the write root outlived its last holder")
}
}
// UNDER REAL CONCURRENCY, which is the shape the parallel tool batch produces:
// eight holders taking and releasing the same root in overlapping windows. The
// root must be present the whole time at least one holder has it, and gone once
// the last releases.
func TestConcurrentHoldersOfOneRoot(t *testing.T) {
scope, _, outside := scopeOutsideRoots(t)
const holders = 8
var wg sync.WaitGroup
start := make(chan struct{})
release := make(chan struct{})
// TWO PER HOLDER, because each one can report at both hasReadRoot checks. At
// one slot per holder the channel fills after eight reports, the ninth send
// blocks forever, wg.Done() is never reached and wg.Wait() hangs — so the
// test would DEADLOCK rather than fail in exactly the case it exists to
// catch, since nothing drains failures until after the wait.
failures := make(chan string, 2*holders)
// Signalled by EVERY holder once its own AddTemporaryRead has returned,
// failure included: the gate below waits for all of them, so a holder that
// returned early without signalling would hang this test rather than fail it.
acquired := make(chan struct{}, holders)
for i := 0; i < holders; i++ {
wg.Add(1)
go func() {
defer wg.Done()
<-start
_, undo, err := scope.AddTemporaryRead(outside)
acquired <- struct{}{}
if err != nil {
failures <- err.Error()
return
}
// Every holder must SEE its own grant for as long as it holds it.
if !hasReadRoot(scope, outside) {
failures <- "a holder could not see the root it had just been granted"
}
<-release
if !hasReadRoot(scope, outside) {
failures <- "a holder lost the root while still holding it"
}
undo()
}()
}
close(start)
// EVERY holder, not the first one. Spinning on hasReadRoot only waited for
// somebody to hold the root, and the spin is tight enough to win that race
// against the goroutines still being scheduled: measured over 200 runs, 194
// of them peaked at a single simultaneous holder and none ever reached
// eight. A test named for concurrent holders was testing one holder at a
// time, and would have passed against an implementation with no refcount at
// all. Counting the acquisitions makes the overlap real instead of hoped for.
for i := 0; i < holders; i++ {
<-acquired
}
close(release)
wg.Wait()
close(failures)
for failure := range failures {
t.Fatal(failure)
}
if hasReadRoot(scope, outside) {
t.Fatal("the root outlived every holder")
}
// AND THE BOOKKEEPING IS EMPTY, not merely invisible. hasReadRoot answers
// what a caller can see; a refcount entry left behind at zero is not visible
// that way and would still leak, one map entry per acquire/release cycle.
scope.mu.Lock()
leftover, stillCounted := scope.tempReads[outside]
scope.mu.Unlock()
if stillCounted {
t.Fatalf("the refcount table kept an entry for %q after every holder released: %v", outside, leftover)
}
}
+2
View File
@@ -94,6 +94,8 @@ type Event struct {
Meta map[string]string `json:"meta,omitempty"`
PromptTokens *int `json:"promptTokens,omitempty"`
CompletionTokens *int `json:"completionTokens,omitempty"`
CachedInputTokens *int `json:"cachedInputTokens,omitempty"`
CacheWriteTokens *int `json:"cacheWriteTokens,omitempty"`
TotalTokens *int `json:"totalTokens,omitempty"`
CostUSD *float64 `json:"costUsd,omitempty"`
Text string `json:"text,omitempty"`
-15
View File
@@ -1,15 +0,0 @@
//go:build !windows
package swarm
import (
"errors"
"os"
)
// isLockContended reports whether an O_EXCL lock-create error means the lock is
// currently held (so the caller should retry/wait) rather than a hard failure.
// On POSIX an existing lock surfaces only as os.ErrExist.
func isLockContended(err error) bool {
return errors.Is(err, os.ErrExist)
}
@@ -0,0 +1,18 @@
//go:build !windows
package swarm
import (
"os"
"path/filepath"
"testing"
)
func makeRedirectedLockPath(t *testing.T, root, target string) string {
t.Helper()
lockPath := filepath.Join(root, "redirected.lock")
if err := os.Symlink(target, lockPath); err != nil {
t.Fatalf("create redirected lock path: %v", err)
}
return lockPath
}
@@ -0,0 +1,20 @@
//go:build windows
package swarm
import (
"os/exec"
"path/filepath"
"strings"
"testing"
)
func makeRedirectedLockPath(t *testing.T, root, target string) string {
t.Helper()
linkDir := filepath.Join(root, "redirected")
out, err := exec.Command("cmd", "/c", "mklink", "/J", linkDir, filepath.Dir(target)).CombinedOutput()
if err != nil {
t.Fatalf("create redirected lock junction: %v %s", err, strings.TrimSpace(string(out)))
}
return filepath.Join(linkDir, filepath.Base(target))
}
-30
View File
@@ -1,30 +0,0 @@
//go:build windows
package swarm
import (
"errors"
"os"
"syscall"
)
// Windows error codes that an O_EXCL create (or a racing delete) can surface
// when another writer holds the lock file open, instead of ERROR_FILE_EXISTS.
// Treating them as "contended" makes the lock retry/wait rather than fail the
// whole send/read spuriously under contention.
const (
errorAccessDenied = syscall.Errno(5) // ERROR_ACCESS_DENIED
errorSharingViolation = syscall.Errno(32) // ERROR_SHARING_VIOLATION
)
// isLockContended reports whether an O_EXCL lock-create error means the lock is
// currently held (so the caller should retry/wait) rather than a hard failure.
// On Windows a held lock can appear as ERROR_FILE_EXISTS (mapped to os.ErrExist)
// but also as ERROR_ACCESS_DENIED / ERROR_SHARING_VIOLATION when the file is
// concurrently open or pending deletion.
func isLockContended(err error) bool {
if errors.Is(err, os.ErrExist) {
return true
}
return errors.Is(err, errorAccessDenied) || errors.Is(err, errorSharingViolation)
}
+12 -58
View File
@@ -8,7 +8,6 @@ import (
"path/filepath"
"regexp"
"strings"
"sync/atomic"
"time"
"github.com/Gitlawb/zero/internal/fsutil"
@@ -17,10 +16,10 @@ import (
// Mailbox is a per-agent, per-team message inbox persisted as a JSON array on
// disk at <baseDir>/<team>/inboxes/<agent>.json, written atomically under a
// lock file. It is hardened for untrusted input:
// advisory lock. It is hardened for untrusted input:
//
// - inbox files and their parent dirs are owner-only (0600/0700);
// - every write takes an exclusive lock file (bounded retry, stale-break);
// - every write takes an exclusive advisory lock (bounded retry);
// - message bodies and whole inboxes are size-capped (fail closed on oversize);
// - team/agent names are sanitized so a name can never escape the base dir;
// - malformed inbox JSON is rejected rather than silently reset (fail closed).
@@ -44,7 +43,6 @@ const (
defaultMaxMessageBytes = 64 * 1024 // 64 KiB per message
defaultMaxMessages = 1000 // per inbox
defaultLockTimeout = 5 * time.Second
lockStaleAfter = 30 * time.Second
inboxFileMode = 0o600
inboxDirMode = 0o700
)
@@ -206,7 +204,7 @@ func (m *Mailbox) Send(team, recipient string, msg Message) error {
if err := m.ensureInboxDir(path); err != nil {
return err
}
release, err := acquireLock(path+".lock", m.lockTimeout())
release, err := acquireLock(m.BaseDir, path+".lock", m.lockTimeout())
if err != nil {
return err
}
@@ -233,12 +231,12 @@ func (m *Mailbox) ReadAndConsume(team, recipient string) ([]Message, error) {
if err != nil {
return nil, err
}
// The lock file lives beside the inbox; create the dir so it can be taken
// The stable lock file lives beside the inbox; create the dir so it can be taken
// even on first read of a not-yet-existing inbox.
if err := m.ensureInboxDir(path); err != nil {
return nil, err
}
release, err := acquireLock(path+".lock", m.lockTimeout())
release, err := acquireLock(m.BaseDir, path+".lock", m.lockTimeout())
if err != nil {
return nil, err
}
@@ -334,63 +332,19 @@ func (m *Mailbox) atomicWriteJSON(path string, data any) error {
return nil
}
// lockSeq makes each lock token unique within the process even when two
// goroutines acquire in the same nanosecond.
var lockSeq atomic.Uint64
// acquireLock takes an exclusive lock by creating lockPath with O_EXCL. It
// retries with a short backoff until timeout, breaking a lock whose file is
// older than lockStaleAfter (so a crashed holder cannot deadlock the swarm).
//
// Each holder writes a unique token into the lock file and release is
// OWNERSHIP-AWARE: it removes the lock only if it still holds our token. After a
// stale-break another writer may legitimately own the lock; unconditionally
// removing it would let a third writer in and create split-brain writers that
// can corrupt or drop mailbox updates.
func acquireLock(lockPath string, timeout time.Duration) (func(), error) {
token := fmt.Sprintf("%d-%d-%d", os.Getpid(), time.Now().UnixNano(), lockSeq.Add(1))
// acquireLock takes an exclusive advisory lock through a stable lock file. It
// retries with a short backoff until timeout. The kernel releases the lock when
// a holder exits, so crash recovery never moves or deletes the canonical path.
func acquireLock(root, lockPath string, timeout time.Duration) (func(), error) {
deadline := time.Now().Add(timeout)
for {
f, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, inboxFileMode)
lock, err := lockutil.TryAcquireFileLockAt(root, lockPath)
if err == nil {
// Record our ownership token; the exclusive create already holds the
// lock, so a write error only weakens diagnostics, not safety.
_, _ = f.WriteString(token)
f.Close()
var released bool
return func() {
if released {
return
}
released = true
if data, rerr := os.ReadFile(lockPath); rerr == nil && string(data) == token {
_ = lockutil.RemoveLockFile(lockPath)
}
}, nil
return func() { _ = lock.Release() }, nil
}
if !isLockContended(err) {
if !errors.Is(err, lockutil.ErrLockHeld) {
return nil, fmt.Errorf("swarm: acquire lock: %w", err)
}
// Lock held: break it if stale via the atomic rename-with-verify in
// lockutil.ReclaimStaleLock (AUDIT-M13), otherwise wait.
if info, statErr := os.Stat(lockPath); statErr == nil && time.Since(info.ModTime()) > lockStaleAfter {
cleared, rerr := lockutil.ReclaimStaleLock(lockPath, token, func(reclaimedPath string) bool {
info, err := os.Stat(reclaimedPath)
return err == nil && time.Since(info.ModTime()) <= lockStaleAfter
})
if rerr != nil {
// Reclaim hit a hard failure: the rename aside failed outright, or a
// live holder's lock could not be put back (the lock path may be
// missing, so re-acquiring would break mutual exclusion). Fail closed
// instead of spinning to the deadline.
return nil, fmt.Errorf("swarm: reclaim stale lock: %w", rerr)
}
if cleared {
continue // cleared a genuinely stale lock; retry the O_EXCL create now
}
// Lost the reclaim race (or it was actually fresh) — fall through to the
// bounded wait instead of hot-spinning on a reclaim that never wins.
}
if time.Now().After(deadline) {
return nil, fmt.Errorf("swarm: timed out acquiring lock %s", filepath.Base(lockPath))
}
+38 -27
View File
@@ -219,64 +219,75 @@ func TestMailboxRejectsSymlinkedInboxDir(t *testing.T) {
}
}
func TestMailboxLockReleaseIsOwnershipAware(t *testing.T) {
func TestMailboxLockReleaseKeepsStablePath(t *testing.T) {
mb := newTestMailbox(t)
path, _ := mb.inboxPath("team", "bob")
if err := mb.ensureInboxDir(path); err != nil {
t.Fatalf("ensureInboxDir: %v", err)
}
lockPath := path + ".lock"
// Writer A acquires the lock.
releaseA, err := acquireLock(lockPath, time.Second)
releaseA, err := acquireLock(mb.BaseDir, lockPath, time.Second)
if err != nil {
t.Fatalf("acquire A: %v", err)
}
// Simulate a stale-break + takeover by writer B: overwrite the lock content
// with B's token (as a fresh acquire after a break would).
if err := os.WriteFile(lockPath, []byte("writer-B-token"), 0o600); err != nil {
t.Fatalf("simulate B takeover: %v", err)
}
// A's release must NOT delete B's lock (ownership-aware).
releaseA()
if _, err := os.Stat(lockPath); err != nil {
t.Fatalf("A's release deleted B's lock (split-brain): %v", err)
t.Fatalf("release removed stable advisory-lock path: %v", err)
}
// B's own release removes it.
os.Remove(lockPath)
}
func TestAcquireLockReclaimsStaleLock(t *testing.T) {
// A crashed holder's stale lock (old mtime) must be reclaimed via the atomic
// rename-with-verify path, leaving no sidelined .stale.* file (AUDIT-M13).
func TestAcquireLockIgnoresOldUnlockedFile(t *testing.T) {
// A crashed holder can leave an old file behind, but the kernel releases its
// advisory lock. Acquisition must succeed without renaming the stable path.
lockPath := filepath.Join(t.TempDir(), "x.lock")
if err := os.WriteFile(lockPath, []byte("dead-holder"), 0o600); err != nil {
t.Fatalf("seed stale lock: %v", err)
t.Fatalf("seed old lock metadata: %v", err)
}
old := time.Now().Add(-2 * lockStaleAfter)
old := time.Now().Add(-time.Hour)
if err := os.Chtimes(lockPath, old, old); err != nil {
t.Fatalf("chtimes: %v", err)
}
release, err := acquireLock(lockPath, time.Second)
release, err := acquireLock(filepath.Dir(lockPath), lockPath, time.Second)
if err != nil {
t.Fatalf("acquireLock should reclaim a genuinely stale lock, got %v", err)
t.Fatalf("acquire over old unlocked file: %v", err)
}
release()
if matches, _ := filepath.Glob(lockPath + ".stale.*"); len(matches) != 0 {
t.Fatalf("reclaim left sidelined files: %v", matches)
if _, err := os.Stat(lockPath); err != nil {
t.Fatalf("stable lock path missing after release: %v", err)
}
}
func TestAcquireLockDoesNotBreakFreshLock(t *testing.T) {
// A fresh, held lock (recent mtime) must never be broken — the stale-break must
// not steal a live lock (AUDIT-M13).
func TestAcquireLockDoesNotBreakHeldLock(t *testing.T) {
// File age is irrelevant while the kernel reports an active holder.
lockPath := filepath.Join(t.TempDir(), "x.lock")
release, err := acquireLock(lockPath, time.Second)
release, err := acquireLock(filepath.Dir(lockPath), lockPath, time.Second)
if err != nil {
t.Fatalf("first acquire: %v", err)
}
defer release()
if _, err := acquireLock(lockPath, 50*time.Millisecond); err == nil {
t.Fatal("acquireLock broke a fresh, held lock (split-brain risk)")
if _, err := acquireLock(filepath.Dir(lockPath), lockPath, 50*time.Millisecond); err == nil {
t.Fatal("acquireLock admitted a second holder")
}
}
func TestAcquireLockRefusesRedirectedPath(t *testing.T) {
root := t.TempDir()
target := filepath.Join(t.TempDir(), "target")
const sentinel = "do not overwrite"
if err := os.WriteFile(target, []byte(sentinel), 0o600); err != nil {
t.Fatal(err)
}
lockPath := makeRedirectedLockPath(t, root, target)
if release, err := acquireLock(root, lockPath, time.Second); err == nil {
release()
t.Fatal("acquireLock followed a redirected lock path")
}
data, err := os.ReadFile(target)
if err != nil {
t.Fatal(err)
}
if string(data) != sentinel {
t.Fatalf("redirected target was modified: %q", data)
}
}
+95
View File
@@ -16,6 +16,101 @@ type applyPatchTool struct {
scope PathScope
}
func (applyPatchTool) isBuiltInApplyPatch() {}
// PrepareFreeformApplyPatchArguments converts native structured-patch input
// into the ordinary apply_patch argument map used by permission checks and tool
// execution. Absolute patch-header paths select a granted write root and are
// rewritten relative to that root; relative-only patches keep workspace-root
// behavior.
func PrepareFreeformApplyPatchArguments(tool Tool, patch string) (map[string]any, error) {
preparer, ok := tool.(interface {
prepareFreeformArguments(string) (map[string]any, error)
})
if !ok {
return nil, fmt.Errorf("tool is not Zero's built-in apply_patch")
}
return preparer.prepareFreeformArguments(patch)
}
func (tool applyPatchTool) prepareFreeformArguments(patch string) (map[string]any, error) {
args := map[string]any{"patch": patch}
if !isStructuredPatch(patch) {
return args, nil
}
lines := strings.Split(strings.ReplaceAll(patch, "\r\n", "\n"), "\n")
selectedRoot := ""
hasRelativeHeader := false
for index, line := range lines {
prefix := ""
switch {
case strings.HasPrefix(line, structuredAddFile):
prefix = structuredAddFile
case strings.HasPrefix(line, structuredDeleteFile):
prefix = structuredDeleteFile
case strings.HasPrefix(line, structuredUpdateFile):
prefix = structuredUpdateFile
case strings.HasPrefix(line, structuredMoveTo):
prefix = structuredMoveTo
default:
continue
}
path := strings.TrimSpace(strings.TrimPrefix(line, prefix))
if !filepath.IsAbs(path) {
hasRelativeHeader = true
continue
}
root, relative, err := tool.freeformPatchRoot(path)
if err != nil {
return nil, err
}
if selectedRoot != "" && selectedRoot != root {
return nil, fmt.Errorf("freeform patch spans multiple write roots")
}
selectedRoot = root
lines[index] = prefix + filepath.ToSlash(relative)
}
if selectedRoot != "" {
roots, err := scopedRoots(tool.workspaceRoot, tool.scope)
if err != nil {
return nil, err
}
if hasRelativeHeader && selectedRoot != roots[0] {
return nil, fmt.Errorf("freeform patch mixes workspace-relative paths with an extra write root")
}
args["cwd"] = selectedRoot
args["patch"] = strings.Join(lines, "\n")
}
return args, nil
}
func (tool applyPatchTool) freeformPatchRoot(path string) (string, string, error) {
roots, err := scopedRoots(tool.workspaceRoot, tool.scope)
if err != nil {
return "", "", err
}
var firstErr error
for _, root := range roots {
candidate := sandbox.NormalizePrefixForRoot(path, root)
candidate, err = filepath.Abs(candidate)
if err != nil {
if firstErr == nil {
firstErr = err
}
continue
}
relative, err := filepath.Rel(root, candidate)
if err == nil && relative != ".." && !strings.HasPrefix(relative, ".."+string(filepath.Separator)) && !filepath.IsAbs(relative) {
return root, filepath.ToSlash(relative), nil
}
if firstErr == nil {
firstErr = outsideWorkspaceError(path)
}
}
return "", "", firstErr
}
func NewScopedApplyPatchTool(workspaceRoot string, scope PathScope) Tool {
return applyPatchTool{
baseTool: baseTool{
+20 -8
View File
@@ -140,7 +140,6 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
return errorResult(fileUnseenMessage(relativePath))
}
previouslySeenWhole := options.FileTracker.SeenWhole(absolutePath)
updated := strings.Replace(content, oldString, newString, 1)
replacedCount := 1
if replaceAll {
@@ -165,15 +164,28 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
// Re-baseline to the content we just wrote so subsequent edits in this session
// 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, trackedLineTotal(updated), trackedLineTotal(updated))
} else {
for _, span := range editedSpans {
options.FileTracker.RecordSeenBytes(absolutePath, span.start, span.end, len(updated))
}
// OUR edit, so we know precisely which lines moved: RecordEdit carries
// across the reads this edit did not disturb instead of dropping them.
//
// Record would drop all of them, and did — a file read in three pieces
// lost every piece to a single two-line edit, and the next six edits into
// regions that had been read were refused as unseen. See RecordEdit.
//
// This SUBSUMES the previouslySeenWhole special-case #956 added here.
// That branch re-recorded 1..total when the file had been read whole;
// RecordEdit does the same thing one level down (its seenWhole arm
// re-baselines as a single covering observation) and additionally keeps
// the partial reads the old else-branch discarded. Two copies of the
// rule would drift, and only one of them sees the pre-edit observation.
options.FileTracker.RecordEdit(absolutePath, []byte(content), []byte(updated), newInfo)
for _, span := range editedSpans {
options.FileTracker.RecordSeenBytes(absolutePath, span.start, span.end, len(updated))
}
} else {
// A formatter rewrote the file after us. We no longer know which line
// holds what was read, so the conservative drop is the right answer here.
options.FileTracker.Record(absolutePath, []byte(updated), newInfo)
}
suffix := ""
+296
View File
@@ -0,0 +1,296 @@
package tools
import (
"context"
"os"
"path/filepath"
"strings"
"testing"
)
// editTrackerFixture writes a numbered file and returns its real path plus a
// tracker that has already recorded the whole content as the current version.
//
// EvalSymlinks matters: on macOS t.TempDir() hands back /var/..., the tools
// resolve it to /private/var/..., and a tracker keyed on the unresolved path
// silently misses every lookup — which makes an edit fail for the very reason
// this file is about, and for entirely the wrong cause.
func editTrackerFixture(t *testing.T, lines []string) (string, *FileTracker, Tool) {
t.Helper()
root, err := filepath.EvalSymlinks(t.TempDir())
if err != nil {
t.Fatal(err)
}
path := filepath.Join(root, "index.go")
if err := os.WriteFile(path, []byte(strings.Join(lines, "\n")), 0o644); err != nil {
t.Fatal(err)
}
tracker := NewFileTracker()
content, _ := os.ReadFile(path)
info, _ := os.Stat(path)
tracker.Record(path, content, info)
return path, tracker, NewScopedEditFileTool(root, nil)
}
func numberedLines(total int, special map[int]string) []string {
lines := make([]string, 0, total)
for i := 1; i <= total; i++ {
if text, ok := special[i]; ok {
lines = append(lines, text)
continue
}
lines = append(lines, "filler")
}
return lines
}
func runTrackedEdit(t *testing.T, tool Tool, tracker *FileTracker, path, oldString, newString string) Result {
t.Helper()
return tool.(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
"path": path, "old_string": oldString, "new_string": newString,
}, RunOptions{FileTracker: tracker})
}
// ONE EDIT MUST NOT ERASE EVERY OTHER READ, and this is the run that made the
// point: a 371-line file read in three pieces (40-45, 85-260, 260-371), one
// two-line edit at line 92, and then six consecutive refusals to edit regions
// that HAD been read and that the edit never touched — until the
// repeated-failure guard halted the run. The error told the model to re-read,
// which its next successful edit would have undone again.
func TestASuccessfulEditKeepsTheReadsItDidNotDisturb(t *testing.T) {
path, tracker, tool := editTrackerFixture(t, numberedLines(371, map[int]string{
92: "cmp := keyCmp(k, n.split)",
200: "prio: l.prio, split: l.split,",
}))
tracker.RecordSeenRange(path, 40, 45, 371)
tracker.RecordSeenRange(path, 85, 260, 371)
tracker.RecordSeenRange(path, 260, 371, 371)
if !tracker.SeenRange(path, 200, 200) {
t.Fatal("setup: line 200 sits inside the 85-260 read and must start out seen")
}
first := runTrackedEdit(t, tool, tracker, path, "cmp := keyCmp(k, n.split)", "cmp := keyCmp(k, n.pivot)")
if strings.Contains(first.Output, "Error") {
t.Fatalf("the first edit failed, so the test never reaches its subject: %s", first.Output)
}
if !tracker.SeenRange(path, 200, 200) {
t.Error("line 200 was read, the edit at line 92 did not touch it, and the file still has 371 lines — but it is no longer considered seen")
}
// The assertion that matters: the SECOND edit is what the run could never do.
second := runTrackedEdit(t, tool, tracker, path, "prio: l.prio, split: l.split,", "prio: l.prio, pivot: l.pivot,")
if strings.Contains(second.Output, "not been read exactly in this session") {
t.Fatalf("a second edit into an already-read region was refused as unseen: %s", second.Output)
}
if strings.Contains(second.Output, "Error") {
t.Fatalf("the second edit failed: %s", second.Output)
}
}
// LINE NUMBERS MOVE WHEN AN EDIT CHANGES THE LINE COUNT, so a range after the
// edit has to be carried across SHIFTED, not left where it was.
func TestReadsAfterAnEditAreShiftedByTheLineDelta(t *testing.T) {
path, tracker, tool := editTrackerFixture(t, numberedLines(100, map[int]string{
10: "short",
80: "target line",
}))
tracker.RecordSeenRange(path, 5, 20, 100)
tracker.RecordSeenRange(path, 70, 90, 100)
// Replace one line with three: everything after line 10 moves down by two.
res := runTrackedEdit(t, tool, tracker, path, "short", "one\ntwo\nthree")
if strings.Contains(res.Output, "Error") {
t.Fatalf("setup edit failed: %s", res.Output)
}
// "target line" was line 80 and is now line 82. It was read either way.
content, _ := os.ReadFile(path)
moved := strings.Count(strings.SplitN(string(content), "target line", 2)[0], "\n") + 1
if moved != 82 {
t.Fatalf("fixture moved the target to line %d, expected 82", moved)
}
if !tracker.SeenRange(path, moved, moved) {
t.Errorf("the read at lines 70-90 was not shifted with the content: line %d reads as unseen", moved)
}
second := runTrackedEdit(t, tool, tracker, path, "target line", "target changed")
if strings.Contains(second.Output, "Error") {
t.Fatalf("editing a line that was read and merely moved was refused: %s", second.Output)
}
}
// THE GUARD STILL GUARDS. Carrying reads across must not credit the model with
// content it never saw — that is the whole reason the check exists.
func TestAnUnreadRegionIsStillRefusedAfterAnEdit(t *testing.T) {
path, tracker, tool := editTrackerFixture(t, numberedLines(300, map[int]string{
50: "seen line",
250: "never read line",
}))
// Only 40-60 is read. Line 250 is not.
tracker.RecordSeenRange(path, 40, 60, 300)
if res := runTrackedEdit(t, tool, tracker, path, "seen line", "seen line edited"); strings.Contains(res.Output, "Error") {
t.Fatalf("setup edit failed: %s", res.Output)
}
res := runTrackedEdit(t, tool, tracker, path, "never read line", "sneaky")
if !strings.Contains(res.Output, "not been read exactly in this session") {
t.Fatalf("an edit into a region that was NEVER read was allowed: %s", res.Output)
}
}
// A RANGE THE EDIT SPANS IS SPLIT, NOT DROPPED — and getting this wrong is not
// hypothetical: dropping on overlap was my first attempt, and it left the
// original defect exactly in place. A read almost always CONTAINS the line it
// is about to edit, because containing it is why it was read.
//
// So the flanks survive and only the rewritten lines stop being known.
func TestARangeTheEditSpansIsSplitAroundIt(t *testing.T) {
path, tracker, tool := editTrackerFixture(t, numberedLines(100, map[int]string{
20: "alpha",
21: "beta",
22: "gamma",
}))
tracker.RecordSeenRange(path, 15, 30, 100)
// Replace a three-line span with one line, entirely inside the read range.
if res := runTrackedEdit(t, tool, tracker, path, "alpha\nbeta\ngamma", "merged"); strings.Contains(res.Output, "Error") {
t.Fatalf("setup edit failed: %s", res.Output)
}
// Before the edit: still known, still at the same line numbers.
if !tracker.SeenRange(path, 15, 19) {
t.Error("lines 15-19 sit before the edit and did not move, but were forgotten")
}
// After it: still known, shifted up by the two lines the edit removed.
if !tracker.SeenRange(path, 21, 28) {
t.Error("lines 23-30 were read and merely moved to 21-28, but were forgotten")
}
// The rewritten span itself is no longer described by that read.
if tracker.SeenRange(path, 15, 30) {
t.Error("the whole original range still reads as seen, so the model is credited with content the edit replaced")
}
}
// AN EXTERNAL CHANGE STILL INVALIDATES EVERYTHING. RecordEdit is only for edits
// we made; a file changed behind our back is exactly the case the blanket drop
// is right for, and it must keep working.
func TestAnExternalChangeStillDropsEveryRead(t *testing.T) {
path, tracker, _ := editTrackerFixture(t, numberedLines(100, map[int]string{50: "hello"}))
tracker.RecordSeenRange(path, 40, 60, 100)
if !tracker.SeenRange(path, 50, 50) {
t.Fatal("setup: line 50 should be seen")
}
// Something outside Zero rewrites the file.
changed := strings.Repeat("different\n", 100)
if err := os.WriteFile(path, []byte(changed), 0o644); err != nil {
t.Fatal(err)
}
info, _ := os.Stat(path)
tracker.Record(path, []byte(changed), info)
if tracker.SeenRange(path, 50, 50) {
t.Error("an external rewrite left the old read ranges in place")
}
}
// A file read in FULL stays read in full: an edit of ours does not make that
// untrue, and re-reading a whole file after every edit is the cost this avoids.
func TestAWhollyReadFileStaysWhollyRead(t *testing.T) {
path, tracker, tool := editTrackerFixture(t, numberedLines(40, map[int]string{5: "one"}))
tracker.RecordSeenRange(path, 1, 40, 40)
if !tracker.SeenWhole(path) {
t.Fatal("setup: the file should read as wholly seen")
}
if res := runTrackedEdit(t, tool, tracker, path, "one", "one\ntwo"); strings.Contains(res.Output, "Error") {
t.Fatalf("edit failed: %s", res.Output)
}
if !tracker.SeenWhole(path) {
t.Error("a wholly-read file stopped being wholly read after an edit")
}
}
// The span helpers, directly: these decide what survives, so their edges are
// worth pinning independently of the tool.
func TestChangedSpanHelpers(t *testing.T) {
t.Run("a single line replaced in place", func(t *testing.T) {
first, last, delta := changedLineSpan("a\nb\nc", "a\nB\nc")
if first != 2 || last != 2 || delta != 0 {
t.Fatalf("got (%d,%d,%d), want (2,2,0)", first, last, delta)
}
})
t.Run("one line becomes three", func(t *testing.T) {
first, last, delta := changedLineSpan("a\nb\nc", "a\nx\ny\nz\nc")
if first != 2 || last != 2 || delta != 2 {
t.Fatalf("got (%d,%d,%d), want (2,2,2)", first, last, delta)
}
})
t.Run("lines removed", func(t *testing.T) {
first, last, delta := changedLineSpan("a\nb\nc\nd", "a\nd")
if first != 2 || last != 3 || delta != -2 {
t.Fatalf("got (%d,%d,%d), want (2,3,-2)", first, last, delta)
}
})
t.Run("bytes", func(t *testing.T) {
first, last, delta := changedByteSpan([]byte("abcdef"), []byte("abXYef"))
if first != 2 || last != 4 || delta != 0 {
t.Fatalf("got (%d,%d,%d), want (2,4,0)", first, last, delta)
}
})
}
// A file read whole IN PIECES must stay seen-whole after an edit. SeenWhole
// derives its answer from the ranges precisely because the raw flag is only set
// by a single covering read; RecordEdit branching on the flag instead sent this
// case down the split path, and the write_file refusal the derived answer exists
// to prevent came back within one edit.
func TestAFileReadWholeInPiecesStaysWholeAfterAnEdit(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "chunked.go")
content := "l1\nl2\nl3\nl4\n"
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
t.Fatal(err)
}
tracker := NewFileTracker()
info, _ := os.Stat(path)
tracker.Record(path, []byte(content), info)
// Two reads that together cover the file, neither covering it alone.
tracker.RecordSeenRange(path, 1, 2, 4)
tracker.RecordSeenRange(path, 3, 4, 4)
if !tracker.SeenWhole(path) {
t.Fatal("two chunked reads did not add up to seen-whole")
}
updated := "l1\nl2X\nl3\nl4\n"
if err := os.WriteFile(path, []byte(updated), 0o600); err != nil {
t.Fatal(err)
}
newInfo, _ := os.Stat(path)
tracker.RecordEdit(path, []byte(content), []byte(updated), newInfo)
if !tracker.SeenWhole(path) {
t.Error("a file read whole in two chunks stopped being seen whole after one edit")
}
}
// observation.total is compared against read ranges, so it has to be the line
// count a READER would give. strings.Split leaves an empty final element for
// content ending in a newline, which counted one line too many — and since
// SeenWhole asks whether the ranges cover 1..total, that made full coverage
// unreachable for a file that had just been read in full.
func TestCountLinesMatchesWhatAReaderWouldSay(t *testing.T) {
for _, testCase := range []struct {
content string
want int
}{
{"a\nb\nc\n", 3}, // the common case: trailing newline opens no line
{"a\nb\nc", 3},
{"\n", 1},
{"", 0},
} {
if got := countLines([]byte(testCase.content)); got != testCase.want {
t.Errorf("countLines(%q) = %d, want %d", testCase.content, got, testCase.want)
}
}
}
+255 -5
View File
@@ -1,11 +1,13 @@
package tools
import (
"bytes"
"crypto/sha256"
"encoding/hex"
"errors"
"os"
"sort"
"strings"
"sync"
"time"
)
@@ -165,6 +167,241 @@ func (tracker *FileTracker) RecordSeenBytes(absPath string, start, end, total in
tracker.seen[absPath] = observation
}
// RecordEdit re-baselines absPath after an edit THIS SESSION made, keeping the
// reads the edit did not disturb.
//
// WHY THIS EXISTS RATHER THAN Record. RecordHash drops every recorded range when
// the content hash moves, which is right for a change we did not make: we cannot
// say which lines still hold what was read. After our own edit we can say
// exactly. The content before the first changed line is byte-identical and sits
// at the same line numbers; the content after the last changed line is
// byte-identical and has moved by a known delta. Only the lines the edit
// actually spans stop describing the file.
//
// Dropping the lot instead cost a real run. A 371-line file was read in three
// pieces (40-45, 85-260, 260-371) and one two-line edit at line 92 erased the
// credit for all three: the next six edits — into regions that had been read,
// that the edit did not touch, in a file whose line count had not changed — were
// each refused as content "not read in this session", and the repeated-failure
// guard halted the run. The error even told the model to re-read, which would
// have been undone by its next successful edit. The guard is right that a model
// must not edit what it has not seen; it was wrong about what it had seen.
func (tracker *FileTracker) RecordEdit(absPath string, before, after []byte, info os.FileInfo) {
if tracker == nil {
return
}
firstLine, lastLineBefore, lineDelta := changedLineSpan(string(before), string(after))
firstByte, lastByteBefore, byteDelta := changedByteSpan(before, after)
tracker.mu.Lock()
defer tracker.mu.Unlock()
version := FileVersion{Hash: HashContent(after)}
if info != nil {
version.Size = info.Size()
version.MTime = info.ModTime()
}
tracker.versions[absPath] = version
observation, tracked := tracker.seen[absPath]
if !tracked {
return
}
// The DERIVED answer, not the raw flag. Branching on observation.whole here
// sent a file read whole in two chunks — seenWhole true, flag false — into
// the split path below, where it stopped being seen whole after one edit.
if seenWhole(observation) {
// Still whole: every line was read, and an edit of ours does not make
// that untrue. Re-baseline as a single covering observation so the
// answer survives regardless of how the ranges shifted.
observation.whole = true
observation.ranges = nil
observation.byteRanges = nil
observation.total = countLines(after)
observation.totalBytes = len(after)
tracker.seen[absPath] = observation
return
}
// SPLIT AROUND THE EDIT, never drop the whole range. A read almost always
// SPANS the line it is about to edit — that is why it was read — so dropping
// on overlap would have thrown away 85-260 to change line 92 and left the
// original defect in place under a longer implementation.
kept := make([]lineRange, 0, len(observation.ranges)+1)
for _, seen := range observation.ranges {
if end := min(seen.end, firstLine-1); seen.start <= end {
kept = append(kept, lineRange{start: seen.start, end: end})
}
if start := max(seen.start, lastLineBefore+1); start <= seen.end {
kept = append(kept, lineRange{start: start + lineDelta, end: seen.end + lineDelta})
}
}
observation.ranges = kept
if observation.total != 0 {
observation.total = countLines(after)
}
// Same split, on half-open byte intervals.
keptBytes := make([]lineRange, 0, len(observation.byteRanges)+1)
for _, seen := range observation.byteRanges {
if end := min(seen.end, firstByte); seen.start < end {
keptBytes = append(keptBytes, lineRange{start: seen.start, end: end})
}
if start := max(seen.start, lastByteBefore); start < seen.end {
keptBytes = append(keptBytes, lineRange{start: start + byteDelta, end: seen.end + byteDelta})
}
}
observation.byteRanges = keptBytes
if observation.totalBytes != 0 {
observation.totalBytes = len(after)
}
tracker.seen[absPath] = observation
}
// changedLineSpan reports the 1-based first line that differs between before and
// after, the 1-based last line of BEFORE that differs, and the line-count delta.
//
// Computed from a common prefix and suffix rather than from the caller's
// replacement spans: one edit_file call with replace_all can rewrite many
// scattered occurrences, and the span between the outermost two is the only
// region that is honestly unknown afterwards.
// The span is derived by scanning BYTES, not by materialising lines.
//
// The first version split both versions into []string. That allocates one string
// header per line for each version — work that scales with the SIZE OF THE FILE
// rather than the size of the edit, and it ran before RecordEdit had even checked
// whether there was a tracked observation to update. @jatmn measured 32,036,640
// bytes for two 2 MB versions of a million short lines, and extrapolated roughly
// 1.6 GB of headers for a 100 MB file before counting the content and hashes
// already live alongside them.
//
// That mattered beyond memory pressure: edit_file writes the updated bytes BEFORE
// calling RecordEdit, so an out-of-memory kill lands after the user's file has
// changed and before the tracker baseline catches up — the file and the record of
// it disagree, and nothing says so.
//
// Bytes give the same answer in bounded memory. Two identical byte prefixes share
// every line that ends inside them, so counting newlines in the common prefix
// counts the unchanged leading lines directly.
func changedLineSpan(before, after string) (firstChanged, lastChangedBefore, delta int) {
beforeCount := countTrackedLines(before)
afterCount := countTrackedLines(after)
prefix := commonLinePrefix(before, after)
// Never more leading lines than the shorter version has. The byte scan can
// otherwise claim a line that only one side possesses — two empty versions
// share a whole "line" that neither actually contains.
prefix = min(prefix, min(beforeCount, afterCount))
// The suffix may not reach back past the lines the prefix already claimed,
// exactly as the line-array loop bounded itself.
suffix := commonLineSuffix(before, after, min(beforeCount-prefix, afterCount-prefix))
return prefix + 1, beforeCount - suffix, afterCount - beforeCount
}
// commonLinePrefix counts the leading lines that are byte-identical in both.
func commonLinePrefix(before, after string) int {
limit := min(len(before), len(after))
at := 0
for at < limit && before[at] == after[at] {
at++
}
lines := strings.Count(before[:at], "\n")
// A LINE THAT ENDS EXACTLY WHERE THE SCAN STOPPED IS STILL SHARED. Scanning
// halts where the bytes diverge OR where the shorter version runs out, and in
// the second case the line containing that point can be complete and equal on
// both sides — "a\nb" against "a\nb\nc" diverges only because the first ended,
// and "b" is a whole matching line in each.
if lineBoundary(before, at) && lineBoundary(after, at) {
lines++
}
return lines
}
// commonLineSuffix counts the trailing lines that are byte-identical in both, up
// to the ceiling the prefix leaves.
func commonLineSuffix(before, after string, ceiling int) int {
if ceiling <= 0 {
return 0
}
limit := min(len(before), len(after))
back := 0
for back < limit && before[len(before)-1-back] == after[len(after)-1-back] {
back++
}
beforeAt, afterAt := len(before)-back, len(after)-back
// Each newline inside the shared tail closes a line that lies wholly within
// it. The line the tail STARTS in is shared only when the tail begins on a
// line boundary in both versions — otherwise its opening bytes differ and it
// is a changed line that merely ends the same way.
lines := strings.Count(before[beforeAt:], "\n")
if lineStart(before, beforeAt) && lineStart(after, afterAt) {
lines++
}
return min(lines, ceiling)
}
// lineBoundary reports whether at is the end of a line: the end of the text, or
// the position of the newline that closes it.
func lineBoundary(text string, at int) bool { return at == len(text) || text[at] == '\n' }
// lineStart reports whether at begins a line.
func lineStart(text string, at int) bool { return at == 0 || text[at-1] == '\n' }
// changedByteSpan is changedLineSpan in bytes: the first differing offset, the
// end offset of the changed region in BEFORE, and the size delta.
func changedByteSpan(before, after []byte) (firstChanged, lastChangedBefore, delta int) {
prefix := 0
for prefix < len(before) && prefix < len(after) && before[prefix] == after[prefix] {
prefix++
}
suffix := 0
for suffix < len(before)-prefix && suffix < len(after)-prefix &&
before[len(before)-1-suffix] == after[len(after)-1-suffix] {
suffix++
}
return prefix, len(before) - suffix, len(after) - len(before)
}
// countLines reports the line count a READER would give the content, which is
// what observation.total is compared against.
//
// The trailing newline does not open a line. strings.Split leaves an empty final
// element for "a\nb\nc\n" and so counted 4 where a read reports 3 — and because
// SeenWhole asks whether the ranges cover 1..total, an inflated total made full
// coverage unreachable for the very file that had just been read in full. Almost
// every text file ends in a newline, so this was the common case rather than an
// edge one.
//
// countTrackedLines keeps the OTHER rule deliberately: the span arithmetic has to
// match the indices strings.Split produced, so it counts the empty element after
// a trailing newline where this one does not. Two rules, two names, both stated —
// they were one function with an if, and that is how they were confused.
func countLines(content []byte) int {
// Counted, not split. This only ever needed a number, and building a slice of
// every line to take its length was the third full per-line allocation in the
// same call path.
if len(content) == 0 {
return 0
}
lines := bytes.Count(content, []byte{'\n'})
if content[len(content)-1] != '\n' {
lines++
}
return lines
}
// countTrackedLines is countLines' rule for the string form: the number of
// elements strings.Split would produce, which counts the empty element after a
// trailing newline. The two differ deliberately — observation.total compares
// against what a READER reports, while the span arithmetic has to match the
// indices the split produced.
func countTrackedLines(text string) int {
if text == "" {
return 0
}
return strings.Count(text, "\n") + 1
}
// 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
@@ -274,14 +511,27 @@ func (tracker *FileTracker) SeenWhole(absPath string) bool {
if !ok {
return false
}
return seenWhole(observation)
}
// seenWhole is the DERIVED answer to "has every line been seen", and the single
// place that decides it. Callers must hold the lock.
//
// Derived from the ranges rather than read off the flag, because the flag is
// only ever set by a SINGLE read covering the file: 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, refused the overwrite with
// advice to read the file that could not change the answer.
//
// IT HAS TO BE ONE FUNCTION. RecordEdit branched on the raw flag while this
// derived the answer, so the very case the derivation exists to rescue fell into
// the wrong arm and lost: a file read whole in two chunks reported SeenWhole
// true, then reported false after a single edit, putting the write_file refusal
// back within one edit of where it was fixed.
func seenWhole(observation fileObservation) bool {
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
}
+83
View File
@@ -3,6 +3,7 @@ package tools
import (
"os"
"path/filepath"
"strings"
"testing"
)
@@ -153,3 +154,85 @@ func TestNilFileTrackerCreatedFilesIsANoop(t *testing.T) {
t.Fatalf("CreatedFiles() on nil tracker = %v, want nil", got)
}
}
// THE SPAN IS DERIVED IN BOUNDED MEMORY, whatever the file size.
//
// changedLineSpan used to split both versions into []string, so it allocated one
// string header per line for each — work scaling with the SIZE OF THE FILE rather
// than the size of the edit, and it ran before RecordEdit had even checked whether
// there was an observation to update. @jatmn measured 32,036,640 bytes for two
// 2 MB versions of a million short lines.
//
// The cost was not only memory. edit_file writes the updated bytes BEFORE calling
// RecordEdit, so an out-of-memory kill lands after the user's file has changed and
// before the tracker baseline catches up — file and record disagree, silently.
//
// This asserts an absolute ceiling rather than a ratio, so a future rewrite that
// reintroduces per-line slices fails here rather than merely getting slower.
func TestTheChangedSpanDoesNotAllocatePerLine(t *testing.T) {
before := strings.Repeat("x\n", 200_000)
after := before[:len(before)-2] + "y\n"
spanBytes := testing.Benchmark(func(b *testing.B) {
b.ReportAllocs()
for i := 0; i < b.N; i++ {
changedLineSpan(before, after)
}
}).AllocedBytesPerOp()
if spanBytes > 4096 {
t.Errorf("changedLineSpan allocated %d bytes for a %d-line pair; per-line slices are back", spanBytes, 200_000)
}
content := []byte(before)
countBytes := testing.Benchmark(func(b *testing.B) {
b.ReportAllocs()
for i := 0; i < b.N; i++ {
countLines(content)
}
}).AllocedBytesPerOp()
if countBytes > 4096 {
t.Errorf("countLines allocated %d bytes to return a number", countBytes)
}
}
// AND IT GIVES THE SAME ANSWERS THE LINE-ARRAY VERSION DID. The rewrite is only
// safe if it is behaviour-preserving, so the replaced implementation is kept here
// as an oracle and both are asked the same questions. This caught 206 mismatches
// in the first attempt at the byte scan — empty versions and partially shared
// trailing lines, neither of which the hand-written cases below would have found.
func TestTheByteScanAgreesWithTheLineSplit(t *testing.T) {
oracle := func(before, after string) (int, int, int) {
split := func(text string) []string {
if text == "" {
return nil
}
return strings.Split(text, "\n")
}
beforeLines, afterLines := split(before), split(after)
prefix := 0
for prefix < len(beforeLines) && prefix < len(afterLines) && beforeLines[prefix] == afterLines[prefix] {
prefix++
}
suffix := 0
for suffix < len(beforeLines)-prefix && suffix < len(afterLines)-prefix &&
beforeLines[len(beforeLines)-1-suffix] == afterLines[len(afterLines)-1-suffix] {
suffix++
}
return prefix + 1, len(beforeLines) - suffix, len(afterLines) - len(beforeLines)
}
versions := []string{
"", "a", "\n", "a\n", "a\nb", "a\nb\n", "a\nb\nc", "a\nb\nc\n",
"\n\n", "x\n\ny", "a\nb\nc\nd\ne", "same\nsame\nsame\n", "ab\nc", "b\nc",
}
for _, before := range versions {
for _, after := range versions {
wantFirst, wantLast, wantDelta := oracle(before, after)
gotFirst, gotLast, gotDelta := changedLineSpan(before, after)
if wantFirst != gotFirst || wantLast != gotLast || wantDelta != gotDelta {
t.Errorf("changedLineSpan(%q, %q) = (%d, %d, %d), want (%d, %d, %d)",
before, after, gotFirst, gotLast, gotDelta, wantFirst, wantLast, wantDelta)
}
}
}
}
+3
View File
@@ -9,6 +9,9 @@ import (
func TestUpdatePlanToolStoresAndFormatsPlan(t *testing.T) {
tool := NewUpdatePlanTool()
if description := strings.ToLower(tool.Description()); !strings.Contains(description, "skip it for bounded changes in one component") {
t.Fatalf("update_plan description must preserve the bounded-task fast path: %q", description)
}
result := tool.Run(context.Background(), map[string]any{
"plan": []any{
+1 -1
View File
@@ -33,7 +33,7 @@ func NewScopedReadFileTool(workspaceRoot string, scope PathScope) Tool {
return readFileTool{
baseTool: baseTool{
name: "read_file",
description: "Read exact file text, each line prefixed with its line number and →. Use for comments, formatting, or edits; prefer read_minified_file for initial code understanding. Use offset and limit for a line range.",
description: "Read exact file text, each line prefixed with its line number and →. Use directly for small files likely to be edited, or when comments, formatting, or line numbers matter. Prefer read_minified_file only for exploratory understanding of large or unfamiliar code. Use offset and limit for a line range.",
parameters: Schema{
Type: "object",
Properties: map[string]PropertySchema{
+1 -1
View File
@@ -28,7 +28,7 @@ func NewScopedReadMinifiedFileTool(workspaceRoot string, scope PathScope) Tool {
return readMinifiedFileTool{
baseTool: baseTool{
name: "read_minified_file",
description: "Read source code in a token-efficient, language-aware form. Prefer this for initial understanding; use read_file for exact text or line numbers.",
description: "Read source code in a token-efficient, language-aware form. Use for exploratory understanding of large or unfamiliar source when no exact edit is planned; use read_file directly for likely edit targets. Do not immediately reread the same file exactly unless a new need for exact text or line numbers appears.",
parameters: Schema{
Type: "object",
Properties: map[string]PropertySchema{
+23
View File
@@ -8,6 +8,29 @@ import (
"testing"
)
func TestReadToolsDescribeExploratoryAndExactUseWithoutRedundantRereads(t *testing.T) {
root := t.TempDir()
minified := strings.ToLower(NewScopedReadMinifiedFileTool(root, nil).Description())
for _, want := range []string{
"exploratory understanding",
"use read_file directly for likely edit targets",
"unless a new need for exact text or line numbers appears",
} {
if !strings.Contains(minified, want) {
t.Fatalf("read_minified_file description missing %q: %s", want, minified)
}
}
exact := strings.ToLower(NewScopedReadFileTool(root, nil).Description())
for _, want := range []string{
"use directly for small files likely to be edited",
"prefer read_minified_file only for exploratory understanding",
} {
if !strings.Contains(exact, want) {
t.Fatalf("read_file description missing %q: %s", want, exact)
}
}
}
func TestReadMinifiedFileStripsCommentsAndLineNumbers(t *testing.T) {
dir := t.TempDir()
src := "package demo\n\nimport \"fmt\"\n\n// secret doc comment\nfunc F() { fmt.Println(\"x\") }\n"
+59 -2
View File
@@ -5,6 +5,7 @@ import (
"fmt"
"sort"
"strings"
"sync"
"github.com/Gitlawb/zero/internal/redaction"
"github.com/Gitlawb/zero/internal/sandbox"
@@ -12,7 +13,17 @@ import (
)
type Registry struct {
tools map[string]Tool
mu sync.RWMutex
tools map[string]Tool
generation uint64
}
// RegistrySnapshot is one immutable registry generation. Tools is a detached,
// name-sorted slice: later registrations cannot change the snapshot observed by
// an in-flight agent turn.
type RegistrySnapshot struct {
Tools []Tool
Generation uint64
}
type RunOptions struct {
@@ -107,23 +118,69 @@ func IsDeferralEligible(t Tool) bool {
}
func (registry *Registry) Register(tool Tool) {
registry.mu.Lock()
defer registry.mu.Unlock()
registry.tools[tool.Name()] = tool
registry.generation++
}
// RegisterBatch publishes a group of tools as one registry generation. Readers
// observe either the complete prior generation or the complete new one.
func (registry *Registry) RegisterBatch(batch []Tool) {
if len(batch) == 0 {
return
}
registry.mu.Lock()
defer registry.mu.Unlock()
for _, tool := range batch {
registry.tools[tool.Name()] = tool
}
registry.generation++
}
func (registry *Registry) Get(name string) (Tool, bool) {
registry.mu.RLock()
defer registry.mu.RUnlock()
tool, ok := registry.tools[name]
return tool, ok
}
func (registry *Registry) All() []Tool {
return registry.Snapshot().Tools
}
// Snapshot returns a stable, sorted view of the registry at one generation.
func (registry *Registry) Snapshot() RegistrySnapshot {
registry.mu.RLock()
tools := make([]Tool, 0, len(registry.tools))
for _, tool := range registry.tools {
tools = append(tools, tool)
}
generation := registry.generation
registry.mu.RUnlock()
sort.Slice(tools, func(left, right int) bool {
return tools[left].Name() < tools[right].Name()
})
return tools
return RegistrySnapshot{Tools: tools, Generation: generation}
}
// Clone returns an independent registry containing exactly one source snapshot.
func (registry *Registry) Clone() *Registry {
clone := NewRegistry()
if registry == nil {
return clone
}
snapshot := registry.Snapshot()
clonedTools := make([]Tool, 0, len(snapshot.Tools))
for _, tool := range snapshot.Tools {
if cloner, ok := tool.(interface{ cloneForRegistry(*Registry) Tool }); ok {
tool = cloner.cloneForRegistry(clone)
}
clonedTools = append(clonedTools, tool)
}
clone.RegisterBatch(clonedTools)
clone.generation = snapshot.Generation
return clone
}
func (registry *Registry) Run(ctx context.Context, name string, args map[string]any) Result {
+77
View File
@@ -2,10 +2,12 @@ package tools
import (
"context"
"fmt"
"os"
"path/filepath"
"slices"
"strings"
"sync"
"testing"
"github.com/Gitlawb/zero/internal/sandbox"
@@ -95,6 +97,81 @@ func TestRegistryAllReturnsToolsSortedByName(t *testing.T) {
}
}
func TestRegistrySnapshotAndCloneStayOnOneGeneration(t *testing.T) {
registry := NewRegistry()
registry.Register(fakePlainTool{baseTool: baseTool{name: "alpha", description: "alpha"}})
first := registry.Snapshot()
registry.RegisterBatch([]Tool{
fakePlainTool{baseTool: baseTool{name: "charlie", description: "charlie"}},
fakePlainTool{baseTool: baseTool{name: "bravo", description: "bravo"}},
})
second := registry.Snapshot()
clone := registry.Clone()
registry.Register(fakePlainTool{baseTool: baseTool{name: "delta", description: "delta"}})
if got := toolNames(first.Tools); !slices.Equal(got, []string{"alpha"}) {
t.Fatalf("first snapshot changed after registration: %v", got)
}
if second.Generation != first.Generation+1 {
t.Fatalf("batch generation = %d, want %d", second.Generation, first.Generation+1)
}
if got := toolNames(second.Tools); !slices.Equal(got, []string{"alpha", "bravo", "charlie"}) {
t.Fatalf("second snapshot tools = %v", got)
}
if got := toolNames(clone.All()); !slices.Equal(got, toolNames(second.Tools)) {
t.Fatalf("clone changed with source registry: %v", got)
}
if clone.Snapshot().Generation != second.Generation {
t.Fatalf("clone generation = %d, want %d", clone.Snapshot().Generation, second.Generation)
}
}
func TestRegistryBatchIsAtomicForConcurrentReaders(t *testing.T) {
registry := NewRegistry()
registry.Register(fakePlainTool{baseTool: baseTool{name: "base", description: "base"}})
batch := make([]Tool, 128)
for index := range batch {
name := fmt.Sprintf("tool_%03d", index)
batch[index] = fakePlainTool{baseTool: baseTool{name: name, description: name}}
}
start := make(chan struct{})
done := make(chan struct{})
var readers sync.WaitGroup
for range 4 {
readers.Add(1)
go func() {
defer readers.Done()
<-start
for {
count := len(registry.Snapshot().Tools)
if count != 1 && count != len(batch)+1 {
t.Errorf("reader observed partial batch with %d tools", count)
return
}
select {
case <-done:
return
default:
}
}
}()
}
close(start)
registry.RegisterBatch(batch)
close(done)
readers.Wait()
}
func toolNames(toolset []Tool) []string {
names := make([]string, 0, len(toolset))
for _, tool := range toolset {
names = append(names, tool.Name())
}
return names
}
func TestCoreNetworkToolsExposeSafetyMetadata(t *testing.T) {
// web_search is only registered when a backend is configured.
t.Setenv("ZERO_WEBSEARCH_BASE_URL", "https://search.example/api")
+5
View File
@@ -58,6 +58,11 @@ func NewToolSearchTool(registry *Registry) Tool {
}
}
func (tool toolSearchTool) cloneForRegistry(registry *Registry) Tool {
tool.registry = registry
return tool
}
// Run satisfies the Tool interface; actual dispatch goes through RunWithOptions.
func (tool toolSearchTool) Run(ctx context.Context, args map[string]any) Result {
return tool.RunWithOptions(ctx, args, RunOptions{})
+31
View File
@@ -73,6 +73,37 @@ func newDeferredFixtureRegistry() *Registry {
return reg
}
func TestRegistryCloneKeepsToolSearchOnFrozenSnapshot(t *testing.T) {
live := newDeferredFixtureRegistry()
live.Register(NewToolSearchTool(live))
runRegistry := live.Clone()
live.Register(searchFakeTool{
name: "late_lookup",
description: "Registered after the run snapshot.",
parameters: Schema{Type: "object", AdditionalProperties: false},
})
result := runRegistry.Run(context.Background(), ToolSearchToolName, map[string]any{
"query": "select:late_lookup",
})
if result.Status != StatusOK {
t.Fatalf("tool_search status = %s: %s", result.Status, result.Output)
}
if loaded := result.Meta["load_tools"]; loaded != "" {
t.Fatalf("snapshot tool_search loaded late tool %q", loaded)
}
if direct := runRegistry.Run(context.Background(), "late_lookup", map[string]any{}); direct.Status != StatusError {
t.Fatalf("late tool unexpectedly executable in run snapshot: %#v", direct)
}
known := runRegistry.Run(context.Background(), ToolSearchToolName, map[string]any{
"query": "select:weather_lookup",
})
if loaded := known.Meta["load_tools"]; loaded != "weather_lookup" {
t.Fatalf("snapshot tool_search lost existing tool: load_tools=%q output=%q", loaded, known.Output)
}
}
func TestToolSearchExposesExpectedSafetyAndSchema(t *testing.T) {
tool := NewToolSearchTool(NewRegistry())
+8
View File
@@ -211,6 +211,14 @@ type Tool interface {
Run(ctx context.Context, args map[string]any) Result
}
// IsBuiltInApplyPatch reports whether tool is Zero's scoped patch tool. The
// marker method is package-private so external and MCP tools cannot claim this
// identity merely by sharing the public name "apply_patch".
func IsBuiltInApplyPatch(tool Tool) bool {
_, ok := tool.(interface{ isBuiltInApplyPatch() })
return ok
}
// ArgsPermissioner is an optional interface a Tool can implement to refine its
// permission for a SPECIFIC call based on its arguments. When a tool implements
// it, the agent loop consults PermissionForArgs(args) instead of the static
+1 -1
View File
@@ -26,7 +26,7 @@ func NewUpdatePlanTool() *updatePlanTool {
return &updatePlanTool{
baseTool: baseTool{
name: "update_plan",
description: "Create or update the in-memory plan for implementation or investigation with at least three meaningful dependent steps. Do not use it for simple lookups, explanations, or code navigation. " +
description: "Create or update the in-memory plan for work spanning multiple components or workstreams, sequencing uncertainty, or sustained tracking across many tool calls. Skip it for bounded changes in one component, simple lookups, explanations, and code navigation. " +
"Pass the full ordered list of steps each call; it replaces the previous plan. " +
"Each item needs a `content` string; `status` defaults to \"pending\" and `id` is " +
"auto-numbered, so you only need to supply `content` (and `status` as the task progresses). " +
+11
View File
@@ -53,6 +53,7 @@ const (
CounterModelSwitches = "model_switches"
CounterInputTokens = "input_tokens"
CounterCachedInputTokens = "cached_input_tokens"
CounterCacheWriteTokens = "cache_write_tokens"
CounterOutputTokens = "output_tokens"
CounterPrewarmAttempts = "prewarm_attempts"
CounterPrefixStable = "prefix_stable"
@@ -64,6 +65,12 @@ const (
CounterPostureEscalations = "posture_escalations"
)
const (
CounterResponseChainReused = "response_chain_reused"
CounterResponseChainReset = "response_chain_reset"
CounterResponsesHTTPFallback = "responses_http_fallback"
)
// Span is one named wall interval attributed to part of a run. Each stamp is
// its own entry (spans are not merged by name): a run that streams two model
// completions has two generation entries, which lets the recorder derive
@@ -290,6 +297,7 @@ func OptionalEventKeys() []string {
"span:" + SpanVerification,
"counter:" + CounterToolCalls,
"counter:" + CounterCachedInputTokens,
"counter:" + CounterCacheWriteTokens,
"counter:" + CounterRetryCount,
"counter:" + CounterReconnectCount,
"counter:" + CounterCompactionCount,
@@ -301,6 +309,9 @@ func OptionalEventKeys() []string {
"counter:" + CounterPrewarmAttempts,
"counter:" + CounterPrefixStable,
"counter:" + CounterPrefixDrift,
"counter:" + CounterResponseChainReused,
"counter:" + CounterResponseChainReset,
"counter:" + CounterResponsesHTTPFallback,
"counter:" + CounterPostureEscalations,
"event:prefix_hash",
"task_state",
+14 -4
View File
@@ -614,11 +614,21 @@ func TestCounterPrecisionRoundTrip(t *testing.T) {
}
}
func TestOptionalEventKeysIncludePostureEscalations(t *testing.T) {
func TestOptionalEventKeysIncludeRuntimeCounters(t *testing.T) {
keys := make(map[string]bool)
for _, key := range OptionalEventKeys() {
if key == "counter:"+CounterPostureEscalations {
return
keys[key] = true
}
for _, counter := range []string{
CounterCacheWriteTokens,
CounterResponseChainReused,
CounterResponseChainReset,
CounterResponsesHTTPFallback,
CounterPostureEscalations,
} {
key := "counter:" + counter
if !keys[key] {
t.Errorf("%s missing from OptionalEventKeys: %v", key, OptionalEventKeys())
}
}
t.Fatalf("posture_escalations missing from OptionalEventKeys: %v", OptionalEventKeys())
}
+1 -1
View File
@@ -109,7 +109,7 @@ var commandDefinitions = []commandDefinition{
name: "/voice",
usage: "/voice",
group: commandGroupRuntime,
description: "Toggle voice mode (hold Space to dictate).",
description: "Toggle voice mode (" + voiceCaptureUsage + ").",
kind: commandVoice,
},
{
+9 -14
View File
@@ -88,7 +88,7 @@ type dictationController struct {
waveBars []int
waveTick int
// voiceModeEnabled repurposes Space into hold-to-record (§13.9).
// voiceModeEnabled repurposes Ctrl+Space into hold-to-record (§13.9).
voiceModeEnabled bool
// eventTypesSupported records whether the terminal confirmed key-release
// reporting (Kitty protocol); eventTypesKnown gates it until the terminal has
@@ -173,7 +173,7 @@ func firstNonEmptyStr(a, b string) string {
func (d dictationController) active() bool { return d.phase != dictIdle }
// toggleDictation starts a recording when idle and stops-and-transcribes when
// recording. Invoked by the voice-mode Space gesture; a call during
// recording. Invoked by the voice-mode Ctrl+Space gesture; a call during
// startup/transcription is ignored (the machine is mid-transition).
func (m model) toggleDictation() (model, tea.Cmd) {
if !m.dictation.available() {
@@ -193,15 +193,15 @@ func (m model) toggleDictation() (model, tea.Cmd) {
}
}
// toggleVoiceMode flips the /voice hold-to-record gesture — the dictation
// trigger. While on, Space records; run /voice again to type spaces normally.
// toggleVoiceMode flips the /voice hold-to-record gesture. While on,
// Ctrl+Space records and ordinary Space continues to type normally.
func (m model) toggleVoiceMode() (model, tea.Cmd) {
if !m.dictation.available() {
return m.appendSystemNotice("Dictation is not configured. See docs/dictation.md to set up a local engine or a Groq/OpenAI key."), nil
}
m.dictation.voiceModeEnabled = !m.dictation.voiceModeEnabled
if m.dictation.voiceModeEnabled {
return m.showTransientNoticeInline("Voice mode on — hold Space to dictate; run /voice again to turn it off.", transientNoticeSuccess), nil
return m.showTransientNoticeInline("Voice mode on — "+voiceCaptureUsage+". Run /voice again to turn it off.", transientNoticeSuccess), nil
}
// Turning voice off is the "done dictating" signal, so release the warm sherpa
// streaming server — otherwise a loaded model keeps idling in RAM (and holding
@@ -374,7 +374,7 @@ func (m model) handleDictationStarted(msg dictationStartedMsg) (model, tea.Cmd)
if m.dictation.phase == dictStarting {
m.dictation.phase = dictRecording
}
// A voice-mode Space release that arrived mid-startup asked us to stop as soon
// A voice-mode Ctrl+Space release that arrived mid-startup asked us to stop as soon
// as recording began.
if m.dictation.voiceStopPending && m.dictation.phase == dictRecording {
m.dictation.voiceStopPending = false
@@ -480,12 +480,8 @@ func needsLeadingSpace(state composerState) bool {
func (m model) dictationStatusChip() string {
switch m.dictation.phase {
case dictStarting, dictRecording:
stop := "release Space to stop"
if !m.dictation.eventTypesSupported {
stop = "press Space to stop" // press-to-toggle fallback
}
wave := zeroTheme.amber.Render("● " + renderWaveBars(m.dictation.waveBars) + " REC")
return wave + zeroTheme.muted.Render(" · "+stop+", Esc to cancel")
return wave + zeroTheme.muted.Render(" · "+m.voiceCaptureStatus()+", Esc to cancel")
case dictTranscribing:
return zeroTheme.accent.Render("●") + " " + zeroTheme.muted.Render("transcribing…")
}
@@ -601,13 +597,12 @@ func (m model) handleRecTick() (model, tea.Cmd) {
return m, nil
}
// voiceModeIndicator renders the persistent "voice mode on" hint shown while
// idle so the user knows Space is repurposed to hold-to-record (§10).
// voiceModeIndicator renders the current voice gesture or transcription state.
func (m model) voiceModeIndicator() string {
if !m.dictation.voiceModeEnabled {
return ""
}
return zeroTheme.accent.Render("🎙 voice") + zeroTheme.muted.Render(" · "+m.dictation.currentModelLabel())
return zeroTheme.accent.Render("🎙 voice") + zeroTheme.muted.Render(" · "+m.voiceCaptureStatus()+" · "+m.dictation.currentModelLabel())
}
// dictationErrorText renders a dictation error for the transcript: a missing-
+41 -10
View File
@@ -4,17 +4,19 @@ import (
tea "charm.land/bubbletea/v2"
)
// Voice mode's Space-hold gesture — the (only) dictation trigger. Two terminal
const voiceCaptureUsage = "hold Ctrl+Space to record; press it to toggle where key releases are unavailable"
// Voice mode's Ctrl+Space gesture — the (only) dictation trigger. Two terminal
// tiers:
//
// 1. The terminal confirms key-release reporting (Ghostty/Kitty/WezTerm, …):
// Space press starts recording, Space release stops it — true hold-to-record,
// no ambiguity.
// 2. The terminal does not confirm release events: Space falls back to
// Ctrl+Space press starts recording and release stops it — true
// hold-to-record, with ordinary Space left available for typing.
// 2. The terminal does not confirm release events: Ctrl+Space falls back to
// press-to-toggle (press to start, press again to stop) — a deliberately
// simpler, robust fallback than inferring release from key-repeat timing
// (racy in a terminal). This works on ANY terminal, so voice mode is always
// usable once enabled.
// (racy in a terminal). Legacy NUL input is decoded as Ctrl+Space by the
// terminal input layer, so the same matcher covers both tiers.
//
// Only active while voice mode is on; the rest of dispatch is built on
// KeyPressMsg and is untouched.
@@ -27,8 +29,37 @@ func (m model) handleKeyboardEnhancements(msg tea.KeyboardEnhancementsMsg) model
return m
}
// handleVoiceSpacePress handles a Space press while voice mode is on.
func (m model) handleVoiceSpacePress(msg tea.KeyMsg) (model, tea.Cmd) {
func voiceCaptureKey(msg tea.KeyMsg) bool {
mod := msg.Key().Mod &^ (tea.ModCapsLock | tea.ModNumLock)
return keyIs(msg, tea.KeySpace) && mod == tea.ModCtrl
}
// voiceCaptureReleaseKey matches either half of the held Ctrl+Space chord. A
// terminal may report Ctrl releasing first, or report the later Space release
// without ModCtrl because Ctrl is no longer down.
func voiceCaptureReleaseKey(msg tea.KeyMsg) bool {
return keyIs(msg, tea.KeySpace) || keyIs(msg, tea.KeyLeftCtrl) || keyIs(msg, tea.KeyRightCtrl)
}
func (m model) voiceCaptureStatus() string {
switch m.dictation.phase {
case dictStarting, dictRecording:
if m.dictation.eventTypesSupported {
return "release Ctrl+Space to stop"
}
return "press Ctrl+Space to stop"
case dictTranscribing:
return "transcribing…"
default:
if m.dictation.eventTypesSupported {
return "hold Ctrl+Space to record"
}
return "press Ctrl+Space to record"
}
}
// handleVoiceCapturePress handles Ctrl+Space while voice mode is on.
func (m model) handleVoiceCapturePress(msg tea.KeyMsg) (model, tea.Cmd) {
if !m.dictation.eventTypesSupported {
// Tier 2: no release events — press-to-toggle.
return m.toggleDictation()
@@ -45,8 +76,8 @@ func (m model) handleVoiceSpacePress(msg tea.KeyMsg) (model, tea.Cmd) {
return m, nil // already recording; the release will stop it
}
// handleVoiceSpaceRelease stops a hold-to-record session when Space is released.
func (m model) handleVoiceSpaceRelease() (model, tea.Cmd) {
// handleVoiceCaptureRelease stops a hold-to-record session when Ctrl+Space is released.
func (m model) handleVoiceCaptureRelease() (model, tea.Cmd) {
if !m.dictation.spaceHeld {
return m, nil
}
+200 -15
View File
@@ -2,6 +2,7 @@ package tui
import (
"context"
"strings"
"testing"
tea "charm.land/bubbletea/v2"
@@ -36,7 +37,7 @@ func TestToggleVoiceModeFlips(t *testing.T) {
if !next.dictation.voiceModeEnabled {
t.Fatal("first /voice should enable voice mode")
}
if next.transientNotice.text != "Voice mode on — hold Space to dictate; run /voice again to turn it off." {
if next.transientNotice.text != "Voice mode on — hold Ctrl+Space to record; press it to toggle where key releases are unavailable. Run /voice again to turn it off." {
t.Errorf("voice-mode-on notice = %q", next.transientNotice.text)
}
if transcriptHasText(next, "Voice mode on") {
@@ -68,15 +69,27 @@ func TestKeyboardEnhancementsRecorded(t *testing.T) {
}
}
func TestVoiceSpaceHoldStartsRecording(t *testing.T) {
func TestVoiceCaptureKeyIgnoresLockModifiers(t *testing.T) {
for _, mod := range []tea.KeyMod{
tea.ModCtrl | tea.ModCapsLock,
tea.ModCtrl | tea.ModNumLock,
tea.ModCtrl | tea.ModCapsLock | tea.ModNumLock,
} {
if !voiceCaptureKey(tea.KeyPressMsg(tea.Key{Code: tea.KeySpace, Mod: mod})) {
t.Fatalf("Ctrl+Space with lock modifiers %v was not recognized", mod)
}
}
}
func TestVoiceCaptureHoldStartsRecording(t *testing.T) {
m := model{dictation: batchOnlyController()}
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = true
press := tea.KeyPressMsg(tea.Key{Code: tea.KeySpace})
next, cmd := m.handleVoiceSpacePress(press)
press := tea.KeyPressMsg(tea.Key{Code: tea.KeySpace, Mod: tea.ModCtrl})
next, cmd := m.handleVoiceCapturePress(press)
if next.dictation.phase != dictStarting {
t.Fatalf("Space press should start recording (phase=%d)", next.dictation.phase)
t.Fatalf("Ctrl+Space press should start recording (phase=%d)", next.dictation.phase)
}
if !next.dictation.spaceHeld {
t.Error("spaceHeld should be set in hold mode")
@@ -86,28 +99,28 @@ func TestVoiceSpaceHoldStartsRecording(t *testing.T) {
}
}
func TestVoiceSpaceHoldIgnoresRepeat(t *testing.T) {
func TestVoiceCaptureHoldIgnoresRepeat(t *testing.T) {
m := model{dictation: batchOnlyController()}
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = true
m.dictation.phase = dictRecording
m.dictation.spaceHeld = true
repeat := tea.KeyPressMsg(tea.Key{Code: tea.KeySpace, IsRepeat: true})
next, _ := m.handleVoiceSpacePress(repeat)
repeat := tea.KeyPressMsg(tea.Key{Code: tea.KeySpace, Mod: tea.ModCtrl, IsRepeat: true})
next, _ := m.handleVoiceCapturePress(repeat)
if next.dictation.phase != dictRecording {
t.Error("auto-repeat while held must not restart or change phase")
}
}
func TestVoiceSpaceReleaseStops(t *testing.T) {
func TestVoiceCaptureReleaseStops(t *testing.T) {
m := model{dictation: batchOnlyController()}
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = true
m.dictation.phase = dictRecording
m.dictation.spaceHeld = true
next, _ := m.handleVoiceSpaceRelease()
next, _ := m.handleVoiceCaptureRelease()
if next.dictation.phase != dictTranscribing {
t.Errorf("release should stop recording → transcribing (phase=%d)", next.dictation.phase)
}
@@ -123,7 +136,7 @@ func TestVoiceReleaseDuringStartupDefersStop(t *testing.T) {
m.dictation.phase = dictStarting
m.dictation.spaceHeld = true
next, _ := m.handleVoiceSpaceRelease()
next, _ := m.handleVoiceCaptureRelease()
if !next.dictation.voiceStopPending {
t.Error("a release during startup should defer the stop")
}
@@ -134,14 +147,186 @@ func TestVoiceReleaseDuringStartupDefersStop(t *testing.T) {
}
}
func TestVoiceSpaceToggleFallback(t *testing.T) {
func TestVoiceCaptureToggleFallback(t *testing.T) {
m := model{dictation: batchOnlyController()}
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = false // no release events → toggle fallback
press := tea.KeyPressMsg(tea.Key{Code: tea.KeySpace})
next, _ := m.handleVoiceSpacePress(press)
press := tea.KeyPressMsg(tea.Key{Code: tea.KeySpace, Mod: tea.ModCtrl})
next, _ := m.handleVoiceCapturePress(press)
if next.dictation.phase != dictStarting {
t.Errorf("toggle-mode Space should start recording (phase=%d)", next.dictation.phase)
t.Errorf("toggle-mode Ctrl+Space should start recording (phase=%d)", next.dictation.phase)
}
}
func TestVoiceModeCtrlSpaceReleaseStopsHoldRecording(t *testing.T) {
m := newModel(t.Context(), Options{})
m.dictation = batchOnlyController()
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = true
m.dictation.phase = dictRecording
m.dictation.spaceHeld = true
updated, _ := m.Update(tea.KeyReleaseMsg(tea.Key{Code: tea.KeySpace, Mod: tea.ModCtrl}))
next := updated.(model)
if next.dictation.phase != dictTranscribing || next.dictation.spaceHeld {
t.Fatalf("Ctrl+Space release did not stop hold recording: phase=%d held=%v", next.dictation.phase, next.dictation.spaceHeld)
}
}
func TestVoiceModeCaptureStopsWhenEitherChordKeyIsReleased(t *testing.T) {
for _, test := range []struct {
name string
key tea.Key
}{
{name: "left Ctrl first", key: tea.Key{Code: tea.KeyLeftCtrl}},
{name: "right Ctrl first", key: tea.Key{Code: tea.KeyRightCtrl}},
{name: "Space after Ctrl", key: tea.Key{Code: tea.KeySpace}},
} {
t.Run(test.name, func(t *testing.T) {
m := newModel(t.Context(), Options{})
m.dictation = batchOnlyController()
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = true
m.dictation.phase = dictRecording
m.dictation.spaceHeld = true
updated, _ := m.Update(tea.KeyReleaseMsg(test.key))
next := updated.(model)
if next.dictation.phase != dictTranscribing || next.dictation.spaceHeld {
t.Fatalf("release did not stop hold recording: phase=%d held=%v", next.dictation.phase, next.dictation.spaceHeld)
}
})
}
}
func TestVoiceModeIndicatorDescribesTierAndPhase(t *testing.T) {
for _, test := range []struct {
name string
releaseEvents bool
phase dictationPhase
spaceHeld bool
want string
}{
{name: "toggle idle", want: "press Ctrl+Space to record"},
{name: "hold idle", releaseEvents: true, want: "hold Ctrl+Space to record"},
{name: "toggle recording", phase: dictRecording, want: "press Ctrl+Space to stop"},
{name: "hold recording", releaseEvents: true, phase: dictRecording, spaceHeld: true, want: "release Ctrl+Space to stop"},
{name: "transcribing", phase: dictTranscribing, want: "transcribing"},
} {
t.Run(test.name, func(t *testing.T) {
m := model{dictation: batchOnlyController()}
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = test.releaseEvents
m.dictation.phase = test.phase
m.dictation.spaceHeld = test.spaceHeld
if got := m.voiceModeIndicator(); !strings.Contains(strings.ToLower(got), strings.ToLower(test.want)) {
t.Fatalf("voice indicator = %q, want %q", got, test.want)
}
})
}
}
func TestDictationStatusChipUsesVoiceCaptureShortcut(t *testing.T) {
for _, test := range []struct {
name string
releaseEvents bool
want string
}{
{name: "toggle", want: "press Ctrl+Space to stop"},
{name: "hold", releaseEvents: true, want: "release Ctrl+Space to stop"},
} {
t.Run(test.name, func(t *testing.T) {
m := model{dictation: batchOnlyController()}
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = test.releaseEvents
m.dictation.phase = dictRecording
if got := m.dictationStatusChip(); !strings.Contains(got, test.want) {
t.Fatalf("dictation status = %q, want %q", got, test.want)
}
})
}
}
func TestVoiceCommandDescriptionExplainsBothTerminalTiers(t *testing.T) {
for _, command := range commandDefinitions {
if command.kind == commandVoice {
for _, want := range []string{"hold Ctrl+Space", "press it to toggle", "key releases"} {
if !strings.Contains(command.description, want) {
t.Fatalf("voice command description = %q, want %q", command.description, want)
}
}
return
}
}
t.Fatal("voice command definition not found")
}
func TestVoiceModePlainSpaceTypesNormally(t *testing.T) {
for _, releaseEvents := range []bool{false, true} {
t.Run(map[bool]string{false: "toggle", true: "hold"}[releaseEvents], func(t *testing.T) {
m := newModel(t.Context(), Options{})
m.dictation = batchOnlyController()
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = releaseEvents
updated, cmd := m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeySpace}))
next := updated.(model)
if got := next.composerValue(); got != " " {
t.Fatalf("plain Space composer = %q, want one space", got)
}
if next.dictation.phase != dictIdle || cmd != nil {
t.Fatalf("plain Space started dictation: phase=%d cmd=%v", next.dictation.phase, cmd != nil)
}
})
}
}
func TestVoiceModeCtrlSpaceStartsAcrossTerminalTiers(t *testing.T) {
for _, releaseEvents := range []bool{false, true} {
t.Run(map[bool]string{false: "toggle", true: "hold"}[releaseEvents], func(t *testing.T) {
m := newModel(t.Context(), Options{})
m.dictation = batchOnlyController()
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = releaseEvents
updated, cmd := m.Update(tea.KeyPressMsg(tea.Key{Code: tea.KeySpace, Mod: tea.ModCtrl}))
next := updated.(model)
if next.dictation.phase != dictStarting || cmd == nil {
t.Fatalf("Ctrl+Space did not start dictation: phase=%d cmd=%v", next.dictation.phase, cmd != nil)
}
if got := next.composerValue(); got != "" {
t.Fatalf("Ctrl+Space typed into composer: %q", got)
}
})
}
}
func TestVoiceModeTypingPreservesSpacesAcrossTerminalTiers(t *testing.T) {
for _, releaseEvents := range []bool{false, true} {
t.Run(map[bool]string{false: "toggle", true: "hold"}[releaseEvents], func(t *testing.T) {
m := newModel(t.Context(), Options{})
m.dictation = batchOnlyController()
m.dictation.voiceModeEnabled = true
m.dictation.eventTypesSupported = releaseEvents
for _, key := range []tea.Key{
{Code: 'f', Text: "f"},
{Code: 'i', Text: "i"},
{Code: 'x', Text: "x"},
{Code: tea.KeySpace},
{Code: 't', Text: "t"},
{Code: 'h', Text: "h"},
{Code: 'e', Text: "e"},
} {
updated, _ := m.Update(tea.KeyPressMsg(key))
m = updated.(model)
}
if got := m.composerValue(); got != "fix the" {
t.Fatalf("composer = %q, want %q", got, "fix the")
}
if m.dictation.phase != dictIdle {
t.Fatalf("typing started dictation: phase=%d", m.dictation.phase)
}
})
}
}
+93 -34
View File
@@ -86,10 +86,12 @@ type model struct {
savedProviders []config.ProviderProfile
provider zeroruntime.Provider
newProvider func(config.ProviderProfile) (zeroruntime.Provider, error)
newTurnSessionProvider func(config.ProviderProfile, zeroruntime.Provider) zeroruntime.TurnSessionProvider
probeProviderHealth func(context.Context, providerhealth.Options) providerhealth.Result
discoverProviderModels func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error)
discoverOllamaContextWindow func(ctx context.Context, baseURL string, model string) (int, error)
registry *tools.Registry
awaitToolReadiness func(context.Context)
// lspManager is created once per session and reused across prompts so gopls (and
// other language servers) stay warm — a fresh manager per run would cold-start
// the server on the first edit of every turn. Nil when cwd is unknown; runs then
@@ -234,17 +236,20 @@ type model struct {
leaderHelpOverlay bool
// leaderPending is true after Ctrl+X until a second key, Esc, or timeout
// resolves the chord (see leader.go). leaderSeq invalidates a stale tick.
leaderPending bool
leaderSeq int
transcriptBodyHeights *transcriptBodyHeightCache
input textinput.Model
composer composerState
composerActive bool
composerCursorVisible bool
composerPastePreviews []composerPastePreview
composerSelection composerSelectionState
dictation dictationController
sttKeyPrompt *sttKeyPromptState
leaderPending bool
leaderSeq int
transcriptBodyHeights *transcriptBodyHeightCache
input textinput.Model
composer composerState
composerActive bool
composerCursorVisible bool
composerBlinkSeq int
composerBlinkIdleTicks int
composerBlinkTicking bool
composerPastePreviews []composerPastePreview
composerSelection composerSelectionState
dictation dictationController
sttKeyPrompt *sttKeyPromptState
// plan holds the sticky plan panel state (steps, expansion, timings)
// synced from the update_plan tool. See plan_panel.go.
plan planPanelState
@@ -492,10 +497,11 @@ type model struct {
// transientNotice is a single, replaceable confirmation shown above the
// composer. Unlike copyStatus it is shared by lightweight slash commands
// and is never persisted into the transcript.
transientNotice transientNotice
transientNoticeSeq int
exitConfirmActive bool
exitConfirmSeq int
transientNotice transientNotice
transientNoticeSeq int
transientNoticeTimerSeq int
exitConfirmActive bool
exitConfirmSeq int
// cancelConfirmActive/cancelConfirmSeq mirror exitConfirmActive/exitConfirmSeq
// (same seq-gated tea.Tick pattern) but guard a DIFFERENT action: Esc
// cancelling a running turn. The two are deliberately separate state (not a
@@ -929,6 +935,13 @@ func newModel(ctx context.Context, options Options) model {
// terminal bracketed paste (Paste: true) is the single paste path.
input.KeyMap.Paste.SetEnabled(false)
input.Focus()
// The main composer paints its own cursor through composerCursorVisible.
// Keep the embedded textinput cursor static so the shared textinput.Blink
// message cannot start a second, invisible 530ms tick chain. Setup inputs
// still use the normal Bubbles cursor and continue blinking independently.
inputStyles := input.Styles()
inputStyles.Cursor.Blink = false
input.SetStyles(inputStyles)
runSpinner := spinner.New(spinner.WithSpinner(spinner.MiniDot))
runSpinner.Spinner.FPS = activeAnimationFrameInterval
@@ -953,6 +966,8 @@ func newModel(ctx context.Context, options Options) model {
userCommands: loadedUserCommands,
loadSkills: options.LoadSkills,
composerCursorVisible: true,
composerBlinkSeq: 1,
composerBlinkTicking: true,
terminalFocused: true,
userConfigPath: options.UserConfigPath,
doctorUserConfigPath: doctorUserConfigPath,
@@ -969,10 +984,12 @@ func newModel(ctx context.Context, options Options) model {
recapsEnabled: options.RecapsEnabled,
provider: options.Provider,
newProvider: options.NewProvider,
newTurnSessionProvider: options.NewTurnSessionProvider,
probeProviderHealth: options.ProbeProviderHealth,
discoverProviderModels: options.DiscoverProviderModels,
discoverOllamaContextWindow: options.DiscoverOllamaContextWindow,
registry: registry,
awaitToolReadiness: options.AwaitToolReadiness,
sessionStore: sessionStore,
peerService: options.PeerService,
sandboxStore: sandboxStore,
@@ -1093,6 +1110,11 @@ const (
// composerCursorBlinkInterval is the on/off period of the composer text cursor.
const composerCursorBlinkInterval = 530 * time.Millisecond
// composerBlinkIdleTicksBeforeSettle bounds the software cursor animation.
// Once the terminal is focused but otherwise idle, the cursor finishes a short
// blink sequence and settles visible instead of waking the TUI forever.
const composerBlinkIdleTicksBeforeSettle = 4
// composerTypingIdleThreshold is how long a typing pause must last before the
// cursor resumes blinking; comfortably above normal inter-keystroke gaps
// (~150-300ms) so it won't flicker mid-sentence.
@@ -1101,16 +1123,27 @@ const composerTypingIdleThreshold = 500 * time.Millisecond
// composerBlinkMsg toggles the composer cursor's visibility each tick. The custom
// composer render draws its own cursor (not textinput's), so it drives its own
// blink rather than relying on textinput.Blink.
type composerBlinkMsg struct{}
type composerBlinkMsg struct{ seq int }
func composerBlinkCmd() tea.Cmd {
func composerBlinkCmd(seq int) tea.Cmd {
return tea.Tick(composerCursorBlinkInterval, func(time.Time) tea.Msg {
return composerBlinkMsg{}
return composerBlinkMsg{seq: seq}
})
}
func (m model) armComposerBlink() (model, tea.Cmd) {
m.composerCursorVisible = true
m.composerBlinkIdleTicks = 0
if m.composerBlinkTicking {
return m, nil
}
m.composerBlinkTicking = true
m.composerBlinkSeq++
return m, composerBlinkCmd(m.composerBlinkSeq)
}
func (m model) Init() tea.Cmd {
cmds := []tea.Cmd{textinput.Blink, composerBlinkCmd()}
cmds := []tea.Cmd{textinput.Blink, composerBlinkCmd(m.composerBlinkSeq)}
if m.petAnimation != nil && !m.reducedMotion {
cmds = append(cmds, petTickCmd(m.petTickSeq, m.petFrameDelay()))
}
@@ -1293,7 +1326,15 @@ 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, recapActivityCmd)
var composerCmd tea.Cmd
switch msg.(type) {
case tea.KeyPressMsg, tea.PasteMsg, tea.FocusMsg:
if nm.terminalFocused {
nm, composerCmd = nm.armComposerBlink()
}
}
nm, noticeCmd := nm.ensureTransientNoticeTimer()
return nm, batchCommands(cmd, mouseCmd, flushCmd, recapActivityCmd, composerCmd, noticeCmd)
}
func batchCommands(cmds ...tea.Cmd) tea.Cmd {
@@ -1356,16 +1397,27 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
}
return m, nil
case composerBlinkMsg:
m = m.expireTransientNotice()
if !m.composerBlinkTicking || msg.seq != m.composerBlinkSeq {
return m, nil
}
switch {
case !m.terminalFocused:
m.composerCursorVisible = false // hidden while unfocused
m.composerCursorVisible = false
m.composerBlinkTicking = false
return m, nil
case m.now().Sub(m.lastCharTime) < composerTypingIdleThreshold:
m.composerCursorVisible = true // solid while actively typing
m.composerCursorVisible = true
m.composerBlinkIdleTicks = 0
default:
m.composerCursorVisible = !m.composerCursorVisible // idle + focused: blink as before
m.composerBlinkIdleTicks++
if m.composerBlinkIdleTicks >= composerBlinkIdleTicksBeforeSettle {
m.composerCursorVisible = true
m.composerBlinkTicking = false
return m, nil
}
m.composerCursorVisible = !m.composerCursorVisible
}
return m, composerBlinkCmd()
return m, composerBlinkCmd(m.composerBlinkSeq)
case tea.BackgroundColorMsg:
// A background reading changes only the contrast direction used for named
// palettes. It never repaints the terminal canvas.
@@ -1487,10 +1539,10 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
case tea.KeyboardEnhancementsMsg:
return m.handleKeyboardEnhancements(msg), nil
case tea.KeyReleaseMsg:
// Voice mode's hold-to-record ends on Space release; every other release
// event is ignored (dispatch elsewhere is press-based).
if m.dictation.voiceModeEnabled && keyIs(msg, tea.KeySpace) {
return m.handleVoiceSpaceRelease()
// Voice mode's hold-to-record ends on Ctrl+Space release; every other
// release event is ignored (dispatch elsewhere is press-based).
if m.dictation.voiceModeEnabled && voiceCaptureReleaseKey(msg) {
return m.handleVoiceCaptureRelease()
}
return m, nil
case tea.KeyPressMsg:
@@ -1624,11 +1676,10 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
return m.appendSystemNotice(fmt.Sprintf("Mouse released — drag to select and copy text. Press %s again to re-enable mouse interaction (clicks, right-click paste).", mouseKey)), nil
}
return m.showTransientNoticeInline("Mouse interaction re-enabled.", transientNoticeSuccess), nil
case m.dictation.voiceModeEnabled && !m.transcriptDetailed && keyIs(msg, tea.KeySpace) && !keyHasMod(msg, tea.ModCtrl) && !keyAlt(msg) && m.noBlockingModal():
// Voice mode (/voice) repurposes Space into the record gesture — the only
// dictation trigger — so it must not also type a space. Turn voice mode
// off (/voice) to type normally.
return m.handleVoiceSpacePress(msg)
case m.dictation.voiceModeEnabled && !m.transcriptDetailed && voiceCaptureKey(msg) && m.noBlockingModal():
// Voice mode reserves Ctrl+Space for recording; ordinary Space continues
// through the composer path below.
return m.handleVoiceCapturePress(msg)
case keyIs(msg, tea.KeyEsc):
// Esc is heavily overloaded below (subchat exit, MCP cancel, ask-user,
// permission deny, wizard/picker/suggestions dismiss, ...) before ever
@@ -2192,6 +2243,8 @@ func (m model) updateModel(msg tea.Msg) (tea.Model, tea.Cmd) {
}
m.terminalFocused = false
m.composerCursorVisible = false
m.composerBlinkTicking = false
m.composerBlinkSeq++
if m.notifier != nil {
m.notifier.SetFocused(false)
}
@@ -5375,6 +5428,9 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
sessionEvents := []pendingSessionEvent{}
usageModelID := m.modelName
var specReview *pendingSpecReviewPrompt
if m.awaitToolReadiness != nil {
m.awaitToolReadiness(runCtx)
}
options := m.agentOptions
options.Registry = cloneToolRegistry(m.registry)
goalAwareRun := !runOptions.specDraft && m.activeLoopID == "" &&
@@ -5417,6 +5473,9 @@ func (m model) runAgentWithOptions(runID int, runCtx context.Context, prompt str
options.ResponseStyle = m.responseStyle
options.Cwd = m.cwd
options.Images = images
if m.newTurnSessionProvider != nil && m.provider != nil {
options.TurnSessionProvider = m.newTurnSessionProvider(m.providerProfile, m.provider)
}
if m.captureRunImages != nil {
m.captureRunImages(images)
}
+156 -11
View File
@@ -109,6 +109,81 @@ func TestPromptSubmitInjectsLiveSessionModelContext(t *testing.T) {
}
}
func TestPromptWaitsForToolReadinessBeforeSnapshot(t *testing.T) {
t.Setenv("XDG_DATA_HOME", t.TempDir())
registry := tools.NewRegistry()
provider := &fakeProvider{events: []zeroruntime.StreamEvent{
{Type: zeroruntime.StreamEventText, Content: "done"},
{Type: zeroruntime.StreamEventDone},
}}
awaited := false
m := newModel(context.Background(), Options{
Cwd: t.TempDir(),
ProviderName: "test",
ModelName: "test-model",
Provider: provider,
Registry: registry,
AwaitToolReadiness: func(context.Context) {
awaited = true
registry.Register(tools.NewScopedReadFileTool(t.TempDir(), nil))
},
})
m.input.SetValue("inspect the workspace")
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if cmd == nil {
t.Fatal("expected prompt submit to start an agent run")
}
next.Update(execCmd(cmd))
if !awaited {
t.Fatal("agent run did not await tool readiness")
}
if len(provider.requests) != 1 {
t.Fatalf("provider requests = %d, want 1", len(provider.requests))
}
for _, definition := range provider.requests[0].Tools {
if definition.Name == "read_file" {
return
}
}
t.Fatal("provider request omitted tool published by readiness callback")
}
func TestPromptBuildsTurnSessionForCurrentProvider(t *testing.T) {
t.Setenv("XDG_DATA_HOME", t.TempDir())
provider := &fakeProvider{events: []zeroruntime.StreamEvent{
{Type: zeroruntime.StreamEventText, Content: "done"},
{Type: zeroruntime.StreamEventDone},
}}
profile := config.ProviderProfile{Name: "chatgpt", Model: "gpt-test"}
called := false
m := newModel(context.Background(), Options{
Cwd: t.TempDir(),
ProviderName: profile.Name,
ModelName: profile.Model,
ProviderProfile: profile,
Provider: provider,
Registry: tools.NewRegistry(),
NewTurnSessionProvider: func(gotProfile config.ProviderProfile, gotProvider zeroruntime.Provider) zeroruntime.TurnSessionProvider {
called = true
if gotProfile.Name != profile.Name || gotProfile.Model != profile.Model || gotProvider != provider {
t.Fatalf("turn session factory inputs = %#v/%#v", gotProfile, gotProvider)
}
return zeroruntime.NewProviderTurnSessionProvider(gotProvider, zeroruntime.ProviderCapabilities{})
},
})
m.input.SetValue("inspect the workspace")
updated, cmd := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
next.Update(execCmd(cmd))
if !called {
t.Fatal("agent run did not build a turn session for the current provider")
}
}
func TestPromptSubmitStoresReasoningSeparatelyFromAnswer(t *testing.T) {
t.Setenv("XDG_DATA_HOME", t.TempDir())
provider := &fakeProvider{events: []zeroruntime.StreamEvent{
@@ -548,10 +623,10 @@ func TestModelCommandSwitchesSessionModel(t *testing.T) {
})
m.input.SetValue("/model gpt-4.1-mini")
updated, cmd := m.Update(testKey(tea.KeyEnter))
updated, _ := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if cmd != nil {
if next.pending {
t.Fatal("expected /model to be handled without starting an agent run")
}
if next.modelName != "gpt-4.1-mini" || next.provider != nextProvider {
@@ -591,10 +666,10 @@ func TestModelCommandAcceptsChatGPTCatalogModelID(t *testing.T) {
}
m.input.SetValue("/model openai/gpt-5.6-sol")
updated, cmd := m.Update(testKey(tea.KeyEnter))
updated, _ := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if cmd != nil {
if next.pending {
t.Fatal("expected /model to be handled without starting an agent run")
}
if next.modelName != "gpt-5.6-sol" || next.provider != nextProvider {
@@ -663,10 +738,10 @@ func TestModelCommandPersistsSelectedModelToUserConfig(t *testing.T) {
})
m.input.SetValue("/model gpt-4.1-mini")
updated, cmd := m.Update(testKey(tea.KeyEnter))
updated, _ := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if cmd != nil {
if next.pending {
t.Fatal("expected /model to be handled without starting an agent run")
}
if next.modelName != "gpt-4.1-mini" {
@@ -2868,6 +2943,8 @@ func TestComposerBlinkStaysSolidWhileTyping(t *testing.T) {
terminalFocused: true,
lastCharTime: base,
composerCursorVisible: true,
composerBlinkSeq: 1,
composerBlinkTicking: true,
}
// Each iteration simulates a keystroke (refreshing lastCharTime) followed by
@@ -2876,7 +2953,7 @@ func TestComposerBlinkStaysSolidWhileTyping(t *testing.T) {
for i := 0; i < 3; i++ {
now = now.Add(200 * time.Millisecond)
m.lastCharTime = now
updated, _ := m.Update(composerBlinkMsg{})
updated, _ := m.Update(composerBlinkMsg{seq: m.composerBlinkSeq})
m = updated.(model)
if !m.composerCursorVisible {
t.Fatalf("tick %d: expected cursor to stay visible while typing, got hidden", i)
@@ -2884,6 +2961,13 @@ func TestComposerBlinkStaysSolidWhileTyping(t *testing.T) {
}
}
func TestMainComposerDisablesEmbeddedCursorBlink(t *testing.T) {
m := newModel(t.Context(), Options{})
if m.input.Styles().Cursor.Blink {
t.Fatal("embedded composer cursor blink must stay disabled")
}
}
func TestComposerBlinkHiddenWhileUnfocused(t *testing.T) {
base := time.Date(2026, 7, 7, 12, 0, 0, 0, time.UTC)
m := model{
@@ -2920,10 +3004,10 @@ func TestComposerBlinkResumesAfterRefocusAndIdle(t *testing.T) {
t.Fatal("expected terminalFocused to be true after FocusMsg")
}
updated, _ = m.Update(composerBlinkMsg{})
updated, _ = m.Update(composerBlinkMsg{seq: m.composerBlinkSeq})
m = updated.(model)
first := m.composerCursorVisible
updated, _ = m.Update(composerBlinkMsg{})
updated, _ = m.Update(composerBlinkMsg{seq: m.composerBlinkSeq})
m = updated.(model)
second := m.composerCursorVisible
if first == second {
@@ -2931,6 +3015,34 @@ func TestComposerBlinkResumesAfterRefocusAndIdle(t *testing.T) {
}
}
func TestComposerBlinkRejectsTickFromBeforeRefocus(t *testing.T) {
base := time.Date(2026, 7, 7, 12, 0, 0, 0, time.UTC)
m := model{
now: func() time.Time { return base },
terminalFocused: true,
lastCharTime: base.Add(-time.Hour),
composerCursorVisible: true,
composerBlinkSeq: 5,
composerBlinkTicking: true,
}
staleSeq := m.composerBlinkSeq
updated, _ := m.Update(tea.BlurMsg{})
m = updated.(model)
updated, refocusCmd := m.Update(tea.FocusMsg{})
m = updated.(model)
if refocusCmd == nil || m.composerBlinkSeq == staleSeq {
t.Fatal("refocus did not re-arm composer blinking with a new sequence")
}
visible := m.composerCursorVisible
updated, staleCmd := m.Update(composerBlinkMsg{seq: staleSeq})
m = updated.(model)
if staleCmd != nil || m.composerCursorVisible != visible {
t.Fatalf("stale tick changed composer state: cmd=%v visible=%v want=%v", staleCmd != nil, m.composerCursorVisible, visible)
}
}
func TestComposerBlinkTogglesWhenIdleAndFocused(t *testing.T) {
base := time.Date(2026, 7, 7, 12, 0, 0, 0, time.UTC)
m := model{
@@ -2938,20 +3050,53 @@ func TestComposerBlinkTogglesWhenIdleAndFocused(t *testing.T) {
terminalFocused: true,
lastCharTime: base.Add(-time.Hour),
composerCursorVisible: true,
composerBlinkSeq: 1,
composerBlinkTicking: true,
}
updated, _ := m.Update(composerBlinkMsg{})
updated, _ := m.Update(composerBlinkMsg{seq: m.composerBlinkSeq})
m = updated.(model)
if m.composerCursorVisible {
t.Fatal("expected cursor to toggle off on first idle+focused tick")
}
updated, _ = m.Update(composerBlinkMsg{})
updated, _ = m.Update(composerBlinkMsg{seq: m.composerBlinkSeq})
m = updated.(model)
if !m.composerCursorVisible {
t.Fatal("expected cursor to toggle back on on second idle+focused tick")
}
}
func TestComposerBlinkSettlesVisibleWithoutRescheduling(t *testing.T) {
m := model{
now: func() time.Time { return time.Unix(100, 0) },
terminalFocused: true,
lastCharTime: time.Unix(0, 0),
composerCursorVisible: true,
composerBlinkSeq: 1,
composerBlinkTicking: true,
}
var cmd tea.Cmd
for range composerBlinkIdleTicksBeforeSettle {
var updated tea.Model
updated, cmd = m.Update(composerBlinkMsg{seq: m.composerBlinkSeq})
m = updated.(model)
}
if m.composerBlinkTicking || cmd != nil || !m.composerCursorVisible {
t.Fatalf("settled cursor = ticking:%v cmd:%v visible:%v", m.composerBlinkTicking, cmd != nil, m.composerCursorVisible)
}
}
func TestComposerInputRestartsSettledBlink(t *testing.T) {
m := newModel(t.Context(), Options{})
m.composerBlinkTicking = false
m.composerCursorVisible = true
updated, cmd := m.Update(tea.KeyPressMsg(tea.Key{Code: 'a', Text: "a"}))
m = updated.(model)
if !m.composerBlinkTicking || cmd == nil || !m.composerCursorVisible {
t.Fatalf("restarted cursor = ticking:%v cmd:%v visible:%v", m.composerBlinkTicking, cmd != nil, m.composerCursorVisible)
}
}
func TestComposerCursorShowsImmediatelyOnKeypress(t *testing.T) {
// The blink phase may have just hidden the caret when the user starts
// typing; the typed character must render with a caret immediately, not
+16 -11
View File
@@ -40,6 +40,7 @@ type Options struct {
CompactionModel string
Provider zeroruntime.Provider
NewProvider func(config.ProviderProfile) (zeroruntime.Provider, error)
NewTurnSessionProvider func(config.ProviderProfile, zeroruntime.Provider) zeroruntime.TurnSessionProvider
ProbeProviderHealth func(context.Context, providerhealth.Options) providerhealth.Result
DiscoverProviderModels func(context.Context, config.ProviderProfile) ([]providermodeldiscovery.Model, error)
DiscoverOllamaContextWindow func(ctx context.Context, baseURL string, model string) (int, error)
@@ -47,17 +48,21 @@ type Options struct {
PrepareRunCompletionWarning func()
RunCompletionWarning func() string
Registry *tools.Registry
SessionStore *sessions.Store
SandboxStore *sandbox.GrantStore
MCPConfig config.MCPConfig
MCPPermissionStore *mcp.PermissionStore
MCPTokenStore *mcp.TokenStore
MCPCommand func(context.Context, []string) MCPCommandResult
SandboxSetupCommand func(context.Context) SandboxSetupCommandResult
UsageTracker *usage.Tracker
SessionCompactor SessionCompactor
PrService *PrService
PeerService *peermsg.Service
// AwaitToolReadiness gives prompt-critical integration startup a bounded
// chance to publish its tools before this turn snapshots the registry. The
// wait runs inside the asynchronous agent command, so the TUI stays usable.
AwaitToolReadiness func(context.Context)
SessionStore *sessions.Store
SandboxStore *sandbox.GrantStore
MCPConfig config.MCPConfig
MCPPermissionStore *mcp.PermissionStore
MCPTokenStore *mcp.TokenStore
MCPCommand func(context.Context, []string) MCPCommandResult
SandboxSetupCommand func(context.Context) SandboxSetupCommandResult
UsageTracker *usage.Tracker
SessionCompactor SessionCompactor
PrService *PrService
PeerService *peermsg.Service
AgentOptions agent.Options
// LoadSkills returns the installed skills (default skills dir merged with any
+4 -4
View File
@@ -521,9 +521,9 @@ func TestModelPickerAppliesActiveProviderCatalogModelID(t *testing.T) {
})
m.input.SetValue("/model openai/gpt-4.1")
updated, cmd := m.Update(testKey(tea.KeyEnter))
updated, _ := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if cmd != nil {
if next.pending {
t.Fatal("expected /model to be handled without starting a run")
}
if captured.Model != "openai/gpt-4.1" {
@@ -867,9 +867,9 @@ func TestModelCommandAcceptsManualModelForCustomProvider(t *testing.T) {
})
m.input.SetValue("/model qwen-custom-latest")
updated, cmd := m.Update(testKey(tea.KeyEnter))
updated, _ := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if cmd != nil {
if next.pending {
t.Fatal("expected /model to be handled without starting a run")
}
if captured.Model != "qwen-custom-latest" {
+6 -6
View File
@@ -46,10 +46,10 @@ func TestEffortCommandListsAndSetsSupportedEffort(t *testing.T) {
}
next.input.SetValue("/effort high")
updated, cmd = next.Update(testKey(tea.KeyEnter))
updated, _ = next.Update(testKey(tea.KeyEnter))
next = updated.(model)
if cmd != nil {
if next.pending {
t.Fatal("expected /effort high to be handled without starting an agent run")
}
if next.reasoningEffort != modelregistry.ReasoningEffortHigh {
@@ -153,10 +153,10 @@ func TestStyleCommandListsAndSetsSessionPreference(t *testing.T) {
}
next.input.SetValue("/style explanatory")
updated, cmd = next.Update(testKey(tea.KeyEnter))
updated, _ = next.Update(testKey(tea.KeyEnter))
next = updated.(model)
if cmd != nil {
if next.pending {
t.Fatal("expected /style explanatory to be handled without starting an agent run")
}
if next.responseStyle != "explanatory" {
@@ -915,10 +915,10 @@ func TestModelSwitchClearsUnsupportedEffortPreference(t *testing.T) {
})
m.input.SetValue("/model gpt-4.1")
updated, cmd := m.Update(testKey(tea.KeyEnter))
updated, _ := m.Update(testKey(tea.KeyEnter))
next := updated.(model)
if cmd != nil {
if next.pending {
t.Fatal("expected /model to be handled without starting an agent run")
}
if next.reasoningEffort != "" {
Loaded 100 of 108 files, more files were not shown because too many files have changed in this diff. Show more