From bc5dbaa1135b14ce6dfe92134fdfd6fe8433f7e5 Mon Sep 17 00:00:00 2001 From: cl-ment Date: Sun, 23 Aug 2026 09:25:58 +0200 Subject: [PATCH 01/14] security(acp): bound frame size and throttle concurrent request handlers (fixes #923) Inbound ACP requests and notifications previously spawned unbound goroutines without rate limiting or concurrency backpressure. This adds maxFrameBytes (64MB) and bounds concurrent in-flight dispatch goroutines via a buffered semaphore (maxConcurrentRequests = 128) in Conn. --- internal/acp/jsonrpc.go | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index 7d9171692..422282da6 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -76,6 +76,11 @@ type HandlerFunc func(ctx context.Context, params json.RawMessage) (any, error) // NotifyFunc handles an inbound notification (no response is sent). type NotifyFunc func(ctx context.Context, params json.RawMessage) +const ( + maxFrameBytes = 64 * 1024 * 1024 + maxConcurrentRequests = 128 +) + // Conn is a JSON-RPC 2.0 peer over a single ndjson stream pair. It both serves // inbound requests/notifications (via registered handlers) and issues outbound // requests/notifications — needed because ACP is bidirectional (the agent calls @@ -95,7 +100,8 @@ type Conn struct { pending map[int64]chan rpcMessage closed bool - wg sync.WaitGroup // tracks in-flight inbound handlers + sem chan struct{} + wg sync.WaitGroup // tracks in-flight inbound handlers } // NewConn builds a peer reading ndjson from r and writing ndjson to w. This @@ -111,6 +117,7 @@ func NewConn(r io.Reader, w io.Writer) *Conn { handlers: make(map[string]HandlerFunc), notifiers: make(map[string]NotifyFunc), pending: make(map[int64]chan rpcMessage), + sem: make(chan struct{}, maxConcurrentRequests), } } @@ -297,14 +304,18 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { c.deliver(msg) case msg.isRequest(): c.wg.Add(1) + c.acquireSem(ctx) go func(m rpcMessage) { + defer c.releaseSem() defer c.wg.Done() c.dispatchRequest(ctx, m) }(msg) case msg.isNotify(): if fn := c.notifiers[msg.Method]; fn != nil { c.wg.Add(1) + c.acquireSem(ctx) go func(m rpcMessage) { + defer c.releaseSem() defer c.wg.Done() fn(ctx, m.Params) }(msg) @@ -317,6 +328,26 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { } } +func (c *Conn) acquireSem(ctx context.Context) { + if c.sem == nil { + return + } + select { + case c.sem <- struct{}{}: + case <-ctx.Done(): + } +} + +func (c *Conn) releaseSem() { + if c.sem == nil { + return + } + select { + case <-c.sem: + default: + } +} + func (c *Conn) dispatchRequest(ctx context.Context, msg rpcMessage) { fn := c.handlers[msg.Method] if fn == nil { From 747f73a6d1dd5de6491864a9610fac8007c632ef Mon Sep 17 00:00:00 2001 From: cl-ment Date: Mon, 24 Aug 2026 12:42:34 +0200 Subject: [PATCH 02/14] fix(acp): unthrottle cancel notifications and isolate request semaphore --- internal/acp/jsonrpc.go | 50 +++++++++++++++--- internal/acp/jsonrpc_test.go | 99 ++++++++++++++++++++++++++++++++++++ 2 files changed, 141 insertions(+), 8 deletions(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index 422282da6..1d9df3aaf 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -77,6 +77,8 @@ type HandlerFunc func(ctx context.Context, params json.RawMessage) (any, error) type NotifyFunc func(ctx context.Context, params json.RawMessage) const ( + // maxFrameBytes limits the maximum size of a single ndjson line including + // the trailing newline delimiter (effective maximum payload is limit - 1). maxFrameBytes = 64 * 1024 * 1024 maxConcurrentRequests = 128 ) @@ -100,6 +102,10 @@ type Conn struct { pending map[int64]chan rpcMessage closed bool + // frameLimit bounds inbound ndjson line size. If zero, maxFrameBytes is used. + // This field has no public setter and is set solely by tests. + frameLimit int64 + sem chan struct{} wg sync.WaitGroup // tracks in-flight inbound handlers } @@ -162,7 +168,7 @@ func (c *Conn) Serve(ctx context.Context) error { // decoder) keeps a single malformed line from making the whole connection // unrecoverable — we report -32700 and continue. // - // generation brackets this SPECIFIC ReadBytes call: bufio may invoke + // generation brackets this SPECIFIC readNDJSONFrame 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 @@ -174,8 +180,12 @@ func (c *Conn) Serve(ctx context.Context) error { // 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. + limit := int64(maxFrameBytes) + if c.frameLimit > 0 { + limit = c.frameLimit + } before := interruptible.generation() - line, err := reader.ReadBytes('\n') + line, err := readNDJSONFrame(reader, limit) interrupted := interruptible.generation() != before if len(bytes.TrimSpace(line)) > 0 { @@ -191,6 +201,28 @@ func (c *Conn) Serve(ctx context.Context) error { } } +// readNDJSONFrame reads a single newline-terminated NDJSON frame from r, +// bounded by limit bytes. The buffer includes the trailing newline delimiter, +// so the maximum payload is limit - 1. If the frame exceeds limit before +// encountering '\n', it returns the accumulated partial line alongside an error. +func readNDJSONFrame(r *bufio.Reader, limit int64) ([]byte, error) { + var buf []byte + for { + chunk, err := r.ReadSlice('\n') + buf = append(buf, chunk...) + if limit > 0 && int64(len(buf)) > limit { + return buf, fmt.Errorf("acp: frame exceeds limit of %d bytes", limit) + } + if err == nil { + return buf, nil + } + if errors.Is(err, bufio.ErrBufferFull) { + continue + } + return buf, 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. @@ -304,18 +336,18 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { c.deliver(msg) case msg.isRequest(): c.wg.Add(1) - c.acquireSem(ctx) go func(m rpcMessage) { - defer c.releaseSem() defer c.wg.Done() + if !c.acquireSem(ctx) { + return + } + defer c.releaseSem() c.dispatchRequest(ctx, m) }(msg) case msg.isNotify(): if fn := c.notifiers[msg.Method]; fn != nil { c.wg.Add(1) - c.acquireSem(ctx) go func(m rpcMessage) { - defer c.releaseSem() defer c.wg.Done() fn(ctx, m.Params) }(msg) @@ -328,13 +360,15 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { } } -func (c *Conn) acquireSem(ctx context.Context) { +func (c *Conn) acquireSem(ctx context.Context) bool { if c.sem == nil { - return + return true } select { case c.sem <- struct{}{}: + return true case <-ctx.Done(): + return false } } diff --git a/internal/acp/jsonrpc_test.go b/internal/acp/jsonrpc_test.go index 7d3b0c9da..6e80d7fd9 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -4,7 +4,9 @@ import ( "context" "encoding/json" "errors" + "fmt" "io" + "strings" "sync" "sync/atomic" "testing" @@ -485,3 +487,100 @@ func asRPCError(err error, target **rpcError) bool { } return ok } + +// TestCancelNotificationDeliveredDuringHandlerSaturation proves that when all 128 +// request handler slots are occupied by blocking work, an inbound notification +// (such as session/cancel) is still delivered and executed immediately rather than +// being dropped, blocked, or starved behind saturated request workers. +func TestCancelNotificationDeliveredDuringHandlerSaturation(t *testing.T) { + clientR, serverW := io.Pipe() + serverR, clientW := io.Pipe() + + server := NewConn(serverR, serverW) + ctx, cancel := context.WithCancel(context.Background()) + defer func() { + cancel() + _ = clientR.Close() + _ = serverW.Close() + _ = clientW.Close() + }() + + releaseHandlers := make(chan struct{}) + handlerStarted := make(chan struct{}, maxConcurrentRequests) + + server.Handle("block", func(ctx context.Context, _ json.RawMessage) (any, error) { + handlerStarted <- struct{}{} + select { + case <-releaseHandlers: + return "ok", nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }) + + cancelDelivered := make(chan struct{}, 1) + server.HandleNotify("session/cancel", func(_ context.Context, _ json.RawMessage) { + cancelDelivered <- struct{}{} + close(releaseHandlers) + }) + + go func() { _ = server.Serve(ctx) }() + + // 1. Fill all 128 request slots + for i := 1; i <= maxConcurrentRequests; i++ { + req := fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":"block"}`+"\n", i) + _, _ = clientW.Write([]byte(req)) + } + + for i := 0; i < maxConcurrentRequests; i++ { + select { + case <-handlerStarted: + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for handler %d to start", i) + } + } + + // 2. Send notification while all request slots are occupied + notify := `{"jsonrpc":"2.0","method":"session/cancel","params":{}}` + "\n" + _, _ = clientW.Write([]byte(notify)) + + // 3. Verify notification is delivered without delay + select { + case <-cancelDelivered: + // Success: notification bypassed request throttling and freed the handlers + case <-time.After(2 * time.Second): + t.Fatal("cancel notification was dropped or blocked by saturated request handlers") + } +} + +// TestReadNDJSONFrameRejectsOversizedTerminatedFrame proves that frames exceeding +// the configured frame limit are rejected with an error rather than buffered unboundedly. +func TestReadNDJSONFrameRejectsOversizedTerminatedFrame(t *testing.T) { + serverR, clientW := io.Pipe() + + server := NewConn(serverR, io.Discard) + server.frameLimit = 64 // 64 bytes frame limit for test injection + server.Handle("ping", func(_ context.Context, _ json.RawMessage) (any, error) { return "pong", nil }) + + ctx, cancel := context.WithCancel(context.Background()) + defer func() { + cancel() + _ = clientW.Close() + }() + + errCh := make(chan error, 1) + go func() { errCh <- server.Serve(ctx) }() + + // Send an oversized frame > 64 bytes terminated by \n + oversized := `{"jsonrpc":"2.0","id":1,"method":"ping","params":{"long_padding":"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"}}` + "\n" + _, _ = clientW.Write([]byte(oversized)) + + select { + case err := <-errCh: + if err == nil || !strings.Contains(err.Error(), "exceeds limit") { + t.Fatalf("expected frame limit error, got: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for Serve to reject oversized frame") + } +} From 9d0e16fc8e46cd008e9129a053b6b1a178f94895 Mon Sep 17 00:00:00 2001 From: cl-ment Date: Mon, 24 Aug 2026 18:32:16 +0200 Subject: [PATCH 03/14] security(acp): bound request admission with immediate -32000 Server Busy error when saturated --- internal/acp/jsonrpc.go | 30 ++++++++++++++++------ internal/acp/jsonrpc_test.go | 48 ++++++++++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 8 deletions(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index 1d9df3aaf..302630f35 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -27,6 +27,7 @@ const ( codeMethodNotFound = -32601 codeInvalidParams = -32602 codeInternalError = -32603 + codeServerBusy = -32000 ) // rpcError is a JSON-RPC 2.0 error object. @@ -335,15 +336,28 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { case msg.isResponse(): c.deliver(msg) case msg.isRequest(): - c.wg.Add(1) - go func(m rpcMessage) { - defer c.wg.Done() - if !c.acquireSem(ctx) { - return + if c.sem != nil { + select { + case c.sem <- struct{}{}: + c.wg.Add(1) + go func(m rpcMessage) { + defer c.wg.Done() + defer c.releaseSem() + c.dispatchRequest(ctx, m) + }(msg) + default: + c.writeError(msg.ID, &rpcError{ + Code: codeServerBusy, + Message: "server busy: max concurrent requests exceeded", + }) } - defer c.releaseSem() - c.dispatchRequest(ctx, m) - }(msg) + } else { + c.wg.Add(1) + go func(m rpcMessage) { + defer c.wg.Done() + c.dispatchRequest(ctx, m) + }(msg) + } case msg.isNotify(): if fn := c.notifiers[msg.Method]; fn != nil { c.wg.Add(1) diff --git a/internal/acp/jsonrpc_test.go b/internal/acp/jsonrpc_test.go index 6e80d7fd9..076ac3bb6 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -1,6 +1,7 @@ package acp import ( + "bufio" "context" "encoding/json" "errors" @@ -426,6 +427,53 @@ func TestConnBidirectionalDuringHandler(t *testing.T) { } } +func TestConnSaturatedRequestsReturnsServerBusy(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + defer inWriter.Close() + defer outReader.Close() + + conn := NewConn(inReader, outWriter) + // Use a small 2-slot semaphore for test + conn.sem = make(chan struct{}, 2) + + handlerBlock := make(chan struct{}) + defer close(handlerBlock) + + conn.Handle("slow", func(ctx context.Context, params json.RawMessage) (any, error) { + <-handlerBlock + return "ok", nil + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { _ = conn.Serve(ctx) }() + + // Send 2 requests that fill semaphore slots + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":"slow","params":{}}` + "\n")) + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":2,"method":"slow","params":{}}` + "\n")) + + // Wait until both slots are occupied + for len(conn.sem) < 2 { + time.Sleep(5 * time.Millisecond) + } + + // Send 3rd request while pool is saturated + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":3,"method":"slow","params":{}}` + "\n")) + + scanner := bufio.NewScanner(outReader) + if !scanner.Scan() { + t.Fatalf("expected response line, got none: %v", scanner.Err()) + } + var resp rpcMessage + if err := json.Unmarshal(scanner.Bytes(), &resp); err != nil { + t.Fatalf("decode response: %v", err) + } + if resp.Error == nil || resp.Error.Code != codeServerBusy { + t.Fatalf("resp.Error = %#v, want code %d (Server Busy)", resp.Error, codeServerBusy) + } +} + // TestConnSurvivesMalformedLine proves a single bad ndjson line yields a -32700 // and does NOT tear down the connection — a following valid request still works. func TestConnSurvivesMalformedLine(t *testing.T) { From 501a70915538d67debe70184d8b7e83d63b755e0 Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Fri, 28 Aug 2026 22:54:01 +0200 Subject: [PATCH 04/14] fix(acp): skip oversized frames and write busy replies off the read loop A frame-limit error no longer reaches handleLine, so a padded valid request cannot run its handler. Saturated -32000 replies are written asynchronously so stdout backpressure cannot stall session/cancel. --- internal/acp/jsonrpc.go | 21 +++++++--- internal/acp/jsonrpc_test.go | 78 ++++++++++++++++++++++++++++++++++-- 2 files changed, 90 insertions(+), 9 deletions(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index 302630f35..f7c3a65a1 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -30,6 +30,8 @@ const ( codeServerBusy = -32000 ) +var errFrameTooLarge = errors.New("acp: frame exceeds limit") + // rpcError is a JSON-RPC 2.0 error object. type rpcError struct { Code int `json:"code"` @@ -189,6 +191,10 @@ func (c *Conn) Serve(ctx context.Context) error { line, err := readNDJSONFrame(reader, limit) interrupted := interruptible.generation() != before + if errors.Is(err, errFrameTooLarge) { + c.failAllPending(err) + return err + } if len(bytes.TrimSpace(line)) > 0 { c.handleLine(ctx, line) } @@ -212,7 +218,7 @@ func readNDJSONFrame(r *bufio.Reader, limit int64) ([]byte, error) { chunk, err := r.ReadSlice('\n') buf = append(buf, chunk...) if limit > 0 && int64(len(buf)) > limit { - return buf, fmt.Errorf("acp: frame exceeds limit of %d bytes", limit) + return buf, fmt.Errorf("%w of %d bytes", errFrameTooLarge, limit) } if err == nil { return buf, nil @@ -346,10 +352,15 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { c.dispatchRequest(ctx, m) }(msg) default: - c.writeError(msg.ID, &rpcError{ - Code: codeServerBusy, - Message: "server busy: max concurrent requests exceeded", - }) + id := append(json.RawMessage(nil), msg.ID...) + c.wg.Add(1) + go func() { + defer c.wg.Done() + c.writeError(id, &rpcError{ + Code: codeServerBusy, + Message: "server busy: max concurrent requests exceeded", + }) + }() } } else { c.wg.Add(1) diff --git a/internal/acp/jsonrpc_test.go b/internal/acp/jsonrpc_test.go index 076ac3bb6..600d7b916 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -608,7 +608,11 @@ func TestReadNDJSONFrameRejectsOversizedTerminatedFrame(t *testing.T) { server := NewConn(serverR, io.Discard) server.frameLimit = 64 // 64 bytes frame limit for test injection - server.Handle("ping", func(_ context.Context, _ json.RawMessage) (any, error) { return "pong", nil }) + var invoked atomic.Bool + server.Handle("ping", func(_ context.Context, _ json.RawMessage) (any, error) { + invoked.Store(true) + return "pong", nil + }) ctx, cancel := context.WithCancel(context.Background()) defer func() { @@ -619,16 +623,82 @@ func TestReadNDJSONFrameRejectsOversizedTerminatedFrame(t *testing.T) { errCh := make(chan error, 1) go func() { errCh <- server.Serve(ctx) }() - // Send an oversized frame > 64 bytes terminated by \n - oversized := `{"jsonrpc":"2.0","id":1,"method":"ping","params":{"long_padding":"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"}}` + "\n" + // Valid request padded with JSON whitespace past the limit: dispatch would + // run ping unless the frame-limit error is checked before handleLine. + oversized := `{"jsonrpc":"2.0","id":1,"method":"ping"}` + strings.Repeat(" ", 64) + "\n" _, _ = clientW.Write([]byte(oversized)) select { case err := <-errCh: - if err == nil || !strings.Contains(err.Error(), "exceeds limit") { + if err == nil || !errors.Is(err, errFrameTooLarge) { t.Fatalf("expected frame limit error, got: %v", err) } case <-time.After(2 * time.Second): t.Fatal("timed out waiting for Serve to reject oversized frame") } + if invoked.Load() { + t.Fatal("handler invoked for an oversized frame") + } +} + +type stallWriter struct { + started chan struct{} + gate chan struct{} + once sync.Once +} + +func (w *stallWriter) Write(p []byte) (int, error) { + w.once.Do(func() { close(w.started) }) + <-w.gate + return len(p), nil +} + +func TestSaturatedBusyWriteDoesNotBlockCancelNotification(t *testing.T) { + inReader, inWriter := io.Pipe() + defer inWriter.Close() + + started := make(chan struct{}) + gate := make(chan struct{}) + defer close(gate) + conn := NewConn(inReader, &stallWriter{started: started, gate: gate}) + conn.sem = make(chan struct{}, 2) + + handlerStarted := make(chan struct{}, 2) + conn.Handle("slow", func(ctx context.Context, _ json.RawMessage) (any, error) { + handlerStarted <- struct{}{} + <-ctx.Done() + return nil, ctx.Err() + }) + cancelSeen := make(chan struct{}) + conn.HandleNotify("session/cancel", func(_ context.Context, _ json.RawMessage) { + close(cancelSeen) + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { _ = conn.Serve(ctx) }() + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":"slow"}` + "\n")) + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":2,"method":"slow"}` + "\n")) + for i := 0; i < 2; i++ { + select { + case <-handlerStarted: + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for slow handler %d", i) + } + } + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":3,"method":"slow"}` + "\n")) + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for busy write to start") + } + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","method":"session/cancel","params":{}}` + "\n")) + select { + case <-cancelSeen: + case <-time.After(2 * time.Second): + t.Fatal("cancel notification stalled behind a blocked busy write") + } } From 22fcb46df0711c6872bab3a47e989fba13063db1 Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Fri, 28 Aug 2026 22:54:59 +0200 Subject: [PATCH 05/14] fix(acp): drop unused acquireSem after non-blocking busy admission --- internal/acp/jsonrpc.go | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index f7c3a65a1..ae6993d7e 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -385,18 +385,6 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { } } -func (c *Conn) acquireSem(ctx context.Context) bool { - if c.sem == nil { - return true - } - select { - case c.sem <- struct{}{}: - return true - case <-ctx.Done(): - return false - } -} - func (c *Conn) releaseSem() { if c.sem == nil { return From c0e48c79dd33a8fbe1587f6f208575645e9abe6e Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sat, 29 Aug 2026 00:17:40 +0200 Subject: [PATCH 06/14] fix(acp): bound busy replies with a single writer and drop the session on overflow Rejected requests share one buffered busy queue and one writer goroutine. When that queue is full the session is cancelled instead of spawning another writer. Admitted handlers skip further writes once overloaded so Serve can exit while stdout is stalled. --- internal/acp/jsonrpc.go | 70 ++++++++++++++++++++++++++++++------ internal/acp/jsonrpc_test.go | 60 +++++++++++++++++++++++++++++++ 2 files changed, 120 insertions(+), 10 deletions(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index ae6993d7e..079a93b52 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -30,7 +30,10 @@ const ( codeServerBusy = -32000 ) -var errFrameTooLarge = errors.New("acp: frame exceeds limit") +var ( + errFrameTooLarge = errors.New("acp: frame exceeds limit") + errBusyOverload = errors.New("acp: busy-reply queue full") +) // rpcError is a JSON-RPC 2.0 error object. type rpcError struct { @@ -84,6 +87,7 @@ const ( // the trailing newline delimiter (effective maximum payload is limit - 1). maxFrameBytes = 64 * 1024 * 1024 maxConcurrentRequests = 128 + maxBusyReplies = 1 ) // Conn is a JSON-RPC 2.0 peer over a single ndjson stream pair. It both serves @@ -111,6 +115,10 @@ type Conn struct { sem chan struct{} wg sync.WaitGroup // tracks in-flight inbound handlers + + busyCh chan json.RawMessage + serveCancel context.CancelFunc + overloaded atomic.Bool } // NewConn builds a peer reading ndjson from r and writing ndjson to w. This @@ -127,6 +135,7 @@ func NewConn(r io.Reader, w io.Writer) *Conn { notifiers: make(map[string]NotifyFunc), pending: make(map[int64]chan rpcMessage), sem: make(chan struct{}, maxConcurrentRequests), + busyCh: make(chan json.RawMessage, maxBusyReplies), } } @@ -154,10 +163,13 @@ func (c *Conn) HandleNotify(method string, fn NotifyFunc) { c.notifiers[method] // blocks the loop from delivering session/cancel or a permission response. func (c *Conn) Serve(ctx context.Context) error { ctx, cancel := context.WithCancel(ctx) + c.serveCancel = cancel + go c.writeBusyLoop(ctx) // On exit, cancel in-flight handlers (so a blocked outbound Call unblocks via // ctx) and then wait for them to finish writing their responses. Without this, // a finite input stream (e.g. piped ndjson that EOFs right after a request) - // would race the dispatch goroutine and drop the response. + // would race the dispatch goroutine and drop the response. The busy-reply + // writer is not on wg: it may be blocked in Write against a stalled peer. defer func() { cancel() c.wg.Wait() @@ -198,6 +210,9 @@ func (c *Conn) Serve(ctx context.Context) error { if len(bytes.TrimSpace(line)) > 0 { c.handleLine(ctx, line) } + if c.overloaded.Load() { + return errBusyOverload + } if err != nil { c.failAllPending(err) if errors.Is(err, io.EOF) || interrupted { @@ -353,14 +368,9 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { }(msg) default: id := append(json.RawMessage(nil), msg.ID...) - c.wg.Add(1) - go func() { - defer c.wg.Done() - c.writeError(id, &rpcError{ - Code: codeServerBusy, - Message: "server busy: max concurrent requests exceeded", - }) - }() + if !c.tryEnqueueBusy(id) { + c.tripOverload() + } } } else { c.wg.Add(1) @@ -515,7 +525,44 @@ func (c *Conn) writeError(id json.RawMessage, e *rpcError) { _ = c.write(rpcMessage{JSONRPC: "2.0", ID: id, Error: e}) } +func (c *Conn) tryEnqueueBusy(id json.RawMessage) bool { + if c.busyCh == nil { + return false + } + select { + case c.busyCh <- id: + return true + default: + return false + } +} + +func (c *Conn) tripOverload() { + if !c.overloaded.CompareAndSwap(false, true) { + return + } + c.failAllPending(errBusyOverload) + if c.serveCancel != nil { + c.serveCancel() + } +} + +func (c *Conn) writeBusyLoop(ctx context.Context) { + busy := &rpcError{Code: codeServerBusy, Message: "server busy: max concurrent requests exceeded"} + for { + select { + case <-ctx.Done(): + return + case id := <-c.busyCh: + c.writeError(id, busy) + } + } +} + func (c *Conn) write(msg rpcMessage) error { + if c.overloaded.Load() { + return errBusyOverload + } msg.JSONRPC = "2.0" data, err := json.Marshal(msg) if err != nil { @@ -523,6 +570,9 @@ func (c *Conn) write(msg rpcMessage) error { } c.writeMu.Lock() defer c.writeMu.Unlock() + if c.overloaded.Load() { + return errBusyOverload + } if _, err := c.w.Write(append(data, '\n')); err != nil { return err } diff --git a/internal/acp/jsonrpc_test.go b/internal/acp/jsonrpc_test.go index 600d7b916..86da790ef 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "io" + "runtime" "strings" "sync" "sync/atomic" @@ -702,3 +703,62 @@ func TestSaturatedBusyWriteDoesNotBlockCancelNotification(t *testing.T) { t.Fatal("cancel notification stalled behind a blocked busy write") } } + +func TestStalledBusyRepliesStayBoundedAndServeExits(t *testing.T) { + inReader, inWriter := io.Pipe() + defer inWriter.Close() + + started := make(chan struct{}) + gate := make(chan struct{}) + defer close(gate) + conn := NewConn(inReader, &stallWriter{started: started, gate: gate}) + conn.sem = make(chan struct{}, 2) + + handlerStarted := make(chan struct{}, 2) + conn.Handle("slow", func(ctx context.Context, _ json.RawMessage) (any, error) { + handlerStarted <- struct{}{} + <-ctx.Done() + return nil, ctx.Err() + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan error, 1) + go func() { done <- conn.Serve(ctx) }() + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":"slow"}` + "\n")) + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":2,"method":"slow"}` + "\n")) + for i := 0; i < 2; i++ { + select { + case <-handlerStarted: + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for slow handler %d", i) + } + } + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":3,"method":"slow"}` + "\n")) + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for busy write to stall") + } + + before := runtime.NumGoroutine() + const extra = 50 + go func() { + for i := 0; i < extra; i++ { + req := fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":"slow"}`+"\n", i+10) + _, _ = inWriter.Write([]byte(req)) + } + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("Serve did not exit after the busy-reply queue overflowed") + } + after := runtime.NumGoroutine() + if delta := after - before; delta > 8 { + t.Fatalf("goroutine growth = %d after %d rejected requests, want bounded", delta, extra) + } +} From c97ce9691b21b9c542458f048f9624f1612d5546 Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sat, 29 Aug 2026 01:17:08 +0200 Subject: [PATCH 07/14] fix(acp): let Call drop writeMu wait when its context is cancelled A stalled peer can hold the serialized writer. The second Call now returns context.Canceled instead of blocking behind that lock. --- internal/acp/jsonrpc.go | 70 +++++++++++++++++++++++++++--------- internal/acp/jsonrpc_test.go | 35 ++++++++++++++++++ 2 files changed, 88 insertions(+), 17 deletions(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index 079a93b52..b123078de 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -344,12 +344,12 @@ func (r *interruptibleReader) Read(p []byte) (int, error) { func (c *Conn) handleLine(ctx context.Context, line []byte) { var msg rpcMessage if err := json.Unmarshal(bytes.TrimSpace(line), &msg); err != nil { - c.writeError(json.RawMessage("null"), &rpcError{Code: codeParseError, Message: "parse error"}) + c.writeError(ctx, json.RawMessage("null"), &rpcError{Code: codeParseError, Message: "parse error"}) return } if msg.JSONRPC != "" && msg.JSONRPC != "2.0" { if len(msg.ID) > 0 { - c.writeError(msg.ID, &rpcError{Code: codeInvalidRequest, Message: "unsupported jsonrpc version"}) + c.writeError(ctx, msg.ID, &rpcError{Code: codeInvalidRequest, Message: "unsupported jsonrpc version"}) } return } @@ -390,7 +390,7 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { default: // Malformed frame; reply only if we can identify a request id. if len(msg.ID) > 0 { - c.writeError(msg.ID, &rpcError{Code: codeInvalidRequest, Message: "invalid request"}) + c.writeError(ctx, msg.ID, &rpcError{Code: codeInvalidRequest, Message: "invalid request"}) } } } @@ -408,20 +408,20 @@ func (c *Conn) releaseSem() { func (c *Conn) dispatchRequest(ctx context.Context, msg rpcMessage) { fn := c.handlers[msg.Method] if fn == nil { - c.writeError(msg.ID, &rpcError{Code: codeMethodNotFound, Message: "method not found: " + msg.Method}) + c.writeError(ctx, msg.ID, &rpcError{Code: codeMethodNotFound, Message: "method not found: " + msg.Method}) return } result, err := fn(ctx, msg.Params) if err != nil { var re *rpcError if errors.As(err, &re) { - c.writeError(msg.ID, re) + c.writeError(ctx, msg.ID, re) } else { - c.writeError(msg.ID, &rpcError{Code: codeInternalError, Message: err.Error()}) + c.writeError(ctx, msg.ID, &rpcError{Code: codeInternalError, Message: err.Error()}) } return } - c.writeResult(msg.ID, result) + c.writeResult(ctx, msg.ID, result) } // Call issues an outbound request and blocks until the response arrives, ctx is @@ -450,7 +450,7 @@ func (c *Conn) Call(ctx context.Context, method string, params any, result any) }() idRaw, _ := json.Marshal(id) - if err := c.write(rpcMessage{JSONRPC: "2.0", ID: idRaw, Method: method, Params: raw}); err != nil { + if err := c.write(ctx, rpcMessage{JSONRPC: "2.0", ID: idRaw, Method: method, Params: raw}); err != nil { return err } @@ -474,7 +474,7 @@ func (c *Conn) Notify(method string, params any) error { if err != nil { return err } - return c.write(rpcMessage{JSONRPC: "2.0", Method: method, Params: raw}) + return c.write(context.Background(), rpcMessage{JSONRPC: "2.0", Method: method, Params: raw}) } func (c *Conn) deliver(msg rpcMessage) { @@ -512,17 +512,17 @@ func (c *Conn) failAllPending(err error) { } } -func (c *Conn) writeResult(id json.RawMessage, result any) { +func (c *Conn) writeResult(ctx context.Context, id json.RawMessage, result any) { raw, err := json.Marshal(result) if err != nil { - c.writeError(id, &rpcError{Code: codeInternalError, Message: err.Error()}) + c.writeError(ctx, id, &rpcError{Code: codeInternalError, Message: err.Error()}) return } - _ = c.write(rpcMessage{JSONRPC: "2.0", ID: id, Result: raw}) + _ = c.write(ctx, rpcMessage{JSONRPC: "2.0", ID: id, Result: raw}) } -func (c *Conn) writeError(id json.RawMessage, e *rpcError) { - _ = c.write(rpcMessage{JSONRPC: "2.0", ID: id, Error: e}) +func (c *Conn) writeError(ctx context.Context, id json.RawMessage, e *rpcError) { + _ = c.write(ctx, rpcMessage{JSONRPC: "2.0", ID: id, Error: e}) } func (c *Conn) tryEnqueueBusy(id json.RawMessage) bool { @@ -554,12 +554,46 @@ func (c *Conn) writeBusyLoop(ctx context.Context) { case <-ctx.Done(): return case id := <-c.busyCh: - c.writeError(id, busy) + c.writeError(ctx, id, busy) } } } -func (c *Conn) write(msg rpcMessage) error { +func (c *Conn) acquireWrite(ctx context.Context) error { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return err + } + done := make(chan struct{}) + go func() { + c.writeMu.Lock() + close(done) + }() + select { + case <-done: + if err := ctx.Err(); err != nil { + c.writeMu.Unlock() + return err + } + return nil + case <-ctx.Done(): + go func() { + <-done + c.writeMu.Unlock() + }() + return ctx.Err() + } +} + +func (c *Conn) write(ctx context.Context, msg rpcMessage) error { + if ctx == nil { + ctx = context.Background() + } + if err := ctx.Err(); err != nil { + return err + } if c.overloaded.Load() { return errBusyOverload } @@ -568,7 +602,9 @@ func (c *Conn) write(msg rpcMessage) error { if err != nil { return err } - c.writeMu.Lock() + if err := c.acquireWrite(ctx); err != nil { + return err + } defer c.writeMu.Unlock() if c.overloaded.Load() { return errBusyOverload diff --git a/internal/acp/jsonrpc_test.go b/internal/acp/jsonrpc_test.go index 86da790ef..5037f3b83 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -762,3 +762,38 @@ func TestStalledBusyRepliesStayBoundedAndServeExits(t *testing.T) { t.Fatalf("goroutine growth = %d after %d rejected requests, want bounded", delta, extra) } } + +func TestCallCancelsWhileWriterHoldsWriteMu(t *testing.T) { + inReader, inWriter := io.Pipe() + defer inWriter.Close() + + started := make(chan struct{}) + gate := make(chan struct{}) + defer close(gate) + conn := NewConn(inReader, &stallWriter{started: started, gate: gate}) + + serveCtx, serveCancel := context.WithCancel(context.Background()) + defer serveCancel() + go func() { _ = conn.Serve(serveCtx) }() + + go func() { _ = conn.Call(context.Background(), "ping", nil, nil) }() + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("first write did not stall") + } + + callCtx, callCancel := context.WithCancel(context.Background()) + errCh := make(chan error, 1) + go func() { errCh <- conn.Call(callCtx, "ping", nil, nil) }() + time.Sleep(50 * time.Millisecond) + callCancel() + select { + case err := <-errCh: + if !errors.Is(err, context.Canceled) { + t.Fatalf("Call returned %v, want context.Canceled", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Call did not return after context cancel while writeMu was held") + } +} From f1a4f5f385595bcbe56167ce04a046c3597cc020 Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sat, 29 Aug 2026 01:42:29 +0200 Subject: [PATCH 08/14] fix(acp): keep in-flight replies after Serve reaches EOF Handler work still sees the cancelled Serve context. The response write uses an independent context so a request that finishes during shutdown is not dropped. write still honors the overload trip. --- internal/acp/jsonrpc.go | 9 +++--- internal/acp/jsonrpc_test.go | 58 ++++++++++++++++++++++++++++++++++-- 2 files changed, 60 insertions(+), 7 deletions(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index b123078de..071a72a89 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -407,21 +407,22 @@ func (c *Conn) releaseSem() { func (c *Conn) dispatchRequest(ctx context.Context, msg rpcMessage) { fn := c.handlers[msg.Method] + writeCtx := context.Background() if fn == nil { - c.writeError(ctx, msg.ID, &rpcError{Code: codeMethodNotFound, Message: "method not found: " + msg.Method}) + c.writeError(writeCtx, msg.ID, &rpcError{Code: codeMethodNotFound, Message: "method not found: " + msg.Method}) return } result, err := fn(ctx, msg.Params) if err != nil { var re *rpcError if errors.As(err, &re) { - c.writeError(ctx, msg.ID, re) + c.writeError(writeCtx, msg.ID, re) } else { - c.writeError(ctx, msg.ID, &rpcError{Code: codeInternalError, Message: err.Error()}) + c.writeError(writeCtx, msg.ID, &rpcError{Code: codeInternalError, Message: err.Error()}) } return } - c.writeResult(ctx, msg.ID, result) + c.writeResult(writeCtx, msg.ID, result) } // Call issues an outbound request and blocks until the response arrives, ctx is diff --git a/internal/acp/jsonrpc_test.go b/internal/acp/jsonrpc_test.go index 5037f3b83..b6aaa3719 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -454,8 +454,11 @@ func TestConnSaturatedRequestsReturnsServerBusy(t *testing.T) { _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":"slow","params":{}}` + "\n")) _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":2,"method":"slow","params":{}}` + "\n")) - // Wait until both slots are occupied + deadline := time.Now().Add(2 * time.Second) for len(conn.sem) < 2 { + if time.Now().After(deadline) { + t.Fatalf("timed out waiting for semaphore saturation, len = %d", len(conn.sem)) + } time.Sleep(5 * time.Millisecond) } @@ -757,12 +760,61 @@ func TestStalledBusyRepliesStayBoundedAndServeExits(t *testing.T) { case <-time.After(2 * time.Second): t.Fatal("Serve did not exit after the busy-reply queue overflowed") } - after := runtime.NumGoroutine() - if delta := after - before; delta > 8 { + settle := time.Now().Add(2 * time.Second) + delta := runtime.NumGoroutine() - before + for delta > 8 && time.Now().Before(settle) { + time.Sleep(10 * time.Millisecond) + delta = runtime.NumGoroutine() - before + } + if delta > 8 { t.Fatalf("goroutine growth = %d after %d rejected requests, want bounded", delta, extra) } } +func TestServeEOFStillWritesInFlightResponse(t *testing.T) { + inReader, inWriter := io.Pipe() + outReader, outWriter := io.Pipe() + defer outWriter.Close() + + conn := NewConn(inReader, outWriter) + started := make(chan struct{}) + conn.Handle("echo", func(ctx context.Context, _ json.RawMessage) (any, error) { + close(started) + <-ctx.Done() + return map[string]string{"ok": "yes"}, nil + }) + + errCh := make(chan error, 1) + go func() { errCh <- conn.Serve(context.Background()) }() + if _, err := inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":"echo"}` + "\n")); err != nil { + t.Fatal(err) + } + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("handler did not start") + } + if err := inWriter.Close(); err != nil { + t.Fatal(err) + } + + got := make(chan string, 1) + go func() { + scanner := bufio.NewScanner(outReader) + if scanner.Scan() { + got <- scanner.Text() + } + }() + select { + case line := <-got: + if !strings.Contains(line, `"ok":"yes"`) && !strings.Contains(line, `"ok": "yes"`) { + t.Fatalf("missing in-flight response: %s", line) + } + case <-time.After(2 * time.Second): + t.Fatal("Serve dropped the in-flight response at EOF") + } +} + func TestCallCancelsWhileWriterHoldsWriteMu(t *testing.T) { inReader, inWriter := io.Pipe() defer inWriter.Close() From 30beba725977eab57a1a6f613264279a55957ac1 Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sun, 30 Aug 2026 17:01:52 +0200 Subject: [PATCH 09/14] fix(acp): cancel stalled writers and bound notification dispatch --- internal/acp/jsonrpc.go | 142 ++++++++++++++++++---- internal/acp/jsonrpc_test.go | 228 +++++++++++++++++++++++++++++++++++ 2 files changed, 345 insertions(+), 25 deletions(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index 071a72a89..7ee141750 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -116,9 +116,15 @@ type Conn struct { sem chan struct{} wg sync.WaitGroup // tracks in-flight inbound handlers - busyCh chan json.RawMessage - serveCancel context.CancelFunc - overloaded atomic.Bool + busyCh chan json.RawMessage + serveCancel context.CancelFunc + overloaded atomic.Bool + writeAbort chan struct{} + writeWaiters atomic.Int32 + + notifyMu sync.Mutex + notifyOn map[string]bool + notifyQ map[string]json.RawMessage } // NewConn builds a peer reading ndjson from r and writing ndjson to w. This @@ -129,13 +135,16 @@ type Conn struct { // lifetime. func NewConn(r io.Reader, w io.Writer) *Conn { return &Conn{ - rawReader: r, - w: w, - handlers: make(map[string]HandlerFunc), - notifiers: make(map[string]NotifyFunc), - pending: make(map[int64]chan rpcMessage), - sem: make(chan struct{}, maxConcurrentRequests), - busyCh: make(chan json.RawMessage, maxBusyReplies), + rawReader: r, + w: w, + handlers: make(map[string]HandlerFunc), + notifiers: make(map[string]NotifyFunc), + pending: make(map[int64]chan rpcMessage), + sem: make(chan struct{}, maxConcurrentRequests), + busyCh: make(chan json.RawMessage, maxBusyReplies), + writeAbort: make(chan struct{}), + notifyOn: make(map[string]bool), + notifyQ: make(map[string]json.RawMessage), } } @@ -380,12 +389,8 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { }(msg) } case msg.isNotify(): - if fn := c.notifiers[msg.Method]; fn != nil { - c.wg.Add(1) - go func(m rpcMessage) { - defer c.wg.Done() - fn(ctx, m.Params) - }(msg) + if !c.overloaded.Load() { + c.dispatchNotify(ctx, msg) } default: // Malformed frame; reply only if we can identify a request id. @@ -395,6 +400,41 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { } } +func (c *Conn) dispatchNotify(ctx context.Context, msg rpcMessage) { + fn := c.notifiers[msg.Method] + if fn == nil { + return + } + params := append(json.RawMessage(nil), msg.Params...) + c.notifyMu.Lock() + if c.notifyOn[msg.Method] { + c.notifyQ[msg.Method] = params + c.notifyMu.Unlock() + return + } + c.notifyOn[msg.Method] = true + c.notifyMu.Unlock() + c.wg.Add(1) + go c.runNotify(ctx, msg.Method, fn, params) +} + +func (c *Conn) runNotify(ctx context.Context, method string, fn NotifyFunc, params json.RawMessage) { + defer c.wg.Done() + for { + fn(ctx, params) + c.notifyMu.Lock() + next, ok := c.notifyQ[method] + if !ok { + delete(c.notifyOn, method) + c.notifyMu.Unlock() + return + } + delete(c.notifyQ, method) + c.notifyMu.Unlock() + params = next + } +} + func (c *Conn) releaseSem() { if c.sem == nil { return @@ -543,6 +583,9 @@ func (c *Conn) tripOverload() { return } c.failAllPending(errBusyOverload) + if c.writeAbort != nil { + close(c.writeAbort) + } if c.serveCancel != nil { c.serveCancel() } @@ -553,49 +596,98 @@ func (c *Conn) writeBusyLoop(ctx context.Context) { for { select { case <-ctx.Done(): + c.flushBusy(busy) return case id := <-c.busyCh: - c.writeError(ctx, id, busy) + c.writeBusy(id, busy) } } } -func (c *Conn) acquireWrite(ctx context.Context) error { +func (c *Conn) flushBusy(busy *rpcError) { + for { + select { + case id := <-c.busyCh: + c.writeBusy(id, busy) + default: + return + } + } +} + +func (c *Conn) writeBusy(id json.RawMessage, busy *rpcError) { + _ = c.writeMsg(context.Background(), rpcMessage{JSONRPC: "2.0", ID: id, Error: busy}, true) +} + +func (c *Conn) lockWrite(ctx context.Context, persist bool) error { if ctx == nil { ctx = context.Background() } if err := ctx.Err(); err != nil { return err } + if !persist && c.overloaded.Load() { + return errBusyOverload + } done := make(chan struct{}) + c.writeWaiters.Add(1) go func() { c.writeMu.Lock() close(done) }() + releaseWait := func() { c.writeWaiters.Add(-1) } + abandon := func() { + go func() { + <-done + c.writeMu.Unlock() + }() + } + if persist { + select { + case <-done: + releaseWait() + return nil + case <-ctx.Done(): + releaseWait() + abandon() + return ctx.Err() + } + } select { case <-done: + releaseWait() if err := ctx.Err(); err != nil { c.writeMu.Unlock() return err } + if c.overloaded.Load() { + c.writeMu.Unlock() + return errBusyOverload + } return nil case <-ctx.Done(): - go func() { - <-done - c.writeMu.Unlock() - }() + releaseWait() + abandon() return ctx.Err() + case <-c.writeAbort: + releaseWait() + abandon() + return errBusyOverload } } func (c *Conn) write(ctx context.Context, msg rpcMessage) error { + return c.writeMsg(ctx, msg, false) +} + +func (c *Conn) writeMsg(ctx context.Context, msg rpcMessage, persist bool) error { if ctx == nil { ctx = context.Background() } if err := ctx.Err(); err != nil { return err } - if c.overloaded.Load() { + if !persist && c.overloaded.Load() { return errBusyOverload } msg.JSONRPC = "2.0" @@ -603,11 +695,11 @@ func (c *Conn) write(ctx context.Context, msg rpcMessage) error { if err != nil { return err } - if err := c.acquireWrite(ctx); err != nil { + if err := c.lockWrite(ctx, persist); err != nil { return err } defer c.writeMu.Unlock() - if c.overloaded.Load() { + if !persist && c.overloaded.Load() { return errBusyOverload } if _, err := c.w.Write(append(data, '\n')); err != nil { diff --git a/internal/acp/jsonrpc_test.go b/internal/acp/jsonrpc_test.go index b6aaa3719..5da118fc8 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -849,3 +849,231 @@ func TestCallCancelsWhileWriterHoldsWriteMu(t *testing.T) { t.Fatal("Call did not return after context cancel while writeMu was held") } } + +func TestOverloadUnblocksHandlerWaitingOnStalledWriter(t *testing.T) { + inReader, inWriter := io.Pipe() + defer inWriter.Close() + + started := make(chan struct{}) + gate := make(chan struct{}) + defer close(gate) + conn := NewConn(inReader, &stallWriter{started: started, gate: gate}) + conn.sem = make(chan struct{}, 1) + + handlerStarted := make(chan struct{}) + release := make(chan struct{}) + conn.Handle("slow", func(ctx context.Context, _ json.RawMessage) (any, error) { + close(handlerStarted) + <-release + return "ok", nil + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan error, 1) + go func() { done <- conn.Serve(ctx) }() + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":"slow"}` + "\n")) + select { + case <-handlerStarted: + case <-time.After(2 * time.Second): + t.Fatal("admitted handler did not start") + } + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":2,"method":"slow"}` + "\n")) + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for busy write to stall") + } + + close(release) + deadline := time.Now().Add(2 * time.Second) + for conn.writeWaiters.Load() < 1 { + if time.Now().After(deadline) { + t.Fatal("admitted handler did not block behind the stalled writer") + } + time.Sleep(5 * time.Millisecond) + } + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":3,"method":"slow"}` + "\n")) + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":4,"method":"slow"}` + "\n")) + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("Serve did not exit on overload while a handler waited behind the stalled writer") + } +} + +func TestNotificationFloodStaysBoundedWhileSaturatedRequestStillCancels(t *testing.T) { + inReader, inWriter := io.Pipe() + defer inWriter.Close() + + conn := NewConn(inReader, io.Discard) + conn.sem = make(chan struct{}, 1) + + started := make(chan struct{}) + unblocked := make(chan struct{}) + conn.Handle("slow", func(ctx context.Context, _ json.RawMessage) (any, error) { + close(started) + select { + case <-unblocked: + return "ok", nil + case <-ctx.Done(): + return nil, ctx.Err() + } + }) + + entered := make(chan struct{}, 1) + hold := make(chan struct{}) + var inFlight atomic.Int64 + conn.HandleNotify("session/cancel", func(_ context.Context, _ json.RawMessage) { + inFlight.Add(1) + select { + case entered <- struct{}{}: + default: + } + <-hold + select { + case <-unblocked: + default: + close(unblocked) + } + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { _ = conn.Serve(ctx) }() + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":"slow"}` + "\n")) + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("saturated request did not start") + } + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","method":"session/cancel","params":{}}` + "\n")) + select { + case <-entered: + case <-time.After(2 * time.Second): + t.Fatal("cancel notification was not delivered while the request slot was full") + } + + before := runtime.NumGoroutine() + const flood = 200 + for i := 0; i < flood; i++ { + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","method":"session/cancel","params":{}}` + "\n")) + } + + settle := time.Now().Add(2 * time.Second) + var delta int + for { + delta = runtime.NumGoroutine() - before + got := inFlight.Load() + if got <= 2 && delta <= 8 { + break + } + if time.Now().After(settle) { + t.Fatalf("notification flood in-flight=%d goroutine growth=%d, want bounded", got, delta) + } + time.Sleep(10 * time.Millisecond) + } + + close(hold) + select { + case <-unblocked: + case <-time.After(2 * time.Second): + t.Fatal("saturated request was not cancelled after coalesced session/cancel") + } +} + +type gatedRecorder struct { + started chan struct{} + gate chan struct{} + once sync.Once + mu sync.Mutex + got string +} + +func (w *gatedRecorder) Write(p []byte) (int, error) { + w.once.Do(func() { close(w.started) }) + <-w.gate + w.mu.Lock() + defer w.mu.Unlock() + w.got += string(p) + return len(p), nil +} + +func (w *gatedRecorder) String() string { + w.mu.Lock() + defer w.mu.Unlock() + return w.got +} + +func TestQueuedBusyReplySurvivesOverloadBurst(t *testing.T) { + inReader, inWriter := io.Pipe() + defer inWriter.Close() + + started := make(chan struct{}) + gate := make(chan struct{}) + rec := &gatedRecorder{started: started, gate: gate} + conn := NewConn(inReader, rec) + conn.sem = make(chan struct{}, 1) + + handlerStarted := make(chan struct{}) + conn.Handle("slow", func(ctx context.Context, _ json.RawMessage) (any, error) { + close(handlerStarted) + <-ctx.Done() + return nil, ctx.Err() + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan error, 1) + go func() { done <- conn.Serve(ctx) }() + + go func() { _ = conn.Notify("hold", nil) }() + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for writer to stall") + } + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":"slow"}` + "\n")) + select { + case <-handlerStarted: + case <-time.After(2 * time.Second): + t.Fatal("handler did not start") + } + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":2,"method":"slow"}` + "\n")) + deadline := time.Now().Add(2 * time.Second) + for conn.writeWaiters.Load() < 1 { + if time.Now().After(deadline) { + t.Fatal("busy writer did not wait for writeMu") + } + time.Sleep(5 * time.Millisecond) + } + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":3,"method":"slow"}` + "\n")) + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":4,"method":"slow"}` + "\n")) + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("Serve did not exit on overload") + } + + close(gate) + waitUntil := time.Now().Add(2 * time.Second) + var got string + for time.Now().Before(waitUntil) { + got = rec.String() + hasQueued := strings.Contains(got, `"id":3`) && strings.Contains(got, "-32000") + hasInflight := strings.Contains(got, `"id":2`) && strings.Contains(got, "-32000") + if hasQueued && hasInflight { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatalf("busy ids 2 (in-flight) and 3 (queued) must receive -32000; wrote %q", got) +} From 8d021507910e532475c66e8f4a2bd4f9d5c4afa7 Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sun, 30 Aug 2026 19:06:10 +0200 Subject: [PATCH 10/14] fix(acp): key cancels by session, bound notify workers, drain Serve on overload --- internal/acp/agent_test.go | 156 ++++++++++++++++++++++++++++++++++- internal/acp/jsonrpc.go | 69 ++++++++++++---- internal/acp/jsonrpc_test.go | 154 ++++++++++++++++++++++++++++++++++ 3 files changed, 360 insertions(+), 19 deletions(-) diff --git a/internal/acp/agent_test.go b/internal/acp/agent_test.go index 4fa97a258..57701845e 100644 --- a/internal/acp/agent_test.go +++ b/internal/acp/agent_test.go @@ -9,6 +9,7 @@ import ( "os" "path/filepath" "strings" + "sync/atomic" "testing" "time" @@ -67,9 +68,11 @@ func testDeps(t *testing.T) Deps { // clientHarness wires a client Conn to an Agent over in-memory pipes and collects // session/update text chunks. type clientHarness struct { - client *Conn - updates chan string - stop func() + agent *Agent + agentConn *Conn + client *Conn + updates chan string + stop func() } func newHarness(t *testing.T, deps Deps) *clientHarness { @@ -80,7 +83,7 @@ func newHarness(t *testing.T, deps Deps) *clientHarness { client := NewConn(br, bw) a := NewAgent(agentConn, deps) - h := &clientHarness{client: client, updates: make(chan string, 128)} + h := &clientHarness{agent: a, agentConn: agentConn, client: client, updates: make(chan string, 128)} client.HandleNotify(MethodSessionUpdate, func(_ context.Context, params json.RawMessage) { var probe struct { Update struct { @@ -599,3 +602,148 @@ func drainTextUntil(t *testing.T, ch <-chan string, done func(string) bool) stri } } } + +func TestACPCancelInterleavedSessionsCancelsBothPrompts(t *testing.T) { + prompt1Entered := make(chan struct{}) + prompt2Entered := make(chan struct{}) + blockPrompt1 := make(chan struct{}) + blockPrompt2 := make(chan struct{}) + + firstCancelStarted := make(chan struct{}) + holdFirstCancel := make(chan struct{}) + var cancelCount atomic.Int32 + + deps := testDeps(t) + deps.RunAgent = func(ctx context.Context, prompt string, _ zeroruntime.Provider, opts agent.Options) (agent.Result, error) { + if opts.SessionID == "sess-1" { + close(prompt1Entered) + select { + case <-blockPrompt1: + return agent.Result{FinalAnswer: "done1"}, nil + case <-ctx.Done(): + return agent.Result{}, ctx.Err() + } + } else if opts.SessionID == "sess-2" { + close(prompt2Entered) + select { + case <-blockPrompt2: + return agent.Result{FinalAnswer: "done2"}, nil + case <-ctx.Done(): + return agent.Result{}, ctx.Err() + } + } + return agent.Result{}, nil + } + + h := newHarness(t, deps) + defer h.stop() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + var initRes InitializeResult + if err := h.client.Call(ctx, MethodInitialize, InitializeParams{ProtocolVersion: ProtocolVersion}, &initRes); err != nil { + t.Fatalf("initialize: %v", err) + } + + root := t.TempDir() + _, _ = deps.Store.Create(sessions.CreateInput{SessionID: "sess-1", Title: "s1", Cwd: root, ModelID: "fake-model"}) + _, _ = deps.Store.Create(sessions.CreateInput{SessionID: "sess-2", Title: "s2", Cwd: root, ModelID: "fake-model"}) + + var load1, load2 LoadSessionResult + if err := h.client.Call(ctx, MethodSessionLoad, LoadSessionParams{SessionID: "sess-1", Cwd: root}, &load1); err != nil { + t.Fatalf("load sess-1: %v", err) + } + if err := h.client.Call(ctx, MethodSessionLoad, LoadSessionParams{SessionID: "sess-2", Cwd: root}, &load2); err != nil { + t.Fatalf("load sess-2: %v", err) + } + + // Intercept cancel notification handler to delay the first cancel execution + origCancel := h.agent.conn.notifiers[MethodSessionCancel] + h.agent.conn.HandleNotify(MethodSessionCancel, func(ctx context.Context, params json.RawMessage) { + if cancelCount.Add(1) == 1 { + close(firstCancelStarted) + <-holdFirstCancel + } + origCancel(ctx, params) + }) + + type promptOutcome struct { + res PromptResult + err error + } + res1Ch := make(chan promptOutcome, 1) + res2Ch := make(chan promptOutcome, 1) + + go func() { + var res PromptResult + err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: "sess-1", Prompt: []ContentBlock{TextBlock("prompt1")}}, &res) + res1Ch <- promptOutcome{res: res, err: err} + }() + + go func() { + var res PromptResult + err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: "sess-2", Prompt: []ContentBlock{TextBlock("prompt2")}}, &res) + res2Ch <- promptOutcome{res: res, err: err} + }() + + select { + case <-prompt1Entered: + case <-time.After(2 * time.Second): + t.Fatal("prompt1 did not enter RunAgent") + } + select { + case <-prompt2Entered: + case <-time.After(2 * time.Second): + t.Fatal("prompt2 did not enter RunAgent") + } + + // Send cancel for sess-1 which starts the worker and waits on holdFirstCancel + _ = h.client.Notify(MethodSessionCancel, CancelParams{SessionID: "sess-1"}) + select { + case <-firstCancelStarted: + case <-time.After(2 * time.Second): + t.Fatal("first cancel did not start") + } + + // Interleave cancel for sess-2, then another cancel for sess-1 + _ = h.client.Notify(MethodSessionCancel, CancelParams{SessionID: "sess-2"}) + _ = h.client.Notify(MethodSessionCancel, CancelParams{SessionID: "sess-1"}) + + deadlineWait := time.Now().Add(2 * time.Second) + for { + h.agent.conn.notifyMu.Lock() + queued := string(h.agent.conn.notifyQ[notifyKey{method: MethodSessionCancel, target: "sess-1"}]) + h.agent.conn.notifyMu.Unlock() + if strings.Contains(queued, "sess-1") || time.Now().After(deadlineWait) { + break + } + time.Sleep(5 * time.Millisecond) + } + + close(holdFirstCancel) + + select { + case out1 := <-res1Ch: + if out1.err != nil { + t.Fatalf("prompt1 error: %v", out1.err) + } + if out1.res.StopReason != StopCancelled { + t.Fatalf("prompt1 StopReason = %q, want %q", out1.res.StopReason, StopCancelled) + } + case <-time.After(2 * time.Second): + t.Fatal("prompt1 did not cancel") + } + + select { + case out2 := <-res2Ch: + if out2.err != nil { + t.Fatalf("prompt2 error: %v", out2.err) + } + if out2.res.StopReason != StopCancelled { + t.Fatalf("prompt2 StopReason = %q, want %q", out2.res.StopReason, StopCancelled) + } + case <-time.After(2 * time.Second): + t.Fatal("prompt2 did not cancel (session cancel was overwritten/coalesced across sessions)") + } +} diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index 7ee141750..6b03aedb5 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -18,6 +18,7 @@ import ( "io" "sync" "sync/atomic" + "time" ) // JSON-RPC 2.0 standard error codes. @@ -88,6 +89,8 @@ const ( maxFrameBytes = 64 * 1024 * 1024 maxConcurrentRequests = 128 maxBusyReplies = 1 + maxNotifyActive = 32 + overloadDrainTimeout = 100 * time.Millisecond ) // Conn is a JSON-RPC 2.0 peer over a single ndjson stream pair. It both serves @@ -123,8 +126,26 @@ type Conn struct { writeWaiters atomic.Int32 notifyMu sync.Mutex - notifyOn map[string]bool - notifyQ map[string]json.RawMessage + notifyOn map[notifyKey]bool + notifyQ map[notifyKey]json.RawMessage +} + +type notifyKey struct { + method string + target string +} + +func notifyTarget(params json.RawMessage) string { + if len(params) == 0 { + return "" + } + var header struct { + SessionID string `json:"sessionId"` + } + if err := json.Unmarshal(params, &header); err == nil && header.SessionID != "" { + return header.SessionID + } + return "" } // NewConn builds a peer reading ndjson from r and writing ndjson to w. This @@ -143,8 +164,8 @@ func NewConn(r io.Reader, w io.Writer) *Conn { sem: make(chan struct{}, maxConcurrentRequests), busyCh: make(chan json.RawMessage, maxBusyReplies), writeAbort: make(chan struct{}), - notifyOn: make(map[string]bool), - notifyQ: make(map[string]json.RawMessage), + notifyOn: make(map[notifyKey]bool), + notifyQ: make(map[notifyKey]json.RawMessage), } } @@ -177,11 +198,24 @@ func (c *Conn) Serve(ctx context.Context) error { // On exit, cancel in-flight handlers (so a blocked outbound Call unblocks via // ctx) and then wait for them to finish writing their responses. Without this, // a finite input stream (e.g. piped ndjson that EOFs right after a request) - // would race the dispatch goroutine and drop the response. The busy-reply - // writer is not on wg: it may be blocked in Write against a stalled peer. + // would race the dispatch goroutine and drop the response. When overloaded, + // skip waiting on wg so a handler already blocked inside a stalled transport + // Write does not retain Serve indefinitely. defer func() { cancel() - c.wg.Wait() + if !c.overloaded.Load() { + c.wg.Wait() + return + } + done := make(chan struct{}) + go func() { + c.wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(overloadDrainTimeout): + } }() interruptible := newInterruptibleReader(ctx, c.rawReader, c.readerCloser) @@ -406,30 +440,35 @@ func (c *Conn) dispatchNotify(ctx context.Context, msg rpcMessage) { return } params := append(json.RawMessage(nil), msg.Params...) + key := notifyKey{method: msg.Method, target: notifyTarget(params)} c.notifyMu.Lock() - if c.notifyOn[msg.Method] { - c.notifyQ[msg.Method] = params + if c.notifyOn[key] { + c.notifyQ[key] = params + c.notifyMu.Unlock() + return + } + if len(c.notifyOn) >= maxNotifyActive { c.notifyMu.Unlock() return } - c.notifyOn[msg.Method] = true + c.notifyOn[key] = true c.notifyMu.Unlock() c.wg.Add(1) - go c.runNotify(ctx, msg.Method, fn, params) + go c.runNotify(ctx, key, fn, params) } -func (c *Conn) runNotify(ctx context.Context, method string, fn NotifyFunc, params json.RawMessage) { +func (c *Conn) runNotify(ctx context.Context, key notifyKey, fn NotifyFunc, params json.RawMessage) { defer c.wg.Done() for { fn(ctx, params) c.notifyMu.Lock() - next, ok := c.notifyQ[method] + next, ok := c.notifyQ[key] if !ok { - delete(c.notifyOn, method) + delete(c.notifyOn, key) c.notifyMu.Unlock() return } - delete(c.notifyQ, method) + delete(c.notifyQ, key) c.notifyMu.Unlock() params = next } diff --git a/internal/acp/jsonrpc_test.go b/internal/acp/jsonrpc_test.go index 5da118fc8..7dc6d4037 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -1077,3 +1077,157 @@ func TestQueuedBusyReplySurvivesOverloadBurst(t *testing.T) { } t.Fatalf("busy ids 2 (in-flight) and 3 (queued) must receive -32000; wrote %q", got) } + +func TestInterleavedSessionCancelsDoNotCoalesceAcrossSessions(t *testing.T) { + inReader, inWriter := io.Pipe() + defer inWriter.Close() + + conn := NewConn(inReader, io.Discard) + + firstRun := make(chan struct{}) + gate := make(chan struct{}) + var ( + mu sync.Mutex + cancelledA int + cancelledB int + ) + + conn.HandleNotify("session/cancel", func(_ context.Context, params json.RawMessage) { + var p struct { + SessionID string `json:"sessionId"` + } + _ = json.Unmarshal(params, &p) + + mu.Lock() + if p.SessionID == "sess-A" { + cancelledA++ + } else if p.SessionID == "sess-B" { + cancelledB++ + } + first := (cancelledA + cancelledB) == 1 + mu.Unlock() + + if first { + close(firstRun) + <-gate + } + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { _ = conn.Serve(ctx) }() + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","method":"session/cancel","params":{"sessionId":"sess-A"}}` + "\n")) + select { + case <-firstRun: + case <-time.After(2 * time.Second): + t.Fatal("first session cancel did not start") + } + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","method":"session/cancel","params":{"sessionId":"sess-B"}}` + "\n")) + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","method":"session/cancel","params":{"sessionId":"sess-A"}}` + "\n")) + + // Wait until worker for sess-B runs or notifyQ receives the second sess-A before releasing gate. + deadlineWait := time.Now().Add(2 * time.Second) + for { + conn.notifyMu.Lock() + queued := string(conn.notifyQ[notifyKey{method: "session/cancel", target: "sess-A"}]) + conn.notifyMu.Unlock() + if strings.Contains(queued, "sess-A") || time.Now().After(deadlineWait) { + break + } + time.Sleep(5 * time.Millisecond) + } + + close(gate) + + deadline := time.Now().Add(2 * time.Second) + for { + mu.Lock() + gotA := cancelledA + gotB := cancelledB + mu.Unlock() + if gotA >= 1 && gotB >= 1 { + break + } + if time.Now().After(deadline) { + t.Fatalf("session cancels not delivered to both sessions: sess-A=%d, sess-B=%d", gotA, gotB) + } + time.Sleep(10 * time.Millisecond) + } +} + +func TestOverloadTerminatesHandlerInsideStalledWrite(t *testing.T) { + inReader, inWriter := io.Pipe() + defer inWriter.Close() + + started := make(chan struct{}) + gate := make(chan struct{}) + defer close(gate) + conn := NewConn(inReader, &stallWriter{started: started, gate: gate}) + conn.sem = make(chan struct{}, 1) + + conn.Handle("echo", func(ctx context.Context, _ json.RawMessage) (any, error) { + return "ok", nil + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan error, 1) + go func() { done <- conn.Serve(ctx) }() + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":"echo"}` + "\n")) + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("handler did not enter stalled write") + } + + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":2,"method":"echo"}` + "\n")) + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":3,"method":"echo"}` + "\n")) + _, _ = inWriter.Write([]byte(`{"jsonrpc":"2.0","id":4,"method":"echo"}` + "\n")) + + select { + case err := <-done: + if !errors.Is(err, errBusyOverload) { + t.Fatalf("Serve returned %v, want %v", err, errBusyOverload) + } + case <-time.After(2 * time.Second): + t.Fatal("Serve did not return on overload while handler was stalled in Write") + } +} + +func TestNotifyActiveIsBounded(t *testing.T) { + inReader, inWriter := io.Pipe() + defer inWriter.Close() + conn := NewConn(inReader, io.Discard) + gate := make(chan struct{}) + conn.HandleNotify("session/cancel", func(context.Context, json.RawMessage) { + <-gate + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { _ = conn.Serve(ctx) }() + + for i := 0; i < maxNotifyActive+40; i++ { + frame := fmt.Sprintf(`{"jsonrpc":"2.0","method":"session/cancel","params":{"sessionId":"s-%d"}}`+"\n", i) + if _, err := inWriter.Write([]byte(frame)); err != nil { + t.Fatalf("write cancel %d: %v", i, err) + } + } + deadline := time.Now().Add(2 * time.Second) + for { + conn.notifyMu.Lock() + n := len(conn.notifyOn) + conn.notifyMu.Unlock() + if n >= maxNotifyActive || time.Now().After(deadline) { + if n > maxNotifyActive { + close(gate) + t.Fatalf("notifyOn = %d, want <= %d", n, maxNotifyActive) + } + break + } + time.Sleep(5 * time.Millisecond) + } + close(gate) +} From b3213554e7342c6a3fbd6df2ce376a886d000ec9 Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sun, 30 Aug 2026 21:16:31 +0200 Subject: [PATCH 11/14] fix(acp): keep session/update ordered and bound admitted request bytes Queue session/update in a FIFO, deliver every session/cancel past the notify cap, account inflight bytes before admit, and join writeBusyLoop. --- internal/acp/jsonrpc.go | 191 ++++++++++++++++++++++++++++------- internal/acp/jsonrpc_test.go | 125 ++++++++++++++++++++++- 2 files changed, 280 insertions(+), 36 deletions(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index 6b03aedb5..a77125ad2 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -88,8 +88,10 @@ const ( // the trailing newline delimiter (effective maximum payload is limit - 1). maxFrameBytes = 64 * 1024 * 1024 maxConcurrentRequests = 128 + maxInflightBytes = 64 * 1024 * 1024 maxBusyReplies = 1 maxNotifyActive = 32 + maxUpdateFIFO = 32 overloadDrainTimeout = 100 * time.Millisecond ) @@ -125,9 +127,15 @@ type Conn struct { writeAbort chan struct{} writeWaiters atomic.Int32 - notifyMu sync.Mutex - notifyOn map[notifyKey]bool - notifyQ map[notifyKey]json.RawMessage + notifyMu sync.Mutex + notifyOn map[notifyKey]bool + notifyQ map[notifyKey]json.RawMessage + cancelOn map[notifyKey]bool + updateOn map[string]bool + sessionUpdateQ map[string][]json.RawMessage + admittedMu sync.Mutex + admittedBytes int64 + inflightLimit int64 } type notifyKey struct { @@ -156,16 +164,19 @@ func notifyTarget(params json.RawMessage) string { // lifetime. func NewConn(r io.Reader, w io.Writer) *Conn { return &Conn{ - rawReader: r, - w: w, - handlers: make(map[string]HandlerFunc), - notifiers: make(map[string]NotifyFunc), - pending: make(map[int64]chan rpcMessage), - sem: make(chan struct{}, maxConcurrentRequests), - busyCh: make(chan json.RawMessage, maxBusyReplies), - writeAbort: make(chan struct{}), - notifyOn: make(map[notifyKey]bool), - notifyQ: make(map[notifyKey]json.RawMessage), + rawReader: r, + w: w, + handlers: make(map[string]HandlerFunc), + notifiers: make(map[string]NotifyFunc), + pending: make(map[int64]chan rpcMessage), + sem: make(chan struct{}, maxConcurrentRequests), + busyCh: make(chan json.RawMessage, maxBusyReplies), + writeAbort: make(chan struct{}), + notifyOn: make(map[notifyKey]bool), + notifyQ: make(map[notifyKey]json.RawMessage), + cancelOn: make(map[notifyKey]bool), + updateOn: make(map[string]bool), + sessionUpdateQ: make(map[string][]json.RawMessage), } } @@ -194,6 +205,7 @@ func (c *Conn) HandleNotify(method string, fn NotifyFunc) { c.notifiers[method] func (c *Conn) Serve(ctx context.Context) error { ctx, cancel := context.WithCancel(ctx) c.serveCancel = cancel + c.wg.Add(1) go c.writeBusyLoop(ctx) // On exit, cancel in-flight handlers (so a blocked outbound Call unblocks via // ctx) and then wait for them to finish writing their responses. Without this, @@ -400,28 +412,20 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { case msg.isResponse(): c.deliver(msg) case msg.isRequest(): - if c.sem != nil { - select { - case c.sem <- struct{}{}: - c.wg.Add(1) - go func(m rpcMessage) { - defer c.wg.Done() - defer c.releaseSem() - c.dispatchRequest(ctx, m) - }(msg) - default: - id := append(json.RawMessage(nil), msg.ID...) - if !c.tryEnqueueBusy(id) { - c.tripOverload() - } + n := int64(len(line)) + if !c.tryAdmit(n) { + id := append(json.RawMessage(nil), msg.ID...) + if !c.tryEnqueueBusy(id) { + c.tripOverload() } - } else { - c.wg.Add(1) - go func(m rpcMessage) { - defer c.wg.Done() - c.dispatchRequest(ctx, m) - }(msg) + break } + c.wg.Add(1) + go func(m rpcMessage, bytes int64) { + defer c.wg.Done() + defer c.releaseAdmit(bytes) + c.dispatchRequest(ctx, m) + }(msg, n) case msg.isNotify(): if !c.overloaded.Load() { c.dispatchNotify(ctx, msg) @@ -440,7 +444,86 @@ func (c *Conn) dispatchNotify(ctx context.Context, msg rpcMessage) { return } params := append(json.RawMessage(nil), msg.Params...) - key := notifyKey{method: msg.Method, target: notifyTarget(params)} + switch msg.Method { + case MethodSessionUpdate: + c.dispatchSessionUpdate(ctx, fn, params) + case MethodSessionCancel: + c.dispatchCancel(ctx, fn, params) + default: + c.dispatchCoalescedNotify(ctx, msg.Method, fn, params) + } +} + +func (c *Conn) dispatchSessionUpdate(ctx context.Context, fn NotifyFunc, params json.RawMessage) { + target := notifyTarget(params) + c.notifyMu.Lock() + if c.updateOn[target] { + q := c.sessionUpdateQ[target] + if len(q) >= maxUpdateFIFO { + c.notifyMu.Unlock() + c.tripOverload() + return + } + c.sessionUpdateQ[target] = append(q, params) + c.notifyMu.Unlock() + return + } + c.updateOn[target] = true + c.notifyMu.Unlock() + c.wg.Add(1) + go c.runSessionUpdate(ctx, target, fn, params) +} + +func (c *Conn) runSessionUpdate(ctx context.Context, target string, fn NotifyFunc, params json.RawMessage) { + defer c.wg.Done() + for { + fn(ctx, params) + c.notifyMu.Lock() + q := c.sessionUpdateQ[target] + if len(q) == 0 { + delete(c.updateOn, target) + c.notifyMu.Unlock() + return + } + params = q[0] + c.sessionUpdateQ[target] = q[1:] + c.notifyMu.Unlock() + } +} + +func (c *Conn) dispatchCancel(ctx context.Context, fn NotifyFunc, params json.RawMessage) { + key := notifyKey{method: MethodSessionCancel, target: notifyTarget(params)} + c.notifyMu.Lock() + if c.cancelOn[key] { + c.notifyQ[key] = params + c.notifyMu.Unlock() + return + } + c.cancelOn[key] = true + c.notifyMu.Unlock() + c.wg.Add(1) + go c.runCancel(ctx, key, fn, params) +} + +func (c *Conn) runCancel(ctx context.Context, key notifyKey, fn NotifyFunc, params json.RawMessage) { + defer c.wg.Done() + for { + fn(ctx, params) + c.notifyMu.Lock() + next, ok := c.notifyQ[key] + if !ok { + delete(c.cancelOn, key) + c.notifyMu.Unlock() + return + } + delete(c.notifyQ, key) + c.notifyMu.Unlock() + params = next + } +} + +func (c *Conn) dispatchCoalescedNotify(ctx context.Context, method string, fn NotifyFunc, params json.RawMessage) { + key := notifyKey{method: method, target: notifyTarget(params)} c.notifyMu.Lock() if c.notifyOn[key] { c.notifyQ[key] = params @@ -457,6 +540,45 @@ func (c *Conn) dispatchNotify(ctx context.Context, msg rpcMessage) { go c.runNotify(ctx, key, fn, params) } +func (c *Conn) tryAdmit(n int64) bool { + limit := int64(maxInflightBytes) + if c.inflightLimit > 0 { + limit = c.inflightLimit + } + c.admittedMu.Lock() + if c.admittedBytes+n > limit { + c.admittedMu.Unlock() + return false + } + c.admittedBytes += n + c.admittedMu.Unlock() + if c.sem == nil { + return true + } + select { + case c.sem <- struct{}{}: + return true + default: + c.admittedMu.Lock() + c.admittedBytes -= n + if c.admittedBytes < 0 { + c.admittedBytes = 0 + } + c.admittedMu.Unlock() + return false + } +} + +func (c *Conn) releaseAdmit(n int64) { + c.releaseSem() + c.admittedMu.Lock() + c.admittedBytes -= n + if c.admittedBytes < 0 { + c.admittedBytes = 0 + } + c.admittedMu.Unlock() +} + func (c *Conn) runNotify(ctx context.Context, key notifyKey, fn NotifyFunc, params json.RawMessage) { defer c.wg.Done() for { @@ -631,6 +753,7 @@ func (c *Conn) tripOverload() { } func (c *Conn) writeBusyLoop(ctx context.Context) { + defer c.wg.Done() busy := &rpcError{Code: codeServerBusy, Message: "server busy: max concurrent requests exceeded"} for { select { diff --git a/internal/acp/jsonrpc_test.go b/internal/acp/jsonrpc_test.go index 7dc6d4037..fa56a9315 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -1202,7 +1202,7 @@ func TestNotifyActiveIsBounded(t *testing.T) { defer inWriter.Close() conn := NewConn(inReader, io.Discard) gate := make(chan struct{}) - conn.HandleNotify("session/cancel", func(context.Context, json.RawMessage) { + conn.HandleNotify("custom/ping", func(context.Context, json.RawMessage) { <-gate }) ctx, cancel := context.WithCancel(context.Background()) @@ -1210,7 +1210,7 @@ func TestNotifyActiveIsBounded(t *testing.T) { go func() { _ = conn.Serve(ctx) }() for i := 0; i < maxNotifyActive+40; i++ { - frame := fmt.Sprintf(`{"jsonrpc":"2.0","method":"session/cancel","params":{"sessionId":"s-%d"}}`+"\n", i) + frame := fmt.Sprintf(`{"jsonrpc":"2.0","method":"custom/ping","params":{"sessionId":"s-%d"}}`+"\n", i) if _, err := inWriter.Write([]byte(frame)); err != nil { t.Fatalf("write cancel %d: %v", i, err) } @@ -1231,3 +1231,124 @@ func TestNotifyActiveIsBounded(t *testing.T) { } close(gate) } + +func TestSessionUpdatePreservesOrder(t *testing.T) { + inReader, inWriter := io.Pipe() + defer inWriter.Close() + conn := NewConn(inReader, io.Discard) + gate := make(chan struct{}) + var got []string + var mu sync.Mutex + conn.HandleNotify(MethodSessionUpdate, func(_ context.Context, params json.RawMessage) { + <-gate + mu.Lock() + got = append(got, string(params)) + mu.Unlock() + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { _ = conn.Serve(ctx) }() + + for i, payload := range []string{`{"sessionId":"s","n":1}`, `{"sessionId":"s","n":2}`, `{"sessionId":"s","n":3}`} { + frame := `{"jsonrpc":"2.0","method":"session/update","params":` + payload + `}` + "\n" + if _, err := inWriter.Write([]byte(frame)); err != nil { + t.Fatalf("write update %d: %v", i, err) + } + } + time.Sleep(50 * time.Millisecond) + close(gate) + deadline := time.Now().Add(2 * time.Second) + for { + mu.Lock() + n := len(got) + mu.Unlock() + if n >= 3 || time.Now().After(deadline) { + break + } + time.Sleep(5 * time.Millisecond) + } + mu.Lock() + defer mu.Unlock() + if len(got) != 3 { + t.Fatalf("got %d updates, want 3 (no coalescing): %v", len(got), got) + } +} + +func TestSessionCancelExceedsNotifyCap(t *testing.T) { + inReader, inWriter := io.Pipe() + defer inWriter.Close() + conn := NewConn(inReader, io.Discard) + var started atomic.Int32 + gate := make(chan struct{}) + conn.HandleNotify(MethodSessionCancel, func(context.Context, json.RawMessage) { + started.Add(1) + <-gate + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { _ = conn.Serve(ctx) }() + + want := maxNotifyActive + 8 + for i := 0; i < want; i++ { + frame := fmt.Sprintf(`{"jsonrpc":"2.0","method":"session/cancel","params":{"sessionId":"s-%d"}}`+"\n", i) + if _, err := inWriter.Write([]byte(frame)); err != nil { + t.Fatalf("write cancel %d: %v", i, err) + } + } + deadline := time.Now().Add(2 * time.Second) + for started.Load() < int32(want) && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + if started.Load() != int32(want) { + close(gate) + t.Fatalf("cancels started = %d, want %d", started.Load(), want) + } + close(gate) +} + +func TestInflightByteBudgetRejects(t *testing.T) { + inReader, inWriter := io.Pipe() + defer inWriter.Close() + var mu sync.Mutex + var out strings.Builder + conn := NewConn(inReader, testWriter(func(p []byte) (int, error) { + mu.Lock() + defer mu.Unlock() + return out.Write(p) + })) + conn.inflightLimit = 80 + conn.sem = make(chan struct{}, 8) + gate := make(chan struct{}) + conn.Handle("session/prompt", func(context.Context, json.RawMessage) (any, error) { + <-gate + return map[string]any{}, nil + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { _ = conn.Serve(ctx) }() + + frame := `{"jsonrpc":"2.0","id":1,"method":"session/prompt","params":{"sessionId":"s","prompt":"xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"}}` + "\n" + if _, err := inWriter.Write([]byte(frame)); err != nil { + t.Fatal(err) + } + time.Sleep(50 * time.Millisecond) + if _, err := inWriter.Write([]byte(`{"jsonrpc":"2.0","id":2,"method":"session/prompt","params":{"sessionId":"s","prompt":"yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy"}}` + "\n")); err != nil { + t.Fatal(err) + } + deadline := time.Now().Add(2 * time.Second) + got := "" + for time.Now().Before(deadline) { + mu.Lock() + got = out.String() + mu.Unlock() + if strings.Contains(got, "-32000") { + break + } + time.Sleep(5 * time.Millisecond) + } + close(gate) + _ = inWriter.Close() + if !strings.Contains(got, "-32000") { + t.Fatalf("expected busy from byte budget, got %q", got) + } +} From 882eafc1b96d70257c723848cea946d0301d664b Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sun, 30 Aug 2026 21:43:08 +0200 Subject: [PATCH 12/14] fix(acp): drain update FIFO, bound special notifies, cap Serve wait Delete emptied sessionUpdateQ keys, admit session/update and cancel under maxSpecialNotify, and always bound writeBusyLoop join on shutdown. --- internal/acp/jsonrpc.go | 20 ++++++++++++++++---- internal/acp/jsonrpc_test.go | 31 ++++++++++++++++++++++++------- 2 files changed, 40 insertions(+), 11 deletions(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index a77125ad2..38a489cff 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -92,6 +92,7 @@ const ( maxBusyReplies = 1 maxNotifyActive = 32 maxUpdateFIFO = 32 + maxSpecialNotify = 256 overloadDrainTimeout = 100 * time.Millisecond ) @@ -215,10 +216,6 @@ func (c *Conn) Serve(ctx context.Context) error { // Write does not retain Serve indefinitely. defer func() { cancel() - if !c.overloaded.Load() { - c.wg.Wait() - return - } done := make(chan struct{}) go func() { c.wg.Wait() @@ -468,12 +465,21 @@ func (c *Conn) dispatchSessionUpdate(ctx context.Context, fn NotifyFunc, params c.notifyMu.Unlock() return } + if c.specialNotifyBusyLocked() { + c.notifyMu.Unlock() + c.tripOverload() + return + } c.updateOn[target] = true c.notifyMu.Unlock() c.wg.Add(1) go c.runSessionUpdate(ctx, target, fn, params) } +func (c *Conn) specialNotifyBusyLocked() bool { + return len(c.updateOn)+len(c.cancelOn) >= maxSpecialNotify +} + func (c *Conn) runSessionUpdate(ctx context.Context, target string, fn NotifyFunc, params json.RawMessage) { defer c.wg.Done() for { @@ -482,6 +488,7 @@ func (c *Conn) runSessionUpdate(ctx context.Context, target string, fn NotifyFun q := c.sessionUpdateQ[target] if len(q) == 0 { delete(c.updateOn, target) + delete(c.sessionUpdateQ, target) c.notifyMu.Unlock() return } @@ -499,6 +506,11 @@ func (c *Conn) dispatchCancel(ctx context.Context, fn NotifyFunc, params json.Ra c.notifyMu.Unlock() return } + if c.specialNotifyBusyLocked() { + c.notifyMu.Unlock() + c.tripOverload() + return + } c.cancelOn[key] = true c.notifyMu.Unlock() c.wg.Add(1) diff --git a/internal/acp/jsonrpc_test.go b/internal/acp/jsonrpc_test.go index fa56a9315..b30cf849b 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -1269,8 +1269,14 @@ func TestSessionUpdatePreservesOrder(t *testing.T) { } mu.Lock() defer mu.Unlock() - if len(got) != 3 { - t.Fatalf("got %d updates, want 3 (no coalescing): %v", len(got), got) + want := []string{`{"sessionId":"s","n":1}`, `{"sessionId":"s","n":2}`, `{"sessionId":"s","n":3}`} + if len(got) != len(want) { + t.Fatalf("got %d updates, want %d: %v", len(got), len(want), got) + } + for i := range want { + if got[i] != want[i] { + t.Fatalf("update %d = %s, want %s", i, got[i], want[i]) + } } } @@ -1316,10 +1322,18 @@ func TestInflightByteBudgetRejects(t *testing.T) { defer mu.Unlock() return out.Write(p) })) - conn.inflightLimit = 80 + frame1 := `{"jsonrpc":"2.0","id":1,"method":"session/prompt","params":{}}` + "\n" + frame2 := `{"jsonrpc":"2.0","id":2,"method":"session/prompt","params":{}}` + "\n" + conn.inflightLimit = int64(len(frame1)) conn.sem = make(chan struct{}, 8) gate := make(chan struct{}) + started := make(chan struct{}) conn.Handle("session/prompt", func(context.Context, json.RawMessage) (any, error) { + select { + case <-started: + default: + close(started) + } <-gate return map[string]any{}, nil }) @@ -1327,12 +1341,15 @@ func TestInflightByteBudgetRejects(t *testing.T) { defer cancel() go func() { _ = conn.Serve(ctx) }() - frame := `{"jsonrpc":"2.0","id":1,"method":"session/prompt","params":{"sessionId":"s","prompt":"xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"}}` + "\n" - if _, err := inWriter.Write([]byte(frame)); err != nil { + if _, err := inWriter.Write([]byte(frame1)); err != nil { t.Fatal(err) } - time.Sleep(50 * time.Millisecond) - if _, err := inWriter.Write([]byte(`{"jsonrpc":"2.0","id":2,"method":"session/prompt","params":{"sessionId":"s","prompt":"yyyyyyyyyyyyyyyyyyyyyyyyyyyyyyyy"}}` + "\n")); err != nil { + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("first request did not start") + } + if _, err := inWriter.Write([]byte(frame2)); err != nil { t.Fatal(err) } deadline := time.Now().Add(2 * time.Second) From d3d6d51e3321b285e60ab51f1fa2c84ef9320e12 Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Sun, 30 Aug 2026 23:20:59 +0200 Subject: [PATCH 13/14] fix(acp): wait for replies on clean Serve exit, persist overload writes Cap wg.Wait at 100ms only when overloaded. Admitted handler replies use persist writes during overload so cancelled turns still emit a frame. --- internal/acp/jsonrpc.go | 30 +++++++++++++++++------------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index 38a489cff..f30de0f8d 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -211,20 +211,24 @@ func (c *Conn) Serve(ctx context.Context) error { // On exit, cancel in-flight handlers (so a blocked outbound Call unblocks via // ctx) and then wait for them to finish writing their responses. Without this, // a finite input stream (e.g. piped ndjson that EOFs right after a request) - // would race the dispatch goroutine and drop the response. When overloaded, - // skip waiting on wg so a handler already blocked inside a stalled transport - // Write does not retain Serve indefinitely. + // would race the dispatch goroutine and drop the response. Clean EOF/cancel + // waits for wg so admitted replies can flush. Overload keeps a bounded drain + // so a stalled Write cannot retain Serve indefinitely. defer func() { cancel() - done := make(chan struct{}) - go func() { - c.wg.Wait() - close(done) - }() - select { - case <-done: - case <-time.After(overloadDrainTimeout): + if c.overloaded.Load() { + done := make(chan struct{}) + go func() { + c.wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(overloadDrainTimeout): + } + return } + c.wg.Wait() }() interruptible := newInterruptibleReader(ctx, c.rawReader, c.readerCloser) @@ -732,11 +736,11 @@ func (c *Conn) writeResult(ctx context.Context, id json.RawMessage, result any) c.writeError(ctx, id, &rpcError{Code: codeInternalError, Message: err.Error()}) return } - _ = c.write(ctx, rpcMessage{JSONRPC: "2.0", ID: id, Result: raw}) + _ = c.writeMsg(ctx, rpcMessage{JSONRPC: "2.0", ID: id, Result: raw}, c.overloaded.Load()) } func (c *Conn) writeError(ctx context.Context, id json.RawMessage, e *rpcError) { - _ = c.write(ctx, rpcMessage{JSONRPC: "2.0", ID: id, Error: e}) + _ = c.writeMsg(ctx, rpcMessage{JSONRPC: "2.0", ID: id, Error: e}, c.overloaded.Load()) } func (c *Conn) tryEnqueueBusy(id json.RawMessage) bool { From a9451402b77af86f8d7cf68c66c2c8b52d5a68e6 Mon Sep 17 00:00:00 2001 From: hazyhaar Date: Mon, 31 Aug 2026 01:27:36 +0200 Subject: [PATCH 14/14] fix(acp): close the writer on overload and drop leftover busy debt Own the write closer, stop admitting busy IDs after overload, and close the writer so a stalled Write cannot hold Serve. Busy JSON is best-effort only while the pipe still accepts bytes. --- internal/acp/jsonrpc.go | 42 ++++++++++----------- internal/acp/jsonrpc_test.go | 72 ++++++++++++++++++++++++++++++------ 2 files changed, 81 insertions(+), 33 deletions(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index f30de0f8d..05c16c8ac 100644 --- a/internal/acp/jsonrpc.go +++ b/internal/acp/jsonrpc.go @@ -101,9 +101,11 @@ const ( // requests/notifications — needed because ACP is bidirectional (the agent calls // the client for session/request_permission, fs/*, terminal/*). type Conn struct { - rawReader io.Reader // wrapped lazily in Serve, once ctx is known — see interruptibleReader - readerCloser io.Closer - w io.Writer + rawReader io.Reader // wrapped lazily in Serve, once ctx is known — see interruptibleReader + readerCloser io.Closer + w io.Writer + writerCloser io.Closer + writeCloseOnce sync.Once writeMu sync.Mutex // serializes all writes to w @@ -190,9 +192,18 @@ func NewConn(r io.Reader, w io.Writer) *Conn { func NewOwnedConn(r io.Reader, w io.Writer) *Conn { c := NewConn(r, w) c.readerCloser, _ = r.(io.Closer) + c.writerCloser, _ = w.(io.Closer) return c } +func (c *Conn) closeWriter() { + c.writeCloseOnce.Do(func() { + if c.writerCloser != nil { + _ = c.writerCloser.Close() + } + }) +} + // Handle registers a request handler for method. func (c *Conn) Handle(method string, fn HandlerFunc) { c.handlers[method] = fn } @@ -217,6 +228,7 @@ func (c *Conn) Serve(ctx context.Context) error { defer func() { cancel() if c.overloaded.Load() { + c.closeWriter() done := make(chan struct{}) go func() { c.wg.Wait() @@ -744,7 +756,7 @@ func (c *Conn) writeError(ctx context.Context, id json.RawMessage, e *rpcError) } func (c *Conn) tryEnqueueBusy(id json.RawMessage) bool { - if c.busyCh == nil { + if c.overloaded.Load() || c.busyCh == nil { return false } select { @@ -774,29 +786,13 @@ func (c *Conn) writeBusyLoop(ctx context.Context) { for { select { case <-ctx.Done(): - c.flushBusy(busy) return case id := <-c.busyCh: - c.writeBusy(id, busy) + _ = c.writeMsg(ctx, rpcMessage{JSONRPC: "2.0", ID: id, Error: busy}, false) } } } -func (c *Conn) flushBusy(busy *rpcError) { - for { - select { - case id := <-c.busyCh: - c.writeBusy(id, busy) - default: - return - } - } -} - -func (c *Conn) writeBusy(id json.RawMessage, busy *rpcError) { - _ = c.writeMsg(context.Background(), rpcMessage{JSONRPC: "2.0", ID: id, Error: busy}, true) -} - func (c *Conn) lockWrite(ctx context.Context, persist bool) error { if ctx == nil { ctx = context.Background() @@ -829,6 +825,10 @@ func (c *Conn) lockWrite(ctx context.Context, persist bool) error { releaseWait() abandon() return ctx.Err() + case <-c.writeAbort: + releaseWait() + abandon() + return errBusyOverload } } select { diff --git a/internal/acp/jsonrpc_test.go b/internal/acp/jsonrpc_test.go index b30cf849b..4e0210594 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -1010,6 +1010,66 @@ func (w *gatedRecorder) String() string { return w.got } +type closeUnblocksWrite struct { + once sync.Once + closed chan struct{} +} + +func (w *closeUnblocksWrite) Write([]byte) (int, error) { + <-w.closed + return 0, io.ErrClosedPipe +} + +func (w *closeUnblocksWrite) Close() error { + w.once.Do(func() { close(w.closed) }) + return nil +} + +func TestOverloadBlockedWriteServeExits(t *testing.T) { + inReader, inWriter := io.Pipe() + defer inWriter.Close() + out := &closeUnblocksWrite{closed: make(chan struct{})} + conn := NewOwnedConn(inReader, out) + conn.sem = make(chan struct{}, 1) + + started := make(chan struct{}) + conn.Handle("slow", func(ctx context.Context, _ json.RawMessage) (any, error) { + close(started) + <-ctx.Done() + return map[string]any{"ok": true}, nil + }) + + done := make(chan error, 1) + go func() { done <- conn.Serve(context.Background()) }() + + if _, err := inWriter.Write([]byte(`{"jsonrpc":"2.0","id":1,"method":"slow"}` + "\n")); err != nil { + t.Fatal(err) + } + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatal("handler did not start") + } + if _, err := inWriter.Write([]byte(`{"jsonrpc":"2.0","id":2,"method":"slow"}` + "\n")); err != nil { + t.Fatal(err) + } + for _, id := range []int{3, 4} { + frame := fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":"slow"}`+"\n", id) + if _, err := inWriter.Write([]byte(frame)); err != nil { + t.Fatal(err) + } + } + + select { + case err := <-done: + if !errors.Is(err, errBusyOverload) { + t.Fatalf("Serve = %v, want %v", err, errBusyOverload) + } + case <-time.After(time.Second): + t.Fatal("Serve must return after closing the writer; do not hang on Write") + } +} + func TestQueuedBusyReplySurvivesOverloadBurst(t *testing.T) { inReader, inWriter := io.Pipe() defer inWriter.Close() @@ -1064,18 +1124,6 @@ func TestQueuedBusyReplySurvivesOverloadBurst(t *testing.T) { } close(gate) - waitUntil := time.Now().Add(2 * time.Second) - var got string - for time.Now().Before(waitUntil) { - got = rec.String() - hasQueued := strings.Contains(got, `"id":3`) && strings.Contains(got, "-32000") - hasInflight := strings.Contains(got, `"id":2`) && strings.Contains(got, "-32000") - if hasQueued && hasInflight { - return - } - time.Sleep(10 * time.Millisecond) - } - t.Fatalf("busy ids 2 (in-flight) and 3 (queued) must receive -32000; wrote %q", got) } func TestInterleavedSessionCancelsDoNotCoalesceAcrossSessions(t *testing.T) {