diff --git a/internal/acp/agent_test.go b/internal/acp/agent_test.go index 4fa97a258..57701845e 100644 --- a/internal/acp/agent_test.go +++ b/internal/acp/agent_test.go @@ -9,6 +9,7 @@ import ( "os" "path/filepath" "strings" + "sync/atomic" "testing" "time" @@ -67,9 +68,11 @@ func testDeps(t *testing.T) Deps { // clientHarness wires a client Conn to an Agent over in-memory pipes and collects // session/update text chunks. type clientHarness struct { - client *Conn - updates chan string - stop func() + agent *Agent + agentConn *Conn + client *Conn + updates chan string + stop func() } func newHarness(t *testing.T, deps Deps) *clientHarness { @@ -80,7 +83,7 @@ func newHarness(t *testing.T, deps Deps) *clientHarness { client := NewConn(br, bw) a := NewAgent(agentConn, deps) - h := &clientHarness{client: client, updates: make(chan string, 128)} + h := &clientHarness{agent: a, agentConn: agentConn, client: client, updates: make(chan string, 128)} client.HandleNotify(MethodSessionUpdate, func(_ context.Context, params json.RawMessage) { var probe struct { Update struct { @@ -599,3 +602,148 @@ func drainTextUntil(t *testing.T, ch <-chan string, done func(string) bool) stri } } } + +func TestACPCancelInterleavedSessionsCancelsBothPrompts(t *testing.T) { + prompt1Entered := make(chan struct{}) + prompt2Entered := make(chan struct{}) + blockPrompt1 := make(chan struct{}) + blockPrompt2 := make(chan struct{}) + + firstCancelStarted := make(chan struct{}) + holdFirstCancel := make(chan struct{}) + var cancelCount atomic.Int32 + + deps := testDeps(t) + deps.RunAgent = func(ctx context.Context, prompt string, _ zeroruntime.Provider, opts agent.Options) (agent.Result, error) { + if opts.SessionID == "sess-1" { + close(prompt1Entered) + select { + case <-blockPrompt1: + return agent.Result{FinalAnswer: "done1"}, nil + case <-ctx.Done(): + return agent.Result{}, ctx.Err() + } + } else if opts.SessionID == "sess-2" { + close(prompt2Entered) + select { + case <-blockPrompt2: + return agent.Result{FinalAnswer: "done2"}, nil + case <-ctx.Done(): + return agent.Result{}, ctx.Err() + } + } + return agent.Result{}, nil + } + + h := newHarness(t, deps) + defer h.stop() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + var initRes InitializeResult + if err := h.client.Call(ctx, MethodInitialize, InitializeParams{ProtocolVersion: ProtocolVersion}, &initRes); err != nil { + t.Fatalf("initialize: %v", err) + } + + root := t.TempDir() + _, _ = deps.Store.Create(sessions.CreateInput{SessionID: "sess-1", Title: "s1", Cwd: root, ModelID: "fake-model"}) + _, _ = deps.Store.Create(sessions.CreateInput{SessionID: "sess-2", Title: "s2", Cwd: root, ModelID: "fake-model"}) + + var load1, load2 LoadSessionResult + if err := h.client.Call(ctx, MethodSessionLoad, LoadSessionParams{SessionID: "sess-1", Cwd: root}, &load1); err != nil { + t.Fatalf("load sess-1: %v", err) + } + if err := h.client.Call(ctx, MethodSessionLoad, LoadSessionParams{SessionID: "sess-2", Cwd: root}, &load2); err != nil { + t.Fatalf("load sess-2: %v", err) + } + + // Intercept cancel notification handler to delay the first cancel execution + origCancel := h.agent.conn.notifiers[MethodSessionCancel] + h.agent.conn.HandleNotify(MethodSessionCancel, func(ctx context.Context, params json.RawMessage) { + if cancelCount.Add(1) == 1 { + close(firstCancelStarted) + <-holdFirstCancel + } + origCancel(ctx, params) + }) + + type promptOutcome struct { + res PromptResult + err error + } + res1Ch := make(chan promptOutcome, 1) + res2Ch := make(chan promptOutcome, 1) + + go func() { + var res PromptResult + err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: "sess-1", Prompt: []ContentBlock{TextBlock("prompt1")}}, &res) + res1Ch <- promptOutcome{res: res, err: err} + }() + + go func() { + var res PromptResult + err := h.client.Call(ctx, MethodSessionPrompt, PromptParams{SessionID: "sess-2", Prompt: []ContentBlock{TextBlock("prompt2")}}, &res) + res2Ch <- promptOutcome{res: res, err: err} + }() + + select { + case <-prompt1Entered: + case <-time.After(2 * time.Second): + t.Fatal("prompt1 did not enter RunAgent") + } + select { + case <-prompt2Entered: + case <-time.After(2 * time.Second): + t.Fatal("prompt2 did not enter RunAgent") + } + + // Send cancel for sess-1 which starts the worker and waits on holdFirstCancel + _ = h.client.Notify(MethodSessionCancel, CancelParams{SessionID: "sess-1"}) + select { + case <-firstCancelStarted: + case <-time.After(2 * time.Second): + t.Fatal("first cancel did not start") + } + + // Interleave cancel for sess-2, then another cancel for sess-1 + _ = h.client.Notify(MethodSessionCancel, CancelParams{SessionID: "sess-2"}) + _ = h.client.Notify(MethodSessionCancel, CancelParams{SessionID: "sess-1"}) + + deadlineWait := time.Now().Add(2 * time.Second) + for { + h.agent.conn.notifyMu.Lock() + queued := string(h.agent.conn.notifyQ[notifyKey{method: MethodSessionCancel, target: "sess-1"}]) + h.agent.conn.notifyMu.Unlock() + if strings.Contains(queued, "sess-1") || time.Now().After(deadlineWait) { + break + } + time.Sleep(5 * time.Millisecond) + } + + close(holdFirstCancel) + + select { + case out1 := <-res1Ch: + if out1.err != nil { + t.Fatalf("prompt1 error: %v", out1.err) + } + if out1.res.StopReason != StopCancelled { + t.Fatalf("prompt1 StopReason = %q, want %q", out1.res.StopReason, StopCancelled) + } + case <-time.After(2 * time.Second): + t.Fatal("prompt1 did not cancel") + } + + select { + case out2 := <-res2Ch: + if out2.err != nil { + t.Fatalf("prompt2 error: %v", out2.err) + } + if out2.res.StopReason != StopCancelled { + t.Fatalf("prompt2 StopReason = %q, want %q", out2.res.StopReason, StopCancelled) + } + case <-time.After(2 * time.Second): + t.Fatal("prompt2 did not cancel (session cancel was overwritten/coalesced across sessions)") + } +} diff --git a/internal/acp/jsonrpc.go b/internal/acp/jsonrpc.go index 7d9171692..05c16c8ac 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. @@ -27,6 +28,12 @@ const ( codeMethodNotFound = -32601 codeInvalidParams = -32602 codeInternalError = -32603 + codeServerBusy = -32000 +) + +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. @@ -76,14 +83,29 @@ 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 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 + maxInflightBytes = 64 * 1024 * 1024 + maxBusyReplies = 1 + maxNotifyActive = 32 + maxUpdateFIFO = 32 + maxSpecialNotify = 256 + overloadDrainTimeout = 100 * time.Millisecond +) + // 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 // 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 @@ -95,7 +117,46 @@ type Conn struct { pending map[int64]chan rpcMessage closed bool - wg sync.WaitGroup // tracks in-flight inbound handlers + // 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 + + busyCh chan json.RawMessage + serveCancel context.CancelFunc + overloaded atomic.Bool + writeAbort chan struct{} + writeWaiters atomic.Int32 + + 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 { + 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 @@ -106,11 +167,19 @@ 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), + 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), } } @@ -123,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 } @@ -138,12 +216,30 @@ 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 + 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, // 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. 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() + if c.overloaded.Load() { + c.closeWriter() + done := make(chan struct{}) + go func() { + c.wg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(overloadDrainTimeout): + } + return + } c.wg.Wait() }() @@ -155,7 +251,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 @@ -167,13 +263,24 @@ 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 errors.Is(err, errFrameTooLarge) { + c.failAllPending(err) + return err + } 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 { @@ -184,6 +291,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("%w of %d bytes", errFrameTooLarge, 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. @@ -283,12 +412,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 } @@ -296,44 +425,233 @@ func (c *Conn) handleLine(ctx context.Context, line []byte) { case msg.isResponse(): c.deliver(msg) case msg.isRequest(): + n := int64(len(line)) + if !c.tryAdmit(n) { + id := append(json.RawMessage(nil), msg.ID...) + if !c.tryEnqueueBusy(id) { + c.tripOverload() + } + break + } c.wg.Add(1) - go func(m rpcMessage) { + go func(m rpcMessage, bytes int64) { defer c.wg.Done() + defer c.releaseAdmit(bytes) c.dispatchRequest(ctx, m) - }(msg) + }(msg, n) 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. 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"}) + } + } +} + +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...) + 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 + } + 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 { + fn(ctx, params) + c.notifyMu.Lock() + q := c.sessionUpdateQ[target] + if len(q) == 0 { + delete(c.updateOn, target) + delete(c.sessionUpdateQ, 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 + } + if c.specialNotifyBusyLocked() { + c.notifyMu.Unlock() + c.tripOverload() + 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 + c.notifyMu.Unlock() + return + } + if len(c.notifyOn) >= maxNotifyActive { + c.notifyMu.Unlock() + return + } + c.notifyOn[key] = true + c.notifyMu.Unlock() + c.wg.Add(1) + 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 { + fn(ctx, params) + c.notifyMu.Lock() + next, ok := c.notifyQ[key] + if !ok { + delete(c.notifyOn, key) + c.notifyMu.Unlock() + return } + delete(c.notifyQ, key) + c.notifyMu.Unlock() + params = next + } +} + +func (c *Conn) releaseSem() { + if c.sem == nil { + return + } + select { + case <-c.sem: + default: } } func (c *Conn) dispatchRequest(ctx context.Context, msg rpcMessage) { fn := c.handlers[msg.Method] + writeCtx := context.Background() if fn == nil { - c.writeError(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(msg.ID, re) + c.writeError(writeCtx, msg.ID, re) } else { - c.writeError(msg.ID, &rpcError{Code: codeInternalError, Message: err.Error()}) + c.writeError(writeCtx, msg.ID, &rpcError{Code: codeInternalError, Message: err.Error()}) } return } - c.writeResult(msg.ID, result) + c.writeResult(writeCtx, msg.ID, result) } // Call issues an outbound request and blocks until the response arrives, ctx is @@ -362,7 +680,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 } @@ -386,7 +704,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) { @@ -424,27 +742,144 @@ 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.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.writeMsg(ctx, rpcMessage{JSONRPC: "2.0", ID: id, Error: e}, c.overloaded.Load()) +} + +func (c *Conn) tryEnqueueBusy(id json.RawMessage) bool { + if c.overloaded.Load() || 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.write(rpcMessage{JSONRPC: "2.0", ID: id, Result: raw}) + c.failAllPending(errBusyOverload) + if c.writeAbort != nil { + close(c.writeAbort) + } + if c.serveCancel != nil { + c.serveCancel() + } } -func (c *Conn) writeError(id json.RawMessage, e *rpcError) { - _ = c.write(rpcMessage{JSONRPC: "2.0", ID: id, Error: e}) +func (c *Conn) writeBusyLoop(ctx context.Context) { + defer c.wg.Done() + busy := &rpcError{Code: codeServerBusy, Message: "server busy: max concurrent requests exceeded"} + for { + select { + case <-ctx.Done(): + return + case id := <-c.busyCh: + _ = c.writeMsg(ctx, rpcMessage{JSONRPC: "2.0", ID: id, Error: busy}, false) + } + } } -func (c *Conn) write(msg rpcMessage) error { +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() + case <-c.writeAbort: + releaseWait() + abandon() + return errBusyOverload + } + } + 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(): + 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 !persist && c.overloaded.Load() { + return errBusyOverload + } msg.JSONRPC = "2.0" data, err := json.Marshal(msg) if err != nil { return err } - c.writeMu.Lock() + if err := c.lockWrite(ctx, persist); err != nil { + return err + } defer c.writeMu.Unlock() + if !persist && 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 7d3b0c9da..4e0210594 100644 --- a/internal/acp/jsonrpc_test.go +++ b/internal/acp/jsonrpc_test.go @@ -1,10 +1,14 @@ package acp import ( + "bufio" "context" "encoding/json" "errors" + "fmt" "io" + "runtime" + "strings" "sync" "sync/atomic" "testing" @@ -424,6 +428,56 @@ 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")) + + 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) + } + + // 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) { @@ -485,3 +539,881 @@ 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 + 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() { + cancel() + _ = clientW.Close() + }() + + errCh := make(chan error, 1) + go func() { errCh <- server.Serve(ctx) }() + + // 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 || !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") + } +} + +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") + } + 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() + + 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") + } +} + +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 +} + +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() + + 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) +} + +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("custom/ping", 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":"custom/ping","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) +} + +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() + 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]) + } + } +} + +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) + })) + 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 + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { _ = conn.Serve(ctx) }() + + if _, err := inWriter.Write([]byte(frame1)); err != nil { + t.Fatal(err) + } + 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) + 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) + } +}