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>

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)
 	}