mcp: send no response for a request the peer cancelled

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.
+
 ```go
 func Example_cancellation() {
 	// For this example, we're going to be collecting observations from the

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.
+
 %include ../../mcp/mcp_example_test.go cancellation -
 
 ### Ping

internal/jsonrpc2/conn.go

@@ -173,6 +173,12 @@ 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 the request context ending for some other
+	// reason. It is set by [Connection.CancelFromPeer] and read by
+	// processResult, both under the connection's stateMu.
+	peerCancelled bool
 }
 
 // Reader abstracts the transport mechanics from the JSON RPC protocol.
@@ -199,6 +205,21 @@ 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 will never be answered.
+//
+// A Connection sends no response to a call the peer cancelled (see
+// [Connection.CancelFromPeer]). For most writers that is simply one write that
+// does not happen, but a writer that maps calls onto transport-level streams
+// has bookkeeping to release: the streamable HTTP transport keeps a POST's
+// stream open until every call it carried has been answered, and without this
+// the stream would hang until the client went away.
+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,10 +473,38 @@ 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, false)
+}
+
+// CancelFromPeer is [Connection.Cancel] for a cancellation the peer requested,
+// such as an MCP "notifications/cancelled" message.
+//
+// In addition to cancelling the handler's Context, it suppresses the response
+// to the inbound call: a peer that asked for a call to be cancelled is not
+// waiting for its result, and the MCP specification says receivers of a
+// cancellation notification should not send a response for the cancelled
+// request.
+func (c *Connection) CancelFromPeer(id ID) {
+	c.cancelIncoming(id, true)
+}
+
+// cancelIncoming cancels the inbound request with the given ID, recording
+// whether the peer is the one that asked for it.
+//
+// A request that has already been responded to is no longer in incomingByID,
+// so a cancellation that loses the race with its own response is a no op, and
+// the response that was already on its way is never retracted.
+func (c *Connection) cancelIncoming(id ID, 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(nil)
@@ -710,17 +759,30 @@ 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. The Writer
+			// may still be holding state for it (the streamable HTTP transport
+			// keeps a POST's stream open until every call in it has been
+			// answered), so tell it 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 {
@@ -806,10 +868,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 {

internal/jsonrpc2/conn_test.go

@@ -5,8 +5,10 @@
 package jsonrpc2
 
 import (
+	"context"
 	"errors"
 	"io"
+	"sync"
 	"testing"
 )
 
@@ -36,3 +38,156 @@ func TestShuttingDownWrapsReadError(t *testing.T) {
 		t.Errorf("shuttingDown() error = %v, want it to wrap io.EOF", err)
 	}
 }
+
+// TestCancelFromPeerSuppressesResponse verifies that a call the peer asked to
+// cancel receives no response, while a call cancelled locally still does.
+//
+// The barrier in both cases is a second call: handlers run one at a time, and
+// a request is only dequeued once the previous one has been responded to, so
+// the answer to the barrier cannot overtake an answer to the first call.
+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))
+			} 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...)
+}

mcp/streamable.go

@@ -1068,6 +1068,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.
 //
@@ -1113,8 +1116,10 @@ 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 {
 			// Flush all pending messages as JSON response.
 			var toWrite []byte
 			if len(s.pendingJSONMessages) == 1 && !s.isBatch {
@@ -1129,7 +1134,7 @@ 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 {
@@ -1944,6 +1949,39 @@ func (c *streamableServerConn) Write(ctx context.Context, msg jsonrpc.Message) e
 	return nil
 }
 
+// DropResponse implements [jsonrpc2.ResponseDropper].
+//
+// The client cancelled this call and gets no response for it, but the stream
+// it arrived on must still be accounted for: a POST's stream hangs until every
+// call it carried has been answered, so without this the HTTP request would
+// stay open until the client went away.
+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()