From 94025827e5248e880190675716f56eebd75524db Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20M=2E=20Requena=20Plens?= Date: Sun, 13 Sep 2026 11:18:08 +0200 Subject: [PATCH 1/6] mcp: send no response for a request the peer cancelled The MCP specification says a receiver of a cancellation notification SHOULD stop processing the request, free its resources and "not send a response for the cancelled request". The SDK always sends one. `processResult` in `internal/jsonrpc2/conn.go` writes the response with `c.write(notDone{req.ctx}, response)`, and `notDone` strips the cancellation from the context so the write goes ahead. Nothing at application level runs between a handler returning and that write, so no server built on this SDK could satisfy the clause. Captured on stdio: a client's `notifications/cancelled` at 6.306s, and the server's response to that same request id at 6.307s. `notDone` is there on purpose and this does not change it. A response must still be written when the handler's context ended for some other reason, and until now the jsonrpc2 layer could not tell the two apart: `Connection.Cancel` is called both by the preempter that saw the peer's notification and by `ServerSession.Close`, which cancels in-flight `subscriptions/listen` handlers to unblock them. So the distinction is recorded rather than guessed. `CancelFromPeer` marks the incoming request as cancelled by the peer, and `processResult` writes no response for such a call. `Cancel` keeps its meaning, and the listen result a session close produces is unaffected. The flag is written and read under the connection's `stateMu`, in the same critical section that removes the request from `incomingByID`, so a cancellation arriving once the response is already on its way finds nothing to mark, and no response is ever retracted. Suppressing the write alone would have traded a spec deviation for a hang. The streamable HTTP transport keeps a POST's stream open until every call it carried has been answered, so a response that never comes leaves that request open until the client goes away. A writer that holds per-call state can now implement `jsonrpc2.ResponseDropper` and be told the response is not coming; `streamableServerConn` implements it by retiring the request through the same accounting a real response goes through, so the stream completes exactly as it would have. Writers with no such state, stdio and the in-memory transport among them, implement nothing and are unaffected. A stream that ends with nothing written on it, because every call it carried was cancelled, answers the POST with 204 No Content: in JSON response mode there is nothing to flush, and in SSE mode no event was written, so in both the header is still unsent and an empty 200 under a Content-Type would claim a body that is not there. A stream that already carried an event or a buffered message ends as before. `TestStreamableCancelledCallGetsNoResponse` in `mcp` drives the whole path over raw HTTP, in SSE and in JSON response mode: initialize, a `tools/call` whose tool parks on `ctx.Done`, `notifications/cancelled` for that id on a second request, then a read of the call's stream to EOF. It asserts that the POST is answered 204 with no Content-Type and no body, and that it ends. On unmodified main it fails with the response the server sent: event: message data: {"jsonrpc":"2.0","id":2,"error":{"code":0,"message":"context canceled"}} It fakes the client with raw HTTP rather than using a `ClientSession` deliberately: an SDK client abandons the POST as soon as it cancels, so it never sees what the server wrote on that stream, and a stream that never completes looks to it exactly like one that did. With the suppression in place but `DropResponse` removed, the same test fails the other way, on the POST never returning. `TestCancelFromPeerSuppressesResponse` in `internal/jsonrpc2` pins the distinction itself. A second call acts as the barrier, since handlers run one at a time: a peer-cancelled call is answered only by the barrier's response and is reported to the dropper, while a locally cancelled one is still answered. Verification on go1.27.1: `gofmt -l .` clean, `go vet ./...` clean, `go test ./...` ok, `go test -race ./internal/jsonrpc2/ ./mcp/` ok, staticcheck clean. --- docs/protocol.md | 5 ++ internal/docs/protocol.src.md | 5 ++ internal/jsonrpc2/conn.go | 64 ++++++++++++-- internal/jsonrpc2/conn_test.go | 157 +++++++++++++++++++++++++++++++++ mcp/mcp_test.go | 4 + mcp/streamable.go | 54 +++++++++++- mcp/streamable_test.go | 103 +++++++++++++++++++++ mcp/transport.go | 14 ++- mcp/transport_test.go | 47 ++++++++++ 9 files changed, 439 insertions(+), 14 deletions(-) diff --git a/docs/protocol.md b/docs/protocol.md index 06c1b77b..3dbab9cb 100644 --- a/docs/protocol.md +++ b/docs/protocol.md @@ -712,6 +712,11 @@ When an RPC exits due to a cancellation error, there's a guarantee that the cancellation notification has been sent, but there's no guarantee that the server has observed it (see [concurrency](#concurrency)). +A receiver of a cancellation notification sends no response for the cancelled +request, as the spec requires. On the streamable HTTP transport the stream that +carried the request is still completed, so the POST it arrived on ends rather +than hanging, with 204 No Content when nothing else was written on it. + ```go func Example_cancellation() { // For this example, we're going to be collecting observations from the diff --git a/internal/docs/protocol.src.md b/internal/docs/protocol.src.md index 7d6b01b5..e367b629 100644 --- a/internal/docs/protocol.src.md +++ b/internal/docs/protocol.src.md @@ -607,6 +607,11 @@ When an RPC exits due to a cancellation error, there's a guarantee that the cancellation notification has been sent, but there's no guarantee that the server has observed it (see [concurrency](#concurrency)). +A receiver of a cancellation notification sends no response for the cancelled +request, as the spec requires. On the streamable HTTP transport the stream that +carried the request is still completed, so the POST it arrived on ends rather +than hanging, with 204 No Content when nothing else was written on it. + %include ../../mcp/mcp_example_test.go cancellation - ### Ping diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 94da6ae3..3bbbf6e4 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -173,6 +173,11 @@ type incomingRequest struct { *Request // the request being processed ctx context.Context cancel context.CancelCauseFunc + + // peerCancelled records that the peer asked for this request to be + // cancelled, as opposed to its context ending for another reason; set by + // [Connection.CancelFromPeer] and read by processResult under stateMu. + peerCancelled bool } // Reader abstracts the transport mechanics from the JSON RPC protocol. @@ -199,6 +204,15 @@ type Writer interface { Write(context.Context, Message) error } +// ResponseDropper is an optional interface for a [Writer] that holds state per +// incoming call, and so has to be told when a call the peer cancelled (see +// [Connection.CancelFromPeer]) will never be answered. +type ResponseDropper interface { + // DropResponse reports that the incoming call with the given ID is + // finished and will receive no response. + DropResponse(id ID) +} + // A ConnectionConfig configures a bidirectional jsonrpc2 connection. type ConnectionConfig struct { Reader Reader // required @@ -452,17 +466,37 @@ func (ac *AsyncCall) Await(ctx context.Context, result any) error { // Cancel will not complain if the ID is not a currently active message, and it // will not cause any messages that have not arrived yet with that ID to be // cancelled. +// +// The inbound call is still answered: use [Connection.CancelFromPeer] when the +// peer itself asked for the cancellation. func (c *Connection) Cancel(id ID) { - c.CancelCause(id, nil) + c.cancelIncoming(id, nil, false) } // CancelCause is like [Connection.Cancel], but records cause as the reason // the Context was cancelled, so that the Handle call can read it back through // [context.Cause]. A nil cause reads as [context.Canceled]. func (c *Connection) CancelCause(id ID, cause error) { + c.cancelIncoming(id, cause, false) +} + +// CancelFromPeer is [Connection.CancelCause] for a cancellation the peer +// requested, such as an MCP "notifications/cancelled": it also suppresses the +// response, which the MCP specification says a cancelled request gets none of. +func (c *Connection) CancelFromPeer(id ID, cause error) { + c.cancelIncoming(id, cause, true) +} + +// cancelIncoming cancels the inbound request with the given ID, recording the +// cause and whether the peer asked for it. A request already responded to is +// no longer in incomingByID, so a late cancellation retracts nothing. +func (c *Connection) cancelIncoming(id ID, cause error, fromPeer bool) { var req *incomingRequest c.updateInFlight(func(s *inFlightState) { req = s.incomingByID[id] + if req != nil && fromPeer { + req.peerCancelled = true + } }) if req != nil { req.cancel(cause) @@ -717,17 +751,29 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e // The caller could theoretically reuse the request's ID as soon as we've // sent the response, so ensure that it is removed from the incoming map - // before sending. + // before sending. Reading peerCancelled here keeps it atomic with that + // removal: a cancellation either arrives before this point and is + // honored, or finds the request gone and does nothing. + var peerCancelled bool c.updateInFlight(func(s *inFlightState) { + peerCancelled = req.peerCancelled delete(s.incomingByID, req.ID) }) - if respErr == nil { + if respErr != nil { + err = c.internalErrorf("%#v returned a malformed result for %q: %w", from, req.Method, respErr) + } + if peerCancelled { + // The peer cancelled this call, so it gets no response; a Writer that + // keeps state per call (a POST's stream stays open until every call + // in it is answered) is told the response is not coming. + if d, ok := c.writer.(ResponseDropper); ok { + d.DropResponse(req.ID) + } + } else if respErr == nil { writeErr := c.write(notDone{req.ctx}, response) if err == nil { err = writeErr } - } else { - err = c.internalErrorf("%#v returned a malformed result for %q: %w", from, req.Method, respErr) } } else { // req is a notification if result != nil { @@ -813,10 +859,10 @@ func (c *Connection) internalErrorf(format string, args ...any) error { // which by default is wrapped in notDone so a transport-level cancellation // does not implicitly cancel every in-flight handler. Cancellation of an // in-flight handler is instead expected to flow only through the jsonrpc2 -// layer's explicit channels: the [Preempter] calling [Connection.Cancel] in -// response to the peer's cancel notification, or the transport itself -// failing (the read loop exits on EOF or a write fails) — both of which -// cancel every in-flight incoming request in turn. +// layer's explicit channels: the [Preempter] calling +// [Connection.CancelFromPeer] in response to the peer's cancel notification, +// or the transport itself failing (the read loop exits on EOF or a write +// fails) — both of which cancel every in-flight incoming request in turn. type notDone struct{ ctx context.Context } func (ic notDone) Value(key any) any { diff --git a/internal/jsonrpc2/conn_test.go b/internal/jsonrpc2/conn_test.go index afeb4c02..60340846 100644 --- a/internal/jsonrpc2/conn_test.go +++ b/internal/jsonrpc2/conn_test.go @@ -5,8 +5,10 @@ package jsonrpc2 import ( + "context" "errors" "io" + "sync" "testing" ) @@ -36,3 +38,158 @@ func TestShuttingDownWrapsReadError(t *testing.T) { t.Errorf("shuttingDown() error = %v, want it to wrap io.EOF", err) } } + +// errPeerAskedToStop stands in for the reason an MCP "notifications/cancelled" +// carries, so the test exercises the signature a real peer cancellation uses +// rather than a nil cause no caller passes. +var errPeerAskedToStop = errors.New("peer asked to stop") + +// TestCancelFromPeerSuppressesResponse verifies that a call the peer asked to +// cancel receives no response, while a call cancelled locally still does; the +// barrier is a second call, which handlers answer only after the first one. +func TestCancelFromPeerSuppressesResponse(t *testing.T) { + tests := []struct { + name string + fromPeer bool + want []ID // the IDs responded to, in order + }{ + {"peer cancellation", true, []ID{Int64ID(2)}}, + {"local cancellation", false, []ID{Int64ID(1), Int64ID(2)}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + incoming := make(chan Message, 2) + writer := &recordingWriter{written: make(chan *Response, 2)} + + started := make(chan struct{}) + handler := HandlerFunc(func(ctx context.Context, req *Request) (any, error) { + if req.Method == "barrier" { + return struct{}{}, nil + } + close(started) + <-ctx.Done() + return nil, context.Cause(ctx) + }) + + conn := NewConnection(context.Background(), ConnectionConfig{ + Reader: &channelReader{messages: incoming}, + Writer: writer, + Closer: &channelCloser{messages: incoming}, + Bind: func(*Connection) Handler { return handler }, + OnDone: func() {}, + OnInternalError: func(err error) { + t.Errorf("internal error: %v", err) + }, + }) + + slow, err := NewCall(Int64ID(1), "slow", nil) + if err != nil { + t.Fatal(err) + } + barrier, err := NewCall(Int64ID(2), "barrier", nil) + if err != nil { + t.Fatal(err) + } + + incoming <- slow + <-started + if test.fromPeer { + conn.CancelFromPeer(Int64ID(1), errPeerAskedToStop) + } else { + conn.Cancel(Int64ID(1)) + } + incoming <- barrier + + var got []ID + for range test.want { + got = append(got, (<-writer.written).ID) + } + if !equalIDs(got, test.want) { + t.Errorf("responded to %v, want %v", got, test.want) + } + + var wantDropped []ID + if test.fromPeer { + wantDropped = []ID{Int64ID(1)} + } + if dropped := writer.dropped(); !equalIDs(dropped, wantDropped) { + t.Errorf("dropped %v, want %v", dropped, wantDropped) + } + + if err := conn.Close(); err != nil { + t.Errorf("Close() = %v", err) + } + }) + } +} + +// equalIDs reports whether two ID slices hold the same IDs in the same order. +func equalIDs(got, want []ID) bool { + if len(got) != len(want) { + return false + } + for i := range got { + if got[i].Raw() != want[i].Raw() { + return false + } + } + return true +} + +// channelReader delivers the messages sent to it, and reports io.EOF once the +// channel is closed. +type channelReader struct { + messages chan Message +} + +func (r *channelReader) Read(ctx context.Context) (Message, error) { + select { + case msg, ok := <-r.messages: + if !ok { + return nil, io.EOF + } + return msg, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +// channelCloser unblocks a [channelReader] reading from the same channel. +type channelCloser struct { + once sync.Once + messages chan Message +} + +func (c *channelCloser) Close() error { + c.once.Do(func() { close(c.messages) }) + return nil +} + +// recordingWriter reports the responses a Connection writes, and the calls it +// is told will get none. +type recordingWriter struct { + written chan *Response + + mu sync.Mutex + drop []ID +} + +func (w *recordingWriter) Write(_ context.Context, msg Message) error { + if resp, ok := msg.(*Response); ok { + w.written <- resp + } + return nil +} + +// DropResponse implements [ResponseDropper]. +func (w *recordingWriter) DropResponse(id ID) { + w.mu.Lock() + defer w.mu.Unlock() + w.drop = append(w.drop, id) +} + +func (w *recordingWriter) dropped() []ID { + w.mu.Lock() + defer w.mu.Unlock() + return append([]ID(nil), w.drop...) +} diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index ccf80b6c..14d5cb55 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -787,6 +787,10 @@ func TestCancellationReason(t *testing.T) { if err := cs.conn.Notify(ctx, notificationCancelled, &CancelledParams{RequestID: call.ID().Raw(), Reason: "user asked"}); err != nil { t.Fatal(err) } + // The receiver sends no response to a cancelled call, so the SDK's own + // cancel path (cancelCall) retires the outgoing call before notifying; + // this test notifies by hand and retires by hand, or synctest deadlocks. + cs.conn.Retire(call, context.Canceled) cause := <-cancelled if cause == nil { diff --git a/mcp/streamable.go b/mcp/streamable.go index 9f8f2234..483d1cfd 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -1237,6 +1237,9 @@ func extractErrorStatus(ctx context.Context, msg jsonrpc.Message) int { // pendingJSONMessages (for JSON mode). The eventID is used for SSE event ID; // pass "" to omit. // +// If data is nil there is nothing to write: the call is only accounting for a +// request that will never be answered (see [streamableServerConn.DropResponse]). +// // If responseTo is valid, it is removed from the requests map. When all // requests have been responded to, the done channel is closed and set to nil. // @@ -1283,8 +1286,12 @@ func (s *stream) deliverLocked(data []byte, eventID string, responseTo jsonrpc.I // there's a brief race between request cancellation and releasing the // stream. if s.pendingJSONMessages != nil { - s.pendingJSONMessages = append(s.pendingJSONMessages, data) - if done { + if data != nil { + s.pendingJSONMessages = append(s.pendingJSONMessages, data) + } + if done && len(s.pendingJSONMessages) == 0 { + s.writeNoContentLocked() + } else if done { // Flush all pending messages as JSON response. var toWrite []byte if len(s.pendingJSONMessages) == 1 && !s.isBatch { @@ -1299,17 +1306,28 @@ func (s *stream) deliverLocked(data []byte, eventID string, responseTo jsonrpc.I return done, err } } - } else { + } else if data != nil { // SSE mode: write event to response writer. s.lastIdx++ if _, err := writeEvent(s.w, Event{Name: "message", Data: data, ID: eventID}); err != nil { return done, err } s.markWrittenLocked() + } else if done && s.lastWrite.IsZero() { + // Nothing was ever written, so the header is still ours to set. + s.writeNoContentLocked() } return done, nil } +// writeNoContentLocked ends a stream that carried nothing: every request the +// POST brought was cancelled before it was answered, so there is no body and +// the status says so instead of an empty 200 under a Content-Type. +func (s *stream) writeNoContentLocked() { + s.w.Header().Del("Content-Type") + s.w.WriteHeader(http.StatusNoContent) +} + // doneLocked reports whether the stream is logically complete. // // s.requests was populated when reading the POST body, requests are deleted as @@ -2141,6 +2159,36 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e return nil } +// DropResponse implements [jsonrpc2.ResponseDropper]: the cancelled call gets +// no response, but a POST's stream hangs until every call it carried has been +// answered, so the call is retired from the stream all the same. +func (c *streamableServerConn) DropResponse(id jsonrpc.ID) { + c.mu.Lock() + var s *stream + if streamID, ok := c.requestStreams[id]; ok { + s = c.streams[streamID] + } + delete(c.requestStreams, id) + c.mu.Unlock() + + if s == nil { + return + } + + s.mu.Lock() + // A nil payload delivers nothing; it only retires the request. An error + // here means the stream is already disconnected, which is not a problem + // when there is nothing to send. + done, _ := s.deliverLocked(nil, "", id, 0) + s.mu.Unlock() + + if done { + c.mu.Lock() + delete(c.streams, s.id) + c.mu.Unlock() + } +} + // Close implements the [Connection] interface. func (c *streamableServerConn) Close() error { c.mu.Lock() diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 7d4782bf..8c5b6e18 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -3752,6 +3752,109 @@ func TestCallCancellation_FastReturn(t *testing.T) { } } +// TestStreamableCancelledCallGetsNoResponse checks that a call the client +// cancelled with notifications/cancelled is answered with 204 and no body, in +// SSE and in JSON mode, through raw HTTP since an SDK client abandons the POST. +func TestStreamableCancelledCallGetsNoResponse(t *testing.T) { + for _, mode := range []struct { + name string + jsonResponse bool + }{ + {name: "sse", jsonResponse: false}, + {name: "json", jsonResponse: true}, + } { + t.Run(mode.name, func(t *testing.T) { + started := make(chan struct{}) + server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil) + server.AddTool( + &Tool{Name: "slow", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { + close(started) + <-ctx.Done() + return nil, ctx.Err() + }) + + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{JSONResponse: mode.jsonResponse}) + defer handler.closeAll() + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + defer httpServer.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + post := func(sessionID string, msg jsonrpc.Message) (*http.Response, error) { + data, err := jsonrpc2.EncodeMessage(msg) + if err != nil { + return nil, err + } + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, httpServer.URL, bytes.NewReader(data)) + if err != nil { + return nil, err + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "application/json, text/event-stream") + if sessionID != "" { + httpReq.Header.Set(sessionIDHeader, sessionID) + } + return http.DefaultClient.Do(httpReq) + } + mustPost := func(sessionID string, msg jsonrpc.Message) *http.Response { + t.Helper() + resp, err := post(sessionID, msg) + if err != nil { + t.Fatal(err) + } + return resp + } + + initResp := mustPost("", req(1, methodInitialize, &InitializeParams{ProtocolVersion: protocolVersion20250618})) + sessionID := initResp.Header.Get(sessionIDHeader) + initResp.Body.Close() + if sessionID == "" { + t.Fatal("initialize response carried no session ID") + } + mustPost(sessionID, req(0, notificationInitialized, &InitializedParams{})).Body.Close() + + // The POST returns only when its stream ends, so call from a goroutine. + type postResult struct { + resp *http.Response + err error + } + call := make(chan postResult, 1) + go func() { + resp, err := post(sessionID, req(2, "tools/call", &CallToolParams{Name: "slow"})) + call <- postResult{resp, err} + }() + <-started + + mustPost(sessionID, req(0, notificationCancelled, &CancelledParams{ + RequestID: int64(2), + Reason: "test cancellation", + })).Body.Close() + + got := <-call + if got.err != nil { + t.Fatalf("the cancelled call's POST: %v", got.err) + } + defer got.resp.Body.Close() + + if got.resp.StatusCode != http.StatusNoContent { + t.Errorf("the cancelled call's POST status = %d, want %d", got.resp.StatusCode, http.StatusNoContent) + } + if ct := got.resp.Header.Get("Content-Type"); ct != "" { + t.Errorf("the cancelled call's POST carried Content-Type %q, want none", ct) + } + body, err := io.ReadAll(got.resp.Body) + if err != nil { + t.Fatalf("reading the cancelled call's stream: %v", err) + } + if len(body) > 0 { + t.Errorf("the cancelled call's stream carried:\n%s\nwant nothing", body) + } + }) + } +} + // TestStreamableStateless_AcceptsNewProtocol is the positive control: // confirms that a stateless server still accepts new-protocol requests // (the rejection in TestStreamableStateful_RejectsNewProtocol must not diff --git a/mcp/transport.go b/mcp/transport.go index 08d3b15f..335e34ea 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -250,7 +250,8 @@ type cancellationPropagator interface { } // A canceller is a jsonrpc2.Preempter that cancels in-flight requests on MCP -// cancelled notifications. +// cancelled notifications. The cancelled request is answered with no response +// at all, as the spec requires. type canceller struct { conn *jsonrpc2.Connection logger *slog.Logger @@ -272,7 +273,7 @@ func (c *canceller) Preempt(ctx context.Context, req *jsonrpc.Request) (result a // travels as the cause of the request's context rather than being // dropped here. c.logger.Debug("request cancelled by the peer", "id", id.Raw(), "reason", params.Reason) - go c.conn.CancelCause(id, &peerCancelledError{reason: params.Reason}) + go c.conn.CancelFromPeer(id, &peerCancelledError{reason: params.Reason}) } return nil, jsonrpc2.ErrNotHandled } @@ -440,6 +441,15 @@ func (s *loggingConn) Write(ctx context.Context, msg jsonrpc.Message) error { return err } +// DropResponse implements [jsonrpc2.ResponseDropper] by forwarding to the +// delegate: the connection asks its writer for the interface, and a wrapper +// without it would silently keep the delegate waiting for the response. +func (s *loggingConn) DropResponse(id jsonrpc.ID) { + if dropper, ok := s.delegate.(jsonrpc2.ResponseDropper); ok { + dropper.DropResponse(id) + } +} + func (s *loggingConn) Close() error { return s.delegate.Close() } diff --git a/mcp/transport_test.go b/mcp/transport_test.go index dadf196b..00c33230 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -364,3 +364,50 @@ func TestIOConnFrameCap(t *testing.T) { }) } } + +// droppingConn is a Connection that records the ids it was told will receive +// no response. It implements [jsonrpc2.ResponseDropper]; plainConn, below, +// deliberately does not. +type droppingConn struct { + dropped []jsonrpc.ID +} + +func (c *droppingConn) SessionID() string { return "" } +func (c *droppingConn) Read(context.Context) (jsonrpc.Message, error) { return nil, io.EOF } +func (c *droppingConn) Write(context.Context, jsonrpc.Message) error { return nil } +func (c *droppingConn) Close() error { return nil } +func (c *droppingConn) DropResponse(id jsonrpc.ID) { c.dropped = append(c.dropped, id) } + +// plainConn is a Connection that holds no per-call state, so it implements no +// ResponseDropper. +type plainConn struct{} + +func (plainConn) SessionID() string { return "" } +func (plainConn) Read(context.Context) (jsonrpc.Message, error) { return nil, io.EOF } +func (plainConn) Write(context.Context, jsonrpc.Message) error { return nil } +func (plainConn) Close() error { return nil } + +// TestLoggingConnDropResponse checks that wrapping a transport for logging +// does not lose the "this call gets no response" signal: loggingConn is the +// one Connection wrapper here, and it used to leave the POST stream open. +func TestLoggingConnDropResponse(t *testing.T) { + t.Run("forwards to a delegate that holds per-call state", func(t *testing.T) { + delegate := &droppingConn{} + conn := &loggingConn{delegate: delegate, w: io.Discard} + + dropper, ok := any(conn).(jsonrpc2.ResponseDropper) + if !ok { + t.Fatal("loggingConn does not implement jsonrpc2.ResponseDropper, so a cancelled call's stream is never released") + } + dropper.DropResponse(jsonrpc.ID{}) + + if len(delegate.dropped) != 1 { + t.Errorf("the delegate was told about %d dropped responses, want 1", len(delegate.dropped)) + } + }) + + t.Run("does nothing for a delegate that holds none", func(t *testing.T) { + conn := &loggingConn{delegate: plainConn{}, w: io.Discard} + conn.DropResponse(jsonrpc.ID{}) + }) +} From 436501f274b9cbffd620c2b4df561e0b09e5873d Mon Sep 17 00:00:00 2001 From: guglielmoc Date: Fri, 9 Oct 2026 12:24:50 +0000 Subject: [PATCH 2/6] feat: implement proper handling for cancelled MCP calls in streamable transports --- docs/protocol.md | 5 +- internal/docs/protocol.src.md | 5 +- internal/jsonrpc2/conn.go | 63 +++---- internal/jsonrpc2/conn_test.go | 179 ++++++++++++------ mcp/streamable.go | 78 ++++---- mcp/streamable_test.go | 325 +++++++++++++++++++++++++++++++++ mcp/transport.go | 56 ++++-- mcp/transport_test.go | 165 ++++++++++++++++- 8 files changed, 725 insertions(+), 151 deletions(-) diff --git a/docs/protocol.md b/docs/protocol.md index 3dbab9cb..f092c600 100644 --- a/docs/protocol.md +++ b/docs/protocol.md @@ -714,8 +714,9 @@ server has observed it (see [concurrency](#concurrency)). A receiver of a cancellation notification sends no response for the cancelled request, as the spec requires. On the streamable HTTP transport the stream that -carried the request is still completed, so the POST it arrived on ends rather -than hanging, with 204 No Content when nothing else was written on it. +carried the request is still completed as soon as the notification arrives, +without waiting for the handler to return, so the POST it arrived on ends +rather than hanging, with 204 No Content when nothing else was written on it. ```go func Example_cancellation() { diff --git a/internal/docs/protocol.src.md b/internal/docs/protocol.src.md index e367b629..4a339def 100644 --- a/internal/docs/protocol.src.md +++ b/internal/docs/protocol.src.md @@ -609,8 +609,9 @@ server has observed it (see [concurrency](#concurrency)). A receiver of a cancellation notification sends no response for the cancelled request, as the spec requires. On the streamable HTTP transport the stream that -carried the request is still completed, so the POST it arrived on ends rather -than hanging, with 204 No Content when nothing else was written on it. +carried the request is still completed as soon as the notification arrives, +without waiting for the handler to return, so the POST it arrived on ends +rather than hanging, with 204 No Content when nothing else was written on it. %include ../../mcp/mcp_example_test.go cancellation - diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 3bbbf6e4..ec68c532 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -173,11 +173,6 @@ type incomingRequest struct { *Request // the request being processed ctx context.Context cancel context.CancelCauseFunc - - // peerCancelled records that the peer asked for this request to be - // cancelled, as opposed to its context ending for another reason; set by - // [Connection.CancelFromPeer] and read by processResult under stateMu. - peerCancelled bool } // Reader abstracts the transport mechanics from the JSON RPC protocol. @@ -204,12 +199,14 @@ type Writer interface { Write(context.Context, Message) error } -// ResponseDropper is an optional interface for a [Writer] that holds state per -// incoming call, and so has to be told when a call the peer cancelled (see -// [Connection.CancelFromPeer]) will never be answered. +// ResponseDropper is an optional interface for a [Writer] that can drop the +// response to a call the peer cancelled (see [Connection.CancelFromPeer]). The +// Connection still writes the response once the call's handler returns, so +// that the writer sees the call finish, and the writer discards it. type ResponseDropper interface { - // DropResponse reports that the incoming call with the given ID is - // finished and will receive no response. + // DropResponse reports, when the cancellation arrives, that the peer + // cancelled the incoming call with the given ID, so its response must not + // be sent. DropResponse(id ID) } @@ -481,26 +478,34 @@ func (c *Connection) CancelCause(id ID, cause error) { } // CancelFromPeer is [Connection.CancelCause] for a cancellation the peer -// requested, such as an MCP "notifications/cancelled": it also suppresses the -// response, which the MCP specification says a cancelled request gets none of. +// requested, such as an MCP "notifications/cancelled": it also tells a writer +// that is a [ResponseDropper] to drop the response, which the MCP +// specification says a cancelled request gets none of. func (c *Connection) CancelFromPeer(id ID, cause error) { c.cancelIncoming(id, cause, true) } // cancelIncoming cancels the inbound request with the given ID, recording the -// cause and whether the peer asked for it. A request already responded to is -// no longer in incomingByID, so a late cancellation retracts nothing. +// cause, and tells the writer to drop the response if the peer asked for the +// cancellation. A request already responded to is no longer in incomingByID, +// so a late cancellation retracts nothing. func (c *Connection) cancelIncoming(id ID, cause error, fromPeer bool) { var req *incomingRequest c.updateInFlight(func(s *inFlightState) { req = s.incomingByID[id] - if req != nil && fromPeer { - req.peerCancelled = true - } }) - if req != nil { - req.cancel(cause) + if req == nil { + return + } + if fromPeer { + // Tell the writer before cancelling, so that it knows by the time a + // handler that honors its context returns. This runs outside stateMu + // since the writer may do I/O. + if d, ok := c.writer.(ResponseDropper); ok { + d.DropResponse(id) + } } + req.cancel(cause) } // Wait blocks until the connection is fully closed, but does not close it. @@ -751,29 +756,17 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e // The caller could theoretically reuse the request's ID as soon as we've // sent the response, so ensure that it is removed from the incoming map - // before sending. Reading peerCancelled here keeps it atomic with that - // removal: a cancellation either arrives before this point and is - // honored, or finds the request gone and does nothing. - var peerCancelled bool + // before sending. c.updateInFlight(func(s *inFlightState) { - peerCancelled = req.peerCancelled delete(s.incomingByID, req.ID) }) - if respErr != nil { - err = c.internalErrorf("%#v returned a malformed result for %q: %w", from, req.Method, respErr) - } - if peerCancelled { - // The peer cancelled this call, so it gets no response; a Writer that - // keeps state per call (a POST's stream stays open until every call - // in it is answered) is told the response is not coming. - if d, ok := c.writer.(ResponseDropper); ok { - d.DropResponse(req.ID) - } - } else if respErr == nil { + if respErr == nil { writeErr := c.write(notDone{req.ctx}, response) if err == nil { err = writeErr } + } else { + err = c.internalErrorf("%#v returned a malformed result for %q: %w", from, req.Method, respErr) } } else { // req is a notification if result != nil { diff --git a/internal/jsonrpc2/conn_test.go b/internal/jsonrpc2/conn_test.go index 60340846..c1818985 100644 --- a/internal/jsonrpc2/conn_test.go +++ b/internal/jsonrpc2/conn_test.go @@ -44,76 +44,62 @@ func TestShuttingDownWrapsReadError(t *testing.T) { // rather than a nil cause no caller passes. var errPeerAskedToStop = errors.New("peer asked to stop") -// TestCancelFromPeerSuppressesResponse verifies that a call the peer asked to -// cancel receives no response, while a call cancelled locally still does; the -// barrier is a second call, which handlers answer only after the first one. -func TestCancelFromPeerSuppressesResponse(t *testing.T) { +// TestCancelFromPeerTellsWriter verifies that a writer that drops responses is +// told about a call the peer cancelled before the handler's context is +// cancelled, and not about a call cancelled locally. The response is written +// either way; dropping it is up to the writer. The barrier is a second call, +// which handlers answer only after the first one. +func TestCancelFromPeerTellsWriter(t *testing.T) { tests := []struct { - name string - fromPeer bool - want []ID // the IDs responded to, in order + name string + fromPeer bool + dropper bool // whether the writer is a ResponseDropper + want []ID // the IDs of the responses written, in order + wantDrops []ID // the IDs DropResponse is called with }{ - {"peer cancellation", true, []ID{Int64ID(2)}}, - {"local cancellation", false, []ID{Int64ID(1), Int64ID(2)}}, + {"peer cancellation, plain writer", true, false, []ID{Int64ID(1), Int64ID(2)}, nil}, + {"peer cancellation, writer drops the response", true, true, []ID{Int64ID(1), Int64ID(2)}, []ID{Int64ID(1)}}, + {"local cancellation", false, true, []ID{Int64ID(1), Int64ID(2)}, nil}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - incoming := make(chan Message, 2) - writer := &recordingWriter{written: make(chan *Response, 2)} - + rec := &recordingWriter{written: make(chan *Response, 2)} + dropper := &dropperWriter{recordingWriter: rec} + var writer Writer = rec + if test.dropper { + writer = dropper + } started := make(chan struct{}) - handler := HandlerFunc(func(ctx context.Context, req *Request) (any, error) { - if req.Method == "barrier" { - return struct{}{}, nil - } + dropsAtCancel := make(chan []ID, 1) + conn, incoming := newCancelTestConn(t, writer, func(ctx context.Context) (any, error) { close(started) <-ctx.Done() + dropsAtCancel <- dropper.drops() return nil, context.Cause(ctx) }) - conn := NewConnection(context.Background(), ConnectionConfig{ - Reader: &channelReader{messages: incoming}, - Writer: writer, - Closer: &channelCloser{messages: incoming}, - Bind: func(*Connection) Handler { return handler }, - OnDone: func() {}, - OnInternalError: func(err error) { - t.Errorf("internal error: %v", err) - }, - }) - - slow, err := NewCall(Int64ID(1), "slow", nil) - if err != nil { - t.Fatal(err) - } - barrier, err := NewCall(Int64ID(2), "barrier", nil) - if err != nil { - t.Fatal(err) - } - - incoming <- slow + incoming <- mustCall(t, 1, "slow") <-started if test.fromPeer { conn.CancelFromPeer(Int64ID(1), errPeerAskedToStop) } else { conn.Cancel(Int64ID(1)) } - incoming <- barrier + incoming <- mustCall(t, 2, "barrier") var got []ID for range test.want { - got = append(got, (<-writer.written).ID) + got = append(got, (<-rec.written).ID) } if !equalIDs(got, test.want) { - t.Errorf("responded to %v, want %v", got, test.want) + t.Errorf("responses written for %v, want %v", got, test.want) } - - var wantDropped []ID - if test.fromPeer { - wantDropped = []ID{Int64ID(1)} + // The writer is told before the handler's context is cancelled. + if drops := <-dropsAtCancel; !equalIDs(drops, test.wantDrops) { + t.Errorf("when the handler's context was cancelled, DropResponse called for %v, want %v", drops, test.wantDrops) } - if dropped := writer.dropped(); !equalIDs(dropped, wantDropped) { - t.Errorf("dropped %v, want %v", dropped, wantDropped) + if drops := dropper.drops(); !equalIDs(drops, test.wantDrops) { + t.Errorf("DropResponse called for %v, want %v", drops, test.wantDrops) } if err := conn.Close(); err != nil { @@ -123,6 +109,83 @@ func TestCancelFromPeerSuppressesResponse(t *testing.T) { } } +// TestCancelFromPeerHandlerIgnoresContext verifies that the writer is told of +// a peer cancellation as soon as it arrives, even when the handler ignores its +// context, that it still receives the response it drops once the handler +// returns, and that a cancellation arriving after that is ignored. +func TestCancelFromPeerHandlerIgnoresContext(t *testing.T) { + rec := &recordingWriter{written: make(chan *Response, 2)} + dropper := &dropperWriter{recordingWriter: rec} + started := make(chan struct{}) + release := make(chan struct{}) + conn, incoming := newCancelTestConn(t, dropper, func(context.Context) (any, error) { + close(started) + <-release + return struct{}{}, nil + }) + + incoming <- mustCall(t, 1, "slow") + <-started + conn.CancelFromPeer(Int64ID(1), errPeerAskedToStop) + + if drops := dropper.drops(); !equalIDs(drops, []ID{Int64ID(1)}) { + t.Errorf("after the cancellation, DropResponse called for %v, want [1]", drops) + } + + close(release) + incoming <- mustCall(t, 2, "barrier") + var got []ID + for range 2 { + got = append(got, (<-rec.written).ID) + } + if want := []ID{Int64ID(1), Int64ID(2)}; !equalIDs(got, want) { + t.Errorf("responses written for %v, want %v", got, want) + } + conn.CancelFromPeer(Int64ID(1), errPeerAskedToStop) + if drops := dropper.drops(); !equalIDs(drops, []ID{Int64ID(1)}) { + t.Errorf("after a late cancellation, DropResponse called for %v, want [1]", drops) + } + + if err := conn.Close(); err != nil { + t.Errorf("Close() = %v", err) + } +} + +// newCancelTestConn returns a Connection reading from the returned channel +// and writing to writer. Its handler answers "barrier" at once and runs slow +// for any other method. Handlers run one at a time, so once the barrier is +// answered, the slow call is finished. +func newCancelTestConn(t *testing.T, writer Writer, slow func(context.Context) (any, error)) (*Connection, chan Message) { + t.Helper() + incoming := make(chan Message, 2) + handler := HandlerFunc(func(ctx context.Context, req *Request) (any, error) { + if req.Method == "barrier" { + return struct{}{}, nil + } + return slow(ctx) + }) + conn := NewConnection(context.Background(), ConnectionConfig{ + Reader: &channelReader{messages: incoming}, + Writer: writer, + Closer: &channelCloser{messages: incoming}, + Bind: func(*Connection) Handler { return handler }, + OnDone: func() {}, + OnInternalError: func(err error) { + t.Errorf("internal error: %v", err) + }, + }) + return conn, incoming +} + +func mustCall(t *testing.T, id int64, method string) *Request { + t.Helper() + call, err := NewCall(Int64ID(id), method, nil) + if err != nil { + t.Fatal(err) + } + return call +} + // equalIDs reports whether two ID slices hold the same IDs in the same order. func equalIDs(got, want []ID) bool { if len(got) != len(want) { @@ -165,13 +228,9 @@ func (c *channelCloser) Close() error { return nil } -// recordingWriter reports the responses a Connection writes, and the calls it -// is told will get none. +// recordingWriter reports the responses a Connection writes. type recordingWriter struct { written chan *Response - - mu sync.Mutex - drop []ID } func (w *recordingWriter) Write(_ context.Context, msg Message) error { @@ -181,15 +240,23 @@ func (w *recordingWriter) Write(_ context.Context, msg Message) error { return nil } -// DropResponse implements [ResponseDropper]. -func (w *recordingWriter) DropResponse(id ID) { +// dropperWriter is a recordingWriter that is also a [ResponseDropper], +// recording the IDs it is told about. +type dropperWriter struct { + *recordingWriter + + mu sync.Mutex + dropped []ID +} + +func (w *dropperWriter) DropResponse(id ID) { w.mu.Lock() defer w.mu.Unlock() - w.drop = append(w.drop, id) + w.dropped = append(w.dropped, id) } -func (w *recordingWriter) dropped() []ID { +func (w *dropperWriter) drops() []ID { w.mu.Lock() defer w.mu.Unlock() - return append([]ID(nil), w.drop...) + return append([]ID(nil), w.dropped...) } diff --git a/mcp/streamable.go b/mcp/streamable.go index 001d929f..850667b7 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -1088,6 +1088,11 @@ type stream struct { // the duration of the subscription, and act as the target for // out-of-band notifications routed through this connection. isListen bool + + // peerTerminated records that the peer cancelled a call on this stream (see + // terminate). From then on deliverLocked discards every message, and only + // keeps the accounting of requests. + peerTerminated bool } // close sends a 'close' event to the client (if protocolVersion >= 2025-11-25 @@ -1114,6 +1119,24 @@ func (s *stream) close(reconnectAfter time.Duration) { s.done = nil } +// terminate ends the stream because the peer cancelled a call on it. The +// hanging request, if any, is released at once, with 204 No Content if +// nothing was written to it. Anything delivered to the stream afterwards, +// including the cancelled call's response, is discarded (see deliverLocked). +func (s *stream) terminate() { + s.mu.Lock() + defer s.mu.Unlock() + s.peerTerminated = true + if s.done == nil { + return // stream not connected or already closed + } + if s.pendingJSONMessages != nil || s.lastWrite.IsZero() { + s.writeNoContentLocked() + } + close(s.done) + s.done = nil +} + // release releases the stream from its HTTP request, allowing it to be // claimed by another request (e.g., for resumption). func (s *stream) release() { @@ -1237,12 +1260,12 @@ func extractErrorStatus(ctx context.Context, msg jsonrpc.Message) int { // pendingJSONMessages (for JSON mode). The eventID is used for SSE event ID; // pass "" to omit. // -// If data is nil there is nothing to write: the call is only accounting for a -// request that will never be answered (see [streamableServerConn.DropResponse]). -// // If responseTo is valid, it is removed from the requests map. When all // requests have been responded to, the done channel is closed and set to nil. // +// If the stream is peerTerminated, data is discarded without error: only the +// accounting above is done. +// // If overrideStatus is non-zero, data is treated as a SEP-2575 protocol-level // error response (>= 2026-07-28): it is written as a single raw JSON-RPC // response body with Content-Type: application/json and HTTP status @@ -1261,6 +1284,9 @@ func (s *stream) deliverLocked(data []byte, eventID string, responseTo jsonrpc.I } // Now, try to deliver the message to the client. done = len(s.requests) == 0 && s.id != "" + if s.peerTerminated { + return done, nil + } if s.done == nil { return done, fmt.Errorf("stream not connected or already closed") } @@ -1286,12 +1312,8 @@ func (s *stream) deliverLocked(data []byte, eventID string, responseTo jsonrpc.I // there's a brief race between request cancellation and releasing the // stream. if s.pendingJSONMessages != nil { - if data != nil { - s.pendingJSONMessages = append(s.pendingJSONMessages, data) - } - if done && len(s.pendingJSONMessages) == 0 { - s.writeNoContentLocked() - } else if done { + s.pendingJSONMessages = append(s.pendingJSONMessages, data) + if done { // Flush all pending messages as JSON response. var toWrite []byte if len(s.pendingJSONMessages) == 1 && !s.isBatch { @@ -1306,22 +1328,19 @@ func (s *stream) deliverLocked(data []byte, eventID string, responseTo jsonrpc.I return done, err } } - } else if data != nil { + } else { // SSE mode: write event to response writer. s.lastIdx++ if _, err := writeEvent(s.w, Event{Name: "message", Data: data, ID: eventID}); err != nil { return done, err } s.markWrittenLocked() - } else if done && s.lastWrite.IsZero() { - // Nothing was ever written, so the header is still ours to set. - s.writeNoContentLocked() } return done, nil } -// writeNoContentLocked ends a stream that carried nothing: every request the -// POST brought was cancelled before it was answered, so there is no body and +// writeNoContentLocked ends a stream that carried nothing because the peer +// cancelled a call on it before anything was written, so there is no body and // the status says so instead of an empty 200 under a Content-Type. func (s *stream) writeNoContentLocked() { s.w.Header().Del("Content-Type") @@ -2120,7 +2139,9 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e delivered := false var errs []error protocolVersion := protocolVersionFromContext(ctx) - if c.eventStore != nil && protocolVersion < protocolVersion20260728 { + // A terminated stream discards messages, so they must not be replayed on + // resumption either. + if c.eventStore != nil && protocolVersion < protocolVersion20260728 && !s.peerTerminated { if err := c.eventStore.Append(ctx, c.sessionID, s.id, data); err != nil { errs = append(errs, err) } else { @@ -2159,33 +2180,20 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e return nil } -// DropResponse implements [jsonrpc2.ResponseDropper]: the cancelled call gets -// no response, but a POST's stream hangs until every call it carried has been -// answered, so the call is retired from the stream all the same. +// DropResponse implements [jsonrpc2.ResponseDropper]: the stream that carries +// the cancelled call is terminated now, rather than when the handler returns. +// The response still goes through Write, which retires the call as usual but +// sends nothing, so the call keeps its ID until then. func (c *streamableServerConn) DropResponse(id jsonrpc.ID) { c.mu.Lock() var s *stream if streamID, ok := c.requestStreams[id]; ok { s = c.streams[streamID] } - delete(c.requestStreams, id) c.mu.Unlock() - if s == nil { - return - } - - s.mu.Lock() - // A nil payload delivers nothing; it only retires the request. An error - // here means the stream is already disconnected, which is not a problem - // when there is nothing to send. - done, _ := s.deliverLocked(nil, "", id, 0) - s.mu.Unlock() - - if done { - c.mu.Lock() - delete(c.streams, s.id) - c.mu.Unlock() + if s != nil { + s.terminate() } } diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 586f8e04..789e6bdd 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -3862,6 +3862,331 @@ func TestStreamableCancelledCallGetsNoResponse(t *testing.T) { } } +// TestStreamableCancelledCallHandlerIgnoresContext checks that the POST of a +// cancelled call is released as soon as the cancellation arrives, even if the +// handler keeps running, and that the call's ID stays reserved until the +// handler returns. +func TestStreamableCancelledCallHandlerIgnoresContext(t *testing.T) { + forEachResponseMode(t, func(t *testing.T, jsonResponse bool) { + started := make(chan struct{}) + release := make(chan struct{}) + server := NewServer(testImpl, nil) + server.AddTool(&Tool{Name: "stubborn", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(context.Context, *CallToolRequest) (*CallToolResult, error) { + close(started) + <-release + return &CallToolResult{}, nil + }) + server.AddTool(&Tool{Name: "noop", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(context.Context, *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{}, nil + }) + s := newCancelTestSession(t, server, &StreamableHTTPOptions{JSONResponse: jsonResponse}, protocolVersion20250618) + // Closing the session waits for the handler, so release it on failure too. + releaseHandler := sync.OnceFunc(func() { close(release) }) + t.Cleanup(releaseHandler) + + call := s.postAsync(req(2, "tools/call", &CallToolParams{Name: "stubborn"})) + <-started + s.cancel(2) + + got := <-call + if got.err != nil { + t.Fatalf("the cancelled call's POST was not released while its handler ran: %v", got.err) + } + if got.status != http.StatusNoContent || len(got.body) > 0 { + t.Errorf("the cancelled call's POST: status %d, body %q; want %d and no body", got.status, got.body, http.StatusNoContent) + } + + reuse := s.mustPost(req(2, "tools/call", &CallToolParams{Name: "noop"})) + if reuse.StatusCode != http.StatusBadRequest { + t.Errorf("reusing the ID while its handler runs: status %d, want %d", reuse.StatusCode, http.StatusBadRequest) + } + + releaseHandler() + s.waitFor("the cancelled call to be forgotten", func(c *streamableServerConn) bool { + return len(c.requestStreams) == 0 && len(c.streams) == 1 // only the standalone stream + }) + if again := s.mustPost(req(2, "tools/call", &CallToolParams{Name: "noop"})); again.StatusCode != http.StatusOK { + t.Errorf("reusing the ID after its handler returned: status %d, want %d", again.StatusCode, http.StatusOK) + } + }) +} + +// TestStreamableCancelledCallAfterEvents checks that cancelling a call whose +// SSE stream already carried a message ends the stream normally, with that +// message and no response. With an event store, resuming the stream must not +// replay the response either. +func TestStreamableCancelledCallAfterEvents(t *testing.T) { + for _, withStore := range []bool{false, true} { + name := "no event store" + if withStore { + name = "event store" + } + t.Run(name, func(t *testing.T) { + started := make(chan struct{}) + server := NewServer(testImpl, nil) + server.AddTool(&Tool{Name: "slow", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { + if err := req.Session.NotifyProgress(ctx, &ProgressNotificationParams{ProgressToken: "tok", Progress: 1}); err != nil { + return nil, err + } + close(started) + <-ctx.Done() + return nil, ctx.Err() + }) + opts := &StreamableHTTPOptions{} + if withStore { + opts.EventStore = NewMemoryEventStore(nil) + } + s := newCancelTestSession(t, server, opts, protocolVersion20250618) + + call := s.postAsync(req(2, "tools/call", &CallToolParams{Name: "slow"})) + <-started + s.cancel(2) + + got := <-call + if got.err != nil { + t.Fatalf("the cancelled call's POST: %v", got.err) + } + if got.status != http.StatusOK { + t.Errorf("the cancelled call's POST: status %d, want %d", got.status, http.StatusOK) + } + if !strings.Contains(got.body, notificationProgress) || strings.Contains(got.body, `"id":2`) { + t.Errorf("the cancelled call's stream carried:\n%s\nwant the progress notification and no response", got.body) + } + s.waitFor("the cancelled call to be forgotten", func(c *streamableServerConn) bool { + return len(c.requestStreams) == 0 && len(c.streams) == 1 + }) + + if !withStore { + return + } + var lastEventID string + for line := range strings.Lines(got.body) { + if id, ok := strings.CutPrefix(strings.TrimSpace(line), "id: "); ok { + lastEventID = id + } + } + if lastEventID == "" { + t.Fatalf("the cancelled call's stream carried no event ID:\n%s", got.body) + } + if replayed := s.resume(lastEventID); strings.Contains(replayed, `"id":2`) { + t.Errorf("resuming the stream replayed the cancelled call's response:\n%s", replayed) + } + }) + } +} + +// TestStreamableCancelledCallInBatch checks that cancelling one call of a +// 2025-03-26 batch ends the batch's POST at once, even with another call still +// running, and that nothing is left behind once every handler has returned. +// Responses to the other calls that were not yet written are dropped. +func TestStreamableCancelledCallInBatch(t *testing.T) { + for _, answeredFirst := range []bool{true, false} { + name := "cancelled first" + if answeredFirst { + name = "answered first" + } + t.Run(name, func(t *testing.T) { + forEachResponseMode(t, func(t *testing.T, jsonResponse bool) { + started := make(chan struct{}) + release := make(chan struct{}) + server := NewServer(testImpl, nil) + server.AddTool(&Tool{Name: "slow", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(ctx context.Context, _ *CallToolRequest) (*CallToolResult, error) { + close(started) + <-ctx.Done() + return nil, ctx.Err() + }) + server.AddTool(&Tool{Name: "gated", InputSchema: &jsonschema.Schema{Type: "object"}}, + func(context.Context, *CallToolRequest) (*CallToolResult, error) { + <-release + return &CallToolResult{Content: []Content{&TextContent{Text: "gated result"}}}, nil + }) + s := newCancelTestSession(t, server, &StreamableHTTPOptions{JSONResponse: jsonResponse}, protocolVersion20250326) + releaseHandler := sync.OnceFunc(func() { close(release) }) + t.Cleanup(releaseHandler) + + call := s.postAsync( + req(2, "tools/call", &CallToolParams{Name: "slow"}), + req(3, "tools/call", &CallToolParams{Name: "gated"})) + <-started + if answeredFirst { + releaseHandler() + s.waitFor("call 3 to be answered", func(c *streamableServerConn) bool { + _, ok := c.requestStreams[jsonrpc2.Int64ID(3)] + return !ok + }) + } + s.cancel(2) + + got := <-call + if got.err != nil { + t.Fatalf("the batch's POST: %v", got.err) + } + if strings.Contains(got.body, `"id":2`) { + t.Errorf("the batch's POST carried a response to the cancelled call:\n%s", got.body) + } + releaseHandler() + s.waitFor("the batch to be forgotten", func(c *streamableServerConn) bool { + return len(c.requestStreams) == 0 && len(c.streams) == 1 + }) + }) + }) + } +} + +func forEachResponseMode(t *testing.T, f func(t *testing.T, jsonResponse bool)) { + t.Run("sse", func(t *testing.T) { f(t, false) }) + t.Run("json", func(t *testing.T) { f(t, true) }) +} + +// cancelTestSession is a raw HTTP client for one session of a streamable +// server, for tests that cancel calls in flight: an SDK client abandons the +// POST of a call it cancels, so it cannot observe how the server ends it. +type cancelTestSession struct { + t *testing.T + ctx context.Context + handler *StreamableHTTPHandler + url string + version string + sessionID string +} + +func newCancelTestSession(t *testing.T, server *Server, opts *StreamableHTTPOptions, version string) *cancelTestSession { + t.Helper() + handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, opts) + t.Cleanup(handler.closeAll) + httpServer := httptest.NewServer(mustNotPanic(t, handler)) + t.Cleanup(httpServer.Close) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + t.Cleanup(cancel) + + s := &cancelTestSession{t: t, ctx: ctx, handler: handler, url: httpServer.URL, version: version} + s.sessionID = s.mustPost(req(1, methodInitialize, &InitializeParams{ProtocolVersion: version})).Header.Get(sessionIDHeader) + if s.sessionID == "" { + t.Fatal("initialize response carried no session ID") + } + s.mustPost(req(0, notificationInitialized, &InitializedParams{})) + return s +} + +// post sends msgs in one POST, as a batch if there is more than one. +func (s *cancelTestSession) post(msgs ...jsonrpc.Message) (*http.Response, error) { + var raw []json.RawMessage + for _, msg := range msgs { + data, err := jsonrpc2.EncodeMessage(msg) + if err != nil { + return nil, err + } + raw = append(raw, data) + } + body := []byte(raw[0]) + if len(raw) > 1 { + body = mustMarshal(raw) + } + httpReq, err := http.NewRequestWithContext(s.ctx, http.MethodPost, s.url, bytes.NewReader(body)) + if err != nil { + return nil, err + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "application/json, text/event-stream") + if s.sessionID != "" { + httpReq.Header.Set(sessionIDHeader, s.sessionID) + httpReq.Header.Set(protocolVersionHeader, s.version) + } + return http.DefaultClient.Do(httpReq) +} + +// mustPost is like post, but discards the response body. +func (s *cancelTestSession) mustPost(msgs ...jsonrpc.Message) *http.Response { + s.t.Helper() + resp, err := s.post(msgs...) + if err != nil { + s.t.Fatal(err) + } + resp.Body.Close() + return resp +} + +type postResult struct { + status int + body string + err error +} + +// postAsync is like post, but returns at once and reports the whole response +// once the server ends it. +func (s *cancelTestSession) postAsync(msgs ...jsonrpc.Message) <-chan postResult { + ch := make(chan postResult, 1) + go func() { + resp, err := s.post(msgs...) + if err != nil { + ch <- postResult{err: err} + return + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + ch <- postResult{status: resp.StatusCode, body: string(body), err: err} + }() + return ch +} + +// resume resumes a stream with a GET carrying lastEventID, and returns the +// events replayed. +func (s *cancelTestSession) resume(lastEventID string) string { + s.t.Helper() + httpReq, err := http.NewRequestWithContext(s.ctx, http.MethodGet, s.url, nil) + if err != nil { + s.t.Fatal(err) + } + httpReq.Header.Set("Accept", "text/event-stream") + httpReq.Header.Set(sessionIDHeader, s.sessionID) + httpReq.Header.Set(protocolVersionHeader, s.version) + httpReq.Header.Set(lastEventIDHeader, lastEventID) + resp, err := http.DefaultClient.Do(httpReq) + if err != nil { + s.t.Fatal(err) + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + s.t.Fatal(err) + } + if resp.StatusCode != http.StatusOK { + s.t.Fatalf("resuming the stream: status %d, body %q", resp.StatusCode, body) + } + return string(body) +} + +func (s *cancelTestSession) cancel(id int64) { + s.t.Helper() + s.mustPost(req(0, notificationCancelled, &CancelledParams{RequestID: id, Reason: "test cancellation"})) +} + +// waitFor waits until cond holds for the session's connection, with its lock +// held. +func (s *cancelTestSession) waitFor(what string, cond func(*streamableServerConn) bool) { + s.t.Helper() + s.handler.mu.Lock() + c := s.handler.sessions[s.sessionID].transport.connection + s.handler.mu.Unlock() + for { + c.mu.Lock() + ok := cond(c) + c.mu.Unlock() + if ok { + return + } + select { + case <-s.ctx.Done(): + s.t.Fatalf("timed out waiting for %s", what) + case <-time.After(5 * time.Millisecond): + } + } +} + // TestStreamableStateless_AcceptsNewProtocol is the positive control: // confirms that a stateless server still accepts new-protocol requests // (the rejection in TestStreamableStateful_RejectsNewProtocol must not diff --git a/mcp/transport.go b/mcp/transport.go index c3ed0894..05bb0ab6 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -13,6 +13,7 @@ import ( "log/slog" "net" "os" + "slices" "sync" "time" @@ -250,8 +251,9 @@ type cancellationPropagator interface { } // A canceller is a jsonrpc2.Preempter that cancels in-flight requests on MCP -// cancelled notifications. The cancelled request is answered with no response -// at all, as the spec requires. +// cancelled notifications. Connections that implement +// [jsonrpc2.ResponseDropper] send no response for the cancelled request, as +// the spec requires. type canceller struct { conn *jsonrpc2.Connection logger *slog.Logger @@ -424,8 +426,8 @@ func (s *loggingConn) Write(ctx context.Context, msg jsonrpc.Message) error { } // DropResponse implements [jsonrpc2.ResponseDropper] by forwarding to the -// delegate: the connection asks its writer for the interface, and a wrapper -// without it would silently keep the delegate waiting for the response. +// delegate: the connection asks its writer for the interface, so a wrapper +// without it would keep the delegate from learning that a call was cancelled. func (s *loggingConn) DropResponse(id jsonrpc.ID) { if dropper, ok := s.delegate.(jsonrpc2.ResponseDropper); ok { dropper.DropResponse(id) @@ -492,8 +494,12 @@ type ioConn struct { // batches correlate incoming requests to the batch in which they arrived. // Since writes may be concurrent to reads, we need to guard this with a mutex. + // + // dropped holds the incoming calls the peer cancelled, whose responses + // Write discards (see [ioConn.DropResponse]). It is guarded by batchMu too. batchMu sync.Mutex batches map[jsonrpc2.ID]*msgBatch // lazily allocated + dropped map[jsonrpc2.ID]bool // lazily allocated closeOnce sync.Once closed chan struct{} @@ -636,30 +642,50 @@ func (t *ioConn) addBatch(batch *msgBatch) error { return nil } +// DropResponse implements [jsonrpc2.ResponseDropper]: the response to the +// cancelled call is discarded by Write, which still accounts for it in the +// call's batch, if any, so that the other responses in the batch are sent. +func (t *ioConn) DropResponse(id jsonrpc.ID) { + t.batchMu.Lock() + defer t.batchMu.Unlock() + if t.dropped == nil { + t.dropped = make(map[jsonrpc2.ID]bool) + } + t.dropped[id] = true +} + // updateBatch records a response in the message batch tracking the -// corresponding incoming call, if any. +// corresponding incoming call, if any. A response to a call the peer +// cancelled (see [ioConn.DropResponse]) is left out of its batch. // -// The second result reports whether resp was part of a batch. If this is true, -// the first result is nil if the batch is still incomplete, or the full set of -// batch responses if resp completed the batch. +// The second result reports whether resp was part of a batch or dropped. If +// this is true, the first result holds the responses to send: none if resp's +// batch is still incomplete or resp was dropped outside a batch, or the +// batch's responses if resp completed the batch. func (t *ioConn) updateBatch(resp *jsonrpc.Response) ([]*jsonrpc.Response, bool) { t.batchMu.Lock() defer t.batchMu.Unlock() + dropped := t.dropped[resp.ID] + delete(t.dropped, resp.ID) if batch, ok := t.batches[resp.ID]; ok { idx, ok := batch.unresolved[resp.ID] if !ok { panic("internal error: inconsistent batches") } - batch.responses[idx] = resp + if !dropped { + batch.responses[idx] = resp + } delete(batch.unresolved, resp.ID) delete(t.batches, resp.ID) if len(batch.unresolved) == 0 { - return batch.responses, true + // nil out dropped responses before returning the batch + // a nil response might be present in case of dropped call by notification/cancel + return slices.DeleteFunc(batch.responses, func(r *jsonrpc.Response) bool { return r == nil }), true } return nil, true } - return nil, false + return nil, dropped } // A msgBatch records information about an incoming batch of jsonrpc.2 calls. @@ -783,10 +809,10 @@ func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error { t.writeMu.Lock() defer t.writeMu.Unlock() - // Batching support: if msg is a Response, it may have completed a batch, so - // check that first. Otherwise, it is a request or notification, and we may - // want to collect it into a batch before sending, if we're configured to use - // outgoing batches. + // Batching support: if msg is a Response, it may have completed a batch, or + // be dropped, so check that first. Otherwise, it is a request or + // notification, and we may want to collect it into a batch before sending, + // if we're configured to use outgoing batches. if resp, ok := msg.(*jsonrpc.Response); ok { if batch, ok := t.updateBatch(resp); ok { if len(batch) > 0 { diff --git a/mcp/transport_test.go b/mcp/transport_test.go index 00c33230..3a5c32e6 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -5,6 +5,7 @@ package mcp import ( + "bufio" "bytes" "context" "errors" @@ -365,6 +366,159 @@ func TestIOConnFrameCap(t *testing.T) { } } +// TestIOConnDropResponse checks that ioConn discards the response to a call +// the peer cancelled, and leaves it out of the batch the call arrived in while +// still sending the batch's other responses. +func TestIOConnDropResponse(t *testing.T) { + call := func(id int) string { return fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":"m"}`, id) } + tests := []struct { + name string + in string // the incoming payload + drop []int64 // the calls the peer cancelled + respond []int64 // the responses written, in order + want [][]any // the IDs of the responses on the wire, by line + }{ + {"single call, dropped", call(1), []int64{1}, []int64{1}, nil}, + {"single call, not dropped", call(1), nil, []int64{1}, [][]any{{int64(1)}}}, + {"batch, one dropped", "[" + call(1) + "," + call(2) + "]", []int64{1}, []int64{1, 2}, [][]any{{int64(2)}}}, + {"batch, dropped last", "[" + call(1) + "," + call(2) + "]", []int64{2}, []int64{1, 2}, [][]any{{int64(1)}}}, + {"batch, all dropped", "[" + call(1) + "," + call(2) + "]", []int64{1, 2}, []int64{2, 1}, nil}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + ctx := context.Background() + var out bytes.Buffer + tr := newIOConn(rwc{ + rc: io.NopCloser(strings.NewReader(test.in + "\n")), + wc: nopCloserWriter{&out}, + }) + t.Cleanup(func() { tr.Close() }) + + for range strings.Count(test.in, `"method"`) { + if _, err := tr.Read(ctx); err != nil { + t.Fatal(err) + } + } + for _, id := range test.drop { + tr.DropResponse(jsonrpc2.Int64ID(id)) + } + for _, id := range test.respond { + if err := tr.Write(ctx, &jsonrpc.Response{ID: jsonrpc2.Int64ID(id), Result: []byte("{}")}); err != nil { + t.Fatalf("Write(response %d): %v", id, err) + } + } + + var got [][]any + for line := range strings.Lines(out.String()) { + msgs, _, err := readBatch([]byte(line)) + if err != nil { + t.Fatalf("decoding %q: %v", line, err) + } + var ids []any + for _, m := range msgs { + ids = append(ids, m.(*jsonrpc.Response).ID.Raw()) + } + got = append(got, ids) + } + if fmt.Sprint(got) != fmt.Sprint(test.want) { + t.Errorf("responses on the wire: %v, want %v", got, test.want) + } + if len(tr.dropped) != 0 { + t.Errorf("after the responses were written, %d dropped IDs are still recorded", len(tr.dropped)) + } + }) + } +} + +// TestIOConnCancelledCallGetsNoResponse checks end to end, over an in-memory +// transport wrapped for logging, that a call the client cancels with +// notifications/cancelled gets no response. +func TestIOConnCancelledCallGetsNoResponse(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + started := make(chan struct{}) + server := NewServer(testImpl, nil) + AddTool(server, &Tool{Name: "slow"}, func(ctx context.Context, _ *CallToolRequest, _ struct{}) (*CallToolResult, any, error) { + close(started) + <-ctx.Done() + return nil, nil, ctx.Err() + }) + serverT, clientT := NewInMemoryTransports() + var log safeBuffer + ss, err := server.Connect(ctx, &LoggingTransport{Transport: serverT, Writer: &log}, nil) + if err != nil { + t.Fatal(err) + } + defer ss.Close() + + peer := clientT.rwc + received := make(chan jsonrpc.Message, 10) + go func() { + defer close(received) + scanner := bufio.NewScanner(peer) + for scanner.Scan() { + msg, err := jsonrpc2.DecodeMessage(scanner.Bytes()) + if err != nil { + t.Errorf("decoding %q: %v", scanner.Bytes(), err) + return + } + received <- msg + } + }() + send := func(msg jsonrpc.Message) { + t.Helper() + data, err := jsonrpc2.EncodeMessage(msg) + if err != nil { + t.Fatal(err) + } + if _, err := peer.Write(append(data, '\n')); err != nil { + t.Fatal(err) + } + } + // awaitResponse returns once the response to id arrives, and fails if a + // response to the cancelled call arrives first. + awaitResponse := func(id int64) { + t.Helper() + for { + select { + case msg := <-received: + resp, ok := msg.(*jsonrpc.Response) + if !ok { + continue + } + switch resp.ID.Raw() { + case int64(2): + t.Fatalf("the cancelled call got a response: %+v", resp) + case id: + return + } + case <-ctx.Done(): + t.Fatalf("timed out waiting for the response to %d", id) + } + } + } + + send(req(1, methodInitialize, &InitializeParams{ProtocolVersion: protocolVersion20250618})) + awaitResponse(1) + send(req(0, notificationInitialized, &InitializedParams{})) + send(req(2, "tools/call", &CallToolParams{Name: "slow"})) + <-started + send(req(0, notificationCancelled, &CancelledParams{RequestID: int64(2)})) + + // The logging wrapper logs the response once ioConn has dropped it; after + // that, a response sent on the wire would precede the ping's. + for !strings.Contains(string(log.Bytes()), `write: {"jsonrpc":"2.0","id":2,`) { + select { + case <-ctx.Done(): + t.Fatalf("timed out waiting for the cancelled call's handler to return; log:\n%s", log.Bytes()) + case <-time.After(5 * time.Millisecond): + } + } + send(req(3, "ping", nil)) + awaitResponse(3) +} + // droppingConn is a Connection that records the ids it was told will receive // no response. It implements [jsonrpc2.ResponseDropper]; plainConn, below, // deliberately does not. @@ -391,22 +545,21 @@ func (plainConn) Close() error { return nil } // does not lose the "this call gets no response" signal: loggingConn is the // one Connection wrapper here, and it used to leave the POST stream open. func TestLoggingConnDropResponse(t *testing.T) { - t.Run("forwards to a delegate that holds per-call state", func(t *testing.T) { + t.Run("forwards to a delegate that drops responses", func(t *testing.T) { delegate := &droppingConn{} - conn := &loggingConn{delegate: delegate, w: io.Discard} + var conn Connection = &loggingConn{delegate: delegate, w: io.Discard} - dropper, ok := any(conn).(jsonrpc2.ResponseDropper) + dropper, ok := conn.(jsonrpc2.ResponseDropper) if !ok { - t.Fatal("loggingConn does not implement jsonrpc2.ResponseDropper, so a cancelled call's stream is never released") + t.Fatal("loggingConn does not implement jsonrpc2.ResponseDropper, so a cancelled call's response is never dropped") } dropper.DropResponse(jsonrpc.ID{}) - if len(delegate.dropped) != 1 { t.Errorf("the delegate was told about %d dropped responses, want 1", len(delegate.dropped)) } }) - t.Run("does nothing for a delegate that holds none", func(t *testing.T) { + t.Run("does nothing for a delegate that does not", func(t *testing.T) { conn := &loggingConn{delegate: plainConn{}, w: io.Discard} conn.DropResponse(jsonrpc.ID{}) }) From 746a459564df749d535ee11bd8603a8c31ebea7f Mon Sep 17 00:00:00 2001 From: guglielmoc Date: Mon, 12 Oct 2026 13:08:06 +0000 Subject: [PATCH 3/6] refactor: remove manual cancellation tests for streamable transport to simplify test suite --- docs/protocol.md | 11 +- internal/docs/protocol.src.md | 11 +- mcp/mcp_test.go | 4 - mcp/streamable.go | 58 +---- mcp/streamable_test.go | 428 ---------------------------------- mcp/transport.go | 80 ++++--- mcp/transport_test.go | 63 +++-- 7 files changed, 88 insertions(+), 567 deletions(-) diff --git a/docs/protocol.md b/docs/protocol.md index f092c600..43c1b65e 100644 --- a/docs/protocol.md +++ b/docs/protocol.md @@ -712,11 +712,12 @@ When an RPC exits due to a cancellation error, there's a guarantee that the cancellation notification has been sent, but there's no guarantee that the server has observed it (see [concurrency](#concurrency)). -A receiver of a cancellation notification sends no response for the cancelled -request, as the spec requires. On the streamable HTTP transport the stream that -carried the request is still completed as soon as the notification arrives, -without waiting for the handler to return, so the POST it arrived on ends -rather than hanging, with 204 No Content when nothing else was written on it. +On a session using protocol version 2026-07-28 or later over a stdio-style +transport (`StdioTransport`, `CommandTransport`, `IOTransport` or +`InMemoryTransport`), a server that receives a cancellation notification sends +no response for the cancelled request, as the spec requires. Over the +streamable HTTP transport, a 2026-07-28 request is cancelled by closing its +HTTP request instead (see `StreamableHTTPOptions.PropagateRequestCancellation`). ```go func Example_cancellation() { diff --git a/internal/docs/protocol.src.md b/internal/docs/protocol.src.md index 4a339def..f0380609 100644 --- a/internal/docs/protocol.src.md +++ b/internal/docs/protocol.src.md @@ -607,11 +607,12 @@ When an RPC exits due to a cancellation error, there's a guarantee that the cancellation notification has been sent, but there's no guarantee that the server has observed it (see [concurrency](#concurrency)). -A receiver of a cancellation notification sends no response for the cancelled -request, as the spec requires. On the streamable HTTP transport the stream that -carried the request is still completed as soon as the notification arrives, -without waiting for the handler to return, so the POST it arrived on ends -rather than hanging, with 204 No Content when nothing else was written on it. +On a session using protocol version 2026-07-28 or later over a stdio-style +transport (`StdioTransport`, `CommandTransport`, `IOTransport` or +`InMemoryTransport`), a server that receives a cancellation notification sends +no response for the cancelled request, as the spec requires. Over the +streamable HTTP transport, a 2026-07-28 request is cancelled by closing its +HTTP request instead (see `StreamableHTTPOptions.PropagateRequestCancellation`). %include ../../mcp/mcp_example_test.go cancellation - diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 4077a017..c53e6d57 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -787,10 +787,6 @@ func TestCancellationReason(t *testing.T) { if err := cs.conn.Notify(ctx, notificationCancelled, &CancelledParams{RequestID: call.ID().Raw(), Reason: "user asked"}); err != nil { t.Fatal(err) } - // The receiver sends no response to a cancelled call, so the SDK's own - // cancel path (cancelCall) retires the outgoing call before notifying; - // this test notifies by hand and retires by hand, or synctest deadlocks. - cs.conn.Retire(call, context.Canceled) cause := <-cancelled if cause == nil { diff --git a/mcp/streamable.go b/mcp/streamable.go index 6cdb7f94..813a3824 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -1088,11 +1088,6 @@ type stream struct { // the duration of the subscription, and act as the target for // out-of-band notifications routed through this connection. isListen bool - - // peerTerminated records that the peer cancelled a call on this stream (see - // terminate). From then on deliverLocked discards every message, and only - // keeps the accounting of requests. - peerTerminated bool } // close sends a 'close' event to the client (if protocolVersion >= 2025-11-25 @@ -1119,24 +1114,6 @@ func (s *stream) close(reconnectAfter time.Duration) { s.done = nil } -// terminate ends the stream because the peer cancelled a call on it. The -// hanging request, if any, is released at once, with 204 No Content if -// nothing was written to it. Anything delivered to the stream afterwards, -// including the cancelled call's response, is discarded (see deliverLocked). -func (s *stream) terminate() { - s.mu.Lock() - defer s.mu.Unlock() - s.peerTerminated = true - if s.done == nil { - return // stream not connected or already closed - } - if s.pendingJSONMessages != nil || s.lastWrite.IsZero() { - s.writeNoContentLocked() - } - close(s.done) - s.done = nil -} - // release releases the stream from its HTTP request, allowing it to be // claimed by another request (e.g., for resumption). func (s *stream) release() { @@ -1263,9 +1240,6 @@ func extractErrorStatus(ctx context.Context, msg jsonrpc.Message) int { // If responseTo is valid, it is removed from the requests map. When all // requests have been responded to, the done channel is closed and set to nil. // -// If the stream is peerTerminated, data is discarded without error: only the -// accounting above is done. -// // If overrideStatus is non-zero, data is treated as a SEP-2575 protocol-level // error response (>= 2026-07-28): it is written as a single raw JSON-RPC // response body with Content-Type: application/json and HTTP status @@ -1284,9 +1258,6 @@ func (s *stream) deliverLocked(data []byte, eventID string, responseTo jsonrpc.I } // Now, try to deliver the message to the client. done = len(s.requests) == 0 && s.id != "" - if s.peerTerminated { - return done, nil - } if s.done == nil { return done, fmt.Errorf("stream not connected or already closed") } @@ -1339,14 +1310,6 @@ func (s *stream) deliverLocked(data []byte, eventID string, responseTo jsonrpc.I return done, nil } -// writeNoContentLocked ends a stream that carried nothing because the peer -// cancelled a call on it before anything was written, so there is no body and -// the status says so instead of an empty 200 under a Content-Type. -func (s *stream) writeNoContentLocked() { - s.w.Header().Del("Content-Type") - s.w.WriteHeader(http.StatusNoContent) -} - // doneLocked reports whether the stream is logically complete. // // s.requests was populated when reading the POST body, requests are deleted as @@ -2139,9 +2102,7 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e delivered := false var errs []error protocolVersion := protocolVersionFromContext(ctx) - // A terminated stream discards messages, so they must not be replayed on - // resumption either. - if c.eventStore != nil && protocolVersion < protocolVersion20260728 && !s.peerTerminated { + if c.eventStore != nil && protocolVersion < protocolVersion20260728 { if err := c.eventStore.Append(ctx, c.sessionID, s.id, data); err != nil { errs = append(errs, err) } else { @@ -2180,23 +2141,6 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e return nil } -// DropResponse implements [jsonrpc2.ResponseDropper]: the stream that carries -// the cancelled call is terminated now, rather than when the handler returns. -// The response still goes through Write, which retires the call as usual but -// sends nothing, so the call keeps its ID until then. -func (c *streamableServerConn) DropResponse(id jsonrpc.ID) { - c.mu.Lock() - var s *stream - if streamID, ok := c.requestStreams[id]; ok { - s = c.streams[streamID] - } - c.mu.Unlock() - - if s != nil { - s.terminate() - } -} - // Close implements the [Connection] interface. func (c *streamableServerConn) Close() error { c.mu.Lock() diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 789e6bdd..a481e7b7 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -3759,434 +3759,6 @@ func TestCallCancellation_FastReturn(t *testing.T) { } } -// TestStreamableCancelledCallGetsNoResponse checks that a call the client -// cancelled with notifications/cancelled is answered with 204 and no body, in -// SSE and in JSON mode, through raw HTTP since an SDK client abandons the POST. -func TestStreamableCancelledCallGetsNoResponse(t *testing.T) { - for _, mode := range []struct { - name string - jsonResponse bool - }{ - {name: "sse", jsonResponse: false}, - {name: "json", jsonResponse: true}, - } { - t.Run(mode.name, func(t *testing.T) { - started := make(chan struct{}) - server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil) - server.AddTool( - &Tool{Name: "slow", InputSchema: &jsonschema.Schema{Type: "object"}}, - func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { - close(started) - <-ctx.Done() - return nil, ctx.Err() - }) - - handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, &StreamableHTTPOptions{JSONResponse: mode.jsonResponse}) - defer handler.closeAll() - httpServer := httptest.NewServer(mustNotPanic(t, handler)) - defer httpServer.Close() - - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - - post := func(sessionID string, msg jsonrpc.Message) (*http.Response, error) { - data, err := jsonrpc2.EncodeMessage(msg) - if err != nil { - return nil, err - } - httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, httpServer.URL, bytes.NewReader(data)) - if err != nil { - return nil, err - } - httpReq.Header.Set("Content-Type", "application/json") - httpReq.Header.Set("Accept", "application/json, text/event-stream") - if sessionID != "" { - httpReq.Header.Set(sessionIDHeader, sessionID) - } - return http.DefaultClient.Do(httpReq) - } - mustPost := func(sessionID string, msg jsonrpc.Message) *http.Response { - t.Helper() - resp, err := post(sessionID, msg) - if err != nil { - t.Fatal(err) - } - return resp - } - - initResp := mustPost("", req(1, methodInitialize, &InitializeParams{ProtocolVersion: protocolVersion20250618})) - sessionID := initResp.Header.Get(sessionIDHeader) - initResp.Body.Close() - if sessionID == "" { - t.Fatal("initialize response carried no session ID") - } - mustPost(sessionID, req(0, notificationInitialized, &InitializedParams{})).Body.Close() - - // The POST returns only when its stream ends, so call from a goroutine. - type postResult struct { - resp *http.Response - err error - } - call := make(chan postResult, 1) - go func() { - resp, err := post(sessionID, req(2, "tools/call", &CallToolParams{Name: "slow"})) - call <- postResult{resp, err} - }() - <-started - - mustPost(sessionID, req(0, notificationCancelled, &CancelledParams{ - RequestID: int64(2), - Reason: "test cancellation", - })).Body.Close() - - got := <-call - if got.err != nil { - t.Fatalf("the cancelled call's POST: %v", got.err) - } - defer got.resp.Body.Close() - - if got.resp.StatusCode != http.StatusNoContent { - t.Errorf("the cancelled call's POST status = %d, want %d", got.resp.StatusCode, http.StatusNoContent) - } - if ct := got.resp.Header.Get("Content-Type"); ct != "" { - t.Errorf("the cancelled call's POST carried Content-Type %q, want none", ct) - } - body, err := io.ReadAll(got.resp.Body) - if err != nil { - t.Fatalf("reading the cancelled call's stream: %v", err) - } - if len(body) > 0 { - t.Errorf("the cancelled call's stream carried:\n%s\nwant nothing", body) - } - }) - } -} - -// TestStreamableCancelledCallHandlerIgnoresContext checks that the POST of a -// cancelled call is released as soon as the cancellation arrives, even if the -// handler keeps running, and that the call's ID stays reserved until the -// handler returns. -func TestStreamableCancelledCallHandlerIgnoresContext(t *testing.T) { - forEachResponseMode(t, func(t *testing.T, jsonResponse bool) { - started := make(chan struct{}) - release := make(chan struct{}) - server := NewServer(testImpl, nil) - server.AddTool(&Tool{Name: "stubborn", InputSchema: &jsonschema.Schema{Type: "object"}}, - func(context.Context, *CallToolRequest) (*CallToolResult, error) { - close(started) - <-release - return &CallToolResult{}, nil - }) - server.AddTool(&Tool{Name: "noop", InputSchema: &jsonschema.Schema{Type: "object"}}, - func(context.Context, *CallToolRequest) (*CallToolResult, error) { - return &CallToolResult{}, nil - }) - s := newCancelTestSession(t, server, &StreamableHTTPOptions{JSONResponse: jsonResponse}, protocolVersion20250618) - // Closing the session waits for the handler, so release it on failure too. - releaseHandler := sync.OnceFunc(func() { close(release) }) - t.Cleanup(releaseHandler) - - call := s.postAsync(req(2, "tools/call", &CallToolParams{Name: "stubborn"})) - <-started - s.cancel(2) - - got := <-call - if got.err != nil { - t.Fatalf("the cancelled call's POST was not released while its handler ran: %v", got.err) - } - if got.status != http.StatusNoContent || len(got.body) > 0 { - t.Errorf("the cancelled call's POST: status %d, body %q; want %d and no body", got.status, got.body, http.StatusNoContent) - } - - reuse := s.mustPost(req(2, "tools/call", &CallToolParams{Name: "noop"})) - if reuse.StatusCode != http.StatusBadRequest { - t.Errorf("reusing the ID while its handler runs: status %d, want %d", reuse.StatusCode, http.StatusBadRequest) - } - - releaseHandler() - s.waitFor("the cancelled call to be forgotten", func(c *streamableServerConn) bool { - return len(c.requestStreams) == 0 && len(c.streams) == 1 // only the standalone stream - }) - if again := s.mustPost(req(2, "tools/call", &CallToolParams{Name: "noop"})); again.StatusCode != http.StatusOK { - t.Errorf("reusing the ID after its handler returned: status %d, want %d", again.StatusCode, http.StatusOK) - } - }) -} - -// TestStreamableCancelledCallAfterEvents checks that cancelling a call whose -// SSE stream already carried a message ends the stream normally, with that -// message and no response. With an event store, resuming the stream must not -// replay the response either. -func TestStreamableCancelledCallAfterEvents(t *testing.T) { - for _, withStore := range []bool{false, true} { - name := "no event store" - if withStore { - name = "event store" - } - t.Run(name, func(t *testing.T) { - started := make(chan struct{}) - server := NewServer(testImpl, nil) - server.AddTool(&Tool{Name: "slow", InputSchema: &jsonschema.Schema{Type: "object"}}, - func(ctx context.Context, req *CallToolRequest) (*CallToolResult, error) { - if err := req.Session.NotifyProgress(ctx, &ProgressNotificationParams{ProgressToken: "tok", Progress: 1}); err != nil { - return nil, err - } - close(started) - <-ctx.Done() - return nil, ctx.Err() - }) - opts := &StreamableHTTPOptions{} - if withStore { - opts.EventStore = NewMemoryEventStore(nil) - } - s := newCancelTestSession(t, server, opts, protocolVersion20250618) - - call := s.postAsync(req(2, "tools/call", &CallToolParams{Name: "slow"})) - <-started - s.cancel(2) - - got := <-call - if got.err != nil { - t.Fatalf("the cancelled call's POST: %v", got.err) - } - if got.status != http.StatusOK { - t.Errorf("the cancelled call's POST: status %d, want %d", got.status, http.StatusOK) - } - if !strings.Contains(got.body, notificationProgress) || strings.Contains(got.body, `"id":2`) { - t.Errorf("the cancelled call's stream carried:\n%s\nwant the progress notification and no response", got.body) - } - s.waitFor("the cancelled call to be forgotten", func(c *streamableServerConn) bool { - return len(c.requestStreams) == 0 && len(c.streams) == 1 - }) - - if !withStore { - return - } - var lastEventID string - for line := range strings.Lines(got.body) { - if id, ok := strings.CutPrefix(strings.TrimSpace(line), "id: "); ok { - lastEventID = id - } - } - if lastEventID == "" { - t.Fatalf("the cancelled call's stream carried no event ID:\n%s", got.body) - } - if replayed := s.resume(lastEventID); strings.Contains(replayed, `"id":2`) { - t.Errorf("resuming the stream replayed the cancelled call's response:\n%s", replayed) - } - }) - } -} - -// TestStreamableCancelledCallInBatch checks that cancelling one call of a -// 2025-03-26 batch ends the batch's POST at once, even with another call still -// running, and that nothing is left behind once every handler has returned. -// Responses to the other calls that were not yet written are dropped. -func TestStreamableCancelledCallInBatch(t *testing.T) { - for _, answeredFirst := range []bool{true, false} { - name := "cancelled first" - if answeredFirst { - name = "answered first" - } - t.Run(name, func(t *testing.T) { - forEachResponseMode(t, func(t *testing.T, jsonResponse bool) { - started := make(chan struct{}) - release := make(chan struct{}) - server := NewServer(testImpl, nil) - server.AddTool(&Tool{Name: "slow", InputSchema: &jsonschema.Schema{Type: "object"}}, - func(ctx context.Context, _ *CallToolRequest) (*CallToolResult, error) { - close(started) - <-ctx.Done() - return nil, ctx.Err() - }) - server.AddTool(&Tool{Name: "gated", InputSchema: &jsonschema.Schema{Type: "object"}}, - func(context.Context, *CallToolRequest) (*CallToolResult, error) { - <-release - return &CallToolResult{Content: []Content{&TextContent{Text: "gated result"}}}, nil - }) - s := newCancelTestSession(t, server, &StreamableHTTPOptions{JSONResponse: jsonResponse}, protocolVersion20250326) - releaseHandler := sync.OnceFunc(func() { close(release) }) - t.Cleanup(releaseHandler) - - call := s.postAsync( - req(2, "tools/call", &CallToolParams{Name: "slow"}), - req(3, "tools/call", &CallToolParams{Name: "gated"})) - <-started - if answeredFirst { - releaseHandler() - s.waitFor("call 3 to be answered", func(c *streamableServerConn) bool { - _, ok := c.requestStreams[jsonrpc2.Int64ID(3)] - return !ok - }) - } - s.cancel(2) - - got := <-call - if got.err != nil { - t.Fatalf("the batch's POST: %v", got.err) - } - if strings.Contains(got.body, `"id":2`) { - t.Errorf("the batch's POST carried a response to the cancelled call:\n%s", got.body) - } - releaseHandler() - s.waitFor("the batch to be forgotten", func(c *streamableServerConn) bool { - return len(c.requestStreams) == 0 && len(c.streams) == 1 - }) - }) - }) - } -} - -func forEachResponseMode(t *testing.T, f func(t *testing.T, jsonResponse bool)) { - t.Run("sse", func(t *testing.T) { f(t, false) }) - t.Run("json", func(t *testing.T) { f(t, true) }) -} - -// cancelTestSession is a raw HTTP client for one session of a streamable -// server, for tests that cancel calls in flight: an SDK client abandons the -// POST of a call it cancels, so it cannot observe how the server ends it. -type cancelTestSession struct { - t *testing.T - ctx context.Context - handler *StreamableHTTPHandler - url string - version string - sessionID string -} - -func newCancelTestSession(t *testing.T, server *Server, opts *StreamableHTTPOptions, version string) *cancelTestSession { - t.Helper() - handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, opts) - t.Cleanup(handler.closeAll) - httpServer := httptest.NewServer(mustNotPanic(t, handler)) - t.Cleanup(httpServer.Close) - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - t.Cleanup(cancel) - - s := &cancelTestSession{t: t, ctx: ctx, handler: handler, url: httpServer.URL, version: version} - s.sessionID = s.mustPost(req(1, methodInitialize, &InitializeParams{ProtocolVersion: version})).Header.Get(sessionIDHeader) - if s.sessionID == "" { - t.Fatal("initialize response carried no session ID") - } - s.mustPost(req(0, notificationInitialized, &InitializedParams{})) - return s -} - -// post sends msgs in one POST, as a batch if there is more than one. -func (s *cancelTestSession) post(msgs ...jsonrpc.Message) (*http.Response, error) { - var raw []json.RawMessage - for _, msg := range msgs { - data, err := jsonrpc2.EncodeMessage(msg) - if err != nil { - return nil, err - } - raw = append(raw, data) - } - body := []byte(raw[0]) - if len(raw) > 1 { - body = mustMarshal(raw) - } - httpReq, err := http.NewRequestWithContext(s.ctx, http.MethodPost, s.url, bytes.NewReader(body)) - if err != nil { - return nil, err - } - httpReq.Header.Set("Content-Type", "application/json") - httpReq.Header.Set("Accept", "application/json, text/event-stream") - if s.sessionID != "" { - httpReq.Header.Set(sessionIDHeader, s.sessionID) - httpReq.Header.Set(protocolVersionHeader, s.version) - } - return http.DefaultClient.Do(httpReq) -} - -// mustPost is like post, but discards the response body. -func (s *cancelTestSession) mustPost(msgs ...jsonrpc.Message) *http.Response { - s.t.Helper() - resp, err := s.post(msgs...) - if err != nil { - s.t.Fatal(err) - } - resp.Body.Close() - return resp -} - -type postResult struct { - status int - body string - err error -} - -// postAsync is like post, but returns at once and reports the whole response -// once the server ends it. -func (s *cancelTestSession) postAsync(msgs ...jsonrpc.Message) <-chan postResult { - ch := make(chan postResult, 1) - go func() { - resp, err := s.post(msgs...) - if err != nil { - ch <- postResult{err: err} - return - } - defer resp.Body.Close() - body, err := io.ReadAll(resp.Body) - ch <- postResult{status: resp.StatusCode, body: string(body), err: err} - }() - return ch -} - -// resume resumes a stream with a GET carrying lastEventID, and returns the -// events replayed. -func (s *cancelTestSession) resume(lastEventID string) string { - s.t.Helper() - httpReq, err := http.NewRequestWithContext(s.ctx, http.MethodGet, s.url, nil) - if err != nil { - s.t.Fatal(err) - } - httpReq.Header.Set("Accept", "text/event-stream") - httpReq.Header.Set(sessionIDHeader, s.sessionID) - httpReq.Header.Set(protocolVersionHeader, s.version) - httpReq.Header.Set(lastEventIDHeader, lastEventID) - resp, err := http.DefaultClient.Do(httpReq) - if err != nil { - s.t.Fatal(err) - } - defer resp.Body.Close() - body, err := io.ReadAll(resp.Body) - if err != nil { - s.t.Fatal(err) - } - if resp.StatusCode != http.StatusOK { - s.t.Fatalf("resuming the stream: status %d, body %q", resp.StatusCode, body) - } - return string(body) -} - -func (s *cancelTestSession) cancel(id int64) { - s.t.Helper() - s.mustPost(req(0, notificationCancelled, &CancelledParams{RequestID: id, Reason: "test cancellation"})) -} - -// waitFor waits until cond holds for the session's connection, with its lock -// held. -func (s *cancelTestSession) waitFor(what string, cond func(*streamableServerConn) bool) { - s.t.Helper() - s.handler.mu.Lock() - c := s.handler.sessions[s.sessionID].transport.connection - s.handler.mu.Unlock() - for { - c.mu.Lock() - ok := cond(c) - c.mu.Unlock() - if ok { - return - } - select { - case <-s.ctx.Done(): - s.t.Fatalf("timed out waiting for %s", what) - case <-time.After(5 * time.Millisecond): - } - } -} - // TestStreamableStateless_AcceptsNewProtocol is the positive control: // confirms that a stateless server still accepts new-protocol requests // (the rejection in TestStreamableStateful_RejectsNewProtocol must not diff --git a/mcp/transport.go b/mcp/transport.go index 05bb0ab6..cce29518 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -13,7 +13,6 @@ import ( "log/slog" "net" "os" - "slices" "sync" "time" @@ -251,9 +250,9 @@ type cancellationPropagator interface { } // A canceller is a jsonrpc2.Preempter that cancels in-flight requests on MCP -// cancelled notifications. Connections that implement -// [jsonrpc2.ResponseDropper] send no response for the cancelled request, as -// the spec requires. +// cancelled notifications. A connection that implements +// [jsonrpc2.ResponseDropper] can then drop the cancelled request's response, +// as the spec requires. type canceller struct { conn *jsonrpc2.Connection logger *slog.Logger @@ -434,6 +433,15 @@ func (s *loggingConn) DropResponse(id jsonrpc.ID) { } } +// sessionUpdated implements [serverConnection] by forwarding to the delegate, +// whose behavior may depend on the negotiated protocol version (see +// [ioConn.DropResponse]). +func (s *loggingConn) sessionUpdated(state ServerSessionState) { + if c, ok := s.delegate.(serverConnection); ok { + c.sessionUpdated(state) + } +} + func (s *loggingConn) Close() error { return s.delegate.Close() } @@ -475,8 +483,12 @@ type ioConn struct { // protocolVersion is the negotiated version of the protocol, // set during session initialization. // Since writes may be concurrent to reads, we need to guard this with a mutex. + // + // dropped holds the incoming calls the peer cancelled, whose responses + // Write discards (see [ioConn.DropResponse]). It is guarded by sessionMu too. sessionMu sync.Mutex protocolVersion string + dropped map[jsonrpc2.ID]bool // lazily allocated writeMu sync.Mutex // guards Write, which must be concurrency safe. rwc io.ReadWriteCloser // the underlying stream @@ -494,12 +506,8 @@ type ioConn struct { // batches correlate incoming requests to the batch in which they arrived. // Since writes may be concurrent to reads, we need to guard this with a mutex. - // - // dropped holds the incoming calls the peer cancelled, whose responses - // Write discards (see [ioConn.DropResponse]). It is guarded by batchMu too. batchMu sync.Mutex batches map[jsonrpc2.ID]*msgBatch // lazily allocated - dropped map[jsonrpc2.ID]bool // lazily allocated closeOnce sync.Once closed chan struct{} @@ -642,50 +650,55 @@ func (t *ioConn) addBatch(batch *msgBatch) error { return nil } -// DropResponse implements [jsonrpc2.ResponseDropper]: the response to the -// cancelled call is discarded by Write, which still accounts for it in the -// call's batch, if any, so that the other responses in the batch are sent. +// DropResponse implements [jsonrpc2.ResponseDropper]: on a session using +// protocol version 2026-07-28 or later, Write discards the response to the +// cancelled call. Sessions on earlier versions keep answering it. func (t *ioConn) DropResponse(id jsonrpc.ID) { - t.batchMu.Lock() - defer t.batchMu.Unlock() + t.sessionMu.Lock() + defer t.sessionMu.Unlock() + if t.protocolVersion < protocolVersion20260728 { + return + } if t.dropped == nil { t.dropped = make(map[jsonrpc2.ID]bool) } t.dropped[id] = true } +// takeDropped reports whether the response to the call with the given ID is +// to be discarded (see [ioConn.DropResponse]), and forgets the call. +func (t *ioConn) takeDropped(id jsonrpc.ID) bool { + t.sessionMu.Lock() + defer t.sessionMu.Unlock() + dropped := t.dropped[id] + delete(t.dropped, id) + return dropped +} + // updateBatch records a response in the message batch tracking the -// corresponding incoming call, if any. A response to a call the peer -// cancelled (see [ioConn.DropResponse]) is left out of its batch. +// corresponding incoming call, if any. // -// The second result reports whether resp was part of a batch or dropped. If -// this is true, the first result holds the responses to send: none if resp's -// batch is still incomplete or resp was dropped outside a batch, or the -// batch's responses if resp completed the batch. +// The second result reports whether resp was part of a batch. If this is true, +// the first result is nil if the batch is still incomplete, or the full set of +// batch responses if resp completed the batch. func (t *ioConn) updateBatch(resp *jsonrpc.Response) ([]*jsonrpc.Response, bool) { t.batchMu.Lock() defer t.batchMu.Unlock() - dropped := t.dropped[resp.ID] - delete(t.dropped, resp.ID) if batch, ok := t.batches[resp.ID]; ok { idx, ok := batch.unresolved[resp.ID] if !ok { panic("internal error: inconsistent batches") } - if !dropped { - batch.responses[idx] = resp - } + batch.responses[idx] = resp delete(batch.unresolved, resp.ID) delete(t.batches, resp.ID) if len(batch.unresolved) == 0 { - // nil out dropped responses before returning the batch - // a nil response might be present in case of dropped call by notification/cancel - return slices.DeleteFunc(batch.responses, func(r *jsonrpc.Response) bool { return r == nil }), true + return batch.responses, true } return nil, true } - return nil, dropped + return nil, false } // A msgBatch records information about an incoming batch of jsonrpc.2 calls. @@ -809,11 +822,14 @@ func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error { t.writeMu.Lock() defer t.writeMu.Unlock() - // Batching support: if msg is a Response, it may have completed a batch, or - // be dropped, so check that first. Otherwise, it is a request or - // notification, and we may want to collect it into a batch before sending, - // if we're configured to use outgoing batches. + // Batching support: if msg is a Response, it may have completed a batch, so + // check that first. Otherwise, it is a request or notification, and we may + // want to collect it into a batch before sending, if we're configured to use + // outgoing batches. if resp, ok := msg.(*jsonrpc.Response); ok { + if t.takeDropped(resp.ID) { + return nil + } if batch, ok := t.updateBatch(resp); ok { if len(batch) > 0 { data, err := marshalMessages(batch) diff --git a/mcp/transport_test.go b/mcp/transport_test.go index 3a5c32e6..90f7cb62 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -367,58 +367,46 @@ func TestIOConnFrameCap(t *testing.T) { } // TestIOConnDropResponse checks that ioConn discards the response to a call -// the peer cancelled, and leaves it out of the batch the call arrived in while -// still sending the batch's other responses. +// the peer cancelled on a session using protocol version 2026-07-28 or later, +// and still sends it on earlier versions. func TestIOConnDropResponse(t *testing.T) { - call := func(id int) string { return fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":"m"}`, id) } tests := []struct { name string - in string // the incoming payload - drop []int64 // the calls the peer cancelled - respond []int64 // the responses written, in order - want [][]any // the IDs of the responses on the wire, by line + version string + drop bool // whether the peer cancelled call 1 + want []any // the IDs of the responses on the wire, in order }{ - {"single call, dropped", call(1), []int64{1}, []int64{1}, nil}, - {"single call, not dropped", call(1), nil, []int64{1}, [][]any{{int64(1)}}}, - {"batch, one dropped", "[" + call(1) + "," + call(2) + "]", []int64{1}, []int64{1, 2}, [][]any{{int64(2)}}}, - {"batch, dropped last", "[" + call(1) + "," + call(2) + "]", []int64{2}, []int64{1, 2}, [][]any{{int64(1)}}}, - {"batch, all dropped", "[" + call(1) + "," + call(2) + "]", []int64{1, 2}, []int64{2, 1}, nil}, + {"new protocol, cancelled", protocolVersion20260728, true, []any{int64(2)}}, + {"new protocol, not cancelled", protocolVersion20260728, false, []any{int64(1), int64(2)}}, + {"earlier protocol, cancelled", protocolVersion20251125, true, []any{int64(1), int64(2)}}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { ctx := context.Background() var out bytes.Buffer tr := newIOConn(rwc{ - rc: io.NopCloser(strings.NewReader(test.in + "\n")), + rc: io.NopCloser(strings.NewReader("")), wc: nopCloserWriter{&out}, }) t.Cleanup(func() { tr.Close() }) + tr.sessionUpdated(ServerSessionState{NegotiatedProtocolVersion: test.version}) - for range strings.Count(test.in, `"method"`) { - if _, err := tr.Read(ctx); err != nil { - t.Fatal(err) - } - } - for _, id := range test.drop { - tr.DropResponse(jsonrpc2.Int64ID(id)) + if test.drop { + tr.DropResponse(jsonrpc2.Int64ID(1)) } - for _, id := range test.respond { + for _, id := range []int64{1, 2} { if err := tr.Write(ctx, &jsonrpc.Response{ID: jsonrpc2.Int64ID(id), Result: []byte("{}")}); err != nil { t.Fatalf("Write(response %d): %v", id, err) } } - var got [][]any + var got []any for line := range strings.Lines(out.String()) { - msgs, _, err := readBatch([]byte(line)) + msg, err := jsonrpc2.DecodeMessage([]byte(line)) if err != nil { t.Fatalf("decoding %q: %v", line, err) } - var ids []any - for _, m := range msgs { - ids = append(ids, m.(*jsonrpc.Response).ID.Raw()) - } - got = append(got, ids) + got = append(got, msg.(*jsonrpc.Response).ID.Raw()) } if fmt.Sprint(got) != fmt.Sprint(test.want) { t.Errorf("responses on the wire: %v, want %v", got, test.want) @@ -431,8 +419,8 @@ func TestIOConnDropResponse(t *testing.T) { } // TestIOConnCancelledCallGetsNoResponse checks end to end, over an in-memory -// transport wrapped for logging, that a call the client cancels with -// notifications/cancelled gets no response. +// transport wrapped for logging, that on a 2026-07-28 session a call the +// client cancels with notifications/cancelled gets no response. func TestIOConnCancelledCallGetsNoResponse(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() @@ -444,6 +432,9 @@ func TestIOConnCancelledCallGetsNoResponse(t *testing.T) { <-ctx.Done() return nil, nil, ctx.Err() }) + AddTool(server, &Tool{Name: "fast"}, func(context.Context, *CallToolRequest, struct{}) (*CallToolResult, any, error) { + return &CallToolResult{}, nil, nil + }) serverT, clientT := NewInMemoryTransports() var log safeBuffer ss, err := server.Connect(ctx, &LoggingTransport{Transport: serverT, Writer: &log}, nil) @@ -499,15 +490,15 @@ func TestIOConnCancelledCallGetsNoResponse(t *testing.T) { } } - send(req(1, methodInitialize, &InitializeParams{ProtocolVersion: protocolVersion20250618})) - awaitResponse(1) - send(req(0, notificationInitialized, &InitializedParams{})) - send(req(2, "tools/call", &CallToolParams{Name: "slow"})) + // A 2026-07-28 session has no initialize: each request declares the + // protocol version in its _meta. + meta := newProtocolMeta("") + send(req(2, "tools/call", &CallToolParams{Meta: meta, Name: "slow"})) <-started send(req(0, notificationCancelled, &CancelledParams{RequestID: int64(2)})) // The logging wrapper logs the response once ioConn has dropped it; after - // that, a response sent on the wire would precede the ping's. + // that, a response sent on the wire would precede the next call's. for !strings.Contains(string(log.Bytes()), `write: {"jsonrpc":"2.0","id":2,`) { select { case <-ctx.Done(): @@ -515,7 +506,7 @@ func TestIOConnCancelledCallGetsNoResponse(t *testing.T) { case <-time.After(5 * time.Millisecond): } } - send(req(3, "ping", nil)) + send(req(3, "tools/call", &CallToolParams{Meta: meta, Name: "fast"})) awaitResponse(3) } From d635fb67f309a101122ea64664d3dd5a1cf4435a Mon Sep 17 00:00:00 2001 From: guglielmoc Date: Tue, 13 Oct 2026 08:34:33 +0000 Subject: [PATCH 4/6] fix: move DropResponse call inside state lock to prevent race conditions during request cancellation --- internal/jsonrpc2/conn.go | 30 ++++++++++++------------------ 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index ec68c532..a40c144f 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -200,13 +200,12 @@ type Writer interface { } // ResponseDropper is an optional interface for a [Writer] that can drop the -// response to a call the peer cancelled (see [Connection.CancelFromPeer]). The -// Connection still writes the response once the call's handler returns, so -// that the writer sees the call finish, and the writer discards it. +// response to a call the peer cancelled (see [Connection.CancelFromPeer]). type ResponseDropper interface { // DropResponse reports, when the cancellation arrives, that the peer // cancelled the incoming call with the given ID, so its response must not - // be sent. + // be sent. It is called with the Connection's internal lock held, so it + // must return quickly and must not call back into the Connection. DropResponse(id ID) } @@ -463,9 +462,6 @@ func (ac *AsyncCall) Await(ctx context.Context, result any) error { // Cancel will not complain if the ID is not a currently active message, and it // will not cause any messages that have not arrived yet with that ID to be // cancelled. -// -// The inbound call is still answered: use [Connection.CancelFromPeer] when the -// peer itself asked for the cancellation. func (c *Connection) Cancel(id ID) { c.cancelIncoming(id, nil, false) } @@ -493,19 +489,17 @@ func (c *Connection) cancelIncoming(id ID, cause error, fromPeer bool) { var req *incomingRequest c.updateInFlight(func(s *inFlightState) { req = s.incomingByID[id] - }) - if req == nil { - return - } - if fromPeer { - // Tell the writer before cancelling, so that it knows by the time a - // handler that honors its context returns. This runs outside stateMu - // since the writer may do I/O. - if d, ok := c.writer.(ResponseDropper); ok { - d.DropResponse(id) + if req != nil && fromPeer { + // Telling the writer under stateMu makes this atomic with processResult + // removing the request before it writes the response. + if d, ok := c.writer.(ResponseDropper); ok { + d.DropResponse(id) + } } + }) + if req != nil { + req.cancel(cause) } - req.cancel(cause) } // Wait blocks until the connection is fully closed, but does not close it. From ffe0f812f251b1ea650f8beff82634084e6a746e Mon Sep 17 00:00:00 2001 From: guglielmoc Date: Tue, 13 Oct 2026 08:58:04 +0000 Subject: [PATCH 5/6] refactor: rename CancelCause to CancelFromPeer in transport layer --- mcp/transport.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mcp/transport.go b/mcp/transport.go index b94bc1c4..77cadc20 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -274,7 +274,7 @@ func (c *canceller) Preempt(ctx context.Context, req *jsonrpc.Request) (result a // travels as the cause of the request's context rather than being // dropped here. c.logger.Debug("request cancelled by the peer", "id", id.Raw(), "reason", params.Reason) - c.conn.CancelCause(id, &peerCancelledError{reason: params.Reason}) + c.conn.CancelFromPeer(id, &peerCancelledError{reason: params.Reason}) } return nil, jsonrpc2.ErrNotHandled } From d62e25890bd1076a2aee1b48fb90b7aa343142de Mon Sep 17 00:00:00 2001 From: guglielmoc Date: Wed, 14 Oct 2026 08:34:08 +0000 Subject: [PATCH 6/6] refactor: simplify and consolidate jsonrpc2 connection cancellation tests --- internal/jsonrpc2/conn_test.go | 228 +++++++++------------------------ 1 file changed, 59 insertions(+), 169 deletions(-) diff --git a/internal/jsonrpc2/conn_test.go b/internal/jsonrpc2/conn_test.go index c1818985..a154b526 100644 --- a/internal/jsonrpc2/conn_test.go +++ b/internal/jsonrpc2/conn_test.go @@ -7,7 +7,9 @@ package jsonrpc2 import ( "context" "errors" + "fmt" "io" + "slices" "sync" "testing" ) @@ -39,69 +41,57 @@ func TestShuttingDownWrapsReadError(t *testing.T) { } } -// errPeerAskedToStop stands in for the reason an MCP "notifications/cancelled" -// carries, so the test exercises the signature a real peer cancellation uses -// rather than a nil cause no caller passes. -var errPeerAskedToStop = errors.New("peer asked to stop") - -// TestCancelFromPeerTellsWriter verifies that a writer that drops responses is -// told about a call the peer cancelled before the handler's context is -// cancelled, and not about a call cancelled locally. The response is written -// either way; dropping it is up to the writer. The barrier is a second call, -// which handlers answer only after the first one. -func TestCancelFromPeerTellsWriter(t *testing.T) { - tests := []struct { - name string - fromPeer bool - dropper bool // whether the writer is a ResponseDropper - want []ID // the IDs of the responses written, in order - wantDrops []ID // the IDs DropResponse is called with - }{ - {"peer cancellation, plain writer", true, false, []ID{Int64ID(1), Int64ID(2)}, nil}, - {"peer cancellation, writer drops the response", true, true, []ID{Int64ID(1), Int64ID(2)}, []ID{Int64ID(1)}}, - {"local cancellation", false, true, []ID{Int64ID(1), Int64ID(2)}, nil}, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - rec := &recordingWriter{written: make(chan *Response, 2)} - dropper := &dropperWriter{recordingWriter: rec} - var writer Writer = rec - if test.dropper { - writer = dropper - } +// TestCancelFromPeer verifies that CancelFromPeer tells a [ResponseDropper] +// writer about the call before the handler's context is cancelled, that a +// local Cancel and a cancellation arriving after the response do not, and that +// the response is written either way: dropping it is up to the writer. +func TestCancelFromPeer(t *testing.T) { + for _, fromPeer := range []bool{true, false} { + t.Run(fmt.Sprintf("fromPeer=%t", fromPeer), func(t *testing.T) { + in := &chanReader{msgs: make(chan Message, 1)} + w := &dropWriter{written: make(chan *Response, 1)} started := make(chan struct{}) dropsAtCancel := make(chan []ID, 1) - conn, incoming := newCancelTestConn(t, writer, func(ctx context.Context) (any, error) { - close(started) - <-ctx.Done() - dropsAtCancel <- dropper.drops() - return nil, context.Cause(ctx) + conn := NewConnection(context.Background(), ConnectionConfig{ + Reader: in, + Writer: w, + Closer: in, + Bind: func(*Connection) Handler { + return HandlerFunc(func(ctx context.Context, _ *Request) (any, error) { + close(started) + <-ctx.Done() + dropsAtCancel <- w.drops() + return nil, context.Cause(ctx) + }) + }, }) - incoming <- mustCall(t, 1, "slow") + call, err := NewCall(Int64ID(1), "slow", nil) + if err != nil { + t.Fatal(err) + } + in.msgs <- call <-started - if test.fromPeer { - conn.CancelFromPeer(Int64ID(1), errPeerAskedToStop) + if fromPeer { + conn.CancelFromPeer(Int64ID(1), errors.New("peer asked to stop")) } else { conn.Cancel(Int64ID(1)) } - incoming <- mustCall(t, 2, "barrier") - var got []ID - for range test.want { - got = append(got, (<-rec.written).ID) + var want []ID + if fromPeer { + want = []ID{Int64ID(1)} } - if !equalIDs(got, test.want) { - t.Errorf("responses written for %v, want %v", got, test.want) + if got := <-dropsAtCancel; !slices.Equal(got, want) { + t.Errorf("when the handler's context was cancelled, DropResponse called for %v, want %v", got, want) } - // The writer is told before the handler's context is cancelled. - if drops := <-dropsAtCancel; !equalIDs(drops, test.wantDrops) { - t.Errorf("when the handler's context was cancelled, DropResponse called for %v, want %v", drops, test.wantDrops) + if resp := <-w.written; resp.ID != Int64ID(1) { + t.Errorf("response written for %v, want 1", resp.ID) } - if drops := dropper.drops(); !equalIDs(drops, test.wantDrops) { - t.Errorf("DropResponse called for %v, want %v", drops, test.wantDrops) + conn.CancelFromPeer(Int64ID(1), nil) // the call is answered: nothing to drop + if got := w.drops(); !slices.Equal(got, want) { + t.Errorf("after a late cancellation, DropResponse called for %v, want %v", got, want) } - if err := conn.Close(); err != nil { t.Errorf("Close() = %v", err) } @@ -109,105 +99,16 @@ func TestCancelFromPeerTellsWriter(t *testing.T) { } } -// TestCancelFromPeerHandlerIgnoresContext verifies that the writer is told of -// a peer cancellation as soon as it arrives, even when the handler ignores its -// context, that it still receives the response it drops once the handler -// returns, and that a cancellation arriving after that is ignored. -func TestCancelFromPeerHandlerIgnoresContext(t *testing.T) { - rec := &recordingWriter{written: make(chan *Response, 2)} - dropper := &dropperWriter{recordingWriter: rec} - started := make(chan struct{}) - release := make(chan struct{}) - conn, incoming := newCancelTestConn(t, dropper, func(context.Context) (any, error) { - close(started) - <-release - return struct{}{}, nil - }) - - incoming <- mustCall(t, 1, "slow") - <-started - conn.CancelFromPeer(Int64ID(1), errPeerAskedToStop) - - if drops := dropper.drops(); !equalIDs(drops, []ID{Int64ID(1)}) { - t.Errorf("after the cancellation, DropResponse called for %v, want [1]", drops) - } - - close(release) - incoming <- mustCall(t, 2, "barrier") - var got []ID - for range 2 { - got = append(got, (<-rec.written).ID) - } - if want := []ID{Int64ID(1), Int64ID(2)}; !equalIDs(got, want) { - t.Errorf("responses written for %v, want %v", got, want) - } - conn.CancelFromPeer(Int64ID(1), errPeerAskedToStop) - if drops := dropper.drops(); !equalIDs(drops, []ID{Int64ID(1)}) { - t.Errorf("after a late cancellation, DropResponse called for %v, want [1]", drops) - } - - if err := conn.Close(); err != nil { - t.Errorf("Close() = %v", err) - } +// chanReader is a Reader and Closer that delivers the messages sent on msgs, +// and reports io.EOF once closed. +type chanReader struct { + msgs chan Message + once sync.Once } -// newCancelTestConn returns a Connection reading from the returned channel -// and writing to writer. Its handler answers "barrier" at once and runs slow -// for any other method. Handlers run one at a time, so once the barrier is -// answered, the slow call is finished. -func newCancelTestConn(t *testing.T, writer Writer, slow func(context.Context) (any, error)) (*Connection, chan Message) { - t.Helper() - incoming := make(chan Message, 2) - handler := HandlerFunc(func(ctx context.Context, req *Request) (any, error) { - if req.Method == "barrier" { - return struct{}{}, nil - } - return slow(ctx) - }) - conn := NewConnection(context.Background(), ConnectionConfig{ - Reader: &channelReader{messages: incoming}, - Writer: writer, - Closer: &channelCloser{messages: incoming}, - Bind: func(*Connection) Handler { return handler }, - OnDone: func() {}, - OnInternalError: func(err error) { - t.Errorf("internal error: %v", err) - }, - }) - return conn, incoming -} - -func mustCall(t *testing.T, id int64, method string) *Request { - t.Helper() - call, err := NewCall(Int64ID(id), method, nil) - if err != nil { - t.Fatal(err) - } - return call -} - -// equalIDs reports whether two ID slices hold the same IDs in the same order. -func equalIDs(got, want []ID) bool { - if len(got) != len(want) { - return false - } - for i := range got { - if got[i].Raw() != want[i].Raw() { - return false - } - } - return true -} - -// channelReader delivers the messages sent to it, and reports io.EOF once the -// channel is closed. -type channelReader struct { - messages chan Message -} - -func (r *channelReader) Read(ctx context.Context) (Message, error) { +func (r *chanReader) Read(ctx context.Context) (Message, error) { select { - case msg, ok := <-r.messages: + case msg, ok := <-r.msgs: if !ok { return nil, io.EOF } @@ -217,46 +118,35 @@ func (r *channelReader) Read(ctx context.Context) (Message, error) { } } -// channelCloser unblocks a [channelReader] reading from the same channel. -type channelCloser struct { - once sync.Once - messages chan Message -} - -func (c *channelCloser) Close() error { - c.once.Do(func() { close(c.messages) }) +func (r *chanReader) Close() error { + r.once.Do(func() { close(r.msgs) }) return nil } -// recordingWriter reports the responses a Connection writes. -type recordingWriter struct { +// dropWriter is a Writer and [ResponseDropper] that reports the responses it +// is given and records the calls it is told to drop. +type dropWriter struct { written chan *Response + + mu sync.Mutex + dropped []ID } -func (w *recordingWriter) Write(_ context.Context, msg Message) error { +func (w *dropWriter) Write(_ context.Context, msg Message) error { if resp, ok := msg.(*Response); ok { w.written <- resp } return nil } -// dropperWriter is a recordingWriter that is also a [ResponseDropper], -// recording the IDs it is told about. -type dropperWriter struct { - *recordingWriter - - mu sync.Mutex - dropped []ID -} - -func (w *dropperWriter) DropResponse(id ID) { +func (w *dropWriter) DropResponse(id ID) { w.mu.Lock() defer w.mu.Unlock() w.dropped = append(w.dropped, id) } -func (w *dropperWriter) drops() []ID { +func (w *dropWriter) drops() []ID { w.mu.Lock() defer w.mu.Unlock() - return append([]ID(nil), w.dropped...) + return slices.Clone(w.dropped) }