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/client.go

@@ -18,13 +18,14 @@ import (
 // A Client is an MCP client, which may be connected to an MCP server
 // using the [Client.Connect] method.
 type Client struct {
-	name           string
-	version        string
-	opts           ClientOptions
-	mu             sync.Mutex
-	roots          *featureSet[*Root]
-	sessions       []*ClientSession
-	methodHandler_ MethodHandler[*ClientSession]
+	name                    string
+	version                 string
+	opts                    ClientOptions
+	mu                      sync.Mutex
+	roots                   *featureSet[*Root]
+	sessions                []*ClientSession
+	sendingMethodHandler_   MethodHandler[*ClientSession]
+	receivingMethodHandler_ MethodHandler[*ClientSession]
 }
 
 // NewClient creates a new Client.
@@ -34,10 +35,11 @@ type Client struct {
 // If non-nil, the provided options configure the Client.
 func NewClient(name, version string, opts *ClientOptions) *Client {
 	c := &Client{
-		name:           name,
-		version:        version,
-		roots:          newFeatureSet(func(r *Root) string { return r.URI }),
-		methodHandler_: defaultMethodHandler[*ClientSession],
+		name:                    name,
+		version:                 version,
+		roots:                   newFeatureSet(func(r *Root) string { return r.URI }),
+		sendingMethodHandler_:   defaultSendingMethodHandler[*ClientSession],
+		receivingMethodHandler_: defaultReceivingMethodHandler[*ClientSession],
 	}
 	if opts != nil {
 		c.opts = *opts
@@ -103,11 +105,13 @@ func (c *Client) Connect(ctx context.Context, t Transport) (cs *ClientSession, e
 		ClientInfo:   &implementation{Name: c.name, Version: c.version},
 		Capabilities: caps,
 	}
-	if err := call(ctx, cs.conn, "initialize", params, &cs.initializeResult); err != nil {
+	res, err := handleSend[*InitializeResult](ctx, cs, methodInitialize, params)
+	if err != nil {
 		_ = cs.Close()
 		return nil, err
 	}
-	if err := cs.conn.Notify(ctx, notificationInitialized, &InitializedParams{}); err != nil {
+	cs.initializeResult = res
+	if err := handleNotify(ctx, cs, notificationInitialized, &InitializedParams{}); err != nil {
 		_ = cs.Close()
 		return nil, err
 	}
@@ -190,16 +194,34 @@ func (c *Client) createMessage(ctx context.Context, cs *ClientSession, params *C
 	return c.opts.CreateMessageHandler(ctx, cs, params)
 }
 
-// AddMiddleware wraps the client's current method handler using the provided
-// middleware. Middleware is applied from right to left, so that the first one
-// is executed first.
+// AddSendingMiddleware wraps the current sending method handler using the provided
+// middleware. Middleware is applied from right to left, so that the first one is
+// executed first.
 //
-// For example, AddMiddleware(m1, m2, m3) augments the client method handler as
+// For example, AddSendingMiddleware(m1, m2, m3) augments the method handler as
 // m1(m2(m3(handler))).
-func (c *Client) AddMiddleware(middleware ...Middleware[*ClientSession]) {
+//
+// Sending middleware is called when a request is sent. It is useful for tasks
+// such as tracing, metrics, and adding progress tokens.
+func (c *Client) AddSendingMiddleware(middleware ...Middleware[*ClientSession]) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	addMiddleware(&c.methodHandler_, middleware)
+	addMiddleware(&c.sendingMethodHandler_, middleware)
+}
+
+// AddReceivingMiddleware wraps the current receiving method handler using
+// the provided middleware. Middleware is applied from right to left, so that the
+// first one is executed first.
+//
+// For example, AddReceivingMiddleware(m1, m2, m3) augments the method handler as
+// m1(m2(m3(handler))).
+//
+// Receiving middleware is called when a request is received. It is useful for tasks
+// such as authentication, request logging and metrics.
+func (c *Client) AddReceivingMiddleware(middleware ...Middleware[*ClientSession]) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	addMiddleware(&c.receivingMethodHandler_, middleware)
 }
 
 // clientMethodInfos maps from the RPC method name to serverMethodInfos.
@@ -213,51 +235,62 @@ var clientMethodInfos = map[string]methodInfo{
 	notificationLoggingMessage:      newMethodInfo(clientMethod((*Client).callLoggingHandler)),
 }
 
-func (cs *ClientSession) methodInfos() map[string]methodInfo {
+func (cs *ClientSession) sendingMethodInfos() map[string]methodInfo {
+	return serverMethodInfos
+}
+
+func (cs *ClientSession) receivingMethodInfos() map[string]methodInfo {
 	return clientMethodInfos
 }
 
 func (cs *ClientSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
-	return handleRequest(ctx, req, cs)
+	return handleReceive(ctx, cs, req)
 }
 
-func (cs *ClientSession) methodHandler() methodHandler {
+func (cs *ClientSession) sendingMethodHandler() methodHandler {
 	cs.client.mu.Lock()
 	defer cs.client.mu.Unlock()
-	return cs.client.methodHandler_
+	return cs.client.sendingMethodHandler_
+}
+
+func (cs *ClientSession) receivingMethodHandler() methodHandler {
+	cs.client.mu.Lock()
+	defer cs.client.mu.Unlock()
+	return cs.client.receivingMethodHandler_
 }
 
 // getConn implements [session.getConn].
 func (cs *ClientSession) getConn() *jsonrpc2.Connection { return cs.conn }
 
-func (cs *ClientSession) ping(ct context.Context, params *PingParams) (Result, error) {
-	return emptyResult{}, nil
+func (*ClientSession) ping(context.Context, *PingParams) (*emptyResult, error) {
+	return &emptyResult{}, nil
 }
 
 // Ping makes an MCP "ping" request to the server.
 func (cs *ClientSession) Ping(ctx context.Context, params *PingParams) error {
-	return call(ctx, cs.conn, methodPing, params, nil)
+	_, err := handleSend[*emptyResult](ctx, cs, methodPing, params)
+	return err
 }
 
 // ListPrompts lists prompts that are currently available on the server.
 func (cs *ClientSession) ListPrompts(ctx context.Context, params *ListPromptsParams) (*ListPromptsResult, error) {
-	return standardCall[ListPromptsResult](ctx, cs.conn, methodListPrompts, params)
+	return handleSend[*ListPromptsResult](ctx, cs, methodListPrompts, params)
 }
 
 // GetPrompt gets a prompt from the server.
 func (cs *ClientSession) GetPrompt(ctx context.Context, params *GetPromptParams) (*GetPromptResult, error) {
-	return standardCall[GetPromptResult](ctx, cs.conn, methodGetPrompt, params)
+	return handleSend[*GetPromptResult](ctx, cs, methodGetPrompt, params)
 }
 
 // ListTools lists tools that are currently available on the server.
 func (cs *ClientSession) ListTools(ctx context.Context, params *ListToolsParams) (*ListToolsResult, error) {
-	return standardCall[ListToolsResult](ctx, cs.conn, methodListTools, params)
+	return handleSend[*ListToolsResult](ctx, cs, methodListTools, params)
 }
 
 // CallTool calls the tool with the given name and arguments.
 // Pass a [CallToolOptions] to provide additional request fields.
 func (cs *ClientSession) CallTool(ctx context.Context, params *CallToolParams[json.RawMessage]) (*CallToolResult, error) {
-	return standardCall[CallToolResult](ctx, cs.conn, methodCallTool, params)
+	return handleSend[*CallToolResult](ctx, cs, methodCallTool, params)
 }
 
 // CallTool is a helper to call a tool with any argument type. It returns an
@@ -286,17 +319,18 @@ func toWireParams[TArgs any](params *CallToolParams[TArgs]) (*CallToolParams[jso
 }
 
 func (cs *ClientSession) SetLevel(ctx context.Context, params *SetLevelParams) error {
-	return call(ctx, cs.conn, methodSetLevel, params, nil)
+	_, err := handleSend[*emptyResult](ctx, cs, methodSetLevel, params)
+	return err
 }
 
 // ListResources lists the resources that are currently available on the server.
 func (cs *ClientSession) ListResources(ctx context.Context, params *ListResourcesParams) (*ListResourcesResult, error) {
-	return standardCall[ListResourcesResult](ctx, cs.conn, methodListResources, params)
+	return handleSend[*ListResourcesResult](ctx, cs, methodListResources, params)
 }
 
 // ReadResource ask the server to read a resource and return its contents.
 func (cs *ClientSession) ReadResource(ctx context.Context, params *ReadResourceParams) (*ReadResourceResult, error) {
-	return standardCall[ReadResourceResult](ctx, cs.conn, methodReadResource, params)
+	return handleSend[*ReadResourceResult](ctx, cs, methodReadResource, params)
 }
 
 func (c *Client) callToolChangedHandler(ctx context.Context, s *ClientSession, params *ToolListChangedParams) (Result, error) {