internal/mcp: implement cancellation

Use the existing jsonrpc2 preemption mechanism to implement MCP
cancellation.

This was mostly straightforward, except where I got confused about ID
unmarshaling. Leave a note about using omitempty when it is available.

Change-Id: I13e073f08c5d5c2cc78d882da4e6ff47f09fb340
Reviewed-on: https://go-review.googlesource.com/c/tools/+/667578
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>

internal/jsonrpc2/messages.go

@@ -10,11 +10,32 @@ import (
 	"fmt"
 )
 
-// ID is a Request identifier.
+// ID is a Request identifier, which is defined by the spec to be a string, integer, or null.
+// https://www.jsonrpc.org/specification#request_object
 type ID struct {
 	value any
 }
 
+// MakeID coerces the given Go value to an ID. The value is assumed to be the
+// default JSON marshaling of a Request identifier -- nil, float64, or string.
+//
+// Returns an error if the value type was a valid Request ID type.
+//
+// TODO: ID can't be a json.Marshaler/Unmarshaler, because we want to omitzero.
+// Simplify this package by making ID json serializable once we can rely on
+// omitzero.
+func MakeID(v any) (ID, error) {
+	switch v := v.(type) {
+	case nil:
+		return ID{}, nil
+	case float64:
+		return Int64ID(int64(v)), nil
+	case string:
+		return StringID(v), nil
+	}
+	return ID{}, fmt.Errorf("%w: invalid ID type %T", ErrParse, v)
+}
+
 // Message is the interface to all jsonrpc2 message types.
 // They share no common functionality, but are a closed set of concrete types
 // that are allowed to implement this interface. The message types are *Request
