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
6 changed files with 274 additions and 121 deletions
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