modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

internal/mcp: dispatcher middleware

Browse source 
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 <rfindley@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
This commit is contained in:
Jonathan Amsterdam 2025-05-15 16:35:46 -04:00
commit 1323fb5216
7 changed files with 249 additions and 125 deletions
Download patch file Download diff file Expand all files Collapse all files

14
design/design.md
View file

@ -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

9
mcp/client.go
View file

@ -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

8
mcp/generate.go
View file

@ -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: "-",

61
mcp/mcp_test.go
View file

@ -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 <initialize
1 <initialize
1 >tools/list
2 >tools/list
2 <tools/list
1 <tools/list
`
if diff := cmp.Diff(want, buf.String()); diff != "" {
t.Errorf("mismatch (-want, +got):\n%s", diff)
}
}
var falseSchema = &jsonschema.Schema{Not: &jsonschema.Schema{}}

90
mcp/protocol.go
View file

@ -54,6 +54,17 @@ type CallToolResult struct {
IsError bool `json:"isError,omitempty"`
}
type CancelledParams struct {
// An optional string describing the reason for the cancellation. This MAY be
// logged or presented to the user.
Reason string `json:"reason,omitempty"`
// The ID of the request to cancel.
//
// This MUST correspond to the ID of a request previously issued in the same
// direction.
RequestID any `json:"requestId"`
}
// Capabilities a client may support. Known capabilities are defined here, in
// this schema, but this is not a closed set: any client can define its own,
// additional capabilities.
@ -88,6 +99,40 @@ type GetPromptResult struct {
Messages []*PromptMessage `json:"messages"`
}
type InitializeParams struct {
Capabilities *ClientCapabilities `json:"capabilities"`
ClientInfo *implementation `json:"clientInfo"`
// The latest version of the Model Context Protocol that the client supports.
// The client MAY decide to support older versions as well.
ProtocolVersion string `json:"protocolVersion"`
}
// After receiving an initialize request from the client, the server sends this
// response.
type InitializeResult struct {
// This result property is reserved by the protocol to allow clients and servers
// to attach additional metadata to their responses.
Meta map[string]json.RawMessage `json:"_meta,omitempty"`
Capabilities *serverCapabilities `json:"capabilities"`
// Instructions describing how to use the server and its features.
//
// This can be used by clients to improve the LLM's understanding of available
// tools, resources, etc. It can be thought of like a "hint" to the model. For
// example, this information MAY be added to the system prompt.
Instructions string `json:"instructions,omitempty"`
// The version of the Model Context Protocol that the server wants to use. This
// may not match the version that the client requested. If the client cannot
// support this version, it MUST disconnect.
ProtocolVersion string `json:"protocolVersion"`
ServerInfo *implementation `json:"serverInfo"`
}
type InitializedParams struct {
// This parameter name is reserved by MCP to allow clients and servers to attach
// additional metadata to their notifications.
Meta map[string]json.RawMessage `json:"_meta,omitempty"`
}
type ListPromptsParams struct {
// An opaque token representing the current pagination position. If provided,
// the server should return results starting after this cursor.
@ -304,57 +349,12 @@ type ToolAnnotations struct {
Title string `json:"title,omitempty"`
}
type cancelledParams struct {
// An optional string describing the reason for the cancellation. This MAY be
// logged or presented to the user.
Reason string `json:"reason,omitempty"`
// The ID of the request to cancel.
//
// This MUST correspond to the ID of a request previously issued in the same
// direction.
RequestID any `json:"requestId"`
}
// Describes the name and version of an MCP implementation.
type implementation struct {
Name string `json:"name"`
Version string `json:"version"`
}
type initializeParams struct {
Capabilities *ClientCapabilities `json:"capabilities"`
ClientInfo *implementation `json:"clientInfo"`
// The latest version of the Model Context Protocol that the client supports.
// The client MAY decide to support older versions as well.
ProtocolVersion string `json:"protocolVersion"`
}
// After receiving an initialize request from the client, the server sends this
// response.
type initializeResult struct {
// This result property is reserved by the protocol to allow clients and servers
// to attach additional metadata to their responses.
Meta map[string]json.RawMessage `json:"_meta,omitempty"`
Capabilities *serverCapabilities `json:"capabilities"`
// Instructions describing how to use the server and its features.
//
// This can be used by clients to improve the LLM's understanding of available
// tools, resources, etc. It can be thought of like a "hint" to the model. For
// example, this information MAY be added to the system prompt.
Instructions string `json:"instructions,omitempty"`
// The version of the Model Context Protocol that the server wants to use. This
// may not match the version that the client requested. If the client cannot
// support this version, it MUST disconnect.
ProtocolVersion string `json:"protocolVersion"`
ServerInfo *implementation `json:"serverInfo"`
}
type initializedParams struct {
// This parameter name is reserved by MCP to allow clients and servers to attach
// additional metadata to their notifications.
Meta map[string]json.RawMessage `json:"_meta,omitempty"`
}
// Present if the server offers any prompt templates.
type promptCapabilities struct {
// Whether this server supports notifications for changes to the prompt list.

188
mcp/server.go
View file

@ -26,11 +26,12 @@ type Server struct {
version string
opts ServerOptions
mu sync.Mutex
prompts *featureSet[*ServerPrompt]
tools *featureSet[*ServerTool]
resources *featureSet[*ServerResource]
sessions []*ServerSession
mu sync.Mutex
prompts *featureSet[*ServerPrompt]
tools *featureSet[*ServerTool]
resources *featureSet[*ServerResource]
sessions []*ServerSession
methodHandler ServerMethodHandler
}
// ServerOptions is used to configure behavior of the server.
@ -50,12 +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 }),
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,
}
}
@ -297,24 +299,111 @@ type ServerSession struct {
conn *jsonrpc2.Connection
mu sync.Mutex
initializeParams *initializeParams
initializeParams *InitializeParams
initialized bool
}
// Ping makes an MCP "ping" request to the client.
// Ping pings the client.
func (ss *ServerSession) Ping(ctx context.Context, _ *PingParams) error {
return call(ctx, ss.conn, "ping", nil, nil)
}
// ListRoots lists the client roots.
func (ss *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) {
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
// m1(m2(m3(handler))).
func (s *Server) AddMiddleware(middleware ...func(ServerMethodHandler) ServerMethodHandler) {
s.mu.Lock()
defer s.mu.Unlock()
for _, m := range slices.Backward(middleware) {
s.methodHandler = m(s.methodHandler)
}
}
// 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{
"initialize": newMethodInfo(sessionMethod((*ServerSession).initialize)),
"ping": newMethodInfo(sessionMethod((*ServerSession).ping)),
"prompts/list": newMethodInfo(serverMethod((*Server).listPrompts)),
"prompts/get": newMethodInfo(serverMethod((*Server).getPrompt)),
"tools/list": newMethodInfo(serverMethod((*Server).listTools)),
"tools/call": newMethodInfo(serverMethod((*Server).callTool)),
"resources/list": newMethodInfo(serverMethod((*Server).listResources)),
"resources/read": newMethodInfo(serverMethod((*Server).readResource)),
// 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)
}
}
// 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)
}
}
// handle invokes the method described by the given JSON RPC request.
func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
ss.mu.Lock()
initialized := ss.initialized
ss.mu.Unlock()
// From the spec:
// "The client SHOULD NOT send requests other than pings before the server
// has responded to the initialize request."
@ -329,39 +418,26 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any
// TODO: 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.
switch req.Method {
case "initialize":
return dispatch(ctx, ss, req, ss.initialize)
case "ping":
// The spec says that 'ping' expects an empty object result.
return struct{}{}, nil
case "prompts/list":
return dispatch(ctx, ss, req, ss.server.listPrompts)
case "prompts/get":
return dispatch(ctx, ss, req, ss.server.getPrompt)
case "tools/list":
return dispatch(ctx, ss, req, ss.server.listTools)
case "tools/call":
return dispatch(ctx, ss, req, ss.server.callTool)
case "resources/list":
return dispatch(ctx, ss, req, ss.server.listResources)
case "resources/read":
return dispatch(ctx, ss, req, ss.server.readResource)
case "notifications/initialized":
info, ok := methodInfos[req.Method]
if !ok {
return nil, jsonrpc2.ErrNotHandled
}
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
}
func (ss *ServerSession) initialize(ctx context.Context, _ *ServerSession, params *initializeParams) (*initializeResult, error) {
func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParams) (*InitializeResult, error) {
ss.mu.Lock()
ss.initializeParams = params
ss.mu.Unlock()
@ -377,7 +453,7 @@ func (ss *ServerSession) initialize(ctx context.Context, _ *ServerSession, param
ss.mu.Unlock()
}()
return &initializeResult{
return &InitializeResult{
// TODO(rfindley): support multiple protocol versions.
ProtocolVersion: "2024-11-05",
Capabilities: &serverCapabilities{
@ -396,6 +472,10 @@ func (ss *ServerSession) initialize(ctx context.Context, _ *ServerSession, param
}, nil
}
func (ss *ServerSession) ping(context.Context, struct{}) (struct{}, error) {
return struct{}{}, nil
}
// Close performs a graceful shutdown of the connection, preventing new
// requests from being handled, and waiting for ongoing requests to return.
// Close then terminates the connection.
@ -407,21 +487,3 @@ func (ss *ServerSession) Close() error {
func (ss *ServerSession) Wait() error {
return ss.conn.Wait()
}
// dispatch turns a strongly type request handler into a jsonrpc2 handler.
//
// Importantly, it returns nil if the handler returned an error, which is a
// requirement of the jsonrpc2 package.
func dispatch[TParams, TResult any](ctx context.Context, conn *ServerSession, req *jsonrpc2.Request, f func(context.Context, *ServerSession, TParams) (TResult, error)) (any, error) {
var params TParams
if err := json.Unmarshal(req.Params, &params); err != nil {
return nil, err
}
// Important: avoid returning a typed nil, as it can't be handled by the
// jsonrpc2 package.
res, err := f(ctx, conn, params)
if err != nil {
return nil, err
}
return res, nil
}

4
mcp/transport.go
View file

@ -126,7 +126,7 @@ type canceller struct {
// Preempt implements jsonrpc2.Preempter.
func (c *canceller) Preempt(ctx context.Context, req *jsonrpc2.Request) (result any, err error) {
if req.Method == "notifications/cancelled" {
var params cancelledParams
var params CancelledParams
if err := json.Unmarshal(req.Params, &params); err != nil {
return nil, err
}
@ -151,7 +151,7 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params,
return fmt.Errorf("calling %q: %w", method, ErrConnectionClosed)
case ctx.Err() != nil:
// Notify the peer of cancellation.
err := conn.Notify(xcontext.Detach(ctx), "notifications/cancelled", &cancelledParams{
err := conn.Notify(xcontext.Detach(ctx), "notifications/cancelled", &CancelledParams{
Reason: ctx.Err().Error(),
RequestID: call.ID().Raw(),
})