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