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
1 changed files with 274 additions and 154 deletions
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
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue