internal/mcp: client middleware

Implement middleware for clients.

Make the middleware types generic so we can factor out code
from both peers.

I was able to make the request handler generic, with one weird trick
(see handleRequest).

The first three declarations of shared.go, and a few others as well,
are copied unchanged from middleware.go.

Change-Id: Ib3239d8fcfda022ddb5b554c38cff4cdbc3704c1
Reviewed-on: https://go-review.googlesource.com/c/tools/+/673635
Reviewed-by: Robert Findley <rfindley@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

mcp/shared.go

@@ -0,0 +1,136 @@
+// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+// This file contains code shared between client and server, including
+// method handler and middleware definitions.
+// TODO: much of this is here so that we can factor out commonalities using
+// generics. Perhaps it can be simplified with reflection.
+
+package mcp
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"slices"
+
+	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
+)
+
+// 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.
+type MethodHandler[S ClientSession | ServerSession] func(
+	ctx context.Context, _ *S, method string, params any) (result any, err error)
+
+// Middleware is a function from MethodHandlers to MethodHandlers.
+type Middleware[S ClientSession | ServerSession] func(MethodHandler[S]) MethodHandler[S]
+
+// addMiddleware wraps the handler in the middleware functions.
+func addMiddleware[S ClientSession | ServerSession](handlerp *MethodHandler[S], middleware []Middleware[S]) {
+	for _, m := range slices.Backward(middleware) {
+		*handlerp = m(*handlerp)
+	}
+}
+
+// session has methods common to both ClientSession and ServerSession.
+type session[S ClientSession | ServerSession] interface {
+	methodHandler() MethodHandler[S]
+	methodInfos() map[string]methodInfo[S]
+}
+
+// toSession[S] converts its argument to a session[S].
+// Note that since S is constrained to ClientSession | ServerSession, and pointers to those
+// types both implement session[S] already, this should be a no-op.
+// That it is not, is due (I believe) to a deficency in generics, possibly related to core types.
+// TODO(jba): revisit in Go 1.26; perhaps the change in spec due to the removal of core types
+// will have resulted by then in a more generous implementation.
+func toSession[S ClientSession | ServerSession](sess *S) session[S] {
+	return any(sess).(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) {
+	info, ok := toSession(sess).methodInfos()[method]
+	if !ok {
+		// This can be called from user code, with an arbitrary value for method.
+		return nil, jsonrpc2.ErrNotHandled
+	}
+	return info.handleMethod(ctx, sess, method, params)
+}
+
+func handleRequest[S ClientSession | ServerSession](ctx context.Context, req *jsonrpc2.Request, sess *S) (any, error) {
+	info, ok := toSession(sess).methodInfos()[req.Method]
+	if !ok {
+		return nil, jsonrpc2.ErrNotHandled
+	}
+	params, err := info.unmarshalParams(req.Params)
+	if err != nil {
+		return nil, fmt.Errorf("handleRequest %q: %w", req.Method, err)
+	}
+
+	// mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol.
+	mh := toSession(sess).methodHandler()
+	res, err := mh(ctx, sess, req.Method, params)
+	if err != nil {
+		return nil, err
+	}
+	return res, nil
+}
+
+// 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)
+	// run the code for the method
+	handleMethod MethodHandler[TSession]
+}
+
+// The following definitions support converting from typed to untyped method handlers.
+// Type parameter meanings:
+// - S: sessions
+// - P: params
+// - 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)
+
+// newMethodInfo creates a methodInfo from a typedMethodHandler.
+func newMethodInfo[S ClientSession | ServerSession, P, R any](d typedMethodHandler[S, P, R]) methodInfo[S] {
+	return methodInfo[S]{
+		unmarshalParams: func(m json.RawMessage) (any, 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) {
+			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] {
+	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] {
+	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] {
+	return func(ctx context.Context, sess *S, p P) (R, error) {
+		return f(sess, ctx, p)
+	}
+}