internal/mcp: meta and progress token

Meta is a struct that contains a map[string]any, and a ProgressToken
field.

Each Params and Result type has a Meta as a value field.

Meta marshals and unmarshals specially, so the ProgressToken field
becomes part of the map. The marshalling code is generalized in
anticipation of every Params and Result growing a map of additional
properties.

Change-Id: I6d2f34270d33501c96437f24f91ef3e51d20d955
Reviewed-on: https://go-review.googlesource.com/c/tools/+/674956
Reviewed-by: Sam Thanawalla <samthanawalla@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Auto-Submit: Jonathan Amsterdam <jba@google.com>
Reviewed-by: Robert Findley <rfindley@google.com>

mcp/shared.go

@@ -14,6 +14,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log"
+	"reflect"
 	"slices"
 	"time"
 
@@ -21,12 +22,10 @@ import (
 )
 
 // A MethodHandler handles MCP messages.
-// The params argument is an XXXParams struct pointer, such as *GetPromptParams.
-// For methods, a MethodHandler must return either an XXResult struct pointer and a nil error, or
-// nil with a non-nil error.
-// For notifications, a MethodHandler must return nil, nil.
+// For methods, exactly one of the return values must be nil.
+// For notifications, both must be nil.
 type MethodHandler[S ClientSession | ServerSession] func(
-	ctx context.Context, _ *S, method string, params any) (result any, err error)
+	ctx context.Context, _ *S, method string, params Params) (result Result, err error)
 
 // Middleware is a function from MethodHandlers to MethodHandlers.
 type Middleware[S ClientSession | ServerSession] func(MethodHandler[S]) MethodHandler[S]
@@ -56,7 +55,7 @@ func toSession[S ClientSession | ServerSession](sess *S) session[S] {
 }
 
 // defaultMethodHandler is the initial MethodHandler for servers and clients, before being wrapped by middleware.