@@ -133,18 +154,9 @@ func DecodeMessage(data []byte) (Message, error) {
 	if msg.VersionTag != wireVersion {
 		return nil, fmt.Errorf("invalid message version tag %s expected %s", msg.VersionTag, wireVersion)
 	}
-	id := ID{}
-	switch v := msg.ID.(type) {
-	case nil:
-	case float64:
-		// coerce the id type to int64 if it is float64, the spec does not allow fractional parts
-		id = Int64ID(int64(v))
-	case int64:
-		id = Int64ID(v)
-	case string:
-		id = StringID(v)
-	default:
-		return nil, fmt.Errorf("invalid message id type <%T>%v", v, v)
+	id, err := MakeID(msg.ID)
+	if err != nil {
+		return nil, err
 	}
 	if msg.Method != "" {
 		// has a method, must be a call

internal/protocol/generate.go

@@ -56,6 +56,9 @@ var declarations = config{
 	"CallToolResult": {
 		Name: "CallToolResult",
 	},
+	"CancelledNotification": {
+		Fields: config{"Params": {Name: "CancelledParams"}},
+	},
 	"ClientCapabilities": {Name: "ClientCapabilities"},
 	"Implementation":     {Name: "Implementation"},
 	"InitializeRequest": {
@@ -73,7 +76,8 @@ var declarations = config{
 	"ListToolsResult": {
 		Name: "ListToolsResult",
 	},
-	"Role": {Name: "Role"},
+	"RequestId": {Substitute: "any"}, // null|number|string
+	"Role":      {Name: "Role"},
 	"ServerCapabilities": {
 		Name: "ServerCapabilities",
 		Fields: config{

internal/protocol/protocol.go

@@ -54,6 +54,17 @@ type CallToolResult struct {
 	IsError bool `json:"isError,omitempty"`
 }
 
+type CancelledParams struct {
+	// An optional string describing the reason for the cancellation. This MAY be
+	// logged or presented to the user.
+	Reason string `json:"reason,omitempty"`
+	// The ID of the request to cancel.
+	//
+	// This MUST correspond to the ID of a request previously issued in the same
+	// direction.
+	RequestId any `json:"requestId"`
+}
+
 // Capabilities a client may support. Known capabilities are defined here, in
 // this schema, but this is not a closed set: any client can define its own,
 // additional capabilities.

mcp/mcp.go

@@ -12,7 +12,6 @@
 // [Client.Connect] or [Server.Connect].
 //
 // TODO:
-//   - Support cancellation.
 //   - Support pagination.
 //   - Support all client/server operations.
 //   - Support Streamable HTTP transport.

mcp/mcp_test.go

@@ -133,14 +133,21 @@ func TestEndToEnd(t *testing.T) {
 	}
 }
 
-func TestServerClosing(t *testing.T) {
+// basicConnection returns a new basic client-server connection configured with
+// the provided tools.
+//
+// The caller should cancel either the client connection or server connection
+// when the connections are no longer needed.
+func basicConnection(t *testing.T, tools ...*Tool) (*ClientConnection, *ServerConnection) {
+	t.Helper()
+
 	ctx := context.Background()
 	ct, st := NewLocalTransport()
 
 	s := NewServer("testServer", "v1.0.0", nil)
 
 	// The 'greet' tool says hi.
-	s.AddTools(MakeTool("greet", "say hi", sayHi))
+	s.AddTools(tools...)
 	cc, err := s.Connect(ctx, st, nil)
 	if err != nil {
 		t.Fatal(err)
@@ -151,7 +158,14 @@ func TestServerClosing(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	return cc, sc
+}
 
+func TestServerClosing(t *testing.T) {
+	cc, sc := basicConnection(t, MakeTool("greet", "say hi", sayHi))
+	defer sc.Close()
+
+	ctx := context.Background()
 	var wg sync.WaitGroup
 	wg.Add(1)
 	go func() {
@@ -209,3 +223,33 @@ func TestBatching(t *testing.T) {
 	}
 
 }
+
+func TestCancellation(t *testing.T) {
+	var (
+		start     = make(chan struct{})
+		cancelled = make(chan struct{}, 1) // don't block the request
+	)
+
+	slowRequest := func(ctx context.Context, cc *ClientConnection, v struct{}) ([]Content, error) {
+		start <- struct{}{}
+		select {
+		case <-ctx.Done():
+			cancelled <- struct{}{}
+		case <-time.After(5 * time.Second):
+			return nil, nil
+		}
+		return nil, nil
+	}
+	_, sc := basicConnection(t, MakeTool("slow", "a slow request", slowRequest))
+	defer sc.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	go sc.CallTool(ctx, "slow", struct{}{})
+	<-start
+	cancel()
+	select {
+	case <-cancelled:
+	case <-time.After(5 * time.Second):
+		t.Fatal("timeout waiting for cancellation")
+	}
+}

mcp/server.go

@@ -240,7 +240,7 @@ func (cc *ClientConnection) Wait() error {
 	return cc.conn.Wait()
 }
 
-// dispatch turns a strongly type handler into a jsonrpc2 handler.
+// dispatch turns a strongly type request handler into a jsonrpc2 handler.
 //
 // Importantly, it returns nil if the handler returned an error, which is a
 // requirement of the jsonrpc2 package.

mcp/transport.go

@@ -15,6 +15,8 @@ import (
 	"sync"
 
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
+	"github.com/modelcontextprotocol/go-sdk/internal/protocol"
+	"github.com/modelcontextprotocol/go-sdk/internal/xcontext"
 )
 
 // A JSONRPC2 error is an error defined by the JSONRPC2 spec.
@@ -103,31 +105,67 @@ func connect[H handler](ctx context.Context, t Transport, opts *ConnectionOption
 		writer = loggingWriter(opts.Logger, writer)
 	}
 
-	var h H
+	var (
+		h         H
+		preempter canceller
+	)
 	bind := func(conn *jsonrpc2.Connection) jsonrpc2.Handler {
 		h = b.bind(conn)
+		preempter.conn = conn
 		return jsonrpc2.HandlerFunc(h.handle)
 	}
 	_ = jsonrpc2.NewConnection(ctx, jsonrpc2.ConnectionConfig{
-		Reader: reader,
-		Writer: writer,
-		Closer: stream,
-		Bind:   bind,
+		Reader:    reader,
+		Writer:    writer,
+		Closer:    stream,
+		Bind:      bind,
+		Preempter: &preempter,
 		OnDone: func() {
 			b.disconnect(h)
 		},
 	})
+	assert(preempter.conn != nil, "unbound preempter")
 	assert(h != zero, "unbound connection")
 	return h, nil
 }
 
+// A canceller is a jsonrpc2.Preempter that cancels in-flight requests on MCP
+// cancelled notifications.
+type canceller struct {
+	conn *jsonrpc2.Connection
+}
+
+// Preempt implements jsonrpc2.Preempter.
+func (c *canceller) Preempt(ctx context.Context, req *jsonrpc2.Request) (result any, err error) {
+	if req.Method == "notifications/cancelled" {
+		var params protocol.CancelledParams
+		if err := json.Unmarshal(req.Params, &params); err != nil {
+			return nil, err
+		}
+		id, err := jsonrpc2.MakeID(params.RequestId)
+		if err != nil {
+			return nil, err
+		}
+		go c.conn.Cancel(id)
+	}
+	return nil, jsonrpc2.ErrNotHandled
+}
+
 // call executes and awaits a jsonrpc2 call on the given connection,
 // translating errors into the mcp domain.
 func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params, result any) error {
-	err := conn.Call(ctx, method, params).Await(ctx, result)
+	call := conn.Call(ctx, method, params)
+	err := call.Await(ctx, result)
 	switch {
 	case errors.Is(err, jsonrpc2.ErrClientClosing), errors.Is(err, jsonrpc2.ErrServerClosing):
 		return fmt.Errorf("calling %q: %w", method, ErrConnectionClosed)
+	case ctx.Err() != nil:
+		// Notify the peer of cancellation.
+		err := conn.Notify(xcontext.Detach(ctx), "notifications/cancelled", &protocol.CancelledParams{
+			Reason:    ctx.Err().Error(),
+			RequestId: call.ID().Raw(),
+		})
+		return errors.Join(ctx.Err(), err)
 	case err != nil:
 		return fmt.Errorf("calling %q: %v", method, err)
 	}