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.