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

mcp: introduce Requests (#267)

Browse source 
Handlers and client/server methods take a single request argument
combining the session and params. This slightly simplifies the signature
for handlers, since there is no longer an often-ignored session
argument.

But more importantly, it opens the door to adding more information in
requests, such as auth info and HTTP request headers.

For #243.
This commit is contained in:
Jonathan Amsterdam 2025-08-12 17:49:50 -04:00 • committed by GitHub
commit 6e03217c83
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 517 additions and 380 deletions
Download patch file Download diff file Expand all files Collapse all files

127
mcp/client.go
View file

@ -24,8 +24,8 @@ type Client struct {
mu sync.Mutex
roots *featureSet[*Root]
sessions []*ClientSession
sendingMethodHandler_ MethodHandler[*ClientSession]
receivingMethodHandler_ MethodHandler[*ClientSession]
sendingMethodHandler_ MethodHandler
receivingMethodHandler_ MethodHandler
}
// NewClient creates a new [Client].
@ -55,14 +55,14 @@ func NewClient(impl *Implementation, opts *ClientOptions) *Client {
type ClientOptions struct {
// Handler for sampling.
// Called when a server calls CreateMessage.
CreateMessageHandler func(context.Context, *ClientSession, *CreateMessageParams) (*CreateMessageResult, error)
CreateMessageHandler func(context.Context, *ClientRequest[*CreateMessageParams]) (*CreateMessageResult, error)
// Handlers for notifications from the server.
ToolListChangedHandler func(context.Context, *ClientSession, *ToolListChangedParams)
PromptListChangedHandler func(context.Context, *ClientSession, *PromptListChangedParams)
ResourceListChangedHandler func(context.Context, *ClientSession, *ResourceListChangedParams)
ResourceUpdatedHandler func(context.Context, *ClientSession, *ResourceUpdatedNotificationParams)
LoggingMessageHandler func(context.Context, *ClientSession, *LoggingMessageParams)
ProgressNotificationHandler func(context.Context, *ClientSession, *ProgressNotificationParams)
ToolListChangedHandler func(context.Context, *ClientRequest[*ToolListChangedParams])
PromptListChangedHandler func(context.Context, *ClientRequest[*PromptListChangedParams])
ResourceListChangedHandler func(context.Context, *ClientRequest[*ResourceListChangedParams])
ResourceUpdatedHandler func(context.Context, *ClientRequest[*ResourceUpdatedNotificationParams])
LoggingMessageHandler func(context.Context, *ClientRequest[*LoggingMessageParams])
ProgressNotificationHandler func(context.Context, *ClientRequest[*ProgressNotificationParams])
// If non-zero, defines an interval for regular "ping" requests.
// If the peer fails to respond to pings originating from the keepalive check,
// the session is automatically closed.
@ -130,7 +130,8 @@ func (c *Client) Connect(ctx context.Context, t Transport, _ *ClientSessionOptio
ClientInfo: c.impl,
Capabilities: caps,
}
res, err := handleSend[*InitializeResult](ctx, cs, methodInitialize, params)
req := &ClientRequest[*InitializeParams]{Session: cs, Params: params}
res, err := handleSend[*InitializeResult](ctx, methodInitialize, req)
if err != nil {
_ = cs.Close()
return nil, err
@ -142,7 +143,8 @@ func (c *Client) Connect(ctx context.Context, t Transport, _ *ClientSessionOptio
if hc, ok := cs.mcpConn.(clientConnection); ok {
hc.sessionUpdated(cs.state)
}
if err := handleNotify(ctx, cs, notificationInitialized, &InitializedParams{}); err != nil {
req2 := &ClientRequest[*InitializedParams]{Session: cs, Params: &InitializedParams{}}
if err := handleNotify(ctx, notificationInitialized, req2); err != nil {
_ = cs.Close()
return nil, err
}
@ -216,7 +218,7 @@ func (c *Client) AddRoots(roots ...*Root) {
if len(roots) == 0 {
return
}
c.changeAndNotify(notificationRootsListChanged, &RootsListChangedParams{},
changeAndNotify(c, notificationRootsListChanged, &RootsListChangedParams{},
func() bool { c.roots.add(roots...); return true })
}
@ -225,14 +227,14 @@ func (c *Client) AddRoots(roots ...*Root) {
// It is not an error to remove a nonexistent root.
// TODO: notification
func (c *Client) RemoveRoots(uris ...string) {
c.changeAndNotify(notificationRootsListChanged, &RootsListChangedParams{},
changeAndNotify(c, notificationRootsListChanged, &RootsListChangedParams{},
func() bool { return c.roots.remove(uris...) })
}
// changeAndNotify is called when a feature is added or removed.
// It calls change, which should do the work and report whether a change actually occurred.
// If there was a change, it notifies a snapshot of the sessions.
func (c *Client) changeAndNotify(notification string, params Params, change func() bool) {
func changeAndNotify[P Params](c *Client, notification string, params P, change func() bool) {
var sessions []*ClientSession
// Lock for the change, but not for the notification.
c.mu.Lock()
@ -243,7 +245,7 @@ func (c *Client) changeAndNotify(notification string, params Params, change func
notifySessions(sessions, notification, params)
}
func (c *Client) listRoots(_ context.Context, _ *ClientSession, _ *ListRootsParams) (*ListRootsResult, error) {
func (c *Client) listRoots(_ context.Context, req *ClientRequest[*ListRootsParams]) (*ListRootsResult, error) {
c.mu.Lock()
defer c.mu.Unlock()
roots := slices.Collect(c.roots.all())
@ -255,12 +257,12 @@ func (c *Client) listRoots(_ context.Context, _ *ClientSession, _ *ListRootsPara
}, nil
}
func (c *Client) createMessage(ctx context.Context, cs *ClientSession, params *CreateMessageParams) (*CreateMessageResult, error) {
func (c *Client) createMessage(ctx context.Context, req *ClientRequest[*CreateMessageParams]) (*CreateMessageResult, error) {
if c.opts.CreateMessageHandler == nil {
// TODO: wrap or annotate this error? Pick a standard code?
return nil, jsonrpc2.NewError(CodeUnsupportedMethod, "client does not support CreateMessage")
}
return c.opts.CreateMessageHandler(ctx, cs, params)
return c.opts.CreateMessageHandler(ctx, req)
}
// AddSendingMiddleware wraps the current sending method handler using the provided
@ -272,7 +274,7 @@ func (c *Client) createMessage(ctx context.Context, cs *ClientSession, params *C
//
// 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]) {
func (c *Client) AddSendingMiddleware(middleware ...Middleware) {
c.mu.Lock()
defer c.mu.Unlock()
addMiddleware(&c.sendingMethodHandler_, middleware)
@ -287,7 +289,7 @@ func (c *Client) AddSendingMiddleware(middleware ...Middleware[*ClientSession])
//
// 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]) {
func (c *Client) AddReceivingMiddleware(middleware ...Middleware) {
c.mu.Lock()
defer c.mu.Unlock()
addMiddleware(&c.receivingMethodHandler_, middleware)
@ -299,16 +301,16 @@ func (c *Client) AddReceivingMiddleware(middleware ...Middleware[*ClientSession]
// TODO(rfindley): actually load and validate the protocol schema, rather than
// curating these method flags.
var clientMethodInfos = map[string]methodInfo{
methodComplete: newMethodInfo(sessionMethod((*ClientSession).Complete), 0),
methodPing: newMethodInfo(sessionMethod((*ClientSession).ping), missingParamsOK),
methodListRoots: newMethodInfo(clientMethod((*Client).listRoots), missingParamsOK),
methodCreateMessage: newMethodInfo(clientMethod((*Client).createMessage), 0),
notificationToolListChanged: newMethodInfo(clientMethod((*Client).callToolChangedHandler), notification|missingParamsOK),
notificationPromptListChanged: newMethodInfo(clientMethod((*Client).callPromptChangedHandler), notification|missingParamsOK),
notificationResourceListChanged: newMethodInfo(clientMethod((*Client).callResourceChangedHandler), notification|missingParamsOK),
notificationResourceUpdated: newMethodInfo(clientMethod((*Client).callResourceUpdatedHandler), notification|missingParamsOK),
notificationLoggingMessage: newMethodInfo(clientMethod((*Client).callLoggingHandler), notification),
notificationProgress: newMethodInfo(sessionMethod((*ClientSession).callProgressNotificationHandler), notification),
methodComplete: newClientMethodInfo(clientSessionMethod((*ClientSession).Complete), 0),
methodPing: newClientMethodInfo(clientSessionMethod((*ClientSession).ping), missingParamsOK),
methodListRoots: newClientMethodInfo(clientMethod((*Client).listRoots), missingParamsOK),
methodCreateMessage: newClientMethodInfo(clientMethod((*Client).createMessage), 0),
notificationToolListChanged: newClientMethodInfo(clientMethod((*Client).callToolChangedHandler), notification|missingParamsOK),
notificationPromptListChanged: newClientMethodInfo(clientMethod((*Client).callPromptChangedHandler), notification|missingParamsOK),
notificationResourceListChanged: newClientMethodInfo(clientMethod((*Client).callResourceChangedHandler), notification|missingParamsOK),
notificationResourceUpdated: newClientMethodInfo(clientMethod((*Client).callResourceUpdatedHandler), notification|missingParamsOK),
notificationLoggingMessage: newClientMethodInfo(clientMethod((*Client).callLoggingHandler), notification),
notificationProgress: newClientMethodInfo(clientSessionMethod((*ClientSession).callProgressNotificationHandler), notification),
}
func (cs *ClientSession) sendingMethodInfos() map[string]methodInfo {
@ -342,25 +344,29 @@ func (*ClientSession) ping(context.Context, *PingParams) (*emptyResult, error) {
return &emptyResult{}, nil
}
func newClientRequest[P Params](cs *ClientSession, params P) *ClientRequest[P] {
return &ClientRequest[P]{Session: cs, Params: params}
}
// Ping makes an MCP "ping" request to the server.
func (cs *ClientSession) Ping(ctx context.Context, params *PingParams) error {
_, err := handleSend[*emptyResult](ctx, cs, methodPing, orZero[Params](params))
_, err := handleSend[*emptyResult](ctx, methodPing, newClientRequest(cs, orZero[Params](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 handleSend[*ListPromptsResult](ctx, cs, methodListPrompts, orZero[Params](params))
return handleSend[*ListPromptsResult](ctx, methodListPrompts, newClientRequest(cs, orZero[Params](params)))
}
// GetPrompt gets a prompt from the server.
func (cs *ClientSession) GetPrompt(ctx context.Context, params *GetPromptParams) (*GetPromptResult, error) {
return handleSend[*GetPromptResult](ctx, cs, methodGetPrompt, orZero[Params](params))
return handleSend[*GetPromptResult](ctx, methodGetPrompt, newClientRequest(cs, orZero[Params](params)))
}
// ListTools lists tools that are currently available on the server.
func (cs *ClientSession) ListTools(ctx context.Context, params *ListToolsParams) (*ListToolsResult, error) {
return handleSend[*ListToolsResult](ctx, cs, methodListTools, orZero[Params](params))
return handleSend[*ListToolsResult](ctx, methodListTools, newClientRequest(cs, orZero[Params](params)))
}
// CallTool calls the tool with the given name and arguments.
@ -373,72 +379,87 @@ func (cs *ClientSession) CallTool(ctx context.Context, params *CallToolParams) (
// Avoid sending nil over the wire.
params.Arguments = map[string]any{}
}
return handleSend[*CallToolResult](ctx, cs, methodCallTool, params)
return handleSend[*CallToolResult](ctx, methodCallTool, newClientRequest(cs, orZero[Params](params)))
}
func (cs *ClientSession) SetLevel(ctx context.Context, params *SetLevelParams) error {
_, err := handleSend[*emptyResult](ctx, cs, methodSetLevel, orZero[Params](params))
_, err := handleSend[*emptyResult](ctx, methodSetLevel, newClientRequest(cs, orZero[Params](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 handleSend[*ListResourcesResult](ctx, cs, methodListResources, orZero[Params](params))
return handleSend[*ListResourcesResult](ctx, methodListResources, newClientRequest(cs, orZero[Params](params)))
}
// ListResourceTemplates lists the resource templates that are currently available on the server.
func (cs *ClientSession) ListResourceTemplates(ctx context.Context, params *ListResourceTemplatesParams) (*ListResourceTemplatesResult, error) {
return handleSend[*ListResourceTemplatesResult](ctx, cs, methodListResourceTemplates, orZero[Params](params))
return handleSend[*ListResourceTemplatesResult](ctx, methodListResourceTemplates, newClientRequest(cs, orZero[Params](params)))
}
// ReadResource asks the server to read a resource and return its contents.
func (cs *ClientSession) ReadResource(ctx context.Context, params *ReadResourceParams) (*ReadResourceResult, error) {
return handleSend[*ReadResourceResult](ctx, cs, methodReadResource, orZero[Params](params))
return handleSend[*ReadResourceResult](ctx, methodReadResource, newClientRequest(cs, orZero[Params](params)))
}
func (cs *ClientSession) Complete(ctx context.Context, params *CompleteParams) (*CompleteResult, error) {
return handleSend[*CompleteResult](ctx, cs, methodComplete, orZero[Params](params))
return handleSend[*CompleteResult](ctx, methodComplete, newClientRequest(cs, orZero[Params](params)))
}
// Subscribe sends a "resources/subscribe" request to the server, asking for
// notifications when the specified resource changes.
func (cs *ClientSession) Subscribe(ctx context.Context, params *SubscribeParams) error {
_, err := handleSend[*emptyResult](ctx, cs, methodSubscribe, orZero[Params](params))
_, err := handleSend[*emptyResult](ctx, methodSubscribe, newClientRequest(cs, orZero[Params](params)))
return err
}
// Unsubscribe sends a "resources/unsubscribe" request to the server, cancelling
// a previous subscription.
func (cs *ClientSession) Unsubscribe(ctx context.Context, params *UnsubscribeParams) error {
_, err := handleSend[*emptyResult](ctx, cs, methodUnsubscribe, orZero[Params](params))
_, err := handleSend[*emptyResult](ctx, methodUnsubscribe, newClientRequest(cs, orZero[Params](params)))
return err
}
func (c *Client) callToolChangedHandler(ctx context.Context, s *ClientSession, params *ToolListChangedParams) (Result, error) {
return callNotificationHandler(ctx, c.opts.ToolListChangedHandler, s, params)
func (c *Client) callToolChangedHandler(ctx context.Context, req *ClientRequest[*ToolListChangedParams]) (Result, error) {
if h := c.opts.ToolListChangedHandler; h != nil {
h(ctx, req)
}
return nil, nil
}
func (c *Client) callPromptChangedHandler(ctx context.Context, s *ClientSession, params *PromptListChangedParams) (Result, error) {
return callNotificationHandler(ctx, c.opts.PromptListChangedHandler, s, params)
func (c *Client) callPromptChangedHandler(ctx context.Context, req *ClientRequest[*PromptListChangedParams]) (Result, error) {
if h := c.opts.PromptListChangedHandler; h != nil {
h(ctx, req)
}
return nil, nil
}
func (c *Client) callResourceChangedHandler(ctx context.Context, s *ClientSession, params *ResourceListChangedParams) (Result, error) {
return callNotificationHandler(ctx, c.opts.ResourceListChangedHandler, s, params)
func (c *Client) callResourceChangedHandler(ctx context.Context, req *ClientRequest[*ResourceListChangedParams]) (Result, error) {
if h := c.opts.ResourceListChangedHandler; h != nil {
h(ctx, req)
}
return nil, nil
}
func (c *Client) callResourceUpdatedHandler(ctx context.Context, s *ClientSession, params *ResourceUpdatedNotificationParams) (Result, error) {
return callNotificationHandler(ctx, c.opts.ResourceUpdatedHandler, s, params)
func (c *Client) callResourceUpdatedHandler(ctx context.Context, req *ClientRequest[*ResourceUpdatedNotificationParams]) (Result, error) {
if h := c.opts.ResourceUpdatedHandler; h != nil {
h(ctx, req)
}
return nil, nil
}
func (c *Client) callLoggingHandler(ctx context.Context, cs *ClientSession, params *LoggingMessageParams) (Result, error) {
func (c *Client) callLoggingHandler(ctx context.Context, req *ClientRequest[*LoggingMessageParams]) (Result, error) {
if h := c.opts.LoggingMessageHandler; h != nil {
h(ctx, cs, params)
h(ctx, req)
}
return nil, nil
}
func (cs *ClientSession) callProgressNotificationHandler(ctx context.Context, params *ProgressNotificationParams) (Result, error) {
return callNotificationHandler(ctx, cs.client.opts.ProgressNotificationHandler, cs, params)
if h := cs.client.opts.ProgressNotificationHandler; h != nil {
h(ctx, clientRequestFor(cs, params))
}
return nil, nil
}
// NotifyProgress sends a progress notification from the client to the server
@ -446,7 +467,7 @@ func (cs *ClientSession) callProgressNotificationHandler(ctx context.Context, pa
// This can be used if the client is performing a long-running task that was
// initiated by the server
func (cs *ClientSession) NotifyProgress(ctx context.Context, params *ProgressNotificationParams) error {
return handleNotify(ctx, cs, notificationProgress, params)
return handleNotify(ctx, notificationProgress, newClientRequest(cs, orZero[Params](params)))
}
// Tools provides an iterator for all tools available on the server,

20
mcp/example_middleware_test.go
View file

@ -29,36 +29,35 @@ func Example_loggingMiddleware() {
},
}))
loggingMiddleware := func(next mcp.MethodHandler[*mcp.ServerSession]) mcp.MethodHandler[*mcp.ServerSession] {
loggingMiddleware := func(next mcp.MethodHandler) mcp.MethodHandler {
return func(
ctx context.Context,
session *mcp.ServerSession,
method string,
params mcp.Params,
req mcp.Request,
) (mcp.Result, error) {
logger.Info("MCP method started",
"method", method,
"session_id", session.ID(),
"has_params", params != nil,
"session_id", req.GetSession().ID(),
"has_params", req.GetParams() != nil,
)
start := time.Now()
result, err := next(ctx, session, method, params)
result, err := next(ctx, method, req)
duration := time.Since(start)
if err != nil {
logger.Error("MCP method failed",
"method", method,
"session_id", session.ID(),
"session_id", req.GetSession().ID(),
"duration_ms", duration.Milliseconds(),
"err", err,
)
} else {
logger.Info("MCP method completed",
"method", method,
"session_id", session.ID(),
"session_id", req.GetSession().ID(),
"duration_ms", duration.Milliseconds(),
"has_result", result != nil,
)
@ -90,10 +89,9 @@ func Example_loggingMiddleware() {
},
func(
ctx context.Context,
ss *mcp.ServerSession,
params *mcp.CallToolParamsFor[map[string]any],
req *mcp.ServerRequest[*mcp.CallToolParamsFor[map[string]any]],
) (*mcp.CallToolResultFor[any], error) {
name, ok := params.Arguments["name"].(string)
name, ok := req.Params.Arguments["name"].(string)
if !ok {
return nil, fmt.Errorf("name parameter is required and must be a string")
}

8
mcp/example_progress_test.go
View file

@ -21,11 +21,11 @@ func Example_progressMiddleware() {
_ = 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) {
if rp, ok := params.(mcp.RequestParams); ok {
func addProgressToken[S mcp.Session](h mcp.MethodHandler) mcp.MethodHandler {
return func(ctx context.Context, method string, req mcp.Request) (result mcp.Result, err error) {
if rp, ok := req.GetParams().(mcp.RequestParams); ok {
rp.SetProgressToken(nextProgressToken.Add(1))
}
return h(ctx, s, method, params)
return h(ctx, method, req)
}
}

62
mcp/mcp_test.go
View file

@ -32,11 +32,11 @@ type hiParams struct {
// TODO(jba): after schemas are stateless (WIP), this can be a variable.
func greetTool() *Tool { return &Tool{Name: "greet", Description: "say hi"} }
func sayHi(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[hiParams]) (*CallToolResultFor[any], error) {
if err := ss.Ping(ctx, nil); err != nil {
func sayHi(ctx context.Context, req *ServerRequest[*CallToolParamsFor[hiParams]]) (*CallToolResultFor[any], error) {
if err := req.Session.Ping(ctx, nil); err != nil {
return nil, fmt.Errorf("ping failed: %v", err)
}
return &CallToolResultFor[any]{Content: []Content{&TextContent{Text: "hi " + params.Arguments.Name}}}, nil
return &CallToolResultFor[any]{Content: []Content{&TextContent{Text: "hi " + req.Params.Arguments.Name}}}, nil
}
var codeReviewPrompt = &Prompt{
@ -73,16 +73,20 @@ func TestEndToEnd(t *testing.T) {
}
sopts := &ServerOptions{
InitializedHandler: func(context.Context, *ServerSession, *InitializedParams) { notificationChans["initialized"] <- 0 },
RootsListChangedHandler: func(context.Context, *ServerSession, *RootsListChangedParams) { notificationChans["roots"] <- 0 },
ProgressNotificationHandler: func(context.Context, *ServerSession, *ProgressNotificationParams) {
InitializedHandler: func(context.Context, *ServerRequest[*InitializedParams]) {
notificationChans["initialized"] <- 0
},
RootsListChangedHandler: func(context.Context, *ServerRequest[*RootsListChangedParams]) {
notificationChans["roots"] <- 0
},
ProgressNotificationHandler: func(context.Context, *ServerRequest[*ProgressNotificationParams]) {
notificationChans["progress_server"] <- 0
},
SubscribeHandler: func(context.Context, *ServerSession, *SubscribeParams) error {
SubscribeHandler: func(context.Context, *ServerRequest[*SubscribeParams]) error {
notificationChans["subscribe"] <- 0
return nil
},
UnsubscribeHandler: func(context.Context, *ServerSession, *UnsubscribeParams) error {
UnsubscribeHandler: func(context.Context, *ServerRequest[*UnsubscribeParams]) error {
notificationChans["unsubscribe"] <- 0
return nil
},
@ -93,7 +97,7 @@ func TestEndToEnd(t *testing.T) {
Description: "say hi",
}, sayHi)
s.AddTool(&Tool{Name: "fail", InputSchema: &jsonschema.Schema{}},
func(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
func(context.Context, *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResult, error) {
return nil, errTestFailure
})
s.AddPrompt(codeReviewPrompt, codReviewPromptHandler)
@ -124,19 +128,25 @@ func TestEndToEnd(t *testing.T) {
loggingMessages := make(chan *LoggingMessageParams, 100) // big enough for all logging
opts := &ClientOptions{
CreateMessageHandler: func(context.Context, *ClientSession, *CreateMessageParams) (*CreateMessageResult, error) {
CreateMessageHandler: func(context.Context, *ClientRequest[*CreateMessageParams]) (*CreateMessageResult, error) {
return &CreateMessageResult{Model: "aModel", Content: &TextContent{}}, nil
},
ToolListChangedHandler: func(context.Context, *ClientSession, *ToolListChangedParams) { notificationChans["tools"] <- 0 },
PromptListChangedHandler: func(context.Context, *ClientSession, *PromptListChangedParams) { notificationChans["prompts"] <- 0 },
ResourceListChangedHandler: func(context.Context, *ClientSession, *ResourceListChangedParams) { notificationChans["resources"] <- 0 },
LoggingMessageHandler: func(_ context.Context, _ *ClientSession, lm *LoggingMessageParams) {
loggingMessages <- lm
ToolListChangedHandler: func(context.Context, *ClientRequest[*ToolListChangedParams]) {
notificationChans["tools"] <- 0
},
ProgressNotificationHandler: func(context.Context, *ClientSession, *ProgressNotificationParams) {
PromptListChangedHandler: func(context.Context, *ClientRequest[*PromptListChangedParams]) {
notificationChans["prompts"] <- 0
},
ResourceListChangedHandler: func(context.Context, *ClientRequest[*ResourceListChangedParams]) {
notificationChans["resources"] <- 0
},
LoggingMessageHandler: func(_ context.Context, req *ClientRequest[*LoggingMessageParams]) {
loggingMessages <- req.Params
},
ProgressNotificationHandler: func(context.Context, *ClientRequest[*ProgressNotificationParams]) {
notificationChans["progress_client"] <- 0
},
ResourceUpdatedHandler: func(context.Context, *ClientSession, *ResourceUpdatedNotificationParams) {
ResourceUpdatedHandler: func(context.Context, *ClientRequest[*ResourceUpdatedNotificationParams]) {
notificationChans["resource_updated"] <- 0
},
}
@ -500,8 +510,8 @@ var embeddedResources = map[string]string{
"info": "This is the MCP test server.",
}
func handleEmbeddedResource(_ context.Context, _ *ServerSession, params *ReadResourceParams) (*ReadResourceResult, error) {
u, err := url.Parse(params.URI)
func handleEmbeddedResource(_ context.Context, req *ServerRequest[*ReadResourceParams]) (*ReadResourceResult, error) {
u, err := url.Parse(req.Params.URI)
if err != nil {
return nil, err
}
@ -515,7 +525,7 @@ func handleEmbeddedResource(_ context.Context, _ *ServerSession, params *ReadRes
}
return &ReadResourceResult{
Contents: []*ResourceContents{
{URI: params.URI, MIMEType: "text/plain", Text: text},
{URI: req.Params.URI, MIMEType: "text/plain", Text: text},
},
}, nil
}
@ -637,7 +647,7 @@ func TestCancellation(t *testing.T) {
cancelled = make(chan struct{}, 1) // don't block the request
)
slowRequest := func(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
slowRequest := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResult, error) {
start <- struct{}{}
select {
case <-ctx.Done():
@ -816,17 +826,17 @@ func TestNoJSONNull(t *testing.T) {
// traceCalls creates a middleware function that prints the method before and after each call
// with the given prefix.
func traceCalls[S Session](w io.Writer, prefix string) Middleware[S] {
return func(h MethodHandler[S]) MethodHandler[S] {
return func(ctx context.Context, sess S, method string, params Params) (Result, error) {
func traceCalls[S Session](w io.Writer, prefix string) Middleware {
return func(h MethodHandler) MethodHandler {
return func(ctx context.Context, method string, req Request) (Result, error) {
fmt.Fprintf(w, "%s >%s\n", prefix, method)
defer fmt.Fprintf(w, "%s <%s\n", prefix, method)
return h(ctx, sess, method, params)
return h(ctx, method, req)
}
}
}
func nopHandler(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
func nopHandler(context.Context, *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResult, error) {
return nil, nil
}

2
mcp/resource.go
View file

@ -35,7 +35,7 @@ type serverResourceTemplate struct {
// A ResourceHandler is a function that reads a resource.
// It will be called when the client calls [ClientSession.ReadResource].
// If it cannot find the resource, it should return the result of calling [ResourceNotFoundError].
type ResourceHandler func(context.Context, *ServerSession, *ReadResourceParams) (*ReadResourceResult, error)
type ResourceHandler func(context.Context, *ServerRequest[*ReadResourceParams]) (*ReadResourceResult, error)
// ResourceNotFoundError returns an error indicating that a resource being read could
// not be found.

194
mcp/server.go
View file

@ -41,8 +41,8 @@ type Server struct {
resources *featureSet[*serverResource]
resourceTemplates *featureSet[*serverResourceTemplate]
sessions []*ServerSession
sendingMethodHandler_ MethodHandler[*ServerSession]
receivingMethodHandler_ MethodHandler[*ServerSession]
sendingMethodHandler_ MethodHandler
receivingMethodHandler_ MethodHandler
resourceSubscriptions map[string]map[*ServerSession]bool // uri -> session -> bool
}
@ -51,24 +51,24 @@ type ServerOptions struct {
// Optional instructions for connected clients.
Instructions string
// If non-nil, called when "notifications/initialized" is received.
InitializedHandler func(context.Context, *ServerSession, *InitializedParams)
InitializedHandler func(context.Context, *ServerRequest[*InitializedParams])
// PageSize is the maximum number of items to return in a single page for
// list methods (e.g. ListTools).
PageSize int
// If non-nil, called when "notifications/roots/list_changed" is received.
RootsListChangedHandler func(context.Context, *ServerSession, *RootsListChangedParams)
RootsListChangedHandler func(context.Context, *ServerRequest[*RootsListChangedParams])
// If non-nil, called when "notifications/progress" is received.
ProgressNotificationHandler func(context.Context, *ServerSession, *ProgressNotificationParams)
ProgressNotificationHandler func(context.Context, *ServerRequest[*ProgressNotificationParams])
// If non-nil, called when "completion/complete" is received.
CompletionHandler func(context.Context, *ServerSession, *CompleteParams) (*CompleteResult, error)
CompletionHandler func(context.Context, *ServerRequest[*CompleteParams]) (*CompleteResult, error)
// If non-zero, defines an interval for regular "ping" requests.
// If the peer fails to respond to pings originating from the keepalive check,
// the session is automatically closed.
KeepAlive time.Duration
// Function called when a client session subscribes to a resource.
SubscribeHandler func(context.Context, *ServerSession, *SubscribeParams) error
SubscribeHandler func(context.Context, *ServerRequest[*SubscribeParams]) error
// Function called when a client session unsubscribes from a resource.
UnsubscribeHandler func(context.Context, *ServerSession, *UnsubscribeParams) error
UnsubscribeHandler func(context.Context, *ServerRequest[*UnsubscribeParams]) error
// If true, advertises the prompts capability during initialization,
// even if no prompts have been registered.
HasPrompts bool
@ -258,11 +258,11 @@ func (s *Server) capabilities() *serverCapabilities {
return caps
}
func (s *Server) complete(ctx context.Context, ss *ServerSession, params *CompleteParams) (Result, error) {
func (s *Server) complete(ctx context.Context, req *ServerRequest[*CompleteParams]) (Result, error) {
if s.opts.CompletionHandler == nil {
return nil, jsonrpc2.ErrMethodNotFound
}
return s.opts.CompletionHandler(ctx, ss, params)
return s.opts.CompletionHandler(ctx, req)
}
// changeAndNotify is called when a feature is added or removed.
@ -287,13 +287,13 @@ func (s *Server) Sessions() iter.Seq[*ServerSession] {
return slices.Values(clients)
}
func (s *Server) listPrompts(_ context.Context, _ *ServerSession, params *ListPromptsParams) (*ListPromptsResult, error) {
func (s *Server) listPrompts(_ context.Context, req *ServerRequest[*ListPromptsParams]) (*ListPromptsResult, error) {
s.mu.Lock()
defer s.mu.Unlock()
if params == nil {
params = &ListPromptsParams{}
if req.Params == nil {
req.Params = &ListPromptsParams{}
}
return paginateList(s.prompts, s.opts.PageSize, params, &ListPromptsResult{}, func(res *ListPromptsResult, prompts []*serverPrompt) {
return paginateList(s.prompts, s.opts.PageSize, req.Params, &ListPromptsResult{}, func(res *ListPromptsResult, prompts []*serverPrompt) {
res.Prompts = []*Prompt{} // avoid JSON null
for _, p := range prompts {
res.Prompts = append(res.Prompts, p.prompt)
@ -301,24 +301,24 @@ func (s *Server) listPrompts(_ context.Context, _ *ServerSession, params *ListPr
})
}
func (s *Server) getPrompt(ctx context.Context, ss *ServerSession, params *GetPromptParams) (*GetPromptResult, error) {
func (s *Server) getPrompt(ctx context.Context, req *ServerRequest[*GetPromptParams]) (*GetPromptResult, error) {
s.mu.Lock()
prompt, ok := s.prompts.get(params.Name)
prompt, ok := s.prompts.get(req.Params.Name)
s.mu.Unlock()
if !ok {
// TODO: surface the error code over the wire, instead of flattening it into the string.
return nil, fmt.Errorf("%s: unknown prompt %q", jsonrpc2.ErrInvalidParams, params.Name)
return nil, fmt.Errorf("%s: unknown prompt %q", jsonrpc2.ErrInvalidParams, req.Params.Name)
}
return prompt.handler(ctx, ss, params)
return prompt.handler(ctx, req.Session, req.Params)
}
func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListToolsParams) (*ListToolsResult, error) {
func (s *Server) listTools(_ context.Context, req *ServerRequest[*ListToolsParams]) (*ListToolsResult, error) {
s.mu.Lock()
defer s.mu.Unlock()
if params == nil {
params = &ListToolsParams{}
if req.Params == nil {
req.Params = &ListToolsParams{}
}
return paginateList(s.tools, s.opts.PageSize, params, &ListToolsResult{}, func(res *ListToolsResult, tools []*serverTool) {
return paginateList(s.tools, s.opts.PageSize, req.Params, &ListToolsResult{}, func(res *ListToolsResult, tools []*serverTool) {
res.Tools = []*Tool{} // avoid JSON null
for _, t := range tools {
res.Tools = append(res.Tools, t.tool)
@ -326,23 +326,23 @@ func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListTool
})
}
func (s *Server) callTool(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) {
func (s *Server) callTool(ctx context.Context, req *ServerRequest[*CallToolParamsFor[json.RawMessage]]) (*CallToolResult, error) {
s.mu.Lock()
st, ok := s.tools.get(params.Name)
st, ok := s.tools.get(req.Params.Name)
s.mu.Unlock()
if !ok {
return nil, fmt.Errorf("%s: unknown tool %q", jsonrpc2.ErrInvalidParams, params.Name)
return nil, fmt.Errorf("%s: unknown tool %q", jsonrpc2.ErrInvalidParams, req.Params.Name)
}
return st.handler(ctx, cc, params)
return st.handler(ctx, req)
}
func (s *Server) listResources(_ context.Context, _ *ServerSession, params *ListResourcesParams) (*ListResourcesResult, error) {
func (s *Server) listResources(_ context.Context, req *ServerRequest[*ListResourcesParams]) (*ListResourcesResult, error) {
s.mu.Lock()
defer s.mu.Unlock()
if params == nil {
params = &ListResourcesParams{}
if req.Params == nil {
req.Params = &ListResourcesParams{}
}
return paginateList(s.resources, s.opts.PageSize, params, &ListResourcesResult{}, func(res *ListResourcesResult, resources []*serverResource) {
return paginateList(s.resources, s.opts.PageSize, req.Params, &ListResourcesResult{}, func(res *ListResourcesResult, resources []*serverResource) {
res.Resources = []*Resource{} // avoid JSON null
for _, r := range resources {
res.Resources = append(res.Resources, r.resource)
@ -350,13 +350,13 @@ func (s *Server) listResources(_ context.Context, _ *ServerSession, params *List
})
}
func (s *Server) listResourceTemplates(_ context.Context, _ *ServerSession, params *ListResourceTemplatesParams) (*ListResourceTemplatesResult, error) {
func (s *Server) listResourceTemplates(_ context.Context, req *ServerRequest[*ListResourceTemplatesParams]) (*ListResourceTemplatesResult, error) {
s.mu.Lock()
defer s.mu.Unlock()
if params == nil {
params = &ListResourceTemplatesParams{}
if req.Params == nil {
req.Params = &ListResourceTemplatesParams{}
}
return paginateList(s.resourceTemplates, s.opts.PageSize, params, &ListResourceTemplatesResult{},
return paginateList(s.resourceTemplates, s.opts.PageSize, req.Params, &ListResourceTemplatesResult{},
func(res *ListResourceTemplatesResult, rts []*serverResourceTemplate) {
res.ResourceTemplates = []*ResourceTemplate{} // avoid JSON null
for _, rt := range rts {
@ -365,8 +365,8 @@ func (s *Server) listResourceTemplates(_ context.Context, _ *ServerSession, para
})
}
func (s *Server) readResource(ctx context.Context, ss *ServerSession, params *ReadResourceParams) (*ReadResourceResult, error) {
uri := params.URI
func (s *Server) readResource(ctx context.Context, req *ServerRequest[*ReadResourceParams]) (*ReadResourceResult, error) {
uri := req.Params.URI
// Look up the resource URI in the lists of resources and resource templates.
// This is a security check as well as an information lookup.
handler, mimeType, ok := s.lookupResourceHandler(uri)
@ -375,7 +375,7 @@ func (s *Server) readResource(ctx context.Context, ss *ServerSession, params *Re
// Treat an unregistered resource the same as a registered one that couldn't be found.
return nil, ResourceNotFoundError(uri)
}
res, err := handler(ctx, ss, params)
res, err := handler(ctx, req)
if err != nil {
return nil, err
}
@ -430,11 +430,11 @@ func fileResourceHandler(dir string) ResourceHandler {
if err != nil {
panic(err)
}
return func(ctx context.Context, ss *ServerSession, params *ReadResourceParams) (_ *ReadResourceResult, err error) {
defer util.Wrapf(&err, "reading resource %s", params.URI)
return func(ctx context.Context, req *ServerRequest[*ReadResourceParams]) (_ *ReadResourceResult, err error) {
defer util.Wrapf(&err, "reading resource %s", req.Params.URI)
// TODO: use a memoizing API here.
rootRes, err := ss.ListRoots(ctx, nil)
rootRes, err := req.Session.ListRoots(ctx, nil)
if err != nil {
return nil, fmt.Errorf("listing roots: %w", err)
}
@ -442,13 +442,13 @@ func fileResourceHandler(dir string) ResourceHandler {
if err != nil {
return nil, err
}
data, err := readFileResource(params.URI, dirFilepath, roots)
data, err := readFileResource(req.Params.URI, dirFilepath, roots)
if err != nil {
return nil, err
}
// TODO(jba): figure out mime type. Omit for now: Server.readResource will fill it in.
return &ReadResourceResult{Contents: []*ResourceContents{
{URI: params.URI, Blob: data},
{URI: req.Params.URI, Blob: data},
}}, nil
}
}
@ -465,39 +465,39 @@ func (s *Server) ResourceUpdated(ctx context.Context, params *ResourceUpdatedNot
return nil
}
func (s *Server) subscribe(ctx context.Context, ss *ServerSession, params *SubscribeParams) (*emptyResult, error) {
func (s *Server) subscribe(ctx context.Context, req *ServerRequest[*SubscribeParams]) (*emptyResult, error) {
if s.opts.SubscribeHandler == nil {
return nil, fmt.Errorf("%w: server does not support resource subscriptions", jsonrpc2.ErrMethodNotFound)
}
if err := s.opts.SubscribeHandler(ctx, ss, params); err != nil {
if err := s.opts.SubscribeHandler(ctx, req); err != nil {
return nil, err
}
s.mu.Lock()
defer s.mu.Unlock()
if s.resourceSubscriptions[params.URI] == nil {
s.resourceSubscriptions[params.URI] = make(map[*ServerSession]bool)
if s.resourceSubscriptions[req.Params.URI] == nil {
s.resourceSubscriptions[req.Params.URI] = make(map[*ServerSession]bool)
}
s.resourceSubscriptions[params.URI][ss] = true
s.resourceSubscriptions[req.Params.URI][req.Session] = true
return &emptyResult{}, nil
}
func (s *Server) unsubscribe(ctx context.Context, ss *ServerSession, params *UnsubscribeParams) (*emptyResult, error) {
func (s *Server) unsubscribe(ctx context.Context, req *ServerRequest[*UnsubscribeParams]) (*emptyResult, error) {
if s.opts.UnsubscribeHandler == nil {
return nil, jsonrpc2.ErrMethodNotFound
}
if err := s.opts.UnsubscribeHandler(ctx, ss, params); err != nil {
if err := s.opts.UnsubscribeHandler(ctx, req); err != nil {
return nil, err
}
s.mu.Lock()
defer s.mu.Unlock()
if subscribedSessions, ok := s.resourceSubscriptions[params.URI]; ok {
delete(subscribedSessions, ss)
if subscribedSessions, ok := s.resourceSubscriptions[req.Params.URI]; ok {
delete(subscribedSessions, req.Session)
if len(subscribedSessions) == 0 {
delete(s.resourceSubscriptions, params.URI)
delete(s.resourceSubscriptions, req.Params.URI)
}
}
@ -611,15 +611,24 @@ func (ss *ServerSession) initialized(ctx context.Context, params *InitializedPar
if wasInitd {
return nil, fmt.Errorf("duplicate %q received", notificationInitialized)
}
return callNotificationHandler(ctx, ss.server.opts.InitializedHandler, ss, params)
if h := ss.server.opts.InitializedHandler; h != nil {
h(ctx, serverRequestFor(ss, params))
}
return nil, nil
}
func (s *Server) callRootsListChangedHandler(ctx context.Context, ss *ServerSession, params *RootsListChangedParams) (Result, error) {
return callNotificationHandler(ctx, s.opts.RootsListChangedHandler, ss, params)
func (s *Server) callRootsListChangedHandler(ctx context.Context, req *ServerRequest[*RootsListChangedParams]) (Result, error) {
if h := s.opts.RootsListChangedHandler; h != nil {
h(ctx, req)
}
return nil, nil
}
func (ss *ServerSession) callProgressNotificationHandler(ctx context.Context, params *ProgressNotificationParams) (Result, error) {
return callNotificationHandler(ctx, ss.server.opts.ProgressNotificationHandler, ss, params)
func (ss *ServerSession) callProgressNotificationHandler(ctx context.Context, p *ProgressNotificationParams) (Result, error) {
if h := ss.server.opts.ProgressNotificationHandler; h != nil {
h(ctx, serverRequestFor(ss, p))
}
return nil, nil
}
// NotifyProgress sends a progress notification from the server to the client
@ -627,7 +636,11 @@ func (ss *ServerSession) callProgressNotificationHandler(ctx context.Context, pa
// This is typically used to report on the status of a long-running request
// that was initiated by the client.
func (ss *ServerSession) NotifyProgress(ctx context.Context, params *ProgressNotificationParams) error {
return handleNotify(ctx, ss, notificationProgress, params)
return handleNotify(ctx, notificationProgress, newServerRequest(ss, orZero[Params](params)))
}
func newServerRequest[P Params](ss *ServerSession, params P) *ServerRequest[P] {
return &ServerRequest[P]{Session: ss, Params: params}
}
// A ServerSession is a logical connection from a single MCP client. Its
@ -665,18 +678,18 @@ func (ss *ServerSession) ID() string {
// Ping pings the client.
func (ss *ServerSession) Ping(ctx context.Context, params *PingParams) error {
_, err := handleSend[*emptyResult](ctx, ss, methodPing, orZero[Params](params))
_, err := handleSend[*emptyResult](ctx, methodPing, newServerRequest(ss, orZero[Params](params)))
return err
}
// ListRoots lists the client roots.
func (ss *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) {
return handleSend[*ListRootsResult](ctx, ss, methodListRoots, orZero[Params](params))
return handleSend[*ListRootsResult](ctx, methodListRoots, newServerRequest(ss, orZero[Params](params)))
}
// CreateMessage sends a sampling request to the client.
func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessageParams) (*CreateMessageResult, error) {
return handleSend[*CreateMessageResult](ctx, ss, methodCreateMessage, orZero[Params](params))
return handleSend[*CreateMessageResult](ctx, methodCreateMessage, newServerRequest(ss, orZero[Params](params)))
}
// Log sends a log message to the client.
@ -695,7 +708,7 @@ func (ss *ServerSession) Log(ctx context.Context, params *LoggingMessageParams)
if compareLevels(params.Level, logLevel) < 0 {
return nil
}
return handleNotify(ctx, ss, notificationLoggingMessage, params)
return handleNotify(ctx, notificationLoggingMessage, newServerRequest(ss, orZero[Params](params)))
}
// AddSendingMiddleware wraps the current sending method handler using the provided
@ -707,7 +720,7 @@ func (ss *ServerSession) Log(ctx context.Context, params *LoggingMessageParams)
//
// 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]) {
func (s *Server) AddSendingMiddleware(middleware ...Middleware) {
s.mu.Lock()
defer s.mu.Unlock()
addMiddleware(&s.sendingMethodHandler_, middleware)
@ -722,7 +735,7 @@ func (s *Server) AddSendingMiddleware(middleware ...Middleware[*ServerSession])
//
// 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]) {
func (s *Server) AddReceivingMiddleware(middleware ...Middleware) {
s.mu.Lock()
defer s.mu.Unlock()
addMiddleware(&s.receivingMethodHandler_, middleware)
@ -734,22 +747,22 @@ func (s *Server) AddReceivingMiddleware(middleware ...Middleware[*ServerSession]
// TODO(rfindley): actually load and validate the protocol schema, rather than
// curating these method flags.
var serverMethodInfos = map[string]methodInfo{
methodComplete: newMethodInfo(serverMethod((*Server).complete), 0),
methodInitialize: newMethodInfo(sessionMethod((*ServerSession).initialize), 0),
methodPing: newMethodInfo(sessionMethod((*ServerSession).ping), missingParamsOK),
methodListPrompts: newMethodInfo(serverMethod((*Server).listPrompts), missingParamsOK),
methodGetPrompt: newMethodInfo(serverMethod((*Server).getPrompt), 0),
methodListTools: newMethodInfo(serverMethod((*Server).listTools), missingParamsOK),
methodCallTool: newMethodInfo(serverMethod((*Server).callTool), 0),
methodListResources: newMethodInfo(serverMethod((*Server).listResources), missingParamsOK),
methodListResourceTemplates: newMethodInfo(serverMethod((*Server).listResourceTemplates), missingParamsOK),
methodReadResource: newMethodInfo(serverMethod((*Server).readResource), 0),
methodSetLevel: newMethodInfo(sessionMethod((*ServerSession).setLevel), 0),
methodSubscribe: newMethodInfo(serverMethod((*Server).subscribe), 0),
methodUnsubscribe: newMethodInfo(serverMethod((*Server).unsubscribe), 0),
notificationInitialized: newMethodInfo(sessionMethod((*ServerSession).initialized), notification|missingParamsOK),
notificationRootsListChanged: newMethodInfo(serverMethod((*Server).callRootsListChangedHandler), notification|missingParamsOK),
notificationProgress: newMethodInfo(sessionMethod((*ServerSession).callProgressNotificationHandler), notification),
methodComplete: newServerMethodInfo(serverMethod((*Server).complete), 0),
methodInitialize: newServerMethodInfo(serverSessionMethod((*ServerSession).initialize), 0),
methodPing: newServerMethodInfo(serverSessionMethod((*ServerSession).ping), missingParamsOK),
methodListPrompts: newServerMethodInfo(serverMethod((*Server).listPrompts), missingParamsOK),
methodGetPrompt: newServerMethodInfo(serverMethod((*Server).getPrompt), 0),
methodListTools: newServerMethodInfo(serverMethod((*Server).listTools), missingParamsOK),
methodCallTool: newServerMethodInfo(serverMethod((*Server).callTool), 0),
methodListResources: newServerMethodInfo(serverMethod((*Server).listResources), missingParamsOK),
methodListResourceTemplates: newServerMethodInfo(serverMethod((*Server).listResourceTemplates), missingParamsOK),
methodReadResource: newServerMethodInfo(serverMethod((*Server).readResource), 0),
methodSetLevel: newServerMethodInfo(serverSessionMethod((*ServerSession).setLevel), 0),
methodSubscribe: newServerMethodInfo(serverMethod((*Server).subscribe), 0),
methodUnsubscribe: newServerMethodInfo(serverMethod((*Server).unsubscribe), 0),
notificationInitialized: newServerMethodInfo(serverSessionMethod((*ServerSession).initialized), notification|missingParamsOK),
notificationRootsListChanged: newServerMethodInfo(serverMethod((*Server).callRootsListChangedHandler), notification|missingParamsOK),
notificationProgress: newServerMethodInfo(serverSessionMethod((*ServerSession).callProgressNotificationHandler), notification),
}
func (ss *ServerSession) sendingMethodInfos() map[string]methodInfo { return clientMethodInfos }
@ -757,15 +770,17 @@ func (ss *ServerSession) sendingMethodInfos() map[string]methodInfo { return cli
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.sendingMethodHandler_
s := ss.server
s.mu.Lock()
defer s.mu.Unlock()
return s.sendingMethodHandler_
}
func (ss *ServerSession) receivingMethodHandler() methodHandler {
ss.server.mu.Lock()
defer ss.server.mu.Unlock()
return ss.server.receivingMethodHandler_
s := ss.server
s.mu.Lock()
defer s.mu.Unlock()
return s.receivingMethodHandler_
}
// getConn implements [session.getConn].
@ -808,13 +823,14 @@ func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParam
version = latestProtocolVersion
}
s := ss.server
return &InitializeResult{
// TODO(rfindley): alter behavior when falling back to an older version:
// reject unsupported features.
ProtocolVersion: version,
Capabilities: ss.server.capabilities(),
Instructions: ss.server.opts.Instructions,
ServerInfo: ss.server.impl,
Capabilities: s.capabilities(),
Instructions: s.opts.Instructions,
ServerInfo: s.impl,
}, nil
}

4
mcp/server_example_test.go
View file

@ -16,10 +16,10 @@ type SayHiParams struct {
Name string `json:"name"`
}
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[SayHiParams]) (*mcp.CallToolResultFor[any], error) {
func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[SayHiParams]]) (*mcp.CallToolResultFor[any], error) {
return &mcp.CallToolResultFor[any]{
Content: []mcp.Content{
&mcp.TextContent{Text: "Hi " + params.Arguments.Name},
&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name},
},
}, nil
}

12
mcp/server_test.go
View file

@ -281,10 +281,10 @@ func TestServerCapabilities(t *testing.T) {
s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil)
},
serverOpts: ServerOptions{
SubscribeHandler: func(ctx context.Context, _ *ServerSession, sp *SubscribeParams) error {
SubscribeHandler: func(context.Context, *ServerRequest[*SubscribeParams]) error {
return nil
},
UnsubscribeHandler: func(ctx context.Context, _ *ServerSession, up *UnsubscribeParams) error {
UnsubscribeHandler: func(context.Context, *ServerRequest[*UnsubscribeParams]) error {
return nil
},
},
@ -307,7 +307,7 @@ func TestServerCapabilities(t *testing.T) {
name: "With completions",
configureServer: func(s *Server) {},
serverOpts: ServerOptions{
CompletionHandler: func(ctx context.Context, ss *ServerSession, params *CompleteParams) (*CompleteResult, error) {
CompletionHandler: func(context.Context, *ServerRequest[*CompleteParams]) (*CompleteResult, error) {
return nil, nil
},
},
@ -325,13 +325,13 @@ func TestServerCapabilities(t *testing.T) {
s.AddTool(tool, nil)
},
serverOpts: ServerOptions{
SubscribeHandler: func(ctx context.Context, _ *ServerSession, sp *SubscribeParams) error {
SubscribeHandler: func(context.Context, *ServerRequest[*SubscribeParams]) error {
return nil
},
UnsubscribeHandler: func(ctx context.Context, _ *ServerSession, up *UnsubscribeParams) error {
UnsubscribeHandler: func(context.Context, *ServerRequest[*UnsubscribeParams]) error {
return nil
},
CompletionHandler: func(ctx context.Context, ss *ServerSession, params *CompleteParams) (*CompleteResult, error) {
CompletionHandler: func(context.Context, *ServerRequest[*CompleteParams]) (*CompleteResult, error) {
return nil, nil
},
},

185
mcp/shared.go
View file

@ -36,7 +36,7 @@ var supportedProtocolVersions = []string{
// A MethodHandler handles MCP messages.
// For methods, exactly one of the return values must be nil.
// For notifications, both must be nil.
type MethodHandler[S Session] func(ctx context.Context, _ S, method string, params Params) (result Result, err error)
type MethodHandler func(ctx context.Context, method string, req Request) (result Result, err error)
// A methodHandler is a MethodHandler[Session] for some session.
// We need to give up type safety here, or we will end up with a type cycle somewhere
@ -46,7 +46,6 @@ type methodHandler any // MethodHandler[*ClientSession] | MethodHandler[*ServerS
// A Session is either a [ClientSession] or a [ServerSession].
type Session interface {
*ClientSession | *ServerSession
// ID returns the session ID, or the empty string if there is none.
ID() string
@ -58,29 +57,29 @@ type Session interface {
}
// Middleware is a function from [MethodHandler] to [MethodHandler].
type Middleware[S Session] func(MethodHandler[S]) MethodHandler[S]
type Middleware func(MethodHandler) MethodHandler
// addMiddleware wraps the handler in the middleware functions.
func addMiddleware[S Session](handlerp *MethodHandler[S], middleware []Middleware[S]) {
func addMiddleware(handlerp *MethodHandler, middleware []Middleware) {
for _, m := range slices.Backward(middleware) {
*handlerp = m(*handlerp)
}
}
func defaultSendingMethodHandler[S Session](ctx context.Context, session S, method string, params Params) (Result, error) {
info, ok := session.sendingMethodInfos()[method]
func defaultSendingMethodHandler[S Session](ctx context.Context, method string, req Request) (Result, error) {
info, ok := req.GetSession().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)
return nil, req.GetSession().getConn().Notify(ctx, method, req.GetParams())
}
// 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 {
if err := call(ctx, req.GetSession().getConn(), method, req.GetParams(), res); err != nil {
return nil, err
}
return res, nil
@ -95,16 +94,16 @@ func orZero[T any, P *U, U any](p P) T {
return any(p).(T)
}
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)
func handleNotify(ctx context.Context, method string, req Request) error {
mh := req.GetSession().sendingMethodHandler().(MethodHandler)
_, err := mh(ctx, method, req)
return err
}
func handleSend[R Result, S Session](ctx context.Context, s S, method string, params Params) (R, error) {
mh := s.sendingMethodHandler().(MethodHandler[S])
func handleSend[R Result](ctx context.Context, method string, req Request) (R, error) {
mh := req.GetSession().sendingMethodHandler().(MethodHandler)
// mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol.
res, err := mh(ctx, s, method, params)
res, err := mh(ctx, method, req)
if err != nil {
var z R
return z, err
@ -113,28 +112,29 @@ func handleSend[R Result, S Session](ctx context.Context, s S, method string, pa
}
// 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]
func defaultReceivingMethodHandler[S Session](ctx context.Context, method string, req Request) (Result, error) {
info, ok := req.GetSession().receivingMethodInfos()[method]
if !ok {
// This can be called from user code, with an arbitrary value for method.
return nil, jsonrpc2.ErrNotHandled
}
return info.handleMethod.(MethodHandler[S])(ctx, session, method, params)
return info.handleMethod.(MethodHandler)(ctx, method, req)
}
func handleReceive[S Session](ctx context.Context, session S, req *jsonrpc.Request) (Result, error) {
info, err := checkRequest(req, session.receivingMethodInfos())
func handleReceive[S Session](ctx context.Context, session S, jreq *jsonrpc.Request) (Result, error) {
info, err := checkRequest(jreq, session.receivingMethodInfos())
if err != nil {
return nil, err
}
params, err := info.unmarshalParams(req.Params)
params, err := info.unmarshalParams(jreq.Params)
if err != nil {
return nil, fmt.Errorf("handling '%s': %w", req.Method, err)
return nil, fmt.Errorf("handling '%s': %w", jreq.Method, err)
}
mh := session.receivingMethodHandler().(MethodHandler[S])
mh := session.receivingMethodHandler().(MethodHandler)
req := info.newRequest(session, params)
// mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol.
res, err := mh(ctx, session, req.Method, params)
res, err := mh(ctx, jreq.Method, req)
if err != nil {
return nil, err
}
@ -179,6 +179,7 @@ type methodInfo struct {
// Unmarshal params from the wire into a Params struct.
// Used on the receive side.
unmarshalParams func(json.RawMessage) (Params, error)
newRequest func(Session, Params) Request
// Run the code when a call to the method is received.
// Used on the receive side.
handleMethod methodHandler
@ -194,7 +195,10 @@ type methodInfo struct {
// - R: results
// A typedMethodHandler is like a MethodHandler, but with type information.
type typedMethodHandler[S Session, P Params, R Result] func(context.Context, S, P) (R, error)
type (
typedClientMethodHandler[P Params, R Result] func(context.Context, *ClientRequest[P]) (R, error)
typedServerMethodHandler[P Params, R Result] func(context.Context, *ServerRequest[P]) (R, error)
)
type paramsPtr[T any] interface {
*T
@ -208,11 +212,45 @@ const (
missingParamsOK // params may be missing or null
)
func newClientMethodInfo[P paramsPtr[T], R Result, T any](d typedClientMethodHandler[P, R], flags methodFlags) methodInfo {
mi := newMethodInfo[P, R](flags)
mi.newRequest = func(s Session, p Params) Request {
r := &ClientRequest[P]{Session: s.(*ClientSession)}
if p != nil {
r.Params = p.(P)
}
return r
}
mi.handleMethod = MethodHandler(func(ctx context.Context, _ string, req Request) (Result, error) {
return d(ctx, req.(*ClientRequest[P]))
})
return mi
}
func newServerMethodInfo[P paramsPtr[T], R Result, T any](d typedServerMethodHandler[P, R], flags methodFlags) methodInfo {
mi := newMethodInfo[P, R](flags)
mi.newRequest = func(s Session, p Params) Request {
r := &ServerRequest[P]{Session: s.(*ServerSession)}
if p != nil {
r.Params = p.(P)
}
return r
}
mi.handleMethod = MethodHandler(func(ctx context.Context, _ string, req Request) (Result, error) {
rf := &ServerRequest[P]{Session: req.GetSession().(*ServerSession)}
if req.GetParams() != nil {
rf.Params = req.GetParams().(P)
}
return d(ctx, rf)
})
return mi
}
// newMethodInfo creates a methodInfo from a typedMethodHandler.
//
// If isRequest is set, the method is treated as a request rather than a
// notification.
func newMethodInfo[S Session, P paramsPtr[T], R Result, T any](d typedMethodHandler[S, P, R], flags methodFlags) methodInfo {
func newMethodInfo[P paramsPtr[T], R Result, T any](flags methodFlags) methodInfo {
return methodInfo{
flags: flags,
unmarshalParams: func(m json.RawMessage) (Params, error) {
@ -234,12 +272,6 @@ func newMethodInfo[S Session, P paramsPtr[T], R Result, T any](d typedMethodHand
}
return orZero[Params](p), nil
},
handleMethod: MethodHandler[S](func(ctx context.Context, session S, _ string, params Params) (Result, error) {
if params == nil {
return d(ctx, session, nil)
}
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
@ -250,26 +282,33 @@ func newMethodInfo[S Session, P paramsPtr[T], R Result, T any](d typedMethodHand
// serverMethod is glue for creating a typedMethodHandler from a method on Server.
func serverMethod[P Params, R Result](
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)
f func(*Server, context.Context, *ServerRequest[P]) (R, error),
) typedServerMethodHandler[P, R] {
return func(ctx context.Context, req *ServerRequest[P]) (R, error) {
return f(req.Session.server, ctx, req)
}
}
// clientMethod is glue for creating a typedMethodHandler from a method on Client.
func clientMethod[P Params, R Result](
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)
f func(*Client, context.Context, *ClientRequest[P]) (R, error),
) typedClientMethodHandler[P, R] {
return func(ctx context.Context, req *ClientRequest[P]) (R, error) {
return f(req.Session.client, ctx, req)
}
}
// sessionMethod is glue for creating a typedMethodHandler from a method on ServerSession.
func sessionMethod[S Session, P Params, R Result](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)
// serverSessionMethod is glue for creating a typedServerMethodHandler from a method on ServerSession.
func serverSessionMethod[P Params, R Result](f func(*ServerSession, context.Context, P) (R, error)) typedServerMethodHandler[P, R] {
return func(ctx context.Context, req *ServerRequest[P]) (R, error) {
return f(req.GetSession().(*ServerSession), ctx, req.Params)
}
}
// clientSessionMethod is glue for creating a typedMethodHandler from a method on ServerSession.
func clientSessionMethod[P Params, R Result](f func(*ClientSession, context.Context, P) (R, error)) typedClientMethodHandler[P, R] {
return func(ctx context.Context, req *ClientRequest[P]) (R, error) {
return f(req.GetSession().(*ClientSession), ctx, req.Params)
}
}
@ -282,16 +321,9 @@ const (
CodeUnsupportedMethod = -31001
)
func callNotificationHandler[S Session, P any](ctx context.Context, h func(context.Context, S, *P), sess S, params *P) (Result, error) {
if h != nil {
h(ctx, sess, params)
}
return nil, nil
}
// notifySessions calls Notify on all the sessions.
// Should be called on a copy of the peer sessions.
func notifySessions[S Session](sessions []S, method string, params Params) {
func notifySessions[S Session, P Params](sessions []S, method string, params P) {
if sessions == nil {
return
}
@ -299,13 +331,25 @@ 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 := handleNotify(ctx, s, method, params); err != nil {
req := newRequest(s, params)
if err := handleNotify(ctx, method, req); err != nil {
// TODO(jba): surface this error better
log.Printf("calling %s: %v", method, err)
}
}
}
func newRequest[S Session, P Params](s S, p P) Request {
switch s := any(s).(type) {
case *ClientSession:
return &ClientRequest[P]{Session: s, Params: p}
case *ServerSession:
return &ServerRequest[P]{Session: s, Params: p}
default:
panic("bad session")
}
}
// Meta is additional metadata for requests, responses and other types.
type Meta map[string]any
@ -335,6 +379,43 @@ func setProgressToken(p Params, pt any) {
m[progressTokenKey] = pt
}
// A Request is a method request with parameters and additional information, such as the session.
// Request is implemented by [*ClientRequest] and [*ServerRequest].
type Request interface {
isRequest()
GetSession() Session
GetParams() Params
}
// A ClientRequest is a request to a client.
type ClientRequest[P Params] struct {
Session *ClientSession
Params P
}
// A ServerRequest is a request to a server.
type ServerRequest[P Params] struct {
Session *ServerSession
Params P
}
func (*ClientRequest[P]) isRequest() {}
func (*ServerRequest[P]) isRequest() {}
func (r *ClientRequest[P]) GetSession() Session { return r.Session }
func (r *ServerRequest[P]) GetSession() Session { return r.Session }
func (r *ClientRequest[P]) GetParams() Params { return r.Params }
func (r *ServerRequest[P]) GetParams() Params { return r.Params }
func serverRequestFor[P Params](s *ServerSession, p P) *ServerRequest[P] {
return &ServerRequest[P]{Session: s, Params: p}
}
func clientRequestFor[P Params](s *ClientSession, p P) *ClientRequest[P] {
return &ClientRequest[P]{Session: s, Params: p}
}
// Params is a parameter (input) type for an MCP call or notification.
type Params interface {
// isParams discourages implementation of Params outside of this package.

19
mcp/shared_test.go
View file

@ -23,7 +23,7 @@ func TestToolValidate(t *testing.T) {
P *int `json:",omitempty"`
}
dummyHandler := func(context.Context, *ServerSession, *CallToolParamsFor[req]) (*CallToolResultFor[any], error) {
dummyHandler := func(context.Context, *ServerRequest[*CallToolParamsFor[req]]) (*CallToolResultFor[any], error) {
return nil, nil
}
@ -73,8 +73,9 @@ func TestToolValidate(t *testing.T) {
if err != nil {
t.Fatal(err)
}
_, err = st.handler(context.Background(), nil,
&CallToolParamsFor[json.RawMessage]{Arguments: json.RawMessage(raw)})
_, err = st.handler(context.Background(), &ServerRequest[*CallToolParamsFor[json.RawMessage]]{
Params: &CallToolParamsFor[json.RawMessage]{Arguments: json.RawMessage(raw)},
})
if err == nil && tt.want != "" {
t.Error("got success, wanted failure")
}
@ -102,12 +103,12 @@ func TestNilParamsHandling(t *testing.T) {
type TestResult = *CallToolResultFor[string]
// Simple test handler
testHandler := func(ctx context.Context, ss *ServerSession, params TestParams) (TestResult, error) {
result := "processed: " + params.Arguments.Name
testHandler := func(ctx context.Context, req *ServerRequest[TestParams]) (TestResult, error) {
result := "processed: " + req.Params.Arguments.Name
return &CallToolResultFor[string]{StructuredContent: result}, nil
}
methodInfo := newMethodInfo(testHandler, missingParamsOK)
methodInfo := newServerMethodInfo(testHandler, missingParamsOK)
// Helper function to test that unmarshalParams doesn't panic and handles nil gracefully
mustNotPanic := func(t *testing.T, rawMsg json.RawMessage, expectNil bool) Params {
@ -183,11 +184,11 @@ func TestNilParamsEdgeCases(t *testing.T) {
}
type TestParams = *CallToolParamsFor[TestArgs]
testHandler := func(ctx context.Context, ss *ServerSession, params TestParams) (*CallToolResultFor[string], error) {
testHandler := func(context.Context, *ServerRequest[TestParams]) (*CallToolResultFor[string], error) {
return &CallToolResultFor[string]{StructuredContent: "test"}, nil
}
methodInfo := newMethodInfo(testHandler, missingParamsOK)
methodInfo := newServerMethodInfo(testHandler, missingParamsOK)
// These should fail normally, not be treated as nil params
invalidCases := []json.RawMessage{
@ -209,7 +210,7 @@ func TestNilParamsEdgeCases(t *testing.T) {
// Test that methods without missingParamsOK flag properly reject nil params
t.Run("reject_when_params_required", func(t *testing.T) {
methodInfoStrict := newMethodInfo(testHandler, 0) // No missingParamsOK flag
methodInfoStrict := newServerMethodInfo(testHandler, 0) // No missingParamsOK flag
testCases := []struct {
name string

4
mcp/sse_example_test.go
View file

@ -18,10 +18,10 @@ type AddParams struct {
X, Y int
}
func Add(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[AddParams]) (*mcp.CallToolResultFor[any], error) {
func Add(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[AddParams]]) (*mcp.CallToolResultFor[any], error) {
return &mcp.CallToolResultFor[any]{
Content: []mcp.Content{
&mcp.TextContent{Text: fmt.Sprintf("%d", params.Arguments.X+params.Arguments.Y)},
&mcp.TextContent{Text: fmt.Sprintf("%d", req.Params.Arguments.X+req.Params.Arguments.Y)},
},
}, nil
}

32
mcp/streamable_test.go
View file

@ -136,12 +136,12 @@ func TestClientReplay(t *testing.T) {
serverReadyToKillProxy := make(chan struct{})
serverClosed := make(chan struct{})
server.AddTool(&Tool{Name: "multiMessageTool", InputSchema: &jsonschema.Schema{}},
func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResult, error) {
go func() {
bgCtx := context.Background()
// Send the first two messages immediately.
ss.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg1"})
ss.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg2"})
req.Session.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg1"})
req.Session.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg2"})
// Signal the test that it can now kill the proxy.
close(serverReadyToKillProxy)
@ -149,8 +149,8 @@ func TestClientReplay(t *testing.T) {
// These messages should be queued for replay by the server after
// the client's connection drops.
ss.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg3"})
ss.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg4"})
req.Session.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg3"})
req.Session.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg4"})
}()
return &CallToolResult{}, nil
})
@ -170,8 +170,8 @@ func TestClientReplay(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
client := NewClient(testImpl, &ClientOptions{
ProgressNotificationHandler: func(ctx context.Context, cc *ClientSession, params *ProgressNotificationParams) {
notifications <- params.Message
ProgressNotificationHandler: func(ctx context.Context, req *ClientRequest[*ProgressNotificationParams]) {
notifications <- req.Params.Message
},
})
clientSession, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: proxy.URL}, nil)
@ -236,9 +236,10 @@ func TestServerInitiatedSSE(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
client := NewClient(testImpl, &ClientOptions{ToolListChangedHandler: func(ctx context.Context, cc *ClientSession, params *ToolListChangedParams) {
notifications <- "toolListChanged"
},
client := NewClient(testImpl, &ClientOptions{
ToolListChangedHandler: func(context.Context, *ClientRequest[*ToolListChangedParams]) {
notifications <- "toolListChanged"
},
})
clientSession, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil)
if err != nil {
@ -246,7 +247,7 @@ func TestServerInitiatedSSE(t *testing.T) {
}
defer clientSession.Close()
server.AddTool(&Tool{Name: "testTool", InputSchema: &jsonschema.Schema{}},
func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
func(context.Context, *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResult, error) {
return &CallToolResult{}, nil
})
receivedNotifications := readNotifications(t, ctx, notifications, 1)
@ -509,9 +510,9 @@ func TestStreamableServerTransport(t *testing.T) {
// Create a server containing a single tool, which runs the test tool
// behavior, if any.
server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil)
AddTool(server, &Tool{Name: "tool"}, func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[any]) (*CallToolResultFor[any], error) {
AddTool(server, &Tool{Name: "tool"}, func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[any]]) (*CallToolResultFor[any], error) {
if test.tool != nil {
test.tool(t, ctx, ss)
test.tool(t, ctx, req.Session)
}
return &CallToolResultFor[any]{}, nil
})
@ -822,14 +823,15 @@ func TestEventID(t *testing.T) {
})
}
}
func TestStreamableStateless(t *testing.T) {
// Test stateless mode behavior
ctx := context.Background()
// This version of sayHi doesn't make a ping request (we can't respond to
// that request from our client).
sayHi := func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[hiParams]) (*CallToolResultFor[any], error) {
return &CallToolResultFor[any]{Content: []Content{&TextContent{Text: "hi " + params.Arguments.Name}}}, nil
sayHi := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[hiParams]]) (*CallToolResultFor[any], error) {
return &CallToolResultFor[any]{Content: []Content{&TextContent{Text: "hi " + req.Params.Arguments.Name}}}, nil
}
server := NewServer(testImpl, nil)
AddTool(server, &Tool{Name: "greet", Description: "say hi"}, sayHi)

23
mcp/tool.go
View file

@ -20,10 +20,12 @@ import (
type ToolHandler = ToolHandlerFor[map[string]any, any]
// A ToolHandlerFor handles a call to tools/call with typed arguments and results.
type ToolHandlerFor[In, Out any] func(context.Context, *ServerSession, *CallToolParamsFor[In]) (*CallToolResultFor[Out], error)
type ToolHandlerFor[In, Out any] func(context.Context, *ServerRequest[*CallToolParamsFor[In]]) (*CallToolResultFor[Out], error)
// A rawToolHandler is like a ToolHandler, but takes the arguments as as json.RawMessage.
type rawToolHandler = func(context.Context, *ServerSession, *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error)
// Second arg is *Request[*ServerSession, *CallToolParamsFor[json.RawMessage]], but that creates
// a cycle.
type rawToolHandler = func(context.Context, any) (*CallToolResult, error)
// A serverTool is a tool definition that is bound to a tool handler.
type serverTool struct {
@ -48,20 +50,25 @@ func newServerTool[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*serverTool
}
}
st.handler = func(ctx context.Context, ss *ServerSession, rparams *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) {
st.handler = func(ctx context.Context, areq any) (*CallToolResult, error) {
req := areq.(*ServerRequest[*CallToolParamsFor[json.RawMessage]])
var args In
if rparams.Arguments != nil {
if err := unmarshalSchema(rparams.Arguments, st.inputResolved, &args); err != nil {
if req.Params.Arguments != nil {
if err := unmarshalSchema(req.Params.Arguments, st.inputResolved, &args); err != nil {
return nil, err
}
}
// TODO(jba): future-proof this copy.
params := &CallToolParamsFor[In]{
Meta: rparams.Meta,
Name: rparams.Name,
Meta: req.Params.Meta,
Name: req.Params.Name,
Arguments: args,
}
res, err := h(ctx, ss, params)
// TODO(jba): improve copy
res, err := h(ctx, &ServerRequest[*CallToolParamsFor[In]]{
Session: req.Session,
Params: params,
})
// TODO(rfindley): investigate why server errors are embedded in this strange way,
// rather than returned as jsonrpc2 server errors.
if err != nil {

2
mcp/tool_test.go
View file

@ -16,7 +16,7 @@ import (
)
// testToolHandler is used for type inference in TestNewServerTool.
func testToolHandler[In, Out any](context.Context, *ServerSession, *CallToolParamsFor[In]) (*CallToolResultFor[Out], error) {
func testToolHandler[In, Out any](context.Context, *ServerRequest[*CallToolParamsFor[In]]) (*CallToolResultFor[Out], error) {
panic("not implemented")
}