diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index 51df655fa..7d9171692 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -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) { diff --git a/internal/acp/jsonrpc_test.go b/internal/acp/jsonrpc_test.go index d9e5d3df3..7d3b0c9da 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -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() diff --git a/internal/cli/acp.go b/internal/cli/acp.go index dfade3b32..c261c19ef 100644 --- a/internal/cli/acp.go +++ b/internal/cli/acp.go @@ -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 diff --git a/internal/cli/acp_test.go b/internal/cli/acp_test.go new file mode 100644 index 000000000..c957432b8 --- /dev/null +++ b/internal/cli/acp_test.go @@ -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) +}