diff --git a/docs/protocol.md b/docs/protocol.md index 06c1b77b..561bdcd5 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. + ```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..4389dc63 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. + %include ../../mcp/mcp_example_test.go cancellation - ### Ping diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 94da6ae3..a71ae346 100644 --- a/internal/jsonrpc2/conn.go +++ b/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,17 +473,49 @@ 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" message. +// +// In addition to cancelling the handler's Context with cause, 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. +// +// It takes a cause rather than cancelling plainly because a peer cancellation +// is the one kind that always arrives with a stated reason, and dropping it +// here would leave the handler unable to tell why it was stopped. +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 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, 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 +770,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 { @@ -813,10 +879,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..e63aadb4 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,161 @@ 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 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), 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 bb76c1cc..4c806bcb 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -787,6 +787,14 @@ 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) } + // A peer that asked for a cancellation is not waiting for a result, and + // the receiver now sends none. The SDK's own cancel path retires the + // outgoing call for exactly that reason before it notifies (cancelCall + // in transport.go); this test sends the notification by hand, so it + // retires the call by hand too. Without it the call stays pending for a + // response that will never come and synctest reports the bubble as + // deadlocked. + cs.conn.Retire(call, context.Canceled) cause := <-cancelled if cause == nil { diff --git a/mcp/streamable.go b/mcp/streamable.go index f24badbb..86e91290 100644 --- a/mcp/streamable.go +++ b/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() diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index ca4d0a74..91b96914 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -3735,6 +3735,102 @@ func TestCallCancellation_FastReturn(t *testing.T) { } } +// TestStreamableCancelledCallGetsNoResponse checks that a call the client +// cancelled with notifications/cancelled is answered with nothing at all, and +// that its stream still completes. +// +// It fakes the client with raw HTTP requests rather than using a +// [ClientSession]: an SDK client abandons the POST as soon as it cancels, so +// it never observes what the server wrote on that stream, and the stream not +// completing would look the same to it as the stream completing. +func TestStreamableCancelledCallGetsNoResponse(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 }, nil) + 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 does not return until the server writes something on the + // stream, and for a cancelled call the only thing it writes is the end of + // the stream, so make the call from its own 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() + + // The stream must end, and carry nothing: no event store is configured, so + // a conforming stream for this call has no events on it at all. + 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 ce841f23..91ebd4f0 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,20 @@ func (s *loggingConn) Write(ctx context.Context, msg jsonrpc.Message) error { return err } +// DropResponse implements [jsonrpc2.ResponseDropper] by forwarding to the +// delegate. +// +// Without it, wrapping a transport for logging would quietly reinstate the +// response to a cancelled call: the connection asks its writer for the +// interface, the wrapper does not have it, and the bookkeeping the delegate +// was waiting for never happens. A delegate that holds no per-call state +// implements nothing and there is nothing to forward. +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 d4d7782c..98f684ae 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -251,3 +251,55 @@ 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. +// +// A connection asks its writer for [jsonrpc2.ResponseDropper] and tells it +// when the peer cancelled a call. loggingConn is the one Connection wrapper in +// this package, and before it forwarded the call, a LoggingTransport around a +// streamable server left the cancelled call's POST stream open until the +// client went away: the response was suppressed, the bookkeeping was not. +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{}) + }) +}