From 1323fb52169b177c2e04d0b7ecf6857a51656ac0 Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Thu, 15 May 2025 16:35:46 -0400 Subject: [PATCH] internal/mcp: dispatcher middleware Implement the dispatcher middleware framework, allowing users to wrap method dispatching with their own code. The Dispatcher functions have an untyped params arg, but expect its value to be a typed, unmarshaled XXXParams struct. That required splitting the elegant dispatch[P, R] function into two parts, one to do the unmarshaling and one to call the method. The nicest way I could find to do that was to map from method names to a pair of functions, one to unmarshal and one to call. We could also write it as two switches on method name, which would require less machinery but would duplicate the list of method names, already a tad fragile because they are literal strings. I also exported the initialize params and result, so user Dispatchers could access them. All Params and Result types will have to be exported for that reason. Change-Id: Ie2c95cc17b4ef3d3d095dccdadb294e5f1590796 Reviewed-on: https://go-review.googlesource.com/c/tools/+/673515 Reviewed-by: Robert Findley LUCI-TryBot-Result: Go LUCI --- design/design.md | 14 ++-- mcp/client.go | 9 +-- mcp/generate.go | 8 +- mcp/mcp_test.go | 61 +++++++++++++++ mcp/protocol.go | 90 +++++++++++------------ mcp/server.go | 188 +++++++++++++++++++++++++++++++---------------- mcp/transport.go | 4 +- 7 files changed, 249 insertions(+), 125 deletions(-) diff --git a/design/design.md b/design/design.md index 713bcc8..5a2b103 100644 --- a/design/design.md +++ b/design/design.md @@ -409,20 +409,20 @@ func (*ClientSession) ResourceTemplates(context.Context, *ListResourceTemplatesP 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. ```go -// A Dispatcher dispatches an MCP message to the appropriate handler. +// 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 Dispatcher func(ctx context.Context, s *ServerSession, method string, params any) (result any, err error) +type ServerMethodHandler func(ctx context.Context, s *ServerSession, method string, params any) (result any, err error) -// AddDispatchers calls each function from right to left on the previous result, beginning +// 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) AddDispatchers(middleware ...func(Dispatcher) Dispatcher)) +func (*Server) AddMiddleware(middleware ...func(ServerMethodHandler) ServerMethodHandler) ``` As an example, this code adds server-side logging: ```go -func withLogging(h mcp.Dispatcher) mcp.Dispatcher { +func withLogging(h mcp.ServerMethodHandler) mcp.ServerMethodHandler{ 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) }() @@ -430,9 +430,11 @@ func withLogging(h mcp.Dispatcher) mcp.Dispatcher { } } -server.AddDispatchers(withLogging) +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 a4ea94f..8f2ee52 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -44,8 +44,7 @@ func NewClient(name, version string, opts *ClientOptions) *Client { } // ClientOptions configures the behavior of the client. -type ClientOptions struct { -} +type ClientOptions struct{} // bind implements the binder[*ClientSession] interface, so that Clients can // be connected using [connect]. @@ -82,14 +81,14 @@ func (c *Client) Connect(ctx context.Context, t Transport) (cs *ClientSession, e if err != nil { return nil, err } - params := &initializeParams{ + params := &InitializeParams{ ClientInfo: &implementation{Name: c.name, Version: c.version}, } if err := call(ctx, cs.conn, "initialize", params, &cs.initializeResult); err != nil { _ = cs.Close() return nil, err } - if err := cs.conn.Notify(ctx, "notifications/initialized", &initializedParams{}); err != nil { + if err := cs.conn.Notify(ctx, "notifications/initialized", &InitializedParams{}); err != nil { _ = cs.Close() return nil, err } @@ -105,7 +104,7 @@ func (c *Client) Connect(ctx context.Context, t Transport) (cs *ClientSession, e type ClientSession struct { conn *jsonrpc2.Connection client *Client - initializeResult *initializeResult + initializeResult *InitializeResult } // Close performs a graceful close of the connection, preventing new requests diff --git a/mcp/generate.go b/mcp/generate.go index cfffef8..feea820 100644 --- a/mcp/generate.go +++ b/mcp/generate.go @@ -66,7 +66,7 @@ var declarations = config{ "CallToolResult": {}, "CancelledNotification": { Name: "-", - Fields: config{"Params": {Name: "cancelledParams"}}, + Fields: config{"Params": {Name: "CancelledParams"}}, }, "ClientCapabilities": {}, "GetPromptRequest": { @@ -77,12 +77,12 @@ var declarations = config{ "Implementation": {Name: "implementation"}, "InitializeRequest": { Name: "-", - Fields: config{"Params": {Name: "initializeParams"}}, + Fields: config{"Params": {Name: "InitializeParams"}}, }, - "InitializeResult": {Name: "initializeResult"}, + "InitializeResult": {Name: "InitializeResult"}, "InitializedNotification": { Name: "-", - Fields: config{"Params": {Name: "initializedParams"}}, + Fields: config{"Params": {Name: "InitializedParams"}}, }, "ListPromptsRequest": { Name: "-", diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 1e8447c..7812350 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -5,6 +5,7 @@ package mcp import ( + "bytes" "context" "errors" "fmt" @@ -397,4 +398,64 @@ func TestCancellation(t *testing.T) { } } +func TestAddMiddleware(t *testing.T) { + ctx := context.Background() + ct, st := NewInMemoryTransports() + s := NewServer("testServer", "v1.0.0", nil) + ss, err := s.Connect(ctx, st) + if err != nil { + t.Fatal(err) + } + // Wait for the server to exit after the client closes its connection. + var clientWG sync.WaitGroup + clientWG.Add(1) + go func() { + if err := ss.Wait(); err != nil { + t.Errorf("server failed: %v", err) + } + 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) + } + } + } + + // "1" is the outer middleware layer, called first; then "2" is called, and finally + // the default dispatcher. + s.AddMiddleware(traceCalls("1"), traceCalls("2")) + + c := NewClient("testClient", "v1.0.0", nil) + cs, err := c.Connect(ctx, ct) + if err != nil { + t.Fatal(err) + } + if _, err := cs.ListTools(ctx, nil); err != nil { + t.Fatal(err) + } + want := ` +1 >initialize +2 >initialize +2 tools/list +2 >tools/list +2