From 14b18b2645fdaa48b0201563e662c2b3068f8e10 Mon Sep 17 00:00:00 2001 From: Rob Findley Date: Wed, 23 Apr 2025 19:29:46 +0000 Subject: [PATCH] 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 Reviewed-by: Jonathan Amsterdam --- internal/jsonrpc2/messages.go | 38 +++++++++++++++++--------- internal/protocol/generate.go | 6 ++++- internal/protocol/protocol.go | 11 ++++++++ mcp/mcp.go | 1 - mcp/mcp_test.go | 48 +++++++++++++++++++++++++++++++-- mcp/server.go | 2 +- mcp/transport.go | 50 ++++++++++++++++++++++++++++++----- 7 files changed, 132 insertions(+), 24 deletions(-) diff --git a/internal/jsonrpc2/messages.go b/internal/jsonrpc2/messages.go index f58e083..e10abfa 100644 --- a/internal/jsonrpc2/messages.go +++ b/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 diff --git a/internal/protocol/generate.go b/internal/protocol/generate.go index 4e903df..d0697ee 100644 --- a/internal/protocol/generate.go +++ b/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{ diff --git a/internal/protocol/protocol.go b/internal/protocol/protocol.go index 08d9515..5b5fdfe 100644 --- a/internal/protocol/protocol.go +++ b/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. diff --git a/mcp/mcp.go b/mcp/mcp.go index 5de696b..bdaa57c 100644 --- a/mcp/mcp.go +++ b/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. diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 444bc7f..654be79 100644 --- a/mcp/mcp_test.go +++ b/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") + } +} diff --git a/mcp/server.go b/mcp/server.go index 1cf0e90..e2a4c27 100644 --- a/mcp/server.go +++ b/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. diff --git a/mcp/transport.go b/mcp/transport.go index 66728b3..17ff5a2 100644 --- a/mcp/transport.go +++ b/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, ¶ms); 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) }