-func defaultMethodHandler[S ClientSession | ServerSession](ctx context.Context, sess *S, method string, params any) (any, error) {
+func defaultMethodHandler[S ClientSession | ServerSession](ctx context.Context, sess *S, method string, params Params) (Result, error) {
 	info, ok := toSession(sess).methodInfos()[method]
 	if !ok {
 		// This can be called from user code, with an arbitrary value for method.
@@ -87,7 +86,7 @@ func handleRequest[S ClientSession | ServerSession](ctx context.Context, req *js
 // methodInfo is information about invoking a method.
 type methodInfo[TSession ClientSession | ServerSession] struct {
 	// unmarshal params from the wire into an XXXParams struct
-	unmarshalParams func(json.RawMessage) (any, error)
+	unmarshalParams func(json.RawMessage) (Params, error)
 	// run the code for the method
 	handleMethod MethodHandler[TSession]
 }
@@ -99,40 +98,40 @@ type methodInfo[TSession ClientSession | ServerSession] struct {
 // - R: results
 
 // A typedMethodHandler is like a MethodHandler, but with type information.
-type typedMethodHandler[S, P, R any] func(context.Context, *S, P) (R, error)
+type typedMethodHandler[S any, P Params, R Result] func(context.Context, *S, P) (R, error)
 
 // newMethodInfo creates a methodInfo from a typedMethodHandler.
-func newMethodInfo[S ClientSession | ServerSession, P, R any](d typedMethodHandler[S, P, R]) methodInfo[S] {
+func newMethodInfo[S ClientSession | ServerSession, P Params, R Result](d typedMethodHandler[S, P, R]) methodInfo[S] {
 	return methodInfo[S]{
-		unmarshalParams: func(m json.RawMessage) (any, error) {
+		unmarshalParams: func(m json.RawMessage) (Params, error) {
 			var p P
 			if err := json.Unmarshal(m, &p); err != nil {
 				return nil, fmt.Errorf("unmarshaling %q into a %T: %w", m, p, err)
 			}
 			return p, nil
 		},
-		handleMethod: func(ctx context.Context, ss *S, _ string, params any) (any, error) {
+		handleMethod: func(ctx context.Context, ss *S, _ string, params Params) (Result, error) {
 			return d(ctx, ss, params.(P))
 		},
 	}
 }
 
 // serverMethod is glue for creating a typedMethodHandler from a method on Server.
-func serverMethod[P, R any](f func(*Server, context.Context, *ServerSession, P) (R, error)) typedMethodHandler[ServerSession, P, R] {
+func serverMethod[P Params, R Result](f func(*Server, context.Context, *ServerSession, P) (R, error)) typedMethodHandler[ServerSession, P, R] {
 	return func(ctx context.Context, ss *ServerSession, p P) (R, error) {
 		return f(ss.server, ctx, ss, p)
 	}
 }
 
 // clientMethod is glue for creating a typedMethodHandler from a method on Server.
-func clientMethod[P, R any](f func(*Client, context.Context, *ClientSession, P) (R, error)) typedMethodHandler[ClientSession, P, R] {
+func clientMethod[P Params, R Result](f func(*Client, context.Context, *ClientSession, P) (R, error)) typedMethodHandler[ClientSession, P, R] {
 	return func(ctx context.Context, cs *ClientSession, p P) (R, error) {
 		return f(cs.client, ctx, cs, p)
 	}
 }
 
 // sessionMethod is glue for creating a typedMethodHandler from a method on ServerSession.
-func sessionMethod[S ClientSession | ServerSession, P, R any](f func(*S, context.Context, P) (R, error)) typedMethodHandler[S, P, R] {
+func sessionMethod[S ClientSession | ServerSession, P Params, R Result](f func(*S, context.Context, P) (R, error)) typedMethodHandler[S, P, R] {
 	return func(ctx context.Context, sess *S, p P) (R, error) {
 		return f(sess, ctx, p)
 	}
@@ -151,7 +150,7 @@ const (
 	CodeUnsupportedMethod = -31001
 )
 
-func callNotificationHandler[S ClientSession | ServerSession, P any](ctx context.Context, h func(context.Context, *S, *P), sess *S, params *P) (any, error) {
+func callNotificationHandler[S ClientSession | ServerSession, P any](ctx context.Context, h func(context.Context, *S, *P), sess *S, params *P) (Result, error) {
 	if h != nil {
 		h(ctx, sess, params)
 	}
@@ -160,7 +159,7 @@ func callNotificationHandler[S ClientSession | ServerSession, P any](ctx context
 
 // notifySessions calls Notify on all the sessions.
 // Should be called on a copy of the peer sessions.
-func notifySessions[S ClientSession | ServerSession](sessions []*S, method string, params any) {
+func notifySessions[S ClientSession | ServerSession](sessions []*S, method string, params Params) {
 	if sessions == nil {
 		return
 	}
@@ -174,3 +173,57 @@ func notifySessions[S ClientSession | ServerSession](sessions []*S, method strin
 		}
 	}
 }
+
+func standardCall[TRes, TParams any](ctx context.Context, conn *jsonrpc2.Connection, method string, params TParams) (*TRes, error) {
+	var result TRes
+	if err := call(ctx, conn, method, params, &result); err != nil {
+		return nil, err
+	}
+	return &result, nil
+}
+
+type Meta struct {
+	Data map[string]any `json:",omitempty"`
+	// For params, the progress token can be nil, a string or an integer.
+	// It should be nil for results.
+	ProgressToken any `json:"progressToken,omitempty"`
+}
+
+type metaSansMethods Meta // avoid infinite recursion during marshaling
+
+func (m Meta) MarshalJSON() ([]byte, error) {
+	if p := m.ProgressToken; p != nil {
+		if k := reflect.ValueOf(p).Kind(); k != reflect.Int && k != reflect.String {
+			return nil, fmt.Errorf("bad type %T for Meta.ProgressToken: must be int or string", p)
+		}
+	}
+	// If ProgressToken is nil, accept Data["progressToken"]. We can't call marshalStructWithMap
+	// in that case because it will complain about duplicate fields. (We'd have to
+	// make it much smarter to avoid that; not worth it.)
+	if m.ProgressToken == nil {
+		return json.Marshal(m.Data)
+	}
+	return marshalStructWithMap((*metaSansMethods)(&m), "Data")
+}
+
+func (m *Meta) UnmarshalJSON(data []byte) error {
+	return unmarshalStructWithMap(data, (*metaSansMethods)(m), "Data")
+}
+
+// Params is a parameter (input) type for an MCP call or notification.
+type Params interface {
+	// Returns a pointer to the params's Meta field.
+	GetMeta() *Meta
+}
+
+// Result is a result of an MCP call.
+type Result interface {
+	// Returns a pointer to the result's Meta field.
+	GetMeta() *Meta
+}
+
+// emptyResult is returned by methods that have no result, like ping.
+// Those methods cannot return nil, because jsonrpc2 cannot handle nils.
+type emptyResult struct{}
+
+func (emptyResult) GetMeta() *Meta { panic("should never be called") }