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>
This commit is contained in:
Rob Findley 2025-04-23 19:29:46 +00:00 • committed by Robert Findley
commit 14b18b2645

View file

@ -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