From a05c3fb8fd2c067a7df3639a5f799fca21cb486e 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 | 82 ++++++++++++++++++++--------------------- 1 file changed, 40 insertions(+), 42 deletions(-) diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index a0852d8b3..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 @@ -89,20 +94,14 @@ type Conn struct { handlers map[string]HandlerFunc notifiers map[string]NotifyFunc - // Notifications for one method are an ordered stream: session/update in - // particular is stateful, so dispatching every frame in an unrelated - // goroutine can make a later chunk overtake an earlier one. Different - // methods retain independent tails, keeping session/cancel responsive while - // an update handler is busy. - notifyMu sync.Mutex - notifyTails map[string]chan struct{} mu sync.Mutex nextID int64 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 @@ -113,12 +112,12 @@ 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), - notifyTails: make(map[string]chan struct{}), - pending: make(map[int64]chan rpcMessage), + 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), } } @@ -305,13 +304,21 @@ 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.dispatchNotification(ctx, msg, fn) + c.wg.Add(1) + c.acquireSem(ctx) + go func(m rpcMessage) { + defer c.releaseSem() + defer c.wg.Done() + fn(ctx, m.Params) + }(msg) } default: // Malformed frame; reply only if we can identify a request id. @@ -321,33 +328,24 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { } } -// dispatchNotification preserves wire order within one method without making -// unrelated notification methods wait for each other. The read loop installs -// each tail before starting its goroutine, so goroutine scheduling cannot -// reorder the chain it observes. -func (c *Conn) dispatchNotification(ctx context.Context, msg rpcMessage, fn NotifyFunc) { - c.notifyMu.Lock() - previous := c.notifyTails[msg.Method] - done := make(chan struct{}) - c.notifyTails[msg.Method] = done - c.notifyMu.Unlock() - - c.wg.Add(1) - go func() { - defer c.wg.Done() - defer func() { - close(done) - c.notifyMu.Lock() - if c.notifyTails[msg.Method] == done { - delete(c.notifyTails, msg.Method) - } - c.notifyMu.Unlock() - }() - if previous != nil { - <-previous - } - fn(ctx, msg.Params) - }() +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) { From ac709b31af2943f2178ff967a6496b3c1dee5d85 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 333854bf4..105be97a6 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" @@ -565,3 +567,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 9489762654f322d57874e0798548239c0404c67c 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 105be97a6..de237a699 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" @@ -506,6 +507,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 fef358aa8a031ac73e0357f45cfd697bd68b9387 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 de237a699..9cc60f465 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -688,7 +688,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() { @@ -699,16 +703,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 b6c85b0db7813f77c14b792858a8027c29638863 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 9563f8e4587e62b45ca4b71035719c66924dd030 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 9cc60f465..1ee00dbf8 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" @@ -782,3 +783,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 f8380b04941622dc9e23832e5c1ac6db6423d883 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 1ee00dbf8..d18359c6b 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -842,3 +842,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 14dab788ff3100e5c137959dae78d4a7ec93ed09 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 d18359c6b..cf41e5d89 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -534,8 +534,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) } @@ -837,12 +840,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 5f252d12f7cf20aaa9a45fe9a153143f629b6c1c 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 cf41e5d89..19b101640 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -929,3 +929,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 d2ec79c1d87cdc80c2fbccd851ff7a1f6ff9ee9d 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 | 146 ++++++++++++++++++++++++++++++++- internal/acp/jsonrpc.go | 69 ++++++++++++---- internal/acp/jsonrpc_test.go | 154 +++++++++++++++++++++++++++++++++++ 3 files changed, 353 insertions(+), 16 deletions(-) diff --git a/internal/acp/agent_test.go b/internal/acp/agent_test.go index bfce788bb..13c2ce93d 100644 --- a/internal/acp/agent_test.go +++ b/internal/acp/agent_test.go @@ -15,7 +15,11 @@ import ( "slices" "sort" "strings" +<<<<<<< HEAD "sync" +======= + "sync/atomic" +>>>>>>> 8d021507 (fix(acp): key cancels by session, bound notify workers, drain Serve on overload) "testing" "time" @@ -104,6 +108,7 @@ func testDeps(t *testing.T) Deps { // session/update text chunks. type clientHarness struct { agent *Agent + agentConn *Conn client *Conn updates chan string notifications chan ContentChunk @@ -122,7 +127,7 @@ func newHarness(t *testing.T, deps Deps) *clientHarness { client := NewConn(br, bw) a := NewAgent(agentConn, deps) - h := &clientHarness{agent: a, client: client, updates: make(chan string, 128), notifications: make(chan ContentChunk, 128), tools: make(chan ToolCallUpdate, 128)} + h := &clientHarness{agent: a, agentConn: agentConn, client: client, updates: make(chan string, 128), notifications: make(chan ContentChunk, 128), tools: make(chan ToolCallUpdate, 128)} client.HandleNotify(MethodSessionUpdate, func(_ context.Context, params json.RawMessage) { // Decode the discriminator ALONE first. ContentChunk and ToolCallUpdate // disagree on the shape of "content" (a block vs an array of them), so @@ -1149,6 +1154,7 @@ func drainTextUntil(t *testing.T, ch <-chan string, done func(string) bool) stri } } +<<<<<<< HEAD // normalisingResolver reproduces what ResolveWorkspaceRoot actually does — abs, // Clean, and a stat that the path exists — WITHOUT resolving symlinks or folding // case, which is the behaviour that makes two spellings of one directory produce @@ -2620,12 +2626,48 @@ func TestACPPromptUnsupportedModelDropsImagesAndNotifies(t *testing.T) { capturedPrompt = prompt capturedOpts = opts return agent.Result{FinalAnswer: "Answered without images"}, nil +======= +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 +>>>>>>> 8d021507 (fix(acp): key cancels by session, bound notify workers, drain Serve on overload) } h := newHarness(t, deps) defer h.stop() +<<<<<<< HEAD ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) +======= + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) +>>>>>>> 8d021507 (fix(acp): key cancels by session, bound notify workers, drain Serve on overload) defer cancel() var initRes InitializeResult @@ -2633,6 +2675,7 @@ func TestACPPromptUnsupportedModelDropsImagesAndNotifies(t *testing.T) { t.Fatalf("initialize: %v", err) } +<<<<<<< HEAD var newRes NewSessionResult if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil { t.Fatalf("session/new: %v", err) @@ -2801,5 +2844,106 @@ func TestACPEndToEndImagePromptSupportedModel(t *testing.T) { } if strings.Contains(capturedPrompt, "does not support image input") { t.Fatalf("prompt unexpectedly contains refusal note: %q", capturedPrompt) +======= + 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)") +>>>>>>> 8d021507 (fix(acp): key cancels by session, bound notify workers, drain Serve on overload) } } 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 19b101640..c14853478 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -1157,3 +1157,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 d30d2f48c8f830460ee6ef3a0d864301a905ce34 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 c14853478..40e2765e2 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -1282,7 +1282,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()) @@ -1290,7 +1290,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) } @@ -1311,3 +1311,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 a7b376d085e65d9416c43a34c071b27db6e8533d 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 40e2765e2..233344111 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -1349,8 +1349,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]) + } } } @@ -1396,10 +1402,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 }) @@ -1407,12 +1421,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 99e7e6b6ee676a3ccac22f984dedf6a770ceea8d 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 d968559b3b931d7746f71beed09b5ef15c0eb27d 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/agent_test.go | 89 +++++++++++++++++++----------------- internal/acp/jsonrpc.go | 42 ++++++++--------- internal/acp/jsonrpc_test.go | 72 ++++++++++++++++++++++++----- 3 files changed, 127 insertions(+), 76 deletions(-) diff --git a/internal/acp/agent_test.go b/internal/acp/agent_test.go index 13c2ce93d..3b9295326 100644 --- a/internal/acp/agent_test.go +++ b/internal/acp/agent_test.go @@ -15,11 +15,8 @@ import ( "slices" "sort" "strings" -<<<<<<< HEAD "sync" -======= "sync/atomic" ->>>>>>> 8d021507 (fix(acp): key cancels by session, bound notify workers, drain Serve on overload) "testing" "time" @@ -1154,7 +1151,6 @@ func drainTextUntil(t *testing.T, ch <-chan string, done func(string) bool) stri } } -<<<<<<< HEAD // normalisingResolver reproduces what ResolveWorkspaceRoot actually does — abs, // Clean, and a stat that the path exists — WITHOUT resolving symlinks or folding // case, which is the behaviour that makes two spellings of one directory produce @@ -2626,48 +2622,12 @@ func TestACPPromptUnsupportedModelDropsImagesAndNotifies(t *testing.T) { capturedPrompt = prompt capturedOpts = opts return agent.Result{FinalAnswer: "Answered without images"}, nil -======= -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 ->>>>>>> 8d021507 (fix(acp): key cancels by session, bound notify workers, drain Serve on overload) } h := newHarness(t, deps) defer h.stop() -<<<<<<< HEAD ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) -======= - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) ->>>>>>> 8d021507 (fix(acp): key cancels by session, bound notify workers, drain Serve on overload) defer cancel() var initRes InitializeResult @@ -2675,7 +2635,6 @@ func TestACPCancelInterleavedSessionsCancelsBothPrompts(t *testing.T) { t.Fatalf("initialize: %v", err) } -<<<<<<< HEAD var newRes NewSessionResult if err := h.client.Call(ctx, MethodSessionNew, NewSessionParams{Cwd: t.TempDir(), McpServers: []McpServer{}}, &newRes); err != nil { t.Fatalf("session/new: %v", err) @@ -2844,7 +2803,52 @@ func TestACPEndToEndImagePromptSupportedModel(t *testing.T) { } if strings.Contains(capturedPrompt, "does not support image input") { t.Fatalf("prompt unexpectedly contains refusal note: %q", capturedPrompt) -======= + } +} + +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"}) @@ -2944,6 +2948,5 @@ func TestACPEndToEndImagePromptSupportedModel(t *testing.T) { } case <-time.After(2 * time.Second): t.Fatal("prompt2 did not cancel (session cancel was overwritten/coalesced across sessions)") ->>>>>>> 8d021507 (fix(acp): key cancels by session, bound notify workers, drain Serve on overload) } } 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 233344111..0583b32e5 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -1090,6 +1090,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() @@ -1144,18 +1204,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) {