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 <rfindley@google.com> LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
This commit is contained in:
Jonathan Amsterdam
parent
8361df724d
commit
bac1b78b4a
5 changed files with 274 additions and 154 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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 <tools/list
|
||||
1 <tools/list
|
||||
`
|
||||
if diff := cmp.Diff(want, buf.String()); diff != "" {
|
||||
if diff := cmp.Diff(want, sbuf.String()); diff != "" {
|
||||
t.Errorf("mismatch (-want, +got):\n%s", diff)
|
||||
}
|
||||
|
||||
_, _ = ss.ListRoots(ctx, nil)
|
||||
|
||||
want = `
|
||||
1 >roots/list
|
||||
2 >roots/list
|
||||
2 <roots/list
|
||||
1 <roots/list
|
||||
`
|
||||
if diff := cmp.Diff(want, cbuf.String()); diff != "" {
|
||||
t.Errorf("mismatch (-want, +got):\n%s", diff)
|
||||
}
|
||||
}
|
||||
|
||||
// traceCalls creates a middleware function that prints the method before and after each call
|
||||
// with the given prefix.
|
||||
func traceCalls[S ClientSession | ServerSession](w io.Writer, prefix string) Middleware[S] {
|
||||
return func(h MethodHandler[S]) MethodHandler[S] {
|
||||
return func(ctx context.Context, sess *S, method string, params any) (any, error) {
|
||||
fmt.Fprintf(w, "%s >%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{}}
|
||||
|
|
|
|||
133
mcp/server.go
133
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
|
||||
}
|
||||
|
||||
|
|
|
|||
136
mcp/shared.go
Normal file
136
mcp/shared.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue