diff --git a/design/design.md b/design/design.md index 9670587..6fcd033 100644 --- a/design/design.md +++ b/design/design.md @@ -408,7 +408,7 @@ func (*ClientSession) ResourceTemplates(context.Context, *ListResourceTemplatesP ### Middleware -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. +We provide a mechanism to add MCP-level middleware on the both the client and server side. Receiving middleware runs after the request has been parsed but before any normal handling. It is analogous to traditional HTTP server middleware. Sending middleware runs after a call to a method but before the request is sent. It is an alternative to transport middleware that exposes MCP types instead of raw JSON-RPC 2.0 messages. It is useful for tracing and setting progress tokens, for example. ```go // A MethodHandler handles MCP messages. @@ -428,8 +428,10 @@ type Middleware[S Session] func(MethodHandler[S]) MethodHandler[S] // // 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]) +func (c *Client) AddSendingMiddleware(middleware ...Middleware[*ClientSession]) +func (c *Client) AddReceivingMiddleware(middleware ...Middleware[*ClientSession]) +func (s *Server) AddSendingMiddleware(middleware ...Middleware[*ServerSession]) +func (s *Server) AddReceivingMiddleware(middleware ...Middleware[*ServerSession]) ``` As an example, this code adds server-side logging: @@ -443,7 +445,7 @@ func withLogging(h mcp.MethodHandler[*ServerSession]) mcp.MethodHandler[*ServerS } } -server.AddMiddleware(withLogging) +server.AddReceivingMiddleware(withLogging) ``` **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. diff --git a/mcp/client.go b/mcp/client.go index 9d712df..258dc31 100644 --- a/mcp/client.go +++ b/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) { diff --git a/mcp/example_progress_test.go b/mcp/example_progress_test.go new file mode 100644 index 0000000..56a5772 --- /dev/null +++ b/mcp/example_progress_test.go @@ -0,0 +1,29 @@ +// 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. + +package mcp_test + +import ( + "context" + "sync/atomic" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +var nextProgressToken atomic.Int64 + +// This middleware function adds a progress token to every outgoing request +// from the client. +func Example_progressMiddleware() { + c := mcp.NewClient("test", "v1", nil) + c.AddSendingMiddleware(addProgressToken[*mcp.ClientSession]) + _ = c +} + +func addProgressToken[S mcp.Session](h mcp.MethodHandler[S]) mcp.MethodHandler[S] { + return func(ctx context.Context, s S, method string, params mcp.Params) (result mcp.Result, err error) { + params.GetMeta().ProgressToken = nextProgressToken.Add(1) + return h(ctx, s, method, params) + } +} diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 198538b..e4409bb 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -117,7 +117,7 @@ func TestEndToEnd(t *testing.T) { t.Run("prompts", func(t *testing.T) { res, err := cs.ListPrompts(ctx, nil) if err != nil { - t.Errorf("prompts/list failed: %v", err) + t.Fatalf("prompts/list failed: %v", err) } wantPrompts := []*Prompt{ { @@ -634,10 +634,12 @@ func TestMiddleware(t *testing.T) { // "1" is the outer middleware layer, called first; then "2" is called, and finally // the default dispatcher. - s.AddMiddleware(traceCalls[*ServerSession](&sbuf, "1"), traceCalls[*ServerSession](&sbuf, "2")) + s.AddSendingMiddleware(traceCalls[*ServerSession](&sbuf, "S1"), traceCalls[*ServerSession](&sbuf, "S2")) + s.AddReceivingMiddleware(traceCalls[*ServerSession](&sbuf, "R1"), traceCalls[*ServerSession](&sbuf, "R2")) c := NewClient("testClient", "v1.0.0", nil) - c.AddMiddleware(traceCalls[*ClientSession](&cbuf, "1"), traceCalls[*ClientSession](&cbuf, "2")) + c.AddSendingMiddleware(traceCalls[*ClientSession](&cbuf, "S1"), traceCalls[*ClientSession](&cbuf, "S2")) + c.AddReceivingMiddleware(traceCalls[*ClientSession](&cbuf, "R1"), traceCalls[*ClientSession](&cbuf, "R2")) cs, err := c.Connect(ctx, ct) if err != nil { @@ -646,34 +648,52 @@ func TestMiddleware(t *testing.T) { if _, err := cs.ListTools(ctx, nil); err != nil { t.Fatal(err) } - want := ` -1 >initialize -2 >initialize -2 notifications/initialized -2 >notifications/initialized -2 tools/list -2 >tools/list -2 roots/list -2 >roots/list -2 initialize +R2 >initialize +R2 notifications/initialized +R2 >notifications/initialized +R2 tools/list +R2 >tools/list +R2 roots/list +S2 >roots/list +S2 initialize +S2 >initialize +S2 notifications/initialized +S2 >notifications/initialized +S2 tools/list +S2 >tools/list +S2 roots/list +R2 >roots/list +R2