fix(acp): interrupt idle reads on cancellation

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,17 +98,34 @@ 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),
 		pending:   make(map[int64]chan rpcMessage),
 	}
 }
 
+// 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,24 +147,137 @@ 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
 		}
 	}
 }
 
+// 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) {