From bac1b78b4a5707dc48ed5d8445e1f7ad0f3e212b Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Fri, 16 May 2025 14:28:28 -0400 Subject: [PATCH] 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 LUCI-TryBot-Result: Go LUCI --- design/design.md | 31 +++++++---- mcp/client.go | 74 +++++++++++++++++--------- mcp/mcp_test.go | 54 ++++++++++++------- mcp/server.go | 133 ++++++++++++--------------------------------- mcp/shared.go | 136 +++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 274 insertions(+), 154 deletions(-) create mode 100644 mcp/shared.go diff --git a/design/design.md b/design/design.md index a5da382..0c5bdfa 100644 --- a/design/design.md +++ b/design/design.md @@ -406,23 +406,34 @@ func (*ClientSession) ResourceTemplates(context.Context, *ListResourceTemplatesP ### Middleware -We provide a mechanism to add MCP-level middleware on the server side, which runs after the request has been parsed but before any normal handling. +We provide a mechanism to add MCP-level middleware on the both the client and server side, which runs after the request has been parsed but before any normal handling. ```go -// A ServerMethodHandler dispatches an MCP message to the appropriate handler. -// The params argument will be an XXXParams struct pointer, such as *GetPromptParams. -// The response if err is non-nil should be an XXXResult struct pointer. -type ServerMethodHandler func(ctx context.Context, s *ServerSession, method string, params any) (result any, err error) +// 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) -// AddMiddlewarecalls each function from right to left on the previous result, beginning -// with the server's current dispatcher, and installs the result as the new dispatcher. -func (*Server) AddMiddleware(middleware ...func(ServerMethodHandler) ServerMethodHandler) +// Middleware is a function from MethodHandlers to MethodHandlers. +type Middleware[S ClientSession | ServerSession] func(MethodHandler[S]) MethodHandler[S] + +// AddMiddleware wraps the client/server's current 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 server method handler as +// m1(m2(m3(handler))). +func (c *Client) AddMiddleware(middleware ...Middleware[ClientSession]) +func (s *Server) AddMiddleware(middleware ...Middleware[ServerSession]) ``` As an example, this code adds server-side logging: ```go -func withLogging(h mcp.ServerMethodHandler) mcp.ServerMethodHandler{ +func withLogging(h mcp.MethodHandler[ServerSession]) mcp.MethodHandler[ServerSession]{ return func(ctx context.Context, s *mcp.ServerSession, method string, params any) (res any, err error) { log.Printf("request: %s %v", method, params) defer func() { log.Printf("response: %v, %v", res, err) }() @@ -433,8 +444,6 @@ func withLogging(h mcp.ServerMethodHandler) mcp.ServerMethodHandler{ server.AddMiddleware(withLogging) ``` -We will provide the same functionality on the client side as well. - **Differences from mcp-go**: Version 0.26.0 of mcp-go defines 24 server hooks. Each hook consists of a field in the `Hooks` struct, a `Hooks.Add` method, and a type for the hook function. These are rarely used. The most common is `OnError`, which occurs fewer than ten times in open-source code. ### Errors diff --git a/mcp/client.go b/mcp/client.go index 8f2ee52..e1602d3 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -18,12 +18,13 @@ 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 + name string + version string + opts ClientOptions + mu sync.Mutex + roots *featureSet[*Root] + sessions []*ClientSession + methodHandler_ MethodHandler[ClientSession] } // NewClient creates a new Client. @@ -33,9 +34,10 @@ 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 }), + name: name, + version: version, + roots: newFeatureSet(func(r *Root) string { return r.URI }), + methodHandler_: defaultMethodHandler[ClientSession], } if opts != nil { c.opts = *opts @@ -140,7 +142,7 @@ func (c *Client) RemoveRoots(uris ...string) { c.roots.remove(uris...) } -func (c *Client) listRoots(_ context.Context, _ *ListRootsParams) (*ListRootsResult, error) { +func (c *Client) listRoots(_ context.Context, _ *ClientSession, _ *ListRootsParams) (*ListRootsResult, error) { c.mu.Lock() defer c.mu.Unlock() return &ListRootsResult{ @@ -148,21 +150,43 @@ func (c *Client) listRoots(_ context.Context, _ *ListRootsParams) (*ListRootsRes }, nil } -func (c *ClientSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { - // TODO: when we switch to ClientSessions, use a copy of the server's dispatch function, or - // maybe just add another type parameter. - // - // No need to check that the connection is initialized, since we initialize - // it in Connect. - switch req.Method { - case "ping": - // The spec says that 'ping' expects an empty object result. - return struct{}{}, nil - case "roots/list": - // ListRootsParams happens to be unused. - return c.client.listRoots(ctx, nil) - } - return nil, jsonrpc2.ErrNotHandled +// 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. +// +// For example, AddMiddleware(m1, m2, m3) augments the client method handler as +// m1(m2(m3(handler))). +func (c *Client) AddMiddleware(middleware ...Middleware[ClientSession]) { + c.mu.Lock() + defer c.mu.Unlock() + addMiddleware(&c.methodHandler_, middleware) +} + +// clientMethodInfos maps from the RPC method name to serverMethodInfos. +var clientMethodInfos = map[string]methodInfo[ClientSession]{ + "ping": newMethodInfo(sessionMethod((*ClientSession).ping)), + "roots/list": newMethodInfo(clientMethod((*Client).listRoots)), + // TODO: notifications +} + +var _ session[ClientSession] = (*ClientSession)(nil) + +func (cs *ClientSession) methodInfos() map[string]methodInfo[ClientSession] { + return clientMethodInfos +} + +func (cs *ClientSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { + return handleRequest(ctx, req, cs) +} + +func (cs *ClientSession) methodHandler() MethodHandler[ClientSession] { + cs.client.mu.Lock() + defer cs.client.mu.Unlock() + return cs.client.methodHandler_ +} + +func (c *ClientSession) ping(ct context.Context, params *PingParams) (struct{}, error) { + return struct{}{}, nil } // Ping makes an MCP "ping" request to the server. diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 1f697b6..a53a4f5 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -9,6 +9,7 @@ import ( "context" "errors" "fmt" + "io" "path/filepath" "runtime" "slices" @@ -27,8 +28,8 @@ type hiParams struct { Name string } -func sayHi(ctx context.Context, cc *ServerSession, v hiParams) ([]*Content, error) { - if err := cc.Ping(ctx, nil); err != nil { +func sayHi(ctx context.Context, ss *ServerSession, v hiParams) ([]*Content, error) { + if err := ss.Ping(ctx, nil); err != nil { return nil, fmt.Errorf("ping failed: %v", err) } return []*Content{NewTextContent("hi " + v.Name)}, nil @@ -398,7 +399,7 @@ func TestCancellation(t *testing.T) { } } -func TestAddMiddleware(t *testing.T) { +func TestMiddleware(t *testing.T) { ctx := context.Background() ct, st := NewInMemoryTransports() s := NewServer("testServer", "v1.0.0", nil) @@ -416,26 +417,17 @@ func TestAddMiddleware(t *testing.T) { clientWG.Done() }() - var buf bytes.Buffer - buf.WriteByte('\n') - - // traceCalls creates a middleware function that prints the method before and after each call - // with the given prefix. - traceCalls := func(prefix string) func(ServerMethodHandler) ServerMethodHandler { - return func(d ServerMethodHandler) ServerMethodHandler { - return func(ctx context.Context, ss *ServerSession, method string, params any) (any, error) { - fmt.Fprintf(&buf, "%s >%s\n", prefix, method) - defer fmt.Fprintf(&buf, "%s <%s\n", prefix, method) - return d(ctx, ss, method, params) - } - } - } + var sbuf, cbuf bytes.Buffer + sbuf.WriteByte('\n') + cbuf.WriteByte('\n') // "1" is the outer middleware layer, called first; then "2" is called, and finally // the default dispatcher. - s.AddMiddleware(traceCalls("1"), traceCalls("2")) + s.AddMiddleware(traceCalls[ServerSession](&sbuf, "1"), traceCalls[ServerSession](&sbuf, "2")) c := NewClient("testClient", "v1.0.0", nil) + c.AddMiddleware(traceCalls[ClientSession](&cbuf, "1"), traceCalls[ClientSession](&cbuf, "2")) + cs, err := c.Connect(ctx, ct) if err != nil { t.Fatal(err) @@ -453,9 +445,33 @@ func TestAddMiddleware(t *testing.T) { 2 roots/list +2 >roots/list +2 %s\n", prefix, method) + defer fmt.Fprintf(w, "%s <%s\n", prefix, method) + return h(ctx, sess, method, params) + } + } +} + var falseSchema = &jsonschema.Schema{Not: &jsonschema.Schema{}} diff --git a/mcp/server.go b/mcp/server.go index 86fbe17..9d35715 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -6,7 +6,6 @@ package mcp import ( "context" - "encoding/json" "fmt" "iter" "net/url" @@ -27,12 +26,12 @@ type Server struct { version string opts ServerOptions - mu sync.Mutex - prompts *featureSet[*ServerPrompt] - tools *featureSet[*ServerTool] - resources *featureSet[*ServerResource] - sessions []*ServerSession - methodHandler ServerMethodHandler + mu sync.Mutex + prompts *featureSet[*ServerPrompt] + tools *featureSet[*ServerTool] + resources *featureSet[*ServerResource] + sessions []*ServerSession + methodHandler_ MethodHandler[ServerSession] } // ServerOptions is used to configure behavior of the server. @@ -52,13 +51,13 @@ func NewServer(name, version string, opts *ServerOptions) *Server { opts = new(ServerOptions) } return &Server{ - name: name, - version: version, - opts: *opts, - prompts: newFeatureSet(func(p *ServerPrompt) string { return p.Prompt.Name }), - tools: newFeatureSet(func(t *ServerTool) string { return t.Tool.Name }), - resources: newFeatureSet(func(r *ServerResource) string { return r.Resource.URI }), - methodHandler: defaulMethodHandler, + name: name, + version: version, + opts: *opts, + prompts: newFeatureSet(func(p *ServerPrompt) string { return p.Prompt.Name }), + tools: newFeatureSet(func(t *ServerTool) string { return t.Tool.Name }), + resources: newFeatureSet(func(r *ServerResource) string { return r.Resource.URI }), + methodHandler_: defaultMethodHandler[ServerSession], } } @@ -321,7 +320,7 @@ type ServerSession struct { // Ping pings the client. func (ss *ServerSession) Ping(ctx context.Context, _ *PingParams) error { - return call(ctx, ss.conn, "ping", nil, nil) + return call(ctx, ss.conn, "ping", (*PingParams)(nil), nil) } // ListRoots lists the client roots. @@ -329,67 +328,20 @@ func (ss *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) return standardCall[ListRootsResult](ctx, ss.conn, "roots/list", params) } -// A ServerMethodHandler handles MCP messages from client to server. -// 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 ServerMethodHandler func(ctx context.Context, _ *ServerSession, method string, params any) (result any, err error) - // AddMiddleware wraps the server's current 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 server handler as +// For example, AddMiddleware(m1, m2, m3) augments the server method handler as // m1(m2(m3(handler))). -func (s *Server) AddMiddleware(middleware ...func(ServerMethodHandler) ServerMethodHandler) { +func (s *Server) AddMiddleware(middleware ...Middleware[ServerSession]) { s.mu.Lock() defer s.mu.Unlock() - for _, m := range slices.Backward(middleware) { - s.methodHandler = m(s.methodHandler) - } + addMiddleware(&s.methodHandler_, middleware) } -// defaulMethodHandler is the initial method handler installed on the server. -func defaulMethodHandler(ctx context.Context, ss *ServerSession, method string, params any) (any, error) { - info, ok := methodInfos[method] - assert(ok, "called with unknown method") - return info.handleMethod(ctx, ss, method, params) -} - -// methodInfo is information about invoking a method. -type methodInfo struct { - // unmarshal params from the wire into an XXXParams struct - unmarshalParams func(json.RawMessage) (any, error) - // run the code for the method - handleMethod ServerMethodHandler -} - -// The following definitions support converting from typed to untyped method handlers. -// Throughout, P is the type parameter for params, and R is the one for result. - -// A typedMethodHandler is like a MethodHandler, but with type information. -type typedMethodHandler[P, R any] func(context.Context, *ServerSession, P) (R, error) - -// newMethodInfo creates a methodInfo from a typedMethodHandler. -func newMethodInfo[P, R any](d typedMethodHandler[P, R]) methodInfo { - return methodInfo{ - unmarshalParams: func(m json.RawMessage) (any, error) { - var p P - if err := json.Unmarshal(m, &p); err != nil { - return nil, err - } - return p, nil - }, - handleMethod: func(ctx context.Context, ss *ServerSession, _ string, params any) (any, error) { - return d(ctx, ss, params.(P)) - }, - } -} - -// methodInfos maps from the RPC method name to methodInfos. -var methodInfos = map[string]methodInfo{ +// serverMethodInfos maps from the RPC method name to serverMethodInfos. +var serverMethodInfos = map[string]methodInfo[ServerSession]{ "initialize": newMethodInfo(sessionMethod((*ServerSession).initialize)), "ping": newMethodInfo(sessionMethod((*ServerSession).ping)), "prompts/list": newMethodInfo(serverMethod((*Server).listPrompts)), @@ -401,18 +353,18 @@ var methodInfos = map[string]methodInfo{ // TODO: notifications } -// 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[P, R] { - return func(ctx context.Context, ss *ServerSession, p P) (R, error) { - return f(ss.server, ctx, ss, p) - } +// *ServerSession implements the session interface. +// See toSession for why this interface seems to be necessary. +var _ session[ServerSession] = (*ServerSession)(nil) + +func (ss *ServerSession) methodInfos() map[string]methodInfo[ServerSession] { + return serverMethodInfos } -// sessionMethod is glue for creating a typedMethodHandler from a method on ServerSession. -func sessionMethod[P, R any](f func(*ServerSession, context.Context, P) (R, error)) typedMethodHandler[P, R] { - return func(ctx context.Context, ss *ServerSession, p P) (R, error) { - return f(ss, ctx, p) - } +func (ss *ServerSession) methodHandler() MethodHandler[ServerSession] { + ss.server.mu.Lock() + defer ss.server.mu.Unlock() + return ss.server.methodHandler_ } // handle invokes the method described by the given JSON RPC request. @@ -430,27 +382,10 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any return nil, fmt.Errorf("method %q is invalid during session initialization", req.Method) } } - - // TODO: embed the incoming request ID in the client context (or, more likely, + // TODO(rfindley): embed the incoming request ID in the client context (or, more likely, // a wrapper around it), so that we can correlate responses and notifications // to the handler; this is required for the new session-based transport. - info, ok := methodInfos[req.Method] - if !ok { - return nil, jsonrpc2.ErrNotHandled - } - params, err := info.unmarshalParams(req.Params) - if err != nil { - return nil, fmt.Errorf("ServerSession:handle %q: %w", req.Method, err) - } - ss.server.mu.Lock() - d := ss.server.methodHandler - ss.server.mu.Unlock() - // d might be user code, so ensure that it returns the right values for the jsonrpc2 protocol. - res, err := d(ctx, ss, req.Method, params) - if err != nil { - return nil, err - } - return res, nil + return handleRequest(ctx, req, ss) } func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParams) (*InitializeResult, error) { @@ -458,8 +393,8 @@ func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParam ss.initializeParams = params ss.mu.Unlock() - // Mark the connection as initialized when this method exits. TODO: - // Technically, the server should not be considered initialized until it has + // Mark the connection as initialized when this method exits. + // TODO: Technically, the server should not be considered initialized until it has // *responded*, but we don't have adequate visibility into the jsonrpc2 // connection to implement that easily. In any case, once we've initialized // here, we can handle requests. @@ -488,7 +423,7 @@ func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParam }, nil } -func (ss *ServerSession) ping(context.Context, struct{}) (struct{}, error) { +func (ss *ServerSession) ping(context.Context, *PingParams) (struct{}, error) { return struct{}{}, nil } diff --git a/mcp/shared.go b/mcp/shared.go new file mode 100644 index 0000000..643fb13 --- /dev/null +++ b/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) + } +}