diff --git a/docs/protocol.md b/docs/protocol.md index 06c1b77b..43c1b65e 100644 --- a/docs/protocol.md +++ b/docs/protocol.md @@ -712,6 +712,13 @@ 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)). +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() { // 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..f0380609 100644 --- a/internal/docs/protocol.src.md +++ b/internal/docs/protocol.src.md @@ -607,6 +607,13 @@ 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)). +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 - ### Ping diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 94da6ae3..a40c144f 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -199,6 +199,16 @@ type Writer interface { Write(context.Context, Message) error } +// ResponseDropper is an optional interface for a [Writer] that can drop the +// 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. 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) +} + // A ConnectionConfig configures a bidirectional jsonrpc2 connection. type ConnectionConfig struct { Reader Reader // required @@ -453,16 +463,39 @@ func (ac *AsyncCall) Await(ctx context.Context, result any) error { // will not cause any messages that have not arrived yet with that ID to be // cancelled. 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 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 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 { + // 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) @@ -813,10 +846,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..a154b526 100644 --- a/internal/jsonrpc2/conn_test.go +++ b/internal/jsonrpc2/conn_test.go @@ -5,8 +5,12 @@ package jsonrpc2 import ( + "context" "errors" + "fmt" "io" + "slices" + "sync" "testing" ) @@ -36,3 +40,113 @@ func TestShuttingDownWrapsReadError(t *testing.T) { t.Errorf("shuttingDown() error = %v, want it to wrap io.EOF", err) } } + +// 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 := 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) + }) + }, + }) + + call, err := NewCall(Int64ID(1), "slow", nil) + if err != nil { + t.Fatal(err) + } + in.msgs <- call + <-started + if fromPeer { + conn.CancelFromPeer(Int64ID(1), errors.New("peer asked to stop")) + } else { + conn.Cancel(Int64ID(1)) + } + + var want []ID + if fromPeer { + want = []ID{Int64ID(1)} + } + if got := <-dropsAtCancel; !slices.Equal(got, want) { + t.Errorf("when the handler's context was cancelled, DropResponse called for %v, want %v", got, want) + } + if resp := <-w.written; resp.ID != Int64ID(1) { + t.Errorf("response written for %v, want 1", resp.ID) + } + 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) + } + }) + } +} + +// 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 +} + +func (r *chanReader) Read(ctx context.Context) (Message, error) { + select { + case msg, ok := <-r.msgs: + if !ok { + return nil, io.EOF + } + return msg, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (r *chanReader) Close() error { + r.once.Do(func() { close(r.msgs) }) + return nil +} + +// 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 *dropWriter) Write(_ context.Context, msg Message) error { + if resp, ok := msg.(*Response); ok { + w.written <- resp + } + return nil +} + +func (w *dropWriter) DropResponse(id ID) { + w.mu.Lock() + defer w.mu.Unlock() + w.dropped = append(w.dropped, id) +} + +func (w *dropWriter) drops() []ID { + w.mu.Lock() + defer w.mu.Unlock() + return slices.Clone(w.dropped) +} diff --git a/mcp/transport.go b/mcp/transport.go index d2423d9d..77cadc20 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -250,7 +250,9 @@ type cancellationPropagator interface { } // A canceller is a jsonrpc2.Preempter that cancels in-flight requests on MCP -// cancelled notifications. +// 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 @@ -272,9 +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) - // Resolve the active request before the read loop can admit another - // invocation with the same ID. A deferred lookup could cancel that one. - c.conn.CancelCause(id, &peerCancelledError{reason: params.Reason}) + c.conn.CancelFromPeer(id, &peerCancelledError{reason: params.Reason}) } return nil, jsonrpc2.ErrNotHandled } @@ -430,6 +430,24 @@ 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, 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) + } +} + +// 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() } @@ -471,8 +489,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 @@ -634,6 +656,31 @@ func (t *ioConn) addBatch(batch *msgBatch) error { return nil } +// 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.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. // @@ -786,6 +833,9 @@ func (t *ioConn) Write(ctx context.Context, msg jsonrpc.Message) error { // 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 dadf196b..90f7cb62 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -5,6 +5,7 @@ package mcp import ( + "bufio" "bytes" "context" "errors" @@ -364,3 +365,193 @@ func TestIOConnFrameCap(t *testing.T) { }) } } + +// TestIOConnDropResponse checks that ioConn discards the response to a call +// 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) { + tests := []struct { + name string + version string + drop bool // whether the peer cancelled call 1 + want []any // the IDs of the responses on the wire, in order + }{ + {"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("")), + wc: nopCloserWriter{&out}, + }) + t.Cleanup(func() { tr.Close() }) + tr.sessionUpdated(ServerSessionState{NegotiatedProtocolVersion: test.version}) + + if test.drop { + tr.DropResponse(jsonrpc2.Int64ID(1)) + } + 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 + for line := range strings.Lines(out.String()) { + msg, err := jsonrpc2.DecodeMessage([]byte(line)) + if err != nil { + t.Fatalf("decoding %q: %v", line, err) + } + 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) + } + 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 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() + + 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() + }) + 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) + 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) + } + } + } + + // 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 next call'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, "tools/call", &CallToolParams{Meta: meta, Name: "fast"})) + 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. +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 drops responses", func(t *testing.T) { + delegate := &droppingConn{} + var conn Connection = &loggingConn{delegate: delegate, w: io.Discard} + + dropper, ok := conn.(jsonrpc2.ResponseDropper) + if !ok { + 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 does not", func(t *testing.T) { + conn := &loggingConn{delegate: plainConn{}, w: io.Discard} + conn.DropResponse(jsonrpc.ID{}) + }) +}