From fe27fa91519a06197cdf03b49861f0565471e726 Mon Sep 17 00:00:00 2001 From: jansen Date: Fri, 21 Aug 2026 16:16:28 +0800 Subject: [PATCH] internal/jsonrpc2: plumb cancellation cause through request context (#1103) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1100. ## What `handleAsync` in `internal/jsonrpc2/conn.go` reported *why* an in-flight request had been canceled by inspecting the connection-level `s.writeErr`, guarded by a TODO waiting on golang/go#51365: ```go // Only deliver to the Handler if not already canceled. if err := req.ctx.Err(); err != nil { c.updateInFlight(func(s *inFlightState) { if s.writeErr != nil { // Assume that req.ctx was canceled due to s.writeErr. // TODO(#51365): use a Context API to plumb this through req.ctx. err = fmt.Errorf("%w: %v", ErrServerClosing, s.writeErr) } }) ... } ``` That API shipped as `context.WithCancelCause` / `context.Cause` in Go 1.21, and this module already targets Go 1.25, so the TODO is now actionable. This PR carries the cancellation cause on each request's context instead of inferring it from global state: - `incomingRequest.cancel` becomes a `context.CancelCauseFunc` (built via `context.WithCancelCause`). - The write-failure and read-side teardowns cancel each in-flight request with a cause wrapping `ErrServerClosing`. - The remaining cancel sites (peer cancel via `Cancel`, normal completion) pass `nil`. - `handleAsync` reads `context.Cause(req.ctx)` instead of looking up `s.writeErr`. ## Why Besides resolving the TODO, this fixes a latent mis-attribution. `s.writeErr` is connection-level state, not per-request: a request canceled for an unrelated reason (e.g. a peer cancel notification via `Cancel`) was reported as `ErrServerClosing` whenever `s.writeErr` happened to be non-nil at that moment. Tying the cause to the specific request's context removes that coupling. `ctx.Err()` still returns `context.Canceled` in all cases, so external behavior is unchanged; only the reported cause becomes accurate and per-request. ## Test plan - `go build ./...` - `go test ./internal/... ./mcp/` (including the longer `mcp` integration tests) — all pass. - Verified `context.WithCancelCause` semantics: with a cause, `Err()` stays `context.Canceled` while `Cause()` returns the wrapped `ErrServerClosing`; with `nil`, both resolve to `context.Canceled` — matching prior behavior. ## Note Per the discussion in #1100, the read-side teardown also carries a descriptive cause (wrapping `ErrServerClosing`) rather than a bare `context.Canceled`. --------- Co-authored-by: Guglielmo Colombo Co-authored-by: Yaroslav --- internal/jsonrpc2/conn.go | 31 ++++++++++++++++--------------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 4994c63..56edb7e 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -172,7 +172,7 @@ func (s *inFlightState) shuttingDown(errClosing error) error { type incomingRequest struct { *Request // the request being processed ctx context.Context - cancel context.CancelFunc + cancel context.CancelCauseFunc } // Reader abstracts the transport mechanics from the JSON RPC protocol. @@ -458,7 +458,7 @@ func (c *Connection) Cancel(id ID) { req = s.incomingByID[id] }) if req != nil { - req.cancel() + req.cancel(nil) } } @@ -552,9 +552,11 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter // Cancel any incoming requests still in flight: with the reader gone we // cannot receive cancellation notifications, and likely cannot write a // response either, so parked handlers have nothing useful left to do. - // Mirrors the equivalent cleanup on write failure. + // The read error (typically io.EOF from the peer disconnecting) is the + // cause, rather than a directional ErrServerClosing/ErrClientClosing + // sentinel: at this layer the connection has no designated end. for _, r := range s.incomingByID { - r.cancel() + r.cancel(err) } }) } @@ -564,7 +566,7 @@ func (c *Connection) readIncoming(ctx context.Context, reader Reader, preempter func (c *Connection) acceptRequest(ctx context.Context, msg *Request, preempter Preempter) { // In theory notifications cannot be cancelled, but we build them a cancel // context anyway. - reqCtx, cancel := context.WithCancel(ctx) + reqCtx, cancel := context.WithCancelCause(ctx) req := &incomingRequest{ Request: msg, ctx: reqCtx, @@ -665,15 +667,14 @@ func (c *Connection) handleAsync() { return } - // Only deliver to the Handler if not already canceled. + // Only deliver to the Handler if not already canceled. If the request + // was canceled with a cause (e.g. a write failure cancels every + // in-flight request), report that cause rather than the bare + // context.Canceled. if err := req.ctx.Err(); err != nil { - c.updateInFlight(func(s *inFlightState) { - if s.writeErr != nil { - // Assume that req.ctx was canceled due to s.writeErr. - // TODO(#51365): use a Context API to plumb this through req.ctx. - err = fmt.Errorf("%w: %v", ErrServerClosing, s.writeErr) - } - }) + if cause := context.Cause(req.ctx); cause != nil { + err = cause + } c.processResult("handleAsync", req, nil, err) continue } @@ -734,7 +735,7 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e } // Cancel the request to free any associated resources. - req.cancel() + req.cancel(nil) c.updateInFlight(func(s *inFlightState) { if s.incoming == 0 { panic("jsonrpc2: processResult called when incoming count is already zero") @@ -777,7 +778,7 @@ func (c *Connection) write(ctx context.Context, msg Message) error { if s.writeErr == nil { s.writeErr = err for _, r := range s.incomingByID { - r.cancel() + r.cancel(fmt.Errorf("%w: %v", ErrServerClosing, err)) } } })