mirror of
https://github.com/Twigpine/zero.git
synced 2026-11-08 09:49:34 -10:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4d9a68b666 | ||
|
|
ad152e0049 | ||
|
|
b4b0418daf | ||
|
|
8e266797a8 | ||
|
|
d37de9214b | ||
|
|
588635487f | ||
|
|
845df6c7d2 | ||
|
|
15eb16f736 | ||
|
|
097c2652cb | ||
|
|
15b7a5475c | ||
|
|
bdbda89f3e | ||
|
|
8acd42a3bd | ||
|
|
1c38e8cba5 | ||
|
|
58daea7a48 | ||
|
|
3e7692d8f0 | ||
|
|
33c94ec7a3 | ||
|
|
c9f9dc02ef | ||
|
|
ec3b822505 | ||
|
|
81a5e6cb71 | ||
|
|
5d1869e7bc | ||
|
|
518d2e53ba |
No files matched your search
@@ -1 +1 @@
|
||||
{".":"0.5.0"}
|
||||
{".":"0.6.0"}
|
||||
@@ -6,6 +6,37 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
|
||||
aims to follow [Semantic Versioning](https://semver.org/spec/v2.0.0.html) once the first release is
|
||||
tagged. Until then, source builds report the version `dev`.
|
||||
|
||||
## [0.6.0](https://github.com/Gitlawb/zero/compare/v0.5.0...v0.6.0) (2026-07-29)
|
||||
|
||||
|
||||
### Features
|
||||
|
||||
* add persistent session goals ([#803](https://github.com/Gitlawb/zero/issues/803)) ([5d1869e](https://github.com/Gitlawb/zero/commit/5d1869e7bc193c97e07418b57b823905c1151d5d))
|
||||
* **providers:** add Fireworks AI provider preset ([#772](https://github.com/Gitlawb/zero/issues/772)) ([58daea7](https://github.com/Gitlawb/zero/commit/58daea7a482927e5ab35afaf97761c60bb35e113))
|
||||
* **skills:** report hash drift in skill info ([#795](https://github.com/Gitlawb/zero/issues/795)) ([9fbdc34](https://github.com/Gitlawb/zero/commit/9fbdc341551601c13434cf7b9b6f5e4ea4c483b1))
|
||||
|
||||
|
||||
### Bug Fixes
|
||||
|
||||
* **acp:** interrupt idle reads on cancellation ([#782](https://github.com/Gitlawb/zero/issues/782)) ([15b7a54](https://github.com/Gitlawb/zero/commit/15b7a5475cfeefddbc59634551451a8872b996b7))
|
||||
* **cli:** include plugins info in shell completions ([#794](https://github.com/Gitlawb/zero/issues/794)) ([fa8734f](https://github.com/Gitlawb/zero/commit/fa8734f17927de8832dfe041e043cfb86f93d12a))
|
||||
* **config:** kill provider-command process tree via job object on Windows ([#690](https://github.com/Gitlawb/zero/issues/690)) ([5c58256](https://github.com/Gitlawb/zero/commit/5c58256abc29c185b449f8ce79647eb3fae61882))
|
||||
* **config:** make the provider-command timeout an upper bound ([#811](https://github.com/Gitlawb/zero/issues/811)) ([#813](https://github.com/Gitlawb/zero/issues/813)) ([1c38e8c](https://github.com/Gitlawb/zero/commit/1c38e8cba50db255f3cda8b7cf4298712b2933e1))
|
||||
* **config:** stop billing provider-command fixtures to the 5s timeout ([#810](https://github.com/Gitlawb/zero/issues/810)) ([f5fa4ff](https://github.com/Gitlawb/zero/commit/f5fa4ffc64537897aa78a56edca36d8f1dacd021))
|
||||
* **cron:** reserve job IDs atomically ([#686](https://github.com/Gitlawb/zero/issues/686)) ([d9b882e](https://github.com/Gitlawb/zero/commit/d9b882eb0a94f38de8447cb54b57860ccd8ed053))
|
||||
* **lint:** resolve static analysis baseline ([#769](https://github.com/Gitlawb/zero/issues/769)) ([583e653](https://github.com/Gitlawb/zero/commit/583e6530bd849a8efcb2b88bbe8c297e8592e62c))
|
||||
* **providers:** retry provably pre-send transport failures ([#447](https://github.com/Gitlawb/zero/issues/447)) ([#750](https://github.com/Gitlawb/zero/issues/750)) ([8acd42a](https://github.com/Gitlawb/zero/commit/8acd42a3bd743a35807d762f23a33fb61c171839))
|
||||
* **sandbox:** add platform-specific write lock around grant state file ([#752](https://github.com/Gitlawb/zero/issues/752)) ([#755](https://github.com/Gitlawb/zero/issues/755)) ([88fb7b6](https://github.com/Gitlawb/zero/commit/88fb7b6192c1ac8173089967d5e9119f92d7a101))
|
||||
* **sandbox:** clarify scoped permission prompts ([#798](https://github.com/Gitlawb/zero/issues/798)) ([9fe02ae](https://github.com/Gitlawb/zero/commit/9fe02ae52e3234cb9dae76e0c5594a287ccbb52a))
|
||||
* **sandbox:** deny reads of git's credential stores ([#815](https://github.com/Gitlawb/zero/issues/815)) ([#816](https://github.com/Gitlawb/zero/issues/816)) ([33c94ec](https://github.com/Gitlawb/zero/commit/33c94ec7a3e7987e6e36b5fd5e037ff4e6e84c26))
|
||||
* **sandbox:** preserve user config and retry denied writes ([#801](https://github.com/Gitlawb/zero/issues/801)) ([81a5e6c](https://github.com/Gitlawb/zero/commit/81a5e6cb716ba4b7ecf7ff3ff261dfa51fd01098))
|
||||
* **sandbox:** scrub the daemon token file pointer ([#677](https://github.com/Gitlawb/zero/issues/677)) ([#818](https://github.com/Gitlawb/zero/issues/818)) ([3e7692d](https://github.com/Gitlawb/zero/commit/3e7692d8f0885d346809869935c81a9cf055d89e))
|
||||
* **swarm:** close lifecycle admission on shutdown ([#776](https://github.com/Gitlawb/zero/issues/776)) ([bdbda89](https://github.com/Gitlawb/zero/commit/bdbda89f3e42d93e8058ab3e1e9ef402659530cd))
|
||||
* **tui:** fix rendering corruption over multipass + Windows Terminal ([#709](https://github.com/Gitlawb/zero/issues/709)) ([5bb1bc1](https://github.com/Gitlawb/zero/commit/5bb1bc11c3188dcea7709ff07c8b25dda3541de8))
|
||||
* **tui:** paste clipboard images with Ctrl+V ([#534](https://github.com/Gitlawb/zero/issues/534)) ([#807](https://github.com/Gitlawb/zero/issues/807)) ([ec3b822](https://github.com/Gitlawb/zero/commit/ec3b8225050eb62f0f41a9582ed2881d2f9dcc8e))
|
||||
* **tui:** reset quiet clock for new runs ([#799](https://github.com/Gitlawb/zero/issues/799)) ([a4a805a](https://github.com/Gitlawb/zero/commit/a4a805a6c82c51cb255e909bd87ca183c8805cf2))
|
||||
* **windows:** use PowerShell for command execution ([#804](https://github.com/Gitlawb/zero/issues/804)) ([518d2e5](https://github.com/Gitlawb/zero/commit/518d2e53ba3fa073090a9b41c278face85b59b64))
|
||||
|
||||
## [0.5.0](https://github.com/Gitlawb/zero/compare/v0.4.0...v0.5.0) (2026-07-22)
|
||||
|
||||
|
||||
|
||||
@@ -131,6 +131,7 @@ export ANTHROPIC_API_KEY=...
|
||||
export GEMINI_API_KEY=...
|
||||
export AIMLAPI_API_KEY=...
|
||||
export LONGCAT_API_KEY=...
|
||||
export FIREWORKS_API_KEY=...
|
||||
export MINIMAX_API_KEY=...
|
||||
export MINIMAXI_API_KEY=...
|
||||
```
|
||||
@@ -147,6 +148,12 @@ To configure Meituan LongCat (LongCat-2.0) directly, run:
|
||||
zero providers setup longcat --set-active
|
||||
```
|
||||
|
||||
To configure Fireworks AI directly, run:
|
||||
|
||||
```bash
|
||||
zero providers setup fireworks --set-active
|
||||
```
|
||||
|
||||
MiniMax presets use the Anthropic-compatible endpoints for the global and China
|
||||
regions:
|
||||
|
||||
|
||||
+138
-6
@@ -17,6 +17,7 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// JSON-RPC 2.0 standard error codes.
|
||||
@@ -80,8 +81,9 @@ type NotifyFunc func(ctx context.Context, params json.RawMessage)
|
||||
// requests/notifications — needed because ACP is bidirectional (the agent calls
|
||||
// the client for session/request_permission, fs/*, terminal/*).
|
||||
type Conn struct {
|
||||
reader *bufio.Reader
|
||||
w io.Writer
|
||||
rawReader io.Reader // wrapped lazily in Serve, once ctx is known — see interruptibleReader
|
||||
readerCloser io.Closer
|
||||
w io.Writer
|
||||
|
||||
writeMu sync.Mutex // serializes all writes to w
|
||||
|
||||
@@ -96,10 +98,15 @@ type Conn struct {
|
||||
wg sync.WaitGroup // tracks in-flight inbound handlers
|
||||
}
|
||||
|
||||
// NewConn builds a peer reading ndjson from r and writing ndjson to w.
|
||||
// NewConn builds a peer reading ndjson from r and writing ndjson to w. This
|
||||
// Conn does not own r: cancelling Serve's context cannot forcibly close it, so
|
||||
// a Read on r that's genuinely idle (never returns on its own) leaves Serve
|
||||
// blocked until r itself produces something — e.g. a test that closes r's own
|
||||
// write side to unblock it. Use NewOwnedConn when this Conn should own r's
|
||||
// lifetime.
|
||||
func NewConn(r io.Reader, w io.Writer) *Conn {
|
||||
return &Conn{
|
||||
reader: bufio.NewReader(r),
|
||||
rawReader: r,
|
||||
w: w,
|
||||
handlers: make(map[string]HandlerFunc),
|
||||
notifiers: make(map[string]NotifyFunc),
|
||||
@@ -107,6 +114,18 @@ func NewConn(r io.Reader, w io.Writer) *Conn {
|
||||
}
|
||||
}
|
||||
|
||||
// NewOwnedConn is like NewConn but declares that this Conn exclusively owns r:
|
||||
// if r implements io.Closer, Serve closes it on context cancellation to
|
||||
// interrupt an otherwise-idle blocking Read. Only pass a reader whose lifetime
|
||||
// this Conn should control — e.g. the process's own stdin — never a reader a
|
||||
// caller retains or shares elsewhere, since cancellation tears it down
|
||||
// permanently for every consumer, not just this Conn.
|
||||
func NewOwnedConn(r io.Reader, w io.Writer) *Conn {
|
||||
c := NewConn(r, w)
|
||||
c.readerCloser, _ = r.(io.Closer)
|
||||
return c
|
||||
}
|
||||
|
||||
// Handle registers a request handler for method.
|
||||
func (c *Conn) Handle(method string, fn HandlerFunc) { c.handlers[method] = fn }
|
||||
|
||||
@@ -128,17 +147,36 @@ func (c *Conn) Serve(ctx context.Context) error {
|
||||
c.wg.Wait()
|
||||
}()
|
||||
|
||||
interruptible := newInterruptibleReader(ctx, c.rawReader, c.readerCloser)
|
||||
reader := bufio.NewReader(interruptible)
|
||||
|
||||
for {
|
||||
// One ndjson value per line. Framing per line (rather than streaming the
|
||||
// decoder) keeps a single malformed line from making the whole connection
|
||||
// unrecoverable — we report -32700 and continue.
|
||||
line, err := c.reader.ReadBytes('\n')
|
||||
//
|
||||
// generation brackets this SPECIFIC ReadBytes call: bufio may invoke
|
||||
// interruptible.Read zero or more times underneath it (zero when the
|
||||
// answer is already sitting in bufio's own buffer — including a
|
||||
// previously-buffered, not-yet-surfaced error: bufio can return a
|
||||
// complete line with a nil error while quietly caching an error it
|
||||
// received ALONGSIDE that line's bytes, which then surfaces on the VERY
|
||||
// NEXT call with no further contact with interruptible.Read at all). The
|
||||
// generation only advances inside interruptible.Read's own cancellation
|
||||
// branch, so "unchanged across this call" is proof this call's outcome
|
||||
// didn't pass through that branch — whether because it was answered from
|
||||
// the buffer or because it raced ctx.Done() and won. Only that proof
|
||||
// makes it safe to call the outcome genuine.
|
||||
before := interruptible.generation()
|
||||
line, err := reader.ReadBytes('\n')
|
||||
interrupted := interruptible.generation() != before
|
||||
|
||||
if len(bytes.TrimSpace(line)) > 0 {
|
||||
c.handleLine(ctx, line)
|
||||
}
|
||||
if err != nil {
|
||||
c.failAllPending(err)
|
||||
if errors.Is(err, io.EOF) || ctx.Err() != nil {
|
||||
if errors.Is(err, io.EOF) || interrupted {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
@@ -146,6 +184,100 @@ func (c *Conn) Serve(ctx context.Context) error {
|
||||
}
|
||||
}
|
||||
|
||||
// interruptibleReader wraps a reader so a context cancellation can unblock a
|
||||
// currently in-flight Read without ever discarding — or misattributing — the
|
||||
// real outcome of a call that wasn't actually blocked.
|
||||
//
|
||||
// Each Read runs the underlying call on its own goroutine and races it against
|
||||
// ctx.Done(). If the underlying call finishes first, its result is returned
|
||||
// exactly as received; ctx being cancelled a moment later has no bearing on a
|
||||
// call that already had its answer. If ctx.Done() fires first, closer (when
|
||||
// non-nil) is closed to nudge a genuinely blocked read, the generation counter
|
||||
// is bumped, and the call STILL waits for and returns the underlying call's own
|
||||
// result — never a synthesized one — since closing a reader's outcome is
|
||||
// whatever the reader itself reports for it, not ours to invent.
|
||||
//
|
||||
// The counter is what lets Serve tell a call that was actually raced against
|
||||
// cancellation apart from one merely answered while cancellation happened to be
|
||||
// pending — the distinction the read loop's own doc comment explains.
|
||||
type interruptibleReader struct {
|
||||
inner io.Reader
|
||||
closer io.Closer
|
||||
ctx context.Context
|
||||
|
||||
closeOnce sync.Once
|
||||
gen atomic.Int64
|
||||
}
|
||||
|
||||
func newInterruptibleReader(ctx context.Context, inner io.Reader, closer io.Closer) *interruptibleReader {
|
||||
return &interruptibleReader{inner: inner, closer: closer, ctx: ctx}
|
||||
}
|
||||
|
||||
func (r *interruptibleReader) generation() int64 { return r.gen.Load() }
|
||||
|
||||
type interruptibleReadResult struct {
|
||||
n int
|
||||
err error
|
||||
}
|
||||
|
||||
// interruptibleReadDecision arbitrates whether an underlying Read completed or
|
||||
// cancellation claimed it first. Completion is recorded before the result is
|
||||
// published so cancellation cannot misclassify a finished read merely because
|
||||
// its goroutine has not sent the result yet.
|
||||
type interruptibleReadDecision struct {
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (d *interruptibleReadDecision) complete() { d.once.Do(func() {}) }
|
||||
|
||||
func (d *interruptibleReadDecision) interrupt() (claimed bool) {
|
||||
d.once.Do(func() { claimed = true })
|
||||
return claimed
|
||||
}
|
||||
|
||||
// readInner records completion in the Read call's deferred epilogue. It also
|
||||
// converts a reader panic into a terminal error so the helper goroutine always
|
||||
// publishes a result for Read to join.
|
||||
func (r *interruptibleReader) readInner(p []byte, decision *interruptibleReadDecision) (res interruptibleReadResult) {
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
res = interruptibleReadResult{err: fmt.Errorf("acp: reader panicked: %v", recovered)}
|
||||
}
|
||||
decision.complete()
|
||||
}()
|
||||
res.n, res.err = r.inner.Read(p)
|
||||
return res
|
||||
}
|
||||
|
||||
func (r *interruptibleReader) Read(p []byte) (int, error) {
|
||||
resultCh := make(chan interruptibleReadResult, 1)
|
||||
var decision interruptibleReadDecision
|
||||
go func() {
|
||||
resultCh <- r.readInner(p, &decision)
|
||||
}()
|
||||
select {
|
||||
case res := <-resultCh:
|
||||
return res.n, res.err
|
||||
case <-r.ctx.Done():
|
||||
// resultCh may not be readable yet even when inner.Read has returned: its
|
||||
// goroutine can be descheduled between recording completion and publishing
|
||||
// the result. Only interrupt a call that cancellation actually claims.
|
||||
if decision.interrupt() {
|
||||
r.gen.Add(1)
|
||||
if r.closer != nil {
|
||||
r.closeOnce.Do(func() { _ = r.closer.Close() })
|
||||
}
|
||||
}
|
||||
// Wait for the real result rather than returning synthetically: p is
|
||||
// still being written by the goroutine above until it does, so returning
|
||||
// early would race that write, and the actual error (whatever closing
|
||||
// the reader produced, or whatever the reader was already about to
|
||||
// report) is what the caller should see.
|
||||
res := <-resultCh
|
||||
return res.n, res.err
|
||||
}
|
||||
}
|
||||
|
||||
// handleLine decodes and dispatches one ndjson frame. A parse failure replies
|
||||
// with -32700 (id null) and keeps the connection alive.
|
||||
func (c *Conn) handleLine(ctx context.Context, line []byte) {
|
||||
|
||||
@@ -3,11 +3,30 @@ package acp
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type testReadCloser struct {
|
||||
read func([]byte) (int, error)
|
||||
close func() error
|
||||
}
|
||||
|
||||
func (r testReadCloser) Read(p []byte) (int, error) { return r.read(p) }
|
||||
func (r testReadCloser) Close() error { return r.close() }
|
||||
|
||||
type testReader func([]byte) (int, error)
|
||||
|
||||
func (r testReader) Read(p []byte) (int, error) { return r(p) }
|
||||
|
||||
type testWriter func([]byte) (int, error)
|
||||
|
||||
func (w testWriter) Write(p []byte) (int, error) { return w(p) }
|
||||
|
||||
// connPair wires two Conns together over in-memory pipes and serves both.
|
||||
func connPair(t *testing.T) (a, b *Conn, stop func()) {
|
||||
t.Helper()
|
||||
@@ -25,6 +44,275 @@ func connPair(t *testing.T) (a, b *Conn, stop func()) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestConnServeCancellationInterruptsIdleRead(t *testing.T) {
|
||||
pipeReader, writer := io.Pipe()
|
||||
t.Cleanup(func() { _ = writer.Close() })
|
||||
// Buffered + non-blocking send, not close: interruptibleReader retries each
|
||||
// underlying Read on its own goroutine, so if bufio ever calls this more than
|
||||
// once, a close() here would panic on the second call.
|
||||
readStarted := make(chan struct{}, 1)
|
||||
closeCalls := make(chan struct{}, 2)
|
||||
reader := testReadCloser{
|
||||
read: func(p []byte) (int, error) {
|
||||
select {
|
||||
case readStarted <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
return pipeReader.Read(p)
|
||||
},
|
||||
close: func() error {
|
||||
closeCalls <- struct{}{}
|
||||
return pipeReader.Close()
|
||||
},
|
||||
}
|
||||
conn := NewOwnedConn(reader, io.Discard)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- conn.Serve(ctx) }()
|
||||
|
||||
<-readStarted
|
||||
cancel()
|
||||
select {
|
||||
case err := <-done:
|
||||
if err != nil {
|
||||
t.Fatalf("Serve returned %v after cancellation, want nil", err)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("Serve did not return after cancelling an idle connection")
|
||||
}
|
||||
if got := len(closeCalls); got != 1 {
|
||||
t.Fatalf("reader Close calls = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConnServePreservesTerminalReadErrorDuringCancellation(t *testing.T) {
|
||||
for _, closable := range []bool{false, true} {
|
||||
name := "non-closable"
|
||||
if closable {
|
||||
name = "closable"
|
||||
}
|
||||
t.Run(name, func(t *testing.T) {
|
||||
wantErr := errors.New("read failed")
|
||||
read := testReader(func(p []byte) (int, error) {
|
||||
return copy(p, "not json"), wantErr
|
||||
})
|
||||
closeCalls := make(chan struct{}, 1)
|
||||
var reader io.Reader = read
|
||||
if closable {
|
||||
reader = testReadCloser{
|
||||
read: read,
|
||||
close: func() error {
|
||||
closeCalls <- struct{}{}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
}
|
||||
// Buffered + non-blocking send, not close: a later change adding a
|
||||
// second write (e.g. an additional error frame) must not panic on a
|
||||
// repeated close of an already-closed channel.
|
||||
writeStarted := make(chan struct{}, 1)
|
||||
releaseWrite := make(chan struct{})
|
||||
writer := testWriter(func(p []byte) (int, error) {
|
||||
select {
|
||||
case writeStarted <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
<-releaseWrite
|
||||
return len(p), nil
|
||||
})
|
||||
newConn := NewConn
|
||||
if closable {
|
||||
// Ownership matters here: this case exists to prove that even a
|
||||
// Conn entitled to close its reader on cancellation still
|
||||
// preserves a terminal error that was already decided.
|
||||
newConn = NewOwnedConn
|
||||
}
|
||||
conn := newConn(reader, writer)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- conn.Serve(ctx) }()
|
||||
|
||||
<-writeStarted
|
||||
cancel()
|
||||
close(releaseWrite)
|
||||
select {
|
||||
case err := <-done:
|
||||
if !errors.Is(err, wantErr) {
|
||||
t.Fatalf("Serve returned %v, want terminal read error %v", err, wantErr)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("Serve did not return after terminal read error")
|
||||
}
|
||||
if got := len(closeCalls); got != 0 {
|
||||
t.Fatalf("reader Close calls = %d, want 0 after terminal read", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestConnServePreservesReadErrorBufferedAlongsideAPriorLine is the deterministic
|
||||
// regression test for jatmn's #782 finding: a real, non-EOF ReadBytes failure can
|
||||
// be DECIDED before cancellation ever happens, yet only SURFACE afterward, and
|
||||
// must still be reported rather than swallowed as a clean shutdown.
|
||||
//
|
||||
// bufio.Reader can obtain data AND a terminal error from ONE underlying Read call
|
||||
// (io.Reader explicitly permits returning (n>0, err) together). If a delimiter is
|
||||
// found within that data, ReadBytes returns the line with a nil error for THIS
|
||||
// call, but bufio caches the error internally and returns it — WITHOUT calling the
|
||||
// underlying reader again — on the very next ReadBytes call. This test forces that
|
||||
// exact sequence: the underlying reader hands back one complete, validly-framed
|
||||
// line together with wantErr in a SINGLE call (asserted below to be its ONLY
|
||||
// call), the loop dispatches that line through a synchronous path that blocks the
|
||||
// read loop, cancellation is asserted to have happened before the loop attempts
|
||||
// its next read, and only then is the block released. By the time the read loop
|
||||
// asks for more input, cancellation is already in effect and the underlying
|
||||
// reader is never touched again — so if Serve attributed the resulting error to
|
||||
// cancellation merely because ctx was already done, it would incorrectly return
|
||||
// nil. It must instead recognize that this read was answered entirely from
|
||||
// bufio's own cache, never raced against ctx.Done(), and report wantErr.
|
||||
func TestConnServePreservesReadErrorBufferedAlongsideAPriorLine(t *testing.T) {
|
||||
wantErr := errors.New("read failed")
|
||||
var readCalls atomic.Int32
|
||||
read := testReader(func(p []byte) (int, error) {
|
||||
readCalls.Add(1)
|
||||
// An unsupported jsonrpc version with an id takes the SYNCHRONOUS
|
||||
// writeError path in handleLine (not a dispatched goroutine), giving this
|
||||
// test a reliable blocking point between processing this line and the
|
||||
// loop's next ReadBytes call.
|
||||
return copy(p, `{"jsonrpc":"1.0","id":1}`+"\n"), wantErr
|
||||
})
|
||||
writeStarted := make(chan struct{}, 1)
|
||||
releaseWrite := make(chan struct{})
|
||||
var releaseWriteOnce sync.Once
|
||||
release := func() { releaseWriteOnce.Do(func() { close(releaseWrite) }) }
|
||||
t.Cleanup(release)
|
||||
writer := testWriter(func(p []byte) (int, error) {
|
||||
select {
|
||||
case writeStarted <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
<-releaseWrite
|
||||
return len(p), nil
|
||||
})
|
||||
conn := NewConn(read, writer)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- conn.Serve(ctx) }()
|
||||
|
||||
// The synchronous write confirms the first (and, per the panic guard above,
|
||||
// only) read has already returned and is being processed — cancelling now
|
||||
// lands squarely in the gap between that read and the loop's next one, never
|
||||
// racing an in-flight interruptibleReader.Read call.
|
||||
select {
|
||||
case <-writeStarted:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("Serve did not reach the synchronous write")
|
||||
}
|
||||
cancel()
|
||||
release()
|
||||
|
||||
select {
|
||||
case err := <-done:
|
||||
if !errors.Is(err, wantErr) {
|
||||
t.Fatalf("Serve returned %v, want the buffered terminal read error %v", err, wantErr)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("Serve did not return after the buffered terminal read error")
|
||||
}
|
||||
if got := readCalls.Load(); got != 1 {
|
||||
t.Fatalf("underlying Read calls = %d, want exactly 1 (bufio must have served the second ReadBytes from its cache)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestInterruptibleReaderPrefersDecidedResultOverCancellation is the regression
|
||||
// test for jatmn's second #782 finding: interruptibleReader.Read's select
|
||||
// between resultCh and ctx.Done() chooses pseudo-randomly when both are ready,
|
||||
// so a real, already-decided outcome (e.g. a terminal transport error) could be
|
||||
// misattributed to cancellation purely because ctx also happened to be done by
|
||||
// the time the select ran. Each iteration lets the underlying Read return (so
|
||||
// resultCh already holds the real outcome) before cancelling, maximizing the
|
||||
// odds of landing exactly in that ambiguous window; the generation counter must
|
||||
// never advance when the real, already-decided result is what gets returned.
|
||||
func TestInterruptibleReaderPrefersDecidedResultOverCancellation(t *testing.T) {
|
||||
wantErr := errors.New("read failed")
|
||||
for i := 0; i < 100; i++ {
|
||||
resultReady := make(chan struct{})
|
||||
read := testReader(func(p []byte) (int, error) {
|
||||
defer close(resultReady)
|
||||
return copy(p, "x"), wantErr
|
||||
})
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
r := newInterruptibleReader(ctx, read, nil)
|
||||
|
||||
readDone := make(chan interruptibleReadResult, 1)
|
||||
go func() {
|
||||
n, err := r.Read(make([]byte, 8))
|
||||
readDone <- interruptibleReadResult{n, err}
|
||||
}()
|
||||
|
||||
<-resultReady
|
||||
time.Sleep(time.Millisecond) // bias resultCh's send ahead of cancel
|
||||
cancel()
|
||||
|
||||
res := <-readDone
|
||||
if !errors.Is(res.err, wantErr) {
|
||||
t.Fatalf("iteration %d: Read returned %v, want %v", i, res.err, wantErr)
|
||||
}
|
||||
if got := r.generation(); got != 0 {
|
||||
t.Fatalf("iteration %d: generation = %d, want 0 (a decided result must never be attributed to cancellation)", i, got)
|
||||
}
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
|
||||
// TestInterruptibleReadCompletionWinsBeforeResultPublication exercises the
|
||||
// actual inner-read wrapper at the handoff between the underlying Read and the
|
||||
// helper goroutine's result publication. Cancellation must not reclassify a
|
||||
// result returned by this wrapper as interrupted.
|
||||
func TestInterruptibleReadCompletionWinsBeforeResultPublication(t *testing.T) {
|
||||
wantErr := errors.New("read failed")
|
||||
read := testReader(func(p []byte) (int, error) {
|
||||
return copy(p, "x"), wantErr
|
||||
})
|
||||
r := newInterruptibleReader(context.Background(), read, nil)
|
||||
var decision interruptibleReadDecision
|
||||
readComplete := make(chan struct{})
|
||||
resultCh := make(chan interruptibleReadResult)
|
||||
go func() {
|
||||
res := r.readInner(make([]byte, 8), &decision)
|
||||
close(readComplete)
|
||||
resultCh <- res
|
||||
}()
|
||||
|
||||
<-readComplete
|
||||
if decision.interrupt() {
|
||||
t.Fatal("cancellation claimed a read that completed before result publication")
|
||||
}
|
||||
res := <-resultCh
|
||||
if res.n != 1 || !errors.Is(res.err, wantErr) {
|
||||
t.Fatalf("inner read result = (%d, %v), want (1, %v)", res.n, res.err, wantErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConnServeReportsReaderPanicAsTerminalError(t *testing.T) {
|
||||
read := testReader(func([]byte) (int, error) {
|
||||
panic("boom")
|
||||
})
|
||||
conn := NewConn(read, io.Discard)
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- conn.Serve(context.Background()) }()
|
||||
|
||||
select {
|
||||
case err := <-done:
|
||||
const want = "acp: reader panicked: boom"
|
||||
if err == nil || err.Error() != want {
|
||||
t.Fatalf("Serve returned %v, want %q", err, want)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("Serve did not return after reader panic")
|
||||
}
|
||||
}
|
||||
|
||||
func TestConnRequestResponse(t *testing.T) {
|
||||
a, b, stop := connPair(t)
|
||||
defer stop()
|
||||
|
||||
@@ -5,9 +5,29 @@ import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/sandbox"
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
)
|
||||
|
||||
type normalizationRecordingShellTool struct {
|
||||
calls []map[string]any
|
||||
}
|
||||
|
||||
func (tool *normalizationRecordingShellTool) Name() string { return "bash" }
|
||||
func (tool *normalizationRecordingShellTool) Description() string {
|
||||
return "record normalized shell arguments"
|
||||
}
|
||||
func (tool *normalizationRecordingShellTool) Parameters() tools.Schema {
|
||||
return tools.Schema{Type: "object", AdditionalProperties: true}
|
||||
}
|
||||
func (tool *normalizationRecordingShellTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectShell, Permission: tools.PermissionAllow}
|
||||
}
|
||||
func (tool *normalizationRecordingShellTool) Run(_ context.Context, args map[string]any) tools.Result {
|
||||
tool.calls = append(tool.calls, cloneArgs(args))
|
||||
return tools.Result{Status: tools.StatusOK}
|
||||
}
|
||||
|
||||
// A shell call that sets sandbox_permissions: with_additional_permissions but
|
||||
// omits (or malforms) additional_permissions can never succeed: the same
|
||||
// validation runs again at grant time regardless of the user's decision. It
|
||||
@@ -144,3 +164,120 @@ func TestExecuteToolCallStillPromptsForValidAdditionalPermissions(t *testing.T)
|
||||
t.Fatalf("OnPermissionRequest called %d times, want exactly 1 for a valid elevation request", promptCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeProviderExpandedEmptyAdditionalPermissions(t *testing.T) {
|
||||
for _, mode := range []string{
|
||||
string(tools.SandboxPermissionsUseDefault),
|
||||
string(tools.SandboxPermissionsWithAdditionalPermissions),
|
||||
} {
|
||||
args := map[string]any{
|
||||
"command": "./zero --version",
|
||||
"sandbox_permissions": mode,
|
||||
"additional_permissions": map[string]any{
|
||||
"file_system": map[string]any{
|
||||
"deny_read": []any{},
|
||||
"entries": []any{},
|
||||
"read": []any{},
|
||||
"write": []any{},
|
||||
},
|
||||
"network": map[string]any{"enabled": false},
|
||||
},
|
||||
}
|
||||
normalizeNoopAdditionalPermissions(args, t.TempDir(), nil)
|
||||
if _, exists := args["additional_permissions"]; exists {
|
||||
t.Fatalf("mode %q retained empty additional_permissions: %#v", mode, args)
|
||||
}
|
||||
if got := args["sandbox_permissions"]; got != string(tools.SandboxPermissionsUseDefault) {
|
||||
t.Fatalf("mode %q normalized to %q, want use_default", mode, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteToolCallNormalizesProviderExpandedEmptyAdditionalPermissions(t *testing.T) {
|
||||
tool := &normalizationRecordingShellTool{}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tool)
|
||||
|
||||
result, err := executeToolCall(
|
||||
context.Background(),
|
||||
registry,
|
||||
ToolCall{
|
||||
ID: "c1",
|
||||
Name: "bash",
|
||||
Arguments: `{"command":"./zero --version","sandbox_permissions":"with_additional_permissions",` +
|
||||
`"additional_permissions":{"file_system":{"deny_read":[],"entries":[],"read":[],"write":[]},"network":{"enabled":false}}}`,
|
||||
},
|
||||
PermissionModeAuto,
|
||||
Options{Cwd: t.TempDir()},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("executeToolCall: %v", err)
|
||||
}
|
||||
if result.Status != tools.StatusOK {
|
||||
t.Fatalf("provider-expanded empty permissions were rejected: %s", result.Output)
|
||||
}
|
||||
if len(tool.calls) != 1 {
|
||||
t.Fatalf("shell tool calls = %d, want 1", len(tool.calls))
|
||||
}
|
||||
args := tool.calls[0]
|
||||
if _, exists := args["additional_permissions"]; exists {
|
||||
t.Fatalf("executeToolCall retained empty additional_permissions: %#v", args)
|
||||
}
|
||||
if got := args["sandbox_permissions"]; got != string(tools.SandboxPermissionsUseDefault) {
|
||||
t.Fatalf("sandbox mode = %q, want use_default", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeAdditionalPermissionsPreservesRealRequest(t *testing.T) {
|
||||
args := map[string]any{
|
||||
"sandbox_permissions": string(tools.SandboxPermissionsWithAdditionalPermissions),
|
||||
"additional_permissions": map[string]any{
|
||||
"network": map[string]any{"enabled": true},
|
||||
},
|
||||
}
|
||||
normalizeNoopAdditionalPermissions(args, t.TempDir(), nil)
|
||||
if _, exists := args["additional_permissions"]; !exists {
|
||||
t.Fatal("non-empty additional_permissions was removed")
|
||||
}
|
||||
if got := args["sandbox_permissions"]; got != string(tools.SandboxPermissionsWithAdditionalPermissions) {
|
||||
t.Fatalf("sandbox mode changed to %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeAdditionalPermissionsRemovesWorkspaceReadAlreadyCoveredBySandbox(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
engine := sandbox.NewEngine(sandbox.EngineOptions{
|
||||
WorkspaceRoot: root,
|
||||
Policy: sandbox.DefaultPolicy(),
|
||||
})
|
||||
args := map[string]any{
|
||||
"command": "./zero --version",
|
||||
"sandbox_permissions": string(tools.SandboxPermissionsWithAdditionalPermissions),
|
||||
"additional_permissions": map[string]any{
|
||||
"file_system": map[string]any{
|
||||
"entries": []any{
|
||||
map[string]any{
|
||||
"access": "read",
|
||||
"path": map[string]any{
|
||||
"type": "path",
|
||||
"path": ".",
|
||||
},
|
||||
},
|
||||
},
|
||||
"deny_read": []any{},
|
||||
"read": []any{},
|
||||
"write": []any{},
|
||||
},
|
||||
"network": map[string]any{"enabled": false},
|
||||
},
|
||||
}
|
||||
|
||||
normalizeNoopAdditionalPermissions(args, root, engine)
|
||||
|
||||
if _, exists := args["additional_permissions"]; exists {
|
||||
t.Fatalf("already-covered additional_permissions was retained: %#v", args)
|
||||
}
|
||||
if got := args["sandbox_permissions"]; got != string(tools.SandboxPermissionsUseDefault) {
|
||||
t.Fatalf("sandbox mode changed to %q", got)
|
||||
}
|
||||
}
|
||||
@@ -413,6 +413,14 @@ func (state *compactionState) maybeCompact(
|
||||
// the bodies of old, large tool results (the model has already acted on
|
||||
// them). If that brings us back under threshold, skip the paid summarizer
|
||||
// entirely and preserve recent turns verbatim.
|
||||
if pruned, reclaimed := pruneSupersededReadResults(messages); reclaimed > 0 {
|
||||
messages = pruned
|
||||
size = state.calibratedTokens(estimateTokens(messages) + toolTokens)
|
||||
if size <= state.threshold {
|
||||
state.lowWaterMark = size
|
||||
return messages
|
||||
}
|
||||
}
|
||||
if pruned, reclaimed := pruneStaleToolOutput(messages, state.preserveLast); reclaimed > 0 {
|
||||
messages = pruned
|
||||
size = state.calibratedTokens(estimateTokens(messages) + toolTokens)
|
||||
|
||||
@@ -379,6 +379,43 @@ func TestRunNoCompactionWhenContextWindowZero(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestMaybeCompactStopsAfterSupersededReadsReachThreshold(t *testing.T) {
|
||||
body := strings.Repeat("same source line\n", 1000)
|
||||
messages := []zeroruntime.Message{
|
||||
{Role: zeroruntime.MessageRoleSystem, Content: "system"},
|
||||
toolCallWithArgs("old", "read_file", `{"path":"a.go"}`),
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "old", Content: body},
|
||||
toolCallWithArgs("new", "read_file", `{"path":"a.go"}`),
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "new", Content: body},
|
||||
}
|
||||
pruned, reclaimed := pruneSupersededReadResults(messages)
|
||||
if reclaimed <= 0 {
|
||||
t.Fatal("test setup did not reclaim superseded read output")
|
||||
}
|
||||
prunedSize := estimateTokens(pruned)
|
||||
if estimateTokens(messages) <= prunedSize {
|
||||
t.Fatal("test setup does not cross the compaction threshold")
|
||||
}
|
||||
|
||||
provider := &summarizeRecordingProvider{}
|
||||
state := &compactionState{
|
||||
enabled: true,
|
||||
threshold: prunedSize,
|
||||
preserveLast: 2,
|
||||
}
|
||||
got := state.maybeCompact(context.Background(), provider, messages, nil)
|
||||
|
||||
if provider.summarizeCalls != 0 {
|
||||
t.Fatalf("cheap pruning should avoid the summarizer, got %d calls", provider.summarizeCalls)
|
||||
}
|
||||
if state.lowWaterMark != prunedSize {
|
||||
t.Fatalf("lowWaterMark = %d, want %d", state.lowWaterMark, prunedSize)
|
||||
}
|
||||
if !strings.Contains(got[2].Content, "superseded identical") || got[4].Content != body {
|
||||
t.Fatalf("unexpected pruned messages: %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// reactiveProvider builds up history with a tool call on turn 1, then errors
|
||||
// with a context-limit message on turn 2 (once the history is large enough that
|
||||
// compaction can shrink it), then succeeds on the same-turn retry. Summary
|
||||
|
||||
@@ -341,15 +341,15 @@ func TestPartitionToolsActiveHidesUnloadedExposesLoaded(t *testing.T) {
|
||||
if search.Description != discovery || !strings.Contains(search.Description, "mcp") {
|
||||
t.Fatalf("tool_search description must carry discovery source, got %q", search.Description)
|
||||
}
|
||||
if strings.Contains(search.Description, "mcp__srv__alpha") {
|
||||
t.Fatalf("tool_search description must not list already-loaded tool names, got %q", search.Description)
|
||||
if !strings.Contains(search.Description, "mcp__srv__alpha") {
|
||||
t.Fatalf("stable tool_search catalog must retain already-loaded tool names, got %q", search.Description)
|
||||
}
|
||||
if !strings.Contains(search.Description, "mcp__srv__beta") {
|
||||
t.Fatalf("tool_search description must list exact hidden tool names, got %q", search.Description)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPartitionToolsActiveNothingHiddenEmptyReminder(t *testing.T) {
|
||||
func TestPartitionToolsActiveNothingHiddenKeepsStableCatalog(t *testing.T) {
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(fakeDeferredTool{name: "mcp__srv__alpha", desc: "alpha"})
|
||||
registry.Register(fakeDeferredTool{name: "mcp__srv__beta", desc: "beta"})
|
||||
@@ -368,9 +368,8 @@ func TestPartitionToolsActiveNothingHiddenEmptyReminder(t *testing.T) {
|
||||
if !exposedNames["tool_search"] {
|
||||
t.Fatalf("expected tool_search exposed on active path, got %#v", exposed)
|
||||
}
|
||||
// No hidden tools means no dynamic discovery text is needed.
|
||||
if reminder != "" {
|
||||
t.Fatalf("expected empty discovery text when nothing is hidden, got %q", reminder)
|
||||
if !strings.Contains(reminder, "mcp__srv__alpha") || !strings.Contains(reminder, "mcp__srv__beta") {
|
||||
t.Fatalf("stable catalog must retain all deferred names after loading, got %q", reminder)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+94
-43
@@ -1076,6 +1076,9 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
if isShellCommandTool(call.Name) {
|
||||
normalizeNoopAdditionalPermissions(args, options.Cwd, options.Sandbox)
|
||||
}
|
||||
// tool_search is the gateway to the allowlisted deferred tools, so a non-empty
|
||||
// EnabledTools allowlist that omits it must NOT reject the call — otherwise the
|
||||
// discovery tool exposed to the model rejects at dispatch time (an inescapable
|
||||
@@ -1354,7 +1357,8 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
|
||||
Cwd: options.Cwd,
|
||||
// Per-session file version tracker so write_file/edit_file refuse to clobber
|
||||
// a file that changed on disk outside Zero since it was last read.
|
||||
FileTracker: options.FileTracker,
|
||||
FileTracker: options.FileTracker,
|
||||
DeferFileObservationCommit: true,
|
||||
// Post-edit diagnostics are NOT forwarded to the tools: they used to run
|
||||
// synchronously inside edit_file/write_file and block every mutating tool
|
||||
// call on the language server (≥300ms debounce floor, 10s cap). The loop
|
||||
@@ -1406,6 +1410,7 @@ func executeToolCall(ctx context.Context, registry *tools.Registry, call ToolCal
|
||||
result = registry.RebudgetAfterHook(call.Name, args, result)
|
||||
}
|
||||
}
|
||||
result = registry.CommitFileObservation(result, options.FileTracker)
|
||||
// Stamp the sandbox risk classification for this EXECUTED call so run-policy
|
||||
// observers can see the risk of an allowed mutation. Prefer the preflight
|
||||
// decision's classification when the sandbox evaluated the call; otherwise
|
||||
@@ -1617,9 +1622,15 @@ func sandboxRestrictedShellRetryCandidate(call ToolCall, args map[string]any, re
|
||||
|
||||
func sandboxDeniedShellResult(result tools.Result) bool {
|
||||
if result.ExecutionOutcome != nil {
|
||||
// Typed denials must be handled by their exact capability path. Never
|
||||
// turn a structured narrow denial into the legacy unrestricted retry.
|
||||
return false
|
||||
// Only a typed platform denial that explicitly requests unrestricted
|
||||
// recovery may enter the legacy retry path. Narrow capability denials
|
||||
// remain on their exact capability path.
|
||||
denial := result.ExecutionOutcome.Denial
|
||||
return denial != nil &&
|
||||
denial.Source == execution.DenialSourcePlatformSandbox &&
|
||||
denial.Capability.Kind == execution.CapabilityUnrestricted &&
|
||||
denial.Recoverable &&
|
||||
denial.NextAction == execution.DenialNextActionRequestApproval
|
||||
}
|
||||
return result.Status == tools.StatusError && result.Meta[tools.SandboxLikelyDeniedMeta] == "true"
|
||||
}
|
||||
@@ -1684,20 +1695,21 @@ func sandboxRestrictionRetryReason(result tools.Result) string {
|
||||
|
||||
func runToolForNetworkRetry(ctx context.Context, registry *tools.Registry, name string, toolCallID string, args map[string]any, permissionMode PermissionMode, options Options, progressCallback func(streamjson.Event)) tools.Result {
|
||||
return registry.RunWithOptions(ctx, name, cloneArgs(args), tools.RunOptions{
|
||||
PermissionGranted: true,
|
||||
PermissionMode: string(permissionMode),
|
||||
Autonomy: options.Autonomy,
|
||||
Sandbox: options.Sandbox,
|
||||
ToolCallID: toolCallID,
|
||||
SessionID: options.SessionID,
|
||||
Model: options.Model,
|
||||
ReasoningEffort: options.ReasoningEffort,
|
||||
Depth: options.Depth,
|
||||
Cwd: options.Cwd,
|
||||
FileTracker: options.FileTracker,
|
||||
EnabledTools: options.EnabledTools,
|
||||
DisabledTools: options.DisabledTools,
|
||||
Progress: progressCallback,
|
||||
PermissionGranted: true,
|
||||
PermissionMode: string(permissionMode),
|
||||
Autonomy: options.Autonomy,
|
||||
Sandbox: options.Sandbox,
|
||||
ToolCallID: toolCallID,
|
||||
SessionID: options.SessionID,
|
||||
Model: options.Model,
|
||||
ReasoningEffort: options.ReasoningEffort,
|
||||
Depth: options.Depth,
|
||||
Cwd: options.Cwd,
|
||||
FileTracker: options.FileTracker,
|
||||
DeferFileObservationCommit: true,
|
||||
EnabledTools: options.EnabledTools,
|
||||
DisabledTools: options.DisabledTools,
|
||||
Progress: progressCallback,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1712,20 +1724,21 @@ func unsandboxedRetryArgs(args map[string]any) map[string]any {
|
||||
|
||||
func runToolForUnsandboxedRetry(ctx context.Context, registry *tools.Registry, name string, toolCallID string, args map[string]any, permissionMode PermissionMode, options Options, progressCallback func(streamjson.Event)) tools.Result {
|
||||
return registry.RunWithOptions(ctx, name, args, tools.RunOptions{
|
||||
PermissionGranted: true,
|
||||
PermissionMode: string(permissionMode),
|
||||
Autonomy: options.Autonomy,
|
||||
Sandbox: options.Sandbox,
|
||||
ToolCallID: toolCallID,
|
||||
SessionID: options.SessionID,
|
||||
Model: options.Model,
|
||||
ReasoningEffort: options.ReasoningEffort,
|
||||
Depth: options.Depth,
|
||||
Cwd: options.Cwd,
|
||||
FileTracker: options.FileTracker,
|
||||
EnabledTools: options.EnabledTools,
|
||||
DisabledTools: options.DisabledTools,
|
||||
Progress: progressCallback,
|
||||
PermissionGranted: true,
|
||||
PermissionMode: string(permissionMode),
|
||||
Autonomy: options.Autonomy,
|
||||
Sandbox: options.Sandbox,
|
||||
ToolCallID: toolCallID,
|
||||
SessionID: options.SessionID,
|
||||
Model: options.Model,
|
||||
ReasoningEffort: options.ReasoningEffort,
|
||||
Depth: options.Depth,
|
||||
Cwd: options.Cwd,
|
||||
FileTracker: options.FileTracker,
|
||||
DeferFileObservationCommit: true,
|
||||
EnabledTools: options.EnabledTools,
|
||||
DisabledTools: options.DisabledTools,
|
||||
Progress: progressCallback,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -2742,6 +2755,45 @@ func shellCommandAdditionalPermissionsRequested(args map[string]any) bool {
|
||||
return strings.TrimSpace(value) == string(tools.SandboxPermissionsWithAdditionalPermissions)
|
||||
}
|
||||
|
||||
// normalizeNoopAdditionalPermissions treats a provider-expanded empty or
|
||||
// already-covered additional_permissions object as if it had been omitted.
|
||||
// Some strict-schema providers materialize every optional property and may add
|
||||
// a workspace read that the default sandbox already grants; rejecting either
|
||||
// no-op makes the model retry an otherwise valid command. Requests that broaden
|
||||
// access and malformed payloads are left untouched so validation stays closed.
|
||||
func normalizeNoopAdditionalPermissions(args map[string]any, basePath string, engine *sandbox.Engine) {
|
||||
raw, exists := args["additional_permissions"]
|
||||
if !exists || raw == nil {
|
||||
return
|
||||
}
|
||||
data, err := json.Marshal(raw)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
var profile sandbox.RequestPermissionProfile
|
||||
if err := json.Unmarshal(data, &profile); err != nil {
|
||||
return
|
||||
}
|
||||
normalized, err := sandbox.NormalizeRequestPermissionProfile(profile, basePath)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
noop := normalized.Empty()
|
||||
if !noop && engine != nil {
|
||||
if grantProfile, grantErr := sandbox.RequestPermissionGrantProfile(normalized); grantErr == nil {
|
||||
noop = engine.CoversRequestPermissions(grantProfile)
|
||||
}
|
||||
}
|
||||
if !noop {
|
||||
return
|
||||
}
|
||||
delete(args, "additional_permissions")
|
||||
if rawMode, ok := args["sandbox_permissions"].(string); ok &&
|
||||
strings.TrimSpace(rawMode) == string(tools.SandboxPermissionsWithAdditionalPermissions) {
|
||||
args["sandbox_permissions"] = string(tools.SandboxPermissionsUseDefault)
|
||||
}
|
||||
}
|
||||
|
||||
func inlineAdditionalPermissionsProfile(args map[string]any, basePath string) (sandbox.RequestPermissionProfile, bool, error) {
|
||||
if !shellCommandAdditionalPermissionsRequested(args) {
|
||||
if _, exists := args["additional_permissions"]; exists {
|
||||
@@ -2932,22 +2984,20 @@ func partitionToolsCached(registry *tools.Registry, permissionMode PermissionMod
|
||||
// 2. tool_search — always present, right after the stable block;
|
||||
// 3. loaded deferred tools, alpha-sorted — APPENDED, so an unloaded->loaded
|
||||
// transition grows the tail instead of inserting into the eager block.
|
||||
// (tool_search's discovery text still shrinks as tools load, so keeping it and
|
||||
// the loaded tail AFTER the eager block preserves the eager tools' cache across a
|
||||
// load; fully stabilizing the loader's own description is a scoped follow-up.)
|
||||
// tool_search's compact catalog also stays byte-stable as tools load, so only
|
||||
// the appended full-schema tail changes.
|
||||
eager := make([]zeroruntime.ToolDefinition, 0, len(visible))
|
||||
loadedTail := make([]zeroruntime.ToolDefinition, 0)
|
||||
var hiddenTools []tools.Tool
|
||||
var deferredCatalog []tools.Tool
|
||||
for _, tool := range visible {
|
||||
name := tool.Name()
|
||||
if name == tools.ToolSearchToolName {
|
||||
continue // added explicitly between the eager block and the loaded tail
|
||||
}
|
||||
if tools.IsDeferred(tool) {
|
||||
deferredCatalog = append(deferredCatalog, tool)
|
||||
if loaded[name] {
|
||||
loadedTail = append(loadedTail, cachedRuntimeToolDefinition(defCache, tool))
|
||||
} else {
|
||||
hiddenTools = append(hiddenTools, tool)
|
||||
}
|
||||
continue
|
||||
}
|
||||
@@ -2960,15 +3010,16 @@ func partitionToolsCached(registry *tools.Registry, permissionMode PermissionMod
|
||||
return loadedTail[left].Name < loadedTail[right].Name
|
||||
})
|
||||
|
||||
discovery := ""
|
||||
if len(hiddenTools) > 0 {
|
||||
discovery = tools.BuildToolSearchDescription(hiddenTools)
|
||||
}
|
||||
// Keep the catalog byte-stable as tools load. tool_search already redirects
|
||||
// requests for an exposed tool back to the current tool list, so retaining
|
||||
// loaded names costs only compact catalog lines and avoids invalidating the
|
||||
// provider cache at the loader definition on every load.
|
||||
discovery := tools.BuildToolSearchDescription(deferredCatalog)
|
||||
|
||||
// tool_search is guaranteed runnable on the active path, so it is ALWAYS exposed
|
||||
// — even when a non-empty EnabledTools allowlist omits it (the operator
|
||||
// allowlisted the deferred tools, not the loader). It sits right after the stable
|
||||
// eager block and carries the discovery list for still-hidden tools.
|
||||
// eager block and carries the stable deferred-tool catalog.
|
||||
description := loader.Description()
|
||||
if discovery != "" {
|
||||
description = discovery
|
||||
|
||||
+67
-17
@@ -64,6 +64,22 @@ func TestTypedExecutionOutcomeOverridesLegacySandboxHeuristics(t *testing.T) {
|
||||
t.Fatal("narrow protected-metadata denial must not become an unrestricted retry")
|
||||
}
|
||||
|
||||
nonRecoverableDenial := protectedDenial
|
||||
nonRecoverableDenial.ExecutionOutcome = &execution.Outcome{
|
||||
State: execution.StateDenied,
|
||||
Kind: execution.OutcomeEnforcementDenied,
|
||||
Denial: &execution.Denial{
|
||||
Capability: execution.Capability{Kind: execution.CapabilityUnrestricted, Scope: "host"},
|
||||
Source: execution.DenialSourcePlatformSandbox,
|
||||
Reason: "sandbox enforcement failed closed",
|
||||
Recoverable: false,
|
||||
NextAction: execution.DenialNextActionRequestApproval,
|
||||
},
|
||||
}
|
||||
if sandboxRestrictedShellRetryCandidate(call, nil, nonRecoverableDenial, options) {
|
||||
t.Fatal("non-recoverable sandbox denial must not become an unrestricted retry")
|
||||
}
|
||||
|
||||
networkDenial := protectedDenial
|
||||
networkDenial.ExecutionOutcome = &execution.Outcome{
|
||||
State: execution.StateDenied,
|
||||
@@ -288,6 +304,18 @@ func (tool *sandboxDeniedRetryTool) Run(_ context.Context, args map[string]any)
|
||||
return tools.Result{
|
||||
Status: tools.StatusError,
|
||||
Output: "touch: cannot touch '/home/user/.npm/cache': Read-only file system",
|
||||
ExecutionOutcome: &execution.Outcome{
|
||||
State: execution.StateDenied,
|
||||
Kind: execution.OutcomeEnforcementDenied,
|
||||
Exit: &execution.Exit{Code: 1},
|
||||
Denial: &execution.Denial{
|
||||
Capability: execution.Capability{Kind: execution.CapabilityUnrestricted, Scope: "host"},
|
||||
Source: execution.DenialSourcePlatformSandbox,
|
||||
Reason: "sandbox blocked command execution",
|
||||
Recoverable: true,
|
||||
NextAction: execution.DenialNextActionRequestApproval,
|
||||
},
|
||||
},
|
||||
Meta: map[string]string{
|
||||
"exit_code": "1",
|
||||
tools.SandboxLikelyDeniedMeta: "true",
|
||||
@@ -304,7 +332,8 @@ type sandboxDeniedExecCommandRetryTool struct {
|
||||
func (tool *sandboxDeniedExecCommandRetryTool) Name() string { return "exec_command" }
|
||||
|
||||
type sandboxNetworkDeniedRetryTool struct {
|
||||
calls []map[string]any
|
||||
calls []map[string]any
|
||||
deferFileObservationCommit []bool
|
||||
}
|
||||
|
||||
func (tool *sandboxNetworkDeniedRetryTool) Name() string { return "bash" }
|
||||
@@ -327,6 +356,7 @@ func (tool *sandboxNetworkDeniedRetryTool) Run(ctx context.Context, args map[str
|
||||
}
|
||||
func (tool *sandboxNetworkDeniedRetryTool) RunWithOptions(ctx context.Context, args map[string]any, options tools.RunOptions) tools.Result {
|
||||
tool.calls = append(tool.calls, cloneArgs(args))
|
||||
tool.deferFileObservationCommit = append(tool.deferFileObservationCommit, options.DeferFileObservationCommit)
|
||||
if shellNetworkAllowed(ctx, options.Sandbox, args) {
|
||||
return tools.Result{Status: tools.StatusOK, Output: "server started"}
|
||||
}
|
||||
@@ -369,7 +399,8 @@ func agentNativeBackendStub() sandbox.Backend {
|
||||
}
|
||||
|
||||
type sandboxNamespaceLimitedRetryTool struct {
|
||||
calls []map[string]any
|
||||
calls []map[string]any
|
||||
deferFileObservationCommit []bool
|
||||
}
|
||||
|
||||
func (tool *sandboxNamespaceLimitedRetryTool) Name() string { return "bash" }
|
||||
@@ -388,8 +419,12 @@ func (tool *sandboxNamespaceLimitedRetryTool) Parameters() tools.Schema {
|
||||
func (tool *sandboxNamespaceLimitedRetryTool) Safety() tools.Safety {
|
||||
return tools.Safety{SideEffect: tools.SideEffectShell, Permission: tools.PermissionPrompt, Reason: "runs shell commands"}
|
||||
}
|
||||
func (tool *sandboxNamespaceLimitedRetryTool) Run(_ context.Context, args map[string]any) tools.Result {
|
||||
func (tool *sandboxNamespaceLimitedRetryTool) Run(ctx context.Context, args map[string]any) tools.Result {
|
||||
return tool.RunWithOptions(ctx, args, tools.RunOptions{})
|
||||
}
|
||||
func (tool *sandboxNamespaceLimitedRetryTool) RunWithOptions(_ context.Context, args map[string]any, options tools.RunOptions) tools.Result {
|
||||
tool.calls = append(tool.calls, cloneArgs(args))
|
||||
tool.deferFileObservationCommit = append(tool.deferFileObservationCommit, options.DeferFileObservationCommit)
|
||||
if args["sandbox_permissions"] == string(tools.SandboxPermissionsRequireEscalated) {
|
||||
return tools.Result{Status: tools.StatusOK, Output: "stdout:\nUSER PID COMMAND\nanaxy 42 firefox\nanaxy 43 Discord"}
|
||||
}
|
||||
@@ -455,6 +490,9 @@ func TestRunRetriesShellUnsandboxedAfterSandboxNamespaceLimitedOutput(t *testing
|
||||
if retryTool.calls[1]["sandbox_permissions"] != string(tools.SandboxPermissionsRequireEscalated) {
|
||||
t.Fatalf("retry args = %#v, want require_escalated", retryTool.calls[1])
|
||||
}
|
||||
if len(retryTool.deferFileObservationCommit) != 2 || !retryTool.deferFileObservationCommit[1] {
|
||||
t.Fatalf("retry defer flags = %#v, want retry observation deferred", retryTool.deferFileObservationCommit)
|
||||
}
|
||||
if len(requests) != 1 {
|
||||
t.Fatalf("permission requests = %#v, want one unsandboxed retry approval", requests)
|
||||
}
|
||||
@@ -606,6 +644,9 @@ func TestRunRetriesNetworkDeniedShellWithNetworkGrant(t *testing.T) {
|
||||
if _, escalated := retryTool.calls[1]["sandbox_permissions"]; escalated {
|
||||
t.Fatalf("network retry must stay sandboxed, got retry args %#v", retryTool.calls[1])
|
||||
}
|
||||
if len(retryTool.deferFileObservationCommit) != 2 || !retryTool.deferFileObservationCommit[1] {
|
||||
t.Fatalf("retry defer flags = %#v, want retry observation deferred", retryTool.deferFileObservationCommit)
|
||||
}
|
||||
if len(requests) != 1 || requests[0].Reason != sandbox.ReasonNetworkBlocked {
|
||||
t.Fatalf("permission requests = %#v, want one network approval", requests)
|
||||
}
|
||||
@@ -2404,7 +2445,11 @@ func TestRunApprovedNetworkBashPromptAppliesTurnNetworkGrant(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
command := "PATH=.:$PATH curl https://example.com"
|
||||
if runtime.GOOS == "windows" {
|
||||
command = "set PATH=.;%PATH% && curl https://example.com"
|
||||
if windowsTestUsesPowerShell() {
|
||||
command = `$env:PATH = '.;' + $env:PATH; curl.cmd https://example.com`
|
||||
} else {
|
||||
command = "set PATH=.;%PATH% && curl https://example.com"
|
||||
}
|
||||
fakeCurl := filepath.Join(root, "curl.cmd")
|
||||
if err := os.WriteFile(fakeCurl, []byte("@echo fake curl %*\r\n"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
@@ -2525,13 +2570,17 @@ func TestRunDoesNotOfferPrefixApprovalForUnsafeBashCommand(t *testing.T) {
|
||||
|
||||
func TestRunPromptsForDestructiveShellInsteadOfSandboxDeny(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
command := "echo rm -rf /"
|
||||
if runtime.GOOS == "windows" && windowsTestUsesPowerShell() {
|
||||
command = `Write-Output 'rm -rf /'`
|
||||
}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.NewScopedBashTool(root, nil))
|
||||
provider := &mockProvider{
|
||||
turns: [][]zeroruntime.StreamEvent{
|
||||
{
|
||||
{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "call-1", ToolName: "bash"},
|
||||
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call-1", ArgumentsFragment: `{"command":"echo rm -rf /"}`},
|
||||
{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "call-1", ArgumentsFragment: `{"command":` + quoteJSONString(command) + `}`},
|
||||
{Type: zeroruntime.StreamEventToolCallEnd, ToolCallID: "call-1"},
|
||||
{Type: zeroruntime.StreamEventDone},
|
||||
},
|
||||
@@ -2627,6 +2676,12 @@ func TestRunAlwaysAllowWithoutSandboxStillAllowsCall(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func windowsTestUsesPowerShell() bool {
|
||||
executable, _ := tools.HostShellCommand("")
|
||||
name := strings.ToLower(filepath.Base(executable))
|
||||
return strings.Contains(name, "powershell") || strings.HasPrefix(name, "pwsh")
|
||||
}
|
||||
|
||||
func containsPermissionDecision(decisions []PermissionDecisionAction, want PermissionDecisionAction) bool {
|
||||
for _, decision := range decisions {
|
||||
if decision == want {
|
||||
@@ -3160,22 +3215,17 @@ func TestBuildSystemPromptInjectsHostShellContext(t *testing.T) {
|
||||
t.Fatalf("expected operating system in environment block, got %q", prompt)
|
||||
}
|
||||
if runtime.GOOS == "windows" {
|
||||
for _, want := range []string{"Windows cmd.exe syntax", "cwd argument", "MSYS binaries", "grep", "require_escalated", "use double quotes around the value"} {
|
||||
if !strings.Contains(prompt, want) {
|
||||
t.Fatalf("expected Windows shell guidance to mention %q, got %q", want, prompt)
|
||||
}
|
||||
wants := []string{"workdir/cwd", "MSYS binaries"}
|
||||
if windowsTestUsesPowerShell() {
|
||||
wants = append(wants, "PowerShell", "Get-ChildItem", "Select-String", "$env:NAME", "require_escalated")
|
||||
} else {
|
||||
wants = append(wants, "cmd.exe", "double quotes")
|
||||
}
|
||||
// The examples must themselves use the safe (double-quoted) form; a
|
||||
// single-quoted example here would teach the model the exact syntax
|
||||
// that fails under cmd.exe.
|
||||
for _, want := range []string{`--jq ".a | b"`, `-run "A|B"`} {
|
||||
for _, want := range wants {
|
||||
if !strings.Contains(prompt, want) {
|
||||
t.Fatalf("expected double-quoted example %q in Windows shell guidance, got %q", want, prompt)
|
||||
t.Fatalf("expected selected Windows shell guidance to mention %q, got %q", want, prompt)
|
||||
}
|
||||
}
|
||||
if strings.Contains(prompt, `'.a | b'`) || strings.Contains(prompt, `'A|B'`) {
|
||||
t.Fatalf("Windows shell guidance must not show the unsafe single-quoted form, got %q", prompt)
|
||||
}
|
||||
} else if !strings.Contains(prompt, "/bin/sh syntax") {
|
||||
t.Fatalf("expected POSIX shell guidance in prompt, got %q", prompt)
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ package agent
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
@@ -19,6 +20,16 @@ type propagationOutputTool struct {
|
||||
|
||||
func TestOutputBudgetHookHelperProcess(t *testing.T) {
|
||||
for index, arg := range os.Args {
|
||||
if arg == "--zero-output-budget-hook-repeat" && index+1 < len(os.Args) {
|
||||
repeats, err := strconv.Atoi(os.Args[index+1])
|
||||
if err != nil {
|
||||
os.Exit(2)
|
||||
}
|
||||
if _, err := os.Stdout.WriteString(strings.Repeat("hook feedback ", repeats)); err != nil {
|
||||
os.Exit(2)
|
||||
}
|
||||
os.Exit(0)
|
||||
}
|
||||
if arg != "--zero-output-budget-hook" || index+1 >= len(os.Args) {
|
||||
continue
|
||||
}
|
||||
@@ -30,25 +41,102 @@ func TestOutputBudgetHookHelperProcess(t *testing.T) {
|
||||
}
|
||||
|
||||
func largeOutputBudgetHookDispatcher() *hooks.Dispatcher {
|
||||
feedback := strings.Repeat("hook feedback ", 200)
|
||||
return outputBudgetHookDispatcherFor("propagation_output", 200)
|
||||
}
|
||||
|
||||
func outputBudgetHookDispatcherFor(toolName string, repeats int) *hooks.Dispatcher {
|
||||
feedback := strings.Repeat("hook feedback ", repeats)
|
||||
hookArgs := []string{
|
||||
"-test.run=TestOutputBudgetHookHelperProcess",
|
||||
"--",
|
||||
"--zero-output-budget-hook",
|
||||
feedback,
|
||||
}
|
||||
if repeats > 1_000 {
|
||||
hookArgs = []string{
|
||||
"-test.run=TestOutputBudgetHookHelperProcess",
|
||||
"--",
|
||||
"--zero-output-budget-hook-repeat",
|
||||
strconv.Itoa(repeats),
|
||||
}
|
||||
}
|
||||
return hooks.NewDispatcher(hooks.DispatcherOptions{Config: hooks.Config{
|
||||
Enabled: true,
|
||||
Hooks: []hooks.Definition{{
|
||||
ID: "large-feedback",
|
||||
Event: hooks.EventAfterTool,
|
||||
Matcher: "propagation_output",
|
||||
Matcher: toolName,
|
||||
Command: os.Args[0],
|
||||
Args: []string{
|
||||
"-test.run=TestOutputBudgetHookHelperProcess",
|
||||
"--",
|
||||
"--zero-output-budget-hook",
|
||||
feedback,
|
||||
},
|
||||
Args: hookArgs,
|
||||
Enabled: true,
|
||||
}},
|
||||
}})
|
||||
}
|
||||
|
||||
func TestExecuteToolCallCommitsReadObservationAfterFinalBoundary(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
targetPath := filepath.Join(root, "target.txt")
|
||||
if err := os.WriteFile(targetPath, []byte("known content\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
linkPath := filepath.Join(root, "small.txt")
|
||||
if err := os.Symlink(targetPath, linkPath); err != nil {
|
||||
t.Skipf("symlink unavailable: %v", err)
|
||||
}
|
||||
resolvedPath, err := filepath.EvalSymlinks(linkPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.NewReadFileTool(root))
|
||||
tracker := tools.NewFileTracker()
|
||||
|
||||
result, abortErr := executeToolCall(context.Background(), registry, ToolCall{
|
||||
ID: "read-small", Name: "read_file", Arguments: `{"path":"small.txt"}`,
|
||||
}, PermissionModeAuto, Options{Cwd: root, FileTracker: tracker})
|
||||
if abortErr != nil || result.Status != tools.StatusOK {
|
||||
t.Fatalf("executeToolCall: result=%#v err=%v", result, abortErr)
|
||||
}
|
||||
if !tracker.SeenWhole(resolvedPath) {
|
||||
t.Fatal("unmodified final read output was not committed as observed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteToolCallDoesNotCommitReadObservationBeforeHookBudget(t *testing.T) {
|
||||
t.Setenv("ZERO_TOOL_OUTPUT_CEILING_TOKENS", "80")
|
||||
root := t.TempDir()
|
||||
targetPath := filepath.Join(root, "target.txt")
|
||||
if err := os.WriteFile(targetPath, []byte("known content\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
linkPath := filepath.Join(root, "small.txt")
|
||||
if err := os.Symlink(targetPath, linkPath); err != nil {
|
||||
t.Skipf("symlink unavailable: %v", err)
|
||||
}
|
||||
resolvedPath, err := filepath.EvalSymlinks(linkPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
registry := tools.NewRegistry()
|
||||
registry.Register(tools.NewReadFileTool(root))
|
||||
tracker := tools.NewFileTracker()
|
||||
|
||||
result, abortErr := executeToolCall(context.Background(), registry, ToolCall{
|
||||
ID: "read-hook", Name: "read_file", Arguments: `{"path":"small.txt"}`,
|
||||
}, PermissionModeAuto, Options{
|
||||
Cwd: root, FileTracker: tracker, Hooks: outputBudgetHookDispatcherFor("read_file", 20_000),
|
||||
})
|
||||
if abortErr != nil {
|
||||
t.Fatalf("executeToolCall: %v", abortErr)
|
||||
}
|
||||
if !result.Truncated {
|
||||
t.Fatalf("precondition: hook feedback was not truncated: %#v", result)
|
||||
}
|
||||
if tracker.SeenWhole(resolvedPath) {
|
||||
t.Fatal("read observation committed before hook output was budgeted")
|
||||
}
|
||||
}
|
||||
|
||||
func (tool propagationOutputTool) Name() string { return "propagation_output" }
|
||||
func (tool propagationOutputTool) Description() string { return "returns output for propagation tests" }
|
||||
func (tool propagationOutputTool) Parameters() tools.Schema {
|
||||
|
||||
@@ -121,3 +121,67 @@ func toolNameForResult(messages []zeroruntime.Message, index int) string {
|
||||
func isPrunedPlaceholder(content string) bool {
|
||||
return strings.HasPrefix(strings.TrimSpace(content), "[pruned ")
|
||||
}
|
||||
|
||||
var supersedableReadTools = map[string]bool{
|
||||
"read_file": true,
|
||||
"read_minified_file": true,
|
||||
"grep": true,
|
||||
"glob": true,
|
||||
"list_directory": true,
|
||||
}
|
||||
|
||||
// pruneSupersededReadResults removes only older byte-identical results of the
|
||||
// same read call. It runs when context pressure already requires a rewrite, so
|
||||
// it does not trade a warm provider-cache prefix for a small apparent saving.
|
||||
func pruneSupersededReadResults(messages []zeroruntime.Message) ([]zeroruntime.Message, int) {
|
||||
type seenResultKey struct {
|
||||
call string
|
||||
content string
|
||||
}
|
||||
seen := map[seenResultKey]struct{}{}
|
||||
out := messages
|
||||
reclaimed := 0
|
||||
copied := false
|
||||
|
||||
for index := len(messages) - 1; index >= 0; index-- {
|
||||
message := messages[index]
|
||||
if message.Role != zeroruntime.MessageRoleTool || isPrunedPlaceholder(message.Content) {
|
||||
continue
|
||||
}
|
||||
call, ok := toolCallForResult(messages, index)
|
||||
if !ok || !supersedableReadTools[call.Name] {
|
||||
continue
|
||||
}
|
||||
key := seenResultKey{
|
||||
call: call.Name + "\x00" + strings.TrimSpace(call.Arguments),
|
||||
content: message.Content,
|
||||
}
|
||||
if _, exists := seen[key]; exists {
|
||||
placeholder := fmt.Sprintf("[superseded identical %s result — the newer result remains in context]", call.Name)
|
||||
if !copied {
|
||||
out = append([]zeroruntime.Message(nil), messages...)
|
||||
copied = true
|
||||
}
|
||||
out[index].Content = placeholder
|
||||
reclaimed += max(0, ApproxTextTokens(message.Content)-ApproxTextTokens(placeholder))
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
}
|
||||
return out, reclaimed
|
||||
}
|
||||
|
||||
func toolCallForResult(messages []zeroruntime.Message, resultIndex int) (zeroruntime.ToolCall, bool) {
|
||||
id := messages[resultIndex].ToolCallID
|
||||
if id == "" {
|
||||
return zeroruntime.ToolCall{}, false
|
||||
}
|
||||
for index := resultIndex - 1; index >= 0; index-- {
|
||||
for _, call := range messages[index].ToolCalls {
|
||||
if call.ID == id {
|
||||
return call, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return zeroruntime.ToolCall{}, false
|
||||
}
|
||||
@@ -25,6 +25,13 @@ func toolCallMsg(id, name string) zeroruntime.Message {
|
||||
}
|
||||
}
|
||||
|
||||
func toolCallWithArgs(id, name, args string) zeroruntime.Message {
|
||||
return zeroruntime.Message{
|
||||
Role: zeroruntime.MessageRoleAssistant,
|
||||
ToolCalls: []zeroruntime.ToolCall{{ID: id, Name: name, Arguments: args}},
|
||||
}
|
||||
}
|
||||
|
||||
func TestPruneSkipsSmallSessions(t *testing.T) {
|
||||
// Total reclaimable is under the gate → no pruning.
|
||||
msgs := []zeroruntime.Message{
|
||||
@@ -120,3 +127,69 @@ func TestPruneDoesNotMutateInput(t *testing.T) {
|
||||
t.Fatal("pruneStaleToolOutput must not mutate the caller's slice")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPruneSupersededReadResultsOnlyDropsIdenticalOlderReads(t *testing.T) {
|
||||
body := strings.Repeat("same source line\n", 1000)
|
||||
messages := []zeroruntime.Message{
|
||||
toolCallWithArgs("old", "read_file", `{"path":"a.go"}`),
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "old", Content: body},
|
||||
toolCallWithArgs("different", "read_file", `{"path":"b.go"}`),
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "different", Content: body},
|
||||
toolCallWithArgs("new", "read_file", `{"path":"a.go"}`),
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "new", Content: body},
|
||||
}
|
||||
out, reclaimed := pruneSupersededReadResults(messages)
|
||||
if reclaimed <= 0 || !strings.Contains(out[1].Content, "superseded identical") {
|
||||
t.Fatalf("older duplicate was not pruned: reclaimed=%d output=%q", reclaimed, out[1].Content)
|
||||
}
|
||||
if out[3].Content != body || out[5].Content != body {
|
||||
t.Fatal("different call and newest result must remain exact")
|
||||
}
|
||||
if messages[1].Content != body {
|
||||
t.Fatal("input slice was mutated")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPruneSupersededReadResultsKeepsChangedAndMutatingResults(t *testing.T) {
|
||||
messages := []zeroruntime.Message{
|
||||
toolCallWithArgs("r1", "read_file", `{"path":"a.go"}`),
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "r1", Content: "version one"},
|
||||
toolCallWithArgs("r2", "read_file", `{"path":"a.go"}`),
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "r2", Content: "version two"},
|
||||
toolCallWithArgs("e1", "edit_file", `{"path":"a.go"}`),
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "e1", Content: "edited"},
|
||||
toolCallWithArgs("e2", "edit_file", `{"path":"a.go"}`),
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "e2", Content: "edited"},
|
||||
}
|
||||
out, reclaimed := pruneSupersededReadResults(messages)
|
||||
if reclaimed != 0 {
|
||||
t.Fatalf("unsafe results were pruned: %d", reclaimed)
|
||||
}
|
||||
for index := range messages {
|
||||
if out[index].Content != messages[index].Content {
|
||||
t.Fatalf("message %d changed unexpectedly", index)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPruneSupersededReadResultsFindsMatchingBodyAcrossChangedRead(t *testing.T) {
|
||||
bodyA := strings.Repeat("version a\n", 100)
|
||||
bodyB := strings.Repeat("version b\n", 100)
|
||||
messages := []zeroruntime.Message{
|
||||
toolCallWithArgs("old-a", "read_file", `{"path":"a.go"}`),
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "old-a", Content: bodyA},
|
||||
toolCallWithArgs("middle-b", "read_file", `{"path":"a.go"}`),
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "middle-b", Content: bodyB},
|
||||
toolCallWithArgs("new-a", "read_file", `{"path":"a.go"}`),
|
||||
{Role: zeroruntime.MessageRoleTool, ToolCallID: "new-a", Content: bodyA},
|
||||
}
|
||||
|
||||
out, reclaimed := pruneSupersededReadResults(messages)
|
||||
|
||||
if reclaimed <= 0 || !strings.Contains(out[1].Content, "superseded identical") {
|
||||
t.Fatalf("older matching body was not pruned across a changed result: reclaimed=%d output=%q", reclaimed, out[1].Content)
|
||||
}
|
||||
if out[3].Content != bodyB || out[5].Content != bodyA {
|
||||
t.Fatal("changed middle result and newest matching result must remain exact")
|
||||
}
|
||||
}
|
||||
@@ -26,6 +26,16 @@ func TestSpecialistDelegationContext(t *testing.T) {
|
||||
if strings.Contains(got, "nameless") {
|
||||
t.Fatalf("nameless entry should be skipped: %q", got)
|
||||
}
|
||||
for _, want := range []string{"reduce total context", "Do not delegate small direct work", "bounded assignment"} {
|
||||
if !strings.Contains(got, want) {
|
||||
t.Fatalf("missing evidence-driven delegation guidance %q: %q", want, got)
|
||||
}
|
||||
}
|
||||
for _, unwanted := range []string{"delegate to it proactively", "launch several specialists in parallel"} {
|
||||
if strings.Contains(got, unwanted) {
|
||||
t.Fatalf("retained broad delegation guidance %q: %q", unwanted, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSystemPromptIncludesDelegationOnlyWithSpecialists(t *testing.T) {
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/repomap"
|
||||
"github.com/Gitlawb/zero/internal/tools"
|
||||
"github.com/Gitlawb/zero/internal/workspaceseed"
|
||||
)
|
||||
|
||||
@@ -178,21 +179,19 @@ func approvedCommandPrefixContext(options Options) string {
|
||||
return "## Approved Command Prefixes\n\nThe following command prefixes have already been approved and do not need another permission prompt:\n" + strings.Join(lines, "\n")
|
||||
}
|
||||
|
||||
// specialistDelegationContext nudges the orchestrator to offload read-heavy or
|
||||
// parallelizable work to a specialist sub-agent via the Task tool, keeping large
|
||||
// tool outputs out of the main context. It renders only when specialists are
|
||||
// known (which is only where the Task tool is actually registered), so a run with
|
||||
// no delegatable specialists produces the previous prompt unchanged.
|
||||
// specialistDelegationContext describes when a specialist earns its extra
|
||||
// context and coordination cost. It renders only when specialists are known
|
||||
// (which is only where the Task tool is actually registered), so a run with no
|
||||
// delegatable specialists produces the previous prompt unchanged.
|
||||
func specialistDelegationContext(options Options) string {
|
||||
if len(options.Specialists) == 0 {
|
||||
return ""
|
||||
}
|
||||
var b strings.Builder
|
||||
b.WriteString("<specialists>\n")
|
||||
b.WriteString("Delegate focused or read-heavy work to a specialist sub-agent with the Task tool instead of doing it inline. ")
|
||||
b.WriteString("When a request matches a specialist's purpose, delegate to it proactively — you do not need the user to ask first. ")
|
||||
b.WriteString("This keeps large tool outputs — searches, file dumps, multi-step exploration — out of your own context, so you stay fast and token-efficient. ")
|
||||
b.WriteString("Prefer delegating codebase search and exploration; for independent subtasks, launch several specialists in parallel. Handle small, direct edits yourself.\n")
|
||||
b.WriteString("Use the Task tool when a specialist's focused context, expertise, or an independent parallel subtask is likely to improve quality or reduce total context. ")
|
||||
b.WriteString("Do not delegate small direct work, duplicate exploration, or split sequential steps merely because a specialist is available. ")
|
||||
b.WriteString("When delegating, give a bounded assignment and consume its concise evidence instead of importing a long transcript.\n")
|
||||
b.WriteString("Available specialists (call Task with the matching name when the task fits its purpose):\n")
|
||||
for _, info := range options.Specialists {
|
||||
name := strings.TrimSpace(info.Name)
|
||||
@@ -310,11 +309,7 @@ func workspaceContext(cwd string) string {
|
||||
b.WriteString("<environment>\n")
|
||||
b.WriteString("Working directory: " + cwd + "\n")
|
||||
b.WriteString("Operating system: " + runtime.GOOS + "\n")
|
||||
if runtime.GOOS == "windows" {
|
||||
b.WriteString("Shell syntax: Windows cmd.exe syntax for exec_command/bash tools. To put | & > < etc inside an arg value, use double quotes around the value, not single quotes (single quotes do not protect metachars in cmd.exe): gh --jq \".a | b\", go test -run \"A|B\". Do not pipe to or invoke POSIX coreutils from Git for Windows (usr\\bin head/grep/tail/cat/...): they are MSYS binaries and fail under the write-restricted sandbox; use native Zero tools (grep, read_file, list_directory, glob) or cmd.exe findstr/more instead, or sandbox_permissions require_escalated only when host-level execution is truly required. Prefer the workdir/cwd argument over cd when changing directories.\n")
|
||||
} else {
|
||||
b.WriteString("Shell syntax: /bin/sh syntax for exec_command/bash tools; prefer the workdir/cwd argument instead of cd when changing directories.\n")
|
||||
}
|
||||
b.WriteString(tools.HostShellEnvironmentGuidance() + "\n")
|
||||
if branch := gitBranchForPrompt(cwd); branch != "" {
|
||||
b.WriteString("Git branch: " + branch + "\n")
|
||||
}
|
||||
|
||||
@@ -31,15 +31,10 @@ work.
|
||||
tests, or config before you modify behavior. Never edit a file you have not
|
||||
read.
|
||||
2. **Plan.** For multi-step work, call update_plan with an ordered checklist and
|
||||
keep it live. The plan bar is the user's progress signal — call update_plan
|
||||
after EACH concrete unit of work (every file written, every command run), not
|
||||
just at coarse milestones: mark the finished item completed and the next one
|
||||
in_progress before you start it. A plan stuck at 0/N while files are landing is
|
||||
a bug, not economy — these calls are cheap, expected, and the update_plan cards
|
||||
are hidden from the transcript, so frequent updates cost the user nothing and
|
||||
never clutter the conversation. Keep at most one item in_progress, and never
|
||||
batch the updates to the end of the turn. Skip the plan for trivial one-step
|
||||
tasks.
|
||||
keep it current at meaningful milestones. Mark completed work and the next
|
||||
in-progress step together when practical; do not spend separate model turns
|
||||
updating the plan after every file or command. Keep at most one item
|
||||
in_progress. Skip the plan for trivial one-step tasks.
|
||||
3. **Implement.** Make focused changes that match the surrounding code's style,
|
||||
naming, and conventions. Prefer the smallest change that fully solves the
|
||||
problem. Avoid broad refactors, unrelated rewrites, dependency churn, and
|
||||
@@ -57,8 +52,9 @@ work.
|
||||
file tools - read_file, list_directory, glob, grep, write_file, edit_file,
|
||||
apply_patch - over shelling out to cat/sed/awk/python for file operations.
|
||||
They are safer, reviewable, and produce clean diffs.
|
||||
- Make one tool call per file. Do not batch multi-file writes into a single
|
||||
shell or script invocation.
|
||||
- Keep edits focused and reviewable. A single patch may update several related
|
||||
files when they form one coherent change; do not hide unrelated edits in a
|
||||
bulk shell or script rewrite.
|
||||
- For edits to existing files, prefer edit_file or apply_patch with minimal,
|
||||
targeted diffs. Match the existing indentation, imports, and idioms. Match the
|
||||
file's comment density: do not add explanatory comments unless the user asks or
|
||||
|
||||
+7
-3
@@ -27,6 +27,8 @@ Usage:
|
||||
Not meant to be run interactively — point your editor's ACP / external-agent
|
||||
setting at "zero acp".`
|
||||
|
||||
var acpSignalContext = signalContext
|
||||
|
||||
// runACP serves ACP over stdio so an editor can drive ZERO's agent core. It
|
||||
// speaks JSON-RPC 2.0 (newline-delimited JSON) on stdin/stdout; stderr stays free
|
||||
// for human-readable diagnostics. The session lifecycle maps onto ZERO's own
|
||||
@@ -42,7 +44,9 @@ func runACP(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int
|
||||
}
|
||||
}
|
||||
|
||||
conn := acp.NewConn(deps.stdin, stdout)
|
||||
// The process's own stdin belongs exclusively to this Conn for the life of
|
||||
// the command, so cancellation may close it to interrupt an idle read.
|
||||
conn := acp.NewOwnedConn(deps.stdin, stdout)
|
||||
acp.NewAgent(conn, acp.Deps{
|
||||
ResolveConfig: deps.resolveConfig,
|
||||
DiscoverModels: func(ctx context.Context, profile config.ProviderProfile) ([]providermodeldiscovery.Model, error) {
|
||||
@@ -73,9 +77,9 @@ func runACP(args []string, stdout io.Writer, stderr io.Writer, deps appDeps) int
|
||||
AgentInfo: acp.Implementation{Name: "zero", Version: version},
|
||||
})
|
||||
|
||||
ctx, stop := signalContext()
|
||||
ctx, stop := acpSignalContext()
|
||||
defer stop()
|
||||
if err := conn.Serve(ctx); err != nil && ctx.Err() == nil {
|
||||
if err := conn.Serve(ctx); err != nil {
|
||||
return writeAppError(stderr, "acp: "+err.Error(), exitCrash)
|
||||
}
|
||||
return exitSuccess
|
||||
|
||||
@@ -0,0 +1,120 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type acpTestReader func([]byte) (int, error)
|
||||
|
||||
func (r acpTestReader) Read(p []byte) (int, error) { return r(p) }
|
||||
|
||||
type acpTestWriter func([]byte) (int, error)
|
||||
|
||||
func (w acpTestWriter) Write(p []byte) (int, error) { return w(p) }
|
||||
|
||||
func installACPTestContext(t *testing.T) context.CancelFunc {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
previous := acpSignalContext
|
||||
acpSignalContext = func() (context.Context, context.CancelFunc) {
|
||||
return ctx, cancel
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
acpSignalContext = previous
|
||||
})
|
||||
return cancel
|
||||
}
|
||||
|
||||
func TestRunACPCancellationPreservesTerminalReadError(t *testing.T) {
|
||||
cancel := installACPTestContext(t)
|
||||
wantErr := errors.New("transport read failed")
|
||||
reader := acpTestReader(func(p []byte) (int, error) {
|
||||
return copy(p, `{"jsonrpc":"1.0","id":1}`+"\n"), wantErr
|
||||
})
|
||||
writeStarted := make(chan struct{}, 1)
|
||||
releaseWrite := make(chan struct{})
|
||||
var stdout bytes.Buffer
|
||||
writer := acpTestWriter(func(p []byte) (int, error) {
|
||||
writeStarted <- struct{}{}
|
||||
<-releaseWrite
|
||||
return stdout.Write(p)
|
||||
})
|
||||
var stderr bytes.Buffer
|
||||
done := make(chan int, 1)
|
||||
go func() {
|
||||
done <- runACP(nil, writer, &stderr, fillAppDeps(appDeps{stdin: reader}))
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-writeStarted:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("ACP did not reach the response write")
|
||||
}
|
||||
// Keep Serve in the synchronous write until cancellation is observable. The
|
||||
// regression requires the genuine read error to win over that cancellation.
|
||||
cancel()
|
||||
close(releaseWrite)
|
||||
|
||||
select {
|
||||
case code := <-done:
|
||||
if code != exitCrash {
|
||||
t.Fatalf("exit code = %d, want crash %d", code, exitCrash)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("ACP did not exit after terminal read error")
|
||||
}
|
||||
if got := stderr.String(); !strings.Contains(got, "acp: "+wantErr.Error()) {
|
||||
t.Fatalf("stderr = %q, want terminal read error", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunACPIdleCancellationExitsCleanly(t *testing.T) {
|
||||
cancel := installACPTestContext(t)
|
||||
pipeReader, pipeWriter := io.Pipe()
|
||||
t.Cleanup(func() { _ = pipeWriter.Close() })
|
||||
readStarted := make(chan struct{}, 1)
|
||||
reader := &acpNotifyingReadCloser{ReadCloser: pipeReader, readStarted: readStarted}
|
||||
var stdout, stderr bytes.Buffer
|
||||
done := make(chan int, 1)
|
||||
go func() {
|
||||
done <- runACP(nil, &stdout, &stderr, fillAppDeps(appDeps{stdin: reader}))
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-readStarted:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("ACP did not begin reading")
|
||||
}
|
||||
cancel()
|
||||
|
||||
select {
|
||||
case code := <-done:
|
||||
if code != exitSuccess {
|
||||
t.Fatalf("exit code = %d, want success %d; stderr: %s", code, exitSuccess, stderr.String())
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("ACP did not exit after cancellation")
|
||||
}
|
||||
if stderr.Len() != 0 {
|
||||
t.Fatalf("stderr = %q, want empty", stderr.String())
|
||||
}
|
||||
}
|
||||
|
||||
type acpNotifyingReadCloser struct {
|
||||
io.ReadCloser
|
||||
readStarted chan<- struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (r *acpNotifyingReadCloser) Read(p []byte) (int, error) {
|
||||
r.once.Do(func() { r.readStarted <- struct{}{} })
|
||||
return r.ReadCloser.Read(p)
|
||||
}
|
||||
@@ -585,6 +585,9 @@ func fillAppDeps(deps appDeps) appDeps {
|
||||
baseProbeProviderHealth := deps.probeProviderHealth
|
||||
deps.probeProviderHealth = func(ctx context.Context, options providerhealth.Options) providerhealth.Result {
|
||||
options.Profile = applyStoredProviderKeyAt(options.Profile, userConfigPath)
|
||||
if options.OAuthResolver == nil {
|
||||
options.OAuthResolver, _ = oauthLoginForProfile(options.Profile)
|
||||
}
|
||||
return baseProbeProviderHealth(ctx, options)
|
||||
}
|
||||
return deps
|
||||
|
||||
@@ -40,10 +40,11 @@ func runContext(args []string, stdout io.Writer, stderr io.Writer, deps appDeps)
|
||||
return writeAppError(stderr, err.Error(), exitCrash)
|
||||
}
|
||||
report, err := contextreport.Build(contextreport.Options{
|
||||
WorkspaceRoot: workspaceRoot,
|
||||
Provider: resolved.Provider,
|
||||
Registry: newCoreRegistry(workspaceRoot),
|
||||
ContextWindow: modelContextWindow(modelRegistry, resolved.Provider.Model),
|
||||
WorkspaceRoot: workspaceRoot,
|
||||
Provider: resolved.Provider,
|
||||
Registry: newCoreRegistry(workspaceRoot),
|
||||
ContextWindow: modelContextWindow(modelRegistry, resolved.Provider.Model),
|
||||
DeferThreshold: resolved.Tools.DeferThreshold,
|
||||
})
|
||||
if err != nil {
|
||||
return writeAppError(stderr, err.Error(), exitCrash)
|
||||
|
||||
@@ -77,11 +77,10 @@ func TestRegisterToolSearchIfEligibleSkipsWhenThresholdZero(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeferredEligibleCountIgnoresCoreTools(t *testing.T) {
|
||||
func TestDeferredEligibleCountIncludesOptionalBuiltins(t *testing.T) {
|
||||
registry := newCoreRegistry(t.TempDir())
|
||||
// newCoreRegistry holds only built-ins; none implement Deferred().
|
||||
if got := deferredEligibleCount(registry, agent.PermissionModeAuto, nil, nil); got != 0 {
|
||||
t.Fatalf("deferredEligibleCount(core) = %d, want 0", got)
|
||||
if got := deferredEligibleCount(registry, agent.PermissionModeAuto, nil, nil); got < 3 {
|
||||
t.Fatalf("deferredEligibleCount(core) = %d, want at least 3 optional built-ins", got)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -5,6 +5,8 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
@@ -12,6 +14,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/config"
|
||||
"github.com/Gitlawb/zero/internal/oauth"
|
||||
"github.com/Gitlawb/zero/internal/providerhealth"
|
||||
"github.com/Gitlawb/zero/internal/sessions"
|
||||
)
|
||||
@@ -109,6 +112,65 @@ func TestRunDoctorConnectivityProbesProvider(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunDoctorConnectivityUsesOAuthLogin(t *testing.T) {
|
||||
tokenPath := filepath.Join(t.TempDir(), "oauth-tokens.json")
|
||||
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
|
||||
t.Setenv("ZERO_OAUTH_STORAGE", "file")
|
||||
t.Setenv("ZERO_OAUTH_TOKENS_PATH", tokenPath)
|
||||
store, err := oauth.NewStore(oauth.StoreOptions{FilePath: tokenPath})
|
||||
if err != nil {
|
||||
t.Fatalf("NewStore: %v", err)
|
||||
}
|
||||
if err := store.Save(oauth.ProviderKey("xai"), oauth.Token{
|
||||
AccessToken: "oauth-doctor-token",
|
||||
ExpiresAt: time.Now().Add(time.Hour),
|
||||
}); err != nil {
|
||||
t.Fatalf("seed OAuth token: %v", err)
|
||||
}
|
||||
|
||||
var gotAuth string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"data":[]}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
profile := config.ProviderProfile{
|
||||
Name: "xai",
|
||||
CatalogID: "xai",
|
||||
ProviderKind: config.ProviderKindOpenAICompatible,
|
||||
BaseURL: server.URL + "/v1",
|
||||
Model: "grok-4",
|
||||
}
|
||||
var stdout bytes.Buffer
|
||||
var stderr bytes.Buffer
|
||||
exitCode := runWithDeps([]string{"doctor", "--connectivity"}, &stdout, &stderr, appDeps{
|
||||
getwd: func() (string, error) {
|
||||
return t.TempDir(), nil
|
||||
},
|
||||
resolveConfig: func(string, config.Overrides) (config.ResolvedConfig, error) {
|
||||
return config.ResolvedConfig{
|
||||
Provider: profile,
|
||||
Providers: []config.ProviderProfile{profile},
|
||||
ActiveProvider: profile.Name,
|
||||
}, nil
|
||||
},
|
||||
probeProviderHealth: providerhealth.Probe,
|
||||
now: fixedCLITime("2026-07-29T20:00:00Z"),
|
||||
})
|
||||
|
||||
if exitCode != exitSuccess {
|
||||
t.Fatalf("doctor exit = %d, want success; stderr=%q\n%s", exitCode, stderr.String(), stdout.String())
|
||||
}
|
||||
if gotAuth != "Bearer oauth-doctor-token" {
|
||||
t.Fatalf("Authorization = %q, want OAuth bearer", gotAuth)
|
||||
}
|
||||
if !strings.Contains(stdout.String(), "[pass] provider.connectivity") {
|
||||
t.Fatalf("doctor did not report passing connectivity:\n%s", stdout.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunDoctorConnectivityReportsProviderHealthFailure(t *testing.T) {
|
||||
var stdout bytes.Buffer
|
||||
var stderr bytes.Buffer
|
||||
|
||||
@@ -19,6 +19,7 @@ func TestProviderCatalogRuntimeProvidersShowOnboardingSetupHints(t *testing.T) {
|
||||
{id: "openai", name: "OpenAI", transport: "openai", defaultModel: "gpt-4.1"},
|
||||
{id: "groq", name: "Groq", transport: "openai-compatible", defaultModel: "llama-3.3-70b-versatile"},
|
||||
{id: "longcat", name: "LongCat", transport: "openai-compatible", defaultModel: "LongCat-2.0"},
|
||||
{id: "fireworks", name: "Fireworks AI", transport: "openai-compatible", defaultModel: "accounts/fireworks/models/kimi-k2p7-code"},
|
||||
{id: "ollama-cloud", name: "Ollama Cloud", transport: "openai-compatible", defaultModel: "qwen3-coder:480b"},
|
||||
{id: "ollama", name: "Ollama Local", transport: "openai-compatible", defaultModel: "llama3.1"},
|
||||
}
|
||||
|
||||
@@ -113,7 +113,7 @@ type PreferencesConfig struct {
|
||||
// ZERO_THEME, so a /theme choice survives restart. Empty = unset (defaults auto).
|
||||
Theme string `json:"theme,omitempty"`
|
||||
// Recaps is a tri-state: nil (unset) defaults to ON; an explicit false means
|
||||
// the user turned post-turn recaps off. A *bool is its own tri-state, so no
|
||||
// the user turned idle recaps off. A *bool is its own tri-state, so no
|
||||
// custom unmarshal is needed (unlike ToolsConfig.DeferThreshold's int).
|
||||
Recaps *bool `json:"recaps,omitempty"`
|
||||
}
|
||||
@@ -130,7 +130,7 @@ type RecentModelEntry struct {
|
||||
// MaxRecentModels caps the persisted/displayed recent-selection history.
|
||||
const MaxRecentModels = 5
|
||||
|
||||
// RecapsEnabled reports whether post-turn recaps are on. Unset defaults to ON.
|
||||
// RecapsEnabled reports whether idle recaps are on. Unset defaults to ON.
|
||||
func (p PreferencesConfig) RecapsEnabled() bool {
|
||||
return p.Recaps == nil || *p.Recaps
|
||||
}
|
||||
|
||||
@@ -546,7 +546,7 @@ func SetRecentModels(path string, entries []RecentModelEntry) (FileConfig, error
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
// SetRecapsEnabled persists the post-turn recap preference, mirroring
|
||||
// SetRecapsEnabled persists the idle recap preference, mirroring
|
||||
// SetFavoriteModels (read-modify-atomic-write).
|
||||
func SetRecapsEnabled(path string, enabled bool) (FileConfig, error) {
|
||||
path = strings.TrimSpace(path)
|
||||
|
||||
@@ -40,6 +40,7 @@ type Options struct {
|
||||
Provider config.ProviderProfile
|
||||
Registry *tools.Registry
|
||||
ContextWindow int
|
||||
DeferThreshold int
|
||||
ProjectContextFiles []string
|
||||
}
|
||||
|
||||
@@ -56,6 +57,8 @@ type Report struct {
|
||||
FreeTokens int `json:"freeTokens"`
|
||||
UsedFraction float64 `json:"usedFraction"`
|
||||
ToolCount int `json:"toolCount"`
|
||||
AvailableToolCount int `json:"availableToolCount,omitempty"`
|
||||
DeferredToolCount int `json:"deferredToolCount,omitempty"`
|
||||
ProjectGuidelineFile string `json:"projectGuidelineFile,omitempty"`
|
||||
Categories []Category `json:"categories"`
|
||||
}
|
||||
@@ -113,8 +116,10 @@ func Build(options Options) (Report, error) {
|
||||
categories = append(categories, category(CategoryWorkspaceMap, "Workspace map", estimateTextTokens(workspaceMap), report.ContextWindow))
|
||||
}
|
||||
|
||||
toolCount, toolTokens := estimateRegistryTools(options.Registry)
|
||||
toolCount, availableTools, deferredTools, toolTokens := estimateRegistryTools(options.Registry, options.DeferThreshold)
|
||||
report.ToolCount = toolCount
|
||||
report.AvailableToolCount = availableTools
|
||||
report.DeferredToolCount = deferredTools
|
||||
if toolTokens > 0 {
|
||||
categories = append(categories, category(CategoryTools, "Tools", toolTokens, report.ContextWindow))
|
||||
}
|
||||
@@ -175,7 +180,11 @@ func Format(report Report) string {
|
||||
lines = append(lines, fmt.Sprintf("usage: %s tokens (context window unknown)", compactNumber(report.UsedTokens)))
|
||||
}
|
||||
if report.ToolCount > 0 {
|
||||
lines = append(lines, fmt.Sprintf("tools: %d", report.ToolCount))
|
||||
toolLine := fmt.Sprintf("tools: %d exposed", report.ToolCount)
|
||||
if report.AvailableToolCount > report.ToolCount {
|
||||
toolLine += fmt.Sprintf(", %d available, %d deferred", report.AvailableToolCount, report.DeferredToolCount)
|
||||
}
|
||||
lines = append(lines, toolLine)
|
||||
}
|
||||
if report.ProjectGuidelineFile != "" {
|
||||
lines = append(lines, "project_guidelines: "+report.ProjectGuidelineFile)
|
||||
@@ -212,29 +221,51 @@ func category(key string, name string, tokens int, contextWindow int) Category {
|
||||
return cat
|
||||
}
|
||||
|
||||
func estimateRegistryTools(registry *tools.Registry) (int, int) {
|
||||
func estimateRegistryTools(registry *tools.Registry, deferThreshold int) (int, int, int, int) {
|
||||
if registry == nil {
|
||||
return 0, 0
|
||||
return 0, 0, 0, 0
|
||||
}
|
||||
all := registry.All()
|
||||
sort.Slice(all, func(left int, right int) bool {
|
||||
return all[left].Name() < all[right].Name()
|
||||
})
|
||||
total := 0
|
||||
count := 0
|
||||
var visible []tools.Tool
|
||||
var deferred []tools.Tool
|
||||
for _, tool := range all {
|
||||
if tool.Safety().Permission == tools.PermissionDeny {
|
||||
continue
|
||||
}
|
||||
count++
|
||||
visible = append(visible, tool)
|
||||
if tools.IsDeferred(tool) {
|
||||
deferred = append(deferred, tool)
|
||||
}
|
||||
}
|
||||
active := deferThreshold > 0 && len(deferred) >= deferThreshold
|
||||
exposed := visible
|
||||
if active {
|
||||
exposed = make([]tools.Tool, 0, len(visible)-len(deferred)+1)
|
||||
for _, tool := range visible {
|
||||
if !tools.IsDeferred(tool) && tool.Name() != tools.ToolSearchToolName {
|
||||
exposed = append(exposed, tool)
|
||||
}
|
||||
}
|
||||
exposed = append(exposed, tools.NewToolSearchTool(registry))
|
||||
}
|
||||
|
||||
total := 0
|
||||
for _, tool := range exposed {
|
||||
total += estimateTextTokens(tool.Name())
|
||||
total += estimateTextTokens(tool.Description())
|
||||
description := tool.Description()
|
||||
if active && tool.Name() == tools.ToolSearchToolName {
|
||||
description = tools.BuildToolSearchDescription(deferred)
|
||||
}
|
||||
total += estimateTextTokens(description)
|
||||
if encoded, err := json.Marshal(tool.Parameters()); err == nil {
|
||||
total += estimateTextTokens(string(encoded))
|
||||
}
|
||||
total += toolDefinitionOverheadTokens
|
||||
}
|
||||
return count, total
|
||||
return len(exposed), len(visible), len(deferred), total
|
||||
}
|
||||
|
||||
func readProjectGuidelines(root string, names []string) (string, string) {
|
||||
|
||||
@@ -130,6 +130,35 @@ func TestBuildHasStableJSONContractAndCategoryMath(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildReportsInitialDeferredToolSurface(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
registry := tools.NewRegistry()
|
||||
for _, tool := range tools.CoreToolsScoped(root, nil) {
|
||||
registry.Register(tool)
|
||||
}
|
||||
|
||||
eager, err := Build(Options{WorkspaceRoot: root, Registry: registry})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
deferred, err := Build(Options{WorkspaceRoot: root, Registry: registry, DeferThreshold: 3})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if deferred.DeferredToolCount < 3 {
|
||||
t.Fatalf("DeferredToolCount = %d, want at least 3", deferred.DeferredToolCount)
|
||||
}
|
||||
if deferred.AvailableToolCount != eager.ToolCount {
|
||||
t.Fatalf("available tools = %d, want eager total %d", deferred.AvailableToolCount, eager.ToolCount)
|
||||
}
|
||||
if deferred.ToolCount >= deferred.AvailableToolCount {
|
||||
t.Fatalf("initial exposed tools = %d, available = %d; expected a smaller initial surface", deferred.ToolCount, deferred.AvailableToolCount)
|
||||
}
|
||||
if categoryByKey(deferred, CategoryTools).Tokens >= categoryByKey(eager, CategoryTools).Tokens {
|
||||
t.Fatalf("deferred tool schemas did not reduce the initial context")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildClampsFreeBudgetWhenOverContextWindow(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
writeTestFile(t, root, "AGENTS.md", strings.Repeat("very large project rule\n", 1000))
|
||||
|
||||
@@ -2,6 +2,7 @@ package imageinput
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
@@ -16,6 +17,13 @@ import (
|
||||
// clipboard, or (nil, "", nil) when the clipboard has no image. Called when
|
||||
// text clipboard is empty (the user pasted a screenshot). The media type is
|
||||
// sniffed from the bytes, not trusted from the clipboard.
|
||||
// ErrClipboardImageUnreadable reports that the clipboard holds an image the host
|
||||
// has no way to extract. It is deliberately distinct from "no image on the
|
||||
// clipboard", which is an ordinary no-op: this one is actionable, and silence is
|
||||
// the wrong response to it.
|
||||
var ErrClipboardImageUnreadable = errors.New(
|
||||
"the clipboard holds an image but no helper on this host can read it; on macOS install pngpaste (brew install pngpaste) or a Python with pyobjc")
|
||||
|
||||
func ReadClipboardImage() ([]byte, string, error) {
|
||||
data, err := readClipboardImageBytes()
|
||||
if err != nil {
|
||||
@@ -95,35 +103,71 @@ func readClipboardImageWindows() ([]byte, error) {
|
||||
|
||||
// readClipboardImageDarwin uses osascript to check for and read a clipboard
|
||||
// image. Returns (nil, nil) when no image is present.
|
||||
func readClipboardImageDarwin() ([]byte, error) {
|
||||
// Check clipboard info for image classes.
|
||||
check := `osascript -e 'clipboard info'`
|
||||
out, err := exec.Command("sh", "-c", check).Output()
|
||||
if err != nil {
|
||||
return nil, nil
|
||||
// darwinClipboardInfo reports what the pasteboard is currently offering, and
|
||||
// darwinClipboardExtract pulls the PNG bytes out of it.
|
||||
//
|
||||
// Both are vars so the two-step logic below can be driven in a test. The step
|
||||
// that matters, "the clipboard holds an image and nothing here can extract it",
|
||||
// only happens when the first succeeds and the second fails, which cannot be
|
||||
// arranged by stubbing the caller: a test that substitutes ReadClipboardImage
|
||||
// asserts the plumbing around this function rather than the decision inside it,
|
||||
// and stays green with that decision reverted.
|
||||
var (
|
||||
darwinClipboardInfo = func() (string, error) {
|
||||
out, err := exec.Command("sh", "-c", `osascript -e 'clipboard info'`).Output()
|
||||
return string(out), err
|
||||
}
|
||||
info := string(out)
|
||||
if !strings.Contains(info, "PNG") && !strings.Contains(info, "JPEG") && !strings.Contains(info, "TIFF") && !strings.Contains(info, "GIF") {
|
||||
return nil, nil
|
||||
}
|
||||
// Write clipboard image to a temp file via AppleScript, then read it.
|
||||
// Using pngpaste if available, falling back to a Python one-liner.
|
||||
cmd := exec.Command("sh", "-c", `pngpaste - 2>/dev/null || python3 -c "
|
||||
darwinClipboardExtract = func() ([]byte, error) {
|
||||
// pngpaste if available, falling back to a Python one-liner. Neither
|
||||
// ships with macOS: pngpaste is Homebrew-only and Apple dropped the
|
||||
// bundled PyObjC, so a stock Mac fails here every time.
|
||||
cmd := exec.Command("sh", "-c", `pngpaste - 2>/dev/null || python3 -c "
|
||||
import AppKit, sys
|
||||
pb = AppKit.NSPasteboard.generalPasteboard()
|
||||
data = pb.dataForType_(AppKit.NSPasteboardTypePNG)
|
||||
if data:
|
||||
sys.stdout.buffer.write(data.bytes())
|
||||
"`)
|
||||
var stdout bytes.Buffer
|
||||
cmd.Stdout = &stdout
|
||||
if err := cmd.Run(); err != nil {
|
||||
var stdout bytes.Buffer
|
||||
cmd.Stdout = &stdout
|
||||
if err := cmd.Run(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return stdout.Bytes(), nil
|
||||
}
|
||||
)
|
||||
|
||||
// darwinClipboardOffersImage reports whether the pasteboard info line names an
|
||||
// image class.
|
||||
func darwinClipboardOffersImage(info string) bool {
|
||||
for _, class := range []string{"PNG", "JPEG", "TIFF", "GIF"} {
|
||||
if strings.Contains(info, class) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func readClipboardImageDarwin() ([]byte, error) {
|
||||
info, err := darwinClipboardInfo()
|
||||
if err != nil {
|
||||
return nil, nil
|
||||
}
|
||||
if stdout.Len() == 0 {
|
||||
if !darwinClipboardOffersImage(info) {
|
||||
return nil, nil
|
||||
}
|
||||
return stdout.Bytes(), nil
|
||||
// Past this point the clipboard is KNOWN to hold an image, because the probe
|
||||
// above said so. A failure here is therefore not "no image", it is "there is
|
||||
// an image and nothing on this host can extract it". Reporting that as
|
||||
// absence is what left users staring at a paste that silently did nothing.
|
||||
data, err := darwinClipboardExtract()
|
||||
if err != nil {
|
||||
return nil, ErrClipboardImageUnreadable
|
||||
}
|
||||
if len(data) == 0 {
|
||||
return nil, ErrClipboardImageUnreadable
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// readClipboardImageLinux tries wl-paste (Wayland) then xclip (X11) to read
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
package imageinput
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// Drives readClipboardImageDarwin's own decision rather than the plumbing around
|
||||
// it. Both shell steps are substituted, so this runs on any platform and, more
|
||||
// importantly, actually reaches the branch under test: a test that stubs
|
||||
// ReadClipboardImage instead stays green with the production change reverted,
|
||||
// which is how the first version of this coverage slipped through.
|
||||
func TestReadClipboardImageDarwinDistinguishesUnreadableFromAbsent(t *testing.T) {
|
||||
restore := func() func() {
|
||||
info, extract := darwinClipboardInfo, darwinClipboardExtract
|
||||
return func() { darwinClipboardInfo, darwinClipboardExtract = info, extract }
|
||||
}()
|
||||
t.Cleanup(restore)
|
||||
|
||||
const pngInfo = `«class PNGf», «class 8BPS»`
|
||||
|
||||
t.Run("image present but no extractor is actionable", func(t *testing.T) {
|
||||
darwinClipboardInfo = func() (string, error) { return pngInfo, nil }
|
||||
darwinClipboardExtract = func() ([]byte, error) { return nil, errors.New("exit status 127") }
|
||||
|
||||
data, err := readClipboardImageDarwin()
|
||||
if !errors.Is(err, ErrClipboardImageUnreadable) {
|
||||
t.Fatalf("err = %v, want ErrClipboardImageUnreadable; a stock Mac gets silence instead of a remedy", err)
|
||||
}
|
||||
if data != nil {
|
||||
t.Fatalf("data = %q, want none", data)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("extractor producing nothing is also actionable", func(t *testing.T) {
|
||||
darwinClipboardInfo = func() (string, error) { return pngInfo, nil }
|
||||
darwinClipboardExtract = func() ([]byte, error) { return nil, nil }
|
||||
|
||||
if _, err := readClipboardImageDarwin(); !errors.Is(err, ErrClipboardImageUnreadable) {
|
||||
t.Fatalf("err = %v, want ErrClipboardImageUnreadable for an extractor that exits 0 with no bytes", err)
|
||||
}
|
||||
})
|
||||
|
||||
// The other half of the distinction, and the reason this cannot simply always
|
||||
// report an error: an empty clipboard is an ordinary no-op, and turning that
|
||||
// into a notice would make every stray Ctrl+V nag the user.
|
||||
t.Run("no image on the clipboard stays silent", func(t *testing.T) {
|
||||
darwinClipboardInfo = func() (string, error) { return `«class utf8»`, nil }
|
||||
darwinClipboardExtract = func() ([]byte, error) {
|
||||
t.Fatal("extractor ran despite the clipboard offering no image")
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
data, err := readClipboardImageDarwin()
|
||||
if err != nil || data != nil {
|
||||
t.Fatalf("data=%q err=%v, want a silent no-op", data, err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("probe failure stays silent", func(t *testing.T) {
|
||||
darwinClipboardInfo = func() (string, error) { return "", errors.New("osascript missing") }
|
||||
darwinClipboardExtract = func() ([]byte, error) {
|
||||
t.Fatal("extractor ran despite the info probe failing")
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
if data, err := readClipboardImageDarwin(); err != nil || data != nil {
|
||||
t.Fatalf("data=%q err=%v, want a silent no-op", data, err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("working extractor returns the bytes", func(t *testing.T) {
|
||||
darwinClipboardInfo = func() (string, error) { return pngInfo, nil }
|
||||
darwinClipboardExtract = func() ([]byte, error) { return []byte("png-bytes"), nil }
|
||||
|
||||
data, err := readClipboardImageDarwin()
|
||||
if err != nil {
|
||||
t.Fatalf("err = %v, want none", err)
|
||||
}
|
||||
if string(data) != "png-bytes" {
|
||||
t.Fatalf("data = %q, want the extractor's bytes", data)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -138,6 +138,7 @@ var descriptors = []Descriptor{
|
||||
openAICompat("groq", "Groq", "https://api.groq.com/openai/v1", "llama-3.3-70b-versatile", []string{"GROQ_API_KEY"}),
|
||||
openAICompat("deepseek", "DeepSeek", "https://api.deepseek.com/v1", "deepseek-chat", []string{"DEEPSEEK_API_KEY"}),
|
||||
openAICompat("together", "Together AI", "https://api.together.xyz/v1", "meta-llama/Llama-3.3-70B-Instruct-Turbo", []string{"TOGETHER_API_KEY"}),
|
||||
openAICompat("fireworks", "Fireworks AI", "https://api.fireworks.ai/inference/v1", "accounts/fireworks/models/kimi-k2p7-code", []string{"FIREWORKS_API_KEY"}, "fireworks-ai"),
|
||||
openAICompat("dashscope", "DashScope", "https://dashscope-intl.aliyuncs.com/compatible-mode/v1", "qwen-plus", []string{"DASHSCOPE_API_KEY", "QWEN_API_KEY"}, "qwen"),
|
||||
openAICompat("moonshot", "Moonshot AI", "https://api.moonshot.ai/v1", "kimi-k2-0905-preview", []string{"MOONSHOT_API_KEY"}, "kimi"),
|
||||
openAICompat("atlascloud", "Atlas Cloud", "https://api.atlascloud.ai/v1", "qwen/qwen3.5-flash", []string{"ATLASCLOUD_API_KEY"}, "atlas cloud", "atlas"),
|
||||
|
||||
@@ -22,6 +22,7 @@ var expectedCatalogIDs = []string{
|
||||
"groq",
|
||||
"deepseek",
|
||||
"together",
|
||||
"fireworks",
|
||||
"dashscope",
|
||||
"moonshot",
|
||||
"atlascloud",
|
||||
@@ -155,6 +156,28 @@ func TestLongCatDescriptor(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestFireworksDescriptor(t *testing.T) {
|
||||
descriptor, err := Require("fireworks")
|
||||
if err != nil {
|
||||
t.Fatalf("Require(fireworks) error = %v", err)
|
||||
}
|
||||
if descriptor.Name != "Fireworks AI" {
|
||||
t.Fatalf("Name = %q, want Fireworks AI", descriptor.Name)
|
||||
}
|
||||
if descriptor.DefaultBaseURL != "https://api.fireworks.ai/inference/v1" {
|
||||
t.Fatalf("DefaultBaseURL = %q, want Fireworks OpenAI-compatible endpoint", descriptor.DefaultBaseURL)
|
||||
}
|
||||
if descriptor.DefaultModel != "accounts/fireworks/models/kimi-k2p7-code" {
|
||||
t.Fatalf("DefaultModel = %q, want accounts/fireworks/models/kimi-k2p7-code", descriptor.DefaultModel)
|
||||
}
|
||||
if descriptor.Transport != TransportOpenAICompatible {
|
||||
t.Fatalf("Transport = %q, want %q", descriptor.Transport, TransportOpenAICompatible)
|
||||
}
|
||||
if !reflect.DeepEqual(descriptor.AuthEnvVars, []string{"FIREWORKS_API_KEY"}) {
|
||||
t.Fatalf("AuthEnvVars = %#v, want FIREWORKS_API_KEY", descriptor.AuthEnvVars)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAtlasCloudDescriptor(t *testing.T) {
|
||||
descriptor, err := Require("atlascloud")
|
||||
if err != nil {
|
||||
@@ -346,7 +369,7 @@ func TestListByTransportPreservesCatalogOrder(t *testing.T) {
|
||||
TransportBedrock: {"bedrock"},
|
||||
TransportVertex: {"vertex"},
|
||||
TransportAnthropicCompat: {"minimax", "minimaxi-cn", "opencode-go-anthropic-compatible", "custom-anthropic-compatible"},
|
||||
TransportOpenAICompat: {"gitlawb-opengateway", "aimlapi", "ollama-cloud", "ollama", "lmstudio", "openrouter", "huggingface", "chatgpt", "groq", "deepseek", "together", "dashscope", "moonshot", "atlascloud", "longcat", "nvidia-nim", "mistral", "github", "xai", "venice", "xiaomi-mimo", "bankr", "zai", "zai-cn", "kilocode", "opencode", "opencode-go", "atomic-chat", "chatgpt-proxy", "custom-openai-compatible"},
|
||||
TransportOpenAICompat: {"gitlawb-opengateway", "aimlapi", "ollama-cloud", "ollama", "lmstudio", "openrouter", "huggingface", "chatgpt", "groq", "deepseek", "together", "fireworks", "dashscope", "moonshot", "atlascloud", "longcat", "nvidia-nim", "mistral", "github", "xai", "venice", "xiaomi-mimo", "bankr", "zai", "zai-cn", "kilocode", "opencode", "opencode-go", "atomic-chat", "chatgpt-proxy", "custom-openai-compatible"},
|
||||
}
|
||||
|
||||
for transport, wantIDs := range cases {
|
||||
|
||||
@@ -46,12 +46,13 @@ const (
|
||||
const defaultTimeout = 5 * time.Second
|
||||
|
||||
type Options struct {
|
||||
Profile config.ProviderProfile
|
||||
Connectivity bool
|
||||
HTTPClient *http.Client
|
||||
Resolver Resolver
|
||||
Timeout time.Duration
|
||||
UserAgent string
|
||||
Profile config.ProviderProfile
|
||||
Connectivity bool
|
||||
HTTPClient *http.Client
|
||||
Resolver Resolver
|
||||
OAuthResolver providerio.TokenResolver
|
||||
Timeout time.Duration
|
||||
UserAgent string
|
||||
}
|
||||
|
||||
type Resolver interface {
|
||||
@@ -207,11 +208,26 @@ func Probe(ctx context.Context, options Options) Result {
|
||||
"providerKind": metadata.ProviderKind,
|
||||
}, profile))
|
||||
|
||||
if credentialRequired(profile, metadata.ProviderKind) && !hasCredential(profile) {
|
||||
oauthConfigured := options.OAuthResolver != nil
|
||||
if oauthConfigured {
|
||||
var resolveErr error
|
||||
profile, oauthConfigured, resolveErr = profileWithOAuthCredential(ctx, profile, options.OAuthResolver, false)
|
||||
if resolveErr != nil {
|
||||
result.add(check("provider.auth", "Provider auth", StatusFail, CategoryAuth, "Provider OAuth credential could not be resolved: "+resolveErr.Error(), map[string]any{
|
||||
"oauthLogin": true,
|
||||
}, profile))
|
||||
return result.finalize()
|
||||
}
|
||||
}
|
||||
if credentialRequired(profile, metadata.ProviderKind) && !hasCredential(profile) && !oauthConfigured {
|
||||
result.add(check("provider.auth", "Provider auth", StatusFail, CategoryAuth, fmt.Sprintf("Provider %s requires API credentials.", providerName(profile)), credentialDetails(profile), profile))
|
||||
return result.finalize()
|
||||
}
|
||||
if hasCredential(profile) {
|
||||
if oauthConfigured {
|
||||
result.add(check("provider.auth", "Provider auth", StatusPass, CategoryAuth, fmt.Sprintf("Provider %s has OAuth credentials configured.", providerName(profile)), map[string]any{
|
||||
"oauthLogin": true,
|
||||
}, profile))
|
||||
} else if hasCredential(profile) {
|
||||
result.add(check("provider.auth", "Provider auth", StatusPass, CategoryAuth, fmt.Sprintf("Provider %s has credentials configured.", providerName(profile)), credentialDetails(profile), profile))
|
||||
} else {
|
||||
result.add(check("provider.auth", "Provider auth", StatusPass, CategoryAuth, fmt.Sprintf("Provider %s does not require API credentials.", providerName(profile)), credentialDetails(profile), profile))
|
||||
@@ -286,9 +302,35 @@ func connectivityCheck(ctx context.Context, profile config.ProviderProfile, kind
|
||||
if err != nil {
|
||||
return classifyTransportError(err, profile)
|
||||
}
|
||||
defer func() {
|
||||
_ = response.Body.Close()
|
||||
}()
|
||||
defer func(body io.ReadCloser) {
|
||||
_ = body.Close()
|
||||
}(response.Body)
|
||||
if response.StatusCode == http.StatusUnauthorized && options.OAuthResolver != nil {
|
||||
refreshed, ok, refreshErr := profileWithOAuthCredential(requestCtx, options.Profile, options.OAuthResolver, true)
|
||||
if refreshErr != nil {
|
||||
return check("provider.connectivity", "Provider connectivity", StatusFail, CategoryAuth, "Provider OAuth credential refresh failed: "+refreshErr.Error(), map[string]any{
|
||||
"statusCode": response.StatusCode,
|
||||
}, profile)
|
||||
}
|
||||
if ok {
|
||||
_ = response.Body.Close()
|
||||
profile = refreshed
|
||||
request, allowLoopbackOrPrivate, err = healthRequest(requestCtx, profile, kind, options)
|
||||
if err != nil {
|
||||
return check("provider.connectivity", "Provider connectivity", StatusFail, CategoryConfig, err.Error(), nil, profile)
|
||||
}
|
||||
if options.HTTPClient == nil {
|
||||
client = newConnectivityClient(timeout, options.Resolver, sensitiveAuthHeaderNames(profile, kind), allowLoopbackOrPrivate)
|
||||
}
|
||||
response, err = client.Do(request)
|
||||
if err != nil {
|
||||
return classifyTransportError(err, profile)
|
||||
}
|
||||
defer func(body io.ReadCloser) {
|
||||
_ = body.Close()
|
||||
}(response.Body)
|
||||
}
|
||||
}
|
||||
|
||||
body, err := io.ReadAll(io.LimitReader(response.Body, 64*1024))
|
||||
if err != nil {
|
||||
@@ -320,6 +362,23 @@ func connectivityCheck(ctx context.Context, profile config.ProviderProfile, kind
|
||||
}, profile)
|
||||
}
|
||||
|
||||
func profileWithOAuthCredential(ctx context.Context, profile config.ProviderProfile, resolver providerio.TokenResolver, forceRefresh bool) (config.ProviderProfile, bool, error) {
|
||||
if resolver == nil {
|
||||
return profile, false, nil
|
||||
}
|
||||
header, value, ok, err := resolver(ctx, forceRefresh)
|
||||
if err != nil || !ok {
|
||||
return profile, false, err
|
||||
}
|
||||
if strings.TrimSpace(header) == "" {
|
||||
header = "Authorization"
|
||||
}
|
||||
profile.AuthHeader = header
|
||||
profile.AuthScheme = ""
|
||||
profile.AuthHeaderValue = value
|
||||
return profile, true, nil
|
||||
}
|
||||
|
||||
func healthRequest(ctx context.Context, profile config.ProviderProfile, kind config.ProviderKind, options Options) (*http.Request, bool, error) {
|
||||
baseURL, err := resolvedBaseURL(profile, kind)
|
||||
if err != nil {
|
||||
@@ -748,6 +807,9 @@ func redact(message string, profile config.ProviderProfile) string {
|
||||
|
||||
func providerSecrets(profile config.ProviderProfile) []string {
|
||||
secrets := []string{profile.APIKey, profile.AuthHeaderValue}
|
||||
if fields := strings.Fields(profile.AuthHeaderValue); len(fields) == 2 {
|
||||
secrets = append(secrets, fields[1])
|
||||
}
|
||||
for _, value := range profile.CustomHeaders {
|
||||
if strings.TrimSpace(value) != "" {
|
||||
secrets = append(secrets, value)
|
||||
|
||||
@@ -180,6 +180,182 @@ func TestProbeResolvedKindConnectivityAuthErrorRedactsSecret(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestProbeConnectivityUsesOAuthCredential(t *testing.T) {
|
||||
var gotAuth string
|
||||
resolverCalled := false
|
||||
client := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
gotAuth = r.Header.Get("Authorization")
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Status: "200 OK",
|
||||
Body: io.NopCloser(strings.NewReader(`{"data":[]}`)),
|
||||
Header: make(http.Header),
|
||||
}, nil
|
||||
})}
|
||||
|
||||
result := Probe(context.Background(), Options{
|
||||
Profile: config.ProviderProfile{
|
||||
Name: "xai",
|
||||
CatalogID: "xai",
|
||||
ProviderKind: config.ProviderKindOpenAICompatible,
|
||||
BaseURL: "https://api.example.com/v1",
|
||||
Model: "grok-4",
|
||||
},
|
||||
Connectivity: true,
|
||||
HTTPClient: client,
|
||||
Resolver: staticResolver{addr: netip.MustParseAddr("93.184.216.34")},
|
||||
OAuthResolver: func(context.Context, bool) (string, string, bool, error) {
|
||||
resolverCalled = true
|
||||
return "Authorization", "Bearer oauth-test-token", true, nil
|
||||
},
|
||||
})
|
||||
|
||||
if result.Status != StatusPass {
|
||||
t.Fatalf("Status = %q, want pass: %#v", result.Status, result.Checks)
|
||||
}
|
||||
if !resolverCalled {
|
||||
t.Fatal("OAuth resolver was not called")
|
||||
}
|
||||
if gotAuth != "Bearer oauth-test-token" {
|
||||
t.Fatalf("Authorization = %q, want OAuth bearer", gotAuth)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProbeConnectivityOAuthCredentialsAreRedacted(t *testing.T) {
|
||||
const (
|
||||
staticKey = "STATICKEYVALUE1234567890abcdef"
|
||||
oauthToken = "hf_VASANTHSENTINEL0011223344"
|
||||
)
|
||||
var authHeaders []string
|
||||
client := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
authHeaders = append(authHeaders, r.Header.Get("Authorization"))
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusUnauthorized,
|
||||
Status: "401 Unauthorized",
|
||||
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"invalid credentials ` + staticKey + ` and ` + oauthToken + `"}}`)),
|
||||
Header: make(http.Header),
|
||||
}, nil
|
||||
})}
|
||||
|
||||
result := Probe(context.Background(), Options{
|
||||
Profile: config.ProviderProfile{
|
||||
Name: "xai",
|
||||
CatalogID: "xai",
|
||||
ProviderKind: config.ProviderKindOpenAICompatible,
|
||||
BaseURL: "https://api.example.com/v1",
|
||||
APIKey: staticKey,
|
||||
Model: "grok-4",
|
||||
},
|
||||
Connectivity: true,
|
||||
HTTPClient: client,
|
||||
Resolver: staticResolver{addr: netip.MustParseAddr("93.184.216.34")},
|
||||
OAuthResolver: func(context.Context, bool) (string, string, bool, error) {
|
||||
return "Authorization", "Bearer " + oauthToken, true, nil
|
||||
},
|
||||
})
|
||||
|
||||
if len(authHeaders) != 2 {
|
||||
t.Fatalf("Authorization headers = %#v, want initial request and one refresh retry", authHeaders)
|
||||
}
|
||||
for _, header := range authHeaders {
|
||||
if header != "Bearer "+oauthToken {
|
||||
t.Fatalf("Authorization = %q, want OAuth bearer", header)
|
||||
}
|
||||
}
|
||||
check := result.Check("provider.connectivity")
|
||||
if check == nil || check.Category != CategoryAuth {
|
||||
t.Fatalf("provider.connectivity = %#v, want auth failure", check)
|
||||
}
|
||||
if strings.Contains(check.Message, staticKey) {
|
||||
t.Fatalf("static API key leaked in message: %q", check.Message)
|
||||
}
|
||||
if strings.Contains(check.Message, oauthToken) {
|
||||
t.Fatalf("OAuth token leaked in message: %q", check.Message)
|
||||
}
|
||||
if strings.Count(check.Message, "[REDACTED]") != 2 {
|
||||
t.Fatalf("expected both credentials to be redacted, got %q", check.Message)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProbeWithoutConnectivityRejectsUnavailableOAuthCredential(t *testing.T) {
|
||||
resolverCalled := false
|
||||
result := Probe(context.Background(), Options{
|
||||
Profile: config.ProviderProfile{
|
||||
Name: "xai",
|
||||
CatalogID: "xai",
|
||||
ProviderKind: config.ProviderKindOpenAICompatible,
|
||||
BaseURL: "https://api.example.com/v1",
|
||||
Model: "grok-4",
|
||||
},
|
||||
OAuthResolver: func(context.Context, bool) (string, string, bool, error) {
|
||||
resolverCalled = true
|
||||
return "", "", false, nil
|
||||
},
|
||||
})
|
||||
|
||||
if !resolverCalled {
|
||||
t.Fatal("OAuth resolver was not called")
|
||||
}
|
||||
if result.Status != StatusFail {
|
||||
t.Fatalf("Status = %q, want fail: %#v", result.Status, result.Checks)
|
||||
}
|
||||
check := result.Check("provider.auth")
|
||||
if check == nil || check.Status != StatusFail || check.Category != CategoryAuth {
|
||||
t.Fatalf("provider.auth = %#v, want auth failure", check)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProbeConnectivityRefreshesOAuthCredentialAfterUnauthorized(t *testing.T) {
|
||||
var authHeaders []string
|
||||
client := &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
authHeaders = append(authHeaders, r.Header.Get("Authorization"))
|
||||
status := http.StatusOK
|
||||
body := `{"data":[]}`
|
||||
if len(authHeaders) == 1 {
|
||||
status = http.StatusUnauthorized
|
||||
body = `{"error":{"message":"expired"}}`
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: status,
|
||||
Status: http.StatusText(status),
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
Header: make(http.Header),
|
||||
}, nil
|
||||
})}
|
||||
|
||||
var refreshFlags []bool
|
||||
result := Probe(context.Background(), Options{
|
||||
Profile: config.ProviderProfile{
|
||||
Name: "xai",
|
||||
CatalogID: "xai",
|
||||
ProviderKind: config.ProviderKindOpenAICompatible,
|
||||
BaseURL: "https://api.example.com/v1",
|
||||
Model: "grok-4",
|
||||
},
|
||||
Connectivity: true,
|
||||
HTTPClient: client,
|
||||
Resolver: staticResolver{addr: netip.MustParseAddr("93.184.216.34")},
|
||||
OAuthResolver: func(_ context.Context, forceRefresh bool) (string, string, bool, error) {
|
||||
refreshFlags = append(refreshFlags, forceRefresh)
|
||||
token := "stale-token"
|
||||
if forceRefresh {
|
||||
token = "fresh-token"
|
||||
}
|
||||
return "Authorization", "Bearer " + token, true, nil
|
||||
},
|
||||
})
|
||||
|
||||
if result.Status != StatusPass {
|
||||
t.Fatalf("Status = %q, want pass after refresh: %#v", result.Status, result.Checks)
|
||||
}
|
||||
if len(refreshFlags) != 2 || refreshFlags[0] || !refreshFlags[1] {
|
||||
t.Fatalf("OAuth refresh flags = %#v, want [false true]", refreshFlags)
|
||||
}
|
||||
if len(authHeaders) != 2 || authHeaders[0] != "Bearer stale-token" || authHeaders[1] != "Bearer fresh-token" {
|
||||
t.Fatalf("Authorization headers = %#v, want stale then fresh bearer", authHeaders)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProbeConnectivityClassifiesTimeout(t *testing.T) {
|
||||
client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||||
return nil, context.DeadlineExceeded
|
||||
|
||||
@@ -122,6 +122,14 @@ var curatedModels = map[string][]Model{
|
||||
{ID: "deepseek-ai/DeepSeek-R1", Description: "reasoning model"},
|
||||
{ID: "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", Description: "multimodal model"},
|
||||
},
|
||||
"fireworks": {
|
||||
// Serverless-only curated set: ordinary API-key setups hit /inference/v1
|
||||
// and cannot use on-demand deployment IDs.
|
||||
{ID: "accounts/fireworks/models/kimi-k2p7-code", Description: "catalog default"},
|
||||
{ID: "accounts/fireworks/models/deepseek-v4-flash", Description: "fast coding model"},
|
||||
{ID: "accounts/fireworks/models/gpt-oss-120b", Description: "general model"},
|
||||
{ID: "accounts/fireworks/models/deepseek-v4-pro", Description: "reasoning model"},
|
||||
},
|
||||
"dashscope": {
|
||||
{ID: "qwen-plus", Description: "catalog default"},
|
||||
{ID: "qwen-max", Description: "strong general model"},
|
||||
|
||||
@@ -28,6 +28,11 @@ func TestModelsAreProviderScoped(t *testing.T) {
|
||||
want: []string{"llama-3.3-70b-versatile", "openai/gpt-oss-120b"},
|
||||
notWant: []string{"gpt-4.1", "claude-sonnet-4.5"},
|
||||
},
|
||||
{
|
||||
provider: "fireworks",
|
||||
want: []string{"accounts/fireworks/models/kimi-k2p7-code", "accounts/fireworks/models/deepseek-v4-flash"},
|
||||
notWant: []string{"gpt-4.1", "llama-3.3-70b-versatile"},
|
||||
},
|
||||
{
|
||||
provider: "chatgpt",
|
||||
want: []string{"gpt-5.5", "gpt-5.4", "gpt-5.4-mini", "gpt-5.3-codex-spark"},
|
||||
|
||||
@@ -132,6 +132,8 @@ func ModelsDevProviderID(provider providercatalog.Descriptor) string {
|
||||
return "zai"
|
||||
case "minimaxi-cn":
|
||||
return "minimax"
|
||||
case "fireworks":
|
||||
return "fireworks-ai"
|
||||
default:
|
||||
return strings.TrimSpace(provider.ID)
|
||||
}
|
||||
|
||||
@@ -131,6 +131,7 @@ func TestModelsDevProviderIDMapsZeroAliases(t *testing.T) {
|
||||
"ollama-cloud": "ollama-cloud",
|
||||
"zai-cn": "zai",
|
||||
"minimaxi-cn": "minimax",
|
||||
"fireworks": "fireworks-ai",
|
||||
}
|
||||
for zeroID, want := range tests {
|
||||
provider, ok := providercatalog.Get(zeroID)
|
||||
|
||||
@@ -25,7 +25,9 @@ type TokenResolver func(ctx context.Context, forceRefresh bool) (header string,
|
||||
//
|
||||
// base carries the API-key auth config (key + default header/scheme). setExtra
|
||||
// (optional) sets the provider's non-auth headers on every attempt. Transient
|
||||
// retries (429/503/529) are handled by the inner SendWithRetry.
|
||||
// retries are handled by the inner SendWithRetry: 429/503/529 on their own
|
||||
// budget, and provably pre-send transport failures (refused or timed-out
|
||||
// connect, DNS, TLS handshake) on a separate, shorter one.
|
||||
func SendWithAuthRetry(
|
||||
ctx context.Context,
|
||||
client *http.Client,
|
||||
|
||||
@@ -3,6 +3,9 @@ package providerio
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
@@ -17,13 +20,22 @@ import (
|
||||
// (previously only the OpenAI provider retried; Anthropic and Gemini surfaced the
|
||||
// first failure).
|
||||
//
|
||||
// What is retried: ONLY 429 (rate limit) and 503 (service unavailable) — the
|
||||
// statuses where the server explicitly did NOT accept the request, so replaying
|
||||
// it cannot duplicate work. Other 5xx (500/502/504) and transport/network errors
|
||||
// are NOT retried: a completion POST is non-idempotent and may have reached and
|
||||
// been processed by the server, so replaying it could duplicate
|
||||
// (billable) work. Only the INITIAL request is ever in scope; once the response
|
||||
// body starts streaming it is never re-issued.
|
||||
// What is retried: 429 (rate limit) and 503 (service unavailable) — the statuses
|
||||
// where the server explicitly did NOT accept the request — plus PROVABLY
|
||||
// pre-send transport failures (DNS resolution, a refused/unreachable dial, a TLS
|
||||
// handshake timeout) where no request bytes ever left this host. In every one of
|
||||
// those cases we KNOW the server did not receive and process the request, so a
|
||||
// replay cannot duplicate work. Other 5xx (500/502/504) and every ambiguous or
|
||||
// post-send transport error (connection reset, broken pipe, EOF, a generic i/o
|
||||
// timeout, context deadline) are NOT retried: a completion POST is
|
||||
// non-idempotent and may already have reached and been processed by the server,
|
||||
// so replaying it could duplicate (billable) work. Only the INITIAL request is
|
||||
// ever in scope; once the response body starts streaming it is never re-issued.
|
||||
// Pre-send classification is by syscall errno (portable across POSIX and
|
||||
// Windows), scoped to an Op=="dial" *net.OpError so the same errno raised on a
|
||||
// post-send read/write can't be misread as pre-send. NXDOMAIN is treated as
|
||||
// permanent (not retried), and the pre-send path fails fast on its own
|
||||
// sub-second schedule rather than the seconds-long 429 backoff.
|
||||
|
||||
const defaultMaxRetryAttempts = 6
|
||||
|
||||
@@ -37,12 +49,32 @@ const maxBackoff = 30 * time.Second
|
||||
// schedule is 2s, 4s, 8s, 16s, then maxBackoff. A var so tests can shrink it.
|
||||
var retryBackoffBase = 2 * time.Second
|
||||
|
||||
// preSendMaxAttempts bounds retries of a provably pre-send transport failure.
|
||||
// Unlike a 429 (whose window is measured in seconds), a refused/unreachable dial
|
||||
// either recovers within a few hundred ms or is a deterministic misconfiguration
|
||||
// (wrong host/port, dead local daemon) that will never recover — so a short,
|
||||
// low-count schedule fails fast instead of stalling the agent ~60s on the
|
||||
// rate-limit schedule. This bounds the pre-send path independently of the caller's
|
||||
// maxAttempts, which is tuned for 429/503.
|
||||
const preSendMaxAttempts = 3
|
||||
|
||||
// preSendBackoffBase is the first wait before replaying a pre-send transport
|
||||
// failure. Modeled on the mid-stream reconnect base (internal/agent/reconnect.go),
|
||||
// it is sub-second so a transient dial blip recovers quickly and a permanent one
|
||||
// fails fast. A var so tests can shrink it.
|
||||
var preSendBackoffBase = 500 * time.Millisecond
|
||||
|
||||
// SendWithRetry issues an HTTP request, retrying ONLY the safe-to-replay server
|
||||
// responses (429 and 503, see ShouldRetryStatus) up to maxAttempts — backing off
|
||||
// between tries and honoring a server Retry-After header and context
|
||||
// cancellation. Other 5xx and transport/network errors are returned immediately,
|
||||
// never replayed (see the package note). The request is rebuilt from body each
|
||||
// attempt; setHeader (if non-nil) sets headers on every attempt.
|
||||
// cancellation. The one transport-error exception is a PROVABLY pre-send failure
|
||||
// (see isPreSendTransportError): a refused/unreachable dial or DNS failure where
|
||||
// no request bytes left this host, retried on its own short schedule
|
||||
// (preSendMaxAttempts / preSendBackoff). Every other 5xx and every ambiguous or
|
||||
// post-send transport error is returned immediately, never replayed, because a
|
||||
// non-idempotent completion POST may already have been received. The request is
|
||||
// rebuilt from body each attempt; setHeader (if non-nil) sets headers on every
|
||||
// attempt.
|
||||
//
|
||||
// It returns the final *http.Response (which the caller inspects for a non-2xx
|
||||
// status, exactly as before) or an error for a network failure / context
|
||||
@@ -60,7 +92,47 @@ func SendWithRetry(
|
||||
if maxAttempts <= 0 {
|
||||
maxAttempts = defaultMaxRetryAttempts
|
||||
}
|
||||
for attempt := 1; ; attempt++ {
|
||||
if client == nil {
|
||||
client = http.DefaultClient
|
||||
}
|
||||
// Redirects are still FOLLOWED, so a redirecting gateway keeps working, but
|
||||
// entering one makes that attempt non-retryable. The hazard is that client.Do
|
||||
// can transmit the original POST to the first host, follow a 307/308, and then
|
||||
// fail the dial to the redirect target: that error is an Op=="dial" error, so
|
||||
// isPreSendTransportError would read it as "nothing was sent" and replay a
|
||||
// completion the first host already received. Observing the redirect rather
|
||||
// than blocking it keeps both properties, since a request that was redirected
|
||||
// is never replayed while redirecting endpoints stay usable. The caller's own
|
||||
// CheckRedirect still runs and still decides whether to proceed.
|
||||
redirected := false
|
||||
callerCheckRedirect := client.CheckRedirect
|
||||
redirectAwareClient := *client
|
||||
redirectAwareClient.CheckRedirect = func(request *http.Request, via []*http.Request) error {
|
||||
redirected = true
|
||||
if callerCheckRedirect != nil {
|
||||
return callerCheckRedirect(request, via)
|
||||
}
|
||||
// Mirror net/http's default policy, which a nil CheckRedirect would have
|
||||
// applied: stop after 10 hops.
|
||||
if len(via) >= 10 {
|
||||
return errors.New("stopped after 10 redirects")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
// The two retry budgets are fully independent, each with its own counter and
|
||||
// its own backoff ordinal, because they answer different questions: a
|
||||
// dial/DNS/TLS failure is a fast local blip (preSendMaxAttempts, sub-second
|
||||
// schedule) while a 429/503 is a server-side window measured in seconds
|
||||
// (maxAttempts, retryBackoffBase). A single shared loop counter coupled them
|
||||
// in both directions: status retries starved the pre-send budget, and
|
||||
// pre-send retries shortened the status window and started its backoff at a
|
||||
// later rung. Neither kind of failure now consumes or advances the other.
|
||||
preSendAttempts := 0
|
||||
statusAttempts := 0
|
||||
for {
|
||||
// Per attempt: only a redirect entered by THIS attempt can make its own
|
||||
// transport error unsafe to replay.
|
||||
redirected = false
|
||||
request, err := http.NewRequestWithContext(ctx, method, url, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -70,32 +142,78 @@ func SendWithRetry(
|
||||
}
|
||||
|
||||
connectSpan := trace.FromContext(ctx).Span(trace.SpanProviderConnect)
|
||||
response, err := client.Do(request)
|
||||
response, err := redirectAwareClient.Do(request)
|
||||
connectSpan.End()
|
||||
if err != nil {
|
||||
// A transport failure on a POST does NOT mean the server didn't receive
|
||||
// it — the request may have arrived and be generating a (billable,
|
||||
// non-idempotent) completion while only the response/connection failed.
|
||||
// Replaying it could duplicate that work, so surface the error instead
|
||||
// of retrying. Only 429/503 responses (below), where we KNOW the request
|
||||
// was not accepted, are safe to retry.
|
||||
// Context cancellation always surfaces as cancellation, never a retry.
|
||||
if ctx.Err() != nil {
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
// A transport failure on a POST does NOT usually mean the server didn't
|
||||
// receive it — the request may have arrived and be generating a
|
||||
// (billable, non-idempotent) completion while only the response or
|
||||
// connection failed. Replaying that could duplicate work, so it is
|
||||
// surfaced immediately. The one safe exception is a PROVABLY pre-send
|
||||
// failure (see isPreSendTransportError): the request never left this
|
||||
// host, so a replay cannot duplicate anything — exactly like a 429/503
|
||||
// where we KNOW the request was not accepted. The pre-send path uses
|
||||
// its own short bound and sub-second schedule (preSendMaxAttempts /
|
||||
// preSendBackoff), not the seconds-long 429 schedule, so a permanent
|
||||
// dial failure fails fast instead of stalling the agent for ~60s.
|
||||
// A redirect means the original POST already left this host, so the
|
||||
// dial that failed belongs to a later hop and replaying would duplicate
|
||||
// a completion the first host may already have processed.
|
||||
if isPreSendTransportError(err) && !redirected {
|
||||
preSendAttempts++
|
||||
if preSendAttempts < preSendMaxAttempts {
|
||||
if r := trace.FromContext(ctx); r != nil {
|
||||
r.Counter(trace.CounterRetryCount, 1)
|
||||
}
|
||||
if preSendBackoff(ctx, preSendAttempts) {
|
||||
continue
|
||||
}
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if ShouldRetryStatus(response.StatusCode) && attempt < maxAttempts {
|
||||
if r := trace.FromContext(ctx); r != nil {
|
||||
r.Counter(trace.CounterRetryCount, 1)
|
||||
// Deliberately NOT guarded on `redirected`, unlike the transport-error
|
||||
// branch above. The two differ in what is known about the request. A
|
||||
// transport failure after a redirect leaves it unknown whether the
|
||||
// redirect target received and began processing the POST, so replaying
|
||||
// could duplicate billable work. Here every hop answered: a 307/308 is the
|
||||
// origin declining and delegating, and a retryable status is the target
|
||||
// declining as well, so nothing accepted the request and a replay
|
||||
// duplicates nothing.
|
||||
//
|
||||
// The exposure is also identical to a retryable status with no redirect in
|
||||
// play, which this package has always retried. Guarding on redirects here
|
||||
// would not close a gap, it would make a rate-limited completion behind a
|
||||
// redirecting gateway fail outright instead of backing off.
|
||||
// The transport hop succeeded, so the connect phase is healthy again and
|
||||
// the pre-send budget starts over. It bounds CONSECUTIVE connect failures,
|
||||
// not the lifetime of the call: without this reset, two early dial retries
|
||||
// followed by a 429 and then one more dial blip would surface a transport
|
||||
// error while the status budget still had room, which is the coupling
|
||||
// these two counters were split apart to remove. Total work stays bounded,
|
||||
// since every reset costs a status attempt and those are capped.
|
||||
preSendAttempts = 0
|
||||
|
||||
if ShouldRetryStatus(response.StatusCode) {
|
||||
statusAttempts++
|
||||
if statusAttempts < maxAttempts {
|
||||
if r := trace.FromContext(ctx); r != nil {
|
||||
r.Counter(trace.CounterRetryCount, 1)
|
||||
}
|
||||
wait := RetryAfter(response)
|
||||
_ = response.Body.Close()
|
||||
if Backoff(ctx, statusAttempts, wait) {
|
||||
continue
|
||||
}
|
||||
// Backoff aborted: the only reason it returns false is ctx cancellation.
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
wait := RetryAfter(response)
|
||||
_ = response.Body.Close()
|
||||
if Backoff(ctx, attempt, wait) {
|
||||
continue
|
||||
}
|
||||
// Backoff aborted: the only reason it returns false is ctx cancellation.
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
// Success, a non-retryable status, or retries exhausted on a retryable
|
||||
@@ -109,6 +227,119 @@ func SendWithRetry(
|
||||
}
|
||||
}
|
||||
|
||||
// isPreSendTransportError reports whether a transport error PROVES no request
|
||||
// bytes reached the server, so replaying the request cannot duplicate a
|
||||
// billable completion. Only narrow, unambiguous pre-connect failures qualify: a
|
||||
// recoverable DNS lookup failure, a refused/unreachable dial, and a TLS
|
||||
// handshake timeout (the handshake completes before any HTTP request bytes are
|
||||
// written).
|
||||
//
|
||||
// The refused/unreachable errnos (ECONNREFUSED/ENETUNREACH/EHOSTUNREACH) are
|
||||
// matched by syscall errno via errors.Is — portable across platforms, since Go
|
||||
// maps the WSA* codes on Windows too — but ONLY when carried by an Op=="dial"
|
||||
// *net.OpError. errors.Is matches an errno anywhere in the chain regardless of
|
||||
// which socket operation raised it, and the kernel raises these same errnos on
|
||||
// an established connection when a route drops mid-generation; scoping to the
|
||||
// dial phase is what makes "no request bytes were sent" structural rather than
|
||||
// argued, so a post-send read/write can never replay a non-idempotent POST.
|
||||
// NXDOMAIN is treated as permanent (deterministic) and is NOT retried.
|
||||
//
|
||||
// Ambiguous or post-send failures are deliberately excluded and checked FIRST,
|
||||
// so an error that is both is never treated as pre-send: a reset or broken pipe
|
||||
// can follow a request that was already sent, and a bare "i/o timeout" covers
|
||||
// read timeouts after send as well as dial timeouts, so none of these prove the
|
||||
// request was not received and a non-idempotent POST must not be replayed on
|
||||
// them.
|
||||
func isPreSendTransportError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
// DNS resolution precedes any connection, so a lookup failure proves nothing
|
||||
// was sent — but only retry a lookup that could plausibly succeed on a replay.
|
||||
// NXDOMAIN (IsNotFound) is authoritative and deterministic (a mistyped or
|
||||
// non-existent host), so retrying it just stalls the agent through the whole
|
||||
// backoff schedule before failing anyway; a timeout or temporary SERVFAIL is
|
||||
// worth another try. errors.As unwraps url.Error/OpError to reach the DNSError.
|
||||
var dnsErr *net.DNSError
|
||||
if errors.As(err, &dnsErr) {
|
||||
return !dnsErr.IsNotFound
|
||||
}
|
||||
// Post-send / ambiguous failures, excluded first. EOF and reset are matched
|
||||
// by identity so wording (or a hostname containing "eof") can't fool them.
|
||||
if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) || isConnResetErrno(err) {
|
||||
return false
|
||||
}
|
||||
// A CONNECT-phase timeout is provably pre-send for the same structural reason
|
||||
// a refused dial is: the connection was never established, so no request
|
||||
// bytes left this host. It has to be admitted BEFORE the wording exclusions
|
||||
// below, because net surfaces it with "i/o timeout" in the message, which the
|
||||
// blanket exclusion would otherwise reject. The gate is what keeps that safe:
|
||||
// a post-send read timeout is also Timeout(), but carries Op "read", so it
|
||||
// never reaches here.
|
||||
if opErr := connectOpError(err); opErr != nil && opErr.Timeout() {
|
||||
return true
|
||||
}
|
||||
msg := strings.ToLower(err.Error())
|
||||
if strings.Contains(msg, "connection reset") ||
|
||||
strings.Contains(msg, "broken pipe") ||
|
||||
strings.Contains(msg, "i/o timeout") {
|
||||
return false
|
||||
}
|
||||
// A refused/unreachable connection is provably pre-send ONLY when it was
|
||||
// raised during the DIAL (connection-establishment) phase. The same errnos
|
||||
// are also raised by the kernel on an ESTABLISHED connection — when a route
|
||||
// drops or an ICMP unreachable arrives mid-generation — which is post-send.
|
||||
// errors.Is matches an errno anywhere in the chain regardless of the
|
||||
// operation, so we first require an Op=="dial" *net.OpError: that makes "no
|
||||
// request bytes left this host" structural, so a post-send read/write carrying
|
||||
// one of these can never replay a non-idempotent POST. dialPreSendErrnos is
|
||||
// platform-specific: on Windows a refused dial carries the raw WSA* errno
|
||||
// (WSAECONNREFUSED etc.), which does NOT satisfy errors.Is against the POSIX
|
||||
// syscall.ECONNREFUSED, so the Windows list adds those codes. The string
|
||||
// markers only catch a POSIX dial error already flattened past its errno.
|
||||
if connectOpError(err) != nil {
|
||||
for _, errno := range dialPreSendErrnos {
|
||||
if errors.Is(err, errno) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case strings.Contains(msg, "connection refused"),
|
||||
strings.Contains(msg, "network is unreachable"),
|
||||
strings.Contains(msg, "no route to host"):
|
||||
return true
|
||||
}
|
||||
}
|
||||
// A TLS handshake times out after the TCP dial connects but before any HTTP
|
||||
// request bytes are written, so it is structurally pre-send. net/http surfaces
|
||||
// it as its own error string, not a *net.OpError, so it is matched by wording
|
||||
// independent of the dial gate.
|
||||
if strings.Contains(msg, "tls handshake timeout") {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// connectOpError returns the outermost *net.OpError in err's chain when it comes
|
||||
// from the connection-establishment phase, else nil. Go tags a direct dial with
|
||||
// Op "dial" and a failed CONNECT through an HTTP proxy with Op "proxyconnect";
|
||||
// a failure on an already-established connection carries Op "read"/"write".
|
||||
// This is how isPreSendTransportError scopes the pre-send match to the connect
|
||||
// phase so a post-send failure can't be misclassified as pre-send.
|
||||
//
|
||||
// "proxyconnect" belongs here for the same reason "dial" does: when HTTPS_PROXY
|
||||
// is set and the CONNECT fails, the tunnel to the model host was never built, so
|
||||
// no request bytes reached it. Omitting it meant every pre-send failure behind a
|
||||
// corporate proxy fell through to no-retry, which is the environment where
|
||||
// transient connect failures are most common.
|
||||
func connectOpError(err error) *net.OpError {
|
||||
var opErr *net.OpError
|
||||
if errors.As(err, &opErr) && (opErr.Op == "dial" || opErr.Op == "proxyconnect") {
|
||||
return opErr
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ShouldRetryStatus reports whether an HTTP status is safe to retry for a
|
||||
// non-idempotent completion POST: 429 (Too Many Requests), 503 (Service
|
||||
// Unavailable), and 529 (Anthropic's "overloaded"). All mean the server
|
||||
@@ -137,6 +368,39 @@ func Backoff(ctx context.Context, attempt int, retryAfter time.Duration) bool {
|
||||
}
|
||||
}
|
||||
|
||||
// preSendBackoff waits before pre-send retry attempt N (1-based) on a short
|
||||
// sub-second schedule (500ms, 1s, ...), returning false if the context is
|
||||
// cancelled during the wait. Separate from Backoff's seconds-long rate-limit
|
||||
// schedule because a pre-send dial failure recovers in milliseconds or never.
|
||||
func preSendBackoff(ctx context.Context, attempt int) bool {
|
||||
timer := time.NewTimer(preSendBackoffWait(attempt))
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return false
|
||||
case <-timer.C:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// preSendBackoffWait computes the pre-send wait before retry attempt N (1-based):
|
||||
// exponential from preSendBackoffBase (500ms, 1s, 2s, ...), capped at maxBackoff.
|
||||
// The exponent is clamped so a large attempt count cannot overflow.
|
||||
func preSendBackoffWait(attempt int) time.Duration {
|
||||
exponent := attempt - 1
|
||||
if exponent > 5 {
|
||||
exponent = 5
|
||||
}
|
||||
if exponent < 0 {
|
||||
exponent = 0
|
||||
}
|
||||
wait := preSendBackoffBase * time.Duration(1<<exponent)
|
||||
if wait > maxBackoff {
|
||||
wait = maxBackoff
|
||||
}
|
||||
return wait
|
||||
}
|
||||
|
||||
// backoffWait computes the wait before retry attempt N (1-based): Retry-After
|
||||
// when supplied, else exponential from retryBackoffBase, both capped at
|
||||
// maxBackoff. The exponent is clamped so a large attempt count cannot overflow.
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
//go:build !windows && !plan9
|
||||
|
||||
package providerio
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
// dialPreSendErrnos are the syscall errnos a refused or unreachable DIAL raises
|
||||
// on this platform. On POSIX the standard syscall constants are exactly what the
|
||||
// net package surfaces, so matching them via errors.Is is sufficient. Plan 9 is
|
||||
// excluded because it does not define these BSD-style errno constants (it uses
|
||||
// string errors), so the constants below would not compile there; Zero targets
|
||||
// linux/darwin/windows, so no Plan 9 variant is needed.
|
||||
var dialPreSendErrnos = []error{
|
||||
syscall.ECONNREFUSED,
|
||||
syscall.ENETUNREACH,
|
||||
syscall.EHOSTUNREACH,
|
||||
// A connect timeout is normally caught by the Timeout() check on the connect
|
||||
// gate; this covers the case where the errno survives but the timeout flag
|
||||
// does not, e.g. an error rebuilt from the errno alone.
|
||||
syscall.ETIMEDOUT,
|
||||
}
|
||||
|
||||
// isConnResetErrno reports whether err carries a connection-reset errno, the
|
||||
// post-send exclusion in isPreSendTransportError. It lives here (per platform)
|
||||
// so retry.go references no syscall constants and stays buildable on Plan 9,
|
||||
// which defines none of them.
|
||||
func isConnResetErrno(err error) bool {
|
||||
return errors.Is(err, syscall.ECONNRESET)
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
//go:build plan9
|
||||
|
||||
package providerio
|
||||
|
||||
// Plan 9 has no BSD-style errno constants (it uses string errors), so the
|
||||
// errno-based classification is unavailable here: isConnResetErrno returns false
|
||||
// and dialPreSendErrnos is empty. retry.go still classifies via its string
|
||||
// markers on this platform, the pre-send inclusions (connection refused /
|
||||
// network unreachable / no route to host) and the separate post-send
|
||||
// "connection reset" exclusion each keep working through the wording checks.
|
||||
// Zero targets linux/darwin/windows, so this variant only keeps the package
|
||||
// building on Plan 9 rather than adding real support.
|
||||
var dialPreSendErrnos = []error{}
|
||||
|
||||
func isConnResetErrno(error) bool { return false }
|
||||
@@ -0,0 +1,40 @@
|
||||
//go:build windows
|
||||
|
||||
package providerio
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"syscall"
|
||||
|
||||
"golang.org/x/sys/windows"
|
||||
)
|
||||
|
||||
// dialPreSendErrnos are the errnos a refused or unreachable DIAL raises on
|
||||
// Windows. A real Windows dial failure carries the raw Winsock error
|
||||
// (WSAECONNREFUSED = 10061, WSAENETUNREACH = 10051, WSAEHOSTUNREACH = 10065), and
|
||||
// errors.Is against the POSIX syscall.ECONNREFUSED returns FALSE there because
|
||||
// Go's Windows syscall.ECONNREFUSED is a distinct value the net package never
|
||||
// produces. So the WSA codes must be matched explicitly, or a refused dial would
|
||||
// silently never be retried on Windows. The POSIX constants stay in the list too
|
||||
// so an error built with them (e.g. a test fixture) still classifies.
|
||||
var dialPreSendErrnos = []error{
|
||||
syscall.ECONNREFUSED,
|
||||
syscall.ENETUNREACH,
|
||||
syscall.EHOSTUNREACH,
|
||||
syscall.ETIMEDOUT,
|
||||
windows.WSAECONNREFUSED,
|
||||
windows.WSAENETUNREACH,
|
||||
windows.WSAEHOSTUNREACH,
|
||||
// Same reasoning as the POSIX ETIMEDOUT entry, with the Winsock code a real
|
||||
// Windows connect timeout actually carries.
|
||||
windows.WSAETIMEDOUT,
|
||||
}
|
||||
|
||||
// isConnResetErrno reports whether err carries a connection-reset errno, the
|
||||
// post-send exclusion in isPreSendTransportError. It matches the same
|
||||
// syscall.ECONNRESET the shared path used before the errno classification was
|
||||
// split out per platform, so behavior on Windows is unchanged; it lives here
|
||||
// so retry.go references no syscall constants and stays buildable on Plan 9.
|
||||
func isConnResetErrno(err error) bool {
|
||||
return errors.Is(err, syscall.ECONNRESET)
|
||||
}
|
||||
@@ -1,15 +1,34 @@
|
||||
//go:build !plan9
|
||||
|
||||
package providerio
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// wrapDialErrno builds the error shape a dial failure has (a *net.OpError
|
||||
// wrapping an *os.SyscallError wrapping the errno) so errors.Is reaches the
|
||||
// errno exactly as in production. These fixtures use the POSIX syscall.Exxx
|
||||
// constants, which do NOT reproduce a real Windows dial: that carries distinct
|
||||
// windows.WSA* errnos (Go's syscall.ECONNREFUSED is a value the net package
|
||||
// never produces on Windows). They still classify on Windows only because
|
||||
// dialPreSendErrnos there also lists the POSIX constants for exactly this
|
||||
// fixture reason; the real Windows dial branch is covered separately by
|
||||
// TestIsPreSendTransportErrorRealRefusedDial, which makes an actual refused dial.
|
||||
func wrapDialErrno(op string, errno syscall.Errno) error {
|
||||
return &net.OpError{Op: op, Net: "tcp", Err: os.NewSyscallError("connectex", errno)}
|
||||
}
|
||||
|
||||
type roundTripperFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
|
||||
@@ -35,6 +54,152 @@ func TestSendWithRetryDoesNotReplayTransportErrors(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// A PROVABLY pre-send transport failure (no request bytes left this host) is
|
||||
// safe to replay and must be retried, bounded by preSendMaxAttempts (its own
|
||||
// short schedule), unlike an ambiguous post-send failure. Real dial failures
|
||||
// arrive as an Op=="dial" *net.OpError, so the errno cases are the production
|
||||
// shape on every platform (Windows included); the string case exercises the
|
||||
// wording fallback for a dial error already flattened past its errno.
|
||||
func TestSendWithRetryReplaysProvablyPreSendErrors(t *testing.T) {
|
||||
shrinkBackoff(t)
|
||||
cases := map[string]error{
|
||||
"errno refused (dial)": wrapDialErrno("dial", syscall.ECONNREFUSED),
|
||||
"errno network unreachable": wrapDialErrno("dial", syscall.ENETUNREACH),
|
||||
"errno host unreachable": wrapDialErrno("dial", syscall.EHOSTUNREACH),
|
||||
"dial string fallback": &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connect: connection refused")},
|
||||
"tls handshake timeout": errors.New("net/http: TLS handshake timeout"),
|
||||
"dns timeout": &net.DNSError{Err: "server misbehaving", Name: "nope.invalid", IsTimeout: true},
|
||||
}
|
||||
for name, transportErr := range cases {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
var calls int32
|
||||
client := &http.Client{Transport: roundTripperFunc(func(*http.Request) (*http.Response, error) {
|
||||
atomic.AddInt32(&calls, 1)
|
||||
return nil, transportErr
|
||||
})}
|
||||
// maxAttempts=6 (the default) proves the pre-send path is bounded by
|
||||
// preSendMaxAttempts, not the caller's 429-tuned maxAttempts.
|
||||
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "http://example.invalid", []byte("{}"), nil, 6)
|
||||
if resp != nil {
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
if err == nil {
|
||||
t.Fatal("expected the transport error to surface after retries are exhausted")
|
||||
}
|
||||
if got := atomic.LoadInt32(&calls); got != preSendMaxAttempts {
|
||||
t.Fatalf("pre-send error tried %d times, want %d (preSendMaxAttempts)", got, preSendMaxAttempts)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// An ambiguous transport failure that could have followed a sent request must
|
||||
// NOT be replayed: a non-idempotent completion POST could duplicate billable
|
||||
// work. This is the safety line the fix must not cross.
|
||||
func TestSendWithRetryDoesNotReplayAmbiguousTransportErrors(t *testing.T) {
|
||||
for name, transportErr := range map[string]error{
|
||||
"generic i/o timeout": errors.New("dial tcp 1.2.3.4:443: i/o timeout"),
|
||||
"broken pipe": errors.New("write tcp: broken pipe"),
|
||||
"unexpected eof": io.ErrUnexpectedEOF,
|
||||
// The pre-send errnos are also raised on an ESTABLISHED connection: a route
|
||||
// dropping mid-generation surfaces on the pending read/write, which is
|
||||
// post-send. Scoping to Op=="dial" must keep these from replaying the POST.
|
||||
"host unreachable on read is post-send": wrapDialErrno("read", syscall.EHOSTUNREACH),
|
||||
"net unreachable on write is post-send": wrapDialErrno("write", syscall.ENETUNREACH),
|
||||
// NXDOMAIN is authoritative and deterministic, so retrying it would only
|
||||
// stall the agent before failing anyway.
|
||||
"dns nxdomain": &net.DNSError{Err: "no such host", Name: "nope.invalid", IsNotFound: true},
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
var calls int32
|
||||
client := &http.Client{Transport: roundTripperFunc(func(*http.Request) (*http.Response, error) {
|
||||
atomic.AddInt32(&calls, 1)
|
||||
return nil, transportErr
|
||||
})}
|
||||
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "http://example.invalid", []byte("{}"), nil, 3)
|
||||
if resp != nil {
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
if err == nil {
|
||||
t.Fatal("expected the transport error to surface")
|
||||
}
|
||||
if got := atomic.LoadInt32(&calls); got != 1 {
|
||||
t.Fatalf("ambiguous transport error replayed %d times, want 1 (no retry)", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// A REAL refused dial (nothing listening on the target) must classify as
|
||||
// pre-send on every platform. This is the regression the errno-constant fixtures
|
||||
// miss: on Windows the kernel raises WSAECONNREFUSED, which errors.Is does NOT
|
||||
// match against syscall.ECONNREFUSED, so before dialPreSendErrnos carried the WSA
|
||||
// codes a refused dial returned false here and Windows never retried it. Using a
|
||||
// real Dial exercises the platform's true error shape.
|
||||
func TestIsPreSendTransportErrorRealRefusedDial(t *testing.T) {
|
||||
// Reserve an ephemeral port, then close it, so a dial there is refused.
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("listen: %v", err)
|
||||
}
|
||||
addr := ln.Addr().String()
|
||||
if cerr := ln.Close(); cerr != nil {
|
||||
t.Fatalf("close listener: %v", cerr)
|
||||
}
|
||||
conn, dialErr := (&net.Dialer{Timeout: 2 * time.Second}).Dial("tcp", addr)
|
||||
if dialErr == nil {
|
||||
_ = conn.Close()
|
||||
t.Skip("expected a refused dial but the port accepted a connection")
|
||||
}
|
||||
if !isPreSendTransportError(dialErr) {
|
||||
t.Fatalf("real refused dial not classified pre-send: %v", dialErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsPreSendTransportError(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
err error
|
||||
want bool
|
||||
}{
|
||||
{"nil", nil, false},
|
||||
// DNS: a lookup that could recover is pre-send; NXDOMAIN is permanent.
|
||||
{"dns timeout", &net.DNSError{Err: "server misbehaving", Name: "x.invalid", IsTimeout: true}, true},
|
||||
{"dns temporary", &net.DNSError{Err: "server misbehaving", Name: "x.invalid", IsTemporary: true}, true},
|
||||
{"dns nxdomain is permanent", &net.DNSError{Err: "no such host", Name: "x.invalid", IsNotFound: true}, false},
|
||||
{"tls handshake timeout", errors.New("net/http: TLS handshake timeout"), true},
|
||||
// Errno-wrapped DIAL failures: how a real refused/unreachable dial arrives
|
||||
// on EVERY platform, including Windows where the wording differs entirely.
|
||||
{"errno refused (dial)", wrapDialErrno("dial", syscall.ECONNREFUSED), true},
|
||||
{"errno network unreachable (dial)", wrapDialErrno("dial", syscall.ENETUNREACH), true},
|
||||
{"errno host unreachable (dial)", wrapDialErrno("dial", syscall.EHOSTUNREACH), true},
|
||||
{"dial operror string fallback", &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connect: connection refused")}, true},
|
||||
// The SAME errnos raised on a post-send read/write are NOT pre-send — the
|
||||
// kernel reports them on an established connection when a route drops.
|
||||
{"host unreachable on read is post-send", wrapDialErrno("read", syscall.EHOSTUNREACH), false},
|
||||
{"net unreachable on write is post-send", wrapDialErrno("write", syscall.ENETUNREACH), false},
|
||||
{"errno reset is post-send", wrapDialErrno("read", syscall.ECONNRESET), false},
|
||||
// A refused/unreachable error already flattened past its *net.OpError can't
|
||||
// be proven pre-send, so it is NOT retried (conservative direction).
|
||||
{"flattened refused, no opError", errors.New("dial tcp 127.0.0.1:1: connect: connection refused"), false},
|
||||
{"connection reset", errors.New("read tcp: connection reset by peer"), false},
|
||||
{"broken pipe", errors.New("write tcp: broken pipe"), false},
|
||||
{"unexpected eof", io.ErrUnexpectedEOF, false},
|
||||
{"eof", io.EOF, false},
|
||||
{"generic io timeout", errors.New("dial tcp: i/o timeout"), false},
|
||||
{"context deadline", context.DeadlineExceeded, false},
|
||||
// Exclusion is checked before inclusion, even for a dial OpError: an "i/o
|
||||
// timeout" in the message wins over the refused wording.
|
||||
{"exclusion wins over inclusion", &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused / i/o timeout")}, false},
|
||||
{"unrelated", errors.New("some other error"), false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := isPreSendTransportError(c.err); got != c.want {
|
||||
t.Errorf("%s: isPreSendTransportError(%v) = %v, want %v", c.name, c.err, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestShouldRetryStatus(t *testing.T) {
|
||||
cases := map[int]bool{
|
||||
http.StatusOK: false,
|
||||
@@ -94,12 +259,17 @@ func TestBackoffWaitsThenReturnsTrue(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// shrinkBackoff makes retry waits negligible for the duration of a test.
|
||||
// shrinkBackoff makes both retry schedules (429 and pre-send) negligible for the
|
||||
// duration of a test.
|
||||
func shrinkBackoff(t *testing.T) {
|
||||
t.Helper()
|
||||
saved := retryBackoffBase
|
||||
savedRetry, savedPreSend := retryBackoffBase, preSendBackoffBase
|
||||
retryBackoffBase = time.Millisecond
|
||||
t.Cleanup(func() { retryBackoffBase = saved })
|
||||
preSendBackoffBase = time.Millisecond
|
||||
t.Cleanup(func() {
|
||||
retryBackoffBase = savedRetry
|
||||
preSendBackoffBase = savedPreSend
|
||||
})
|
||||
}
|
||||
|
||||
func TestBackoffWaitSchedule(t *testing.T) {
|
||||
@@ -126,6 +296,30 @@ func TestBackoffWaitSchedule(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// The pre-send schedule is sub-second and doubles, far shorter than the 429
|
||||
// schedule: a permanent dial failure fails in ~1.5s across preSendMaxAttempts
|
||||
// (500ms + 1s) instead of stalling the agent ~60s on 2/4/8/16/30s. This is the
|
||||
// second half of the fix for a mistyped host / dead local daemon hanging a turn.
|
||||
func TestPreSendBackoffWaitSchedule(t *testing.T) {
|
||||
cases := []struct {
|
||||
attempt int
|
||||
want time.Duration
|
||||
}{
|
||||
{1, 500 * time.Millisecond},
|
||||
{2, 1 * time.Second},
|
||||
{3, 2 * time.Second},
|
||||
// The exponent clamps at 5 (500ms<<5 = 16s) so a large attempt count can't
|
||||
// overflow; in practice preSendMaxAttempts caps retries at attempt 2, so
|
||||
// only the 500ms/1s rungs are ever reached.
|
||||
{50, 16 * time.Second},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := preSendBackoffWait(c.attempt); got != c.want {
|
||||
t.Errorf("preSendBackoffWait(%d) = %v, want %v", c.attempt, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSendWithRetryRetriesThenSucceeds(t *testing.T) {
|
||||
shrinkBackoff(t)
|
||||
var hits int32
|
||||
@@ -193,3 +387,316 @@ func TestSendWithRetryReturnsLastResponseAfterMaxAttempts(t *testing.T) {
|
||||
t.Fatalf("server hit %d times, want 2 (maxAttempts)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// Redirects are still followed, so a gateway that answers the completion endpoint
|
||||
// with a 307/308 keeps working: the caller gets the final response, not the 3xx.
|
||||
func TestSendWithRetryFollowsRedirects(t *testing.T) {
|
||||
var calls int32
|
||||
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
||||
if atomic.AddInt32(&calls, 1) == 1 {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusTemporaryRedirect,
|
||||
Header: http.Header{"Location": {"https://redirect-target.invalid/v1"}},
|
||||
Body: http.NoBody,
|
||||
Request: r,
|
||||
}, nil
|
||||
}
|
||||
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Request: r}, nil
|
||||
})}
|
||||
|
||||
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://origin.invalid/v1", []byte("{}"), nil, 3)
|
||||
if err != nil {
|
||||
t.Fatalf("a redirected completion must succeed: %v", err)
|
||||
}
|
||||
_ = resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200 (the redirect should have been followed)", resp.StatusCode)
|
||||
}
|
||||
if got := atomic.LoadInt32(&calls); got != 2 {
|
||||
t.Fatalf("transport called %d times, want 2 (origin then redirect target)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// Once an attempt has entered a redirect, the original POST has already left this
|
||||
// host, so a dial failure on the redirect hop must NOT be replayed even though it
|
||||
// arrives as an Op=="dial" error that isPreSendTransportError would otherwise
|
||||
// treat as provably pre-send. Replaying it would re-bill a completion the first
|
||||
// host may already have processed.
|
||||
func TestSendWithRetryDoesNotReplayAfterRedirect(t *testing.T) {
|
||||
shrinkBackoff(t)
|
||||
var calls int32
|
||||
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
||||
if atomic.AddInt32(&calls, 1) == 1 {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusTemporaryRedirect,
|
||||
Header: http.Header{"Location": {"https://redirect-target.invalid/v1"}},
|
||||
Body: http.NoBody,
|
||||
Request: r,
|
||||
}, nil
|
||||
}
|
||||
// The dial to the redirect target fails: pre-send in shape, but post-send
|
||||
// in fact, because the origin already received the POST.
|
||||
return nil, wrapDialErrno("dial", syscall.ECONNREFUSED)
|
||||
})}
|
||||
|
||||
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://origin.invalid/v1", []byte("{}"), nil, 6)
|
||||
if resp != nil {
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
if err == nil {
|
||||
t.Fatal("expected the redirect-hop dial failure to surface as an error")
|
||||
}
|
||||
// Exactly two: the origin request and the single redirect hop. A third would
|
||||
// mean the redirected POST was replayed.
|
||||
if got := atomic.LoadInt32(&calls); got != 2 {
|
||||
t.Fatalf("redirected POST was replayed: transport called %d times, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A caller that supplies its own CheckRedirect keeps it: the retry path observes
|
||||
// redirects, it does not take the decision away from the caller.
|
||||
func TestSendWithRetryPreservesCallerRedirectPolicy(t *testing.T) {
|
||||
var policyCalls int32
|
||||
client := &http.Client{
|
||||
CheckRedirect: func(*http.Request, []*http.Request) error {
|
||||
atomic.AddInt32(&policyCalls, 1)
|
||||
return http.ErrUseLastResponse
|
||||
},
|
||||
Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusTemporaryRedirect,
|
||||
Header: http.Header{"Location": {"https://redirect-target.invalid/v1"}},
|
||||
Body: http.NoBody,
|
||||
Request: r,
|
||||
}, nil
|
||||
}),
|
||||
}
|
||||
|
||||
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://origin.invalid/v1", []byte("{}"), nil, 3)
|
||||
if err != nil {
|
||||
t.Fatalf("caller policy asked to stop at the 3xx, so it must surface: %v", err)
|
||||
}
|
||||
_ = resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusTemporaryRedirect {
|
||||
t.Fatalf("status = %d, want the 307 surfaced per the caller's policy", resp.StatusCode)
|
||||
}
|
||||
if got := atomic.LoadInt32(&policyCalls); got != 1 {
|
||||
t.Fatalf("caller CheckRedirect called %d times, want 1 (it must not be discarded)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// The pre-send retry budget is independent of the 429/503 status retries: two
|
||||
// rate-limit responses must not consume the pre-send allowance, so the first
|
||||
// refused dial after them still gets its own preSendMaxAttempts tries.
|
||||
func TestSendWithRetryPreSendBudgetSurvivesStatusRetries(t *testing.T) {
|
||||
shrinkBackoff(t)
|
||||
var calls int32
|
||||
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
||||
if atomic.AddInt32(&calls, 1) <= 2 {
|
||||
return &http.Response{StatusCode: http.StatusTooManyRequests, Body: http.NoBody, Request: r}, nil
|
||||
}
|
||||
return nil, wrapDialErrno("dial", syscall.ECONNREFUSED)
|
||||
})}
|
||||
|
||||
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://x.invalid/v1", []byte("{}"), nil, 6)
|
||||
if resp != nil {
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
if err == nil {
|
||||
t.Fatal("expected the exhausted pre-send failure to surface as an error")
|
||||
}
|
||||
// Two status responses (calls 1-2), then the pre-send failure gets its own
|
||||
// budget: preSendMaxAttempts=3 means it retries twice more before returning.
|
||||
if got := atomic.LoadInt32(&calls); got != 2+preSendMaxAttempts {
|
||||
t.Fatalf("pre-send budget eaten by status retries: transport called %d times, want %d (2 status + %d pre-send)", got, 2+preSendMaxAttempts, preSendMaxAttempts)
|
||||
}
|
||||
}
|
||||
|
||||
// The inverse sequence: pre-send retries must not consume the 429/503 budget or
|
||||
// advance its backoff rung either. Two safely replayed refused dials followed by
|
||||
// rate limiting must still leave the status path its full maxAttempts window,
|
||||
// otherwise a couple of dial blips silently shorten every later rate-limit wait.
|
||||
func TestSendWithRetryStatusBudgetSurvivesPreSendRetries(t *testing.T) {
|
||||
shrinkBackoff(t)
|
||||
const maxAttempts = 3
|
||||
var calls int32
|
||||
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
||||
if atomic.AddInt32(&calls, 1) <= 2 {
|
||||
return nil, wrapDialErrno("dial", syscall.ECONNREFUSED)
|
||||
}
|
||||
return &http.Response{StatusCode: http.StatusTooManyRequests, Body: http.NoBody, Request: r}, nil
|
||||
})}
|
||||
|
||||
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://x.invalid/v1", []byte("{}"), nil, maxAttempts)
|
||||
if err != nil {
|
||||
t.Fatalf("exhausted status retries surface the response, not an error: %v", err)
|
||||
}
|
||||
_ = resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusTooManyRequests {
|
||||
t.Fatalf("status = %d, want 429", resp.StatusCode)
|
||||
}
|
||||
// Calls 1-2 are the retried dial failures; the status path then gets its own
|
||||
// full window (maxAttempts requests) before the 429 is surfaced.
|
||||
if got := atomic.LoadInt32(&calls); got != 2+maxAttempts {
|
||||
t.Fatalf("status budget eaten by pre-send retries: transport called %d times, want %d (2 pre-send + %d status)", got, 2+maxAttempts, maxAttempts)
|
||||
}
|
||||
}
|
||||
|
||||
// A redirect followed by a RETRYABLE STATUS is still retried, and deliberately
|
||||
// so. This is the case a reviewer flagged as a possible replay of billable work,
|
||||
// and the asymmetry with the transport-error guard above is intentional rather
|
||||
// than an oversight, so it is pinned here.
|
||||
//
|
||||
// The two situations differ in what is known about the request:
|
||||
//
|
||||
// - Transport error after a redirect: the original POST reached the origin and
|
||||
// the dial to the redirect target then failed, so whether the target received
|
||||
// and began processing it is UNKNOWN. Replaying could duplicate work, which
|
||||
// is why that path refuses to retry.
|
||||
// - Retryable status after a redirect: every hop answered. A 307/308 is the
|
||||
// origin declining to process and delegating, and a 429/503/529 is the target
|
||||
// declining as well. Nothing accepted the request, so replaying duplicates
|
||||
// nothing.
|
||||
//
|
||||
// The risk here is also exactly the risk of a 429 without any redirect, which
|
||||
// this package has always retried; the redirect adds no new ambiguity. Refusing
|
||||
// to retry would instead make a rate-limited completion behind a redirecting
|
||||
// gateway fail outright rather than back off.
|
||||
func TestSendWithRetryRetriesRedirectedRetryableStatus(t *testing.T) {
|
||||
var calls int32
|
||||
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
||||
switch atomic.AddInt32(&calls, 1) {
|
||||
case 1, 3:
|
||||
// The origin declines and delegates, both times.
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusTemporaryRedirect,
|
||||
Header: http.Header{"Location": {"https://redirect-target.invalid/v1"}},
|
||||
Body: http.NoBody,
|
||||
Request: r,
|
||||
}, nil
|
||||
case 2:
|
||||
// The target declines too: rate limited, so nothing was accepted.
|
||||
return &http.Response{StatusCode: http.StatusTooManyRequests, Body: http.NoBody, Request: r}, nil
|
||||
default:
|
||||
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Request: r}, nil
|
||||
}
|
||||
})}
|
||||
|
||||
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://origin.invalid/v1", []byte("{}"), nil, 3)
|
||||
if err != nil {
|
||||
t.Fatalf("a rate-limited completion behind a redirect must still retry: %v", err)
|
||||
}
|
||||
_ = resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200 after the retry succeeded", resp.StatusCode)
|
||||
}
|
||||
// origin, target(429), origin, target(200).
|
||||
if got := atomic.LoadInt32(&calls); got != 4 {
|
||||
t.Fatalf("transport called %d times, want 4 (redirect, 429, redirect, 200)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A CONNECT-phase timeout is provably pre-send: the connection was never
|
||||
// established, so no request bytes left this host. It reaches the classifier
|
||||
// carrying "i/o timeout", which the blanket wording exclusion rejects, so
|
||||
// admitting it depends on the connect gate running first.
|
||||
//
|
||||
// The post-send counterpart is the point of the table: a read timeout is also
|
||||
// Timeout(), and must stay excluded, or a reply that timed out mid-generation
|
||||
// would replay a completion the server already billed.
|
||||
func TestPreSendClassifiesConnectTimeouts(t *testing.T) {
|
||||
timeoutErr := &timeoutError{}
|
||||
for name, testCase := range map[string]struct {
|
||||
err error
|
||||
want bool
|
||||
}{
|
||||
"dial timeout": {
|
||||
err: &net.OpError{Op: "dial", Net: "tcp", Err: timeoutErr},
|
||||
want: true,
|
||||
},
|
||||
"proxy connect timeout": {
|
||||
err: &net.OpError{Op: "proxyconnect", Net: "tcp", Err: timeoutErr},
|
||||
want: true,
|
||||
},
|
||||
"read timeout stays post-send": {
|
||||
err: &net.OpError{Op: "read", Net: "tcp", Err: timeoutErr},
|
||||
want: false,
|
||||
},
|
||||
"write timeout stays post-send": {
|
||||
err: &net.OpError{Op: "write", Net: "tcp", Err: timeoutErr},
|
||||
want: false,
|
||||
},
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
if got := isPreSendTransportError(testCase.err); got != testCase.want {
|
||||
t.Fatalf("isPreSendTransportError(%v) = %v, want %v", testCase.err, got, testCase.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// A failed CONNECT through an HTTP proxy is pre-send for the same structural
|
||||
// reason a refused dial is, but net tags it Op "proxyconnect", so a gate
|
||||
// matching only "dial" left every pre-send failure behind a corporate proxy
|
||||
// unretried.
|
||||
func TestPreSendClassifiesProxyConnectRefusal(t *testing.T) {
|
||||
err := &net.OpError{Op: "proxyconnect", Net: "tcp", Err: syscall.ECONNREFUSED}
|
||||
if !isPreSendTransportError(err) {
|
||||
t.Fatal("a refused proxy CONNECT was not classified pre-send; no request bytes reached the model host")
|
||||
}
|
||||
// The op still has to be a connect phase. The same errno on an established
|
||||
// connection is post-send and must not replay a POST.
|
||||
post := &net.OpError{Op: "read", Net: "tcp", Err: syscall.ECONNREFUSED}
|
||||
if isPreSendTransportError(post) {
|
||||
t.Fatal("a read-phase error was classified pre-send")
|
||||
}
|
||||
}
|
||||
|
||||
// timeoutError is a net.Error whose Timeout() is true, matching the shape
|
||||
// net/http surfaces for a connect deadline.
|
||||
type timeoutError struct{}
|
||||
|
||||
func (timeoutError) Error() string { return "i/o timeout" }
|
||||
func (timeoutError) Timeout() bool { return true }
|
||||
func (timeoutError) Temporary() bool { return true }
|
||||
|
||||
// The pre-send budget bounds CONSECUTIVE connect failures, not the lifetime of
|
||||
// the call. Without a reset on a successful hop, an early run of dial blips
|
||||
// permanently spends the budget, so a later blip surfaces as a transport error
|
||||
// while the status budget still has room. That is the coupling the two counters
|
||||
// were split apart to remove, in the one ordering the earlier tests did not
|
||||
// cover: pre-send, then status, then pre-send again.
|
||||
func TestPreSendBudgetResetsAfterSuccessfulHop(t *testing.T) {
|
||||
refused := &net.OpError{Op: "dial", Net: "tcp", Err: syscall.ECONNREFUSED}
|
||||
var calls int32
|
||||
client := &http.Client{Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) {
|
||||
switch atomic.AddInt32(&calls, 1) {
|
||||
case 1, 2:
|
||||
// Two connect failures, spending the pre-send budget under the old
|
||||
// lifetime accounting.
|
||||
return nil, refused
|
||||
case 3:
|
||||
// A hop that reaches the server. It declines, so the request was not
|
||||
// accepted and retrying stays safe.
|
||||
return &http.Response{StatusCode: http.StatusTooManyRequests, Body: http.NoBody, Request: r}, nil
|
||||
case 4:
|
||||
// One more connect blip. With a lifetime cap this surfaces as an
|
||||
// error; the call should still recover.
|
||||
return nil, refused
|
||||
default:
|
||||
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Request: r}, nil
|
||||
}
|
||||
})}
|
||||
|
||||
resp, err := SendWithRetry(context.Background(), client, http.MethodPost, "https://origin.invalid/v1", []byte("{}"), nil, 3)
|
||||
if err != nil {
|
||||
t.Fatalf("interleaved connect and status failures returned %v; the pre-send budget did not reset after the hop that reached the server", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200", resp.StatusCode)
|
||||
}
|
||||
if got := atomic.LoadInt32(&calls); got != 5 {
|
||||
t.Fatalf("transport called %d times, want 5 (dial, dial, 429, dial, 200)", got)
|
||||
}
|
||||
}
|
||||
@@ -27,11 +27,22 @@ var destructivePrograms = map[string]bool{
|
||||
"mkfs": true, "fdisk": true, "shred": true, "dd": true, "parted": true,
|
||||
}
|
||||
|
||||
var powerShellRemoveItemPrograms = map[string]bool{
|
||||
"remove-item": true,
|
||||
"ri": true,
|
||||
"rd": true,
|
||||
"rmdir": true,
|
||||
"del": true,
|
||||
"erase": true,
|
||||
"rm": true,
|
||||
}
|
||||
|
||||
// networkPrograms are commands that perform network egress/ingress.
|
||||
var networkPrograms = map[string]bool{
|
||||
"curl": true, "wget": true, "ssh": true, "scp": true, "sftp": true,
|
||||
"rsync": true, "nc": true, "ncat": true, "netcat": true, "telnet": true,
|
||||
"ftp": true,
|
||||
"ftp": true, "iwr": true, "irm": true, "invoke-webrequest": true,
|
||||
"invoke-restmethod": true,
|
||||
}
|
||||
|
||||
var localServerPrograms = map[string]bool{
|
||||
@@ -172,7 +183,10 @@ func analyzeInto(script string, result *AnalysisResult, seen map[string]bool, de
|
||||
if commandUsesNetwork(prog, rest) {
|
||||
result.Network = true
|
||||
}
|
||||
if destructivePrograms[prog] || (prog == "rm" && hasRecursiveForce(rest)) || (prog == "find" && hasFindDelete(rest)) {
|
||||
if destructivePrograms[prog] ||
|
||||
(prog == "rm" && hasRecursiveForce(rest)) ||
|
||||
(powerShellRemoveItemPrograms[prog] && hasPowerShellRecursiveForce(rest)) ||
|
||||
(prog == "find" && hasFindDelete(rest)) {
|
||||
result.Destructive = true
|
||||
}
|
||||
return true
|
||||
@@ -479,6 +493,19 @@ func hasRecursiveForce(args []*syntax.Word) bool {
|
||||
return recursive && force
|
||||
}
|
||||
|
||||
func hasPowerShellRecursiveForce(args []*syntax.Word) bool {
|
||||
recursive, force := false, false
|
||||
for _, arg := range args {
|
||||
switch strings.ToLower(wordText(arg)) {
|
||||
case "-recurse", "-r":
|
||||
recursive = true
|
||||
case "-force":
|
||||
force = true
|
||||
}
|
||||
}
|
||||
return recursive && force
|
||||
}
|
||||
|
||||
func hasFindDelete(args []*syntax.Word) bool {
|
||||
for _, arg := range args {
|
||||
if wordText(arg) == "-delete" {
|
||||
|
||||
@@ -45,6 +45,27 @@ func TestAnalyzeCommand(t *testing.T) {
|
||||
{name: "sudo dynamic flag then rm -rf", script: `sudo "$maybe" rm -rf /tmp/x`, destructive: true},
|
||||
|
||||
{name: "curl", script: "curl https://example.com", network: true},
|
||||
{name: "PowerShell iwr alias", script: "iwr https://example.com", network: true},
|
||||
{name: "PowerShell irm alias", script: "irm https://example.com", network: true},
|
||||
{name: "PowerShell Invoke-WebRequest", script: "Invoke-WebRequest https://example.com", network: true},
|
||||
{name: "PowerShell Invoke-RestMethod", script: "Invoke-RestMethod https://example.com", network: true},
|
||||
{name: "PowerShell Remove-Item recursive force", script: `Remove-Item -Recurse -Force 'C:\temp\x'`, destructive: true},
|
||||
{name: "PowerShell rm recursive force", script: `rm -Recurse -Force 'C:\temp\x'`, destructive: true},
|
||||
{name: "PowerShell ri recursive force", script: `ri -Recurse -Force 'C:\temp\x'`, destructive: true},
|
||||
{name: "PowerShell rd recursive force", script: `rd -Recurse -Force 'C:\temp\x'`, destructive: true},
|
||||
{name: "PowerShell rmdir recursive force", script: `rmdir -Recurse -Force 'C:\temp\x'`, destructive: true},
|
||||
{name: "PowerShell del recursive force", script: `del -Recurse -Force 'C:\temp\x'`, destructive: true},
|
||||
{name: "PowerShell erase recursive force", script: `erase -Recurse -Force 'C:\temp\x'`, destructive: true},
|
||||
{name: "PowerShell ambiguous force abbreviation", script: `Remove-Item -Recurse -f 'C:\temp\x'`, destructive: false},
|
||||
{name: "PowerShell Remove-Item without recurse", script: `Remove-Item -Force 'C:\temp\x'`, destructive: false},
|
||||
{name: "Windows curl cmd", script: "curl.cmd https://example.com", network: true},
|
||||
{name: "Windows curl exe", script: "curl.exe https://example.com", network: true},
|
||||
{name: "Windows drive path curl exe", script: `'C:\tools\curl.exe' https://example.com`, network: true},
|
||||
{name: "Windows UNC path curl exe", script: `'\\server\share\curl.exe' https://example.com`, network: true},
|
||||
{name: "Windows dot relative path curl exe", script: `'.\curl.exe' https://example.com`, network: true},
|
||||
{name: "Windows relative path curl exe", script: `'tools\curl.exe' https://example.com`, network: true},
|
||||
{name: "Windows drive relative path curl exe", script: `'C:curl.exe' https://example.com`, network: true},
|
||||
{name: "Windows npm cmd", script: "npm.cmd install", network: true},
|
||||
{name: "wget piped to shell", script: "wget -qO- https://x.test | sh", network: true},
|
||||
{name: "python http server", script: "python3 -m http.server 8000", network: true},
|
||||
{name: "python pip install", script: "python3 -m pip install requests", network: true},
|
||||
|
||||
@@ -74,6 +74,25 @@ func TestEngineBashAllowGrantDoesNotBypassNetworkPrompt(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestEngineClassifiesPowerShellCurlCommandAsNetwork(t *testing.T) {
|
||||
engine := NewEngine(EngineOptions{
|
||||
WorkspaceRoot: t.TempDir(),
|
||||
Policy: DefaultPolicy(),
|
||||
})
|
||||
decision := engine.Evaluate(context.Background(), Request{
|
||||
ToolName: "bash",
|
||||
SideEffect: SideEffectShell,
|
||||
Permission: PermissionPrompt,
|
||||
PermissionMode: PermissionModeAsk,
|
||||
Args: map[string]any{
|
||||
"command": `$env:PATH = '.;' + $env:PATH; curl.cmd https://example.com`,
|
||||
},
|
||||
})
|
||||
if decision.Action != ActionPrompt || decision.Reason != ReasonNetworkBlocked || !HasRiskCategory(decision.Risk, "network") {
|
||||
t.Fatalf("PowerShell curl command decision = %#v, want network prompt", decision)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnsandboxedExecutionAllowedPreservesDeniedReads(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
policy := DefaultPolicy()
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
@@ -378,7 +379,31 @@ func TestCredentialDenyReadPathsIn(t *testing.T) {
|
||||
home := t.TempDir()
|
||||
awsDir := filepath.Join(home, ".aws")
|
||||
gcloudDir := filepath.Join(home, ".config", "gcloud")
|
||||
if err := mkdirAll(awsDir, gcloudDir); err != nil {
|
||||
npmrc := filepath.Join(home, ".npmrc")
|
||||
ghHosts := filepath.Join(home, ".config", "gh", "hosts.yml")
|
||||
netrc := filepath.Join(home, ".netrc")
|
||||
dockerConfig := filepath.Join(home, ".docker", "config.json")
|
||||
kubeConfig := filepath.Join(home, ".kube", "config")
|
||||
zeroConfig := filepath.Join(home, ".config", "zero")
|
||||
if err := mkdirAll(
|
||||
awsDir,
|
||||
gcloudDir,
|
||||
filepath.Dir(ghHosts),
|
||||
filepath.Dir(dockerConfig),
|
||||
filepath.Dir(kubeConfig),
|
||||
filepath.Dir(zeroConfig),
|
||||
); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, path := range []string{npmrc, ghHosts, netrc, dockerConfig, kubeConfig, zeroConfig} {
|
||||
if path == zeroConfig {
|
||||
continue
|
||||
}
|
||||
if err := os.WriteFile(path, []byte("{}"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := os.MkdirAll(zeroConfig, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
keyFile := filepath.Join(home, "sa-key.json")
|
||||
@@ -387,7 +412,17 @@ func TestCredentialDenyReadPathsIn(t *testing.T) {
|
||||
}
|
||||
|
||||
paths := credentialDenyReadPathsIn(home, keyFile, nil)
|
||||
for _, want := range normalizeProfilePaths([]string{awsDir, gcloudDir, keyFile}) {
|
||||
for _, want := range normalizeProfilePaths([]string{
|
||||
awsDir,
|
||||
gcloudDir,
|
||||
keyFile,
|
||||
npmrc,
|
||||
ghHosts,
|
||||
netrc,
|
||||
dockerConfig,
|
||||
kubeConfig,
|
||||
zeroConfig,
|
||||
}) {
|
||||
if !stringSliceContains(paths, want) {
|
||||
t.Errorf("credential deny paths = %#v, want %q included", paths, want)
|
||||
}
|
||||
@@ -418,3 +453,58 @@ func TestCredentialDenyReadPathsIn(t *testing.T) {
|
||||
t.Errorf("credential deny paths without home = %#v, want key file included", homeless)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCredentialDenyReadPathsForEnvironmentHonorsConfigOverrides(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
configHome := filepath.Join(root, "xdg")
|
||||
npmrc := filepath.Join(root, "npm", "userconfig")
|
||||
ghConfigDir := filepath.Join(root, "gh")
|
||||
cloudSDKConfig := filepath.Join(root, "gcloud")
|
||||
netrc := filepath.Join(root, "netrc")
|
||||
dockerConfigDir := filepath.Join(root, "docker")
|
||||
kubeConfigA := filepath.Join(root, "kube", "a")
|
||||
kubeConfigB := filepath.Join(root, "kube", "b")
|
||||
zeroConfig := filepath.Join(configHome, "zero")
|
||||
filePaths := []string{
|
||||
npmrc,
|
||||
filepath.Join(ghConfigDir, "hosts.yml"),
|
||||
netrc,
|
||||
filepath.Join(dockerConfigDir, "config.json"),
|
||||
kubeConfigA,
|
||||
kubeConfigB,
|
||||
zeroConfig,
|
||||
}
|
||||
if err := os.MkdirAll(cloudSDKConfig, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, path := range filePaths {
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if path == zeroConfig {
|
||||
if err := os.MkdirAll(path, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if err := os.WriteFile(path, []byte("{}"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
got := credentialDenyReadPathsForEnvironment(credentialPathEnvironment{
|
||||
ConfigHome: configHome,
|
||||
CloudSDKConfig: cloudSDKConfig,
|
||||
NPMUserConfig: npmrc,
|
||||
GHConfigDir: ghConfigDir,
|
||||
Netrc: netrc,
|
||||
DockerConfigDir: dockerConfigDir,
|
||||
KubeConfig: strings.Join([]string{kubeConfigA, kubeConfigB}, string(filepath.ListSeparator)),
|
||||
}, nil)
|
||||
wantPaths := append(filePaths, cloudSDKConfig)
|
||||
for _, want := range normalizeProfilePaths(wantPaths) {
|
||||
if !stringSliceContains(got, want) {
|
||||
t.Errorf("credential deny paths = %#v, want override %q included", got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
+105
-12
@@ -145,9 +145,10 @@ func permissionProfileReadRoots(workspaceRoot string, policy Policy, scope *Scop
|
||||
}
|
||||
|
||||
// credentialDenyReadPaths returns default deny-read entries for well-known
|
||||
// credential stores (~/.aws, ~/.config/gcloud, ~/.azure, and the file
|
||||
// GOOGLE_APPLICATION_CREDENTIALS points to) so sandboxed commands cannot read
|
||||
// cloud secrets under the read-all workspace posture. Two deliberate limits:
|
||||
// credential stores so sandboxed commands cannot read secrets under the
|
||||
// read-all workspace posture. This includes tool configuration files that can
|
||||
// carry credentials and are now discoverable through the preserved caller
|
||||
// environment. Two deliberate limits:
|
||||
//
|
||||
// - Windows is skipped: a non-empty profile DenyRead switches the Windows
|
||||
// runner onto the capability-SID/ACL deny path and away from the
|
||||
@@ -166,23 +167,114 @@ func credentialDenyReadPaths(policy Policy) []string {
|
||||
// A failed home lookup only drops the home-based candidates; the
|
||||
// GOOGLE_APPLICATION_CREDENTIALS target must be protected regardless.
|
||||
home, _ := os.UserHomeDir()
|
||||
return credentialDenyReadPathsIn(home, os.Getenv("GOOGLE_APPLICATION_CREDENTIALS"), policy.AllowRead)
|
||||
return credentialDenyReadPathsForEnvironment(credentialPathEnvironment{
|
||||
Home: home,
|
||||
ConfigHome: os.Getenv("XDG_CONFIG_HOME"),
|
||||
CloudSDKConfig: os.Getenv("CLOUDSDK_CONFIG"),
|
||||
GoogleCredentials: os.Getenv("GOOGLE_APPLICATION_CREDENTIALS"),
|
||||
NPMUserConfig: firstNonEmpty(os.Getenv("NPM_CONFIG_USERCONFIG"), os.Getenv("npm_config_userconfig")),
|
||||
GHConfigDir: os.Getenv("GH_CONFIG_DIR"),
|
||||
Netrc: os.Getenv("NETRC"),
|
||||
DockerConfigDir: os.Getenv("DOCKER_CONFIG"),
|
||||
KubeConfig: os.Getenv("KUBECONFIG"),
|
||||
}, policy.AllowRead)
|
||||
}
|
||||
|
||||
// credentialDenyReadPathsIn is the pure core of credentialDenyReadPaths,
|
||||
// separated so tests can exercise it against a synthetic home directory.
|
||||
func credentialDenyReadPathsIn(home string, googleCredentials string, allowRead []string) []string {
|
||||
return credentialDenyReadPathsForEnvironment(credentialPathEnvironment{
|
||||
Home: home,
|
||||
GoogleCredentials: googleCredentials,
|
||||
}, allowRead)
|
||||
}
|
||||
|
||||
type credentialPathEnvironment struct {
|
||||
Home string
|
||||
ConfigHome string
|
||||
CloudSDKConfig string
|
||||
GoogleCredentials string
|
||||
NPMUserConfig string
|
||||
GHConfigDir string
|
||||
Netrc string
|
||||
DockerConfigDir string
|
||||
KubeConfig string
|
||||
}
|
||||
|
||||
func credentialDenyReadPathsForEnvironment(env credentialPathEnvironment, allowRead []string) []string {
|
||||
var candidates []string
|
||||
if home = strings.TrimSpace(home); home != "" {
|
||||
home := strings.TrimSpace(env.Home)
|
||||
configHome := strings.TrimSpace(env.ConfigHome)
|
||||
if configHome == "" && home != "" {
|
||||
configHome = filepath.Join(home, ".config")
|
||||
}
|
||||
if home != "" {
|
||||
candidates = append(candidates,
|
||||
filepath.Join(home, ".aws"),
|
||||
filepath.Join(home, ".config", "gcloud"),
|
||||
filepath.Join(home, ".azure"),
|
||||
// git's credential store backend, which holds host passwords and
|
||||
// personal access tokens in cleartext. Denied rather than the whole
|
||||
// of ~/.ssh, because these cost nothing functionally: git reads them
|
||||
// through a credential helper for authentication, not for identity,
|
||||
// so a sandboxed git still works and simply cannot authenticate as
|
||||
// the user. SSH key material is a harder trade and is tracked
|
||||
// separately (#815).
|
||||
filepath.Join(home, ".git-credentials"),
|
||||
)
|
||||
}
|
||||
if target := strings.TrimSpace(googleCredentials); target != "" {
|
||||
npmUserConfig := strings.TrimSpace(env.NPMUserConfig)
|
||||
if npmUserConfig == "" && home != "" {
|
||||
npmUserConfig = filepath.Join(home, ".npmrc")
|
||||
}
|
||||
netrc := strings.TrimSpace(env.Netrc)
|
||||
if netrc == "" && home != "" {
|
||||
netrc = filepath.Join(home, ".netrc")
|
||||
}
|
||||
dockerConfigDir := strings.TrimSpace(env.DockerConfigDir)
|
||||
if dockerConfigDir == "" && home != "" {
|
||||
dockerConfigDir = filepath.Join(home, ".docker")
|
||||
}
|
||||
candidates = append(candidates, npmUserConfig, netrc)
|
||||
if dockerConfigDir != "" {
|
||||
candidates = append(candidates, filepath.Join(dockerConfigDir, "config.json"))
|
||||
}
|
||||
if configHome != "" {
|
||||
candidates = append(candidates,
|
||||
filepath.Join(configHome, "zero"),
|
||||
// The XDG location of the same git credential store. Denying the file
|
||||
// rather than the directory is deliberate: ~/.config/git also holds
|
||||
// the global git config, which userGitConfigReadPaths grants on
|
||||
// purpose so a sandboxed git can read user.name and aliases instead
|
||||
// of failing outright.
|
||||
filepath.Join(configHome, "git", "credentials"),
|
||||
)
|
||||
}
|
||||
cloudSDKConfig := strings.TrimSpace(env.CloudSDKConfig)
|
||||
if cloudSDKConfig == "" && configHome != "" {
|
||||
cloudSDKConfig = filepath.Join(configHome, "gcloud")
|
||||
}
|
||||
if cloudSDKConfig != "" {
|
||||
candidates = append(candidates, cloudSDKConfig)
|
||||
}
|
||||
ghConfigDir := strings.TrimSpace(env.GHConfigDir)
|
||||
if ghConfigDir == "" && configHome != "" {
|
||||
ghConfigDir = filepath.Join(configHome, "gh")
|
||||
}
|
||||
if ghConfigDir != "" {
|
||||
candidates = append(candidates, filepath.Join(ghConfigDir, "hosts.yml"))
|
||||
}
|
||||
if target := strings.TrimSpace(env.GoogleCredentials); target != "" {
|
||||
candidates = append(candidates, target)
|
||||
}
|
||||
kubeConfig := strings.TrimSpace(env.KubeConfig)
|
||||
if kubeConfig == "" && home != "" {
|
||||
kubeConfig = filepath.Join(home, ".kube", "config")
|
||||
}
|
||||
for _, path := range filepath.SplitList(kubeConfig) {
|
||||
if path = strings.TrimSpace(path); path != "" {
|
||||
candidates = append(candidates, path)
|
||||
}
|
||||
}
|
||||
allowRoots := normalizeProfilePaths(allowRead)
|
||||
out := make([]string, 0, len(candidates))
|
||||
for _, path := range normalizeProfilePaths(candidates) {
|
||||
@@ -206,11 +298,12 @@ func credentialDenyReadPathsIn(home string, googleCredentials string, allowRead
|
||||
|
||||
// userGitConfigReadPaths returns the user's global git config FILES so a
|
||||
// sandboxed git can read identity and config (user.name/email, aliases) instead
|
||||
// of failing with "unable to access ~/.gitconfig". It is deliberately the config
|
||||
// files only — not the ~/.config/git directory, which can hold an XDG credential
|
||||
// store — so credentials and the rest of HOME stay unreadable. Granted at the
|
||||
// macOS-seatbelt read rule (not the cross-platform PermissionProfile) so the
|
||||
// HOME-dependent paths don't leak into the platform-agnostic policy snapshot.
|
||||
// of failing with "unable to access ~/.gitconfig". On macOS, where reads are
|
||||
// allow-listed, granting only these files avoids exposing the surrounding
|
||||
// configuration directory. Linux uses a read-all profile with explicit
|
||||
// credential deny rules. The paths are granted at the macOS seatbelt read rule
|
||||
// rather than the cross-platform PermissionProfile so HOME-dependent paths
|
||||
// don't leak into the platform-agnostic policy snapshot.
|
||||
func userGitConfigReadPaths() []string {
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil || strings.TrimSpace(home) == "" {
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
package sandbox
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// git's credential store holds host passwords and personal access tokens in
|
||||
// cleartext, in one of two locations depending on whether the user is on the
|
||||
// XDG layout. Neither was denied, so a sandboxed command could read them
|
||||
// (#815).
|
||||
//
|
||||
// Scoped to the credential files on purpose. Denying ~/.ssh as well would stop
|
||||
// a sandboxed git push over SSH from working, which is a functional trade that
|
||||
// issue tracks separately; these two cost nothing, because git reads them for
|
||||
// authentication rather than identity.
|
||||
func TestCredentialDenyReadPathsCoversGitCredentialStores(t *testing.T) {
|
||||
home := t.TempDir()
|
||||
configHome := filepath.Join(home, ".config")
|
||||
if err := os.MkdirAll(filepath.Join(configHome, "git"), 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// The builder only emits paths that exist, so these have to be real.
|
||||
gitCredentials := filepath.Join(home, ".git-credentials")
|
||||
xdgCredentials := filepath.Join(configHome, "git", "credentials")
|
||||
gitConfig := filepath.Join(configHome, "git", "config")
|
||||
for _, path := range []string{gitCredentials, xdgCredentials, gitConfig} {
|
||||
if err := os.WriteFile(path, []byte("x"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
denied := credentialDenyReadPathsForEnvironment(credentialPathEnvironment{
|
||||
Home: home,
|
||||
ConfigHome: configHome,
|
||||
}, nil)
|
||||
covered := func(target string) bool {
|
||||
norm := normalizeProfilePath(target)
|
||||
for _, entry := range denied {
|
||||
if entry == norm || pathWithinRoot(entry, norm) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
if !covered(gitCredentials) {
|
||||
t.Fatalf("~/.git-credentials is readable; deny list = %v", denied)
|
||||
}
|
||||
if !covered(xdgCredentials) {
|
||||
t.Fatalf("~/.config/git/credentials is readable; deny list = %v", denied)
|
||||
}
|
||||
// The global git config must stay readable. userGitConfigReadPaths grants it
|
||||
// deliberately so a sandboxed git can read user.name and aliases rather than
|
||||
// failing with "unable to access", so denying the directory instead of the
|
||||
// credential file would break that on purpose-built setups.
|
||||
if covered(gitConfig) {
|
||||
t.Fatalf("~/.config/git/config was denied; a sandboxed git loses its identity config")
|
||||
}
|
||||
}
|
||||
|
||||
// An explicit read grant still wins, matching how every other entry behaves: a
|
||||
// user who deliberately scopes a workspace over one of these paths is not
|
||||
// overridden by the default deny.
|
||||
//
|
||||
// Asserted from both sides on purpose. Checking only that the path is absent
|
||||
// under an allowRead is satisfied just as well by a deny rule that was never
|
||||
// added, so the run without allowRead has to establish that there is something
|
||||
// there for the grant to override.
|
||||
func TestGitCredentialDenyYieldsToExplicitAllowRead(t *testing.T) {
|
||||
home := t.TempDir()
|
||||
gitCredentials := filepath.Join(home, ".git-credentials")
|
||||
if err := os.WriteFile(gitCredentials, []byte("x"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
target := normalizeProfilePath(gitCredentials)
|
||||
listed := func(entries []string) bool {
|
||||
for _, entry := range entries {
|
||||
if entry == target {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
if !listed(credentialDenyReadPathsForEnvironment(credentialPathEnvironment{Home: home}, nil)) {
|
||||
t.Fatalf("~/.git-credentials is not denied without an allowRead; nothing for the grant to override")
|
||||
}
|
||||
if listed(credentialDenyReadPathsForEnvironment(credentialPathEnvironment{Home: home}, []string{home})) {
|
||||
t.Fatalf("explicit allowRead %q did not re-include the credential store", home)
|
||||
}
|
||||
}
|
||||
@@ -1027,6 +1027,13 @@ func scrubSensitiveEnv(env []string, additionalKeys ...string) []string {
|
||||
"GH_TOKEN",
|
||||
"ZERO_WEBSEARCH_API_KEY",
|
||||
"ZERO_DAEMON_REMOTE_TOKEN",
|
||||
// The file form of the same bridge token. TokenFromEnv accepts either,
|
||||
// so scrubbing only the inline variable left the pointer readable, and
|
||||
// the default sandbox posture is read-all. There is no fallback
|
||||
// location to guess at, the path comes from this variable alone, so
|
||||
// removing it closes the leak rather than half of it. Same reasoning as
|
||||
// GOOGLE_APPLICATION_CREDENTIALS below.
|
||||
"ZERO_DAEMON_REMOTE_TOKEN_FILE",
|
||||
}
|
||||
for _, descriptor := range providercatalog.All() {
|
||||
for _, key := range descriptor.AuthEnvVars {
|
||||
|
||||
@@ -699,6 +699,11 @@ func TestScrubSensitiveEnv(t *testing.T) {
|
||||
"XAI_API_KEY=xai-12345",
|
||||
"HUGGINGFACE_API_KEY=hf_12345",
|
||||
"GOOGLE_APPLICATION_CREDENTIALS=/home/user/sa-key.json",
|
||||
// Both forms of the daemon bridge token. The inline one was already
|
||||
// scrubbed; the file form was not, which left a sandboxed command the
|
||||
// path to read it from (#677).
|
||||
"ZERO_DAEMON_REMOTE_TOKEN=bridge-token-inline",
|
||||
"ZERO_DAEMON_REMOTE_TOKEN_FILE=/home/user/.zero/remote-token",
|
||||
"COMPANY_LLM_SECRET=custom-secret",
|
||||
"ZERO_OAUTH_MY_SVC_CLIENT_SECRET=oauth-secret",
|
||||
"zero_oauth_second_client_secret=case-insensitive-secret",
|
||||
|
||||
@@ -266,6 +266,74 @@ func TestWindowsRestrictedTokenNestedPipeCapture(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestWindowsRestrictedTokenPowerShell exercises the default Windows shell
|
||||
// shape end to end under the write-restricted token. It covers PowerShell
|
||||
// initialization, a native pipeline (which creates IPC objects), UTF-8 output,
|
||||
// and workspace write enforcement.
|
||||
func TestWindowsRestrictedTokenPowerShell(t *testing.T) {
|
||||
if os.Getenv("ZERO_SANDBOX_REAL_SMOKE") != "1" {
|
||||
t.Skip("set ZERO_SANDBOX_REAL_SMOKE=1 to run real Windows sandbox smoke tests")
|
||||
}
|
||||
runnerExe := realSmokeExecutable(t, "ZERO_WINDOWS_COMMAND_RUNNER_EXE", WindowsSandboxCommandRunnerName)
|
||||
powerShell := realSmokePowerShell(t)
|
||||
|
||||
root := t.TempDir()
|
||||
outside := t.TempDir()
|
||||
sandboxHome := filepath.Join(root, ".zero-sandbox")
|
||||
profile := PermissionProfile{
|
||||
FileSystem: FileSystemPolicy{
|
||||
Kind: FileSystemRestricted,
|
||||
ReadRoots: []string{root},
|
||||
WriteRoots: []WritableRoot{{Root: root}},
|
||||
IncludePlatformRoots: true,
|
||||
AllowTemp: true,
|
||||
},
|
||||
Network: NetworkPolicy{Mode: NetworkDeny},
|
||||
}
|
||||
config := WindowsSandboxCommandArgsOptions{
|
||||
SandboxHome: sandboxHome,
|
||||
CommandCWD: root,
|
||||
WorkspaceRoots: []string{root},
|
||||
PermissionProfile: profile,
|
||||
SandboxLevel: WindowsSandboxLevelUnelevated,
|
||||
}
|
||||
|
||||
insideMarker := filepath.Join(root, "powershell-ok.txt")
|
||||
script := "$ErrorActionPreference='Stop'; " +
|
||||
"[Console]::OutputEncoding=[System.Text.Encoding]::UTF8; " +
|
||||
"$values = 1,2,3 | ForEach-Object { $_ * 2 }; " +
|
||||
"[IO.File]::WriteAllText(" + powershellSingleQuote(insideMarker) + ", (($values -join ',') + ' ✓'))"
|
||||
runWindowsRealSmokeCommand(t, runnerExe, config, []string{
|
||||
powerShell, "-NoLogo", "-NoProfile", "-Command", script,
|
||||
}, 0)
|
||||
if bytes, err := os.ReadFile(insideMarker); err != nil || strings.TrimSpace(string(bytes)) != "2,4,6 ✓" {
|
||||
t.Fatalf("PowerShell sandbox marker = %q, %v; want pipeline output", bytes, err)
|
||||
}
|
||||
|
||||
outsideMarker := filepath.Join(outside, "powershell-denied.txt")
|
||||
deniedScript := "$ErrorActionPreference='Stop'; " +
|
||||
"[IO.File]::WriteAllText(" + powershellSingleQuote(outsideMarker) + ", 'leaked')"
|
||||
runWindowsRealSmokeCommand(t, runnerExe, config, []string{
|
||||
powerShell, "-NoLogo", "-NoProfile", "-Command", deniedScript,
|
||||
}, 1)
|
||||
if _, err := os.Stat(outsideMarker); err == nil {
|
||||
t.Fatal("sandboxed PowerShell wrote outside every granted root")
|
||||
} else if !os.IsNotExist(err) {
|
||||
t.Fatalf("stat outside PowerShell marker: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func realSmokePowerShell(t *testing.T) string {
|
||||
t.Helper()
|
||||
for _, candidate := range []string{"pwsh.exe", "pwsh", "powershell.exe", "powershell"} {
|
||||
if path, err := exec.LookPath(candidate); err == nil {
|
||||
return path
|
||||
}
|
||||
}
|
||||
t.Skip("PowerShell is unavailable")
|
||||
return ""
|
||||
}
|
||||
|
||||
func realSmokeExecutable(t *testing.T, envKey string, fallbackName string) string {
|
||||
t.Helper()
|
||||
if path := os.Getenv(envKey); path != "" {
|
||||
|
||||
@@ -27,13 +27,10 @@ var fallbackSandboxRuntimes = struct {
|
||||
}{roots: make(map[string]string)}
|
||||
|
||||
type SandboxRuntime struct {
|
||||
Root string `json:"root,omitempty"`
|
||||
Home string `json:"home,omitempty"`
|
||||
Cache string `json:"cache,omitempty"`
|
||||
Config string `json:"config,omitempty"`
|
||||
Data string `json:"data,omitempty"`
|
||||
State string `json:"state,omitempty"`
|
||||
Temp string `json:"temp,omitempty"`
|
||||
Root string `json:"root,omitempty"`
|
||||
Cache string `json:"cache,omitempty"`
|
||||
Data string `json:"data,omitempty"`
|
||||
Temp string `json:"temp,omitempty"`
|
||||
}
|
||||
|
||||
func prepareSandboxRuntime(workspaceRoot string) (SandboxRuntime, func(), error) {
|
||||
@@ -75,21 +72,15 @@ func prepareSandboxRuntime(workspaceRoot string) (SandboxRuntime, func(), error)
|
||||
}
|
||||
}()
|
||||
runtimeState := SandboxRuntime{
|
||||
Root: root,
|
||||
Home: filepath.Join(root, "home"),
|
||||
Cache: filepath.Join(root, "cache"),
|
||||
Config: filepath.Join(root, "config"),
|
||||
Data: filepath.Join(root, "data"),
|
||||
State: filepath.Join(root, "state"),
|
||||
Temp: filepath.Join(root, "tmp"),
|
||||
Root: root,
|
||||
Cache: filepath.Join(root, "cache"),
|
||||
Data: filepath.Join(root, "data"),
|
||||
Temp: filepath.Join(root, "tmp"),
|
||||
}
|
||||
directories := []string{
|
||||
runtimeState.Root,
|
||||
runtimeState.Home,
|
||||
runtimeState.Cache,
|
||||
runtimeState.Config,
|
||||
runtimeState.Data,
|
||||
runtimeState.State,
|
||||
runtimeState.Temp,
|
||||
filepath.Join(runtimeState.Cache, "npm"),
|
||||
filepath.Join(runtimeState.Cache, "yarn"),
|
||||
@@ -207,16 +198,11 @@ func sandboxRuntimeEnvironment(env []string, runtimeState *SandboxRuntime) []str
|
||||
return env
|
||||
}
|
||||
overrides := []string{
|
||||
"HOME=" + runtimeState.Home,
|
||||
"XDG_CACHE_HOME=" + runtimeState.Cache,
|
||||
"XDG_CONFIG_HOME=" + runtimeState.Config,
|
||||
"XDG_DATA_HOME=" + runtimeState.Data,
|
||||
"XDG_STATE_HOME=" + runtimeState.State,
|
||||
"TMPDIR=" + runtimeState.Temp,
|
||||
"TMP=" + runtimeState.Temp,
|
||||
"TEMP=" + runtimeState.Temp,
|
||||
"npm_config_cache=" + filepath.Join(runtimeState.Cache, "npm"),
|
||||
"NPM_CONFIG_USERCONFIG=" + filepath.Join(runtimeState.Config, "npmrc"),
|
||||
"YARN_CACHE_FOLDER=" + filepath.Join(runtimeState.Cache, "yarn"),
|
||||
"COREPACK_HOME=" + filepath.Join(runtimeState.Cache, "corepack"),
|
||||
"PIP_CACHE_DIR=" + filepath.Join(runtimeState.Cache, "pip"),
|
||||
|
||||
@@ -2,7 +2,9 @@ package sandbox
|
||||
|
||||
import (
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
@@ -22,12 +24,18 @@ func TestPrepareSandboxRuntimeStaysOutsideWorkspace(t *testing.T) {
|
||||
if pathWithinRoot(workspace, runtimeState.Root) {
|
||||
t.Fatalf("runtime root %q must stay outside workspace %q", runtimeState.Root, workspace)
|
||||
}
|
||||
for _, path := range []string{runtimeState.Home, runtimeState.Cache, runtimeState.Config, runtimeState.Data, runtimeState.State, runtimeState.Temp} {
|
||||
for _, path := range []string{runtimeState.Cache, runtimeState.Data, runtimeState.Temp, filepath.Join(runtimeState.Data, "cargo")} {
|
||||
info, err := os.Stat(path)
|
||||
if err != nil || !info.IsDir() {
|
||||
t.Fatalf("managed runtime directory %q was not prepared: %v", path, err)
|
||||
}
|
||||
}
|
||||
for _, name := range []string{"home", "config", "state"} {
|
||||
path := filepath.Join(runtimeState.Root, name)
|
||||
if _, err := os.Stat(path); !os.IsNotExist(err) {
|
||||
t.Fatalf("obsolete synthetic runtime directory %q exists: %v", path, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrepareSandboxRuntimeCleansExpiredSibling(t *testing.T) {
|
||||
@@ -116,34 +124,63 @@ func TestCleanupSandboxRuntimeSkipsActiveLease(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSandboxRuntimeEnvironmentUsesManagedState(t *testing.T) {
|
||||
func TestSandboxRuntimeEnvironmentPreservesUserConfiguration(t *testing.T) {
|
||||
root := filepath.Join(t.TempDir(), "runtime")
|
||||
runtimeState := SandboxRuntime{
|
||||
Root: root,
|
||||
Home: filepath.Join(root, "home"),
|
||||
Cache: filepath.Join(root, "cache"),
|
||||
Config: filepath.Join(root, "config"),
|
||||
Data: filepath.Join(root, "data"),
|
||||
State: filepath.Join(root, "state"),
|
||||
Temp: filepath.Join(root, "tmp"),
|
||||
Root: root,
|
||||
Cache: filepath.Join(root, "cache"),
|
||||
Data: filepath.Join(root, "data"),
|
||||
Temp: filepath.Join(root, "tmp"),
|
||||
}
|
||||
env := sandboxRuntimeEnvironment([]string{
|
||||
"HOME=/workspace",
|
||||
"HOME=/home/user",
|
||||
"XDG_CACHE_HOME=/host/cache",
|
||||
"XDG_CONFIG_HOME=/home/user/.config",
|
||||
"XDG_DATA_HOME=/home/user/.local/share",
|
||||
"XDG_STATE_HOME=/home/user/.local/state",
|
||||
"NPM_CONFIG_USERCONFIG=/home/user/.npmrc",
|
||||
"PATH=/usr/bin",
|
||||
}, &runtimeState)
|
||||
|
||||
for key, want := range map[string]string{
|
||||
"HOME": runtimeState.Home,
|
||||
"XDG_CACHE_HOME": runtimeState.Cache,
|
||||
"XDG_CONFIG_HOME": runtimeState.Config,
|
||||
"XDG_DATA_HOME": runtimeState.Data,
|
||||
"XDG_STATE_HOME": runtimeState.State,
|
||||
"TMPDIR": runtimeState.Temp,
|
||||
"npm_config_cache": filepath.Join(runtimeState.Cache, "npm"),
|
||||
"NPM_CONFIG_USERCONFIG": filepath.Join(runtimeState.Config, "npmrc"),
|
||||
"YARN_CACHE_FOLDER": filepath.Join(runtimeState.Cache, "yarn"),
|
||||
"COREPACK_HOME": filepath.Join(runtimeState.Cache, "corepack"),
|
||||
"HOME": "/home/user",
|
||||
"XDG_CONFIG_HOME": "/home/user/.config",
|
||||
"XDG_DATA_HOME": "/home/user/.local/share",
|
||||
"XDG_STATE_HOME": "/home/user/.local/state",
|
||||
"NPM_CONFIG_USERCONFIG": "/home/user/.npmrc",
|
||||
} {
|
||||
if got := envListValue(env, key, ""); got != want {
|
||||
t.Fatalf("%s = %q, want %q; env=%#v", key, got, want, env)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSandboxRuntimeEnvironmentUsesManagedWritableState(t *testing.T) {
|
||||
root := filepath.Join(t.TempDir(), "runtime")
|
||||
runtimeState := SandboxRuntime{
|
||||
Root: root,
|
||||
Cache: filepath.Join(root, "cache"),
|
||||
Data: filepath.Join(root, "data"),
|
||||
Temp: filepath.Join(root, "tmp"),
|
||||
}
|
||||
env := sandboxRuntimeEnvironment([]string{
|
||||
"XDG_CACHE_HOME=/host/cache",
|
||||
"TMPDIR=/host/tmp",
|
||||
"PATH=/usr/bin",
|
||||
}, &runtimeState)
|
||||
|
||||
for key, want := range map[string]string{
|
||||
"XDG_CACHE_HOME": runtimeState.Cache,
|
||||
"TMPDIR": runtimeState.Temp,
|
||||
"TMP": runtimeState.Temp,
|
||||
"TEMP": runtimeState.Temp,
|
||||
"npm_config_cache": filepath.Join(runtimeState.Cache, "npm"),
|
||||
"YARN_CACHE_FOLDER": filepath.Join(runtimeState.Cache, "yarn"),
|
||||
"COREPACK_HOME": filepath.Join(runtimeState.Cache, "corepack"),
|
||||
"PIP_CACHE_DIR": filepath.Join(runtimeState.Cache, "pip"),
|
||||
"GOCACHE": filepath.Join(runtimeState.Cache, "go-build"),
|
||||
"GOMODCACHE": filepath.Join(runtimeState.Data, "go-mod"),
|
||||
"CARGO_HOME": filepath.Join(runtimeState.Data, "cargo"),
|
||||
} {
|
||||
if got := envListValue(env, key, ""); got != want {
|
||||
t.Fatalf("%s = %q, want %q; env=%#v", key, got, want, env)
|
||||
@@ -154,9 +191,93 @@ func TestSandboxRuntimeEnvironmentUsesManagedState(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSandboxRuntimeEnvironmentPreservesGitGlobalIgnore(t *testing.T) {
|
||||
git, err := exec.LookPath("git")
|
||||
if err != nil {
|
||||
t.Skip("git is unavailable")
|
||||
}
|
||||
root := t.TempDir()
|
||||
home := filepath.Join(root, "home")
|
||||
configHome := filepath.Join(home, ".config")
|
||||
repository := filepath.Join(root, "repository")
|
||||
runtimeRoot := filepath.Join(root, "runtime")
|
||||
for _, directory := range []string{
|
||||
filepath.Join(configHome, "git"),
|
||||
repository,
|
||||
filepath.Join(runtimeRoot, "cache"),
|
||||
filepath.Join(runtimeRoot, "data"),
|
||||
filepath.Join(runtimeRoot, "tmp"),
|
||||
} {
|
||||
if err := os.MkdirAll(directory, 0o700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(configHome, "git", "ignore"), []byte("ignored.txt\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(repository, "ignored.txt"), []byte("ignored\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(repository, "visible.txt"), []byte("visible\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
runtimeState := SandboxRuntime{
|
||||
Root: runtimeRoot,
|
||||
Cache: filepath.Join(runtimeRoot, "cache"),
|
||||
Data: filepath.Join(runtimeRoot, "data"),
|
||||
Temp: filepath.Join(runtimeRoot, "tmp"),
|
||||
}
|
||||
inherited := append(os.Environ(),
|
||||
"GIT_CONFIG_GLOBAL="+filepath.Join(root, "missing-global-config"),
|
||||
"GIT_DIR="+filepath.Join(root, "wrong-repository"),
|
||||
"GIT_CONFIG_COUNT=1",
|
||||
"GIT_CONFIG_KEY_0=core.excludesFile",
|
||||
"GIT_CONFIG_VALUE_0="+filepath.Join(root, "missing-ignore"),
|
||||
)
|
||||
env := sandboxRuntimeEnvironment(upsertEnvList(withoutGitEnvironmentOverrides(inherited),
|
||||
"HOME="+home,
|
||||
"XDG_CONFIG_HOME="+configHome,
|
||||
"GIT_CONFIG_NOSYSTEM=1",
|
||||
"GIT_TERMINAL_PROMPT=0",
|
||||
), &runtimeState)
|
||||
runGit := func(args ...string) string {
|
||||
t.Helper()
|
||||
command := exec.Command(git, args...)
|
||||
command.Dir = repository
|
||||
command.Env = env
|
||||
output, err := command.CombinedOutput()
|
||||
if err != nil {
|
||||
t.Fatalf("git %s: %v\n%s", strings.Join(args, " "), err, output)
|
||||
}
|
||||
return string(output)
|
||||
}
|
||||
runGit("init", "--quiet")
|
||||
status := runGit("status", "--short", "--untracked-files=all")
|
||||
if strings.Contains(status, "ignored.txt") {
|
||||
t.Fatalf("global ignore was not honored:\n%s", status)
|
||||
}
|
||||
if !strings.Contains(status, "visible.txt") {
|
||||
t.Fatalf("non-ignored file is missing from status:\n%s", status)
|
||||
}
|
||||
}
|
||||
|
||||
func withoutGitEnvironmentOverrides(env []string) []string {
|
||||
out := make([]string, 0, len(env))
|
||||
for _, value := range env {
|
||||
key, _, ok := strings.Cut(value, "=")
|
||||
if ok && strings.HasPrefix(strings.ToUpper(key), "GIT_") {
|
||||
continue
|
||||
}
|
||||
out = append(out, value)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestEngineCommandPlanCarriesManagedRuntime(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
cacheRoot := t.TempDir()
|
||||
t.Setenv("HOME", filepath.Join(t.TempDir(), "home"))
|
||||
original := sandboxUserCacheDir
|
||||
sandboxUserCacheDir = func() (string, error) { return cacheRoot, nil }
|
||||
t.Cleanup(func() { sandboxUserCacheDir = original })
|
||||
@@ -189,8 +310,8 @@ func TestEngineCommandPlanCarriesManagedRuntime(t *testing.T) {
|
||||
} else if !inUse {
|
||||
t.Fatal("command plan runtime must be marked in use")
|
||||
}
|
||||
if got := envListValue(plan.Env, "HOME", ""); got != plan.PermissionProfile.Runtime.Home {
|
||||
t.Fatalf("HOME = %q, want managed home %q", got, plan.PermissionProfile.Runtime.Home)
|
||||
if got := envListValue(plan.Env, "HOME", ""); got != os.Getenv("HOME") {
|
||||
t.Fatalf("HOME = %q, want caller home %q", got, os.Getenv("HOME"))
|
||||
}
|
||||
foundWriteRoot := false
|
||||
for _, root := range plan.PermissionProfile.FileSystem.WriteRoots {
|
||||
|
||||
@@ -430,25 +430,70 @@ func shellDashCPayload(program string, fields []string) string {
|
||||
// name: it strips shell quoting/escaping characters (", ', `, \) wherever they
|
||||
// appear in the token (including embedded ones like `vi\m` or `v"i"m`), strips
|
||||
// leading command-substitution markers, removes any directory prefix (so
|
||||
// /usr/bin/vim and C:\tools\vim.exe match "vim"), and lowercases. This closes
|
||||
// path/quote/substitution evasions of the detector.
|
||||
// /usr/bin/vim and C:\tools\vim.exe match "vim"), removes Windows executable
|
||||
// suffixes, and lowercases. This closes path/quote/substitution evasions of the
|
||||
// detector and keeps curl.exe/npm.cmd equivalent to curl/npm for risk analysis.
|
||||
func normalizeProgramToken(field string) string {
|
||||
token := strings.TrimSpace(field)
|
||||
token = strings.TrimLeft(token, "$(")
|
||||
token = strings.TrimRight(token, ")")
|
||||
// Strip shell quoting/escaping characters (", ', `, \) wherever they appear
|
||||
// in the token — surrounding, embedded, or as a mid-word escape — so
|
||||
// "vim", v"i"m, 'v'im, and vi\m all collapse to the program name. This is
|
||||
// done BEFORE the directory-prefix trim so an escape can't masquerade as a
|
||||
// path separator (e.g. vi\m must become vim, not m).
|
||||
token = stripChars(token, "\"'`\\")
|
||||
// Strip a directory prefix so /usr/bin/vim reduces to the basename. (A
|
||||
// Windows-style backslash path separator is already removed above, so only
|
||||
// the POSIX separator remains to split on.)
|
||||
if i := strings.LastIndex(token, "/"); i >= 0 {
|
||||
token = token[i+1:]
|
||||
// PowerShell commonly invokes drive-letter, UNC, and relative paths. Extract
|
||||
// their basename while backslashes still carry path-separator meaning;
|
||||
// removing them first would turn C:\tools\curl.exe into c:toolscurl.exe and
|
||||
// bypass the canonical curl risk entry.
|
||||
pathToken := stripChars(token, "\"'`")
|
||||
if windowsPathBasename, ok := windowsExecutablePathBasename(pathToken); ok {
|
||||
token = windowsPathBasename
|
||||
} else {
|
||||
// Strip shell quoting/escaping characters (", ', `, \) wherever they
|
||||
// appear in the token — surrounding, embedded, or as a mid-word escape
|
||||
// — so "vim", v"i"m, 'v'im, and vi\m all collapse to the program name.
|
||||
token = stripChars(token, "\"'`\\")
|
||||
// Strip a directory prefix so /usr/bin/vim reduces to the basename.
|
||||
if i := strings.LastIndex(token, "/"); i >= 0 {
|
||||
token = token[i+1:]
|
||||
}
|
||||
}
|
||||
return strings.ToLower(token)
|
||||
token = strings.ToLower(token)
|
||||
for _, suffix := range []string{".exe", ".cmd", ".bat", ".com"} {
|
||||
if strings.HasSuffix(token, suffix) {
|
||||
return strings.TrimSuffix(token, suffix)
|
||||
}
|
||||
}
|
||||
return token
|
||||
}
|
||||
|
||||
func windowsExecutablePathBasename(token string) (string, bool) {
|
||||
drivePath := len(token) >= 3 &&
|
||||
((token[0] >= 'a' && token[0] <= 'z') || (token[0] >= 'A' && token[0] <= 'Z')) &&
|
||||
token[1] == ':'
|
||||
uncPath := strings.HasPrefix(token, `\\`)
|
||||
explicitRelativePath := strings.HasPrefix(token, `.\`) || strings.HasPrefix(token, `..\`)
|
||||
backslashRelativePath := strings.ContainsRune(token, '\\') && hasWindowsExecutableSuffix(token)
|
||||
if !drivePath && !uncPath && !explicitRelativePath && !backslashRelativePath {
|
||||
return "", false
|
||||
}
|
||||
index := strings.LastIndexAny(token, `\/`)
|
||||
if index >= 0 {
|
||||
if index+1 >= len(token) {
|
||||
return "", false
|
||||
}
|
||||
return token[index+1:], true
|
||||
}
|
||||
if drivePath && len(token) > 2 {
|
||||
return token[2:], true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func hasWindowsExecutableSuffix(token string) bool {
|
||||
token = strings.ToLower(token)
|
||||
for _, suffix := range []string{".exe", ".cmd", ".bat", ".com"} {
|
||||
if strings.HasSuffix(token, suffix) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// stripChars returns s with every rune in cutset removed.
|
||||
|
||||
@@ -0,0 +1,422 @@
|
||||
package sessions
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const (
|
||||
GoalObjectiveMaxLength = 4_000
|
||||
GoalMaxConsecutiveContinuations = 20
|
||||
goalContinuationLimitStatusReason = "automatic continuation limit reached"
|
||||
)
|
||||
|
||||
func validateGoalObjective(objective string) (string, error) {
|
||||
objective = strings.TrimSpace(objective)
|
||||
if objective == "" {
|
||||
return "", fmt.Errorf("goal objective is required")
|
||||
}
|
||||
if utf8.RuneCountInString(objective) > GoalObjectiveMaxLength {
|
||||
return "", fmt.Errorf("goal objective cannot exceed %d characters", GoalObjectiveMaxLength)
|
||||
}
|
||||
return objective, nil
|
||||
}
|
||||
|
||||
// CreateGoal persists a new active goal for a session. It refuses to replace an
|
||||
// existing goal so callers must make replacement an explicit user decision.
|
||||
func (store *Store) CreateGoal(sessionID, objective string, tokenBudget int) (Metadata, Event, error) {
|
||||
if !ValidSessionID(sessionID) {
|
||||
return Metadata{}, Event{}, fmt.Errorf("invalid zero session id %q", sessionID)
|
||||
}
|
||||
var err error
|
||||
objective, err = validateGoalObjective(objective)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
if tokenBudget < 0 {
|
||||
return Metadata{}, Event{}, fmt.Errorf("goal token budget cannot be negative")
|
||||
}
|
||||
|
||||
unlock, err := store.lockSession(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
defer unlock()
|
||||
|
||||
session, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
if session.Goal != nil {
|
||||
return Metadata{}, Event{}, fmt.Errorf("session already has a goal")
|
||||
}
|
||||
now := store.timestamp()
|
||||
session.Goal = &Goal{
|
||||
Objective: objective,
|
||||
Status: GoalStatusActive,
|
||||
TokenBudget: tokenBudget,
|
||||
ContinuationLimit: GoalMaxConsecutiveContinuations,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
if err := store.writeMetadata(session); err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
event, err := store.appendEventLocked(sessionID, AppendEventInput{
|
||||
Type: EventGoalCreated,
|
||||
Payload: session.Goal,
|
||||
})
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
loaded, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
return loaded, event, nil
|
||||
}
|
||||
|
||||
// UpdateGoal changes a goal's lifecycle state. Terminal goals remain available
|
||||
// for status/history until the user explicitly clears them.
|
||||
func (store *Store) UpdateGoal(sessionID string, status GoalStatus, reason string) (Metadata, Event, error) {
|
||||
if !ValidSessionID(sessionID) {
|
||||
return Metadata{}, Event{}, fmt.Errorf("invalid zero session id %q", sessionID)
|
||||
}
|
||||
if !validGoalStatus(status) {
|
||||
return Metadata{}, Event{}, fmt.Errorf("invalid goal status %q", status)
|
||||
}
|
||||
|
||||
unlock, err := store.lockSession(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
defer unlock()
|
||||
|
||||
session, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
if session.Goal == nil {
|
||||
return Metadata{}, Event{}, fmt.Errorf("session has no goal")
|
||||
}
|
||||
session.Goal.Status = status
|
||||
session.Goal.StatusReason = strings.TrimSpace(reason)
|
||||
if status == GoalStatusActive {
|
||||
session.Goal.ContinuationCount = 0
|
||||
if session.Goal.ContinuationLimit <= 0 {
|
||||
session.Goal.ContinuationLimit = GoalMaxConsecutiveContinuations
|
||||
}
|
||||
}
|
||||
session.Goal.UpdatedAt = store.timestamp()
|
||||
if err := store.writeMetadata(session); err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
event, err := store.appendEventLocked(sessionID, AppendEventInput{
|
||||
Type: EventGoalUpdated,
|
||||
Payload: session.Goal,
|
||||
})
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
loaded, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
return loaded, event, nil
|
||||
}
|
||||
|
||||
// EditGoal replaces the objective and optional budget without resetting usage or
|
||||
// lifecycle timestamps. User-facing command handling owns confirmation.
|
||||
func (store *Store) EditGoal(sessionID, objective string, tokenBudget int) (Metadata, Event, error) {
|
||||
if !ValidSessionID(sessionID) {
|
||||
return Metadata{}, Event{}, fmt.Errorf("invalid zero session id %q", sessionID)
|
||||
}
|
||||
var err error
|
||||
objective, err = validateGoalObjective(objective)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
if tokenBudget < 0 {
|
||||
return Metadata{}, Event{}, fmt.Errorf("goal token budget cannot be negative")
|
||||
}
|
||||
unlock, err := store.lockSession(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
defer unlock()
|
||||
|
||||
session, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
if session.Goal == nil {
|
||||
return Metadata{}, Event{}, fmt.Errorf("session has no goal")
|
||||
}
|
||||
session.Goal.Objective = objective
|
||||
session.Goal.TokenBudget = tokenBudget
|
||||
session.Goal.Status = GoalStatusActive
|
||||
session.Goal.StatusReason = ""
|
||||
session.Goal.ContinuationCount = 0
|
||||
if session.Goal.ContinuationLimit <= 0 {
|
||||
session.Goal.ContinuationLimit = GoalMaxConsecutiveContinuations
|
||||
}
|
||||
if tokenBudget > 0 && session.Goal.TokensUsed >= tokenBudget {
|
||||
session.Goal.Status = GoalStatusBudgetLimited
|
||||
session.Goal.StatusReason = "token budget reached"
|
||||
}
|
||||
session.Goal.UpdatedAt = store.timestamp()
|
||||
if err := store.writeMetadata(session); err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
event, err := store.appendEventLocked(sessionID, AppendEventInput{
|
||||
Type: EventGoalUpdated,
|
||||
Payload: session.Goal,
|
||||
})
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
loaded, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
return loaded, event, nil
|
||||
}
|
||||
|
||||
// ResetGoalContinuations starts a fresh autonomous-run allowance after explicit
|
||||
// user input. It makes the safety bound consecutive rather than lifetime-wide.
|
||||
func (store *Store) ResetGoalContinuations(sessionID string) (Metadata, error) {
|
||||
if !ValidSessionID(sessionID) {
|
||||
return Metadata{}, fmt.Errorf("invalid zero session id %q", sessionID)
|
||||
}
|
||||
unlock, err := store.lockSession(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, err
|
||||
}
|
||||
defer unlock()
|
||||
|
||||
session, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, err
|
||||
}
|
||||
if session.Goal == nil || session.Goal.Status != GoalStatusActive {
|
||||
return session, nil
|
||||
}
|
||||
changed := false
|
||||
if session.Goal.ContinuationCount != 0 {
|
||||
session.Goal.ContinuationCount = 0
|
||||
changed = true
|
||||
}
|
||||
if session.Goal.ContinuationLimit <= 0 {
|
||||
session.Goal.ContinuationLimit = GoalMaxConsecutiveContinuations
|
||||
changed = true
|
||||
}
|
||||
if !changed {
|
||||
return session, nil
|
||||
}
|
||||
session.Goal.UpdatedAt = store.timestamp()
|
||||
if err := store.writeMetadata(session); err != nil {
|
||||
return Metadata{}, err
|
||||
}
|
||||
return session, nil
|
||||
}
|
||||
|
||||
// ReserveGoalContinuation atomically reserves one automatic continuation. Once
|
||||
// the persisted consecutive-run limit is exhausted it pauses the goal before
|
||||
// another provider request can start, even when the provider reports no usage.
|
||||
func (store *Store) ReserveGoalContinuation(sessionID string) (Metadata, *Event, bool, error) {
|
||||
if !ValidSessionID(sessionID) {
|
||||
return Metadata{}, nil, false, fmt.Errorf("invalid zero session id %q", sessionID)
|
||||
}
|
||||
unlock, err := store.lockSession(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, nil, false, err
|
||||
}
|
||||
defer unlock()
|
||||
|
||||
session, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, nil, false, err
|
||||
}
|
||||
if session.Goal == nil || session.Goal.Status != GoalStatusActive {
|
||||
return session, nil, false, nil
|
||||
}
|
||||
limit := session.Goal.ContinuationLimit
|
||||
if limit <= 0 {
|
||||
limit = GoalMaxConsecutiveContinuations
|
||||
session.Goal.ContinuationLimit = limit
|
||||
}
|
||||
if session.Goal.ContinuationCount < limit {
|
||||
session.Goal.ContinuationCount++
|
||||
session.Goal.UpdatedAt = store.timestamp()
|
||||
if err := store.writeMetadata(session); err != nil {
|
||||
return Metadata{}, nil, false, err
|
||||
}
|
||||
return session, nil, true, nil
|
||||
}
|
||||
|
||||
session.Goal.Status = GoalStatusPaused
|
||||
session.Goal.StatusReason = goalContinuationLimitStatusReason
|
||||
session.Goal.UpdatedAt = store.timestamp()
|
||||
if err := store.writeMetadata(session); err != nil {
|
||||
return Metadata{}, nil, false, err
|
||||
}
|
||||
event, err := store.appendEventLocked(sessionID, AppendEventInput{
|
||||
Type: EventGoalUpdated,
|
||||
Payload: session.Goal,
|
||||
})
|
||||
if err != nil {
|
||||
return Metadata{}, nil, false, err
|
||||
}
|
||||
loaded, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, nil, false, err
|
||||
}
|
||||
return loaded, &event, false, nil
|
||||
}
|
||||
|
||||
// AddGoalUsage accounts tokens consumed while a goal is active. Reaching the
|
||||
// optional budget pauses the goal before another autonomous turn can start.
|
||||
func (store *Store) AddGoalUsage(sessionID string, tokens int) (Metadata, *Event, error) {
|
||||
if !ValidSessionID(sessionID) {
|
||||
return Metadata{}, nil, fmt.Errorf("invalid zero session id %q", sessionID)
|
||||
}
|
||||
if tokens < 0 {
|
||||
return Metadata{}, nil, fmt.Errorf("goal token usage cannot be negative")
|
||||
}
|
||||
unlock, err := store.lockSession(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, nil, err
|
||||
}
|
||||
defer unlock()
|
||||
|
||||
session, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, nil, err
|
||||
}
|
||||
if session.Goal == nil || tokens == 0 {
|
||||
return session, nil, nil
|
||||
}
|
||||
session.Goal.TokensUsed += tokens
|
||||
session.Goal.UpdatedAt = store.timestamp()
|
||||
budgetReached := false
|
||||
if session.Goal.Status == GoalStatusActive &&
|
||||
session.Goal.TokenBudget > 0 &&
|
||||
session.Goal.TokensUsed >= session.Goal.TokenBudget {
|
||||
session.Goal.Status = GoalStatusBudgetLimited
|
||||
session.Goal.StatusReason = "token budget reached"
|
||||
budgetReached = true
|
||||
}
|
||||
if err := store.writeMetadata(session); err != nil {
|
||||
return Metadata{}, nil, err
|
||||
}
|
||||
if !budgetReached {
|
||||
return session, nil, nil
|
||||
}
|
||||
event, err := store.appendEventLocked(sessionID, AppendEventInput{
|
||||
Type: EventGoalUpdated,
|
||||
Payload: session.Goal,
|
||||
})
|
||||
if err != nil {
|
||||
return Metadata{}, nil, err
|
||||
}
|
||||
loaded, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, nil, err
|
||||
}
|
||||
return loaded, &event, nil
|
||||
}
|
||||
|
||||
// PauseGoalIfActive pauses a goal only while it is still active. It is used by
|
||||
// cancellation paths where an in-flight agent may have completed or blocked the
|
||||
// goal immediately before the cancellation reached the runtime.
|
||||
func (store *Store) PauseGoalIfActive(sessionID, reason string) (Metadata, *Event, error) {
|
||||
if !ValidSessionID(sessionID) {
|
||||
return Metadata{}, nil, fmt.Errorf("invalid zero session id %q", sessionID)
|
||||
}
|
||||
unlock, err := store.lockSession(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, nil, err
|
||||
}
|
||||
defer unlock()
|
||||
|
||||
session, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, nil, err
|
||||
}
|
||||
if session.Goal == nil || session.Goal.Status != GoalStatusActive {
|
||||
return session, nil, nil
|
||||
}
|
||||
session.Goal.Status = GoalStatusPaused
|
||||
session.Goal.StatusReason = strings.TrimSpace(reason)
|
||||
session.Goal.UpdatedAt = store.timestamp()
|
||||
if err := store.writeMetadata(session); err != nil {
|
||||
return Metadata{}, nil, err
|
||||
}
|
||||
event, err := store.appendEventLocked(sessionID, AppendEventInput{
|
||||
Type: EventGoalUpdated,
|
||||
Payload: session.Goal,
|
||||
})
|
||||
if err != nil {
|
||||
return Metadata{}, nil, err
|
||||
}
|
||||
loaded, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, nil, err
|
||||
}
|
||||
return loaded, &event, nil
|
||||
}
|
||||
|
||||
// ClearGoal removes the current goal while preserving an audit event.
|
||||
func (store *Store) ClearGoal(sessionID string) (Metadata, Event, error) {
|
||||
if !ValidSessionID(sessionID) {
|
||||
return Metadata{}, Event{}, fmt.Errorf("invalid zero session id %q", sessionID)
|
||||
}
|
||||
unlock, err := store.lockSession(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
defer unlock()
|
||||
|
||||
session, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
if session.Goal == nil {
|
||||
return Metadata{}, Event{}, fmt.Errorf("session has no goal")
|
||||
}
|
||||
cleared := cloneGoal(session.Goal)
|
||||
session.Goal = nil
|
||||
if err := store.writeMetadata(session); err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
event, err := store.appendEventLocked(sessionID, AppendEventInput{
|
||||
Type: EventGoalCleared,
|
||||
Payload: cleared,
|
||||
})
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
loaded, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, Event{}, err
|
||||
}
|
||||
return loaded, event, nil
|
||||
}
|
||||
|
||||
func validGoalStatus(status GoalStatus) bool {
|
||||
switch status {
|
||||
case GoalStatusActive, GoalStatusPaused, GoalStatusBlocked, GoalStatusBudgetLimited, GoalStatusUsageLimited, GoalStatusComplete:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func cloneGoal(goal *Goal) *Goal {
|
||||
if goal == nil {
|
||||
return nil
|
||||
}
|
||||
copy := *goal
|
||||
return ©
|
||||
}
|
||||
@@ -0,0 +1,382 @@
|
||||
package sessions
|
||||
|
||||
import (
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestGoalLifecyclePersistsInSessionMetadata(t *testing.T) {
|
||||
now := time.Date(2026, 7, 24, 10, 0, 0, 0, time.UTC)
|
||||
store := NewStore(StoreOptions{
|
||||
RootDir: t.TempDir(),
|
||||
Now: func() time.Time {
|
||||
now = now.Add(time.Second)
|
||||
return now
|
||||
},
|
||||
})
|
||||
session, err := store.Create(CreateInput{SessionID: "goal_session"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
created, event, err := store.CreateGoal(session.SessionID, "Ship the release", 1_000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if event.Type != EventGoalCreated {
|
||||
t.Fatalf("create event = %q, want %q", event.Type, EventGoalCreated)
|
||||
}
|
||||
if created.Goal == nil || created.Goal.Objective != "Ship the release" ||
|
||||
created.Goal.Status != GoalStatusActive || created.Goal.TokenBudget != 1_000 {
|
||||
t.Fatalf("created goal = %#v", created.Goal)
|
||||
}
|
||||
|
||||
accounted, _, err := store.AddGoalUsage(session.SessionID, 250)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if accounted.Goal.TokensUsed != 250 || accounted.Goal.Status != GoalStatusActive {
|
||||
t.Fatalf("accounted goal = %#v", accounted.Goal)
|
||||
}
|
||||
|
||||
paused, event, err := store.UpdateGoal(session.SessionID, GoalStatusPaused, "user interrupted")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if event.Type != EventGoalUpdated || paused.Goal.Status != GoalStatusPaused ||
|
||||
paused.Goal.StatusReason != "user interrupted" {
|
||||
t.Fatalf("paused goal/event = %#v / %#v", paused.Goal, event)
|
||||
}
|
||||
|
||||
loaded, err := store.Get(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if loaded == nil || loaded.Goal == nil || loaded.Goal.Status != GoalStatusPaused {
|
||||
t.Fatalf("reloaded session = %#v", loaded)
|
||||
}
|
||||
|
||||
cleared, event, err := store.ClearGoal(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if event.Type != EventGoalCleared || cleared.Goal != nil {
|
||||
t.Fatalf("cleared goal/event = %#v / %#v", cleared.Goal, event)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGoalBudgetPausesAtLimit(t *testing.T) {
|
||||
store := NewStore(StoreOptions{RootDir: t.TempDir()})
|
||||
session, err := store.Create(CreateInput{SessionID: "budget_session"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := store.CreateGoal(session.SessionID, "Stay bounded", 100); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
updated, event, err := store.AddGoalUsage(session.SessionID, 100)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if updated.Goal.Status != GoalStatusBudgetLimited || updated.Goal.StatusReason != "token budget reached" {
|
||||
t.Fatalf("budgeted goal = %#v", updated.Goal)
|
||||
}
|
||||
if event == nil || event.Type != EventGoalUpdated {
|
||||
t.Fatalf("budget transition event = %#v", event)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGoalContinuationLimitStopsWithoutProviderUsage(t *testing.T) {
|
||||
store := NewStore(StoreOptions{RootDir: t.TempDir()})
|
||||
session, err := store.Create(CreateInput{SessionID: "continuation_limit"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := store.CreateGoal(session.SessionID, "Stay bounded without usage events", 0); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
for continuation := 1; continuation <= GoalMaxConsecutiveContinuations; continuation++ {
|
||||
updated, event, reserved, err := store.ReserveGoalContinuation(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !reserved || event != nil {
|
||||
t.Fatalf("continuation %d: reserved=%v event=%#v", continuation, reserved, event)
|
||||
}
|
||||
if updated.Goal.ContinuationCount != continuation {
|
||||
t.Fatalf("continuation count = %d, want %d", updated.Goal.ContinuationCount, continuation)
|
||||
}
|
||||
if updated.Goal.TokensUsed != 0 {
|
||||
t.Fatalf("provider-independent guard unexpectedly recorded tokens: %#v", updated.Goal)
|
||||
}
|
||||
}
|
||||
|
||||
stopped, event, reserved, err := store.ReserveGoalContinuation(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if reserved || event == nil || stopped.Goal.Status != GoalStatusPaused ||
|
||||
stopped.Goal.StatusReason != goalContinuationLimitStatusReason {
|
||||
t.Fatalf("unbounded continuation was not stopped: goal=%#v event=%#v reserved=%v", stopped.Goal, event, reserved)
|
||||
}
|
||||
reloaded, err := store.Get(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if reloaded.Goal.ContinuationCount != GoalMaxConsecutiveContinuations ||
|
||||
reloaded.Goal.Status != GoalStatusPaused {
|
||||
t.Fatalf("continuation guard did not persist: %#v", reloaded.Goal)
|
||||
}
|
||||
|
||||
resumed, _, err := store.UpdateGoal(session.SessionID, GoalStatusActive, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resumed.Goal.ContinuationCount != 0 {
|
||||
t.Fatalf("explicit resume did not reset consecutive continuations: %#v", resumed.Goal)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResetGoalContinuationsPersistsOnlyActiveChanges(t *testing.T) {
|
||||
newStore := func(t *testing.T) (*Store, Metadata) {
|
||||
t.Helper()
|
||||
now := time.Date(2026, 7, 26, 10, 0, 0, 0, time.UTC)
|
||||
store := NewStore(StoreOptions{
|
||||
RootDir: t.TempDir(),
|
||||
Now: func() time.Time {
|
||||
now = now.Add(time.Second)
|
||||
return now
|
||||
},
|
||||
})
|
||||
session, err := store.Create(CreateInput{SessionID: "reset_goal"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return store, session
|
||||
}
|
||||
pinMetadataTime := func(t *testing.T, store *Store, sessionID string) time.Time {
|
||||
t.Helper()
|
||||
path := store.metadataPath(sessionID)
|
||||
sentinel := time.Date(2001, 2, 3, 4, 5, 6, 0, time.UTC)
|
||||
if err := os.Chtimes(path, sentinel, sentinel); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
info, err := os.Stat(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return info.ModTime()
|
||||
}
|
||||
assertMetadataTime := func(t *testing.T, store *Store, sessionID string, want time.Time) {
|
||||
t.Helper()
|
||||
info, err := os.Stat(store.metadataPath(sessionID))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !info.ModTime().Equal(want) {
|
||||
t.Fatalf("metadata was rewritten: modtime=%s want=%s", info.ModTime(), want)
|
||||
}
|
||||
}
|
||||
|
||||
t.Run("no goal is a no-op", func(t *testing.T) {
|
||||
store, session := newStore(t)
|
||||
modTime := pinMetadataTime(t, store, session.SessionID)
|
||||
updated, err := store.ResetGoalContinuations(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if updated.Goal != nil {
|
||||
t.Fatalf("no-goal reset created a goal: %#v", updated.Goal)
|
||||
}
|
||||
assertMetadataTime(t, store, session.SessionID, modTime)
|
||||
})
|
||||
|
||||
t.Run("inactive goal is a no-op", func(t *testing.T) {
|
||||
store, session := newStore(t)
|
||||
if _, _, err := store.CreateGoal(session.SessionID, "Pause safely", 0); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, reserved, err := store.ReserveGoalContinuation(session.SessionID); err != nil || !reserved {
|
||||
t.Fatalf("reserve continuation: reserved=%v err=%v", reserved, err)
|
||||
}
|
||||
paused, _, err := store.UpdateGoal(session.SessionID, GoalStatusPaused, "user paused")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
modTime := pinMetadataTime(t, store, session.SessionID)
|
||||
updated, err := store.ResetGoalContinuations(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if updated.Goal.ContinuationCount != 1 || updated.Goal.UpdatedAt != paused.Goal.UpdatedAt {
|
||||
t.Fatalf("inactive goal changed during reset: before=%#v after=%#v", paused.Goal, updated.Goal)
|
||||
}
|
||||
assertMetadataTime(t, store, session.SessionID, modTime)
|
||||
})
|
||||
|
||||
t.Run("already reset active goal is a no-op", func(t *testing.T) {
|
||||
store, session := newStore(t)
|
||||
created, _, err := store.CreateGoal(session.SessionID, "Stay reset", 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
modTime := pinMetadataTime(t, store, session.SessionID)
|
||||
updated, err := store.ResetGoalContinuations(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if updated.Goal.ContinuationCount != 0 ||
|
||||
updated.Goal.ContinuationLimit != GoalMaxConsecutiveContinuations ||
|
||||
updated.Goal.UpdatedAt != created.Goal.UpdatedAt {
|
||||
t.Fatalf("already-reset goal changed: before=%#v after=%#v", created.Goal, updated.Goal)
|
||||
}
|
||||
assertMetadataTime(t, store, session.SessionID, modTime)
|
||||
})
|
||||
|
||||
t.Run("active goal reset persists", func(t *testing.T) {
|
||||
store, session := newStore(t)
|
||||
if _, _, err := store.CreateGoal(session.SessionID, "Reset progress", 0); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
reserved, _, ok, err := store.ReserveGoalContinuation(session.SessionID)
|
||||
if err != nil || !ok {
|
||||
t.Fatalf("reserve continuation: reserved=%v err=%v", ok, err)
|
||||
}
|
||||
updated, err := store.ResetGoalContinuations(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if updated.Goal.ContinuationCount != 0 ||
|
||||
updated.Goal.ContinuationLimit != GoalMaxConsecutiveContinuations ||
|
||||
updated.Goal.UpdatedAt == reserved.Goal.UpdatedAt {
|
||||
t.Fatalf("active reset did not persist: before=%#v after=%#v", reserved.Goal, updated.Goal)
|
||||
}
|
||||
reloaded, err := store.Get(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if reloaded.Goal == nil || *reloaded.Goal != *updated.Goal {
|
||||
t.Fatalf("active reset was not durable: updated=%#v reloaded=%#v", updated.Goal, reloaded.Goal)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestGoalObjectiveLengthIsBounded(t *testing.T) {
|
||||
store := NewStore(StoreOptions{RootDir: t.TempDir()})
|
||||
session, err := store.Create(CreateInput{SessionID: "objective_limit"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tooLong := strings.Repeat("g", GoalObjectiveMaxLength+1)
|
||||
if _, _, err := store.CreateGoal(session.SessionID, tooLong, 0); err == nil ||
|
||||
!strings.Contains(err.Error(), "cannot exceed") {
|
||||
t.Fatalf("oversized objective error = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditGoalUpdatesStateAndRejectsInvalidInputWithoutMutation(t *testing.T) {
|
||||
store := NewStore(StoreOptions{RootDir: t.TempDir()})
|
||||
session, err := store.Create(CreateInput{SessionID: "edit_goal"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := store.CreateGoal(session.SessionID, "Original objective", 1_000); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := store.AddGoalUsage(session.SessionID, 250); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
edited, event, err := store.EditGoal(session.SessionID, "Updated objective", 500)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if event.Type != EventGoalUpdated {
|
||||
t.Fatalf("edit event = %q, want %q", event.Type, EventGoalUpdated)
|
||||
}
|
||||
if edited.Goal == nil ||
|
||||
edited.Goal.Objective != "Updated objective" ||
|
||||
edited.Goal.TokenBudget != 500 ||
|
||||
edited.Goal.TokensUsed != 250 ||
|
||||
edited.Goal.Status != GoalStatusActive ||
|
||||
edited.Goal.StatusReason != "" {
|
||||
t.Fatalf("edited goal = %#v", edited.Goal)
|
||||
}
|
||||
|
||||
limited, _, err := store.EditGoal(session.SessionID, "Stay within budget", 200)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if limited.Goal == nil ||
|
||||
limited.Goal.Status != GoalStatusBudgetLimited ||
|
||||
limited.Goal.StatusReason != "token budget reached" ||
|
||||
limited.Goal.TokensUsed != 250 {
|
||||
t.Fatalf("budget-limited goal = %#v", limited.Goal)
|
||||
}
|
||||
|
||||
before := *limited.Goal
|
||||
eventsBefore, err := store.ReadEvents(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := store.EditGoal(session.SessionID, " ", 200); err == nil {
|
||||
t.Fatal("EditGoal should reject an empty objective")
|
||||
}
|
||||
if _, _, err := store.EditGoal(session.SessionID, "Invalid budget", -1); err == nil {
|
||||
t.Fatal("EditGoal should reject a negative token budget")
|
||||
}
|
||||
after, err := store.Get(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if after == nil || after.Goal == nil || *after.Goal != before {
|
||||
t.Fatalf("invalid edit mutated goal: before=%#v after=%#v", before, after)
|
||||
}
|
||||
eventsAfter, err := store.ReadEvents(session.SessionID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(eventsAfter) != len(eventsBefore) {
|
||||
t.Fatalf("invalid edit appended events: before=%d after=%d", len(eventsBefore), len(eventsAfter))
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateGoalRefusesImplicitReplacement(t *testing.T) {
|
||||
store := NewStore(StoreOptions{RootDir: t.TempDir()})
|
||||
session, err := store.Create(CreateInput{SessionID: "replace_session"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := store.CreateGoal(session.SessionID, "First", 0); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := store.CreateGoal(session.SessionID, "Second", 0); err == nil {
|
||||
t.Fatal("CreateGoal should require an explicit clear before replacement")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPauseGoalIfActiveDoesNotOverwriteTerminalState(t *testing.T) {
|
||||
store := NewStore(StoreOptions{RootDir: t.TempDir()})
|
||||
session, err := store.Create(CreateInput{SessionID: "terminal_session"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := store.CreateGoal(session.SessionID, "Finish", 0); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := store.UpdateGoal(session.SessionID, GoalStatusComplete, ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
updated, event, err := store.PauseGoalIfActive(session.SessionID, "cancelled")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if event != nil || updated.Goal.Status != GoalStatusComplete {
|
||||
t.Fatalf("terminal goal changed during cancellation: goal=%#v event=%#v", updated.Goal, event)
|
||||
}
|
||||
}
|
||||
@@ -43,7 +43,7 @@ func TestUpdateTitle(t *testing.T) {
|
||||
t.Fatalf("title not trimmed/stored: %q", updated.Title)
|
||||
}
|
||||
if updated.UpdatedAt != before.UpdatedAt {
|
||||
t.Fatalf("UpdatedAt must not change on retitle: before=%q after=%q", before.UpdatedAt, updated.UpdatedAt)
|
||||
t.Fatalf("UpdatedAt must not change on rename: before=%q after=%q", before.UpdatedAt, updated.UpdatedAt)
|
||||
}
|
||||
if updated.EventCount != before.EventCount {
|
||||
t.Fatalf("EventCount changed: before=%d after=%d", before.EventCount, updated.EventCount)
|
||||
@@ -67,7 +67,7 @@ func TestUpdateTitle(t *testing.T) {
|
||||
|
||||
// An unchanged title is a no-op (still succeeds).
|
||||
if _, err := store.UpdateTitle(session.SessionID, "Clean Generated Title"); err != nil {
|
||||
t.Fatalf("no-op retitle should succeed: %v", err)
|
||||
t.Fatalf("no-op rename should succeed: %v", err)
|
||||
}
|
||||
|
||||
// An invalid session id is rejected.
|
||||
@@ -76,6 +76,30 @@ func TestUpdateTitle(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateTitleIfCurrent(t *testing.T) {
|
||||
store := newTitleTestStore(t)
|
||||
session, err := store.Create(CreateInput{Title: "Automatic title"})
|
||||
if err != nil {
|
||||
t.Fatalf("create: %v", err)
|
||||
}
|
||||
|
||||
updated, applied, err := store.UpdateTitleIfCurrent(session.SessionID, "Automatic title", "Generated title")
|
||||
if err != nil || !applied || updated.Title != "Generated title" {
|
||||
t.Fatalf("matching update = (%#v, %v, %v)", updated, applied, err)
|
||||
}
|
||||
|
||||
if _, err := store.UpdateTitle(session.SessionID, "Manual title"); err != nil {
|
||||
t.Fatalf("manual rename: %v", err)
|
||||
}
|
||||
updated, applied, err = store.UpdateTitleIfCurrent(session.SessionID, "Generated title", "Late generated title")
|
||||
if err != nil {
|
||||
t.Fatalf("stale update: %v", err)
|
||||
}
|
||||
if applied || updated.Title != "Manual title" {
|
||||
t.Fatalf("stale update should preserve manual title: (%#v, %v)", updated, applied)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateModel(t *testing.T) {
|
||||
store := newTitleTestStore(t)
|
||||
session, err := store.Create(CreateInput{ModelID: "model-a"})
|
||||
|
||||
@@ -45,6 +45,9 @@ const (
|
||||
EventSpecDraft EventType = "spec_draft"
|
||||
EventSpecApproved EventType = "spec_approved"
|
||||
EventSpecRejected EventType = "spec_rejected"
|
||||
EventGoalCreated EventType = "goal_created"
|
||||
EventGoalUpdated EventType = "goal_updated"
|
||||
EventGoalCleared EventType = "goal_cleared"
|
||||
)
|
||||
|
||||
type SessionKind string
|
||||
@@ -65,6 +68,29 @@ const (
|
||||
SpecStatusRejected SpecStatus = "rejected"
|
||||
)
|
||||
|
||||
type GoalStatus string
|
||||
|
||||
const (
|
||||
GoalStatusActive GoalStatus = "active"
|
||||
GoalStatusPaused GoalStatus = "paused"
|
||||
GoalStatusBlocked GoalStatus = "blocked"
|
||||
GoalStatusBudgetLimited GoalStatus = "budget_limited"
|
||||
GoalStatusUsageLimited GoalStatus = "usage_limited"
|
||||
GoalStatusComplete GoalStatus = "complete"
|
||||
)
|
||||
|
||||
type Goal struct {
|
||||
Objective string `json:"objective"`
|
||||
Status GoalStatus `json:"status"`
|
||||
StatusReason string `json:"statusReason,omitempty"`
|
||||
TokenBudget int `json:"tokenBudget,omitempty"`
|
||||
TokensUsed int `json:"tokensUsed,omitempty"`
|
||||
ContinuationCount int `json:"continuationCount,omitempty"`
|
||||
ContinuationLimit int `json:"continuationLimit,omitempty"`
|
||||
CreatedAt string `json:"createdAt"`
|
||||
UpdatedAt string `json:"updatedAt"`
|
||||
}
|
||||
|
||||
type Metadata struct {
|
||||
SessionID string `json:"sessionId"`
|
||||
SessionKind SessionKind `json:"sessionKind,omitempty"`
|
||||
@@ -91,6 +117,7 @@ type Metadata struct {
|
||||
SpecRejectReason string `json:"specRejectReason,omitempty"`
|
||||
SpecSourceSessionID string `json:"specSourceSessionId,omitempty"`
|
||||
SpecImplSessionID string `json:"specImplSessionId,omitempty"`
|
||||
Goal *Goal `json:"goal,omitempty"`
|
||||
CreatedAt string `json:"createdAt"`
|
||||
UpdatedAt string `json:"updatedAt"`
|
||||
EventCount int `json:"eventCount"`
|
||||
@@ -123,6 +150,7 @@ type CreateInput struct {
|
||||
SpecRejectReason string
|
||||
SpecSourceSessionID string
|
||||
SpecImplSessionID string
|
||||
Goal *Goal
|
||||
}
|
||||
|
||||
type ForkInput struct {
|
||||
@@ -279,6 +307,7 @@ func (store *Store) Create(input CreateInput) (Metadata, error) {
|
||||
SpecRejectReason: strings.TrimSpace(input.SpecRejectReason),
|
||||
SpecSourceSessionID: strings.TrimSpace(input.SpecSourceSessionID),
|
||||
SpecImplSessionID: strings.TrimSpace(input.SpecImplSessionID),
|
||||
Goal: cloneGoal(input.Goal),
|
||||
CreatedAt: timestamp,
|
||||
UpdatedAt: timestamp,
|
||||
EventCount: 0,
|
||||
@@ -690,7 +719,7 @@ func (store *Store) appendPreparedEventsLocked(sessionID string, inputs []prepar
|
||||
// serialized under the same per-session lock as AppendEvent and re-reads the
|
||||
// latest metadata under that lock before rewriting, so a concurrent append can't
|
||||
// clobber the new title (nor the title clobber a concurrent append's event
|
||||
// count/timestamp). UpdatedAt is deliberately left untouched: a retitle is not
|
||||
// count/timestamp). UpdatedAt is deliberately left untouched: a rename is not
|
||||
// activity, so it must not reorder the session in the resumable list. A blank
|
||||
// title is rejected so a failed model generation can never erase a useful
|
||||
// first-message title, and an unchanged title is a no-op (no rewrite/fsync).
|
||||
@@ -722,6 +751,40 @@ func (store *Store) UpdateTitle(sessionID string, title string) (Metadata, error
|
||||
return session, nil
|
||||
}
|
||||
|
||||
// UpdateTitleIfCurrent replaces a session title only when it still matches
|
||||
// expected. It lets background automatic naming avoid overwriting a newer manual
|
||||
// rename while keeping the check and write under the same per-session lock.
|
||||
func (store *Store) UpdateTitleIfCurrent(sessionID string, expected string, title string) (Metadata, bool, error) {
|
||||
if !ValidSessionID(sessionID) {
|
||||
return Metadata{}, false, fmt.Errorf("invalid zero session id %q", sessionID)
|
||||
}
|
||||
trimmed := strings.TrimSpace(title)
|
||||
if trimmed == "" {
|
||||
return Metadata{}, false, fmt.Errorf("zero session title is required")
|
||||
}
|
||||
unlock, err := store.lockSession(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, false, err
|
||||
}
|
||||
defer unlock()
|
||||
|
||||
session, err := store.readMetadata(sessionID)
|
||||
if err != nil {
|
||||
return Metadata{}, false, err
|
||||
}
|
||||
if session.Title != strings.TrimSpace(expected) {
|
||||
return session, false, nil
|
||||
}
|
||||
if session.Title == trimmed {
|
||||
return session, true, nil
|
||||
}
|
||||
session.Title = trimmed
|
||||
if err := store.writeMetadata(session); err != nil {
|
||||
return Metadata{}, false, err
|
||||
}
|
||||
return session, true, nil
|
||||
}
|
||||
|
||||
// UpdateModel replaces a session's selected model without changing its activity
|
||||
// timestamp or event counters. An empty model clears the session override.
|
||||
func (store *Store) UpdateModel(sessionID string, modelID string) (Metadata, error) {
|
||||
|
||||
+180
-27
@@ -12,6 +12,12 @@ import (
|
||||
// team is at its slot cap the member is queued and launches when a slot frees
|
||||
// (it stays pending in the coordinator until then).
|
||||
func (s *Swarm) Spawn(pol Policy, teamName, agentType, task, cwd string) (string, error) {
|
||||
release, err := s.beginLifecycleAdmission()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer release()
|
||||
|
||||
def, err := s.registry.Lookup(agentType)
|
||||
if err != nil {
|
||||
return "", err
|
||||
@@ -23,15 +29,44 @@ func (s *Swarm) Spawn(pol Policy, teamName, agentType, task, cwd string) (string
|
||||
}
|
||||
s.rememberCwd(id, cwd)
|
||||
spec := s.buildSpec(pol, id, id, team, def, task, cwd)
|
||||
s.dispatch(spec)
|
||||
s.dispatchAdmitted(spec)
|
||||
return id, nil
|
||||
}
|
||||
|
||||
// beginLifecycleAdmission admits one unit of lifecycle work if shutdown has not
|
||||
// begun. On success it returns a release func the caller MUST call exactly once
|
||||
// (a deferred call is the norm) — Close blocks until every admitted unit has
|
||||
// released. On shutdown it returns a nil func and ErrSwarmClosed, and there is
|
||||
// nothing to release.
|
||||
//
|
||||
// It deliberately does NOT keep lifecycleMu held for the caller's duration, which
|
||||
// an earlier revision did. Close needs the write lock before it can cancel, so a
|
||||
// launcher that blocks in Launch until its context is cancelled would deadlock
|
||||
// against every Close caller: the launch holds the read lock while Close waits
|
||||
// for the write lock to issue the cancellation neither side can reach. Holding the
|
||||
// lock only across the flag check and the counter increment keeps the ordering
|
||||
// guarantee (no admission once closed is set) without letting external launcher
|
||||
// code sit inside the critical section.
|
||||
//
|
||||
// The read lock is enough to make the WaitGroup safe: Close flips closed under
|
||||
// the write lock, so an Add here either completes before Close acquires it (and is
|
||||
// therefore visible to the later Wait) or observes closed and is refused.
|
||||
func (s *Swarm) beginLifecycleAdmission() (func(), error) {
|
||||
s.lifecycleMu.RLock()
|
||||
defer s.lifecycleMu.RUnlock()
|
||||
if s.closed {
|
||||
return nil, ErrSwarmClosed
|
||||
}
|
||||
s.lifecycleWork.Add(1)
|
||||
return s.lifecycleWork.Done, nil
|
||||
}
|
||||
|
||||
// dispatch admits a spec to its team (launching now or queuing for a slot).
|
||||
func (s *Swarm) dispatch(spec MemberSpec) {
|
||||
// The caller must hold an admission ticket from beginLifecycleAdmission.
|
||||
func (s *Swarm) dispatchAdmitted(spec MemberSpec) {
|
||||
t := s.team(spec.Team)
|
||||
if t.admit(spec) {
|
||||
s.launch(t, spec)
|
||||
s.launchAdmitted(t, spec)
|
||||
}
|
||||
// Otherwise the spec is queued; the coordinator task stays pending until a
|
||||
// slot frees and afterExit launches it.
|
||||
@@ -39,17 +74,73 @@ func (s *Swarm) dispatch(spec MemberSpec) {
|
||||
|
||||
// launch starts a member for spec and supervises it. A synchronous launch
|
||||
// failure fails the task and frees the slot.
|
||||
func (s *Swarm) launch(t *Team, spec MemberSpec) {
|
||||
handle, err := s.launcher.Launch(s.baseCtx, spec)
|
||||
// The caller must hold an admission ticket from beginLifecycleAdmission.
|
||||
func (s *Swarm) launchAdmitted(t *Team, spec MemberSpec) {
|
||||
handle, err := s.launchMemberAdmitted(spec)
|
||||
if err != nil {
|
||||
_ = s.coord.Fail(spec.TaskID, "launch: "+err.Error())
|
||||
s.afterExit(t)
|
||||
if err == ErrSwarmClosed {
|
||||
t.releaseSlot()
|
||||
return
|
||||
}
|
||||
s.afterExitAdmitted(t)
|
||||
return
|
||||
}
|
||||
m := &Member{ID: spec.ID, AgentType: spec.AgentType, TaskID: spec.TaskID, handle: handle}
|
||||
t.addMember(m)
|
||||
_ = s.coord.SetStatus(spec.TaskID, StatusRunning)
|
||||
go s.watch(t, m, spec)
|
||||
|
||||
committed := s.commitLaunch(handle, func() {
|
||||
m := &Member{ID: spec.ID, AgentType: spec.AgentType, TaskID: spec.TaskID, handle: handle}
|
||||
t.addMember(m)
|
||||
_ = s.coord.SetStatus(spec.TaskID, StatusRunning)
|
||||
s.watchers.Add(1)
|
||||
go func() {
|
||||
defer s.watchers.Done()
|
||||
s.watch(t, m, spec)
|
||||
}()
|
||||
})
|
||||
if !committed {
|
||||
_ = s.coord.Fail(spec.TaskID, "launch: "+ErrSwarmClosed.Error())
|
||||
t.releaseSlot()
|
||||
}
|
||||
}
|
||||
|
||||
// commitLaunch closes the window between MemberLauncher.Launch returning and the
|
||||
// swarm taking ownership of the handle it produced. Launch is external code and
|
||||
// therefore runs without lifecycleMu held, so Close may have won after
|
||||
// launchMemberAdmitted's pre-check but before Launch returned. adopt runs under
|
||||
// the read lock so the shutdown check and everything the caller does to take
|
||||
// ownership (registration, status transition, watcher Add, handle swap) are
|
||||
// atomic with respect to Close setting closed and beginning its waits.
|
||||
//
|
||||
// It reports whether the handle was adopted. When shutdown won, the handle is
|
||||
// reaped before returning false: its launch context was already cancelled by
|
||||
// Close, and MemberHandle has no separate cancellation operation — Launch's
|
||||
// context is its cancellation contract — so waiting is how a
|
||||
// successfully-created process/goroutine avoids being abandoned. The caller
|
||||
// records the terminal outcome for the task it was launching.
|
||||
func (s *Swarm) commitLaunch(handle MemberHandle, adopt func()) bool {
|
||||
s.lifecycleMu.RLock()
|
||||
if s.closed {
|
||||
s.lifecycleMu.RUnlock()
|
||||
_, _ = handle.Wait()
|
||||
return false
|
||||
}
|
||||
defer s.lifecycleMu.RUnlock()
|
||||
adopt()
|
||||
return true
|
||||
}
|
||||
|
||||
// launchMemberAdmitted re-checks shutdown at the last internal boundary before
|
||||
// invoking the external launcher. A lifecycle ticket proves the operation won
|
||||
// admission before Close, but the caller may not reach its launch until after
|
||||
// Close has set closed; in that case no new member should be started.
|
||||
func (s *Swarm) launchMemberAdmitted(spec MemberSpec) (MemberHandle, error) {
|
||||
s.lifecycleMu.RLock()
|
||||
closed := s.closed
|
||||
s.lifecycleMu.RUnlock()
|
||||
if closed {
|
||||
return nil, ErrSwarmClosed
|
||||
}
|
||||
return s.launcher.Launch(s.baseCtx, spec)
|
||||
}
|
||||
|
||||
// watch awaits a member, applies bounded relaunch on temporary failures, records
|
||||
@@ -60,42 +151,98 @@ func (s *Swarm) launch(t *Team, spec MemberSpec) {
|
||||
// life (Member.ID == MemberSpec.ID); a retry never reuses the struct for a
|
||||
// different spec.
|
||||
func (s *Swarm) watch(t *Team, m *Member, spec MemberSpec) {
|
||||
res, err := m.handle.Wait()
|
||||
if err != nil {
|
||||
if isRetryable(err) && m.restarts < maxMemberRestarts {
|
||||
if nh, relErr := s.launcher.Launch(s.baseCtx, spec); relErr == nil {
|
||||
m.restarts++
|
||||
m.handle = nh
|
||||
go s.watch(t, m, spec)
|
||||
return
|
||||
for {
|
||||
res, err := m.handle.Wait()
|
||||
if err != nil {
|
||||
if isRetryable(err) && m.restarts < maxMemberRestarts && s.relaunchAdmitted(m, spec) {
|
||||
continue
|
||||
}
|
||||
// fall through to record the original error if relaunch fails
|
||||
// Fall through if shutdown started or the relaunch failed.
|
||||
// res.SessionID is preserved by the handle even on error, so a member that
|
||||
// ran then failed stays drillable; a pure launch error carries an empty id.
|
||||
_ = s.coord.FailWithSession(m.TaskID, memberError(err), res.SessionID)
|
||||
} else {
|
||||
_ = s.coord.CompleteWithSession(m.TaskID, res.Result, res.SessionID)
|
||||
}
|
||||
// res.SessionID is preserved by the handle even on error, so a member that
|
||||
// ran then failed stays drillable; a pure launch error carries an empty id.
|
||||
_ = s.coord.FailWithSession(m.TaskID, memberError(err), res.SessionID)
|
||||
} else {
|
||||
_ = s.coord.CompleteWithSession(m.TaskID, res.Result, res.SessionID)
|
||||
break
|
||||
}
|
||||
t.removeMember(m.ID)
|
||||
s.afterExit(t)
|
||||
}
|
||||
|
||||
// relaunchAdmitted performs one bounded relaunch of an existing member's spec,
|
||||
// reporting whether the member was relaunched and its new handle adopted. A
|
||||
// relaunch is lifecycle work like any other, so it runs the same three gates an
|
||||
// initial launch does: admission (refused outright once shutdown began), the
|
||||
// pre-Launch re-check, and the post-Launch commit — the last of which matters
|
||||
// because the production FuncLauncher returns a live handle immediately without
|
||||
// consulting its context, so a Close landing between the pre-check and the return
|
||||
// would otherwise resume supervising a member started after shutdown. When any
|
||||
// gate refuses, the caller records the member's terminal failure instead.
|
||||
func (s *Swarm) relaunchAdmitted(m *Member, spec MemberSpec) bool {
|
||||
release, err := s.beginLifecycleAdmission()
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
defer release()
|
||||
handle, err := s.launchMemberAdmitted(spec)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return s.commitLaunch(handle, func() {
|
||||
m.restarts++
|
||||
m.handle = handle
|
||||
})
|
||||
}
|
||||
|
||||
// afterExit releases the just-vacated slot and launches the next queued spec, if
|
||||
// any. Each exit drains at most one queued member; that member's own exit drains
|
||||
// the next, so the queue empties one-per-slot without unbounded recursion.
|
||||
func (s *Swarm) afterExit(t *Team) {
|
||||
release, err := s.beginLifecycleAdmission()
|
||||
if err != nil {
|
||||
t.releaseSlot()
|
||||
return
|
||||
}
|
||||
defer release()
|
||||
s.afterExitAdmitted(t)
|
||||
}
|
||||
|
||||
// afterExitAdmitted drains at most one queued spec while lifecycle admission is
|
||||
// held open. The caller must hold an admission ticket from beginLifecycleAdmission.
|
||||
func (s *Swarm) afterExitAdmitted(t *Team) {
|
||||
next, ok := t.onExit()
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
s.launch(t, next)
|
||||
// This call's own ticket only proves shutdown hadn't begun when it was
|
||||
// granted; Close may have set closed in the meantime, and t.onExit above
|
||||
// already removed next from t.queue, so Close's clearQueue sweep can no
|
||||
// longer see it (nor will it wait for one, since this ticket is what Close's
|
||||
// lifecycleWork.Wait() is blocked on). Re-check before launching so a dequeue
|
||||
// that straddles the shutdown boundary fails the task in place of a launch a
|
||||
// shutting-down swarm never wanted started.
|
||||
s.lifecycleMu.RLock()
|
||||
closed := s.closed
|
||||
s.lifecycleMu.RUnlock()
|
||||
if closed {
|
||||
t.releaseSlot()
|
||||
_ = s.coord.Fail(next.TaskID, ErrSwarmClosed.Error())
|
||||
return
|
||||
}
|
||||
s.launchAdmitted(t, next)
|
||||
}
|
||||
|
||||
// Handoff transfers a task to a fresh member of toAgentType, delivering a note to
|
||||
// the new member's inbox and marking the original task handed-off. It returns the
|
||||
// new task id. A handoff of an already-terminal task is rejected (fail closed).
|
||||
func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string) (string, error) {
|
||||
release, err := s.beginLifecycleAdmission()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer release()
|
||||
|
||||
task, ok := s.coord.Get(taskID)
|
||||
if !ok {
|
||||
return "", fmt.Errorf("%w: %s", ErrUnknownTask, taskID)
|
||||
@@ -131,7 +278,7 @@ func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string)
|
||||
// Retire the original task (it has been re-delegated).
|
||||
_ = s.coord.SetStatus(taskID, StatusHandedOff)
|
||||
spec := s.buildSpec(pol, newID, newID, team, def, handoffTask, cwd)
|
||||
s.dispatch(spec)
|
||||
s.dispatchAdmitted(spec)
|
||||
return newID, nil
|
||||
}
|
||||
|
||||
@@ -139,6 +286,12 @@ func (s *Swarm) Handoff(pol Policy, teamName, taskID, toAgentType, note string)
|
||||
// (e.g. a crashed worker) onto fresh members of toAgentType, returning the
|
||||
// adopted task ids. Terminal tasks and tasks with a live owner are left alone.
|
||||
func (s *Swarm) AdoptOrphans(pol Policy, teamName, toAgentType string) ([]string, error) {
|
||||
release, err := s.beginLifecycleAdmission()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer release()
|
||||
|
||||
def, err := s.registry.Lookup(toAgentType)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -161,7 +314,7 @@ func (s *Swarm) AdoptOrphans(pol Policy, teamName, toAgentType string) ([]string
|
||||
cwd := s.cwdFor(task.ID)
|
||||
s.rememberCwd(task.ID, cwd)
|
||||
spec := s.buildSpec(pol, newAgent, task.ID, team, def, task.Description, cwd)
|
||||
s.dispatch(spec)
|
||||
s.dispatchAdmitted(spec)
|
||||
adopted = append(adopted, task.ID)
|
||||
}
|
||||
return adopted, nil
|
||||
|
||||
@@ -2,6 +2,7 @@ package swarm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
@@ -217,6 +218,141 @@ func TestPermanentErrorNoRetry(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestLifecycleAdmissionRejectedAfterClose(t *testing.T) {
|
||||
l := newLauncher(okFor)
|
||||
sw := newSwarmFor(t, l)
|
||||
sw.Close()
|
||||
|
||||
if _, err := sw.Spawn(Policy{}, "team", "teammate", "task", ""); !errors.Is(err, ErrSwarmClosed) {
|
||||
t.Fatalf("Spawn after Close error = %v, want ErrSwarmClosed", err)
|
||||
}
|
||||
if _, err := sw.Handoff(Policy{}, "team", "task", "teammate", "note"); !errors.Is(err, ErrSwarmClosed) {
|
||||
t.Fatalf("Handoff after Close error = %v, want ErrSwarmClosed", err)
|
||||
}
|
||||
if _, err := sw.AdoptOrphans(Policy{}, "team", "teammate"); !errors.Is(err, ErrSwarmClosed) {
|
||||
t.Fatalf("AdoptOrphans after Close error = %v, want ErrSwarmClosed", err)
|
||||
}
|
||||
if got := len(l.recorded()); got != 0 {
|
||||
t.Fatalf("launches after Close = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCloseDoesNotLaunchQueuedMembers(t *testing.T) {
|
||||
gate := make(chan struct{})
|
||||
l := newLauncher(okFor)
|
||||
l.gate = gate
|
||||
sw := newSwarmFor(t, l)
|
||||
|
||||
var ids []string
|
||||
for i := 0; i < 3; i++ {
|
||||
id, err := sw.Spawn(Policy{}, "team", "teammate", "task", "")
|
||||
if err != nil {
|
||||
t.Fatalf("Spawn %d: %v", i, err)
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
if got := len(l.recorded()); got != 2 {
|
||||
t.Fatalf("initial launches = %d, want 2 with one queued", got)
|
||||
}
|
||||
|
||||
sw.Close()
|
||||
if got := len(l.recorded()); got != 2 {
|
||||
t.Fatalf("launches after Close = %d, want queued member not launched", got)
|
||||
}
|
||||
for _, id := range ids {
|
||||
task, ok := sw.Coordinator().Get(id)
|
||||
if !ok || !task.Status.terminal() {
|
||||
t.Fatalf("task %s after Close = %+v, want terminal", id, task)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestClosePreventsMemberRetry(t *testing.T) {
|
||||
started := make(chan struct{}, maxMemberRestarts+1)
|
||||
release := make(chan struct{})
|
||||
l := FuncLauncher{Run: func(context.Context, MemberSpec) (MemberResult, error) {
|
||||
started <- struct{}{}
|
||||
<-release
|
||||
return MemberResult{}, ErrMemberTemporary
|
||||
}}
|
||||
sw := newSwarmFor(t, l)
|
||||
_, err := sw.Spawn(Policy{}, "team", "teammate", "task", "")
|
||||
if err != nil {
|
||||
t.Fatalf("Spawn: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("initial member did not start")
|
||||
}
|
||||
|
||||
closed := make(chan struct{})
|
||||
go func() {
|
||||
sw.Close()
|
||||
close(closed)
|
||||
}()
|
||||
waitFor(t, "swarm closed state", func() bool {
|
||||
sw.lifecycleMu.RLock()
|
||||
defer sw.lifecycleMu.RUnlock()
|
||||
return sw.closed
|
||||
})
|
||||
close(release)
|
||||
select {
|
||||
case <-closed:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("Close did not return after retryable member exit")
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
t.Fatal("member retried after Close")
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func TestCloseWaitsForMemberWatchers(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
l := FuncLauncher{Run: func(context.Context, MemberSpec) (MemberResult, error) {
|
||||
<-release
|
||||
return MemberResult{Result: "done"}, nil
|
||||
}}
|
||||
sw := newSwarmFor(t, l)
|
||||
id, err := sw.Spawn(Policy{}, "team", "teammate", "task", "")
|
||||
if err != nil {
|
||||
t.Fatalf("Spawn: %v", err)
|
||||
}
|
||||
waitFor(t, "task running", func() bool {
|
||||
task, ok := sw.Coordinator().Get(id)
|
||||
return ok && task.Status == StatusRunning
|
||||
})
|
||||
|
||||
const callers = 3
|
||||
closing := make(chan struct{}, callers)
|
||||
closed := make(chan struct{}, callers)
|
||||
for i := 0; i < callers; i++ {
|
||||
go func() {
|
||||
closing <- struct{}{}
|
||||
sw.Close()
|
||||
closed <- struct{}{}
|
||||
}()
|
||||
}
|
||||
for i := 0; i < callers; i++ {
|
||||
<-closing
|
||||
}
|
||||
select {
|
||||
case <-closed:
|
||||
t.Fatal("a Close caller returned before the member watcher exited")
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
close(release)
|
||||
for i := 0; i < callers; i++ {
|
||||
select {
|
||||
case <-closed:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("all Close callers did not return after the member watcher exited")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandoffDeliversNoteAndRetiresOriginal(t *testing.T) {
|
||||
gate := make(chan struct{})
|
||||
l := newLauncher(okFor)
|
||||
@@ -364,3 +500,587 @@ func TestSpawnUnknownAgentType(t *testing.T) {
|
||||
t.Fatal("Spawn with unknown agent type must error")
|
||||
}
|
||||
}
|
||||
|
||||
// blockingLauncher blocks inside Launch until its context is cancelled — the
|
||||
// shape of a real launcher that waits on a slot, a daemon connection, or a
|
||||
// sandbox handshake before it can return a handle.
|
||||
type blockingLauncher struct {
|
||||
entered chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (l *blockingLauncher) Launch(ctx context.Context, spec MemberSpec) (MemberHandle, error) {
|
||||
l.once.Do(func() { close(l.entered) })
|
||||
<-ctx.Done()
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
// TestCloseCancelsLauncherBlockedInLaunch is the regression test for holding
|
||||
// lifecycle admission across MemberLauncher.Launch. When admission was a read lock
|
||||
// held for the caller's whole duration, this deadlocked: the spawn sat inside
|
||||
// Launch waiting for its context, while Close waited for the write lock it needed
|
||||
// before it could cancel that very context. Neither side could progress, so both
|
||||
// the spawn and every Close caller hung.
|
||||
//
|
||||
// Admission is now a counted ticket, so Close can flip the flag, cancel, and then
|
||||
// wait the ticket out.
|
||||
func TestCloseCancelsLauncherBlockedInLaunch(t *testing.T) {
|
||||
launcher := &blockingLauncher{entered: make(chan struct{})}
|
||||
sw := newSwarmFor(t, launcher)
|
||||
|
||||
spawned := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := sw.Spawn(Policy{}, "team", "teammate", "task", "")
|
||||
spawned <- err
|
||||
}()
|
||||
select {
|
||||
case <-launcher.entered:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("launcher never entered Launch")
|
||||
}
|
||||
|
||||
closed := make(chan struct{})
|
||||
go func() {
|
||||
sw.Close()
|
||||
close(closed)
|
||||
}()
|
||||
select {
|
||||
case <-closed:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("Close deadlocked against a launcher blocked in Launch")
|
||||
}
|
||||
select {
|
||||
case <-spawned:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("Spawn never returned after Close cancelled the launch context")
|
||||
}
|
||||
}
|
||||
|
||||
// TestCloseWaitsForAdmittedSpawn pins the other half of the contract: Close is a
|
||||
// barrier, so it must not return while an admitted spawn is still running. Without
|
||||
// the lifecycleWork wait, Close could finish while this launch was mid-flight.
|
||||
func TestCloseWaitsForAdmittedSpawn(t *testing.T) {
|
||||
entered := make(chan struct{})
|
||||
finish := make(chan struct{})
|
||||
launcher := &gatedLaunchLauncher{entered: entered, finish: finish}
|
||||
sw := newSwarmFor(t, launcher)
|
||||
|
||||
go func() {
|
||||
_, _ = sw.Spawn(Policy{}, "team", "teammate", "task", "")
|
||||
}()
|
||||
select {
|
||||
case <-entered:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("launcher never entered Launch")
|
||||
}
|
||||
|
||||
closed := make(chan struct{})
|
||||
go func() {
|
||||
sw.Close()
|
||||
close(closed)
|
||||
}()
|
||||
select {
|
||||
case <-closed:
|
||||
t.Fatal("Close returned while an admitted spawn was still in flight")
|
||||
case <-time.After(200 * time.Millisecond):
|
||||
}
|
||||
close(finish)
|
||||
select {
|
||||
case <-closed:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("Close did not return after the admitted spawn finished")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdmittedDispatchDoesNotLaunchAfterClose covers an operation that won a
|
||||
// lifecycle ticket before shutdown but did not reach dispatch until after Close
|
||||
// set closed. The ticket keeps Close from returning early, but it must not permit
|
||||
// a new external launch with the already-cancelled swarm context.
|
||||
func TestAdmittedDispatchDoesNotLaunchAfterClose(t *testing.T) {
|
||||
l := newLauncher(okFor)
|
||||
sw := newSwarmFor(t, l)
|
||||
sw.maxTeamSize = 1
|
||||
|
||||
if _, err := sw.coord.Register("late-task", "late-agent", "team", "late work"); err != nil {
|
||||
t.Fatalf("Register late task: %v", err)
|
||||
}
|
||||
release, err := sw.beginLifecycleAdmission()
|
||||
if err != nil {
|
||||
t.Fatalf("beginLifecycleAdmission: %v", err)
|
||||
}
|
||||
|
||||
closed := make(chan struct{})
|
||||
go func() {
|
||||
sw.Close()
|
||||
close(closed)
|
||||
}()
|
||||
waitFor(t, "swarm closed state", func() bool {
|
||||
sw.lifecycleMu.RLock()
|
||||
defer sw.lifecycleMu.RUnlock()
|
||||
return sw.closed
|
||||
})
|
||||
|
||||
sw.dispatchAdmitted(MemberSpec{
|
||||
ID: "late-agent", TaskID: "late-task", AgentType: "teammate", Team: "team",
|
||||
})
|
||||
release()
|
||||
|
||||
select {
|
||||
case <-closed:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("Close did not return after the admitted dispatch finished")
|
||||
}
|
||||
if got := len(l.recorded()); got != 0 {
|
||||
t.Fatalf("launches after closed was set = %d, want 0", got)
|
||||
}
|
||||
team := sw.team("team")
|
||||
if got := team.Running(); got != 0 {
|
||||
t.Fatalf("team running after rejected launch = %d, want 0", got)
|
||||
}
|
||||
task, ok := sw.Coordinator().Get("late-task")
|
||||
if !ok {
|
||||
t.Fatal("late task vanished")
|
||||
}
|
||||
if task.Status != StatusFailed || !strings.Contains(task.Err, ErrSwarmClosed.Error()) {
|
||||
t.Fatalf("late task after Close = %+v, want failed with ErrSwarmClosed", task)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCloseBetweenLaunchPrecheckAndReturn rejects and reaps a handle created by
|
||||
// a Launch that straddles shutdown. The launch has passed the pre-check when it
|
||||
// enters the gate; Close then sets closed and cancels its context before the
|
||||
// launcher is allowed to return successfully.
|
||||
func TestCloseBetweenLaunchPrecheckAndReturn(t *testing.T) {
|
||||
launcher := &closeRaceLauncher{
|
||||
entered: make(chan struct{}),
|
||||
finish: make(chan struct{}),
|
||||
waited: make(chan struct{}),
|
||||
}
|
||||
sw := newSwarmFor(t, launcher)
|
||||
|
||||
spawned := make(chan string, 1)
|
||||
go func() {
|
||||
id, _ := sw.Spawn(Policy{}, "team", "teammate", "task", "")
|
||||
spawned <- id
|
||||
}()
|
||||
select {
|
||||
case <-launcher.entered:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("launcher never entered Launch")
|
||||
}
|
||||
|
||||
closed := make(chan struct{})
|
||||
go func() {
|
||||
sw.Close()
|
||||
close(closed)
|
||||
}()
|
||||
waitFor(t, "swarm closed state", func() bool {
|
||||
sw.lifecycleMu.RLock()
|
||||
defer sw.lifecycleMu.RUnlock()
|
||||
return sw.closed
|
||||
})
|
||||
close(launcher.finish)
|
||||
|
||||
var id string
|
||||
select {
|
||||
case id = <-spawned:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("Spawn did not return after Launch")
|
||||
}
|
||||
select {
|
||||
case <-closed:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("Close did not return after late handle was reaped")
|
||||
}
|
||||
select {
|
||||
case <-launcher.waited:
|
||||
default:
|
||||
t.Fatal("late handle was not reaped")
|
||||
}
|
||||
|
||||
team := sw.team("team")
|
||||
if got := team.Running(); got != 0 {
|
||||
t.Fatalf("team running after close-raced Launch = %d, want 0", got)
|
||||
}
|
||||
if got := len(team.liveAgents()); got != 0 {
|
||||
t.Fatalf("registered members after close-raced Launch = %d, want 0", got)
|
||||
}
|
||||
task, ok := sw.Coordinator().Get(id)
|
||||
if !ok {
|
||||
t.Fatal("task vanished")
|
||||
}
|
||||
if task.Status != StatusFailed || !strings.Contains(task.Err, ErrSwarmClosed.Error()) {
|
||||
t.Fatalf("task after close-raced Launch = %+v, want failed with ErrSwarmClosed", task)
|
||||
}
|
||||
}
|
||||
|
||||
type closeRaceLauncher struct {
|
||||
entered chan struct{}
|
||||
finish chan struct{}
|
||||
waited chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (l *closeRaceLauncher) Launch(ctx context.Context, spec MemberSpec) (MemberHandle, error) {
|
||||
l.once.Do(func() { close(l.entered) })
|
||||
<-l.finish
|
||||
return &closeRaceHandle{id: spec.ID, ctx: ctx, waited: l.waited}, nil
|
||||
}
|
||||
|
||||
type closeRaceHandle struct {
|
||||
id string
|
||||
ctx context.Context
|
||||
waited chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (h *closeRaceHandle) ID() string { return h.id }
|
||||
|
||||
func (h *closeRaceHandle) Wait() (MemberResult, error) {
|
||||
<-h.ctx.Done()
|
||||
h.once.Do(func() { close(h.waited) })
|
||||
return MemberResult{}, h.ctx.Err()
|
||||
}
|
||||
|
||||
// TestCloseBetweenRetryLaunchPrecheckAndReturn is the retry counterpart of
|
||||
// TestCloseBetweenLaunchPrecheckAndReturn: a bounded relaunch wins admission and
|
||||
// passes the pre-check, then Close lands while the launcher is inside Launch. The
|
||||
// relaunched handle must be reaped rather than adopted, so the watcher never
|
||||
// resumes supervising a member started after shutdown began — a relaunch whose
|
||||
// member goes on to succeed must not be recorded as the task's outcome.
|
||||
func TestCloseBetweenRetryLaunchPrecheckAndReturn(t *testing.T) {
|
||||
launcher := &retryCloseRaceLauncher{
|
||||
entered: make(chan struct{}),
|
||||
finish: make(chan struct{}),
|
||||
late: &retryLateHandle{id: "teammate-1"},
|
||||
}
|
||||
sw := newSwarmFor(t, launcher)
|
||||
|
||||
id, err := sw.Spawn(Policy{}, "team", "teammate", "task", "")
|
||||
if err != nil {
|
||||
t.Fatalf("Spawn: %v", err)
|
||||
}
|
||||
// The first member exits retryable, so the watcher is now inside the relaunch.
|
||||
select {
|
||||
case <-launcher.entered:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("launcher never entered the relaunch Launch")
|
||||
}
|
||||
|
||||
closed := make(chan struct{})
|
||||
go func() {
|
||||
sw.Close()
|
||||
close(closed)
|
||||
}()
|
||||
waitFor(t, "swarm closed state", func() bool {
|
||||
sw.lifecycleMu.RLock()
|
||||
defer sw.lifecycleMu.RUnlock()
|
||||
return sw.closed
|
||||
})
|
||||
close(launcher.finish)
|
||||
|
||||
select {
|
||||
case <-closed:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("Close did not return after the late relaunch handle was reaped")
|
||||
}
|
||||
if got := launcher.launchCount(); got != 2 {
|
||||
t.Fatalf("launch attempts = %d, want 2 (initial + one relaunch)", got)
|
||||
}
|
||||
if got := launcher.late.waitCount(); got != 1 {
|
||||
t.Fatalf("late relaunch handle waits = %d, want 1 (reaped, not supervised)", got)
|
||||
}
|
||||
|
||||
team := sw.team("team")
|
||||
if got := team.Running(); got != 0 {
|
||||
t.Fatalf("team running after close-raced relaunch = %d, want 0", got)
|
||||
}
|
||||
if got := len(team.liveAgents()); got != 0 {
|
||||
t.Fatalf("registered members after close-raced relaunch = %d, want 0", got)
|
||||
}
|
||||
task, ok := sw.Coordinator().Get(id)
|
||||
if !ok {
|
||||
t.Fatal("task vanished")
|
||||
}
|
||||
if task.Status != StatusFailed {
|
||||
t.Fatalf("task after close-raced relaunch = %+v, want failed, not the late member's outcome", task)
|
||||
}
|
||||
if task.Result != "" {
|
||||
t.Fatalf("task result after close-raced relaunch = %q, want the late member's result discarded", task.Result)
|
||||
}
|
||||
}
|
||||
|
||||
// retryCloseRaceLauncher fails its first member with a retryable error, then
|
||||
// blocks inside the relaunch's Launch until the test releases it.
|
||||
type retryCloseRaceLauncher struct {
|
||||
mu sync.Mutex
|
||||
calls int
|
||||
entered chan struct{}
|
||||
finish chan struct{}
|
||||
late *retryLateHandle
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (l *retryCloseRaceLauncher) Launch(_ context.Context, spec MemberSpec) (MemberHandle, error) {
|
||||
l.mu.Lock()
|
||||
l.calls++
|
||||
first := l.calls == 1
|
||||
l.mu.Unlock()
|
||||
if first {
|
||||
return &temporaryHandle{id: spec.ID}, nil
|
||||
}
|
||||
l.once.Do(func() { close(l.entered) })
|
||||
<-l.finish
|
||||
return l.late, nil
|
||||
}
|
||||
|
||||
func (l *retryCloseRaceLauncher) launchCount() int {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
return l.calls
|
||||
}
|
||||
|
||||
// temporaryHandle exits immediately with a retryable error, driving the watcher
|
||||
// into its bounded relaunch path.
|
||||
type temporaryHandle struct{ id string }
|
||||
|
||||
func (h *temporaryHandle) ID() string { return h.id }
|
||||
|
||||
func (h *temporaryHandle) Wait() (MemberResult, error) {
|
||||
return MemberResult{}, ErrMemberTemporary
|
||||
}
|
||||
|
||||
// retryLateHandle models a member whose work does not consult the cancelled
|
||||
// launch context and reports success. Adopting it after shutdown would record a
|
||||
// post-close result on the task, so a correct shutdown reaps it exactly once and
|
||||
// discards what it returns.
|
||||
type retryLateHandle struct {
|
||||
id string
|
||||
mu sync.Mutex
|
||||
waits int
|
||||
}
|
||||
|
||||
func (h *retryLateHandle) ID() string { return h.id }
|
||||
|
||||
func (h *retryLateHandle) Wait() (MemberResult, error) {
|
||||
h.mu.Lock()
|
||||
h.waits++
|
||||
h.mu.Unlock()
|
||||
return MemberResult{Result: "late result", SessionID: "late-session"}, nil
|
||||
}
|
||||
|
||||
func (h *retryLateHandle) waitCount() int {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
return h.waits
|
||||
}
|
||||
|
||||
func TestCloseRacesLifecycleAdmission(t *testing.T) {
|
||||
gate := make(chan struct{})
|
||||
l := newLauncher(okFor)
|
||||
l.gate = gate
|
||||
sw := newSwarmFor(t, l)
|
||||
|
||||
const spawns = 64
|
||||
start := make(chan struct{})
|
||||
results := make(chan struct {
|
||||
id string
|
||||
err error
|
||||
}, spawns)
|
||||
var callers sync.WaitGroup
|
||||
callers.Add(spawns)
|
||||
for i := 0; i < spawns; i++ {
|
||||
go func() {
|
||||
defer callers.Done()
|
||||
<-start
|
||||
id, err := sw.Spawn(Policy{}, "team", "teammate", "racing task", "")
|
||||
results <- struct {
|
||||
id string
|
||||
err error
|
||||
}{id: id, err: err}
|
||||
}()
|
||||
}
|
||||
closed := make(chan struct{})
|
||||
go func() {
|
||||
<-start
|
||||
sw.Close()
|
||||
close(closed)
|
||||
}()
|
||||
|
||||
close(start)
|
||||
callers.Wait()
|
||||
<-closed
|
||||
close(results)
|
||||
|
||||
for result := range results {
|
||||
if result.err != nil {
|
||||
if !errors.Is(result.err, ErrSwarmClosed) {
|
||||
t.Fatalf("Spawn racing Close error = %v, want ErrSwarmClosed", result.err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
task, ok := sw.Coordinator().Get(result.id)
|
||||
if !ok || !task.Status.terminal() {
|
||||
t.Fatalf("admitted task %q after Close = %+v, want terminal", result.id, task)
|
||||
}
|
||||
}
|
||||
for _, task := range sw.Coordinator().List() {
|
||||
if !task.Status.terminal() {
|
||||
t.Fatalf("task %q remained %s after Close", task.ID, task.Status)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestCloseFailsQueuedSpecOnLateCreatedTeam is the regression test for queue
|
||||
// sweeping after shutdown begins. Spawn/Handoff/AdoptOrphans calls that obtained
|
||||
// tickets before closed flipped can still create a team and append to its queue.
|
||||
// Close must wait for those calls before both snapshotting teams and sweeping
|
||||
// their queues, or a late-created team's queued task remains pending forever.
|
||||
//
|
||||
// This drives that interleaving directly: two operations obtain tickets before
|
||||
// Close, then the first creates and fills a previously unseen one-slot team and
|
||||
// the second queues behind it after Close has flipped closed.
|
||||
func TestCloseFailsQueuedSpecOnLateCreatedTeam(t *testing.T) {
|
||||
sw := newSwarmFor(t, newLauncher(okFor))
|
||||
sw.maxTeamSize = 1
|
||||
|
||||
if _, err := sw.coord.Register("late-task", "late-agent", "team", "late work"); err != nil {
|
||||
t.Fatalf("Register late task: %v", err)
|
||||
}
|
||||
firstRelease, err := sw.beginLifecycleAdmission()
|
||||
if err != nil {
|
||||
t.Fatalf("first beginLifecycleAdmission: %v", err)
|
||||
}
|
||||
secondRelease, err := sw.beginLifecycleAdmission()
|
||||
if err != nil {
|
||||
firstRelease()
|
||||
t.Fatalf("second beginLifecycleAdmission: %v", err)
|
||||
}
|
||||
runningSpec := MemberSpec{ID: "running-agent", TaskID: "running-task", AgentType: "teammate", Team: "team"}
|
||||
queuedSpec := MemberSpec{ID: "late-agent", TaskID: "late-task", AgentType: "teammate", Team: "team"}
|
||||
|
||||
closed := make(chan struct{})
|
||||
go func() {
|
||||
sw.Close()
|
||||
close(closed)
|
||||
}()
|
||||
waitFor(t, "swarm closed state", func() bool {
|
||||
sw.lifecycleMu.RLock()
|
||||
defer sw.lifecycleMu.RUnlock()
|
||||
return sw.closed
|
||||
})
|
||||
select {
|
||||
case <-closed:
|
||||
t.Fatal("Close returned while an admitted ticket was still open")
|
||||
default:
|
||||
}
|
||||
|
||||
// Neither operation touched a team before Close took its old, buggy snapshot.
|
||||
// Finish the team-admission portion of both operations while their tickets are
|
||||
// still held: the first creates the team and fills its only slot; the second
|
||||
// appends to that late-created team's queue.
|
||||
team := sw.team("team")
|
||||
if !team.admit(runningSpec) {
|
||||
t.Fatal("first spec should take the late-created team's only slot")
|
||||
}
|
||||
firstRelease()
|
||||
if team.admit(queuedSpec) {
|
||||
t.Fatal("second spec should queue on the late-created team")
|
||||
}
|
||||
secondRelease()
|
||||
|
||||
select {
|
||||
case <-closed:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("Close did not return after the admitted ticket released")
|
||||
}
|
||||
|
||||
if got := team.QueueDepth(); got != 0 {
|
||||
t.Fatalf("queue depth after Close = %d, want 0 (late spec swept)", got)
|
||||
}
|
||||
task, ok := sw.Coordinator().Get("late-task")
|
||||
if !ok {
|
||||
t.Fatal("late-queued task vanished")
|
||||
}
|
||||
if task.Status != StatusFailed {
|
||||
t.Fatalf("late-queued task status = %v, want %v (failed, not stranded)", task.Status, StatusFailed)
|
||||
}
|
||||
if !strings.Contains(task.Err, ErrSwarmClosed.Error()) {
|
||||
t.Fatalf("late-queued task err = %q, want it to mention %q", task.Err, ErrSwarmClosed.Error())
|
||||
}
|
||||
}
|
||||
|
||||
// TestAfterExitAdmittedFailsDequeuedSpecAfterClose is the regression test for
|
||||
// jatmn's second #776 finding: afterExitAdmitted used to dequeue a queued spec
|
||||
// (t.onExit) and launch it unconditionally once its own ticket was granted,
|
||||
// even though Close can flip closed in the gap between that dequeue — which
|
||||
// makes the spec invisible to Close's clearQueue sweep — and the launch. A
|
||||
// launcher that performs work before honoring its cancelled context could then
|
||||
// start a brand new worker after shutdown had already begun.
|
||||
func TestAfterExitAdmittedFailsDequeuedSpecAfterClose(t *testing.T) {
|
||||
l := newLauncher(okFor)
|
||||
sw := newSwarmFor(t, l)
|
||||
|
||||
if _, err := sw.coord.Register("queued-task", "queued-agent", "team", "queued work"); err != nil {
|
||||
t.Fatalf("Register: %v", err)
|
||||
}
|
||||
spec := MemberSpec{ID: "queued-agent", TaskID: "queued-task", AgentType: "teammate", Team: "team"}
|
||||
|
||||
// Occupy the only slot, then queue spec behind it.
|
||||
team := &Team{Name: "team", members: map[string]*Member{}, maxSize: 1}
|
||||
if !team.admit(MemberSpec{ID: "running-agent", TaskID: "running-task"}) {
|
||||
t.Fatal("first spec should take the only slot")
|
||||
}
|
||||
if team.admit(spec) {
|
||||
t.Fatal("second spec should queue, not launch, over the cap")
|
||||
}
|
||||
|
||||
// Simulate shutdown racing in between this call's ticket admission (already
|
||||
// granted, hence no beginLifecycleAdmission call here) and its dequeue+launch,
|
||||
// exactly like the finding describes: closed is already set by the time
|
||||
// afterExitAdmitted's post-dequeue check runs.
|
||||
sw.lifecycleMu.Lock()
|
||||
sw.closed = true
|
||||
sw.lifecycleMu.Unlock()
|
||||
|
||||
sw.afterExitAdmitted(team) // running-agent "exits" -> dequeues spec
|
||||
|
||||
if got := team.Running(); got != 0 {
|
||||
t.Fatalf("team running after close-raced drain = %d, want 0 (slot released, not held for a launch)", got)
|
||||
}
|
||||
task, ok := sw.Coordinator().Get("queued-task")
|
||||
if !ok {
|
||||
t.Fatal("dequeued task vanished")
|
||||
}
|
||||
if task.Status != StatusFailed {
|
||||
t.Fatalf("dequeued task status = %v, want %v (failed, not launched)", task.Status, StatusFailed)
|
||||
}
|
||||
if !strings.Contains(task.Err, ErrSwarmClosed.Error()) {
|
||||
t.Fatalf("dequeued task err = %q, want it to mention %q", task.Err, ErrSwarmClosed.Error())
|
||||
}
|
||||
for _, s := range l.recorded() {
|
||||
if s.ID == "queued-agent" {
|
||||
t.Fatal("a spec dequeued after shutdown began must never launch")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// gatedLaunchLauncher blocks in Launch until the test releases it, ignoring
|
||||
// context cancellation so the test controls exactly when the admitted work ends.
|
||||
type gatedLaunchLauncher struct {
|
||||
entered chan struct{}
|
||||
finish chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (l *gatedLaunchLauncher) Launch(_ context.Context, spec MemberSpec) (MemberHandle, error) {
|
||||
l.once.Do(func() { close(l.entered) })
|
||||
<-l.finish
|
||||
return &funcHandle{id: spec.ID, done: closedChan()}, nil
|
||||
}
|
||||
|
||||
func closedChan() chan struct{} {
|
||||
ch := make(chan struct{})
|
||||
close(ch)
|
||||
return ch
|
||||
}
|
||||
+12
-13
@@ -121,11 +121,12 @@ type Scheduler struct {
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
|
||||
mu sync.Mutex
|
||||
jobs map[string]*scheduledJob
|
||||
closed bool
|
||||
seq atomic.Uint64
|
||||
wg sync.WaitGroup
|
||||
mu sync.Mutex
|
||||
jobs map[string]*scheduledJob
|
||||
closed bool
|
||||
seq atomic.Uint64
|
||||
wg sync.WaitGroup
|
||||
closeOnce sync.Once
|
||||
}
|
||||
|
||||
// newScheduler builds a Scheduler bound to sw. Its context derives from the
|
||||
@@ -229,15 +230,13 @@ func (s *Scheduler) List() []JobStatus {
|
||||
// Close cancels every job and waits for their loops to exit. Safe to call more
|
||||
// than once.
|
||||
func (s *Scheduler) Close() {
|
||||
s.mu.Lock()
|
||||
if s.closed {
|
||||
s.closeOnce.Do(func() {
|
||||
s.mu.Lock()
|
||||
s.closed = true
|
||||
s.cancel()
|
||||
s.mu.Unlock()
|
||||
return
|
||||
}
|
||||
s.closed = true
|
||||
s.cancel()
|
||||
s.mu.Unlock()
|
||||
s.wg.Wait()
|
||||
s.wg.Wait()
|
||||
})
|
||||
}
|
||||
|
||||
// run is one job's timing loop. It requests a fresh one-shot ticker per interval
|
||||
|
||||
@@ -164,6 +164,25 @@ func TestSchedulerCloseStopsJobs(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerCreatedAfterSwarmCloseIsClosed(t *testing.T) {
|
||||
sw := newSwarmFor(t, newLauncher(okFor))
|
||||
sw.Close()
|
||||
|
||||
sched := sw.Scheduler()
|
||||
sched.mu.Lock()
|
||||
closed := sched.closed
|
||||
sched.mu.Unlock()
|
||||
if !closed {
|
||||
t.Fatal("scheduler created after Swarm.Close was not closed")
|
||||
}
|
||||
if _, err := sched.Add(Policy{}, "team", "teammate", "t", "", Schedule{Every: time.Hour}); err == nil {
|
||||
t.Fatal("Add on a scheduler created after Swarm.Close should fail")
|
||||
}
|
||||
if got := len(sched.List()); got != 0 {
|
||||
t.Fatalf("scheduler created after Swarm.Close has %d jobs, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerAddRejectedAfterContextCancel(t *testing.T) {
|
||||
sw := newSwarmFor(t, newLauncher(okFor))
|
||||
sched := sw.Scheduler()
|
||||
|
||||
+103
-13
@@ -64,6 +64,18 @@ type Swarm struct {
|
||||
baseCtx context.Context
|
||||
cancel context.CancelFunc
|
||||
|
||||
// lifecycleMu guards closed and orders admission against shutdown. It is held
|
||||
// only for the flag check plus the lifecycleWork increment — never across a
|
||||
// launcher call, which would deadlock Close (see beginLifecycleAdmission).
|
||||
lifecycleMu sync.RWMutex
|
||||
closed bool
|
||||
// lifecycleWork counts admitted-but-unfinished lifecycle operations (spawn,
|
||||
// handoff, adoption, retry, queue drain). Close waits it out so no admitted
|
||||
// operation is still touching the swarm when Close returns.
|
||||
lifecycleWork sync.WaitGroup
|
||||
watchers sync.WaitGroup
|
||||
closeOnce sync.Once
|
||||
|
||||
mu sync.Mutex
|
||||
teams map[string]*Team
|
||||
taskCwd map[string]string // taskID -> cwd, for handoff/adoption relaunch
|
||||
@@ -137,31 +149,93 @@ func New(opts Options) (*Swarm, error) {
|
||||
}
|
||||
|
||||
// Close cancels every running member's context and releases resources. It is
|
||||
// safe to call more than once. The scheduler is closed first so no new spawn
|
||||
// fires after shutdown begins.
|
||||
// safe to call more than once. It closes lifecycle admission before canceling
|
||||
// members, fails queued tasks, and waits for every member watcher to exit.
|
||||
func (s *Swarm) Close() {
|
||||
s.mu.Lock()
|
||||
sched := s.scheduler
|
||||
s.mu.Unlock()
|
||||
if sched != nil {
|
||||
sched.Close()
|
||||
}
|
||||
if s.cancel != nil {
|
||||
s.cancel()
|
||||
}
|
||||
s.closeOnce.Do(func() {
|
||||
// Closing admission and cancelling are separate steps on purpose. The flag
|
||||
// flips under lifecycleMu so no further work is admitted, and the lock is
|
||||
// released BEFORE cancelling: a member already being launched holds no lock
|
||||
// here, but it may be blocked inside MemberLauncher.Launch waiting for the
|
||||
// context this cancel provides, and it can only release its admission ticket
|
||||
// once that cancellation lands.
|
||||
s.lifecycleMu.Lock()
|
||||
s.closed = true
|
||||
s.mu.Lock()
|
||||
sched := s.scheduler
|
||||
s.mu.Unlock()
|
||||
s.lifecycleMu.Unlock()
|
||||
|
||||
if s.cancel != nil {
|
||||
s.cancel()
|
||||
}
|
||||
if sched != nil {
|
||||
sched.Close()
|
||||
}
|
||||
// Wait for every already-admitted lifecycle operation (spawn, handoff,
|
||||
// adoption, retry, queue drain) to finish BEFORE sweeping queues below.
|
||||
// Such work can create teams as well as append to their queues, so snapshot
|
||||
// the teams only after every ticket is back. At that point no team or queue
|
||||
// can gain a new entry, and the sweep cannot miss late-admitted work.
|
||||
s.lifecycleWork.Wait()
|
||||
s.mu.Lock()
|
||||
teams := make([]*Team, 0, len(s.teams))
|
||||
for _, team := range s.teams {
|
||||
teams = append(teams, team)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
for _, team := range teams {
|
||||
for _, spec := range team.clearQueue() {
|
||||
_ = s.coord.Fail(spec.TaskID, ErrSwarmClosed.Error())
|
||||
}
|
||||
}
|
||||
// Every watcher is added while its admitting lifecycle ticket is still held,
|
||||
// so lifecycleWork.Wait above also orders every Add before this Wait. The
|
||||
// watcher itself can outlive that ticket, which is why both waits are needed.
|
||||
s.watchers.Wait()
|
||||
})
|
||||
}
|
||||
|
||||
// Scheduler returns the swarm's recurring-spawn scheduler, creating it on first
|
||||
// use. Scheduling is opt-in: until a job is added the scheduler does nothing.
|
||||
//
|
||||
// The post-close handoff between this method and Close relies on the ORDER in
|
||||
// which each reads/writes s.scheduler and s.closed, not on holding one lock across
|
||||
// both: Close takes lifecycleMu.Lock() (excluding this RLock), sets closed, reads
|
||||
// whatever s.scheduler currently is, and only THEN releases lifecycleMu and calls
|
||||
// sched.Close() on what it captured. So a call here either:
|
||||
// - runs entirely before Close's write lock — sees closed == false, and the
|
||||
// scheduler it creates/returns is exactly the one Close's later snapshot will
|
||||
// capture and close; or
|
||||
// - runs entirely after Close's write lock — sees closed == true (memory model:
|
||||
// the RLock/Lock pair orders the read after Close's write) and closes the
|
||||
// scheduler itself here, because Close already captured its own snapshot
|
||||
// (possibly nil, if this call is what first creates one) and will not look
|
||||
// again.
|
||||
//
|
||||
// Either way exactly one of the two call sites closes the scheduler that ends up
|
||||
// installed, and neither can observe a "created after close but never closed"
|
||||
// scheduler — that would require reading closed == false here while also reading
|
||||
// a scheduler Close had already passed over, which the lock ordering rules out.
|
||||
func (s *Swarm) Scheduler() *Scheduler {
|
||||
s.lifecycleMu.RLock()
|
||||
closed := s.closed
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.scheduler == nil {
|
||||
s.scheduler = newScheduler(s)
|
||||
}
|
||||
return s.scheduler
|
||||
sched := s.scheduler
|
||||
s.mu.Unlock()
|
||||
s.lifecycleMu.RUnlock()
|
||||
if closed {
|
||||
sched.Close()
|
||||
}
|
||||
return sched
|
||||
}
|
||||
|
||||
// ErrSwarmClosed reports lifecycle work submitted after shutdown begins.
|
||||
var ErrSwarmClosed = errors.New("swarm: closed")
|
||||
|
||||
// rememberCwd records a task's working dir so a handoff/adoption relaunch keeps it.
|
||||
func (s *Swarm) rememberCwd(taskID, cwd string) {
|
||||
s.mu.Lock()
|
||||
@@ -288,6 +362,22 @@ func (t *Team) onExit() (MemberSpec, bool) {
|
||||
return MemberSpec{}, false
|
||||
}
|
||||
|
||||
func (t *Team) releaseSlot() {
|
||||
t.mu.Lock()
|
||||
if t.running > 0 {
|
||||
t.running--
|
||||
}
|
||||
t.mu.Unlock()
|
||||
}
|
||||
|
||||
func (t *Team) clearQueue() []MemberSpec {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
queued := t.queue
|
||||
t.queue = nil
|
||||
return queued
|
||||
}
|
||||
|
||||
func (t *Team) addMember(m *Member) {
|
||||
t.mu.Lock()
|
||||
t.members[m.ID] = m
|
||||
|
||||
@@ -80,8 +80,19 @@ func (tool applyPatchTool) RunWithOptions(ctx context.Context, args map[string]a
|
||||
return errorResult("Error applying patch: " + err.Error())
|
||||
}
|
||||
var createdTargets []string
|
||||
var fullySuppliedTargets []string
|
||||
wholeBefore := map[string]bool{}
|
||||
if options.FileTracker != nil {
|
||||
createdTargets = missingPatchTargets(applyRoot, patch)
|
||||
fullySuppliedTargets = completeCreatedPatchTargets(applyRoot, patch)
|
||||
for _, path := range patchHeaderPaths(patch) {
|
||||
if path == "" || path == "/dev/null" {
|
||||
continue
|
||||
}
|
||||
if absolute, _, rerr := resolveWorkspaceTargetPath(applyRoot, path); rerr == nil {
|
||||
wholeBefore[absolute] = options.FileTracker.SeenWhole(absolute)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
command := exec.CommandContext(ctx, "git", "apply", "--whitespace=nowarn", patchPath)
|
||||
@@ -102,13 +113,28 @@ func (tool applyPatchTool) RunWithOptions(ctx context.Context, args map[string]a
|
||||
result := okResult(summary)
|
||||
result.ChangedFiles = changedFilesFromPatch(relativeRoot, patch)
|
||||
result.Display = Display{Summary: summary, Kind: "diff", Preview: capPreviewDiff(patch)}
|
||||
// git apply already rejects a patch whose context drifted, so it has its own
|
||||
// staleness guard. Drop any tracked baseline for the files it rewrote so a
|
||||
// subsequent edit_file/write_file re-reads instead of false-flagging the
|
||||
// patch's own change as an external modification.
|
||||
fullySupplied := make(map[string]bool, len(fullySuppliedTargets))
|
||||
for _, absolute := range fullySuppliedTargets {
|
||||
fullySupplied[absolute] = true
|
||||
}
|
||||
// Re-baseline files changed by this tool. When the model had already seen the
|
||||
// whole input (or supplied a complete new file), the exact patch plus that
|
||||
// input determines the whole output, so a follow-up edit needs no wasted read.
|
||||
// Partial observations remain conservative and are cleared by Record.
|
||||
for _, changed := range result.ChangedFiles {
|
||||
if absolute, _, rerr := resolveScopedPath(tool.workspaceRoot, tool.scope, changed); rerr == nil {
|
||||
options.FileTracker.Forget(absolute)
|
||||
content, readErr := os.ReadFile(absolute)
|
||||
if readErr != nil {
|
||||
options.FileTracker.Forget(absolute)
|
||||
continue
|
||||
}
|
||||
info, _ := os.Stat(absolute)
|
||||
wasWhole := wholeBefore[absolute] || fullySupplied[absolute]
|
||||
options.FileTracker.Record(absolute, content, info)
|
||||
if wasWhole {
|
||||
lines := lineCount(string(content))
|
||||
options.FileTracker.RecordSeenRange(absolute, 1, lines, lines)
|
||||
}
|
||||
}
|
||||
}
|
||||
recordCreatedPatchTargets(options.FileTracker, createdTargets)
|
||||
@@ -134,6 +160,60 @@ func missingPatchTargets(root string, patch string) []string {
|
||||
return missing
|
||||
}
|
||||
|
||||
// completeCreatedPatchTargets returns only files whose full contents are
|
||||
// supplied by a /dev/null creation patch. A missing rename/copy destination is
|
||||
// created by git too, but its bytes come from an unread source and must not gain
|
||||
// whole-file observation credit.
|
||||
func completeCreatedPatchTargets(root string, patch string) []string {
|
||||
seen := map[string]bool{}
|
||||
var created []string
|
||||
oldRemaining, newRemaining := 0, 0
|
||||
inHunk := false
|
||||
fromDevNull := false
|
||||
for _, line := range strings.Split(strings.ReplaceAll(patch, "\r\n", "\n"), "\n") {
|
||||
if inHunk && (oldRemaining > 0 || newRemaining > 0) {
|
||||
switch {
|
||||
case strings.HasPrefix(line, "-"):
|
||||
oldRemaining--
|
||||
case strings.HasPrefix(line, "+"):
|
||||
newRemaining--
|
||||
case strings.HasPrefix(line, "\\"):
|
||||
default:
|
||||
oldRemaining--
|
||||
newRemaining--
|
||||
}
|
||||
continue
|
||||
}
|
||||
inHunk = false
|
||||
switch {
|
||||
case strings.HasPrefix(line, "diff --git "):
|
||||
fromDevNull = false
|
||||
case strings.HasPrefix(line, "@@"):
|
||||
oldRemaining, newRemaining = parseHunkCounts(line)
|
||||
inHunk = oldRemaining > 0 || newRemaining > 0
|
||||
case strings.HasPrefix(line, "--- "):
|
||||
fromDevNull = patchFileHeaderPath(line) == "/dev/null"
|
||||
case strings.HasPrefix(line, "+++ "):
|
||||
path := patchFileHeaderPath(line)
|
||||
if !fromDevNull || path == "" || path == "/dev/null" {
|
||||
fromDevNull = false
|
||||
continue
|
||||
}
|
||||
fromDevNull = false
|
||||
absolute, _, err := resolveWorkspaceTargetPath(root, stripPatchPrefix(path))
|
||||
if err != nil || seen[absolute] {
|
||||
continue
|
||||
}
|
||||
if _, err := os.Stat(absolute); !os.IsNotExist(err) {
|
||||
continue
|
||||
}
|
||||
seen[absolute] = true
|
||||
created = append(created, absolute)
|
||||
}
|
||||
}
|
||||
return created
|
||||
}
|
||||
|
||||
func recordCreatedPatchTargets(tracker *FileTracker, missingBefore []string) {
|
||||
if tracker == nil {
|
||||
return
|
||||
@@ -240,15 +320,7 @@ func patchHeaderPaths(patch string) []string {
|
||||
oldRemaining, newRemaining = parseHunkCounts(line)
|
||||
inHunk = oldRemaining > 0 || newRemaining > 0
|
||||
case strings.HasPrefix(line, "--- "), strings.HasPrefix(line, "+++ "):
|
||||
// Take everything after the fixed 4-char prefix instead of splitting on
|
||||
// spaces, so a path that contains spaces survives. Drop a trailing
|
||||
// tab-delimited timestamp (unified-diff convention) and unquote a
|
||||
// C-quoted git path (git quotes names with spaces/specials) (L18).
|
||||
rest := line[len("--- "):] // "--- " and "+++ " are both 4 bytes
|
||||
if tab := strings.IndexByte(rest, '\t'); tab >= 0 {
|
||||
rest = rest[:tab]
|
||||
}
|
||||
if p := strings.TrimSpace(unquoteGitPath(rest)); p != "" && p != "/dev/null" {
|
||||
if p := patchFileHeaderPath(line); p != "" && p != "/dev/null" {
|
||||
paths = append(paths, stripPatchPrefix(p))
|
||||
}
|
||||
}
|
||||
@@ -256,6 +328,17 @@ func patchHeaderPaths(patch string) []string {
|
||||
return paths
|
||||
}
|
||||
|
||||
func patchFileHeaderPath(line string) string {
|
||||
if len(line) < len("--- ") {
|
||||
return ""
|
||||
}
|
||||
rest := line[len("--- "):] // "--- " and "+++ " are both 4 bytes
|
||||
if tab := strings.IndexByte(rest, '\t'); tab >= 0 {
|
||||
rest = rest[:tab]
|
||||
}
|
||||
return strings.TrimSpace(unquoteGitPath(rest))
|
||||
}
|
||||
|
||||
// parseHunkCounts reads the old/new line counts from a "@@ -a,b +c,d @@" header.
|
||||
// A missing count (e.g. "@@ -a +c @@") means 1 per unified-diff convention.
|
||||
//
|
||||
|
||||
+10
-27
@@ -43,6 +43,7 @@ func NewScopedBashTool(workspaceRoot string, scope PathScope) Tool {
|
||||
baseTool: baseTool{
|
||||
name: "bash",
|
||||
description: "Execute a shell command inside the workspace (or an explicitly granted extra directory) after permission is granted. " + shellGuidance,
|
||||
deferred: true,
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]PropertySchema{
|
||||
@@ -106,7 +107,7 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS
|
||||
// unsandboxed) can actually bypass the MSYS guard instead of being
|
||||
// hard-blocked by the same check it was meant to escalate past.
|
||||
commandEngine := commandEngineForSandboxPermissions(engine, sandboxPermissions)
|
||||
if issue := detectShellCommandIssue(commandText, runtime.GOOS); issue != nil && !msysGuardBypassed(issue, commandEngine) {
|
||||
if issue := detectShellCommandIssueForRuntime(commandText, detectShellRuntime(runtime.GOOS)); issue != nil && !msysGuardBypassed(issue, commandEngine) {
|
||||
return shellIssueBlockResult(*issue)
|
||||
}
|
||||
|
||||
@@ -293,23 +294,19 @@ func shellIssueBlockResult(issue shellIssue) Result {
|
||||
}
|
||||
|
||||
// buildBashCommand returns the exec.Cmd and the sandbox plan for running
|
||||
// commandText. On Windows, when the command is not wrapped by the sandbox
|
||||
// engine (plan.Wrapped == false), it also overrides the child's raw command
|
||||
// line so commandText reaches cmd.exe unescaped; see
|
||||
// zeroSandbox.WindowsShellCommandLine for why that matters. The wrapped case
|
||||
// gets the same treatment inside the sandboxed runner process itself
|
||||
// (internal/sandbox/windows_process_windows.go), since that command line is
|
||||
// built there, not here.
|
||||
// commandText. PowerShell is preferred on Windows, with cmd.exe retained as
|
||||
// the fallback. Only cmd.exe needs the raw command-line override.
|
||||
func buildBashCommand(ctx context.Context, commandText string, absoluteCwd string, engine *zeroSandbox.Engine) (*exec.Cmd, zeroSandbox.CommandPlan, error) {
|
||||
hostShell := detectShellRuntime(runtime.GOOS)
|
||||
spec := zeroSandbox.CommandSpec{
|
||||
Name: shellExecutable(),
|
||||
Args: shellArguments(commandText),
|
||||
Name: hostShell.Executable,
|
||||
Args: hostShell.arguments(commandText),
|
||||
Dir: absoluteCwd,
|
||||
}
|
||||
if engine != nil {
|
||||
command, plan, err := engine.CommandContext(ctx, spec)
|
||||
if err == nil {
|
||||
applyWindowsShellCommandLine(command, commandText, plan.Wrapped)
|
||||
applyWindowsShellCommandLine(command, commandText, plan.Wrapped, hostShell.Kind == shellKindCmd)
|
||||
}
|
||||
return command, plan, err
|
||||
}
|
||||
@@ -333,7 +330,7 @@ func buildBashCommand(ctx context.Context, commandText string, absoluteCwd strin
|
||||
}
|
||||
command := exec.CommandContext(ctx, spec.Name, spec.Args...)
|
||||
command.Dir = spec.Dir
|
||||
applyWindowsShellCommandLine(command, commandText, plan.Wrapped)
|
||||
applyWindowsShellCommandLine(command, commandText, plan.Wrapped, hostShell.Kind == shellKindCmd)
|
||||
return command, plan, nil
|
||||
}
|
||||
|
||||
@@ -385,20 +382,6 @@ func interactiveBlockResult(detection zeroSandbox.InteractiveCommandResult) Resu
|
||||
}
|
||||
}
|
||||
|
||||
func shellExecutable() string {
|
||||
if runtime.GOOS == "windows" {
|
||||
return "cmd.exe"
|
||||
}
|
||||
return "/bin/sh"
|
||||
}
|
||||
|
||||
func shellArguments(command string) []string {
|
||||
if runtime.GOOS == "windows" {
|
||||
return zeroSandbox.WindowsShellArgs(command)
|
||||
}
|
||||
return []string{"-c", command}
|
||||
}
|
||||
|
||||
func commandExitCode(err error) int {
|
||||
if err == nil {
|
||||
return 0
|
||||
@@ -621,7 +604,7 @@ func truncateHeadTailWithTotal(value string, total, maxBytes int) (string, int,
|
||||
|
||||
func formatBashOutputWithShellHint(stdout string, stderr string, exitCode int, meta map[string]string) string {
|
||||
output := formatBashOutput(stdout, stderr, exitCode)
|
||||
if issue := detectShellOutputIssue(stdout+"\n"+stderr, runtime.GOOS); issue != nil {
|
||||
if issue := detectShellOutputIssueForRuntime(stdout+"\n"+stderr, detectShellRuntime(runtime.GOOS)); issue != nil {
|
||||
meta["shell_issue"] = issue.Kind
|
||||
output = appendShellIssueHint(output, *issue)
|
||||
}
|
||||
|
||||
@@ -36,6 +36,6 @@ func hardenProcessLifetime(command *exec.Cmd) {
|
||||
}
|
||||
}
|
||||
|
||||
// applyWindowsShellCommandLine is a no-op outside Windows; there is no
|
||||
// cmd.exe command-line-parsing quirk to work around here.
|
||||
func applyWindowsShellCommandLine(command *exec.Cmd, commandText string, wrapped bool) {}
|
||||
// applyWindowsShellCommandLine is a no-op outside Windows.
|
||||
func applyWindowsShellCommandLine(command *exec.Cmd, commandText string, wrapped bool, cmdFallback bool) {
|
||||
}
|
||||
@@ -31,15 +31,12 @@ func hardenProcessLifetime(command *exec.Cmd) {
|
||||
}
|
||||
}
|
||||
|
||||
// applyWindowsShellCommandLine overrides command's raw child command line so
|
||||
// commandText reaches cmd.exe unescaped instead of auto-quoted the way
|
||||
// exec.Cmd would normally encode a single Args element. Skipped when wrapped
|
||||
// is true: the sandbox engine then routes execution through a separate
|
||||
// zero-windows-command-runner process, which builds its own child command
|
||||
// line from scratch (internal/sandbox/windows_process_windows.go) rather than
|
||||
// inheriting whatever this outer exec.Cmd is configured with.
|
||||
func applyWindowsShellCommandLine(command *exec.Cmd, commandText string, wrapped bool) {
|
||||
if wrapped {
|
||||
// applyWindowsShellCommandLine overrides the cmd.exe fallback's raw child
|
||||
// command line so commandText reaches cmd.exe unescaped. PowerShell consumes
|
||||
// ordinary argv quoting and must not take this path. Wrapped commands are
|
||||
// handled after unwrapping by the Windows sandbox runner.
|
||||
func applyWindowsShellCommandLine(command *exec.Cmd, commandText string, wrapped bool, cmdFallback bool) {
|
||||
if wrapped || !cmdFallback {
|
||||
return
|
||||
}
|
||||
command.SysProcAttr = &syscall.SysProcAttr{CmdLine: zeroSandbox.WindowsShellCommandLine(commandText)}
|
||||
|
||||
@@ -124,11 +124,15 @@ func TestBashToolDescribesHostShellSyntax(t *testing.T) {
|
||||
}
|
||||
|
||||
if runtime.GOOS == "windows" {
|
||||
if !strings.Contains(description, "cmd.exe") || !strings.Contains(description, "cwd") {
|
||||
t.Fatalf("expected Windows cmd.exe and cwd guidance in bash description, got %q", description)
|
||||
}
|
||||
if !strings.Contains(description, "double quotes") || !strings.Contains(description, `--jq ".a | b"`) {
|
||||
t.Fatalf("expected the double-quote metacharacter rule in bash description, got %q", description)
|
||||
shell := detectShellRuntime(runtime.GOOS)
|
||||
if shell.Kind == shellKindPowerShell {
|
||||
for _, want := range []string{"powershell", "cwd", "get-childitem", "select-string", "$env:name"} {
|
||||
if !strings.Contains(description, want) {
|
||||
t.Fatalf("expected Windows PowerShell guidance %q in bash description, got %q", want, description)
|
||||
}
|
||||
}
|
||||
} else if !strings.Contains(description, "cmd.exe") || !strings.Contains(description, "cwd") {
|
||||
t.Fatalf("expected Windows cmd.exe fallback guidance in bash description, got %q", description)
|
||||
}
|
||||
return
|
||||
}
|
||||
@@ -455,13 +459,14 @@ func TestBashToolReturnsNonzeroExitAsError(t *testing.T) {
|
||||
if result.Status != StatusError {
|
||||
t.Fatalf("expected error status, got %s", result.Status)
|
||||
}
|
||||
for _, want := range []string{"stdout:\nbefore failure", "stderr:\nfailure details", "exit_code: 7"} {
|
||||
wantExitCode := strconv.Itoa(helperFailureExitCode())
|
||||
for _, want := range []string{"stdout:\nbefore failure", "stderr:\nfailure details", "exit_code: " + wantExitCode} {
|
||||
if !strings.Contains(result.Output, want) {
|
||||
t.Fatalf("expected output to contain %q, got %q", want, result.Output)
|
||||
}
|
||||
}
|
||||
if result.Meta["exit_code"] != "7" {
|
||||
t.Fatalf("expected exit_code metadata 7, got %q", result.Meta["exit_code"])
|
||||
if result.Meta["exit_code"] != wantExitCode {
|
||||
t.Fatalf("expected exit_code metadata %s, got %q", wantExitCode, result.Meta["exit_code"])
|
||||
}
|
||||
if result.ExecutionOutcome == nil || result.ExecutionOutcome.State != execution.StateFailed || result.ExecutionOutcome.Kind != execution.OutcomeApplicationFailure {
|
||||
t.Fatalf("execution outcome = %#v, want failed/application_failure", result.ExecutionOutcome)
|
||||
@@ -578,10 +583,16 @@ func TestBashToolRequireEscalatedMsysGuard(t *testing.T) {
|
||||
}
|
||||
registry := NewRegistry()
|
||||
registry.Register(NewScopedBashTool(root, nil))
|
||||
msysCommand := "cat somefile.txt"
|
||||
if detectShellRuntime(runtime.GOOS).Kind == shellKindPowerShell {
|
||||
// cat is a native Get-Content alias in PowerShell. Use a name that
|
||||
// still resolves to an incompatible Git-for-Windows MSYS executable.
|
||||
msysCommand = "grep pattern somefile.txt"
|
||||
}
|
||||
|
||||
t.Run("default sandboxing still blocks an MSYS-prone command", func(t *testing.T) {
|
||||
result := registry.RunWithOptions(context.Background(), "bash", map[string]any{
|
||||
"command": "cat somefile.txt",
|
||||
"command": msysCommand,
|
||||
}, RunOptions{
|
||||
PermissionGranted: true,
|
||||
Sandbox: newEngine(),
|
||||
@@ -594,7 +605,7 @@ func TestBashToolRequireEscalatedMsysGuard(t *testing.T) {
|
||||
|
||||
t.Run("approved require_escalated bypasses the MSYS guard", func(t *testing.T) {
|
||||
result := registry.RunWithOptions(context.Background(), "bash", map[string]any{
|
||||
"command": "cat somefile.txt",
|
||||
"command": msysCommand,
|
||||
"sandbox_permissions": string(SandboxPermissionsRequireEscalated),
|
||||
}, RunOptions{
|
||||
PermissionGranted: true,
|
||||
@@ -614,7 +625,14 @@ func TestBashToolRequireEscalatedMsysGuard(t *testing.T) {
|
||||
|
||||
t.Run("approved require_escalated still blocks an unrelated syntax issue", func(t *testing.T) {
|
||||
result := registry.RunWithOptions(context.Background(), "bash", map[string]any{
|
||||
"command": `cd /d/tmp/zero-pr-158 && dir`,
|
||||
// No `&&` here on purpose. The POSIX-style path is what this case is
|
||||
// about, and on Windows PowerShell 5.1 a chained command trips the
|
||||
// separate windows_powershell_version preflight first, which shadows
|
||||
// the block being asserted. That only shows up on 5.1: the runners
|
||||
// have PowerShell 7, where && is valid and the version check does not
|
||||
// fire, so the substitution kept this green in CI while failing on the
|
||||
// shell a stock Windows box actually gets.
|
||||
"command": `cd /d/tmp/zero-pr-158`,
|
||||
"sandbox_permissions": string(SandboxPermissionsRequireEscalated),
|
||||
}, RunOptions{
|
||||
PermissionGranted: true,
|
||||
@@ -648,8 +666,9 @@ func TestBashToolIgnoresMsysMarkersInCommandArgumentsAfterFailure(t *testing.T)
|
||||
"command": command,
|
||||
})
|
||||
|
||||
if result.Status != StatusError || result.Meta["exit_code"] != "7" {
|
||||
t.Fatalf("expected the helper's real exit 7 failure, got %s: %#v", result.Status, result)
|
||||
wantExitCode := strconv.Itoa(helperFailureExitCode())
|
||||
if result.Status != StatusError || result.Meta["exit_code"] != wantExitCode {
|
||||
t.Fatalf("expected the helper failure exit %s, got %s: %#v", wantExitCode, result.Status, result)
|
||||
}
|
||||
if result.Meta["shell_issue"] == "windows_msys_sandbox" {
|
||||
t.Fatalf("expected the MSYS marker in the command's own argument text to be ignored, got %#v", result)
|
||||
@@ -860,6 +879,9 @@ func TestBashToolRunsCommandLineForLoopSyntax(t *testing.T) {
|
||||
if runtime.GOOS != "windows" {
|
||||
t.Skip("cmd.exe command-line vs batch-file parsing is Windows-specific")
|
||||
}
|
||||
if detectShellRuntime(runtime.GOOS).Kind != shellKindCmd {
|
||||
t.Skip("cmd.exe fallback is not selected")
|
||||
}
|
||||
root := t.TempDir()
|
||||
|
||||
result := NewScopedBashTool(root, nil).Run(context.Background(), map[string]any{
|
||||
@@ -881,6 +903,10 @@ func helperCommand(name string) string {
|
||||
return executable + " --zero-bash-helper " + name
|
||||
}
|
||||
|
||||
func helperFailureExitCode() int {
|
||||
return 7
|
||||
}
|
||||
|
||||
func shellQuote(value string) string {
|
||||
if runtime.GOOS == "windows" {
|
||||
return value
|
||||
|
||||
@@ -74,6 +74,11 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
|
||||
if cerr := options.FileTracker.CheckConflict(absolutePath, contentBytes); cerr != nil {
|
||||
return errorResult(fileConflictMessage(relativePath))
|
||||
}
|
||||
if options.FileTracker != nil {
|
||||
if _, tracked := options.FileTracker.Version(absolutePath); !tracked {
|
||||
return errorResult(fileUnseenMessage(relativePath))
|
||||
}
|
||||
}
|
||||
content := string(contentBytes)
|
||||
occurrences := strings.Count(content, oldString)
|
||||
|
||||
@@ -131,7 +136,11 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
|
||||
if !replaceAll && occurrences > 1 {
|
||||
return errorResult(fmt.Sprintf("Error: old_string matches %d locations in %s. Either make old_string more specific, or pass replace_all: true to replace every occurrence.", occurrences, relativePath))
|
||||
}
|
||||
if options.FileTracker != nil && !allOccurrencesSeen(options.FileTracker, absolutePath, content, oldString, replaceAll) {
|
||||
return errorResult(fileUnseenMessage(relativePath))
|
||||
}
|
||||
|
||||
previouslySeenWhole := options.FileTracker.SeenWhole(absolutePath)
|
||||
updated := strings.Replace(content, oldString, newString, 1)
|
||||
replacedCount := 1
|
||||
if replaceAll {
|
||||
@@ -141,12 +150,14 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
|
||||
if updated == content {
|
||||
return okResult("No changes: new_string is identical to old_string.")
|
||||
}
|
||||
editedSpans := replacementByteSpans(content, oldString, newString, replaceAll)
|
||||
if err := recheckScopedWriteTarget(tool.workspaceRoot, tool.scope, requestedPath); err != nil {
|
||||
return errorResult("Error writing " + relativePath + ": " + err.Error())
|
||||
}
|
||||
if err := os.WriteFile(absolutePath, []byte(updated), 0o644); err != nil {
|
||||
return errorResult("Error writing " + relativePath + ": " + err.Error())
|
||||
}
|
||||
modelKnownContent := updated
|
||||
// Optional format-on-write (ZERO_FORMAT_ON_WRITE). Must run BEFORE the
|
||||
// FileTracker re-baseline: recording pre-format content would make the very
|
||||
// next edit look like an external modification and trip the conflict guard.
|
||||
@@ -155,6 +166,15 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
|
||||
// compare against the current on-disk state, not the pre-edit version.
|
||||
newInfo, _ := os.Stat(absolutePath)
|
||||
options.FileTracker.Record(absolutePath, []byte(updated), newInfo)
|
||||
if updated == modelKnownContent {
|
||||
if previouslySeenWhole {
|
||||
options.FileTracker.RecordSeenRange(absolutePath, 1, lineCount(updated), lineCount(updated))
|
||||
} else {
|
||||
for _, span := range editedSpans {
|
||||
options.FileTracker.RecordSeenBytes(absolutePath, span.start, span.end, len(updated))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
suffix := ""
|
||||
if replacedCount != 1 {
|
||||
@@ -169,3 +189,54 @@ func (tool editFileTool) RunWithOptions(ctx context.Context, args map[string]any
|
||||
result.Display = Display{Summary: fmt.Sprintf("Edited %s", relativePath), Kind: "diff", Preview: boundedUnifiedDiff(relativePath, content, updated)}
|
||||
return result
|
||||
}
|
||||
|
||||
// replacementByteSpans returns every half-open post-edit byte range whose exact
|
||||
// contents came from newString. Offsets are translated from the original with a
|
||||
// cumulative delta so replace_all remains correct after earlier replacements.
|
||||
func replacementByteSpans(content, oldString, newString string, replaceAll bool) []lineRange {
|
||||
spans := make([]lineRange, 0, 1)
|
||||
offset := 0
|
||||
delta := 0
|
||||
for {
|
||||
index := strings.Index(content[offset:], oldString)
|
||||
if index < 0 {
|
||||
return spans
|
||||
}
|
||||
index += offset
|
||||
start := index + delta
|
||||
spans = append(spans, lineRange{start: start, end: start + len(newString)})
|
||||
if !replaceAll {
|
||||
return spans
|
||||
}
|
||||
offset = index + len(oldString)
|
||||
delta += len(newString) - len(oldString)
|
||||
}
|
||||
}
|
||||
|
||||
func allOccurrencesSeen(tracker *FileTracker, path, content, oldString string, replaceAll bool) bool {
|
||||
offset := 0
|
||||
for {
|
||||
index := strings.Index(content[offset:], oldString)
|
||||
if index < 0 {
|
||||
return true
|
||||
}
|
||||
index += offset
|
||||
start, end := lineSpanForOffset(content, index, len(oldString))
|
||||
if !tracker.SeenBytes(path, index, index+len(oldString)) && !tracker.SeenRange(path, start, end) {
|
||||
return false
|
||||
}
|
||||
if !replaceAll {
|
||||
return true
|
||||
}
|
||||
offset = index + len(oldString)
|
||||
}
|
||||
}
|
||||
|
||||
func lineSpanForOffset(content string, offset, length int) (int, int) {
|
||||
if offset < 0 {
|
||||
return 1, 1
|
||||
}
|
||||
start := strings.Count(content[:offset], "\n") + 1
|
||||
end := start + strings.Count(content[offset:min(len(content), offset+length)], "\n")
|
||||
return start, end
|
||||
}
|
||||
@@ -173,7 +173,7 @@ func (tool execCommandTool) run(ctx context.Context, args map[string]any, engine
|
||||
// unsandboxed) can actually bypass the MSYS guard instead of being
|
||||
// hard-blocked by the same check it was meant to escalate past.
|
||||
commandEngine := commandEngineForSandboxPermissions(engine, sandboxPermissions)
|
||||
if issue := detectShellCommandIssue(commandText, runtimeGOOS()); issue != nil && !msysGuardBypassed(issue, commandEngine) {
|
||||
if issue := detectShellCommandIssueForRuntime(commandText, detectShellRuntime(runtimeGOOS())); issue != nil && !msysGuardBypassed(issue, commandEngine) {
|
||||
return shellIssueBlockResult(*issue)
|
||||
}
|
||||
if interactive := zeroSandbox.DetectInteractiveCommand(commandText, runtimeGOOS()); interactive.Interactive {
|
||||
@@ -459,7 +459,7 @@ func execToolResultWithBudget(input execToolResultInput, directBudget bool) Resu
|
||||
}
|
||||
body := formatExecCommandOutput(output, input.sessionID, input.exited, input.exitCode, input.interrupted)
|
||||
if status == StatusError && input.exited && !input.interrupted {
|
||||
if issue := detectShellOutputIssue(output, runtimeGOOS()); issue != nil {
|
||||
if issue := detectShellOutputIssueForRuntime(output, detectShellRuntime(runtimeGOOS())); issue != nil {
|
||||
meta["shell_issue"] = issue.Kind
|
||||
body = appendShellIssueHint(body, *issue)
|
||||
}
|
||||
|
||||
@@ -18,7 +18,7 @@ import (
|
||||
)
|
||||
|
||||
func TestIndependentExecCommandConstructorsShareDefaultManager(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
root := execTestRoot(t)
|
||||
execTool := NewScopedExecCommandTool(root, nil, nil)
|
||||
writeTool := NewWriteStdinTool(nil)
|
||||
|
||||
@@ -76,17 +76,21 @@ func TestExecCommandToolDescribesHostShellSyntax(t *testing.T) {
|
||||
description := strings.ToLower(strings.Join(descriptionParts, " "))
|
||||
|
||||
if runtime.GOOS == "windows" {
|
||||
if !strings.Contains(description, "cmd.exe") || !strings.Contains(description, "cwd") {
|
||||
t.Fatalf("expected Windows cmd.exe and cwd guidance in exec_command description, got %q", description)
|
||||
}
|
||||
if !strings.Contains(description, "double quotes") || !strings.Contains(description, `--jq ".a | b"`) {
|
||||
t.Fatalf("expected the double-quote metacharacter rule in exec_command description, got %q", description)
|
||||
shell := detectShellRuntime(runtime.GOOS)
|
||||
if shell.Kind == shellKindPowerShell {
|
||||
for _, want := range []string{"powershell", "cwd", "get-childitem", "select-string", "$env:name"} {
|
||||
if !strings.Contains(description, want) {
|
||||
t.Fatalf("expected Windows PowerShell guidance %q in exec_command description, got %q", want, description)
|
||||
}
|
||||
}
|
||||
} else if !strings.Contains(description, "cmd.exe") || !strings.Contains(description, "cwd") {
|
||||
t.Fatalf("expected Windows cmd.exe fallback guidance in exec_command description, got %q", description)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecCommandReturnsSessionAndWriteStdinPollsCompletion(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
root := execTestRoot(t)
|
||||
manager := newExecSessionManager()
|
||||
execTool := NewScopedExecCommandTool(root, nil, manager)
|
||||
writeTool := NewWriteStdinTool(manager)
|
||||
@@ -134,7 +138,7 @@ func TestExecCommandReturnsSessionAndWriteStdinPollsCompletion(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestExecCommandRequireEscalatedBypassesNativeSandboxAfterApproval(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
root := execTestRoot(t)
|
||||
manager := newExecSessionManager()
|
||||
registry := NewRegistry()
|
||||
registry.Register(NewScopedExecCommandTool(root, nil, manager))
|
||||
@@ -183,7 +187,7 @@ func TestExecCommandRequireEscalatedBypassesMsysGuardAfterApproval(t *testing.T)
|
||||
if runtime.GOOS != "windows" {
|
||||
t.Skip("windows-only MSYS sandbox guard")
|
||||
}
|
||||
root := t.TempDir()
|
||||
root := execTestRoot(t)
|
||||
manager := newExecSessionManager()
|
||||
registry := NewRegistry()
|
||||
registry.Register(NewScopedExecCommandTool(root, nil, manager))
|
||||
@@ -192,9 +196,15 @@ func TestExecCommandRequireEscalatedBypassesMsysGuardAfterApproval(t *testing.T)
|
||||
Policy: sandbox.DefaultPolicy(),
|
||||
Backend: sandbox.Backend{Name: sandbox.BackendUnavailable, Message: "native sandbox unavailable"},
|
||||
})
|
||||
msysCommand := "cat somefile.txt"
|
||||
if detectShellRuntime(runtime.GOOS).Kind == shellKindPowerShell {
|
||||
// cat is a native Get-Content alias in PowerShell, so use an executable
|
||||
// name that still exercises the MSYS guard.
|
||||
msysCommand = "grep pattern somefile.txt"
|
||||
}
|
||||
|
||||
result := registry.RunWithOptions(context.Background(), ExecCommandToolName, map[string]any{
|
||||
"cmd": "cat somefile.txt",
|
||||
"cmd": msysCommand,
|
||||
"sandbox_permissions": string(SandboxPermissionsRequireEscalated),
|
||||
}, RunOptions{
|
||||
PermissionGranted: true,
|
||||
@@ -204,8 +214,8 @@ func TestExecCommandRequireEscalatedBypassesMsysGuardAfterApproval(t *testing.T)
|
||||
|
||||
// Assert on the preflight block sentinel (exit_code "-1", set only by
|
||||
// shellIssueBlockResult) rather than shell_issue: once the guard is
|
||||
// bypassed, "cat somefile.txt" actually runs, and its real,
|
||||
// PATH-dependent output could otherwise trip the unrelated
|
||||
// bypassed, the command actually runs, and its real, PATH-dependent output
|
||||
// could otherwise trip the unrelated
|
||||
// post-execution detectShellOutputIssue heuristic and make this
|
||||
// assertion flaky for reasons unrelated to the guard under test.
|
||||
if result.Meta["exit_code"] == "-1" {
|
||||
@@ -214,7 +224,7 @@ func TestExecCommandRequireEscalatedBypassesMsysGuardAfterApproval(t *testing.T)
|
||||
}
|
||||
|
||||
func TestExecCommandReturnsExitCodeWhenCommandCompletesDuringInitialYield(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
root := execTestRoot(t)
|
||||
manager := newExecSessionManager()
|
||||
execTool := NewScopedExecCommandTool(root, nil, manager)
|
||||
|
||||
@@ -253,7 +263,7 @@ func TestExecCommandReportsWorkspaceChanges(t *testing.T) {
|
||||
if runtime.GOOS == "windows" {
|
||||
t.Skip("uses POSIX shell syntax")
|
||||
}
|
||||
root := t.TempDir()
|
||||
root := execTestRoot(t)
|
||||
result := NewScopedExecCommandTool(root, nil, newExecSessionManager()).Run(context.Background(), map[string]any{
|
||||
"cmd": "mkdir -p src node_modules/pkg && printf 'export {}' > src/main.ts && printf generated > node_modules/pkg/index.js",
|
||||
"yield_time_ms": 30000,
|
||||
@@ -284,8 +294,9 @@ func TestExecCommandApplicationFailureHasTypedOutcome(t *testing.T) {
|
||||
if result.ExecutionOutcome == nil || result.ExecutionOutcome.State != execution.StateFailed || result.ExecutionOutcome.Kind != execution.OutcomeApplicationFailure {
|
||||
t.Fatalf("execution outcome = %#v, want failed/application_failure", result.ExecutionOutcome)
|
||||
}
|
||||
if result.ExecutionOutcome.Exit == nil || result.ExecutionOutcome.Exit.Code != 7 {
|
||||
t.Fatalf("execution exit = %#v, want code 7", result.ExecutionOutcome.Exit)
|
||||
wantExitCode := helperFailureExitCode()
|
||||
if result.ExecutionOutcome.Exit == nil || result.ExecutionOutcome.Exit.Code != wantExitCode {
|
||||
t.Fatalf("execution exit = %#v, want code %d", result.ExecutionOutcome.Exit, wantExitCode)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -316,6 +327,26 @@ func TestExecCommandUsesStructuredAdapterDenial(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecCommandDoesNotInferNativeSandboxDenialFromOutput(t *testing.T) {
|
||||
result := execToolResult(execToolResultInput{
|
||||
commandText: "touch \"$HOME/probe\"",
|
||||
output: "touch: cannot touch '/home/user/probe': Read-only file system",
|
||||
exited: true,
|
||||
exitCode: 1,
|
||||
enforcement: execution.Enforcement{Backend: string(sandbox.BackendLinuxBwrap), Level: string(sandbox.EnforcementNative)},
|
||||
})
|
||||
|
||||
if result.ExecutionOutcome == nil || result.ExecutionOutcome.State != execution.StateFailed || result.ExecutionOutcome.Kind != execution.OutcomeApplicationFailure {
|
||||
t.Fatalf("execution outcome = %#v, want failed/application_failure", result.ExecutionOutcome)
|
||||
}
|
||||
if result.ExecutionOutcome.Denial != nil {
|
||||
t.Fatalf("command output created a typed denial: %#v", result.ExecutionOutcome.Denial)
|
||||
}
|
||||
if result.Meta[SandboxLikelyDeniedMeta] == "true" {
|
||||
t.Fatalf("command output created sandbox-denial metadata: %#v", result.Meta)
|
||||
}
|
||||
}
|
||||
|
||||
func executionRequestHasCapability(request execution.Request, kind execution.CapabilityKind, scope string) bool {
|
||||
for _, capability := range request.Capabilities {
|
||||
if capability.Kind == kind && (scope == "" || capability.Scope == scope) {
|
||||
@@ -326,7 +357,7 @@ func executionRequestHasCapability(request execution.Request, kind execution.Cap
|
||||
}
|
||||
|
||||
func TestExecCommandForegroundServerReturnsSessionAndServesHTTP(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
root := execTestRoot(t)
|
||||
manager := newExecSessionManager()
|
||||
execTool := NewScopedExecCommandTool(root, nil, manager)
|
||||
writeTool := NewWriteStdinTool(manager)
|
||||
@@ -380,7 +411,7 @@ func parseListeningAddress(output string) string {
|
||||
}
|
||||
|
||||
func TestExecCommandReapsFinishedUnpolledSession(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
root := execTestRoot(t)
|
||||
manager := execution.NewProcessManager(execution.ProcessManagerOptions{CompletedRetention: 10 * time.Millisecond})
|
||||
execTool := NewScopedExecCommandTool(root, nil, manager)
|
||||
|
||||
@@ -506,7 +537,7 @@ func TestWriteStdinInterruptTerminatesSession(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestWriteStdinRejectsInputForNonTTYSession(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
root := execTestRoot(t)
|
||||
manager := newExecSessionManager()
|
||||
execTool := NewScopedExecCommandTool(root, nil, manager)
|
||||
writeTool := NewWriteStdinTool(manager)
|
||||
@@ -537,9 +568,43 @@ func TestWriteStdinRejectsInputForNonTTYSession(t *testing.T) {
|
||||
manager.Stop(sessionID)
|
||||
}
|
||||
|
||||
// execTestRoot is t.TempDir with a cleanup that tolerates Windows still holding
|
||||
// the directory.
|
||||
//
|
||||
// A terminated process does not release its handles the instant the session
|
||||
// reports it exited, and Windows refuses to remove a directory with an open
|
||||
// handle where POSIX does not. t.TempDir's own cleanup removes once and fails
|
||||
// the test on a sharing violation, which turns ordinary teardown timing into a
|
||||
// test failure with nothing actually wrong underneath it.
|
||||
func execTestRoot(t *testing.T) string {
|
||||
t.Helper()
|
||||
|
||||
root, err := os.MkdirTemp("", "zero-exec-test-")
|
||||
if err != nil {
|
||||
t.Fatalf("create test root: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
deadline := time.Now().Add(10 * time.Second)
|
||||
for {
|
||||
if err := os.RemoveAll(root); err == nil {
|
||||
return
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
// Deliberately not a failure. The directory is under the OS temp
|
||||
// root and the process holding it has already been terminated, so
|
||||
// the worst case is a stale directory, not a broken test.
|
||||
t.Logf("test root %s still held after the cleanup deadline; leaving it", root)
|
||||
return
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
})
|
||||
return root
|
||||
}
|
||||
|
||||
func TestWriteStdinStopIntentTerminatesNonTTYSession(t *testing.T) {
|
||||
for _, chars := range []string{`\u0003`, "exit\n"} {
|
||||
root := t.TempDir()
|
||||
root := execTestRoot(t)
|
||||
manager := newExecSessionManager()
|
||||
execTool := NewScopedExecCommandTool(root, nil, manager)
|
||||
writeTool := NewWriteStdinTool(manager)
|
||||
@@ -556,16 +621,46 @@ func TestWriteStdinStopIntentTerminatesNonTTYSession(t *testing.T) {
|
||||
t.Fatalf("session_id is not numeric: %v", err)
|
||||
}
|
||||
|
||||
result := writeTool.Run(context.Background(), map[string]any{
|
||||
"session_id": sessionID,
|
||||
"chars": chars,
|
||||
"yield_time_ms": 1000,
|
||||
})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("stop input %q status = %s: %s", chars, result.Status, result.Output)
|
||||
}
|
||||
if result.Meta["session_id"] != "" {
|
||||
t.Fatalf("stop input %q should not leave session running, meta=%#v output=%q", chars, result.Meta, result.Output)
|
||||
// Re-send the stop until the session reports it has exited, rather than
|
||||
// waiting a fixed slice and asserting on whatever is true by then.
|
||||
//
|
||||
// A single 1s yield used to be enough because the session ran under
|
||||
// cmd.exe. Under PowerShell the tree is slower to stand up and tear down,
|
||||
// and the FIRST stop in a fresh process pays the interpreter's cold start,
|
||||
// so it routinely returns with termination already under way but not
|
||||
// finished. The old assertion read that as "the stop did not work" and
|
||||
// failed, intermittently, on whichever of the two inputs happened to run
|
||||
// first. CI never saw it because the runners have PowerShell 7 while a
|
||||
// stock Windows box only has 5.1.
|
||||
//
|
||||
// Repeating the same input rather than polling with empty chars is
|
||||
// deliberate: Terminate is a no-op against an already-dead tree, so it is
|
||||
// safe to repeat, and it keeps the interrupt flag that an empty-chars poll
|
||||
// would drop, which is what the interrupted assertion below needs.
|
||||
var result Result
|
||||
// Every poll's output is kept, not just the last one. Each Continue
|
||||
// returns only what the session produced since the previous collection,
|
||||
// so a marker printed just before the poll that observes the exit lands
|
||||
// in the earlier result and would be invisible to a check against the
|
||||
// final one alone.
|
||||
var collected strings.Builder
|
||||
deadline := time.Now().Add(30 * time.Second)
|
||||
for {
|
||||
result = writeTool.Run(context.Background(), map[string]any{
|
||||
"session_id": sessionID,
|
||||
"chars": chars,
|
||||
"yield_time_ms": 500,
|
||||
})
|
||||
collected.WriteString(result.Output)
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("stop input %q status = %s: %s", chars, result.Status, result.Output)
|
||||
}
|
||||
if result.Meta["session_id"] == "" {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("stop input %q left the session running past the deadline, meta=%#v output=%q", chars, result.Meta, result.Output)
|
||||
}
|
||||
}
|
||||
if result.Meta["exit_code"] == "" {
|
||||
t.Fatalf("stop input %q should report exit_code, meta=%#v output=%q", chars, result.Meta, result.Output)
|
||||
@@ -573,6 +668,17 @@ func TestWriteStdinStopIntentTerminatesNonTTYSession(t *testing.T) {
|
||||
if result.Meta["interrupted"] != "true" {
|
||||
t.Fatalf("stop input %q should report interrupted metadata, meta=%#v output=%q", chars, result.Meta, result.Output)
|
||||
}
|
||||
// The session has to have been killed rather than allowed to finish.
|
||||
// Without this the retry loop makes the test unfalsifiable: the helper
|
||||
// sleeps 5s, well inside the deadline, so a terminate that did nothing at
|
||||
// all would let the sleep end on its own and every assertion above would
|
||||
// still hold. The interrupted flag does not close the hole either, since
|
||||
// it echoes the request rather than the outcome. The helper prints this
|
||||
// line only on the way out of a completed sleep, so its absence is what
|
||||
// separates the two, and it says nothing about how long the loop took.
|
||||
if strings.Contains(collected.String(), "long sleep finished") {
|
||||
t.Fatalf("stop input %q let the session run to completion instead of terminating it, output=%q", chars, collected.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -612,7 +718,7 @@ func TestExecCommandTTYSessionAcceptsInputOnLinux(t *testing.T) {
|
||||
if runtime.GOOS != "linux" {
|
||||
t.Skip("pty transport is currently implemented for linux")
|
||||
}
|
||||
root := t.TempDir()
|
||||
root := execTestRoot(t)
|
||||
manager := newExecSessionManager()
|
||||
execTool := NewScopedExecCommandTool(root, nil, manager)
|
||||
writeTool := NewWriteStdinTool(manager)
|
||||
@@ -650,7 +756,7 @@ func TestExecCommandTTYSessionAcceptsInputOnLinux(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestExecSessionSnapshotsAndStopAll(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
root := execTestRoot(t)
|
||||
manager := newExecSessionManager()
|
||||
execTool := NewScopedExecCommandTool(root, nil, manager).(execCommandTool)
|
||||
|
||||
|
||||
@@ -79,6 +79,147 @@ func TestEditFileSucceedsWhenUnchangedAndRebaselines(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyPatchPreservesWholeFileObservationForFollowUpEdit(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "f.txt")
|
||||
if err := os.WriteFile(path, []byte("alpha\nbeta\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
registry := safetyRegistry(t, dir)
|
||||
registry.Register(NewScopedApplyPatchTool(dir, nil))
|
||||
tracker := NewFileTracker()
|
||||
opts := grantedOpts(tracker)
|
||||
|
||||
if result := registry.RunWithOptions(context.Background(), "read_file", map[string]any{"path": "f.txt"}, opts); result.Status != StatusOK {
|
||||
t.Fatalf("read_file failed: %s", result.Output)
|
||||
}
|
||||
patch := strings.Join([]string{
|
||||
"diff --git a/f.txt b/f.txt",
|
||||
"--- a/f.txt",
|
||||
"+++ b/f.txt",
|
||||
"@@ -1,2 +1,2 @@",
|
||||
"-alpha",
|
||||
"+gamma",
|
||||
" beta",
|
||||
"",
|
||||
}, "\n")
|
||||
if result := registry.RunWithOptions(context.Background(), "apply_patch", map[string]any{"patch": patch}, opts); result.Status != StatusOK {
|
||||
if gitApplyUnavailable(result.Output) {
|
||||
t.Skipf("git binary unavailable: %s", result.Output)
|
||||
}
|
||||
t.Fatalf("apply_patch failed: %s", result.Output)
|
||||
}
|
||||
|
||||
result := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
|
||||
"path": "f.txt", "old_string": "beta", "new_string": "delta",
|
||||
}, opts)
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("follow-up edit should not require a wasted re-read: %s", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyPatchRenameOrCopyDoesNotCreditUnreadDestination(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
patch string
|
||||
}{
|
||||
{
|
||||
name: "rename",
|
||||
patch: strings.Join([]string{
|
||||
"diff --git a/source.txt b/destination.txt",
|
||||
"similarity index 100%",
|
||||
"rename from source.txt",
|
||||
"rename to destination.txt",
|
||||
"",
|
||||
}, "\n"),
|
||||
},
|
||||
{
|
||||
name: "copy",
|
||||
patch: strings.Join([]string{
|
||||
"diff --git a/source.txt b/destination.txt",
|
||||
"similarity index 100%",
|
||||
"copy from source.txt",
|
||||
"copy to destination.txt",
|
||||
"",
|
||||
}, "\n"),
|
||||
},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(dir, "source.txt"), []byte("unread source\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
registry := safetyRegistry(t, dir)
|
||||
registry.Register(NewScopedApplyPatchTool(dir, nil))
|
||||
tracker := NewFileTracker()
|
||||
opts := grantedOpts(tracker)
|
||||
|
||||
if result := registry.RunWithOptions(context.Background(), "apply_patch", map[string]any{"patch": tc.patch}, opts); result.Status != StatusOK {
|
||||
if gitApplyUnavailable(result.Output) {
|
||||
t.Skipf("git binary unavailable: %s", result.Output)
|
||||
}
|
||||
t.Fatalf("apply_patch failed: %s", result.Output)
|
||||
}
|
||||
destination, err := filepath.EvalSymlinks(filepath.Join(dir, "destination.txt"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if tracker.SeenWhole(destination) {
|
||||
t.Fatal("rename/copy destination was credited as fully seen without a read")
|
||||
}
|
||||
result := registry.RunWithOptions(context.Background(), "write_file", map[string]any{
|
||||
"path": "destination.txt", "content": "blind overwrite\n", "overwrite": true,
|
||||
}, opts)
|
||||
if result.Status != StatusError || !strings.Contains(result.Output, "has not been read exactly") {
|
||||
t.Fatalf("blind destination overwrite was not refused: %#v", result)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyPatchCompleteCreationCreditsSuppliedFile(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
registry := safetyRegistry(t, dir)
|
||||
registry.Register(NewScopedApplyPatchTool(dir, nil))
|
||||
tracker := NewFileTracker()
|
||||
opts := grantedOpts(tracker)
|
||||
patch := strings.Join([]string{
|
||||
"diff --git a/new.txt b/new.txt",
|
||||
"new file mode 100644",
|
||||
"--- /dev/null",
|
||||
"+++ b/new.txt",
|
||||
"@@ -0,0 +1 @@",
|
||||
"+supplied content",
|
||||
"",
|
||||
}, "\n")
|
||||
if result := registry.RunWithOptions(context.Background(), "apply_patch", map[string]any{"patch": patch}, opts); result.Status != StatusOK {
|
||||
if gitApplyUnavailable(result.Output) {
|
||||
t.Skipf("git binary unavailable: %s", result.Output)
|
||||
}
|
||||
t.Fatalf("apply_patch failed: %s", result.Output)
|
||||
}
|
||||
created, err := filepath.EvalSymlinks(filepath.Join(dir, "new.txt"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !tracker.SeenWhole(created) {
|
||||
t.Fatal("complete /dev/null creation was not credited as supplied content")
|
||||
}
|
||||
result := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
|
||||
"path": "new.txt", "old_string": "supplied", "new_string": "updated",
|
||||
}, opts)
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("follow-up edit of supplied file failed: %s", result.Output)
|
||||
}
|
||||
updated, err := os.ReadFile(filepath.Join(dir, "new.txt"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, want := strings.ReplaceAll(string(updated), "\r\n", "\n"), "updated content\n"; got != want {
|
||||
t.Fatalf("follow-up edit content = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteFileOverwriteRefusedWhenChangedOnDisk(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "f.txt")
|
||||
@@ -103,7 +244,7 @@ func TestWriteFileOverwriteRefusedWhenChangedOnDisk(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteFileOverwriteAllowedForUntrackedFile(t *testing.T) {
|
||||
func TestWriteFileOverwriteRequiresExactRead(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "f.txt")
|
||||
if err := os.WriteFile(path, []byte("v1\n"), 0o644); err != nil {
|
||||
@@ -112,15 +253,127 @@ func TestWriteFileOverwriteAllowedForUntrackedFile(t *testing.T) {
|
||||
registry := safetyRegistry(t, dir)
|
||||
tracker := NewFileTracker()
|
||||
|
||||
// Never read through a tool → no baseline → overwrite is the model's call.
|
||||
// Never read through a tool: an overwrite would discard unseen content.
|
||||
res := registry.RunWithOptions(context.Background(), "write_file", map[string]any{
|
||||
"path": "f.txt", "content": "v3\n", "overwrite": true,
|
||||
}, grantedOpts(tracker))
|
||||
if res.Status != StatusOK {
|
||||
t.Fatalf("untracked overwrite should be allowed, got %q", res.Output)
|
||||
if res.Status != StatusError || !strings.Contains(res.Output, "has not been read exactly") {
|
||||
t.Fatalf("unseen overwrite should be refused, got %q", res.Output)
|
||||
}
|
||||
if got, _ := os.ReadFile(path); string(got) != "v3\n" {
|
||||
t.Fatalf("file = %q, want v3", got)
|
||||
if got, _ := os.ReadFile(path); string(got) != "v1\n" {
|
||||
t.Fatalf("file = %q, want original content", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFileAllowsSeenRangeAndRejectsUnseenRange(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "f.txt")
|
||||
if err := os.WriteFile(path, []byte("alpha\nbeta\ngamma\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
registry := safetyRegistry(t, dir)
|
||||
tracker := NewFileTracker()
|
||||
opts := grantedOpts(tracker)
|
||||
|
||||
if res := registry.RunWithOptions(context.Background(), "read_file", map[string]any{
|
||||
"path": "f.txt", "start_line": 2, "end_line": 2,
|
||||
}, opts); res.Status != StatusOK {
|
||||
t.Fatalf("partial read failed: %s", res.Output)
|
||||
}
|
||||
if res := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
|
||||
"path": "f.txt", "old_string": "gamma", "new_string": "delta",
|
||||
}, opts); res.Status != StatusError || !strings.Contains(res.Output, "has not been read exactly") {
|
||||
t.Fatalf("unseen edit should be refused, got %q", res.Output)
|
||||
}
|
||||
if res := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
|
||||
"path": "f.txt", "old_string": "beta", "new_string": "bravo",
|
||||
}, opts); res.Status != StatusOK {
|
||||
t.Fatalf("seen-range edit should succeed, got %q", res.Output)
|
||||
}
|
||||
if content, err := os.ReadFile(path); err != nil || !strings.Contains(string(content), "bravo") {
|
||||
t.Fatalf("seen-range edit did not update the file: content=%q err=%v", content, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEditFileReplaceAllRecordsEveryShiftedReplacementSpan(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "f.txt")
|
||||
if err := os.WriteFile(path, []byte("before\nneedle\nmiddle\nneedle\nafter\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
registry := safetyRegistry(t, dir)
|
||||
tracker := NewFileTracker()
|
||||
opts := grantedOpts(tracker)
|
||||
|
||||
if res := registry.RunWithOptions(context.Background(), "read_file", map[string]any{
|
||||
"path": "f.txt", "start_line": 2, "end_line": 4,
|
||||
}, opts); res.Status != StatusOK {
|
||||
t.Fatalf("partial read failed: %q", res.Output)
|
||||
}
|
||||
if res := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
|
||||
"path": "f.txt",
|
||||
"old_string": "needle",
|
||||
"new_string": "left\nright",
|
||||
"replace_all": true,
|
||||
}, opts); res.Status != StatusOK {
|
||||
t.Fatalf("replace_all failed: %q", res.Output)
|
||||
}
|
||||
// This follow-up targets both generated "right" lines. It succeeds only if
|
||||
// replace_all recorded the shifted second replacement as seen.
|
||||
if res := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
|
||||
"path": "f.txt",
|
||||
"old_string": "right",
|
||||
"new_string": "done",
|
||||
"replace_all": true,
|
||||
}, opts); res.Status != StatusOK {
|
||||
t.Fatalf("follow-up edit of every replacement should succeed: %q", res.Output)
|
||||
}
|
||||
if content, err := os.ReadFile(path); err != nil || strings.Count(string(content), "done") != 2 {
|
||||
t.Fatalf("unexpected follow-up content: content=%q err=%v", content, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteFileCanOverwriteNewlyCreatedEmptyFile(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
registry := safetyRegistry(t, dir)
|
||||
tracker := NewFileTracker()
|
||||
opts := grantedOpts(tracker)
|
||||
|
||||
if res := registry.RunWithOptions(context.Background(), "write_file", map[string]any{
|
||||
"path": "empty.txt", "content": "",
|
||||
}, opts); res.Status != StatusOK {
|
||||
t.Fatalf("empty create failed: %q", res.Output)
|
||||
}
|
||||
if res := registry.RunWithOptions(context.Background(), "write_file", map[string]any{
|
||||
"path": "empty.txt", "content": "now populated\n", "overwrite": true,
|
||||
}, opts); res.Status != StatusOK {
|
||||
t.Fatalf("empty overwrite should succeed: %q", res.Output)
|
||||
}
|
||||
if content, err := os.ReadFile(filepath.Join(dir, "empty.txt")); err != nil || string(content) != "now populated\n" {
|
||||
t.Fatalf("unexpected overwritten content: content=%q err=%v", content, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMinifiedReadDoesNotAuthorizeExactEdit(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "f.go")
|
||||
if err := os.WriteFile(path, []byte("package p\n\nvar value = 1\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
registry := NewRegistry()
|
||||
registry.Register(NewScopedReadMinifiedFileTool(dir, nil))
|
||||
registry.Register(NewScopedEditFileTool(dir, nil))
|
||||
tracker := NewFileTracker()
|
||||
opts := grantedOpts(tracker)
|
||||
|
||||
if res := registry.RunWithOptions(context.Background(), "read_minified_file", map[string]any{"path": "f.go"}, opts); res.Status != StatusOK {
|
||||
t.Fatalf("minified read failed: %s", res.Output)
|
||||
}
|
||||
res := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
|
||||
"path": "f.go", "old_string": "var value = 1", "new_string": "var value = 2",
|
||||
}, opts)
|
||||
if res.Status != StatusError || !strings.Contains(res.Output, "has not been read exactly") {
|
||||
t.Fatalf("minified view must not authorize exact edit, got %q", res.Output)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"os"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
@@ -34,6 +35,7 @@ type FileVersion struct {
|
||||
type FileTracker struct {
|
||||
mu sync.Mutex
|
||||
versions map[string]FileVersion
|
||||
seen map[string]fileObservation
|
||||
// created preserves the order in which brand-new files (paths that did not
|
||||
// exist before a write tool created them) were first written this session.
|
||||
// A map would lose that order and complicates de-duplication, so a slice
|
||||
@@ -43,7 +45,36 @@ type FileTracker struct {
|
||||
}
|
||||
|
||||
func NewFileTracker() *FileTracker {
|
||||
return &FileTracker{versions: make(map[string]FileVersion)}
|
||||
return &FileTracker{
|
||||
versions: make(map[string]FileVersion),
|
||||
seen: make(map[string]fileObservation),
|
||||
}
|
||||
}
|
||||
|
||||
type lineRange struct {
|
||||
start int
|
||||
end int
|
||||
}
|
||||
|
||||
type fileObservation struct {
|
||||
whole bool
|
||||
// total is the file's line count as of the last recorded read, so coverage
|
||||
// can be judged from the ranges alone. Without it a file read in pieces
|
||||
// could never be known to have been seen in full.
|
||||
total int
|
||||
ranges []lineRange
|
||||
totalBytes int
|
||||
byteRanges []lineRange // zero-based, half-open byte intervals
|
||||
}
|
||||
|
||||
type pendingFileObservation struct {
|
||||
path string
|
||||
output string
|
||||
hash string
|
||||
start int
|
||||
end int
|
||||
total int
|
||||
byteMode bool
|
||||
}
|
||||
|
||||
// Record stores the version of absPath given its content and optional stat info.
|
||||
@@ -63,10 +94,200 @@ func (tracker *FileTracker) RecordHash(absPath string, hash string, info os.File
|
||||
version.MTime = info.ModTime()
|
||||
}
|
||||
tracker.mu.Lock()
|
||||
if previous, ok := tracker.versions[absPath]; !ok || previous.Hash != hash {
|
||||
delete(tracker.seen, absPath)
|
||||
}
|
||||
tracker.versions[absPath] = version
|
||||
tracker.mu.Unlock()
|
||||
}
|
||||
|
||||
// RecordSeenRange records exact, non-elided source lines returned to the model.
|
||||
// Ranges accumulate while the content hash remains unchanged.
|
||||
func (tracker *FileTracker) RecordSeenRange(absPath string, start, end, total int) {
|
||||
if tracker == nil || start < 1 {
|
||||
return
|
||||
}
|
||||
tracker.mu.Lock()
|
||||
defer tracker.mu.Unlock()
|
||||
observation := tracker.seen[absPath]
|
||||
if total == 0 {
|
||||
observation.whole = true
|
||||
observation.ranges = nil
|
||||
tracker.seen[absPath] = observation
|
||||
return
|
||||
}
|
||||
if end < start {
|
||||
return
|
||||
}
|
||||
// A new line count means the file changed shape since the last read, so
|
||||
// previously recorded ranges no longer describe it.
|
||||
if observation.total != 0 && observation.total != total {
|
||||
observation = fileObservation{}
|
||||
}
|
||||
observation.total = total
|
||||
if start == 1 && end >= total {
|
||||
observation.whole = true
|
||||
observation.ranges = nil
|
||||
tracker.seen[absPath] = observation
|
||||
return
|
||||
}
|
||||
observation.ranges = append(observation.ranges, lineRange{start: start, end: end})
|
||||
tracker.seen[absPath] = observation
|
||||
}
|
||||
|
||||
// RecordSeenBytes records an exact, non-elided byte interval returned to the
|
||||
// model. Byte intervals make oversized single-line files recoverable without
|
||||
// crediting the model for the omitted middle of a truncated line read.
|
||||
func (tracker *FileTracker) RecordSeenBytes(absPath string, start, end, total int) {
|
||||
if tracker == nil || start < 0 || end < start || total < 0 || end > total {
|
||||
return
|
||||
}
|
||||
tracker.mu.Lock()
|
||||
defer tracker.mu.Unlock()
|
||||
observation := tracker.seen[absPath]
|
||||
if total == 0 {
|
||||
observation.whole = true
|
||||
observation.byteRanges = nil
|
||||
tracker.seen[absPath] = observation
|
||||
return
|
||||
}
|
||||
if observation.totalBytes != 0 && observation.totalBytes != total {
|
||||
observation = fileObservation{}
|
||||
}
|
||||
observation.totalBytes = total
|
||||
if start == 0 && end >= total {
|
||||
observation.whole = true
|
||||
observation.byteRanges = nil
|
||||
tracker.seen[absPath] = observation
|
||||
return
|
||||
}
|
||||
observation.byteRanges = append(observation.byteRanges, lineRange{start: start, end: end})
|
||||
tracker.seen[absPath] = observation
|
||||
}
|
||||
|
||||
// coversFully reports whether ranges together cover every line in [start, end].
|
||||
//
|
||||
// Ranges are merged rather than scanned line by line: a caller asking about a
|
||||
// large file would otherwise walk every line against every recorded range, and
|
||||
// the answer is the same.
|
||||
func coversFully(ranges []lineRange, start int, end int) bool {
|
||||
if start > end {
|
||||
return true
|
||||
}
|
||||
if len(ranges) == 0 {
|
||||
return false
|
||||
}
|
||||
sorted := make([]lineRange, len(ranges))
|
||||
copy(sorted, ranges)
|
||||
sort.Slice(sorted, func(left int, right int) bool {
|
||||
if sorted[left].start != sorted[right].start {
|
||||
return sorted[left].start < sorted[right].start
|
||||
}
|
||||
return sorted[left].end < sorted[right].end
|
||||
})
|
||||
next := start
|
||||
for _, seen := range sorted {
|
||||
if seen.start > next {
|
||||
return false
|
||||
}
|
||||
if seen.end >= next {
|
||||
next = seen.end + 1
|
||||
}
|
||||
if next > end {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return next > end
|
||||
}
|
||||
|
||||
// coversBytes reports whether half-open byte ranges cover [start, end).
|
||||
func coversBytes(ranges []lineRange, start, end int) bool {
|
||||
if start >= end {
|
||||
return true
|
||||
}
|
||||
if len(ranges) == 0 {
|
||||
return false
|
||||
}
|
||||
sorted := make([]lineRange, len(ranges))
|
||||
copy(sorted, ranges)
|
||||
sort.Slice(sorted, func(left, right int) bool {
|
||||
if sorted[left].start != sorted[right].start {
|
||||
return sorted[left].start < sorted[right].start
|
||||
}
|
||||
return sorted[left].end < sorted[right].end
|
||||
})
|
||||
next := start
|
||||
for _, seen := range sorted {
|
||||
if seen.start > next {
|
||||
return false
|
||||
}
|
||||
if seen.end > next {
|
||||
next = seen.end
|
||||
}
|
||||
if next >= end {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// SeenRange reports whether every line in the requested range was returned
|
||||
// exactly to the model for the currently tracked version.
|
||||
func (tracker *FileTracker) SeenRange(absPath string, start, end int) bool {
|
||||
if tracker == nil {
|
||||
return true
|
||||
}
|
||||
tracker.mu.Lock()
|
||||
defer tracker.mu.Unlock()
|
||||
observation, ok := tracker.seen[absPath]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
if observation.whole {
|
||||
return true
|
||||
}
|
||||
return coversFully(observation.ranges, start, end)
|
||||
}
|
||||
|
||||
// SeenBytes reports whether every byte in [start, end) was returned exactly to
|
||||
// the model for the currently tracked version.
|
||||
func (tracker *FileTracker) SeenBytes(absPath string, start, end int) bool {
|
||||
if tracker == nil {
|
||||
return true
|
||||
}
|
||||
tracker.mu.Lock()
|
||||
defer tracker.mu.Unlock()
|
||||
observation, ok := tracker.seen[absPath]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return observation.whole || coversBytes(observation.byteRanges, start, end)
|
||||
}
|
||||
|
||||
func (tracker *FileTracker) SeenWhole(absPath string) bool {
|
||||
if tracker == nil {
|
||||
return true
|
||||
}
|
||||
tracker.mu.Lock()
|
||||
defer tracker.mu.Unlock()
|
||||
observation, ok := tracker.seen[absPath]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
if observation.whole {
|
||||
return true
|
||||
}
|
||||
// Derived from the ranges rather than tracked as its own flag. The flag was
|
||||
// only ever set by a SINGLE read covering the file, so a file read in two
|
||||
// halves stayed "not seen whole" forever even though SeenRange agreed every
|
||||
// line had been seen — and write_file, which gates on this, then refused the
|
||||
// overwrite with advice to read the file that could not change the answer.
|
||||
if observation.total > 0 && coversFully(observation.ranges, 1, observation.total) {
|
||||
return true
|
||||
}
|
||||
return observation.totalBytes > 0 && coversBytes(observation.byteRanges, 0, observation.totalBytes)
|
||||
}
|
||||
|
||||
// Version returns the recorded version for absPath and whether one exists.
|
||||
func (tracker *FileTracker) Version(absPath string) (FileVersion, bool) {
|
||||
if tracker == nil {
|
||||
@@ -119,6 +340,7 @@ func (tracker *FileTracker) Forget(absPath string) {
|
||||
}
|
||||
tracker.mu.Lock()
|
||||
delete(tracker.versions, absPath)
|
||||
delete(tracker.seen, absPath)
|
||||
tracker.mu.Unlock()
|
||||
}
|
||||
|
||||
@@ -153,3 +375,7 @@ func fileConflictMessage(relativePath string) string {
|
||||
return "Error writing " + relativePath + ": " + ErrFileChangedOnDisk.Error() +
|
||||
" (it may have been edited outside Zero). Re-read it with read_file, then re-apply your change so you do not overwrite the newer content."
|
||||
}
|
||||
|
||||
func fileUnseenMessage(relativePath string) string {
|
||||
return "Error writing " + relativePath + ": the intended change depends on content that has not been read exactly in this session. Read the affected lines with read_file, or use byte_offset/byte_limit for an oversized single line, then retry the edit."
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// SeenWhole used to be a flag set only by a SINGLE read that covered the file,
|
||||
// while SeenRange judged coverage from the accumulated ranges. The two then
|
||||
// disagreed about the same reads, and write_file — which gates on SeenWhole —
|
||||
// refused overwrites it should have allowed.
|
||||
func TestSeenWholeIsDerivedFromAccumulatedRanges(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "f.go")
|
||||
|
||||
for name, testCase := range map[string]struct {
|
||||
reads [][3]int // start, end, total
|
||||
want bool
|
||||
}{
|
||||
"one read covering the file": {
|
||||
reads: [][3]int{{1, 100, 100}},
|
||||
want: true,
|
||||
},
|
||||
"two adjacent reads covering it together": {
|
||||
reads: [][3]int{{1, 50, 100}, {51, 100, 100}},
|
||||
want: true,
|
||||
},
|
||||
"overlapping reads covering it together": {
|
||||
reads: [][3]int{{1, 60, 100}, {40, 100, 100}},
|
||||
want: true,
|
||||
},
|
||||
"out-of-order reads covering it together": {
|
||||
reads: [][3]int{{51, 100, 100}, {1, 50, 100}},
|
||||
want: true,
|
||||
},
|
||||
"a gap in the middle": {
|
||||
reads: [][3]int{{1, 40, 100}, {60, 100, 100}},
|
||||
want: false,
|
||||
},
|
||||
"missing the first line": {
|
||||
reads: [][3]int{{2, 100, 100}},
|
||||
want: false,
|
||||
},
|
||||
"missing the last line": {
|
||||
reads: [][3]int{{1, 99, 100}},
|
||||
want: false,
|
||||
},
|
||||
"nothing read at all": {
|
||||
reads: nil,
|
||||
want: false,
|
||||
},
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
tracker := NewFileTracker()
|
||||
for _, read := range testCase.reads {
|
||||
tracker.RecordSeenRange(path, read[0], read[1], read[2])
|
||||
}
|
||||
if got := tracker.SeenWhole(path); got != testCase.want {
|
||||
t.Errorf("SeenWhole = %v, want %v after reads %v", got, testCase.want, testCase.reads)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// SeenWhole and SeenRange must not disagree about the same reads.
|
||||
func TestSeenWholeAgreesWithSeenRange(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "f.go")
|
||||
tracker := NewFileTracker()
|
||||
tracker.RecordSeenRange(path, 1, 50, 100)
|
||||
tracker.RecordSeenRange(path, 51, 100, 100)
|
||||
|
||||
if !tracker.SeenRange(path, 1, 100) {
|
||||
t.Fatal("SeenRange says the whole file was not seen")
|
||||
}
|
||||
if !tracker.SeenWhole(path) {
|
||||
t.Error("SeenRange covers every line but SeenWhole disagrees")
|
||||
}
|
||||
}
|
||||
|
||||
// A file that changed length since the recorded reads must not stay "seen":
|
||||
// the old ranges describe a shape the file no longer has.
|
||||
func TestSeenWholeResetsWhenTheFileChangesLength(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "f.go")
|
||||
tracker := NewFileTracker()
|
||||
tracker.RecordSeenRange(path, 1, 100, 100)
|
||||
if !tracker.SeenWhole(path) {
|
||||
t.Fatal("precondition: the whole file was read")
|
||||
}
|
||||
// The file grew; only the first 100 lines of 200 have been seen.
|
||||
tracker.RecordSeenRange(path, 1, 100, 200)
|
||||
if tracker.SeenWhole(path) {
|
||||
t.Error("file grew to 200 lines but is still reported as seen whole")
|
||||
}
|
||||
if !tracker.SeenRange(path, 1, 100) {
|
||||
t.Error("the lines that were read should still count as seen")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,263 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// A file past the read byte budget cannot be read whole in one call: the read is
|
||||
// truncated and RecordSeenRange is skipped, which is correct — the model did not
|
||||
// see those lines. But SeenWhole was a flag only a single whole-file read could
|
||||
// set, so no sequence of scoped reads could ever clear it either, and write_file
|
||||
// refused the overwrite forever while telling the model to read the file again.
|
||||
func TestLargeFileBecomesWritableAfterScopedReadsCoverIt(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
// read_file resolves the workspace root through EvalSymlinks before it
|
||||
// records anything, so the tracker is keyed on the resolved path. On macOS
|
||||
// t.TempDir() sits under /var, a symlink to /private/var, and asserting
|
||||
// against the raw path would look up a key that never existed — passing on
|
||||
// Windows and failing there.
|
||||
resolvedDir, err := filepath.EvalSymlinks(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("resolve temp dir: %v", err)
|
||||
}
|
||||
path := filepath.Join(resolvedDir, "big.go")
|
||||
|
||||
const lines = 2400
|
||||
var builder strings.Builder
|
||||
for index := 0; index < lines; index++ {
|
||||
fmt.Fprintf(&builder, "line %06d %s\n", index, strings.Repeat("x", 60))
|
||||
}
|
||||
if writeErr := os.WriteFile(path, []byte(builder.String()), 0o600); writeErr != nil {
|
||||
t.Fatalf("seed file: %v", writeErr)
|
||||
}
|
||||
|
||||
tracker := NewFileTracker()
|
||||
options := RunOptions{FileTracker: tracker}
|
||||
ctx := context.Background()
|
||||
read := NewReadFileTool(dir).(interface {
|
||||
RunWithOptions(context.Context, map[string]any, RunOptions) Result
|
||||
})
|
||||
write := NewScopedWriteFileTool(dir, nil).(interface {
|
||||
RunWithOptions(context.Context, map[string]any, RunOptions) Result
|
||||
})
|
||||
|
||||
// One unbounded read is byte-truncated, so nothing is recorded.
|
||||
full := read.RunWithOptions(ctx, map[string]any{"path": "big.go"}, options)
|
||||
if full.Meta["truncation_reason"] != "byte_budget" {
|
||||
t.Fatalf("precondition: expected the file to exceed the read byte budget, got reason %q", full.Meta["truncation_reason"])
|
||||
}
|
||||
if tracker.SeenWhole(path) {
|
||||
t.Fatal("a truncated read must not credit the model with the whole file")
|
||||
}
|
||||
|
||||
// Scoped reads that together cover it must.
|
||||
for _, span := range [][2]int{{1, 1200}, {1201, 2400}} {
|
||||
result := read.RunWithOptions(ctx, map[string]any{
|
||||
"path": "big.go", "start_line": span[0], "end_line": span[1],
|
||||
}, options)
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("scoped read %v: %s", span, result.Output)
|
||||
}
|
||||
if result.Meta["truncation_reason"] == "byte_budget" {
|
||||
t.Fatalf("scoped read %v was itself truncated; pick smaller spans", span)
|
||||
}
|
||||
}
|
||||
if !tracker.SeenWhole(path) {
|
||||
t.Fatal("scoped reads covered every line but the file is still not seen whole")
|
||||
}
|
||||
|
||||
overwrite := write.RunWithOptions(ctx, map[string]any{
|
||||
"path": "big.go", "content": "replaced\n", "overwrite": true,
|
||||
}, options)
|
||||
if overwrite.Status != StatusOK {
|
||||
t.Fatalf("overwrite refused after the file was fully read in scoped reads: %s", overwrite.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSingleLongLineBecomesWritableAfterExactByteReads(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "bundle.min.js")
|
||||
content := strings.Repeat("x", readOutputBudgetBytes*3)
|
||||
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
|
||||
t.Fatalf("seed file: %v", err)
|
||||
}
|
||||
resolvedPath, err := filepath.EvalSymlinks(path)
|
||||
if err != nil {
|
||||
t.Fatalf("resolve file: %v", err)
|
||||
}
|
||||
|
||||
tracker := NewFileTracker()
|
||||
options := RunOptions{FileTracker: tracker}
|
||||
ctx := context.Background()
|
||||
read := NewReadFileTool(dir).(interface {
|
||||
RunWithOptions(context.Context, map[string]any, RunOptions) Result
|
||||
})
|
||||
write := NewScopedWriteFileTool(dir, nil).(interface {
|
||||
RunWithOptions(context.Context, map[string]any, RunOptions) Result
|
||||
})
|
||||
|
||||
full := read.RunWithOptions(ctx, map[string]any{"path": "bundle.min.js"}, options)
|
||||
if full.Meta["truncation_reason"] != "byte_budget" {
|
||||
t.Fatalf("precondition: expected byte truncation, got %#v", full.Meta)
|
||||
}
|
||||
if tracker.SeenWhole(resolvedPath) {
|
||||
t.Fatal("a truncated line read must not credit the whole file")
|
||||
}
|
||||
|
||||
const chunkSize = 64 * 1024
|
||||
for offset := 0; offset < len(content); offset += chunkSize {
|
||||
result := read.RunWithOptions(ctx, map[string]any{
|
||||
"path": "bundle.min.js",
|
||||
"byte_offset": offset,
|
||||
"byte_limit": chunkSize,
|
||||
}, options)
|
||||
if result.Status != StatusOK || result.Truncated {
|
||||
t.Fatalf("byte read at %d was not exact: status=%s truncated=%v output=%q", offset, result.Status, result.Truncated, result.Output)
|
||||
}
|
||||
}
|
||||
if !tracker.SeenWhole(resolvedPath) {
|
||||
t.Fatal("exact byte reads covered the file but it is not seen whole")
|
||||
}
|
||||
|
||||
overwrite := write.RunWithOptions(ctx, map[string]any{
|
||||
"path": "bundle.min.js", "content": "replaced\n", "overwrite": true,
|
||||
}, options)
|
||||
if overwrite.Status != StatusOK {
|
||||
t.Fatalf("overwrite refused after exact byte reads: %s", overwrite.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExactByteReadsAdvanceAlongUTF8Boundaries(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "unicode.min.js")
|
||||
content := strings.Repeat("😀", 40_000)
|
||||
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resolvedPath, err := filepath.EvalSymlinks(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
tracker := NewFileTracker()
|
||||
options := RunOptions{FileTracker: tracker}
|
||||
read := NewReadFileTool(dir).(interface {
|
||||
RunWithOptions(context.Context, map[string]any, RunOptions) Result
|
||||
})
|
||||
offset := 0
|
||||
for offset < len(content) {
|
||||
result := read.RunWithOptions(context.Background(), map[string]any{
|
||||
"path": "unicode.min.js", "byte_offset": offset, "byte_limit": 65_535,
|
||||
}, options)
|
||||
if result.Status != StatusOK || result.Truncated || !utf8.ValidString(result.Output) {
|
||||
t.Fatalf("invalid exact UTF-8 read at %d: %#v", offset, result)
|
||||
}
|
||||
next, parseErr := strconv.Atoi(result.Meta["next_byte_offset"])
|
||||
if parseErr != nil || next <= offset {
|
||||
t.Fatalf("invalid continuation offset %q after %d", result.Meta["next_byte_offset"], offset)
|
||||
}
|
||||
offset = next
|
||||
}
|
||||
if !tracker.SeenWhole(resolvedPath) {
|
||||
t.Fatal("UTF-8-safe byte reads covered the file but it is not seen whole")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryTruncationDoesNotCreditUndeliveredByteRanges(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "unicode.min.js")
|
||||
content := strings.Repeat("界", 84_000)
|
||||
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resolvedPath, err := filepath.EvalSymlinks(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
registry := safetyRegistry(t, dir)
|
||||
tracker := NewFileTracker()
|
||||
result := registry.RunWithOptions(context.Background(), "read_file", map[string]any{
|
||||
"path": "unicode.min.js", "byte_offset": 0, "byte_limit": readFileByteChunkMax,
|
||||
}, grantedOpts(tracker))
|
||||
if !result.Truncated {
|
||||
t.Fatalf("precondition: multibyte chunk was not truncated at the registry boundary: %#v", result.Meta)
|
||||
}
|
||||
next, err := strconv.Atoi(result.Meta["next_byte_offset"])
|
||||
if err != nil || next <= 0 {
|
||||
t.Fatalf("invalid continuation offset %q", result.Meta["next_byte_offset"])
|
||||
}
|
||||
if tracker.SeenBytes(resolvedPath, 0, next) || tracker.SeenWhole(resolvedPath) {
|
||||
t.Fatal("registry-truncated bytes were credited as delivered to the model")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryTruncationDoesNotCreditWholeLineRead(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "unicode.go")
|
||||
content := strings.Repeat("界界界界界界界界\n", 4_000)
|
||||
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resolvedPath, err := filepath.EvalSymlinks(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
registry := safetyRegistry(t, dir)
|
||||
tracker := NewFileTracker()
|
||||
result := registry.RunWithOptions(context.Background(), "read_file", map[string]any{
|
||||
"path": "unicode.go",
|
||||
}, grantedOpts(tracker))
|
||||
if !result.Truncated {
|
||||
t.Fatalf("precondition: multibyte line read was not truncated at the registry boundary: %#v", result.Meta)
|
||||
}
|
||||
if tracker.SeenWhole(resolvedPath) {
|
||||
t.Fatal("registry-truncated line output credited the whole file")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPartialByteEditDoesNotCreditAnEntireLongLine(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "bundle.min.js")
|
||||
content := "UNIQUE_TARGET" + strings.Repeat("x", readOutputBudgetBytes*2)
|
||||
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resolvedPath, err := filepath.EvalSymlinks(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
registry := safetyRegistry(t, dir)
|
||||
tracker := NewFileTracker()
|
||||
opts := grantedOpts(tracker)
|
||||
read := registry.RunWithOptions(context.Background(), "read_file", map[string]any{
|
||||
"path": "bundle.min.js", "byte_offset": 0, "byte_limit": 64,
|
||||
}, opts)
|
||||
if read.Status != StatusOK || read.Truncated {
|
||||
t.Fatalf("exact byte read failed: %#v", read)
|
||||
}
|
||||
edit := registry.RunWithOptions(context.Background(), "edit_file", map[string]any{
|
||||
"path": "bundle.min.js", "old_string": "UNIQUE_TARGET", "new_string": "UPDATED_TARGET",
|
||||
}, opts)
|
||||
if edit.Status != StatusOK {
|
||||
t.Fatalf("edit of exactly observed bytes failed: %s", edit.Output)
|
||||
}
|
||||
if tracker.SeenWhole(resolvedPath) {
|
||||
t.Fatal("editing observed bytes credited the unread remainder of the long line")
|
||||
}
|
||||
overwrite := registry.RunWithOptions(context.Background(), "write_file", map[string]any{
|
||||
"path": "bundle.min.js", "content": "replacement", "overwrite": true,
|
||||
}, opts)
|
||||
if overwrite.Status != StatusError {
|
||||
t.Fatalf("whole overwrite should remain blocked after a partial byte read, got %#v", overwrite)
|
||||
}
|
||||
}
|
||||
@@ -61,15 +61,37 @@ func TestFormatOnWriteFormatsAndKeepsTrackerConsistent(t *testing.T) {
|
||||
t.Fatalf("expected gofmt-formatted content, got %q", onDisk)
|
||||
}
|
||||
|
||||
// The tracker must have been re-baselined to the POST-format content: a
|
||||
// follow-up edit must not trip the external-modification conflict guard.
|
||||
// The tracker is re-baselined to the post-format bytes, but those bytes were
|
||||
// not returned exactly to the model, so an edit must require an exact read.
|
||||
edit := NewScopedEditFileTool(dir, nil).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
|
||||
"path": "a.go",
|
||||
"old_string": "func A() {",
|
||||
"new_string": "func B() {",
|
||||
}, RunOptions{FileTracker: tracker})
|
||||
if edit.Status != StatusError || !strings.Contains(edit.Output, "has not been read exactly") {
|
||||
t.Fatalf("formatter-modified content must require an exact read: %q", edit.Output)
|
||||
}
|
||||
read := NewScopedReadFileTool(dir, nil).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
|
||||
"path": "a.go",
|
||||
}, RunOptions{FileTracker: tracker})
|
||||
if read.Status != StatusOK {
|
||||
t.Fatalf("exact read failed: %q", read.Output)
|
||||
}
|
||||
edit = NewScopedEditFileTool(dir, nil).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
|
||||
"path": "a.go",
|
||||
"old_string": "func A() {",
|
||||
"new_string": "func B( ) {",
|
||||
}, RunOptions{FileTracker: tracker})
|
||||
if edit.Status != StatusOK {
|
||||
t.Fatalf("follow-up edit must not conflict after formatting: %q", edit.Output)
|
||||
t.Fatalf("follow-up edit after exact read failed: %q", edit.Output)
|
||||
}
|
||||
edit = NewScopedEditFileTool(dir, nil).(optionsAwareTool).RunWithOptions(context.Background(), map[string]any{
|
||||
"path": "a.go",
|
||||
"old_string": "func B() {",
|
||||
"new_string": "func C() {",
|
||||
}, RunOptions{FileTracker: tracker})
|
||||
if edit.Status != StatusError || !strings.Contains(edit.Output, "has not been read exactly") {
|
||||
t.Fatalf("formatter-modified edit must require an exact read: %q", edit.Output)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,187 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/sessions"
|
||||
)
|
||||
|
||||
type goalTool struct {
|
||||
baseTool
|
||||
store *sessions.Store
|
||||
sessionID string
|
||||
action string
|
||||
}
|
||||
|
||||
// NewGoalTools returns goal lifecycle tools bound to one captured session. The
|
||||
// binding prevents a late tool call from mutating a different session after a
|
||||
// resume or session switch.
|
||||
func NewGoalTools(store *sessions.Store, sessionID string) []Tool {
|
||||
return []Tool{
|
||||
newGetGoalTool(store, sessionID),
|
||||
newCreateGoalTool(store, sessionID),
|
||||
newUpdateGoalTool(store, sessionID),
|
||||
}
|
||||
}
|
||||
|
||||
func newGetGoalTool(store *sessions.Store, sessionID string) Tool {
|
||||
return &goalTool{
|
||||
baseTool: baseTool{
|
||||
name: "get_goal",
|
||||
description: "Get the current persistent goal and its status, token budget, and usage for this session.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
AdditionalProperties: false,
|
||||
},
|
||||
safety: readOnlySafety("Reads goal state for the current session."),
|
||||
capabilities: ToolCapabilities{Effect: EffectReadOnly, ThreadSafe: true},
|
||||
},
|
||||
store: store,
|
||||
sessionID: sessionID,
|
||||
action: "get",
|
||||
}
|
||||
}
|
||||
|
||||
func newCreateGoalTool(store *sessions.Store, sessionID string) Tool {
|
||||
minimum := 0
|
||||
maximum := 1_000_000_000
|
||||
maxObjectiveLength := sessions.GoalObjectiveMaxLength
|
||||
return &goalTool{
|
||||
baseTool: baseTool{
|
||||
name: "create_goal",
|
||||
description: "Create one persistent goal for this session when the user explicitly asks to pursue an ongoing objective. " +
|
||||
"Do not infer a goal from an ordinary one-turn task. This fails while any goal is still stored.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]PropertySchema{
|
||||
"objective": {
|
||||
Type: "string",
|
||||
Description: "The concrete objective to keep pursuing.",
|
||||
MaxLength: &maxObjectiveLength,
|
||||
},
|
||||
"token_budget": {
|
||||
Type: "integer",
|
||||
Description: "Optional maximum total tokens. Zero means no goal-specific limit.",
|
||||
Minimum: &minimum,
|
||||
Maximum: &maximum,
|
||||
},
|
||||
},
|
||||
Required: []string{"objective"},
|
||||
AdditionalProperties: false,
|
||||
},
|
||||
safety: Safety{
|
||||
SideEffect: SideEffectWrite,
|
||||
Permission: PermissionPrompt,
|
||||
Reason: "Creates persistent goal state and enables automatic follow-up runs for the current session.",
|
||||
AdvertiseInAuto: true,
|
||||
},
|
||||
capabilities: ToolCapabilities{Effect: EffectInteractive, ThreadSafe: false, ResourceKeys: sessionResourceKeys},
|
||||
},
|
||||
store: store,
|
||||
sessionID: sessionID,
|
||||
action: "create",
|
||||
}
|
||||
}
|
||||
|
||||
func newUpdateGoalTool(store *sessions.Store, sessionID string) Tool {
|
||||
return &goalTool{
|
||||
baseTool: baseTool{
|
||||
name: "update_goal",
|
||||
description: "Finish the current goal or mark it blocked. Use complete only when the objective is achieved and no required work remains. " +
|
||||
"Use blocked only when progress cannot continue without user input or an external state change.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]PropertySchema{
|
||||
"status": {
|
||||
Type: "string",
|
||||
Description: "The terminal state to apply.",
|
||||
Enum: []string{string(sessions.GoalStatusComplete), string(sessions.GoalStatusBlocked)},
|
||||
},
|
||||
"reason": {
|
||||
Type: "string",
|
||||
Description: "Why the goal is blocked. Optional for completed goals.",
|
||||
},
|
||||
},
|
||||
Required: []string{"status"},
|
||||
AdditionalProperties: false,
|
||||
},
|
||||
safety: Safety{
|
||||
SideEffect: SideEffectNone,
|
||||
Permission: PermissionAllow,
|
||||
Reason: "Updates persistent goal state for the current session.",
|
||||
},
|
||||
capabilities: ToolCapabilities{Effect: EffectInteractive, ThreadSafe: false, ResourceKeys: sessionResourceKeys},
|
||||
},
|
||||
store: store,
|
||||
sessionID: sessionID,
|
||||
action: "update",
|
||||
}
|
||||
}
|
||||
|
||||
func (tool *goalTool) Run(_ context.Context, args map[string]any) Result {
|
||||
if tool.store == nil || !sessions.ValidSessionID(tool.sessionID) {
|
||||
return errorResult("Error: Goal state is unavailable for this run.")
|
||||
}
|
||||
switch tool.action {
|
||||
case "get":
|
||||
session, err := tool.store.Get(tool.sessionID)
|
||||
if err != nil {
|
||||
return errorResult("Error: Read goal: " + err.Error())
|
||||
}
|
||||
if session == nil || session.Goal == nil {
|
||||
return okResult("No goal is set for this session.")
|
||||
}
|
||||
return goalResult(session.Goal)
|
||||
case "create":
|
||||
objective, err := stringArg(args, "objective", "", true)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for create_goal: " + err.Error())
|
||||
}
|
||||
tokenBudget, err := intArg(args, "token_budget", 0, 0, 1_000_000_000)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for create_goal: " + err.Error())
|
||||
}
|
||||
session, _, err := tool.store.CreateGoal(tool.sessionID, objective, tokenBudget)
|
||||
if err != nil {
|
||||
return errorResult("Error: Create goal: " + err.Error())
|
||||
}
|
||||
return goalResult(session.Goal)
|
||||
case "update":
|
||||
statusText, err := stringArg(args, "status", "", true)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for update_goal: " + err.Error())
|
||||
}
|
||||
status := sessions.GoalStatus(strings.ToLower(statusText))
|
||||
if status != sessions.GoalStatusComplete && status != sessions.GoalStatusBlocked {
|
||||
return errorResult(`Error: Invalid arguments for update_goal: status must be "complete" or "blocked"`)
|
||||
}
|
||||
reason, err := stringArgWithEmpty(args, "reason", "", false, true)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for update_goal: " + err.Error())
|
||||
}
|
||||
if status == sessions.GoalStatusBlocked && strings.TrimSpace(reason) == "" {
|
||||
return errorResult("Error: Invalid arguments for update_goal: blocked goals require a reason")
|
||||
}
|
||||
session, _, err := tool.store.UpdateGoal(tool.sessionID, status, reason)
|
||||
if err != nil {
|
||||
return errorResult("Error: Update goal: " + err.Error())
|
||||
}
|
||||
return goalResult(session.Goal)
|
||||
default:
|
||||
return errorResult("Error: Unknown goal action.")
|
||||
}
|
||||
}
|
||||
|
||||
func goalResult(goal *sessions.Goal) Result {
|
||||
if goal == nil {
|
||||
return okResult("No goal is set for this session.")
|
||||
}
|
||||
data, err := json.MarshalIndent(goal, "", " ")
|
||||
if err != nil {
|
||||
return errorResult(fmt.Sprintf("Error: Encode goal: %v", err))
|
||||
}
|
||||
return okResult(string(data))
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/sessions"
|
||||
)
|
||||
|
||||
func TestGoalToolsAreBoundToTheirSession(t *testing.T) {
|
||||
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
|
||||
first, err := store.Create(sessions.CreateInput{SessionID: "first"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
second, err := store.Create(sessions.CreateInput{SessionID: "second"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
goalTools := NewGoalTools(store, first.SessionID)
|
||||
|
||||
result := goalTools[1].Run(context.Background(), map[string]any{"objective": "Finish first"})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("create result = %#v", result)
|
||||
}
|
||||
loadedFirst, _ := store.Get(first.SessionID)
|
||||
loadedSecond, _ := store.Get(second.SessionID)
|
||||
if loadedFirst.Goal == nil || loadedFirst.Goal.Objective != "Finish first" {
|
||||
t.Fatalf("first goal = %#v", loadedFirst.Goal)
|
||||
}
|
||||
if loadedSecond.Goal != nil {
|
||||
t.Fatalf("second goal unexpectedly changed = %#v", loadedSecond.Goal)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateGoalToolDeclaresTokenBudgetMaximum(t *testing.T) {
|
||||
create := NewGoalTools(nil, "goal")[1]
|
||||
maximum := create.Parameters().Properties["token_budget"].Maximum
|
||||
if maximum == nil || *maximum != 1_000_000_000 {
|
||||
t.Fatalf("token_budget maximum = %v, want 1000000000", maximum)
|
||||
}
|
||||
maxLength := create.Parameters().Properties["objective"].MaxLength
|
||||
if maxLength == nil || *maxLength != sessions.GoalObjectiveMaxLength {
|
||||
t.Fatalf("objective maxLength = %v, want %d", maxLength, sessions.GoalObjectiveMaxLength)
|
||||
}
|
||||
safety := create.Safety()
|
||||
if safety.SideEffect != SideEffectWrite || safety.Permission != PermissionPrompt || !safety.AdvertiseInAuto {
|
||||
t.Fatalf("create_goal safety = %#v, want prompted persistent write advertised in auto mode", safety)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateGoalToolRestrictsAgentTransitions(t *testing.T) {
|
||||
store := sessions.NewStore(sessions.StoreOptions{RootDir: t.TempDir()})
|
||||
session, err := store.Create(sessions.CreateInput{SessionID: "goal"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := store.CreateGoal(session.SessionID, "Finish it", 0); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
update := NewGoalTools(store, session.SessionID)[2]
|
||||
|
||||
result := update.Run(context.Background(), map[string]any{"status": "paused"})
|
||||
if result.Status != StatusError || !strings.Contains(result.Output, "complete") {
|
||||
t.Fatalf("paused result = %#v", result)
|
||||
}
|
||||
result = update.Run(context.Background(), map[string]any{"status": "blocked"})
|
||||
if result.Status != StatusError || !strings.Contains(result.Output, "require a reason") {
|
||||
t.Fatalf("reasonless blocked result = %#v", result)
|
||||
}
|
||||
result = update.Run(context.Background(), map[string]any{"status": "complete"})
|
||||
if result.Status != StatusOK {
|
||||
t.Fatalf("complete result = %#v", result)
|
||||
}
|
||||
}
|
||||
@@ -45,6 +45,7 @@ func newBrowserInstallTool(options localcontrol.BrowserOptions) Tool {
|
||||
return browserInstallTool{
|
||||
baseTool: baseTool{
|
||||
name: "browser_install",
|
||||
deferred: true,
|
||||
description: "Install the local browser runtime used by browser automation.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
@@ -86,6 +87,7 @@ func newBrowserLaunchToolWithLauncher(options localcontrol.BrowserOptions, launc
|
||||
return browserLaunchTool{
|
||||
baseTool: baseTool{
|
||||
name: "browser_launch",
|
||||
deferred: true,
|
||||
description: "Launch a supported local Chromium/Electron app with Chrome DevTools enabled and attach browser automation to it. Use this for installed Electron apps such as Discord instead of launching GUI apps through shell commands.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
@@ -165,6 +167,7 @@ func newBrowserConnectTool(options localcontrol.BrowserOptions) Tool {
|
||||
return browserConnectTool{
|
||||
baseTool: baseTool{
|
||||
name: "browser_connect",
|
||||
deferred: true,
|
||||
description: "Attach local browser automation to an existing Chrome/Electron DevTools endpoint. For Electron apps, first launch the app with --remote-debugging-port=<port>, then connect to that port instead of using desktop control.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
@@ -205,6 +208,7 @@ func newBrowserOpenTool(options localcontrol.BrowserOptions) Tool {
|
||||
return browserOpenTool{
|
||||
baseTool: baseTool{
|
||||
name: "browser_open",
|
||||
deferred: true,
|
||||
description: "Open a URL in the local browser automation session.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
@@ -245,6 +249,7 @@ func newBrowserSnapshotTool(options localcontrol.BrowserOptions) Tool {
|
||||
return browserSnapshotTool{
|
||||
baseTool: baseTool{
|
||||
name: "browser_snapshot",
|
||||
deferred: true,
|
||||
description: "Return an accessibility snapshot from the current local browser automation session.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
@@ -288,6 +293,7 @@ func newBrowserClickTool(options localcontrol.BrowserOptions) Tool {
|
||||
return browserClickTool{
|
||||
baseTool: baseTool{
|
||||
name: "browser_click",
|
||||
deferred: true,
|
||||
description: "Click a ref from browser_snapshot in the current local browser automation session.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
@@ -328,6 +334,7 @@ func newBrowserTypeTool(options localcontrol.BrowserOptions) Tool {
|
||||
return browserTypeTool{
|
||||
baseTool: baseTool{
|
||||
name: "browser_type",
|
||||
deferred: true,
|
||||
description: "Type text into a ref from browser_snapshot in the current local browser automation session.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
@@ -369,6 +376,7 @@ func newBrowserPressTool(options localcontrol.BrowserOptions) Tool {
|
||||
return browserPressTool{
|
||||
baseTool: baseTool{
|
||||
name: "browser_press",
|
||||
deferred: true,
|
||||
description: "Press a keyboard key in the current local browser automation session.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
@@ -454,6 +462,7 @@ func newBrowserActionTool(options localcontrol.BrowserOptions) Tool {
|
||||
return browserActionTool{
|
||||
baseTool: baseTool{
|
||||
name: "browser_action",
|
||||
deferred: true,
|
||||
description: "Run an allowed action against the current local browser automation session.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
|
||||
@@ -40,6 +40,7 @@ func newCaptureArtifactTool(options LocalControlArtifactOptions) Tool {
|
||||
return captureArtifactTool{
|
||||
baseTool: baseTool{
|
||||
name: "capture_artifact",
|
||||
deferred: true,
|
||||
description: "Capture local browser, desktop, or terminal state into the configured artifact directory.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
|
||||
@@ -27,6 +27,7 @@ func newDesktopWindowsTool(options localcontrol.DesktopOptions) Tool {
|
||||
return desktopWindowsTool{
|
||||
baseTool: baseTool{
|
||||
name: "desktop_windows",
|
||||
deferred: true,
|
||||
description: "List native desktop windows through the local desktop automation helper.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
@@ -70,6 +71,7 @@ func newDesktopSnapshotTool(options localcontrol.DesktopOptions) Tool {
|
||||
return desktopSnapshotTool{
|
||||
baseTool: baseTool{
|
||||
name: "desktop_snapshot",
|
||||
deferred: true,
|
||||
description: "Read a native desktop window accessibility snapshot through the local desktop automation helper.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
@@ -117,6 +119,7 @@ func newDesktopActionTool(options localcontrol.DesktopOptions) Tool {
|
||||
return desktopActionTool{
|
||||
baseTool: baseTool{
|
||||
name: "desktop_action",
|
||||
deferred: true,
|
||||
description: "Run an allowed native desktop action through the local desktop automation helper.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
|
||||
@@ -23,6 +23,7 @@ func newTerminalSessionTool(options localcontrol.TerminalOptions) Tool {
|
||||
return terminalSessionTool{
|
||||
baseTool: baseTool{
|
||||
name: "terminal_session",
|
||||
deferred: true,
|
||||
description: "Launch or control a local virtual terminal session through the local terminal automation helper.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
|
||||
@@ -32,7 +32,8 @@ func NewScopedLSPNavigateTool(workspaceRoot string, scope PathScope) Tool {
|
||||
root := normalizeWorkspaceRoot(workspaceRoot)
|
||||
return lspNavigateTool{
|
||||
baseTool: baseTool{
|
||||
name: "lsp_navigate",
|
||||
name: "lsp_navigate",
|
||||
deferred: true,
|
||||
description: "Navigate code semantically via the language server: jump to a symbol's " +
|
||||
"definition, find all references, find implementations of an interface/method, or " +
|
||||
"search workspace symbols by name. Use this instead of grep when you need precise " +
|
||||
|
||||
@@ -57,6 +57,9 @@ func applySelfManagedOutputBudget(tool Tool, toolName string, args map[string]an
|
||||
|
||||
category := resolveOutputCategory(tool, toolName, args)
|
||||
budgeted := budgetSemanticOutput(result.Output, category, budget)
|
||||
if result.Truncated && budgeted.spillPath == "" {
|
||||
budgeted.spillPath = result.Meta["spill_path"]
|
||||
}
|
||||
if result.Truncated && !budgeted.truncated {
|
||||
budgeted.truncated = true
|
||||
budgeted.reason = result.Meta["truncation_reason"]
|
||||
@@ -64,7 +67,7 @@ func applySelfManagedOutputBudget(tool Tool, toolName string, args map[string]an
|
||||
budgeted.reason = "upstream_tool_budget"
|
||||
}
|
||||
}
|
||||
if budgeted.truncated {
|
||||
if budgeted.truncated && budgeted.spillPath == "" {
|
||||
budgeted = attachExistingSpill(toolName, result.Output, budget, budgeted)
|
||||
}
|
||||
result.Output = budgeted.text
|
||||
|
||||
+121
-9
@@ -10,8 +10,11 @@ import (
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const readFileByteChunkMax = 64 * 1024
|
||||
|
||||
type readFileTool struct {
|
||||
baseTool
|
||||
workspaceRoot string
|
||||
@@ -28,14 +31,16 @@ func NewScopedReadFileTool(workspaceRoot string, scope PathScope) Tool {
|
||||
return readFileTool{
|
||||
baseTool: baseTool{
|
||||
name: "read_file",
|
||||
description: "Read a file with optional 1-based inclusive line range and max line cap.",
|
||||
description: "Read a file by 1-based inclusive line range, or by exact byte chunks for oversized single-line files.",
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]PropertySchema{
|
||||
"path": {Type: "string", Description: "Path of the file to read."},
|
||||
"start_line": {Type: "integer", Description: "1-based inclusive line number to start reading from.", Minimum: intPtr(1)},
|
||||
"end_line": {Type: "integer", Description: "1-based inclusive line number to stop reading at.", Minimum: intPtr(1)},
|
||||
"max_lines": {Type: "integer", Description: "Maximum number of lines to return.", Minimum: intPtr(1)},
|
||||
"path": {Type: "string", Description: "Path of the file to read."},
|
||||
"start_line": {Type: "integer", Description: "1-based inclusive line number to start reading from.", Minimum: intPtr(1)},
|
||||
"end_line": {Type: "integer", Description: "1-based inclusive line number to stop reading at.", Minimum: intPtr(1)},
|
||||
"max_lines": {Type: "integer", Description: "Maximum number of lines to return.", Minimum: intPtr(1)},
|
||||
"byte_offset": {Type: "integer", Description: "Zero-based byte offset for an exact chunk read. Use with byte_limit for oversized single-line files; cannot be combined with line ranges.", Minimum: intPtr(0)},
|
||||
"byte_limit": {Type: "integer", Description: "Maximum source bytes returned in exact byte mode. Defaults to 65536.", Default: readFileByteChunkMax, Minimum: intPtr(1), Maximum: intPtr(readFileByteChunkMax)},
|
||||
},
|
||||
Required: []string{"path"},
|
||||
AdditionalProperties: false,
|
||||
@@ -76,6 +81,20 @@ func (tool readFileTool) run(args map[string]any, options RunOptions, directBudg
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for read_file: " + err.Error())
|
||||
}
|
||||
_, hasByteOffset := args["byte_offset"]
|
||||
_, hasByteLimit := args["byte_limit"]
|
||||
byteMode := hasByteOffset || hasByteLimit
|
||||
byteOffset, err := intArg(args, "byte_offset", 0, 0, 0)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for read_file: " + err.Error())
|
||||
}
|
||||
byteLimit, err := intArg(args, "byte_limit", readFileByteChunkMax, 1, readFileByteChunkMax)
|
||||
if err != nil {
|
||||
return errorResult("Error: Invalid arguments for read_file: " + err.Error())
|
||||
}
|
||||
if byteMode && (args["start_line"] != nil || args["end_line"] != nil || args["max_lines"] != nil) {
|
||||
return errorResult("Error: Invalid arguments for read_file: byte_offset/byte_limit cannot be combined with line ranges")
|
||||
}
|
||||
|
||||
absolutePath, relativePath, err := resolveScopedReadPath(tool.workspaceRoot, tool.scope, requestedPath)
|
||||
if err != nil {
|
||||
@@ -91,12 +110,61 @@ func (tool readFileTool) run(args map[string]any, options RunOptions, directBudg
|
||||
// Stat is best-effort: a missing FileInfo only drops the diagnostic size/mtime,
|
||||
// not the authoritative content hash.
|
||||
options.FileTracker.RecordHash(absolutePath, stats.hash, stats.info)
|
||||
if byteMode {
|
||||
result, seenStart, seenEnd := renderReadFileBytes(absolutePath, relativePath, stats.bytes, byteOffset, byteLimit)
|
||||
if result.Status == StatusOK && !result.Truncated && seenEnd > seenStart {
|
||||
if options.deferFileObservation {
|
||||
result.pendingFileObservation = &pendingFileObservation{
|
||||
path: absolutePath, output: result.Output, hash: stats.hash,
|
||||
start: seenStart, end: seenEnd, total: stats.bytes, byteMode: true,
|
||||
}
|
||||
} else {
|
||||
options.FileTracker.RecordSeenBytes(absolutePath, seenStart, seenEnd, stats.bytes)
|
||||
result.Meta["file_version"] = stats.hash
|
||||
result.Meta["seen_bytes"] = fmt.Sprintf("%d-%d", seenStart, seenEnd)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
maxBytes := 0
|
||||
if directBudget {
|
||||
maxBytes = readOutputBudgetBytes
|
||||
}
|
||||
return renderReadFileRange(absolutePath, relativePath, stats.lines, startLine, endLine, maxLines, maxBytes)
|
||||
result := renderReadFileRange(absolutePath, relativePath, stats.lines, startLine, endLine, maxLines, maxBytes)
|
||||
if result.Status == StatusOK && result.Meta["truncation_reason"] != "byte_budget" {
|
||||
seenStart, seenEnd := renderedReadRange(stats.lines, startLine, endLine, maxLines)
|
||||
if options.deferFileObservation {
|
||||
result.pendingFileObservation = &pendingFileObservation{
|
||||
path: absolutePath, output: result.Output, hash: stats.hash,
|
||||
start: seenStart, end: seenEnd, total: stats.lines,
|
||||
}
|
||||
} else {
|
||||
options.FileTracker.RecordSeenRange(absolutePath, seenStart, seenEnd, stats.lines)
|
||||
if result.Meta == nil {
|
||||
result.Meta = map[string]string{}
|
||||
}
|
||||
result.Meta["file_version"] = stats.hash
|
||||
result.Meta["seen_lines"] = fmt.Sprintf("%d-%d", seenStart, seenEnd)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func renderedReadRange(total, start, end, maxLines int) (int, int) {
|
||||
if start > total {
|
||||
return start, start
|
||||
}
|
||||
if end == 0 || end > total {
|
||||
end = total
|
||||
}
|
||||
if end < start {
|
||||
end = start
|
||||
}
|
||||
if maxLines > 0 && end-start+1 > maxLines {
|
||||
end = start + maxLines - 1
|
||||
}
|
||||
return start, end
|
||||
}
|
||||
|
||||
func renderReadFileRange(absolutePath string, relativePath string, total int, startLine int, endLine int, maxLines int, maxBytes int) Result {
|
||||
@@ -140,9 +208,9 @@ func renderReadFileRange(absolutePath string, relativePath string, total int, st
|
||||
// Tool.Run calls preserve the legacy prefix-only byte budget.
|
||||
var budgetedOutput *outputBudgetBuilder
|
||||
if maxBytes > 0 {
|
||||
budgetedOutput = newOutputBudgetBuilder(maxBytes, "use start_line/end_line or max_lines to continue with a smaller range")
|
||||
budgetedOutput = newOutputBudgetBuilder(maxBytes, "use start_line/end_line or max_lines for normal files; use byte_offset/byte_limit for an oversized single line")
|
||||
} else {
|
||||
budgetedOutput = newHeadTailOutputBudgetBuilder(readOutputBudgetBytes, "use start_line/end_line or max_lines to continue with a smaller range")
|
||||
budgetedOutput = newHeadTailOutputBudgetBuilder(readOutputBudgetBytes, "use start_line/end_line or max_lines for normal files; use byte_offset/byte_limit for an oversized single line")
|
||||
}
|
||||
budgetedOutput.WriteString(header)
|
||||
budgetedOutput.WriteString("\n")
|
||||
@@ -186,6 +254,7 @@ func renderReadFileRange(absolutePath string, relativePath string, total int, st
|
||||
|
||||
type readFileStats struct {
|
||||
lines int
|
||||
bytes int
|
||||
hash string
|
||||
info os.FileInfo
|
||||
}
|
||||
@@ -200,6 +269,7 @@ func scanReadFileStats(path string) (readFileStats, error) {
|
||||
hasher := sha256.New()
|
||||
reader := bufio.NewReader(file)
|
||||
lines := 0
|
||||
bytes := 0
|
||||
for {
|
||||
raw, _, err := readRawLine(reader)
|
||||
if err == io.EOF {
|
||||
@@ -211,13 +281,55 @@ func scanReadFileStats(path string) (readFileStats, error) {
|
||||
if _, err := hasher.Write(raw); err != nil {
|
||||
return readFileStats{}, err
|
||||
}
|
||||
bytes += len(raw)
|
||||
lines++
|
||||
}
|
||||
if lines == 0 {
|
||||
lines = 1
|
||||
}
|
||||
info, _ := file.Stat()
|
||||
return readFileStats{lines: lines, hash: hex.EncodeToString(hasher.Sum(nil)), info: info}, nil
|
||||
return readFileStats{lines: lines, bytes: bytes, hash: hex.EncodeToString(hasher.Sum(nil)), info: info}, nil
|
||||
}
|
||||
|
||||
func renderReadFileBytes(path, relativePath string, total, requestedStart, limit int) (Result, int, int) {
|
||||
if requestedStart >= total {
|
||||
return okResult(fmt.Sprintf("File: %s\n(byte_offset %d is past the end of the file, which has %d bytes)", relativePath, requestedStart, total)), 0, 0
|
||||
}
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return errorResult("Error reading file " + relativePath + ": " + err.Error()), 0, 0
|
||||
}
|
||||
defer file.Close()
|
||||
|
||||
want := min(limit, total-requestedStart)
|
||||
data := make([]byte, want)
|
||||
read, err := file.ReadAt(data, int64(requestedStart))
|
||||
if err != nil && err != io.EOF {
|
||||
return errorResult("Error reading file " + relativePath + ": " + err.Error()), 0, 0
|
||||
}
|
||||
data = data[:read]
|
||||
start := requestedStart
|
||||
for len(data) > 0 && data[0]&0xc0 == 0x80 {
|
||||
data = data[1:]
|
||||
start++
|
||||
}
|
||||
for trim := 0; !utf8.Valid(data) && trim < utf8.UTFMax && len(data) > 0; trim++ {
|
||||
data = data[:len(data)-1]
|
||||
}
|
||||
if !utf8.Valid(data) {
|
||||
return errorResult("Error reading file " + relativePath + ": exact byte mode requires UTF-8 text"), 0, 0
|
||||
}
|
||||
end := start + len(data)
|
||||
if end <= start {
|
||||
return errorResult("Error reading file " + relativePath + ": byte_limit is too small to return the next UTF-8 character"), 0, 0
|
||||
}
|
||||
|
||||
header := fmt.Sprintf("File: %s (bytes %d-%d of %d)", relativePath, start, end-1, total)
|
||||
output := header + "\n\n" + string(data)
|
||||
if end < total {
|
||||
output += fmt.Sprintf("\n\n[continue with byte_offset=%d]", end)
|
||||
}
|
||||
return Result{Status: StatusOK, Output: output, Meta: map[string]string{"next_byte_offset": strconv.Itoa(end)}}, start, end
|
||||
}
|
||||
|
||||
func appendReadFileRange(output *outputBudgetBuilder, path string, startLine int, selectedLines int, width int) error {
|
||||
|
||||
@@ -29,6 +29,7 @@ func NewScopedReadMinifiedFileTool(workspaceRoot string, scope PathScope) Tool {
|
||||
baseTool: baseTool{
|
||||
name: "read_minified_file",
|
||||
description: "Read a source file in a dense, token-cheap form: comments and redundant whitespace removed, no line numbers. Use it to scan or understand code for far fewer tokens than read_file. For exact text, comments, line numbers, or before editing, use read_file instead.",
|
||||
deferred: true,
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]PropertySchema{
|
||||
@@ -112,7 +113,7 @@ func (tool readMinifiedFileTool) run(args map[string]any, options RunOptions, di
|
||||
meta["estimated_tokens_saved"] = strconv.Itoa(savedTokens)
|
||||
toolResult := Result{Status: StatusOK, Output: output, Meta: meta}
|
||||
if directBudget {
|
||||
return applyLegacyByteBudgetToResult(toolResult, readOutputBudgetBytes, "use read_file with start_line/end_line or max_lines for a smaller exact range")
|
||||
return applyLegacyByteBudgetToResult(toolResult, readOutputBudgetBytes, "use read_file with start_line/end_line or max_lines for normal files, or byte_offset/byte_limit for an oversized single line")
|
||||
}
|
||||
return toolResult
|
||||
}
|
||||
|
||||
@@ -2,7 +2,9 @@ package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"github.com/Gitlawb/zero/internal/redaction"
|
||||
"github.com/Gitlawb/zero/internal/sandbox"
|
||||
@@ -29,6 +31,10 @@ type RunOptions struct {
|
||||
// disk outside Zero since it was last read. nil disables the feature entirely
|
||||
// (the read/write tools behave exactly as before).
|
||||
FileTracker *FileTracker
|
||||
// DeferFileObservationCommit lets a caller with an additional output layer
|
||||
// commit exact-read credit after that final layer has run.
|
||||
DeferFileObservationCommit bool
|
||||
deferFileObservation bool
|
||||
// EnabledTools / DisabledTools carry the run's operator tool filters so a
|
||||
// filter-aware tool (tool_search) never discloses or loads an operator-hidden
|
||||
// tool. They use the same allow/deny semantics as the agent's filter gate:
|
||||
@@ -130,17 +136,23 @@ func (registry *Registry) RunWithOptions(ctx context.Context, name string, args
|
||||
// args/paths) are redacted at the boundary just like tool output. The output
|
||||
// ceiling runs after the scrub so the transcript and the spill file agree on
|
||||
// what was hidden.
|
||||
commitFileObservation := !options.DeferFileObservationCommit
|
||||
options.deferFileObservation = true
|
||||
selfManagedOutput := false
|
||||
var tool Tool
|
||||
var ok bool
|
||||
defer func() {
|
||||
result = scrubResultSecrets(result)
|
||||
result = reduceCommandOutput(name, args, result)
|
||||
if selfManagedOutput {
|
||||
result = applySelfManagedOutputBudget(tool, name, args, result)
|
||||
} else {
|
||||
result = applyRegistryOutputBudget(tool, name, args, result)
|
||||
result = enforceOutputCeiling(name, result)
|
||||
}
|
||||
if commitFileObservation {
|
||||
result = registry.CommitFileObservation(result, options.FileTracker)
|
||||
}
|
||||
}()
|
||||
|
||||
tool, ok = registry.Get(name)
|
||||
@@ -220,6 +232,32 @@ func (registry *Registry) RunWithOptions(ctx context.Context, name string, args
|
||||
return res
|
||||
}
|
||||
|
||||
// CommitFileObservation grants exact-read credit only when the tool's complete
|
||||
// output remains intact at the final boundary that will be sent to the model.
|
||||
func (registry *Registry) CommitFileObservation(result Result, tracker *FileTracker) Result {
|
||||
observation := result.pendingFileObservation
|
||||
result.pendingFileObservation = nil
|
||||
if observation == nil || tracker == nil || result.Status != StatusOK || result.Truncated ||
|
||||
!strings.HasPrefix(result.Output, observation.output) {
|
||||
return result
|
||||
}
|
||||
if observation.byteMode {
|
||||
tracker.RecordSeenBytes(observation.path, observation.start, observation.end, observation.total)
|
||||
} else {
|
||||
tracker.RecordSeenRange(observation.path, observation.start, observation.end, observation.total)
|
||||
}
|
||||
if result.Meta == nil {
|
||||
result.Meta = map[string]string{}
|
||||
}
|
||||
result.Meta["file_version"] = observation.hash
|
||||
if observation.byteMode {
|
||||
result.Meta["seen_bytes"] = fmt.Sprintf("%d-%d", observation.start, observation.end)
|
||||
} else {
|
||||
result.Meta["seen_lines"] = fmt.Sprintf("%d-%d", observation.start, observation.end)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func effectiveToolPermission(tool Tool, args map[string]any) Permission {
|
||||
if permissioner, ok := tool.(ArgsPermissioner); ok {
|
||||
return permissioner.PermissionForArgs(args)
|
||||
|
||||
@@ -515,3 +515,44 @@ func TestIsDeferralEligibleDecouplesFromDeferred(t *testing.T) {
|
||||
t.Fatal("IsDeferralEligible(optedOut) = true, want false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOptionalBuiltinsUseDeferredDiscovery(t *testing.T) {
|
||||
wantDeferred := map[string]bool{
|
||||
"bash": true,
|
||||
"browser_action": true,
|
||||
"browser_click": true,
|
||||
"browser_connect": true,
|
||||
"browser_install": true,
|
||||
"browser_launch": true,
|
||||
"browser_open": true,
|
||||
"browser_press": true,
|
||||
"browser_snapshot": true,
|
||||
"browser_type": true,
|
||||
"capture_artifact": true,
|
||||
"desktop_action": true,
|
||||
"desktop_snapshot": true,
|
||||
"desktop_windows": true,
|
||||
"lsp_navigate": true,
|
||||
"read_minified_file": true,
|
||||
"terminal_session": true,
|
||||
"web_fetch": true,
|
||||
"web_search": true,
|
||||
}
|
||||
for _, tool := range BuiltinCatalog(t.TempDir()) {
|
||||
if _, listed := wantDeferred[tool.Name()]; listed {
|
||||
if !IsDeferred(tool) {
|
||||
t.Errorf("%s should be deferred", tool.Name())
|
||||
}
|
||||
delete(wantDeferred, tool.Name())
|
||||
} else if tool.Name() == "read_file" || tool.Name() == "grep" || tool.Name() == ExecCommandToolName || tool.Name() == "edit_file" {
|
||||
if IsDeferred(tool) {
|
||||
t.Errorf("essential tool %s must stay eager", tool.Name())
|
||||
}
|
||||
} else if IsDeferred(tool) {
|
||||
t.Errorf("unexpected deferred builtin %s", tool.Name())
|
||||
}
|
||||
}
|
||||
for missing := range wantDeferred {
|
||||
t.Errorf("deferred builtin %s missing from catalog", missing)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,290 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const commandReducerMinPassingLines = 12
|
||||
|
||||
type commandOutputReducer struct {
|
||||
kind string
|
||||
label string
|
||||
reduce func(string) (string, int)
|
||||
}
|
||||
|
||||
// reduceCommandOutput removes repetitive success-only lines from recognized
|
||||
// simple commands. Unsupported or compound shell input passes through exactly,
|
||||
// and every changed result retains the pre-reduction output as an artifact.
|
||||
func reduceCommandOutput(toolName string, args map[string]any, result Result) Result {
|
||||
if toolName != ExecCommandToolName && toolName != "bash" {
|
||||
return result
|
||||
}
|
||||
command, err := commandTextForReducer(toolName, args)
|
||||
if err != nil {
|
||||
return result
|
||||
}
|
||||
reducer, ok := commandOutputReducerFor(command)
|
||||
if !ok {
|
||||
return result
|
||||
}
|
||||
|
||||
reduced, omitted := reducer.reduce(result.Output)
|
||||
if omitted < commandReducerMinPassingLines || len(reduced) >= len(result.Output) {
|
||||
return result
|
||||
}
|
||||
originalBytes := len(result.Output)
|
||||
originalTokens := estimateOutputTokens(result.Output)
|
||||
spillPath := spillTruncatedOutput(toolName, result.Output)
|
||||
if spillPath == "" {
|
||||
return result
|
||||
}
|
||||
|
||||
notice := fmt.Sprintf("[zero] compacted %d %s; exact output saved to %s", omitted, reducer.label, spillPath)
|
||||
result.Output = strings.TrimRight(reduced, "\n") + "\n" + notice
|
||||
result.Truncated = true
|
||||
if result.Meta == nil {
|
||||
result.Meta = map[string]string{}
|
||||
}
|
||||
result.Meta["command_output_reduced"] = "true"
|
||||
result.Meta["command_output_reducer"] = reducer.kind
|
||||
result.Meta["command_output_omitted_lines"] = strconv.Itoa(omitted)
|
||||
result.Meta["command_output_original_bytes"] = strconv.Itoa(originalBytes)
|
||||
result.Meta["command_output_original_tokens"] = strconv.Itoa(originalTokens)
|
||||
result.Meta["command_output_retained_tokens"] = strconv.Itoa(estimateOutputTokens(result.Output))
|
||||
result.Meta["spill_path"] = spillPath
|
||||
result.Meta["truncation_reason"] = "command_aware_reduction"
|
||||
return result
|
||||
}
|
||||
|
||||
func commandTextForReducer(toolName string, args map[string]any) (string, error) {
|
||||
if toolName == "bash" {
|
||||
return bashCommandArg(args)
|
||||
}
|
||||
return execCommandArg(args)
|
||||
}
|
||||
|
||||
func simpleCommandFields(command string) ([]string, bool) {
|
||||
command = strings.TrimSpace(command)
|
||||
if separator := strings.LastIndexAny(command, " \t"); separator >= 0 && command[separator+1:] == "2>&1" {
|
||||
command = strings.TrimSpace(command[:separator])
|
||||
}
|
||||
if command == "" || strings.ContainsAny(command, "|;&<>\n\r`") || strings.Contains(command, "$(") {
|
||||
return nil, false
|
||||
}
|
||||
fields := strings.Fields(command)
|
||||
if len(fields) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
return fields, true
|
||||
}
|
||||
|
||||
func commandOutputReducerFor(command string) (commandOutputReducer, bool) {
|
||||
fields, ok := simpleCommandFields(command)
|
||||
if !ok {
|
||||
return commandOutputReducer{}, false
|
||||
}
|
||||
if len(fields) >= 2 && fields[0] == "go" && fields[1] == "test" {
|
||||
return commandOutputReducer{kind: "go_test", label: "passing package lines", reduce: reduceGoTestPassingPackages}, true
|
||||
}
|
||||
if len(fields) >= 2 && fields[0] == "cargo" {
|
||||
switch fields[1] {
|
||||
case "test":
|
||||
return commandOutputReducer{kind: "cargo_test", label: "passing Rust test lines", reduce: reduceCargoTestPassingLines}, true
|
||||
case "build", "check":
|
||||
return commandOutputReducer{kind: "cargo_build", label: "successful Rust build progress lines", reduce: reduceCargoBuildProgressLines}, true
|
||||
}
|
||||
}
|
||||
if isPytestCommand(fields) {
|
||||
return commandOutputReducer{kind: "pytest", label: "passing pytest progress lines", reduce: reducePytestPassingLines}, true
|
||||
}
|
||||
if isJavaScriptTestCommand(fields) {
|
||||
return commandOutputReducer{kind: "javascript_test", label: "passing JavaScript test-file lines", reduce: reduceJavaScriptTestPassingLines}, true
|
||||
}
|
||||
return commandOutputReducer{}, false
|
||||
}
|
||||
|
||||
func isPytestCommand(fields []string) bool {
|
||||
if fields[0] == "pytest" || fields[0] == "py.test" {
|
||||
return true
|
||||
}
|
||||
return len(fields) >= 3 && (fields[0] == "python" || fields[0] == "python3") && fields[1] == "-m" && fields[2] == "pytest"
|
||||
}
|
||||
|
||||
func isJavaScriptTestCommand(fields []string) bool {
|
||||
switch fields[0] {
|
||||
case "jest", "vitest":
|
||||
return true
|
||||
case "npx":
|
||||
return len(fields) >= 2 && (fields[1] == "jest" || fields[1] == "vitest")
|
||||
case "npm", "pnpm", "yarn":
|
||||
if len(fields) >= 2 && fields[1] == "test" {
|
||||
return true
|
||||
}
|
||||
if len(fields) >= 3 && fields[1] == "run" && fields[2] == "test" {
|
||||
return true
|
||||
}
|
||||
return len(fields) >= 3 && fields[1] == "exec" && (fields[2] == "jest" || fields[2] == "vitest")
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func reduceGoTestPassingPackages(output string) (string, int) {
|
||||
return reduceRepeatedOutputLines(output, "passing package lines", func(line string) bool {
|
||||
packageSuccess := strings.HasPrefix(line, "ok \t") || strings.HasPrefix(line, "? \t") ||
|
||||
strings.HasPrefix(line, "ok ") || strings.HasPrefix(line, "? ")
|
||||
return packageSuccess && !strings.Contains(line, "coverage:") && !strings.Contains(line, "[no test files]")
|
||||
})
|
||||
}
|
||||
|
||||
func reduceCargoTestPassingLines(output string) (string, int) {
|
||||
if !strings.Contains(output, "test result:") {
|
||||
return output, 0
|
||||
}
|
||||
return reduceRepeatedOutputLines(output, "passing Rust test lines", func(line string) bool {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
return strings.HasPrefix(trimmed, "test ") && strings.HasSuffix(trimmed, " ... ok")
|
||||
})
|
||||
}
|
||||
|
||||
func reduceCargoBuildProgressLines(output string) (string, int) {
|
||||
if !outputHasCargoFinishedSummary(output) {
|
||||
return output, 0
|
||||
}
|
||||
return reduceRepeatedOutputLines(output, "successful Rust build progress lines", func(line string) bool {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
return strings.HasPrefix(trimmed, "Checking ") || strings.HasPrefix(trimmed, "Compiling ")
|
||||
})
|
||||
}
|
||||
|
||||
func reducePytestPassingLines(output string) (string, int) {
|
||||
if !outputHasPytestSummary(output) {
|
||||
return output, 0
|
||||
}
|
||||
return reduceRepeatedOutputLines(output, "passing pytest progress lines", isPassingPytestProgressLine)
|
||||
}
|
||||
|
||||
func isPassingPytestProgressLine(line string) bool {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
open := strings.LastIndexByte(trimmed, '[')
|
||||
if open < 0 || !strings.HasSuffix(trimmed, "%]") {
|
||||
return false
|
||||
}
|
||||
progress := strings.TrimSpace(trimmed[:open])
|
||||
fields := strings.Fields(progress)
|
||||
if len(fields) == 0 {
|
||||
return false
|
||||
}
|
||||
marks := fields[len(fields)-1]
|
||||
return marks != "" && strings.Trim(marks, ".") == ""
|
||||
}
|
||||
|
||||
func reduceJavaScriptTestPassingLines(output string) (string, int) {
|
||||
if !outputHasJavaScriptTestSummary(output) {
|
||||
return output, 0
|
||||
}
|
||||
return reduceRepeatedOutputLines(output, "passing JavaScript test-file lines", func(line string) bool {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
return strings.HasPrefix(line, "PASS ") || strings.HasPrefix(trimmed, "✓ ") || strings.HasPrefix(trimmed, "✔ ")
|
||||
})
|
||||
}
|
||||
|
||||
func outputHasCargoFinishedSummary(output string) bool {
|
||||
for _, line := range strings.Split(output, "\n") {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if strings.HasPrefix(trimmed, "Finished `") && strings.Contains(trimmed, "` profile ") &&
|
||||
strings.Contains(trimmed, " target(s) in ") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func outputHasPytestSummary(output string) bool {
|
||||
for _, line := range strings.Split(output, "\n") {
|
||||
summary := strings.Trim(strings.TrimSpace(line), "= ")
|
||||
if summary == "short test summary info" || isPytestResultSummary(summary) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func isPytestResultSummary(summary string) bool {
|
||||
fields := strings.Fields(summary)
|
||||
if len(fields) < 4 {
|
||||
return false
|
||||
}
|
||||
first := strings.Trim(fields[0], ",;:()")
|
||||
if _, err := strconv.Atoi(first); err != nil || !summaryHasCountedStatus(summary, "passed") {
|
||||
return false
|
||||
}
|
||||
for index := 1; index+1 < len(fields); index++ {
|
||||
if fields[index] != "in" {
|
||||
continue
|
||||
}
|
||||
duration := strings.TrimSuffix(strings.Trim(fields[index+1], ",;:()"), "s")
|
||||
if _, err := strconv.ParseFloat(duration, 64); err == nil {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func outputHasJavaScriptTestSummary(output string) bool {
|
||||
for _, line := range strings.Split(output, "\n") {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if strings.HasPrefix(trimmed, "Test Suites:") && summaryHasCountedStatus(trimmed, "passed") {
|
||||
return true
|
||||
}
|
||||
if strings.HasPrefix(trimmed, "Test Files ") && summaryHasCountedStatus(trimmed, "passed") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func summaryHasCountedStatus(summary string, status string) bool {
|
||||
fields := strings.Fields(summary)
|
||||
for index := 1; index < len(fields); index++ {
|
||||
if strings.Trim(fields[index], ",;:()") != status {
|
||||
continue
|
||||
}
|
||||
count := strings.Trim(fields[index-1], ",;:()")
|
||||
if _, err := strconv.Atoi(count); err == nil {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func reduceRepeatedOutputLines(output string, label string, omit func(string) bool) (string, int) {
|
||||
lines := strings.Split(output, "\n")
|
||||
kept := make([]string, 0, len(lines))
|
||||
omitted := 0
|
||||
for _, line := range lines {
|
||||
if omit(line) {
|
||||
omitted++
|
||||
continue
|
||||
}
|
||||
kept = append(kept, line)
|
||||
}
|
||||
if omitted == 0 {
|
||||
return output, 0
|
||||
}
|
||||
|
||||
insertAt := len(kept)
|
||||
for index, line := range kept {
|
||||
if strings.HasPrefix(line, "exit_code:") || strings.HasPrefix(line, "session_id:") {
|
||||
insertAt = index
|
||||
break
|
||||
}
|
||||
}
|
||||
summary := fmt.Sprintf("[zero] %d %s omitted", omitted, label)
|
||||
kept = append(kept, "")
|
||||
copy(kept[insertAt+1:], kept[insertAt:])
|
||||
kept[insertAt] = summary
|
||||
return strings.Join(kept, "\n"), omitted
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type commandReducerCorpusCase struct {
|
||||
name string
|
||||
command string
|
||||
output string
|
||||
want string
|
||||
}
|
||||
|
||||
func commandReducerCorpus() []commandReducerCorpusCase {
|
||||
goLines := []string{"output:"}
|
||||
cargoTestLines := []string{"output:", "running 24 tests"}
|
||||
cargoCheckLines := []string{"output:"}
|
||||
pytestLines := []string{"output:"}
|
||||
vitestLines := []string{"output:", " RUN v3.2.4 /workspace"}
|
||||
jestLines := []string{"output:"}
|
||||
for index := 0; index < 24; index++ {
|
||||
goLines = append(goLines, fmt.Sprintf("ok \texample.test/pkg%d\t0.01s", index))
|
||||
cargoTestLines = append(cargoTestLines, fmt.Sprintf("test module::case_%02d ... ok", index))
|
||||
cargoCheckLines = append(cargoCheckLines, fmt.Sprintf(" Checking dependency_%02d v1.0.%d", index, index))
|
||||
pytestLines = append(pytestLines, fmt.Sprintf("tests/test_module_%02d.py ........ [%3d%%]", index, (index+1)*100/24))
|
||||
vitestLines = append(vitestLines, fmt.Sprintf(" ✓ src/module_%02d.test.ts (8 tests) 12ms", index))
|
||||
jestLines = append(jestLines, fmt.Sprintf("PASS src/module_%02d.test.ts", index))
|
||||
}
|
||||
goLines = append(goLines, "exit_code: 0")
|
||||
cargoTestLines = append(cargoTestLines,
|
||||
"test result: ok. 24 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.08s",
|
||||
"exit_code: 0",
|
||||
)
|
||||
cargoCheckLines = append(cargoCheckLines,
|
||||
" Finished `dev` profile [unoptimized + debuginfo] target(s) in 2.41s",
|
||||
"exit_code: 0",
|
||||
)
|
||||
pytestLines = append(pytestLines,
|
||||
"============================= 192 passed in 3.21s =============================",
|
||||
"exit_code: 0",
|
||||
)
|
||||
vitestLines = append(vitestLines,
|
||||
" Test Files 24 passed (24)",
|
||||
" Tests 192 passed (192)",
|
||||
" Duration 3.21s",
|
||||
"exit_code: 0",
|
||||
)
|
||||
jestLines = append(jestLines,
|
||||
"Test Suites: 24 passed, 24 total",
|
||||
"Tests: 192 passed, 192 total",
|
||||
"Time: 3.21 s",
|
||||
"exit_code: 0",
|
||||
)
|
||||
return []commandReducerCorpusCase{
|
||||
{name: "go_test", command: "go test ./...", output: strings.Join(goLines, "\n"), want: "exit_code: 0"},
|
||||
{name: "cargo_test", command: "cargo test --workspace", output: strings.Join(cargoTestLines, "\n"), want: "24 passed"},
|
||||
{name: "cargo_check", command: "cargo check --workspace", output: strings.Join(cargoCheckLines, "\n"), want: "Finished `dev` profile"},
|
||||
{name: "pytest", command: "pytest -q", output: strings.Join(pytestLines, "\n"), want: "192 passed"},
|
||||
{name: "npm_vitest", command: "npm test", output: strings.Join(vitestLines, "\n"), want: "192 passed"},
|
||||
{name: "pnpm_vitest", command: "pnpm test", output: strings.Join(vitestLines, "\n"), want: "192 passed"},
|
||||
{name: "yarn_jest", command: "yarn test", output: strings.Join(jestLines, "\n"), want: "192 passed"},
|
||||
}
|
||||
}
|
||||
|
||||
func TestCommandOutputReductionCorpusMetrics(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
totalRawTokens := 0
|
||||
totalModelTokens := 0
|
||||
for _, testCase := range commandReducerCorpus() {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
result := reduceCommandOutput(ExecCommandToolName, map[string]any{"cmd": testCase.command}, Result{
|
||||
Status: StatusOK,
|
||||
Output: testCase.output,
|
||||
})
|
||||
if !strings.Contains(result.Output, testCase.want) {
|
||||
t.Fatalf("decisive summary %q was lost: %q", testCase.want, result.Output)
|
||||
}
|
||||
if spillPath := result.Meta["spill_path"]; spillPath != "" {
|
||||
spilled, err := os.ReadFile(spillPath)
|
||||
if err != nil {
|
||||
t.Fatalf("read raw artifact: %v", err)
|
||||
}
|
||||
if string(spilled) != testCase.output {
|
||||
t.Fatal("raw artifact differs from original output")
|
||||
}
|
||||
}
|
||||
rawTokens := estimateOutputTokens(testCase.output)
|
||||
modelTokens := estimateOutputTokens(result.Output)
|
||||
totalRawTokens += rawTokens
|
||||
totalModelTokens += modelTokens
|
||||
reduction := 100 * (rawTokens - modelTokens) / rawTokens
|
||||
t.Logf("raw_bytes=%d model_bytes=%d raw_tokens=%d model_tokens=%d reduction_pct=%d reduced=%t",
|
||||
len(testCase.output), len(result.Output), rawTokens, modelTokens, reduction, result.Meta["command_output_reduced"] == "true")
|
||||
})
|
||||
}
|
||||
if totalRawTokens == 0 {
|
||||
t.Fatal("empty reducer corpus")
|
||||
}
|
||||
if totalModelTokens >= totalRawTokens {
|
||||
t.Fatalf("aggregate reducer output did not save tokens: raw=%d model=%d", totalRawTokens, totalModelTokens)
|
||||
}
|
||||
t.Logf("aggregate raw_tokens=%d model_tokens=%d reduction_pct=%d",
|
||||
totalRawTokens, totalModelTokens, 100*(totalRawTokens-totalModelTokens)/totalRawTokens)
|
||||
}
|
||||
@@ -0,0 +1,317 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type commandReducerFixtureTool struct {
|
||||
name string
|
||||
output string
|
||||
}
|
||||
|
||||
func (tool commandReducerFixtureTool) Name() string { return tool.name }
|
||||
func (tool commandReducerFixtureTool) Description() string {
|
||||
return "returns command output for reducer integration tests"
|
||||
}
|
||||
func (tool commandReducerFixtureTool) Parameters() Schema {
|
||||
return Schema{Type: "object", AdditionalProperties: true}
|
||||
}
|
||||
func (tool commandReducerFixtureTool) Safety() Safety {
|
||||
return Safety{SideEffect: SideEffectRead, Permission: PermissionAllow, Reason: "test output"}
|
||||
}
|
||||
func (tool commandReducerFixtureTool) Run(context.Context, map[string]any) Result {
|
||||
return Result{Status: StatusOK, Output: tool.output}
|
||||
}
|
||||
|
||||
func TestReduceCommandOutputCompactsPassingGoPackagesAndKeepsRawArtifact(t *testing.T) {
|
||||
var lines []string
|
||||
lines = append(lines, "output:")
|
||||
for index := 0; index < commandReducerMinPassingLines+3; index++ {
|
||||
lines = append(lines, fmt.Sprintf("ok \texample.test/pkg%d\t0.01s", index))
|
||||
}
|
||||
lines = append(lines, "exit_code: 0")
|
||||
original := strings.Join(lines, "\n")
|
||||
|
||||
result := reduceCommandOutput(ExecCommandToolName, map[string]any{"cmd": "go test ./..."}, Result{
|
||||
Status: StatusOK,
|
||||
Output: original,
|
||||
})
|
||||
|
||||
if result.Output == original || !strings.Contains(result.Output, "passing package lines omitted") {
|
||||
t.Fatalf("expected compact result, got %q", result.Output)
|
||||
}
|
||||
if !strings.Contains(result.Output, "exit_code: 0") {
|
||||
t.Fatalf("exit status was lost: %q", result.Output)
|
||||
}
|
||||
spillPath := result.Meta["spill_path"]
|
||||
spilled, err := os.ReadFile(spillPath)
|
||||
if err != nil {
|
||||
t.Fatalf("read raw artifact: %v", err)
|
||||
}
|
||||
if string(spilled) != original {
|
||||
t.Fatalf("raw artifact differs from original output")
|
||||
}
|
||||
|
||||
budgeted := applySelfManagedOutputBudget(
|
||||
NewExecCommandTool(t.TempDir(), newExecSessionManager()),
|
||||
ExecCommandToolName,
|
||||
map[string]any{"cmd": "go test ./..."},
|
||||
result,
|
||||
)
|
||||
if budgeted.Meta["spill_path"] != spillPath {
|
||||
t.Fatalf("budget boundary replaced the exact raw artifact: got %q want %q", budgeted.Meta["spill_path"], spillPath)
|
||||
}
|
||||
spilled, err = os.ReadFile(budgeted.Meta["spill_path"])
|
||||
if err != nil || string(spilled) != original {
|
||||
t.Fatalf("budgeted raw artifact was not preserved: err=%v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceCommandOutputFailsOpenForCompoundAndFailedOutput(t *testing.T) {
|
||||
failure := "output:\n--- FAIL: TestThing\npanic: broken\nFAIL\texample.test/pkg\nexit_code: 1"
|
||||
for _, command := range []string{"go test ./... | tee test.log", "go test ./... && echo done"} {
|
||||
result := reduceCommandOutput(ExecCommandToolName, map[string]any{"cmd": command}, Result{Status: StatusError, Output: failure})
|
||||
if result.Output != failure || result.Meta != nil {
|
||||
t.Fatalf("compound command %q should pass through, got %#v", command, result)
|
||||
}
|
||||
}
|
||||
result := reduceCommandOutput(ExecCommandToolName, map[string]any{"cmd": "go test ./..."}, Result{Status: StatusError, Output: failure})
|
||||
if result.Output != failure {
|
||||
t.Fatalf("failure evidence should pass through when there is no repetitive success bulk: %q", result.Output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceGoTestPassingPackagesPreservesCoverageResults(t *testing.T) {
|
||||
output := strings.Join([]string{
|
||||
"ok \texample.test/a\t0.01s\tcoverage: 82.4% of statements",
|
||||
"? \texample.test/b\t[no test files]",
|
||||
"ok \texample.test/c\t0.01s",
|
||||
}, "\n")
|
||||
|
||||
reduced, omitted := reduceGoTestPassingPackages(output)
|
||||
if !strings.Contains(reduced, "coverage: 82.4% of statements") {
|
||||
t.Fatalf("coverage result was removed: %q", reduced)
|
||||
}
|
||||
if !strings.Contains(reduced, "[no test files]") {
|
||||
t.Fatalf("no-test-files result was removed: %q", reduced)
|
||||
}
|
||||
if omitted != 1 {
|
||||
t.Fatalf("omitted=%d, want only the plain passing line omitted", omitted)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceGoTestPassingPackagesPreservesIndentedTestLogs(t *testing.T) {
|
||||
output := strings.Join([]string{
|
||||
" ok user-provided test log",
|
||||
" ? another test log",
|
||||
"ok \texample.test/pkg\t0.01s",
|
||||
}, "\n")
|
||||
|
||||
reduced, omitted := reduceGoTestPassingPackages(output)
|
||||
if !strings.Contains(reduced, " ok user-provided test log") ||
|
||||
!strings.Contains(reduced, " ? another test log") {
|
||||
t.Fatalf("indented test log was removed: %q", reduced)
|
||||
}
|
||||
if omitted != 1 {
|
||||
t.Fatalf("omitted=%d, want only the package result omitted", omitted)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelfManagedBudgetPreservesRawSpillWhenReducedOutputAlsoTruncates(t *testing.T) {
|
||||
raw := "raw output\n" + strings.Repeat("ok \texample.test/raw\t0.01s\n", 1000)
|
||||
rawSpill := spillTruncatedOutput(ExecCommandToolName, raw)
|
||||
if rawSpill == "" {
|
||||
t.Fatal("failed to create raw spill")
|
||||
}
|
||||
reduced := strings.Repeat("reduced but still large\n", 1000)
|
||||
result := applySelfManagedOutputBudget(
|
||||
NewExecCommandTool(t.TempDir(), newExecSessionManager()),
|
||||
ExecCommandToolName,
|
||||
map[string]any{"cmd": "go test ./...", "max_output_tokens": 100},
|
||||
Result{
|
||||
Status: StatusOK,
|
||||
Output: reduced,
|
||||
Truncated: true,
|
||||
Meta: map[string]string{
|
||||
"spill_path": rawSpill,
|
||||
"truncation_reason": "command_aware_reduction",
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
if result.Meta["spill_path"] != rawSpill {
|
||||
t.Fatalf("budget boundary replaced raw spill: got %q want %q", result.Meta["spill_path"], rawSpill)
|
||||
}
|
||||
spilled, err := os.ReadFile(result.Meta["spill_path"])
|
||||
if err != nil || string(spilled) != raw {
|
||||
t.Fatalf("raw spill was not preserved: err=%v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceCommandOutputCompactsRecognizedTestAndBuildNoise(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
for _, testCase := range commandReducerCorpus() {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
result := reduceCommandOutput(ExecCommandToolName, map[string]any{"cmd": testCase.command}, Result{
|
||||
Status: StatusOK,
|
||||
Output: testCase.output,
|
||||
})
|
||||
if result.Meta["command_output_reduced"] != "true" {
|
||||
t.Fatalf("recognized command output was not reduced: %#v", result)
|
||||
}
|
||||
if !strings.Contains(result.Output, testCase.want) || !strings.Contains(result.Output, "exit_code: 0") {
|
||||
t.Fatalf("decisive summary or exit status was lost: %q", result.Output)
|
||||
}
|
||||
spillPath := result.Meta["spill_path"]
|
||||
spilled, err := os.ReadFile(spillPath)
|
||||
if err != nil {
|
||||
t.Fatalf("read raw artifact: %v", err)
|
||||
}
|
||||
if string(spilled) != testCase.output {
|
||||
t.Fatal("raw artifact differs from original output")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceCommandOutputPreservesFailureEvidence(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
passing := make([]string, 0, commandReducerMinPassingLines+2)
|
||||
for index := 0; index < commandReducerMinPassingLines+2; index++ {
|
||||
passing = append(passing, fmt.Sprintf("test module::case_%02d ... ok", index))
|
||||
}
|
||||
original := strings.Join(append(passing,
|
||||
"test module::broken ... FAILED",
|
||||
"failures:",
|
||||
"---- module::broken stdout ----",
|
||||
"assertion failed: expected zero",
|
||||
"test result: FAILED. 14 passed; 1 failed; finished in 0.08s",
|
||||
"exit_code: 101",
|
||||
), "\n")
|
||||
|
||||
result := reduceCommandOutput(ExecCommandToolName, map[string]any{"cmd": "cargo test --workspace"}, Result{
|
||||
Status: StatusError,
|
||||
Output: original,
|
||||
})
|
||||
for _, evidence := range []string{"module::broken ... FAILED", "assertion failed: expected zero", "14 passed; 1 failed", "exit_code: 101"} {
|
||||
if !strings.Contains(result.Output, evidence) {
|
||||
t.Fatalf("failure evidence %q was lost: %q", evidence, result.Output)
|
||||
}
|
||||
}
|
||||
if result.Meta["command_output_reduced"] != "true" {
|
||||
t.Fatalf("repetitive passing output was not reduced: %#v", result)
|
||||
}
|
||||
spillPath := result.Meta["spill_path"]
|
||||
if spillPath == "" {
|
||||
t.Fatal("reduced failure output did not retain a raw artifact")
|
||||
}
|
||||
spilled, err := os.ReadFile(spillPath)
|
||||
if err != nil {
|
||||
t.Fatalf("read raw artifact: %v", err)
|
||||
}
|
||||
if string(spilled) != original {
|
||||
t.Fatal("raw artifact differs from original output")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceCommandOutputPassesThroughUnsupportedAndCompoundCommands(t *testing.T) {
|
||||
outputs := make(map[string]string)
|
||||
for _, fixture := range commandReducerCorpus() {
|
||||
outputs[fixture.name] = fixture.output
|
||||
}
|
||||
testCases := []struct {
|
||||
command string
|
||||
output string
|
||||
}{
|
||||
{command: "npm test | tee test.log", output: outputs["npm_vitest"]},
|
||||
{command: "cargo test && echo done", output: outputs["cargo_test"]},
|
||||
{command: "pytest; notify-send done", output: outputs["pytest"]},
|
||||
{command: "cargo test 1>test.log", output: outputs["cargo_test"]},
|
||||
{command: "cargo test 2>>errors.log", output: outputs["cargo_test"]},
|
||||
{command: "git status --short", output: outputs["npm_vitest"]},
|
||||
}
|
||||
for _, testCase := range testCases {
|
||||
result := reduceCommandOutput(ExecCommandToolName, map[string]any{"cmd": testCase.command}, Result{Status: StatusOK, Output: testCase.output})
|
||||
if result.Output != testCase.output || result.Meta != nil {
|
||||
t.Fatalf("command %q should pass through exactly, got %#v", testCase.command, result)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceCommandOutputAllowsTrailingStderrMerge(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
fixture := commandReducerCorpus()[1]
|
||||
result := reduceCommandOutput(ExecCommandToolName, map[string]any{"cmd": fixture.command + " 2>&1"}, Result{
|
||||
Status: StatusOK,
|
||||
Output: fixture.output,
|
||||
})
|
||||
if result.Meta["command_output_reducer"] != "cargo_test" {
|
||||
t.Fatalf("trailing stderr merge prevented reduction: %#v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceCommandOutputRequiresAuthoritativeRunnerSummary(t *testing.T) {
|
||||
lookalikes := map[string]string{
|
||||
"cargo test": "test user::message ... ok",
|
||||
"cargo check": "Checking a custom deployment state",
|
||||
"pytest": "application log ........ [100%]",
|
||||
"npm test": "PASS user-provided log message",
|
||||
}
|
||||
for command, line := range lookalikes {
|
||||
original := strings.Repeat(line+"\n", commandReducerMinPassingLines+2) + "exit_code: 0"
|
||||
result := reduceCommandOutput(ExecCommandToolName, map[string]any{"cmd": command}, Result{Status: StatusOK, Output: original})
|
||||
if result.Output != original || result.Meta != nil {
|
||||
t.Fatalf("summary-free output for %q should pass through, got %#v", command, result)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceCommandOutputRejectsLookalikeRunnerSummaries(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
testCases := []struct {
|
||||
command string
|
||||
line string
|
||||
summary string
|
||||
}{
|
||||
{command: "cargo check", line: "Compiling dependency v1.0.0", summary: "Finished generating bindings"},
|
||||
{command: "pytest", line: "tests/test_module.py ........ [100%]", summary: "application log: 14 passed validation in staging"},
|
||||
{command: "npm test", line: "PASS src/example.test.ts", summary: "Test Files discovered: 14"},
|
||||
}
|
||||
for _, testCase := range testCases {
|
||||
original := strings.Repeat(testCase.line+"\n", commandReducerMinPassingLines+2) + testCase.summary + "\nexit_code: 1"
|
||||
result := reduceCommandOutput(ExecCommandToolName, map[string]any{"cmd": testCase.command}, Result{Status: StatusError, Output: original})
|
||||
if result.Output != original || result.Meta != nil {
|
||||
t.Fatalf("lookalike summary for %q should pass through exactly, got %#v", testCase.command, result)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryAppliesCommandReducerToExecAndBashTools(t *testing.T) {
|
||||
t.Setenv("TMPDIR", t.TempDir())
|
||||
fixture := commandReducerCorpus()[1]
|
||||
for _, toolName := range []string{ExecCommandToolName, "bash"} {
|
||||
t.Run(toolName, func(t *testing.T) {
|
||||
registry := NewRegistry()
|
||||
registry.Register(commandReducerFixtureTool{name: toolName, output: fixture.output})
|
||||
args := map[string]any{"cmd": fixture.command}
|
||||
if toolName == "bash" {
|
||||
args = map[string]any{"command": fixture.command}
|
||||
}
|
||||
result := registry.Run(context.Background(), toolName, args)
|
||||
if result.Status != StatusOK || result.Meta["command_output_reducer"] != "cargo_test" {
|
||||
t.Fatalf("registry result was not reduced: %#v", result)
|
||||
}
|
||||
if !strings.Contains(result.Output, fixture.want) {
|
||||
t.Fatalf("registry reduction lost summary: %q", result.Output)
|
||||
}
|
||||
spilled, err := os.ReadFile(result.Meta["spill_path"])
|
||||
if err != nil || string(spilled) != fixture.output {
|
||||
t.Fatalf("registry raw artifact mismatch: err=%v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+301
-14
@@ -1,15 +1,33 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
"unicode"
|
||||
|
||||
zeroSandbox "github.com/Gitlawb/zero/internal/sandbox"
|
||||
)
|
||||
|
||||
type shellKind string
|
||||
|
||||
const (
|
||||
shellKindPOSIX shellKind = "posix"
|
||||
shellKindPowerShell shellKind = "powershell"
|
||||
shellKindCmd shellKind = "cmd"
|
||||
)
|
||||
|
||||
type shellRuntime struct {
|
||||
GOOS string
|
||||
Executable string
|
||||
Syntax string
|
||||
Kind shellKind
|
||||
}
|
||||
|
||||
type shellIssue struct {
|
||||
@@ -20,7 +38,18 @@ type shellIssue struct {
|
||||
|
||||
const windowsMsysSandboxKind = "windows_msys_sandbox"
|
||||
|
||||
const windowsMsysSandboxSuggestion = "MSYS/Cygwin coreutils and shells (bash, sh) from Git for Windows cannot run under Zero's write-restricted Windows sandbox, and the WSL bash launcher cannot reach the WSL service from it either. This also hits native commands that spawn Git Bash internally: git hooks and git/gh credential helpers can fail this way even though git and gh themselves run fine. Prefer Zero native tools (grep, read_file with offset/limit, list_directory, glob), cmd.exe findstr, or PowerShell Select-Object -First/-Last. If host-level execution is truly required, rerun with sandbox_permissions: \"require_escalated\" and a narrow justification."
|
||||
const windowsMsysSandboxSuggestion = "MSYS/Cygwin executables and shells (bash, sh) from Git for Windows cannot run under Zero's write-restricted Windows sandbox, and the WSL bash launcher cannot reach the WSL service from it either. This also hits native commands that spawn Git Bash internally: git hooks and git/gh credential helpers can fail this way even though git and gh themselves run fine. Prefer native PowerShell cmdlets or Zero tools (grep, read_file with offset/limit, list_directory, glob). If host-level execution is truly required, rerun with sandbox_permissions: \"require_escalated\" and a narrow justification."
|
||||
|
||||
const windowsPowerShellUTF8Prefix = "try { [Console]::OutputEncoding = [System.Text.Encoding]::UTF8 } catch {}\n"
|
||||
|
||||
const windowsPowerShellFailurePrefix = "$ErrorActionPreference = 'Stop'\n"
|
||||
|
||||
const windowsPowerShellExitSuffix = "\nif ($null -ne $LASTEXITCODE) { exit $LASTEXITCODE }\n"
|
||||
|
||||
var (
|
||||
hostShellOnce sync.Once
|
||||
hostShell shellRuntime
|
||||
)
|
||||
|
||||
// windowsMsysProneNames is the single source of truth for POSIX coreutil and
|
||||
// shell names that commonly resolve to a Git-for-Windows MSYS/Cygwin binary
|
||||
@@ -54,24 +83,139 @@ var (
|
||||
)
|
||||
|
||||
func detectShellRuntime(goos string) shellRuntime {
|
||||
if goos == "windows" {
|
||||
return shellRuntime{GOOS: goos, Executable: "cmd.exe", Syntax: "Windows cmd.exe"}
|
||||
if goos != runtime.GOOS {
|
||||
return detectShellRuntimeWithLookup(goos, exec.LookPath, os.Getenv)
|
||||
}
|
||||
return shellRuntime{GOOS: goos, Executable: "/bin/sh", Syntax: "/bin/sh"}
|
||||
hostShellOnce.Do(func() {
|
||||
hostShell = detectShellRuntimeWithProbe(goos, exec.LookPath, os.Getenv, powerShellExecutableUsable)
|
||||
})
|
||||
return hostShell
|
||||
}
|
||||
|
||||
func detectShellRuntimeWithLookup(goos string, lookPath func(string) (string, error), getenv func(string) string) shellRuntime {
|
||||
return detectShellRuntimeWithProbe(goos, lookPath, getenv, func(string) bool { return true })
|
||||
}
|
||||
|
||||
func detectShellRuntimeWithProbe(goos string, lookPath func(string) (string, error), getenv func(string) string, usable func(string) bool) shellRuntime {
|
||||
if goos == "windows" {
|
||||
for _, candidate := range windowsPowerShellCandidates(getenv) {
|
||||
if path, err := lookPath(candidate); err == nil && strings.TrimSpace(path) != "" && usable(path) {
|
||||
return shellRuntime{
|
||||
GOOS: goos,
|
||||
Executable: path,
|
||||
Syntax: "PowerShell",
|
||||
Kind: shellKindPowerShell,
|
||||
}
|
||||
}
|
||||
}
|
||||
return shellRuntime{GOOS: goos, Executable: "cmd.exe", Syntax: "Windows cmd.exe", Kind: shellKindCmd}
|
||||
}
|
||||
return shellRuntime{GOOS: goos, Executable: "/bin/sh", Syntax: "/bin/sh", Kind: shellKindPOSIX}
|
||||
}
|
||||
|
||||
func powerShellExecutableUsable(path string) bool {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
command := exec.CommandContext(ctx, path, "-NoLogo", "-NoProfile", "-Command", "exit 0")
|
||||
return command.Run() == nil
|
||||
}
|
||||
|
||||
func windowsPowerShellCandidates(getenv func(string) string) []string {
|
||||
candidates := []string{"pwsh.exe", "pwsh"}
|
||||
if programFiles := strings.TrimSpace(getenv("ProgramFiles")); programFiles != "" {
|
||||
candidates = append(candidates, filepath.Join(programFiles, "PowerShell", "7", "pwsh.exe"))
|
||||
}
|
||||
candidates = append(candidates, "powershell.exe", "powershell")
|
||||
if systemRoot := strings.TrimSpace(getenv("SystemRoot")); systemRoot != "" {
|
||||
candidates = append(candidates, filepath.Join(systemRoot, "System32", "WindowsPowerShell", "v1.0", "powershell.exe"))
|
||||
}
|
||||
return candidates
|
||||
}
|
||||
|
||||
func (shell shellRuntime) arguments(command string) []string {
|
||||
switch shell.Kind {
|
||||
case shellKindPowerShell:
|
||||
script := windowsPowerShellUTF8Prefix + windowsPowerShellFailurePrefix + command + windowsPowerShellExitSuffix
|
||||
return []string{"-NoLogo", "-NoProfile", "-Command", script}
|
||||
case shellKindCmd:
|
||||
return zeroSandbox.WindowsShellArgs(command)
|
||||
default:
|
||||
return []string{"-c", command}
|
||||
}
|
||||
}
|
||||
|
||||
// HostShellCommand returns the shell executable and argv used for a command on
|
||||
// this host. Keeping this in one place makes agent tools and the user-typed TUI
|
||||
// shell escape use the same Windows shell and fallback behavior.
|
||||
func HostShellCommand(command string) (string, []string) {
|
||||
shell := detectShellRuntime(runtime.GOOS)
|
||||
return shell.Executable, shell.arguments(command)
|
||||
}
|
||||
|
||||
// NewHostShellCommandContext constructs a direct host-shell command. The
|
||||
// cmd.exe fallback needs a Windows-specific raw command line, while PowerShell
|
||||
// and POSIX shells use normal argv handling.
|
||||
func NewHostShellCommandContext(ctx context.Context, commandText string) *exec.Cmd {
|
||||
shell := detectShellRuntime(runtime.GOOS)
|
||||
command := exec.CommandContext(ctx, shell.Executable, shell.arguments(commandText)...)
|
||||
applyWindowsShellCommandLine(command, commandText, false, shell.Kind == shellKindCmd)
|
||||
return command
|
||||
}
|
||||
|
||||
func shellGuidanceForGOOS(goos string) string {
|
||||
runtime := detectShellRuntime(goos)
|
||||
if goos == "windows" {
|
||||
return "Uses " + runtime.Syntax + " syntax on Windows; prefer cwd over cd when changing directories. To include | & > < or other metacharacters in an argument value, wrap the value in double quotes (e.g. --jq \".a | b\"); single quotes do not protect metacharacters in cmd.exe. MSYS/Cygwin coreutils on PATH (Git for Windows usr\\bin) are not sandbox-compatible; prefer native Zero file tools."
|
||||
return shellGuidanceForRuntime(detectShellRuntime(goos))
|
||||
}
|
||||
|
||||
func shellGuidanceForRuntime(shell shellRuntime) string {
|
||||
if shell.GOOS == "windows" && shell.Kind == shellKindPowerShell {
|
||||
guidance := "Uses PowerShell syntax on Windows; prefer cwd/workdir over Set-Location when changing directories. Examples: Get-ChildItem -Force; Get-ChildItem -Recurse -Filter *.go; Get-ChildItem -Recurse | Select-String -Pattern 'TODO'; Get-Process | Where-Object { $_.ProcessName -like '*node*' }; $env:NAME='value'; @'\nprint('hello')\n'@ | python -. Do not invoke Git-for-Windows MSYS/Cygwin executables (bash, sh, grep.exe, sed.exe, head.exe, and similar) inside the restricted sandbox; prefer PowerShell cmdlets or native Zero tools."
|
||||
return guidance + legacyPowerShellChainGuidance(shell)
|
||||
}
|
||||
guidance := "Uses " + runtime.Syntax + " syntax."
|
||||
if goos == "darwin" {
|
||||
if shell.GOOS == "windows" {
|
||||
return "Uses " + shell.Syntax + " syntax on Windows because PowerShell was unavailable; prefer cwd/workdir over cd when changing directories. To include | & > < or other metacharacters in an argument value, wrap the value in double quotes (e.g. --jq \".a | b\"); single quotes do not protect metacharacters in cmd.exe. MSYS/Cygwin coreutils on PATH (Git for Windows usr\\bin) are not sandbox-compatible; prefer native Zero file tools."
|
||||
}
|
||||
guidance := "Uses " + shell.Syntax + " syntax."
|
||||
if shell.GOOS == "darwin" {
|
||||
guidance += " To find or stop a process, use `lsof -i :PORT` (or `lsof -nP -iTCP -sTCP:LISTEN`) for the PID then `kill <pid>`; `ps` and `pgrep` do not work under the sandbox."
|
||||
}
|
||||
return guidance
|
||||
}
|
||||
|
||||
// HostShellEnvironmentGuidance returns the concise, model-facing shell rule
|
||||
// for the current host.
|
||||
func HostShellEnvironmentGuidance() string {
|
||||
return hostShellEnvironmentGuidanceForRuntime(detectShellRuntime(runtime.GOOS))
|
||||
}
|
||||
|
||||
func hostShellEnvironmentGuidanceForRuntime(shell shellRuntime) string {
|
||||
if shell.GOOS == "windows" && shell.Kind == shellKindPowerShell {
|
||||
guidance := "Shell syntax: PowerShell for exec_command/bash tools. Use PowerShell cmdlets and pipelines (Get-ChildItem, Get-Content, Select-String, Select-Object), use $env:NAME='value' for environment variables, and prefer the workdir/cwd argument over Set-Location. Do not invoke Git-for-Windows MSYS binaries (bash, sh, grep.exe, sed.exe, head.exe, and similar) inside the restricted sandbox; use native PowerShell cmdlets or Zero tools instead, or sandbox_permissions require_escalated only when host-level execution is truly required."
|
||||
return guidance + legacyPowerShellChainGuidance(shell)
|
||||
}
|
||||
if shell.GOOS == "windows" {
|
||||
return "Shell syntax: Windows cmd.exe syntax for exec_command/bash tools because PowerShell is unavailable. To put | & > < etc inside an arg value, use double quotes around the value, not single quotes. Do not invoke Git-for-Windows MSYS binaries inside the restricted sandbox; use native Zero tools instead. Prefer the workdir/cwd argument over cd."
|
||||
}
|
||||
return "Shell syntax: /bin/sh syntax for exec_command/bash tools; prefer the workdir/cwd argument instead of cd when changing directories."
|
||||
}
|
||||
|
||||
func legacyPowerShellChainGuidance(shell shellRuntime) string {
|
||||
if !shell.isWindowsPowerShell() {
|
||||
return ""
|
||||
}
|
||||
return " This host uses Windows PowerShell 5.1, which does not support && or ||. Use separate statements when execution is unconditional, or if ($?) { ... } and if (-not $?) { ... } for conditional chaining."
|
||||
}
|
||||
|
||||
func (shell shellRuntime) isWindowsPowerShell() bool {
|
||||
if shell.Kind != shellKindPowerShell {
|
||||
return false
|
||||
}
|
||||
executable := strings.TrimSpace(shell.Executable)
|
||||
if index := strings.LastIndexAny(executable, `\/`); index >= 0 {
|
||||
executable = executable[index+1:]
|
||||
}
|
||||
return strings.EqualFold(executable, "powershell") || strings.EqualFold(executable, "powershell.exe")
|
||||
}
|
||||
|
||||
// MsysProneCommandName reports whether a bare command name commonly resolves to
|
||||
// a Git-for-Windows MSYS binary that fails under the Windows restricted sandbox.
|
||||
func MsysProneCommandName(name string) bool {
|
||||
@@ -172,7 +316,15 @@ func msysProneCommandWord(word string) bool {
|
||||
}
|
||||
|
||||
func detectShellCommandIssue(command string, goos string) *shellIssue {
|
||||
if goos != "windows" {
|
||||
shell := shellRuntime{GOOS: goos, Executable: "/bin/sh", Syntax: "/bin/sh", Kind: shellKindPOSIX}
|
||||
if goos == "windows" {
|
||||
shell = shellRuntime{GOOS: goos, Executable: "cmd.exe", Syntax: "Windows cmd.exe", Kind: shellKindCmd}
|
||||
}
|
||||
return detectShellCommandIssueForRuntime(command, shell)
|
||||
}
|
||||
|
||||
func detectShellCommandIssueForRuntime(command string, shell shellRuntime) *shellIssue {
|
||||
if shell.GOOS != "windows" {
|
||||
return nil
|
||||
}
|
||||
trimmed := strings.TrimSpace(command)
|
||||
@@ -184,7 +336,24 @@ func detectShellCommandIssue(command string, goos string) *shellIssue {
|
||||
// cat), so an escaped metachar can't stand in as a fake segment boundary
|
||||
// either.
|
||||
unquoted := stripCmdCaretEscapes(stripDoubleQuotedSpans(trimmed))
|
||||
if shell.Kind == shellKindPowerShell {
|
||||
unquoted = stripPowerShellQuotedSpans(trimmed)
|
||||
if shell.isWindowsPowerShell() && (strings.Contains(unquoted, "&&") || strings.Contains(unquoted, "||")) {
|
||||
return &shellIssue{
|
||||
Kind: "windows_powershell_version",
|
||||
Message: "Command uses && or ||, but this host runs Windows PowerShell 5.1, which does not support those operators.",
|
||||
Suggestion: "Use separate statements when execution is unconditional, or if ($?) { ... } and if (-not $?) { ... } for conditional chaining.",
|
||||
}
|
||||
}
|
||||
}
|
||||
if windowsBashStyleCDPattern.MatchString(unquoted) {
|
||||
if shell.Kind == shellKindPowerShell {
|
||||
return &shellIssue{
|
||||
Kind: "windows_shell_syntax",
|
||||
Message: "Command looks like POSIX/Bash syntax, but Zero runs PowerShell commands on this Windows host.",
|
||||
Suggestion: "Use the cwd/workdir argument instead of cd and use native PowerShell syntax or Zero tools such as list_directory, read_file, grep, and glob.",
|
||||
}
|
||||
}
|
||||
return &shellIssue{
|
||||
Kind: "windows_shell_syntax",
|
||||
Message: "Command looks like POSIX/Bash syntax, but Zero runs bash tool commands through Windows cmd.exe on this host.",
|
||||
@@ -198,18 +367,128 @@ func detectShellCommandIssue(command string, goos string) *shellIssue {
|
||||
// Being segment/word anchored rather than a whole-string regex or scan,
|
||||
// neither check matches text that only appears inside a quoted argument
|
||||
// (e.g. a commit message or PR comment body discussing head.exe).
|
||||
for _, segment := range windowsCommandSegments(trimmed) {
|
||||
segments := windowsCommandSegments(trimmed)
|
||||
if shell.Kind == shellKindPowerShell {
|
||||
segments = powerShellCommandSegments(trimmed)
|
||||
}
|
||||
for _, segment := range segments {
|
||||
word := firstCommandWord(segment)
|
||||
if windowsMsysBinaryPathPattern.MatchString(word) {
|
||||
return windowsMsysSandboxIssue("Command invokes an MSYS/Cygwin binary path that cannot run under Zero's Windows sandbox.")
|
||||
}
|
||||
if msysProneCommandWord(word) {
|
||||
if shell.Kind == shellKindPowerShell && powerShellAliasWord(word) {
|
||||
continue
|
||||
}
|
||||
return windowsMsysSandboxIssue("Command uses a POSIX coreutil (head/tail/grep/cat/...) that commonly resolves to Git-for-Windows MSYS binaries incompatible with the Windows sandbox.")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// powerShellCommandSegments splits the subset of PowerShell syntax needed for
|
||||
// command-position checks. It respects single/double quoted strings and the
|
||||
// backtick escape, then splits on pipeline/statement operators. Complex script
|
||||
// interpretation remains PowerShell's job; this scanner only determines
|
||||
// whether a known-incompatible MSYS executable is actually invoked.
|
||||
func powerShellCommandSegments(command string) []string {
|
||||
var segments []string
|
||||
var current strings.Builder
|
||||
var quote rune
|
||||
escaped := false
|
||||
for _, c := range command {
|
||||
if escaped {
|
||||
current.WriteRune(c)
|
||||
escaped = false
|
||||
continue
|
||||
}
|
||||
if c == '`' {
|
||||
current.WriteRune(c)
|
||||
escaped = true
|
||||
continue
|
||||
}
|
||||
if quote != 0 {
|
||||
current.WriteRune(c)
|
||||
if c == quote {
|
||||
quote = 0
|
||||
}
|
||||
continue
|
||||
}
|
||||
if c == '\'' || c == '"' {
|
||||
quote = c
|
||||
current.WriteRune(c)
|
||||
continue
|
||||
}
|
||||
if c == ';' || c == '|' || c == '&' {
|
||||
if segment := strings.TrimSpace(current.String()); segment != "" {
|
||||
segments = append(segments, segment)
|
||||
}
|
||||
current.Reset()
|
||||
continue
|
||||
}
|
||||
current.WriteRune(c)
|
||||
}
|
||||
if segment := strings.TrimSpace(current.String()); segment != "" {
|
||||
segments = append(segments, segment)
|
||||
}
|
||||
return segments
|
||||
}
|
||||
|
||||
func powerShellAliasWord(word string) bool {
|
||||
trimmed := strings.Trim(word, `"'`)
|
||||
if strings.ContainsAny(trimmed, `\/`) || strings.HasSuffix(strings.ToLower(trimmed), ".exe") {
|
||||
return false
|
||||
}
|
||||
switch strings.ToLower(trimmed) {
|
||||
case "cat", "ls":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func stripPowerShellQuotedSpans(command string) string {
|
||||
var b strings.Builder
|
||||
b.Grow(len(command))
|
||||
var quote rune
|
||||
escaped := false
|
||||
runes := []rune(command)
|
||||
for index := 0; index < len(runes); index++ {
|
||||
c := runes[index]
|
||||
if escaped {
|
||||
b.WriteRune(' ')
|
||||
escaped = false
|
||||
continue
|
||||
}
|
||||
if c == '`' {
|
||||
b.WriteRune(' ')
|
||||
escaped = true
|
||||
continue
|
||||
}
|
||||
if quote != 0 {
|
||||
b.WriteRune(' ')
|
||||
if c == quote {
|
||||
// PowerShell escapes a single quote inside a single-quoted
|
||||
// string by doubling it.
|
||||
if quote == '\'' && index+1 < len(runes) && runes[index+1] == '\'' {
|
||||
index++
|
||||
b.WriteRune(' ')
|
||||
continue
|
||||
}
|
||||
quote = 0
|
||||
}
|
||||
continue
|
||||
}
|
||||
if c == '\'' || c == '"' {
|
||||
quote = c
|
||||
b.WriteRune(' ')
|
||||
continue
|
||||
}
|
||||
b.WriteRune(c)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// stripDoubleQuotedSpans replaces the contents of every double-quoted span
|
||||
// (quotes included) with spaces, preserving the string's length and the
|
||||
// position of unquoted text. cmd.exe treats a double-quoted span as a single
|
||||
@@ -256,7 +535,15 @@ func stripCmdCaretEscapes(command string) string {
|
||||
// quoted-text false positives the preflight command-position check exists to
|
||||
// avoid, just after execution instead of before it.
|
||||
func detectShellOutputIssue(output string, goos string) *shellIssue {
|
||||
if goos != "windows" {
|
||||
shell := shellRuntime{GOOS: goos, Executable: "/bin/sh", Syntax: "/bin/sh", Kind: shellKindPOSIX}
|
||||
if goos == "windows" {
|
||||
shell = shellRuntime{GOOS: goos, Executable: "cmd.exe", Syntax: "Windows cmd.exe", Kind: shellKindCmd}
|
||||
}
|
||||
return detectShellOutputIssueForRuntime(output, shell)
|
||||
}
|
||||
|
||||
func detectShellOutputIssueForRuntime(output string, shell shellRuntime) *shellIssue {
|
||||
if shell.GOOS != "windows" {
|
||||
return nil
|
||||
}
|
||||
lower := strings.ToLower(output)
|
||||
@@ -266,8 +553,8 @@ func detectShellOutputIssue(output string, goos string) *shellIssue {
|
||||
if wslServiceDeniedInOutput(lower) {
|
||||
return windowsMsysSandboxIssue("WSL bash could not connect to the WSL service under Zero's Windows sandbox (Bash/Service/CreateInstance/E_ACCESSDENIED).")
|
||||
}
|
||||
if strings.Contains(lower, "the syntax of the command is incorrect") ||
|
||||
strings.Contains(lower, "is not recognized as an internal or external command") {
|
||||
if shell.Kind == shellKindCmd && (strings.Contains(lower, "the syntax of the command is incorrect") ||
|
||||
strings.Contains(lower, "is not recognized as an internal or external command")) {
|
||||
return &shellIssue{
|
||||
Kind: "windows_shell_syntax",
|
||||
Message: "Windows cmd.exe rejected the command syntax.",
|
||||
|
||||
@@ -1,10 +1,218 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os/exec"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestDetectShellRuntimePrefersPowerShellSevenOnWindows(t *testing.T) {
|
||||
lookPath := func(name string) (string, error) {
|
||||
if name == "pwsh.exe" {
|
||||
return `C:\Program Files\PowerShell\7\pwsh.exe`, nil
|
||||
}
|
||||
return "", errors.New("not found")
|
||||
}
|
||||
shell := detectShellRuntimeWithLookup("windows", lookPath, func(string) string { return "" })
|
||||
if shell.Kind != shellKindPowerShell || !strings.HasSuffix(strings.ToLower(shell.Executable), `\pwsh.exe`) {
|
||||
t.Fatalf("shell = %#v, want PowerShell 7", shell)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectShellRuntimeFallsBackThroughWindowsPowerShellToCmd(t *testing.T) {
|
||||
t.Run("windows powershell", func(t *testing.T) {
|
||||
lookPath := func(name string) (string, error) {
|
||||
if name == "powershell.exe" {
|
||||
return `C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe`, nil
|
||||
}
|
||||
return "", errors.New("not found")
|
||||
}
|
||||
shell := detectShellRuntimeWithLookup("windows", lookPath, func(string) string { return "" })
|
||||
if shell.Kind != shellKindPowerShell || !strings.HasSuffix(strings.ToLower(shell.Executable), `\powershell.exe`) {
|
||||
t.Fatalf("shell = %#v, want Windows PowerShell", shell)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("cmd", func(t *testing.T) {
|
||||
shell := detectShellRuntimeWithLookup(
|
||||
"windows",
|
||||
func(string) (string, error) { return "", errors.New("not found") },
|
||||
func(string) string { return "" },
|
||||
)
|
||||
if shell.Kind != shellKindCmd || shell.Executable != "cmd.exe" {
|
||||
t.Fatalf("shell = %#v, want cmd.exe fallback", shell)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestDetectShellRuntimeSkipsUnusablePowerShell(t *testing.T) {
|
||||
lookPath := func(name string) (string, error) {
|
||||
switch name {
|
||||
case "pwsh.exe":
|
||||
return `C:\broken\pwsh.exe`, nil
|
||||
case "powershell.exe":
|
||||
return `C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe`, nil
|
||||
default:
|
||||
return "", errors.New("not found")
|
||||
}
|
||||
}
|
||||
shell := detectShellRuntimeWithProbe(
|
||||
"windows",
|
||||
lookPath,
|
||||
func(string) string { return "" },
|
||||
func(path string) bool { return !strings.Contains(path, `\broken\`) },
|
||||
)
|
||||
if shell.Kind != shellKindPowerShell || strings.Contains(shell.Executable, `\broken\`) {
|
||||
t.Fatalf("shell = %#v, want usable Windows PowerShell fallback", shell)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPowerShellArgumentsDisableProfilesAndRequestUTF8(t *testing.T) {
|
||||
shell := shellRuntime{GOOS: "windows", Executable: "pwsh.exe", Syntax: "PowerShell", Kind: shellKindPowerShell}
|
||||
args := shell.arguments("Write-Output 'hello'")
|
||||
joined := strings.Join(args, "\n")
|
||||
for _, want := range []string{
|
||||
"-NoLogo",
|
||||
"-NoProfile",
|
||||
"-Command",
|
||||
windowsPowerShellUTF8Prefix,
|
||||
"$ErrorActionPreference = 'Stop'",
|
||||
"Write-Output 'hello'",
|
||||
"exit $LASTEXITCODE",
|
||||
} {
|
||||
if !strings.Contains(joined, want) {
|
||||
t.Fatalf("PowerShell args missing %q: %#v", want, args)
|
||||
}
|
||||
}
|
||||
script := args[len(args)-1]
|
||||
preferenceIndex := strings.Index(script, "$ErrorActionPreference")
|
||||
commandIndex := strings.Index(script, "Write-Output 'hello'")
|
||||
exitIndex := strings.Index(script, "exit $LASTEXITCODE")
|
||||
if preferenceIndex < 0 || commandIndex <= preferenceIndex || exitIndex <= commandIndex {
|
||||
t.Fatalf("PowerShell failure handling is not ordered around the command: %q", script)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWindowsPowerShellArgumentsPropagateFailures(t *testing.T) {
|
||||
if runtime.GOOS != "windows" {
|
||||
t.Skip("Windows PowerShell integration test")
|
||||
}
|
||||
executable, err := exec.LookPath("powershell.exe")
|
||||
if err != nil {
|
||||
t.Skipf("Windows PowerShell unavailable: %v", err)
|
||||
}
|
||||
shell := shellRuntime{
|
||||
GOOS: "windows",
|
||||
Executable: executable,
|
||||
Syntax: "PowerShell",
|
||||
Kind: shellKindPowerShell,
|
||||
}
|
||||
|
||||
t.Run("native exit code after cmdlet", func(t *testing.T) {
|
||||
command := exec.Command(executable, shell.arguments("cmd /c exit 5; Write-Output done")...)
|
||||
output, err := command.CombinedOutput()
|
||||
var exitErr *exec.ExitError
|
||||
if !errors.As(err, &exitErr) || exitErr.ExitCode() != 5 {
|
||||
t.Fatalf("exit = %v, output = %q; want exact native exit code 5", err, output)
|
||||
}
|
||||
if !strings.Contains(string(output), "done") {
|
||||
t.Fatalf("output = %q, want command output before propagated failure", output)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("cmdlet error stops script", func(t *testing.T) {
|
||||
missing := strings.ReplaceAll(t.TempDir()+`\missing`, `'`, `''`)
|
||||
commandText := "Get-Item -LiteralPath '" + missing + "'; Write-Output after"
|
||||
command := exec.Command(executable, shell.arguments(commandText)...)
|
||||
output, err := command.CombinedOutput()
|
||||
if err == nil {
|
||||
t.Fatalf("cmdlet failure reported success: %q", output)
|
||||
}
|
||||
if strings.Contains(string(output), "after") {
|
||||
t.Fatalf("script continued after cmdlet failure: %q", output)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestWindowsPowerShellGuidanceAvoidsUnsupportedChainOperators(t *testing.T) {
|
||||
shell := shellRuntime{
|
||||
GOOS: "windows",
|
||||
Executable: `C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe`,
|
||||
Syntax: "PowerShell",
|
||||
Kind: shellKindPowerShell,
|
||||
}
|
||||
for _, guidance := range []string{
|
||||
shellGuidanceForRuntime(shell),
|
||||
hostShellEnvironmentGuidanceForRuntime(shell),
|
||||
} {
|
||||
for _, want := range []string{"Windows PowerShell 5.1", "&&", "||", "if ($?)"} {
|
||||
if !strings.Contains(guidance, want) {
|
||||
t.Fatalf("legacy PowerShell guidance missing %q: %q", want, guidance)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWindowsPowerShellCommandIssueFlagsUnsupportedChainOperators(t *testing.T) {
|
||||
legacy := shellRuntime{
|
||||
GOOS: "windows",
|
||||
Executable: "powershell.exe",
|
||||
Syntax: "PowerShell",
|
||||
Kind: shellKindPowerShell,
|
||||
}
|
||||
for _, command := range []string{
|
||||
"Write-Output one && Write-Output two",
|
||||
"Write-Output one || Write-Output two",
|
||||
} {
|
||||
issue := detectShellCommandIssueForRuntime(command, legacy)
|
||||
if issue == nil || issue.Kind != "windows_powershell_version" {
|
||||
t.Fatalf("legacy PowerShell command %q issue = %#v, want version issue", command, issue)
|
||||
}
|
||||
}
|
||||
for _, command := range []string{
|
||||
`Write-Output "one && two"`,
|
||||
"Write-Output 'one || two'",
|
||||
"Write-Output one `&`& Write-Output two",
|
||||
} {
|
||||
if issue := detectShellCommandIssueForRuntime(command, legacy); issue != nil {
|
||||
t.Fatalf("literal legacy PowerShell operators in %q produced issue %#v", command, issue)
|
||||
}
|
||||
}
|
||||
|
||||
modern := shellRuntime{GOOS: "windows", Executable: "pwsh.exe", Syntax: "PowerShell", Kind: shellKindPowerShell}
|
||||
if issue := detectShellCommandIssueForRuntime("Write-Output one && Write-Output two", modern); issue != nil {
|
||||
t.Fatalf("PowerShell 7 chain operator produced issue %#v", issue)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPowerShellCommandIssueAllowsAliasesAndBlocksMsysExecutables(t *testing.T) {
|
||||
shell := shellRuntime{GOOS: "windows", Executable: "pwsh.exe", Syntax: "PowerShell", Kind: shellKindPowerShell}
|
||||
for _, command := range []string{
|
||||
`ls -Force`,
|
||||
`cat README.md`,
|
||||
`Get-ChildItem | Select-Object -First 10`,
|
||||
`Write-Output 'grep README.md; head file.txt'`,
|
||||
`Write-Output 'cd /tmp'`,
|
||||
} {
|
||||
if issue := detectShellCommandIssueForRuntime(command, shell); issue != nil {
|
||||
t.Fatalf("PowerShell-native command %q was blocked: %#v", command, issue)
|
||||
}
|
||||
}
|
||||
for _, command := range []string{
|
||||
`grep README.md`,
|
||||
`Get-Content README.md | head -10`,
|
||||
`cat.exe README.md`,
|
||||
`bash -lc "make test"`,
|
||||
`Write-Output ok; sed -n '1,5p' README.md`,
|
||||
} {
|
||||
if issue := detectShellCommandIssueForRuntime(command, shell); issue == nil || issue.Kind != windowsMsysSandboxKind {
|
||||
t.Fatalf("MSYS-prone command %q was not blocked: %#v", command, issue)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectShellCommandIssueFlagsMsysBinaryPaths(t *testing.T) {
|
||||
for _, command := range []string{
|
||||
`for /F %i in ('whoami') do echo %i | "C:\Program Files\Git\usr\bin\head.exe" -1`,
|
||||
|
||||
@@ -83,6 +83,7 @@ type PropertySchema struct {
|
||||
Minimum *int `json:"minimum,omitempty"`
|
||||
Maximum *int `json:"maximum,omitempty"`
|
||||
MinLength *int `json:"minLength,omitempty"`
|
||||
MaxLength *int `json:"maxLength,omitempty"`
|
||||
MinItems *int `json:"minItems,omitempty"`
|
||||
// Properties/Required describe nested object fields (for Type "object" or an
|
||||
// object-typed Items).
|
||||
@@ -114,6 +115,9 @@ type Result struct {
|
||||
ChangeSummaries []execution.Change
|
||||
// Display carries a short, structured summary for the TUI / stream.
|
||||
Display Display
|
||||
// pendingFileObservation is proposed by read_file and committed only after
|
||||
// the final model-visible output boundary confirms the exact content survived.
|
||||
pendingFileObservation *pendingFileObservation
|
||||
}
|
||||
|
||||
// Display carries a short, structured summary of a tool result for the TUI/stream.
|
||||
@@ -158,6 +162,7 @@ type baseTool struct {
|
||||
parameters Schema
|
||||
safety Safety
|
||||
capabilities ToolCapabilities // zero value = EffectUnknown, not thread-safe
|
||||
deferred bool
|
||||
}
|
||||
|
||||
func (tool baseTool) Name() string {
|
||||
@@ -176,6 +181,11 @@ func (tool baseTool) Safety() Safety {
|
||||
return tool.safety
|
||||
}
|
||||
|
||||
// Deferred reports whether this built-in is discoverable on demand.
|
||||
func (tool baseTool) Deferred() bool {
|
||||
return tool.deferred
|
||||
}
|
||||
|
||||
func okResult(output string) Result {
|
||||
return Result{Status: StatusOK, Output: output}
|
||||
}
|
||||
|
||||
@@ -107,6 +107,7 @@ func newWebFetchToolWithClientAndResolver(client *http.Client, resolver webFetch
|
||||
baseTool: baseTool{
|
||||
name: "web_fetch",
|
||||
description: "Fetch text from a public remote HTTP or HTTPS URL after network permission is granted. Do not use for localhost, private network URLs, or local dev servers; use bash with curl for those.",
|
||||
deferred: true,
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]PropertySchema{
|
||||
|
||||
@@ -58,6 +58,7 @@ func newWebSearchToolWithBackend(backend searchBackend) Tool {
|
||||
baseTool: baseTool{
|
||||
name: "web_search",
|
||||
description: "Search the web for a query and return ranked results (title, URL, snippet). Complements web_fetch, which retrieves a single known URL.",
|
||||
deferred: true,
|
||||
parameters: Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]PropertySchema{
|
||||
|
||||
@@ -75,6 +75,9 @@ func (tool writeFileTool) RunWithOptions(ctx context.Context, args map[string]an
|
||||
// stale view. Only read current bytes when there is a baseline to compare,
|
||||
// so a first-touch create/overwrite stays a single write with no extra read.
|
||||
if existed {
|
||||
if options.FileTracker != nil && !options.FileTracker.SeenWhole(absolutePath) {
|
||||
return errorResult(fileUnseenMessage(relativePath))
|
||||
}
|
||||
if _, tracked := options.FileTracker.Version(absolutePath); tracked {
|
||||
// Fail CLOSED: if the tracked file can't be re-read to verify it, refuse
|
||||
// the overwrite rather than clobbering a file whose current state is
|
||||
@@ -107,6 +110,7 @@ func (tool writeFileTool) RunWithOptions(ctx context.Context, args map[string]an
|
||||
if err := os.WriteFile(absolutePath, []byte(content), 0o644); err != nil {
|
||||
return errorResult("Error writing file " + relativePath + ": " + err.Error())
|
||||
}
|
||||
modelKnownContent := content
|
||||
// Optional format-on-write (ZERO_FORMAT_ON_WRITE). Must run BEFORE the
|
||||
// FileTracker baseline: recording pre-format content would make the very
|
||||
// next edit look like an external modification and trip the conflict guard.
|
||||
@@ -115,6 +119,9 @@ func (tool writeFileTool) RunWithOptions(ctx context.Context, args map[string]an
|
||||
// session compares against what is now on disk.
|
||||
newInfo, _ := os.Stat(absolutePath)
|
||||
options.FileTracker.Record(absolutePath, []byte(content), newInfo)
|
||||
if content == modelKnownContent {
|
||||
options.FileTracker.RecordSeenRange(absolutePath, 1, lineCount(content), lineCount(content))
|
||||
}
|
||||
if !existed {
|
||||
options.FileTracker.RecordCreated(absolutePath)
|
||||
}
|
||||
|
||||
+6
-2
@@ -83,6 +83,7 @@ func (m model) handleBTWCommand(question string) (model, tea.Cmd) {
|
||||
})
|
||||
}
|
||||
parent.btw = btwState{}
|
||||
parent.goalContinuationsSuspended = true
|
||||
// A scrollback print that was already scheduled may acknowledge after the
|
||||
// side surface is active. The hidden model does not receive that unscoped
|
||||
// acknowledgement, so clear its print latch and rebuild on return.
|
||||
@@ -175,6 +176,7 @@ func (m model) leaveBTW() (model, tea.Cmd) {
|
||||
}
|
||||
m, _ = m.clearLoopsForSessionSwitch()
|
||||
parent := *m.btw.parent
|
||||
parent.goalContinuationsSuspended = false
|
||||
parent.btwRunIDSeq = maxInt(parent.btwRunIDSeq, m.runID)
|
||||
parent.btw = btwState{}
|
||||
// A hidden parent completion may have scheduled an unscoped git-sweep result
|
||||
@@ -196,13 +198,15 @@ func (m model) leaveBTW() (model, tea.Cmd) {
|
||||
text: "Returned from the isolated BTW conversation. Its messages were not added to this session.",
|
||||
})
|
||||
parent.resetFlushFrontier("· returned from btw ·")
|
||||
return parent, batchCommands(sweepCmd, spinnerCmd)
|
||||
var goalCmd tea.Cmd
|
||||
parent, goalCmd = parent.launchGoalContinuationIfReady()
|
||||
return parent, batchCommands(sweepCmd, spinnerCmd, goalCmd)
|
||||
}
|
||||
|
||||
func btwCommandUnavailable(command parsedCommand) bool {
|
||||
arg := strings.ToLower(strings.TrimSpace(command.text))
|
||||
switch command.kind {
|
||||
case commandNew, commandResume, commandRetitle, commandSpec, commandLoop,
|
||||
case commandNew, commandResume, commandRename, commandSpec, commandLoop, commandGoal,
|
||||
commandRewind, commandCompact, commandSTTModel, commandMCP:
|
||||
return true
|
||||
case commandModel:
|
||||
|
||||
Loaded 100 of 130 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user