internal/mcp: middleware on the sending side

Implement sending middleware.

Sending middleware is useful for adding tracing, progress tokens,
and so on.

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

mcp/shared.go

@@ -16,6 +16,7 @@ import (
 	"log"
 	"reflect"
 	"slices"
+	"strings"
 	"time"
 
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
@@ -36,8 +37,10 @@ type methodHandler any // MethodHandler[*ClientSession] | MethodHandler[*ServerS
 // A Session is either a ClientSession or a ServerSession.
 type Session interface {
 	*ClientSession | *ServerSession
-	methodHandler() methodHandler
-	methodInfos() map[string]methodInfo
+	sendingMethodInfos() map[string]methodInfo
+	receivingMethodInfos() map[string]methodInfo
+	sendingMethodHandler() methodHandler
+	receivingMethodHandler() methodHandler
 	getConn() *jsonrpc2.Connection
 }
 
@@ -51,9 +54,45 @@ func addMiddleware[S Session](handlerp *MethodHandler[S], middleware []Middlewar
 	}
 }
 
-// defaultMethodHandler is the initial MethodHandler for servers and clients, before being wrapped by middleware.
-func defaultMethodHandler[S Session](ctx context.Context, session S, method string, params Params) (Result, error) {
-	info, ok := session.methodInfos()[method]
+func defaultSendingMethodHandler[S Session](ctx context.Context, session S, method string, params Params) (Result, error) {
+	info, ok := session.sendingMethodInfos()[method]
+	if !ok {
+		// This can be called from user code, with an arbitrary value for method.
+		return nil, jsonrpc2.ErrNotHandled
+	}
+	// Notifications don't have results.
+	if strings.HasPrefix(method, "notifications/") {
+		return nil, session.getConn().Notify(ctx, method, params)
+	}
+	// Create the result to unmarshal into.
+	// The concrete type of the result is the return type of the receiving function.
+	res := info.newResult()
+	if err := call(ctx, session.getConn(), method, params, res); err != nil {
+		return nil, err
+	}
+	return res, nil
+}
+
+func handleNotify[S Session](ctx context.Context, session S, method string, params Params) error {
+	mh := session.sendingMethodHandler().(MethodHandler[S])
+	_, err := mh(ctx, session, method, params)
+	return err
+}
+
+func handleSend[R Result, S Session](ctx context.Context, s S, method string, params Params) (R, error) {
+	mh := s.sendingMethodHandler().(MethodHandler[S])
+	// mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol.
+	res, err := mh(ctx, s, method, params)
+	if err != nil {
+		var z R
+		return z, err
+	}
+	return res.(R), nil
+}
+
+// defaultReceivingMethodHandler is the initial MethodHandler for servers and clients, before being wrapped by middleware.
+func defaultReceivingMethodHandler[S Session](ctx context.Context, session S, method string, params Params) (Result, error) {
+	info, ok := session.receivingMethodInfos()[method]
 	if !ok {
 		// This can be called from user code, with an arbitrary value for method.
 		return nil, jsonrpc2.ErrNotHandled
@@ -61,8 +100,8 @@ func defaultMethodHandler[S Session](ctx context.Context, session S, method stri
 	return info.handleMethod.(MethodHandler[S])(ctx, session, method, params)
 }
 
-func handleRequest[S Session](ctx context.Context, req *jsonrpc2.Request, session S) (any, error) {
-	info, ok := session.methodInfos()[req.Method]
+func handleReceive[S Session](ctx context.Context, session S, req *jsonrpc2.Request) (Result, error) {
+	info, ok := session.receivingMethodInfos()[req.Method]
 	if !ok {
 		return nil, jsonrpc2.ErrNotHandled
 	}
@@ -71,8 +110,8 @@ func handleRequest[S Session](ctx context.Context, req *jsonrpc2.Request, sessio
 		return nil, fmt.Errorf("handleRequest %q: %w", req.Method, err)
 	}
 
+	mh := session.receivingMethodHandler().(MethodHandler[S])
 	// mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol.
-	mh := session.methodHandler().(MethodHandler[S])
 	res, err := mh(ctx, session, req.Method, params)
 	if err != nil {
 		return nil, err
@@ -80,12 +119,17 @@ func handleRequest[S Session](ctx context.Context, req *jsonrpc2.Request, sessio
 	return res, nil
 }
 
-// methodInfo is information about invoking a method.
+// methodInfo is information about sending and receiving a method.
 type methodInfo struct {
-	// unmarshal params from the wire into an XXXParams struct
+	// Unmarshal params from the wire into a Params struct.
+	// Used on the receive side.
 	unmarshalParams func(json.RawMessage) (Params, error)
-	// run the code for the method
+	// Run the code when a call to the method is received.
+	// Used on the receive side.
 	handleMethod methodHandler
+	// Create a pointer to a Result struct.
+	// Used on the send side.
+	newResult func() Result
 }
 
 // The following definitions support converting from typed to untyped method handlers.
@@ -112,6 +156,11 @@ func newMethodInfo[S Session, P Params, R Result](d typedMethodHandler[S, P, R])
 		handleMethod: MethodHandler[S](func(ctx context.Context, session S, _ string, params Params) (Result, error) {
 			return d(ctx, session, params.(P))
 		}),
+		// newResult is used on the send side, to construct the value to unmarshal the result into.
+		// R is a pointer to a result struct. There is no way to "unpointer" it without reflection.
+		// TODO(jba): explore generic approaches to this, perhaps by treating R in
+		// the signature as the unpointered type.
+		newResult: func() Result { return reflect.New(reflect.TypeFor[R]().Elem()).Interface().(R) },
 	}
 }
 
@@ -170,21 +219,13 @@ func notifySessions[S Session](sessions []S, method string, params Params) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 	for _, s := range sessions {
-		if err := s.getConn().Notify(ctx, method, params); err != nil {
+		if err := handleNotify(ctx, s, method, params); err != nil {
 			// TODO(jba): surface this error better
 			log.Printf("calling %s: %v", method, err)
 		}
 	}
 }
 
-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.
@@ -229,4 +270,4 @@ type Result interface {
 // Those methods cannot return nil, because jsonrpc2 cannot handle nils.
 type emptyResult struct{}
 
-func (emptyResult) GetMeta() *Meta { panic("should never be called") }
+func (*emptyResult) GetMeta() *Meta { panic("should never be called") }