internal/mcp: middleware on the sending side
Implement sending middleware. Sending middleware is useful for adding tracing, progress tokens, and so on. Change-Id: I0acab2550f0212d549c4adbcf1ebf425a9486df6 Reviewed-on: https://go-review.googlesource.com/c/tools/+/676355 Reviewed-by: Sam Thanawalla <samthanawalla@google.com> 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
19d58e6955
commit
0d0c48a468
7 changed files with 274 additions and 121 deletions
|
|
@ -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.
|
||||
|
|
|
|||
100
mcp/client.go
100
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) {
|
||||
|
|
|
|||
29
mcp/example_progress_test.go
Normal file
29
mcp/example_progress_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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 <initialize
|
||||
1 <initialize
|
||||
1 >notifications/initialized
|
||||
2 >notifications/initialized
|
||||
2 <notifications/initialized
|
||||
1 <notifications/initialized
|
||||
1 >tools/list
|
||||
2 >tools/list
|
||||
2 <tools/list
|
||||
1 <tools/list
|
||||
`
|
||||
if diff := cmp.Diff(want, sbuf.String()); diff != "" {
|
||||
t.Errorf("mismatch (-want, +got):\n%s", diff)
|
||||
if _, err := ss.ListRoots(ctx, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, _ = ss.ListRoots(ctx, nil)
|
||||
|
||||
want = `
|
||||
1 >roots/list
|
||||
2 >roots/list
|
||||
2 <roots/list
|
||||
1 <roots/list
|
||||
wantServer := `
|
||||
R1 >initialize
|
||||
R2 >initialize
|
||||
R2 <initialize
|
||||
R1 <initialize
|
||||
R1 >notifications/initialized
|
||||
R2 >notifications/initialized
|
||||
R2 <notifications/initialized
|
||||
R1 <notifications/initialized
|
||||
R1 >tools/list
|
||||
R2 >tools/list
|
||||
R2 <tools/list
|
||||
R1 <tools/list
|
||||
S1 >roots/list
|
||||
S2 >roots/list
|
||||
S2 <roots/list
|
||||
S1 <roots/list
|
||||
`
|
||||
if diff := cmp.Diff(want, cbuf.String()); diff != "" {
|
||||
t.Errorf("mismatch (-want, +got):\n%s", diff)
|
||||
if diff := cmp.Diff(wantServer, sbuf.String()); diff != "" {
|
||||
t.Errorf("server mismatch (-want, +got):\n%s", diff)
|
||||
}
|
||||
|
||||
wantClient := `
|
||||
S1 >initialize
|
||||
S2 >initialize
|
||||
S2 <initialize
|
||||
S1 <initialize
|
||||
S1 >notifications/initialized
|
||||
S2 >notifications/initialized
|
||||
S2 <notifications/initialized
|
||||
S1 <notifications/initialized
|
||||
S1 >tools/list
|
||||
S2 >tools/list
|
||||
S2 <tools/list
|
||||
S1 <tools/list
|
||||
R1 >roots/list
|
||||
R2 >roots/list
|
||||
R2 <roots/list
|
||||
R1 <roots/list
|
||||
`
|
||||
if diff := cmp.Diff(wantClient, cbuf.String()); diff != "" {
|
||||
t.Errorf("client mismatch (-want, +got):\n%s", diff)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -32,12 +32,13 @@ type Server struct {
|
|||
version string
|
||||
opts ServerOptions
|
||||
|
||||
mu sync.Mutex
|
||||
prompts *featureSet[*ServerPrompt]
|
||||
tools *featureSet[*ServerTool]
|
||||
resources *featureSet[*ServerResource]
|
||||
sessions []*ServerSession
|
||||
methodHandler_ MethodHandler[*ServerSession]
|
||||
mu sync.Mutex
|
||||
prompts *featureSet[*ServerPrompt]
|
||||
tools *featureSet[*ServerTool]
|
||||
resources *featureSet[*ServerResource]
|
||||
sessions []*ServerSession
|
||||
sendingMethodHandler_ MethodHandler[*ServerSession]
|
||||
receivingMethodHandler_ MethodHandler[*ServerSession]
|
||||
}
|
||||
|
||||
// ServerOptions is used to configure behavior of the server.
|
||||
|
|
@ -71,13 +72,14 @@ func NewServer(name, version string, opts *ServerOptions) *Server {
|
|||
opts.PageSize = DefaultPageSize
|
||||
}
|
||||
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_: defaultMethodHandler[*ServerSession],
|
||||
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 }),
|
||||
sendingMethodHandler_: defaultSendingMethodHandler[*ServerSession],
|
||||
receivingMethodHandler_: defaultReceivingMethodHandler[*ServerSession],
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -401,18 +403,19 @@ type ServerSession struct {
|
|||
}
|
||||
|
||||
// Ping pings the client.
|
||||
func (ss *ServerSession) Ping(ctx context.Context, _ *PingParams) error {
|
||||
return call(ctx, ss.conn, methodPing, (*PingParams)(nil), nil)
|
||||
func (ss *ServerSession) Ping(ctx context.Context, params *PingParams) error {
|
||||
_, err := handleSend[*emptyResult](ctx, ss, methodPing, params)
|
||||
return err
|
||||
}
|
||||
|
||||
// ListRoots lists the client roots.
|
||||
func (ss *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) {
|
||||
return standardCall[ListRootsResult](ctx, ss.conn, methodListRoots, params)
|
||||
return handleSend[*ListRootsResult](ctx, ss, methodListRoots, params)
|
||||
}
|
||||
|
||||
// CreateMessage sends a sampling request to the client.
|
||||
func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessageParams) (*CreateMessageResult, error) {
|
||||
return standardCall[CreateMessageResult](ctx, ss.conn, methodCreateMessage, params)
|
||||
return handleSend[*CreateMessageResult](ctx, ss, methodCreateMessage, params)
|
||||
}
|
||||
|
||||
// LoggingMessage sends a logging message to the client.
|
||||
|
|
@ -431,19 +434,37 @@ func (ss *ServerSession) LoggingMessage(ctx context.Context, params *LoggingMess
|
|||
if compareLevels(params.Level, logLevel) < 0 {
|
||||
return nil
|
||||
}
|
||||
return ss.conn.Notify(ctx, notificationLoggingMessage, params)
|
||||
return handleNotify(ctx, ss, notificationLoggingMessage, params)
|
||||
}
|
||||
|
||||
// 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.
|
||||
// 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 server method handler as
|
||||
// For example, AddSendingMiddleware(m1, m2, m3) augments the method handler as
|
||||
// m1(m2(m3(handler))).
|
||||
func (s *Server) AddMiddleware(middleware ...Middleware[*ServerSession]) {
|
||||
//
|
||||
// Sending middleware is called when a request is sent. It is useful for tasks
|
||||
// such as tracing, metrics, and adding progress tokens.
|
||||
func (s *Server) AddSendingMiddleware(middleware ...Middleware[*ServerSession]) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
addMiddleware(&s.methodHandler_, middleware)
|
||||
addMiddleware(&s.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 (s *Server) AddReceivingMiddleware(middleware ...Middleware[*ServerSession]) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
addMiddleware(&s.receivingMethodHandler_, middleware)
|
||||
}
|
||||
|
||||
// serverMethodInfos maps from the RPC method name to serverMethodInfos.
|
||||
|
|
@ -461,14 +482,20 @@ var serverMethodInfos = map[string]methodInfo{
|
|||
notificationRootsListChanged: newMethodInfo(serverMethod((*Server).callRootsListChangedHandler)),
|
||||
}
|
||||
|
||||
func (ss *ServerSession) methodInfos() map[string]methodInfo {
|
||||
return serverMethodInfos
|
||||
}
|
||||
func (ss *ServerSession) sendingMethodInfos() map[string]methodInfo { return clientMethodInfos }
|
||||
|
||||
func (ss *ServerSession) methodHandler() methodHandler {
|
||||
func (ss *ServerSession) receivingMethodInfos() map[string]methodInfo { return serverMethodInfos }
|
||||
|
||||
func (ss *ServerSession) sendingMethodHandler() methodHandler {
|
||||
ss.server.mu.Lock()
|
||||
defer ss.server.mu.Unlock()
|
||||
return ss.server.methodHandler_
|
||||
return ss.server.sendingMethodHandler_
|
||||
}
|
||||
|
||||
func (ss *ServerSession) receivingMethodHandler() methodHandler {
|
||||
ss.server.mu.Lock()
|
||||
defer ss.server.mu.Unlock()
|
||||
return ss.server.receivingMethodHandler_
|
||||
}
|
||||
|
||||
// getConn implements [session.getConn].
|
||||
|
|
@ -492,7 +519,7 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any
|
|||
// 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.
|
||||
return handleRequest(ctx, req, ss)
|
||||
return handleReceive(ctx, ss, req)
|
||||
}
|
||||
|
||||
func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParams) (*InitializeResult, error) {
|
||||
|
|
@ -534,15 +561,15 @@ func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParam
|
|||
}, nil
|
||||
}
|
||||
|
||||
func (ss *ServerSession) ping(context.Context, *PingParams) (Result, error) {
|
||||
return emptyResult{}, nil
|
||||
func (ss *ServerSession) ping(context.Context, *PingParams) (*emptyResult, error) {
|
||||
return &emptyResult{}, nil
|
||||
}
|
||||
|
||||
func (ss *ServerSession) setLevel(_ context.Context, params *SetLevelParams) (Result, error) {
|
||||
func (ss *ServerSession) setLevel(_ context.Context, params *SetLevelParams) (*emptyResult, error) {
|
||||
ss.mu.Lock()
|
||||
defer ss.mu.Unlock()
|
||||
ss.logLevel = params.Level
|
||||
return emptyResult{}, nil
|
||||
return &emptyResult{}, nil
|
||||
}
|
||||
|
||||
// Close performs a graceful shutdown of the connection, preventing new
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ import (
|
|||
"log"
|
||||
"reflect"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
||||
|
|
@ -36,8 +37,10 @@ type methodHandler any // MethodHandler[*ClientSession] | MethodHandler[*ServerS
|
|||
// A Session is either a ClientSession or a ServerSession.
|
||||
type Session interface {
|
||||
*ClientSession | *ServerSession
|
||||
methodHandler() methodHandler
|
||||
methodInfos() map[string]methodInfo
|
||||
sendingMethodInfos() map[string]methodInfo
|
||||
receivingMethodInfos() map[string]methodInfo
|
||||
sendingMethodHandler() methodHandler
|
||||
receivingMethodHandler() methodHandler
|
||||
getConn() *jsonrpc2.Connection
|
||||
}
|
||||
|
||||
|
|
@ -51,9 +54,45 @@ func addMiddleware[S Session](handlerp *MethodHandler[S], middleware []Middlewar
|
|||
}
|
||||
}
|
||||
|
||||
// defaultMethodHandler is the initial MethodHandler for servers and clients, before being wrapped by middleware.
|
||||
func defaultMethodHandler[S Session](ctx context.Context, session S, method string, params Params) (Result, error) {
|
||||
info, ok := session.methodInfos()[method]
|
||||
func defaultSendingMethodHandler[S Session](ctx context.Context, session S, method string, params Params) (Result, error) {
|
||||
info, ok := session.sendingMethodInfos()[method]
|
||||
if !ok {
|
||||
// This can be called from user code, with an arbitrary value for method.
|
||||
return nil, jsonrpc2.ErrNotHandled
|
||||
}
|
||||
// Notifications don't have results.
|
||||
if strings.HasPrefix(method, "notifications/") {
|
||||
return nil, session.getConn().Notify(ctx, method, params)
|
||||
}
|
||||
// Create the result to unmarshal into.
|
||||
// The concrete type of the result is the return type of the receiving function.
|
||||
res := info.newResult()
|
||||
if err := call(ctx, session.getConn(), method, params, res); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func handleNotify[S Session](ctx context.Context, session S, method string, params Params) error {
|
||||
mh := session.sendingMethodHandler().(MethodHandler[S])
|
||||
_, err := mh(ctx, session, method, params)
|
||||
return err
|
||||
}
|
||||
|
||||
func handleSend[R Result, S Session](ctx context.Context, s S, method string, params Params) (R, error) {
|
||||
mh := s.sendingMethodHandler().(MethodHandler[S])
|
||||
// mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol.
|
||||
res, err := mh(ctx, s, method, params)
|
||||
if err != nil {
|
||||
var z R
|
||||
return z, err
|
||||
}
|
||||
return res.(R), nil
|
||||
}
|
||||
|
||||
// defaultReceivingMethodHandler is the initial MethodHandler for servers and clients, before being wrapped by middleware.
|
||||
func defaultReceivingMethodHandler[S Session](ctx context.Context, session S, method string, params Params) (Result, error) {
|
||||
info, ok := session.receivingMethodInfos()[method]
|
||||
if !ok {
|
||||
// This can be called from user code, with an arbitrary value for method.
|
||||
return nil, jsonrpc2.ErrNotHandled
|
||||
|
|
@ -61,8 +100,8 @@ func defaultMethodHandler[S Session](ctx context.Context, session S, method stri
|
|||
return info.handleMethod.(MethodHandler[S])(ctx, session, method, params)
|
||||
}
|
||||
|
||||
func handleRequest[S Session](ctx context.Context, req *jsonrpc2.Request, session S) (any, error) {
|
||||
info, ok := session.methodInfos()[req.Method]
|
||||
func handleReceive[S Session](ctx context.Context, session S, req *jsonrpc2.Request) (Result, error) {
|
||||
info, ok := session.receivingMethodInfos()[req.Method]
|
||||
if !ok {
|
||||
return nil, jsonrpc2.ErrNotHandled
|
||||
}
|
||||
|
|
@ -71,8 +110,8 @@ func handleRequest[S Session](ctx context.Context, req *jsonrpc2.Request, sessio
|
|||
return nil, fmt.Errorf("handleRequest %q: %w", req.Method, err)
|
||||
}
|
||||
|
||||
mh := session.receivingMethodHandler().(MethodHandler[S])
|
||||
// mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol.
|
||||
mh := session.methodHandler().(MethodHandler[S])
|
||||
res, err := mh(ctx, session, req.Method, params)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
|
@ -80,12 +119,17 @@ func handleRequest[S Session](ctx context.Context, req *jsonrpc2.Request, sessio
|
|||
return res, nil
|
||||
}
|
||||
|
||||
// methodInfo is information about invoking a method.
|
||||
// methodInfo is information about sending and receiving a method.
|
||||
type methodInfo struct {
|
||||
// unmarshal params from the wire into an XXXParams struct
|
||||
// Unmarshal params from the wire into a Params struct.
|
||||
// Used on the receive side.
|
||||
unmarshalParams func(json.RawMessage) (Params, error)
|
||||
// run the code for the method
|
||||
// Run the code when a call to the method is received.
|
||||
// Used on the receive side.
|
||||
handleMethod methodHandler
|
||||
// Create a pointer to a Result struct.
|
||||
// Used on the send side.
|
||||
newResult func() Result
|
||||
}
|
||||
|
||||
// The following definitions support converting from typed to untyped method handlers.
|
||||
|
|
@ -112,6 +156,11 @@ func newMethodInfo[S Session, P Params, R Result](d typedMethodHandler[S, P, R])
|
|||
handleMethod: MethodHandler[S](func(ctx context.Context, session S, _ string, params Params) (Result, error) {
|
||||
return d(ctx, session, params.(P))
|
||||
}),
|
||||
// newResult is used on the send side, to construct the value to unmarshal the result into.
|
||||
// R is a pointer to a result struct. There is no way to "unpointer" it without reflection.
|
||||
// TODO(jba): explore generic approaches to this, perhaps by treating R in
|
||||
// the signature as the unpointered type.
|
||||
newResult: func() Result { return reflect.New(reflect.TypeFor[R]().Elem()).Interface().(R) },
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -170,21 +219,13 @@ func notifySessions[S Session](sessions []S, method string, params Params) {
|
|||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
for _, s := range sessions {
|
||||
if err := s.getConn().Notify(ctx, method, params); err != nil {
|
||||
if err := handleNotify(ctx, s, method, params); err != nil {
|
||||
// TODO(jba): surface this error better
|
||||
log.Printf("calling %s: %v", method, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func standardCall[TRes, TParams any](ctx context.Context, conn *jsonrpc2.Connection, method string, params TParams) (*TRes, error) {
|
||||
var result TRes
|
||||
if err := call(ctx, conn, method, params, &result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
type Meta struct {
|
||||
Data map[string]any `json:",omitempty"`
|
||||
// For params, the progress token can be nil, a string or an integer.
|
||||
|
|
@ -229,4 +270,4 @@ type Result interface {
|
|||
// Those methods cannot return nil, because jsonrpc2 cannot handle nils.
|
||||
type emptyResult struct{}
|
||||
|
||||
func (emptyResult) GetMeta() *Meta { panic("should never be called") }
|
||||
func (*emptyResult) GetMeta() *Meta { panic("should never be called") }
|
||||
|
|
|
|||
|
|
@ -141,7 +141,7 @@ func (c *canceller) Preempt(ctx context.Context, req *jsonrpc2.Request) (result
|
|||
|
||||
// call executes and awaits a jsonrpc2 call on the given connection,
|
||||
// translating errors into the mcp domain.
|
||||
func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params, result any) error {
|
||||
func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params Params, result Result) error {
|
||||
// TODO: the "%w"s in this function effectively make jsonrpc2.WireError part of the API.
|
||||
// Consider alternatives.
|
||||
call := conn.Call(ctx, method, params)
|
||||
|
|
@ -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), notificationCancelled, &CancelledParams{
|
||||
Reason: ctx.Err().Error(),
|
||||
RequestID: call.ID().Raw(),
|
||||
})
|
||||
|
|
|
|||
Loading…
Reference in a new issue