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
25 changed files with 517 additions and 380 deletions
Download patch file Download diff file Expand all files Collapse all files

4
README.md
View file

@ -115,9 +115,9 @@ type HiParams struct {
Name string `json:"name" jsonschema:"the name of the person to greet"` Name string `json:"name" jsonschema:"the name of the person to greet"`
} }
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[HiParams]) (*mcp.CallToolResultFor[any], error) { func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[HiParams]]) (*mcp.CallToolResultFor[any], error) {
return &mcp.CallToolResultFor[any]{ return &mcp.CallToolResultFor[any]{
Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + params.Arguments.Name}}, Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}},
}, nil }, nil
} }

53
design/design.md
View file

@ -408,15 +408,12 @@ We provide a mechanism to add MCP-level middleware on the both the client and se
```go ```go
// A MethodHandler handles MCP messages. // A MethodHandler handles MCP messages.
// The params argument is an XXXParams struct pointer, such as *GetPromptParams. // For methods, exactly one of the return values must be nil.
// For methods, a MethodHandler must return either an XXResult struct pointer and a nil error, or // For notifications, both must be nil.
// nil with a non-nil error. type MethodHandler func(ctx context.Context, method string, req Request) (result Result, err error)
// For notifications, a MethodHandler must return nil, nil.
type MethodHandler[S Session] func(
ctx context.Context, _ *S, method string, params Params) (result Result, err error)
// Middleware is a function from MethodHandlers to MethodHandlers. // Middleware is a function from MethodHandlers to MethodHandlers.
type Middleware[S Session] func(MethodHandler[S]) MethodHandler[S] type Middleware func(MethodHandler) MethodHandler
// AddMiddleware wraps the client/server's current method handler using the provided // AddMiddleware wraps the client/server's current method handler using the provided
// middleware. Middleware is applied from right to left, so that the first one // middleware. Middleware is applied from right to left, so that the first one
@ -424,17 +421,17 @@ type Middleware[S Session] func(MethodHandler[S]) MethodHandler[S]
// //
// For example, AddMiddleware(m1, m2, m3) augments the server method handler as // For example, AddMiddleware(m1, m2, m3) augments the server method handler as
// m1(m2(m3(handler))). // m1(m2(m3(handler))).
func (c *Client) AddSendingMiddleware(middleware ...Middleware[*ClientSession]) func (c *Client) AddSendingMiddleware(middleware ...Middleware)
func (c *Client) AddReceivingMiddleware(middleware ...Middleware[*ClientSession]) func (c *Client) AddReceivingMiddleware(middleware ...Middleware)
func (s *Server) AddSendingMiddleware(middleware ...Middleware[*ServerSession]) func (s *Server) AddSendingMiddleware(middleware ...Middleware)
func (s *Server) AddReceivingMiddleware(middleware ...Middleware[*ServerSession]) func (s *Server) AddReceivingMiddleware(middleware ...Middleware)
``` ```
As an example, this code adds server-side logging: As an example, this code adds server-side logging:
```go ```go
func withLogging(h mcp.MethodHandler[*mcp.ServerSession]) mcp.MethodHandler[*mcp.ServerSession]{ func withLogging(h mcp.MethodHandler) mcp.MethodHandler{
return func(ctx context.Context, s *mcp.ServerSession, method string, params mcp.Params) (res mcp.Result, err error) { return func(ctx context.Context, method string, req mcp.Request) (res mcp.Result, err error) {
log.Printf("request: %s %v", method, params) log.Printf("request: %s %v", method, params)
defer func() { log.Printf("response: %v, %v", res, err) }() defer func() { log.Printf("response: %v, %v", res, err) }()
return h(ctx, s , method, params) return h(ctx, s , method, params)
@ -597,8 +594,9 @@ type Tool struct {
Name string `json:"name"` Name string `json:"name"`
} }
type ToolHandlerFor[In, Out any] func(context.Context, *ServerSession, *CallToolParamsFor[In]) (*CallToolResultFor[Out], error) // A ToolHandlerFor handles a call to tools/call with typed arguments and results.
type ToolHandler = ToolHandlerFor[map[string]any, any] type ToolHandlerFor[In, Out any] func(context.Context, *ServerRequest[*CallToolParamsFor[In]]) (*CallToolResultFor[Out], error)
``` ```
Add tools to a server with the `AddTool` method or function. The function is generic and infers schemas from the handler Add tools to a server with the `AddTool` method or function. The function is generic and infers schemas from the handler
@ -648,8 +646,8 @@ type AddParams struct {
Y int `json:"y"` Y int `json:"y"`
} }
func addHandler(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[AddParams]) (*mcp.CallToolResultFor[int], error) { func addHandler(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[AddParams]]) (*mcp.CallToolResultFor[int], error) {
return &mcp.CallToolResultFor[int]{StructuredContent: params.Arguments.X + params.Arguments.Y}, nil return &mcp.CallToolResultFor[int]{StructuredContent: req.Params.Arguments.X + req.Params.Arguments.Y}, nil
} }
``` ```
@ -665,8 +663,6 @@ Client sessions can call the spec method `ListTools` or an iterator method `Tool
```go ```go
func (cs *ClientSession) CallTool(context.Context, *CallToolParams[json.RawMessage]) (*CallToolResult, error) func (cs *ClientSession) CallTool(context.Context, *CallToolParams[json.RawMessage]) (*CallToolResult, error)
func CallTool[TArgs any](context.Context, *ClientSession, *CallToolParams[TArgs]) (*CallToolResult, error)
``` ```
**Differences from mcp-go**: We provide a full JSON Schema implementation for validating tool input schemas against incoming arguments. The `jsonschema.Schema` type provides exported features for all keywords in the JSON Schema draft2020-12 spec. Tool definers can use it to construct any schema they want. The `jsonschema.For[T]` function can infer a schema from a Go struct. These combined features eliminate the need for variadic arguments to construct tool schemas. **Differences from mcp-go**: We provide a full JSON Schema implementation for validating tool input schemas against incoming arguments. The `jsonschema.Schema` type provides exported features for all keywords in the JSON Schema draft2020-12 spec. Tool definers can use it to construct any schema they want. The `jsonschema.For[T]` function can infer a schema from a Go struct. These combined features eliminate the need for variadic arguments to construct tool schemas.
@ -752,10 +748,10 @@ If a server author wants to support resource subscriptions, they must provide ha
```go ```go
type ServerOptions struct { type ServerOptions struct {
... ...
// Function called when a client session subscribes to a resource. // Function called when a client session subscribes to a resource.
SubscribeHandler func(context.Context, ss *ServerSession, *SubscribeParams) error SubscribeHandler func(context.Context, *ServerRequest[*SubscribeParams]) error
// Function called when a client session unsubscribes from a resource. // Function called when a client session unsubscribes from a resource.
UnsubscribeHandler func(context.Context, ss *ServerSession, *UnsubscribeParams) error UnsubscribeHandler func(context.Context, *ServerRequest[*UnsubscribeParams]) error
} }
``` ```
@ -774,10 +770,10 @@ When a list of tools, prompts or resources changes as the result of an AddXXX or
```go ```go
type ClientOptions struct { type ClientOptions struct {
... ...
ToolListChangedHandler func(context.Context, *ClientSession, *ToolListChangedParams) ToolListChangedHandler func(context.Context, *ClientRequest[*ToolListChangedParams])
PromptListChangedHandler func(context.Context, *ClientSession, *PromptListChangedParams) PromptListChangedHandler func(context.Context, *ClientRequest[*PromptListChangedParams])
// For both resources and resource templates. // For both resources and resource templates.
ResourceListChangedHandler func(context.Context, *ClientSession, *ResourceListChangedParams) ResourceListChangedHandler func(context.Context, *ClientRequest[*ResourceListChangedParams])
} }
``` ```
@ -788,13 +784,10 @@ type ClientOptions struct {
Clients call the spec method `Complete` to request completions. If a server installs a `CompletionHandler`, it will be called when the client sends a completion request. Clients call the spec method `Complete` to request completions. If a server installs a `CompletionHandler`, it will be called when the client sends a completion request.
```go ```go
// A CompletionHandler handles a call to completion/complete.
type CompletionHandler func(context.Context, *ServerSession, *CompleteParams) (*CompleteResult, error)
type ServerOptions struct { type ServerOptions struct {
... ...
// If non-nil, called when a client sends a completion request. // If non-nil, called when a client sends a completion request.
CompletionHandler CompletionHandler CompletionHandler func(context.Context, *ServerRequest[*CompleteParams]) (*CompleteResult, error)
} }
``` ```

6
examples/server/completion/main.go
View file

@ -16,17 +16,17 @@ import (
// a CompletionHandler to an MCP Server's options. // a CompletionHandler to an MCP Server's options.
func main() { func main() {
// Define your custom CompletionHandler logic. // Define your custom CompletionHandler logic.
myCompletionHandler := func(_ context.Context, _ *mcp.ServerSession, params *mcp.CompleteParams) (*mcp.CompleteResult, error) { myCompletionHandler := func(_ context.Context, req *mcp.ServerRequest[*mcp.CompleteParams]) (*mcp.CompleteResult, error) {
// In a real application, you'd implement actual completion logic here. // In a real application, you'd implement actual completion logic here.
// For this example, we return a fixed set of suggestions. // For this example, we return a fixed set of suggestions.
var suggestions []string var suggestions []string
switch params.Ref.Type { switch req.Params.Ref.Type {
case "ref/prompt": case "ref/prompt":
suggestions = []string{"suggestion1", "suggestion2", "suggestion3"} suggestions = []string{"suggestion1", "suggestion2", "suggestion3"}
case "ref/resource": case "ref/resource":
suggestions = []string{"suggestion4", "suggestion5", "suggestion6"} suggestions = []string{"suggestion4", "suggestion5", "suggestion6"}
default: default:
return nil, fmt.Errorf("unrecognized content type %s", params.Ref.Type) return nil, fmt.Errorf("unrecognized content type %s", req.Params.Ref.Type)
} }
return &mcp.CompleteResult{ return &mcp.CompleteResult{

4
examples/server/custom-transport/main.go
View file

@ -85,10 +85,10 @@ type HiArgs struct {
} }
// SayHi is a tool handler that responds with a greeting. // SayHi is a tool handler that responds with a greeting.
func SayHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[HiArgs]) (*mcp.CallToolResultFor[struct{}], error) { func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[HiArgs]]) (*mcp.CallToolResultFor[struct{}], error) {
return &mcp.CallToolResultFor[struct{}]{ return &mcp.CallToolResultFor[struct{}]{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{Text: "Hi " + params.Arguments.Name}, &mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name},
}, },
}, nil }, nil
} }

10
examples/server/hello/main.go
View file

@ -22,10 +22,10 @@ type HiArgs struct {
Name string `json:"name" jsonschema:"the name to say hi to"` Name string `json:"name" jsonschema:"the name to say hi to"`
} }
func SayHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[HiArgs]) (*mcp.CallToolResultFor[struct{}], error) { func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[HiArgs]]) (*mcp.CallToolResultFor[struct{}], error) {
return &mcp.CallToolResultFor[struct{}]{ return &mcp.CallToolResultFor[struct{}]{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{Text: "Hi " + params.Arguments.Name}, &mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name},
}, },
}, nil }, nil
} }
@ -69,8 +69,8 @@ var embeddedResources = map[string]string{
"info": "This is the hello example server.", "info": "This is the hello example server.",
} }
func handleEmbeddedResource(_ context.Context, _ *mcp.ServerSession, params *mcp.ReadResourceParams) (*mcp.ReadResourceResult, error) { func handleEmbeddedResource(_ context.Context, req *mcp.ServerRequest[*mcp.ReadResourceParams]) (*mcp.ReadResourceResult, error) {
u, err := url.Parse(params.URI) u, err := url.Parse(req.Params.URI)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -84,7 +84,7 @@ func handleEmbeddedResource(_ context.Context, _ *mcp.ServerSession, params *mcp
} }
return &mcp.ReadResourceResult{ return &mcp.ReadResourceResult{
Contents: []*mcp.ResourceContents{ Contents: []*mcp.ResourceContents{
{URI: params.URI, MIMEType: "text/plain", Text: text}, {URI: req.Params.URI, MIMEType: "text/plain", Text: text},
}, },
}, nil }, nil
} }

36
examples/server/memory/kb.go
View file

@ -84,7 +84,7 @@ func (fs *fileStore) Read() ([]byte, error) {
// Write saves data to file with 0600 permissions. // Write saves data to file with 0600 permissions.
func (fs *fileStore) Write(data []byte) error { func (fs *fileStore) Write(data []byte) error {
if err := os.WriteFile(fs.path, data, 0600); err != nil { if err := os.WriteFile(fs.path, data, 0o600); err != nil {
return fmt.Errorf("failed to write file %s: %w", fs.path, err) return fmt.Errorf("failed to write file %s: %w", fs.path, err)
} }
return nil return nil
@ -431,10 +431,10 @@ func (k knowledgeBase) openNodes(names []string) (KnowledgeGraph, error) {
}, nil }, nil
} }
func (k knowledgeBase) CreateEntities(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[CreateEntitiesArgs]) (*mcp.CallToolResultFor[CreateEntitiesResult], error) { func (k knowledgeBase) CreateEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[CreateEntitiesArgs]]) (*mcp.CallToolResultFor[CreateEntitiesResult], error) {
var res mcp.CallToolResultFor[CreateEntitiesResult] var res mcp.CallToolResultFor[CreateEntitiesResult]
entities, err := k.createEntities(params.Arguments.Entities) entities, err := k.createEntities(req.Params.Arguments.Entities)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -450,10 +450,10 @@ func (k knowledgeBase) CreateEntities(ctx context.Context, ss *mcp.ServerSession
return &res, nil return &res, nil
} }
func (k knowledgeBase) CreateRelations(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[CreateRelationsArgs]) (*mcp.CallToolResultFor[CreateRelationsResult], error) { func (k knowledgeBase) CreateRelations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[CreateRelationsArgs]]) (*mcp.CallToolResultFor[CreateRelationsResult], error) {
var res mcp.CallToolResultFor[CreateRelationsResult] var res mcp.CallToolResultFor[CreateRelationsResult]
relations, err := k.createRelations(params.Arguments.Relations) relations, err := k.createRelations(req.Params.Arguments.Relations)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -469,10 +469,10 @@ func (k knowledgeBase) CreateRelations(ctx context.Context, ss *mcp.ServerSessio
return &res, nil return &res, nil
} }
func (k knowledgeBase) AddObservations(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[AddObservationsArgs]) (*mcp.CallToolResultFor[AddObservationsResult], error) { func (k knowledgeBase) AddObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[AddObservationsArgs]]) (*mcp.CallToolResultFor[AddObservationsResult], error) {
var res mcp.CallToolResultFor[AddObservationsResult] var res mcp.CallToolResultFor[AddObservationsResult]
observations, err := k.addObservations(params.Arguments.Observations) observations, err := k.addObservations(req.Params.Arguments.Observations)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -488,10 +488,10 @@ func (k knowledgeBase) AddObservations(ctx context.Context, ss *mcp.ServerSessio
return &res, nil return &res, nil
} }
func (k knowledgeBase) DeleteEntities(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[DeleteEntitiesArgs]) (*mcp.CallToolResultFor[struct{}], error) { func (k knowledgeBase) DeleteEntities(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[DeleteEntitiesArgs]]) (*mcp.CallToolResultFor[struct{}], error) {
var res mcp.CallToolResultFor[struct{}] var res mcp.CallToolResultFor[struct{}]
err := k.deleteEntities(params.Arguments.EntityNames) err := k.deleteEntities(req.Params.Arguments.EntityNames)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -503,10 +503,10 @@ func (k knowledgeBase) DeleteEntities(ctx context.Context, ss *mcp.ServerSession
return &res, nil return &res, nil
} }
func (k knowledgeBase) DeleteObservations(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[DeleteObservationsArgs]) (*mcp.CallToolResultFor[struct{}], error) { func (k knowledgeBase) DeleteObservations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[DeleteObservationsArgs]]) (*mcp.CallToolResultFor[struct{}], error) {
var res mcp.CallToolResultFor[struct{}] var res mcp.CallToolResultFor[struct{}]
err := k.deleteObservations(params.Arguments.Deletions) err := k.deleteObservations(req.Params.Arguments.Deletions)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -518,10 +518,10 @@ func (k knowledgeBase) DeleteObservations(ctx context.Context, ss *mcp.ServerSes
return &res, nil return &res, nil
} }
func (k knowledgeBase) DeleteRelations(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[DeleteRelationsArgs]) (*mcp.CallToolResultFor[struct{}], error) { func (k knowledgeBase) DeleteRelations(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[DeleteRelationsArgs]]) (*mcp.CallToolResultFor[struct{}], error) {
var res mcp.CallToolResultFor[struct{}] var res mcp.CallToolResultFor[struct{}]
err := k.deleteRelations(params.Arguments.Relations) err := k.deleteRelations(req.Params.Arguments.Relations)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -533,7 +533,7 @@ func (k knowledgeBase) DeleteRelations(ctx context.Context, ss *mcp.ServerSessio
return &res, nil return &res, nil
} }
func (k knowledgeBase) ReadGraph(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[struct{}]) (*mcp.CallToolResultFor[KnowledgeGraph], error) { func (k knowledgeBase) ReadGraph(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[struct{}]]) (*mcp.CallToolResultFor[KnowledgeGraph], error) {
var res mcp.CallToolResultFor[KnowledgeGraph] var res mcp.CallToolResultFor[KnowledgeGraph]
graph, err := k.loadGraph() graph, err := k.loadGraph()
@ -549,10 +549,10 @@ func (k knowledgeBase) ReadGraph(ctx context.Context, ss *mcp.ServerSession, par
return &res, nil return &res, nil
} }
func (k knowledgeBase) SearchNodes(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[SearchNodesArgs]) (*mcp.CallToolResultFor[KnowledgeGraph], error) { func (k knowledgeBase) SearchNodes(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[SearchNodesArgs]]) (*mcp.CallToolResultFor[KnowledgeGraph], error) {
var res mcp.CallToolResultFor[KnowledgeGraph] var res mcp.CallToolResultFor[KnowledgeGraph]
graph, err := k.searchNodes(params.Arguments.Query) graph, err := k.searchNodes(req.Params.Arguments.Query)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -565,10 +565,10 @@ func (k knowledgeBase) SearchNodes(ctx context.Context, ss *mcp.ServerSession, p
return &res, nil return &res, nil
} }
func (k knowledgeBase) OpenNodes(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[OpenNodesArgs]) (*mcp.CallToolResultFor[KnowledgeGraph], error) { func (k knowledgeBase) OpenNodes(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[OpenNodesArgs]]) (*mcp.CallToolResultFor[KnowledgeGraph], error) {
var res mcp.CallToolResultFor[KnowledgeGraph] var res mcp.CallToolResultFor[KnowledgeGraph]
graph, err := k.openNodes(params.Arguments.Names) graph, err := k.openNodes(req.Params.Arguments.Names)
if err != nil { if err != nil {
return nil, err return nil, err
} }

30
examples/server/memory/kb_test.go
View file

@ -230,7 +230,7 @@ func TestSaveAndLoadGraph(t *testing.T) {
// Test malformed data handling // Test malformed data handling
if fs, ok := s.(*fileStore); ok { if fs, ok := s.(*fileStore); ok {
err := os.WriteFile(fs.path, []byte("invalid json"), 0600) err := os.WriteFile(fs.path, []byte("invalid json"), 0o600)
if err != nil { if err != nil {
t.Fatalf("failed to write invalid json: %v", err) t.Fatalf("failed to write invalid json: %v", err)
} }
@ -450,7 +450,7 @@ func TestMCPServerIntegration(t *testing.T) {
}, },
} }
createResult, err := kb.CreateEntities(ctx, serverSession, createEntitiesParams) createResult, err := kb.CreateEntities(ctx, requestFor(serverSession, createEntitiesParams))
if err != nil { if err != nil {
t.Fatalf("MCP CreateEntities failed: %v", err) t.Fatalf("MCP CreateEntities failed: %v", err)
} }
@ -463,7 +463,7 @@ func TestMCPServerIntegration(t *testing.T) {
// Test ReadGraph through MCP // Test ReadGraph through MCP
readParams := &mcp.CallToolParamsFor[struct{}]{} readParams := &mcp.CallToolParamsFor[struct{}]{}
readResult, err := kb.ReadGraph(ctx, serverSession, readParams) readResult, err := kb.ReadGraph(ctx, requestFor(serverSession, readParams))
if err != nil { if err != nil {
t.Fatalf("MCP ReadGraph failed: %v", err) t.Fatalf("MCP ReadGraph failed: %v", err)
} }
@ -487,7 +487,7 @@ func TestMCPServerIntegration(t *testing.T) {
}, },
} }
relationsResult, err := kb.CreateRelations(ctx, serverSession, createRelationsParams) relationsResult, err := kb.CreateRelations(ctx, requestFor(serverSession, createRelationsParams))
if err != nil { if err != nil {
t.Fatalf("MCP CreateRelations failed: %v", err) t.Fatalf("MCP CreateRelations failed: %v", err)
} }
@ -510,7 +510,7 @@ func TestMCPServerIntegration(t *testing.T) {
}, },
} }
obsResult, err := kb.AddObservations(ctx, serverSession, addObsParams) obsResult, err := kb.AddObservations(ctx, requestFor(serverSession, addObsParams))
if err != nil { if err != nil {
t.Fatalf("MCP AddObservations failed: %v", err) t.Fatalf("MCP AddObservations failed: %v", err)
} }
@ -528,7 +528,7 @@ func TestMCPServerIntegration(t *testing.T) {
}, },
} }
searchResult, err := kb.SearchNodes(ctx, serverSession, searchParams) searchResult, err := kb.SearchNodes(ctx, requestFor(serverSession, searchParams))
if err != nil { if err != nil {
t.Fatalf("MCP SearchNodes failed: %v", err) t.Fatalf("MCP SearchNodes failed: %v", err)
} }
@ -546,7 +546,7 @@ func TestMCPServerIntegration(t *testing.T) {
}, },
} }
openResult, err := kb.OpenNodes(ctx, serverSession, openParams) openResult, err := kb.OpenNodes(ctx, requestFor(serverSession, openParams))
if err != nil { if err != nil {
t.Fatalf("MCP OpenNodes failed: %v", err) t.Fatalf("MCP OpenNodes failed: %v", err)
} }
@ -569,7 +569,7 @@ func TestMCPServerIntegration(t *testing.T) {
}, },
} }
deleteObsResult, err := kb.DeleteObservations(ctx, serverSession, deleteObsParams) deleteObsResult, err := kb.DeleteObservations(ctx, requestFor(serverSession, deleteObsParams))
if err != nil { if err != nil {
t.Fatalf("MCP DeleteObservations failed: %v", err) t.Fatalf("MCP DeleteObservations failed: %v", err)
} }
@ -590,7 +590,7 @@ func TestMCPServerIntegration(t *testing.T) {
}, },
} }
deleteRelResult, err := kb.DeleteRelations(ctx, serverSession, deleteRelParams) deleteRelResult, err := kb.DeleteRelations(ctx, requestFor(serverSession, deleteRelParams))
if err != nil { if err != nil {
t.Fatalf("MCP DeleteRelations failed: %v", err) t.Fatalf("MCP DeleteRelations failed: %v", err)
} }
@ -605,7 +605,7 @@ func TestMCPServerIntegration(t *testing.T) {
}, },
} }
deleteEntResult, err := kb.DeleteEntities(ctx, serverSession, deleteEntParams) deleteEntResult, err := kb.DeleteEntities(ctx, requestFor(serverSession, deleteEntParams))
if err != nil { if err != nil {
t.Fatalf("MCP DeleteEntities failed: %v", err) t.Fatalf("MCP DeleteEntities failed: %v", err)
} }
@ -614,7 +614,7 @@ func TestMCPServerIntegration(t *testing.T) {
} }
// Verify final state // Verify final state
finalRead, err := kb.ReadGraph(ctx, serverSession, readParams) finalRead, err := kb.ReadGraph(ctx, requestFor(serverSession, readParams))
if err != nil { if err != nil {
t.Fatalf("Final MCP ReadGraph failed: %v", err) t.Fatalf("Final MCP ReadGraph failed: %v", err)
} }
@ -647,7 +647,7 @@ func TestMCPErrorHandling(t *testing.T) {
}, },
} }
_, err := kb.AddObservations(ctx, serverSession, addObsParams) _, err := kb.AddObservations(ctx, requestFor(serverSession, addObsParams))
if err == nil { if err == nil {
t.Errorf("expected MCP AddObservations to return error for non-existent entity") t.Errorf("expected MCP AddObservations to return error for non-existent entity")
} else { } else {
@ -678,7 +678,7 @@ func TestMCPResponseFormat(t *testing.T) {
}, },
} }
result, err := kb.CreateEntities(ctx, serverSession, createParams) result, err := kb.CreateEntities(ctx, requestFor(serverSession, createParams))
if err != nil { if err != nil {
t.Fatalf("CreateEntities failed: %v", err) t.Fatalf("CreateEntities failed: %v", err)
} }
@ -701,3 +701,7 @@ func TestMCPResponseFormat(t *testing.T) {
t.Errorf("expected Content[0] to be TextContent") t.Errorf("expected Content[0] to be TextContent")
} }
} }
func requestFor[P mcp.Params](ss *mcp.ServerSession, p P) *mcp.ServerRequest[P] {
return &mcp.ServerRequest[P]{Session: ss, Params: p}
}

22
examples/server/sequentialthinking/main.go
View file

@ -231,8 +231,8 @@ func deepCopyThoughts(thoughts []*Thought) []*Thought {
} }
// StartThinking begins a new sequential thinking session for a complex problem. // StartThinking begins a new sequential thinking session for a complex problem.
func StartThinking(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[StartThinkingArgs]) (*mcp.CallToolResultFor[any], error) { func StartThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[StartThinkingArgs]]) (*mcp.CallToolResultFor[any], error) {
args := params.Arguments args := req.Params.Arguments
sessionID := args.SessionID sessionID := args.SessionID
if sessionID == "" { if sessionID == "" {
@ -266,8 +266,8 @@ func StartThinking(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallT
} }
// ContinueThinking adds the next thought step, revises a previous step, or creates a branch in the thinking process. // ContinueThinking adds the next thought step, revises a previous step, or creates a branch in the thinking process.
func ContinueThinking(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[ContinueThinkingArgs]) (*mcp.CallToolResultFor[any], error) { func ContinueThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[ContinueThinkingArgs]]) (*mcp.CallToolResultFor[any], error) {
args := params.Arguments args := req.Params.Arguments
// Handle revision of existing thought // Handle revision of existing thought
if args.ReviseStep != nil { if args.ReviseStep != nil {
@ -395,8 +395,8 @@ func ContinueThinking(ctx context.Context, ss *mcp.ServerSession, params *mcp.Ca
} }
// ReviewThinking provides a complete review of the thinking process for a session. // ReviewThinking provides a complete review of the thinking process for a session.
func ReviewThinking(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[ReviewThinkingArgs]) (*mcp.CallToolResultFor[any], error) { func ReviewThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[ReviewThinkingArgs]]) (*mcp.CallToolResultFor[any], error) {
args := params.Arguments args := req.Params.Arguments
// Get a snapshot of the session to avoid race conditions // Get a snapshot of the session to avoid race conditions
sessionSnapshot, exists := store.SessionSnapshot(args.SessionID) sessionSnapshot, exists := store.SessionSnapshot(args.SessionID)
@ -434,11 +434,11 @@ func ReviewThinking(ctx context.Context, ss *mcp.ServerSession, params *mcp.Call
} }
// ThinkingHistory handles resource requests for thinking session data and history. // ThinkingHistory handles resource requests for thinking session data and history.
func ThinkingHistory(ctx context.Context, ss *mcp.ServerSession, params *mcp.ReadResourceParams) (*mcp.ReadResourceResult, error) { func ThinkingHistory(ctx context.Context, req *mcp.ServerRequest[*mcp.ReadResourceParams]) (*mcp.ReadResourceResult, error) {
// Extract session ID from URI (e.g., "thinking://session_123") // Extract session ID from URI (e.g., "thinking://session_123")
u, err := url.Parse(params.URI) u, err := url.Parse(req.Params.URI)
if err != nil { if err != nil {
return nil, fmt.Errorf("invalid thinking resource URI: %s", params.URI) return nil, fmt.Errorf("invalid thinking resource URI: %s", req.Params.URI)
} }
if u.Scheme != "thinking" { if u.Scheme != "thinking" {
return nil, fmt.Errorf("invalid thinking resource URI scheme: %s", u.Scheme) return nil, fmt.Errorf("invalid thinking resource URI scheme: %s", u.Scheme)
@ -456,7 +456,7 @@ func ThinkingHistory(ctx context.Context, ss *mcp.ServerSession, params *mcp.Rea
return &mcp.ReadResourceResult{ return &mcp.ReadResourceResult{
Contents: []*mcp.ResourceContents{ Contents: []*mcp.ResourceContents{
{ {
URI: params.URI, URI: req.Params.URI,
MIMEType: "application/json", MIMEType: "application/json",
Text: string(data), Text: string(data),
}, },
@ -478,7 +478,7 @@ func ThinkingHistory(ctx context.Context, ss *mcp.ServerSession, params *mcp.Rea
return &mcp.ReadResourceResult{ return &mcp.ReadResourceResult{
Contents: []*mcp.ResourceContents{ Contents: []*mcp.ResourceContents{
{ {
URI: params.URI, URI: req.Params.URI,
MIMEType: "application/json", MIMEType: "application/json",
Text: string(data), Text: string(data),
}, },

30
examples/server/sequentialthinking/main_test.go
View file

@ -31,7 +31,7 @@ func TestStartThinking(t *testing.T) {
Arguments: args, Arguments: args,
} }
result, err := StartThinking(ctx, nil, params) result, err := StartThinking(ctx, requestFor(params))
if err != nil { if err != nil {
t.Fatalf("StartThinking() error = %v", err) t.Fatalf("StartThinking() error = %v", err)
} }
@ -89,7 +89,7 @@ func TestContinueThinking(t *testing.T) {
Arguments: startArgs, Arguments: startArgs,
} }
_, err := StartThinking(ctx, nil, startParams) _, err := StartThinking(ctx, requestFor(startParams))
if err != nil { if err != nil {
t.Fatalf("StartThinking() error = %v", err) t.Fatalf("StartThinking() error = %v", err)
} }
@ -105,7 +105,7 @@ func TestContinueThinking(t *testing.T) {
Arguments: continueArgs, Arguments: continueArgs,
} }
result, err := ContinueThinking(ctx, nil, continueParams) result, err := ContinueThinking(ctx, requestFor(continueParams))
if err != nil { if err != nil {
t.Fatalf("ContinueThinking() error = %v", err) t.Fatalf("ContinueThinking() error = %v", err)
} }
@ -158,7 +158,7 @@ func TestContinueThinkingWithCompletion(t *testing.T) {
Arguments: startArgs, Arguments: startArgs,
} }
_, err := StartThinking(ctx, nil, startParams) _, err := StartThinking(ctx, requestFor(startParams))
if err != nil { if err != nil {
t.Fatalf("StartThinking() error = %v", err) t.Fatalf("StartThinking() error = %v", err)
} }
@ -176,7 +176,7 @@ func TestContinueThinkingWithCompletion(t *testing.T) {
Arguments: continueArgs, Arguments: continueArgs,
} }
result, err := ContinueThinking(ctx, nil, continueParams) result, err := ContinueThinking(ctx, requestFor(continueParams))
if err != nil { if err != nil {
t.Fatalf("ContinueThinking() error = %v", err) t.Fatalf("ContinueThinking() error = %v", err)
} }
@ -233,7 +233,7 @@ func TestContinueThinkingRevision(t *testing.T) {
Arguments: continueArgs, Arguments: continueArgs,
} }
result, err := ContinueThinking(ctx, nil, continueParams) result, err := ContinueThinking(ctx, requestFor(continueParams))
if err != nil { if err != nil {
t.Fatalf("ContinueThinking() error = %v", err) t.Fatalf("ContinueThinking() error = %v", err)
} }
@ -289,7 +289,7 @@ func TestContinueThinkingBranching(t *testing.T) {
Arguments: continueArgs, Arguments: continueArgs,
} }
result, err := ContinueThinking(ctx, nil, continueParams) result, err := ContinueThinking(ctx, requestFor(continueParams))
if err != nil { if err != nil {
t.Fatalf("ContinueThinking() error = %v", err) t.Fatalf("ContinueThinking() error = %v", err)
} }
@ -356,7 +356,7 @@ func TestReviewThinking(t *testing.T) {
Arguments: reviewArgs, Arguments: reviewArgs,
} }
result, err := ReviewThinking(ctx, nil, reviewParams) result, err := ReviewThinking(ctx, requestFor(reviewParams))
if err != nil { if err != nil {
t.Fatalf("ReviewThinking() error = %v", err) t.Fatalf("ReviewThinking() error = %v", err)
} }
@ -431,7 +431,7 @@ func TestThinkingHistory(t *testing.T) {
URI: "thinking://sessions", URI: "thinking://sessions",
} }
result, err := ThinkingHistory(ctx, nil, listParams) result, err := ThinkingHistory(ctx, requestFor(listParams))
if err != nil { if err != nil {
t.Fatalf("ThinkingHistory() error = %v", err) t.Fatalf("ThinkingHistory() error = %v", err)
} }
@ -461,7 +461,7 @@ func TestThinkingHistory(t *testing.T) {
URI: "thinking://session1", URI: "thinking://session1",
} }
result, err = ThinkingHistory(ctx, nil, sessionParams) result, err = ThinkingHistory(ctx, requestFor(sessionParams))
if err != nil { if err != nil {
t.Fatalf("ThinkingHistory() error = %v", err) t.Fatalf("ThinkingHistory() error = %v", err)
} }
@ -496,7 +496,7 @@ func TestInvalidOperations(t *testing.T) {
Arguments: continueArgs, Arguments: continueArgs,
} }
_, err := ContinueThinking(ctx, nil, continueParams) _, err := ContinueThinking(ctx, requestFor(continueParams))
if err == nil { if err == nil {
t.Error("Expected error for non-existent session") t.Error("Expected error for non-existent session")
} }
@ -511,7 +511,7 @@ func TestInvalidOperations(t *testing.T) {
Arguments: reviewArgs, Arguments: reviewArgs,
} }
_, err = ReviewThinking(ctx, nil, reviewParams) _, err = ReviewThinking(ctx, requestFor(reviewParams))
if err == nil { if err == nil {
t.Error("Expected error for non-existent session in review") t.Error("Expected error for non-existent session in review")
} }
@ -541,8 +541,12 @@ func TestInvalidOperations(t *testing.T) {
Arguments: invalidReviseArgs, Arguments: invalidReviseArgs,
} }
_, err = ContinueThinking(ctx, nil, invalidReviseParams) _, err = ContinueThinking(ctx, requestFor(invalidReviseParams))
if err == nil { if err == nil {
t.Error("Expected error for invalid revision step") t.Error("Expected error for invalid revision step")
} }
} }
func requestFor[P mcp.Params](p P) *mcp.ServerRequest[P] {
return &mcp.ServerRequest[P]{Params: p}
}

4
examples/server/sse/main.go
View file

@ -24,10 +24,10 @@ type SayHiParams struct {
Name string `json:"name"` 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]{ return &mcp.CallToolResultFor[any]{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{Text: "Hi " + params.Arguments.Name}, &mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name},
}, },
}, nil }, nil
} }

4
internal/readme/server/server.go
View file

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

127
mcp/client.go
View file

@ -24,8 +24,8 @@ type Client struct {
mu sync.Mutex mu sync.Mutex
roots *featureSet[*Root] roots *featureSet[*Root]
sessions []*ClientSession sessions []*ClientSession
sendingMethodHandler_ MethodHandler[*ClientSession] sendingMethodHandler_ MethodHandler
receivingMethodHandler_ MethodHandler[*ClientSession] receivingMethodHandler_ MethodHandler
} }
// NewClient creates a new [Client]. // NewClient creates a new [Client].
@ -55,14 +55,14 @@ func NewClient(impl *Implementation, opts *ClientOptions) *Client {
type ClientOptions struct { type ClientOptions struct {
// Handler for sampling. // Handler for sampling.
// Called when a server calls CreateMessage. // 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. // Handlers for notifications from the server.
ToolListChangedHandler func(context.Context, *ClientSession, *ToolListChangedParams) ToolListChangedHandler func(context.Context, *ClientRequest[*ToolListChangedParams])
PromptListChangedHandler func(context.Context, *ClientSession, *PromptListChangedParams) PromptListChangedHandler func(context.Context, *ClientRequest[*PromptListChangedParams])
ResourceListChangedHandler func(context.Context, *ClientSession, *ResourceListChangedParams) ResourceListChangedHandler func(context.Context, *ClientRequest[*ResourceListChangedParams])
ResourceUpdatedHandler func(context.Context, *ClientSession, *ResourceUpdatedNotificationParams) ResourceUpdatedHandler func(context.Context, *ClientRequest[*ResourceUpdatedNotificationParams])
LoggingMessageHandler func(context.Context, *ClientSession, *LoggingMessageParams) LoggingMessageHandler func(context.Context, *ClientRequest[*LoggingMessageParams])
ProgressNotificationHandler func(context.Context, *ClientSession, *ProgressNotificationParams) ProgressNotificationHandler func(context.Context, *ClientRequest[*ProgressNotificationParams])
// If non-zero, defines an interval for regular "ping" requests. // If non-zero, defines an interval for regular "ping" requests.
// If the peer fails to respond to pings originating from the keepalive check, // If the peer fails to respond to pings originating from the keepalive check,
// the session is automatically closed. // the session is automatically closed.
@ -130,7 +130,8 @@ func (c *Client) Connect(ctx context.Context, t Transport, _ *ClientSessionOptio
ClientInfo: c.impl, ClientInfo: c.impl,
Capabilities: caps, 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 { if err != nil {
_ = cs.Close() _ = cs.Close()
return nil, err return nil, err
@ -142,7 +143,8 @@ func (c *Client) Connect(ctx context.Context, t Transport, _ *ClientSessionOptio
if hc, ok := cs.mcpConn.(clientConnection); ok { if hc, ok := cs.mcpConn.(clientConnection); ok {
hc.sessionUpdated(cs.state) 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() _ = cs.Close()
return nil, err return nil, err
} }
@ -216,7 +218,7 @@ func (c *Client) AddRoots(roots ...*Root) {
if len(roots) == 0 { if len(roots) == 0 {
return return
} }
c.changeAndNotify(notificationRootsListChanged, &RootsListChangedParams{}, changeAndNotify(c, notificationRootsListChanged, &RootsListChangedParams{},
func() bool { c.roots.add(roots...); return true }) 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. // It is not an error to remove a nonexistent root.
// TODO: notification // TODO: notification
func (c *Client) RemoveRoots(uris ...string) { func (c *Client) RemoveRoots(uris ...string) {
c.changeAndNotify(notificationRootsListChanged, &RootsListChangedParams{}, changeAndNotify(c, notificationRootsListChanged, &RootsListChangedParams{},
func() bool { return c.roots.remove(uris...) }) func() bool { return c.roots.remove(uris...) })
} }
// changeAndNotify is called when a feature is added or removed. // 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. // 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. // 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 var sessions []*ClientSession
// Lock for the change, but not for the notification. // Lock for the change, but not for the notification.
c.mu.Lock() c.mu.Lock()
@ -243,7 +245,7 @@ func (c *Client) changeAndNotify(notification string, params Params, change func
notifySessions(sessions, notification, params) 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() c.mu.Lock()
defer c.mu.Unlock() defer c.mu.Unlock()
roots := slices.Collect(c.roots.all()) roots := slices.Collect(c.roots.all())
@ -255,12 +257,12 @@ func (c *Client) listRoots(_ context.Context, _ *ClientSession, _ *ListRootsPara
}, nil }, 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 { if c.opts.CreateMessageHandler == nil {
// TODO: wrap or annotate this error? Pick a standard code? // TODO: wrap or annotate this error? Pick a standard code?
return nil, jsonrpc2.NewError(CodeUnsupportedMethod, "client does not support CreateMessage") 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 // 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 // Sending middleware is called when a request is sent. It is useful for tasks
// such as tracing, metrics, and adding progress tokens. // such as tracing, metrics, and adding progress tokens.
func (c *Client) AddSendingMiddleware(middleware ...Middleware[*ClientSession]) { func (c *Client) AddSendingMiddleware(middleware ...Middleware) {
c.mu.Lock() c.mu.Lock()
defer c.mu.Unlock() defer c.mu.Unlock()
addMiddleware(&c.sendingMethodHandler_, middleware) 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 // Receiving middleware is called when a request is received. It is useful for tasks
// such as authentication, request logging and metrics. // such as authentication, request logging and metrics.
func (c *Client) AddReceivingMiddleware(middleware ...Middleware[*ClientSession]) { func (c *Client) AddReceivingMiddleware(middleware ...Middleware) {
c.mu.Lock() c.mu.Lock()
defer c.mu.Unlock() defer c.mu.Unlock()
addMiddleware(&c.receivingMethodHandler_, middleware) 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 // TODO(rfindley): actually load and validate the protocol schema, rather than
// curating these method flags. // curating these method flags.
var clientMethodInfos = map[string]methodInfo{ var clientMethodInfos = map[string]methodInfo{
methodComplete: newMethodInfo(sessionMethod((*ClientSession).Complete), 0), methodComplete: newClientMethodInfo(clientSessionMethod((*ClientSession).Complete), 0),
methodPing: newMethodInfo(sessionMethod((*ClientSession).ping), missingParamsOK), methodPing: newClientMethodInfo(clientSessionMethod((*ClientSession).ping), missingParamsOK),
methodListRoots: newMethodInfo(clientMethod((*Client).listRoots), missingParamsOK), methodListRoots: newClientMethodInfo(clientMethod((*Client).listRoots), missingParamsOK),
methodCreateMessage: newMethodInfo(clientMethod((*Client).createMessage), 0), methodCreateMessage: newClientMethodInfo(clientMethod((*Client).createMessage), 0),
notificationToolListChanged: newMethodInfo(clientMethod((*Client).callToolChangedHandler), notification|missingParamsOK), notificationToolListChanged: newClientMethodInfo(clientMethod((*Client).callToolChangedHandler), notification|missingParamsOK),
notificationPromptListChanged: newMethodInfo(clientMethod((*Client).callPromptChangedHandler), notification|missingParamsOK), notificationPromptListChanged: newClientMethodInfo(clientMethod((*Client).callPromptChangedHandler), notification|missingParamsOK),
notificationResourceListChanged: newMethodInfo(clientMethod((*Client).callResourceChangedHandler), notification|missingParamsOK), notificationResourceListChanged: newClientMethodInfo(clientMethod((*Client).callResourceChangedHandler), notification|missingParamsOK),
notificationResourceUpdated: newMethodInfo(clientMethod((*Client).callResourceUpdatedHandler), notification|missingParamsOK), notificationResourceUpdated: newClientMethodInfo(clientMethod((*Client).callResourceUpdatedHandler), notification|missingParamsOK),
notificationLoggingMessage: newMethodInfo(clientMethod((*Client).callLoggingHandler), notification), notificationLoggingMessage: newClientMethodInfo(clientMethod((*Client).callLoggingHandler), notification),
notificationProgress: newMethodInfo(sessionMethod((*ClientSession).callProgressNotificationHandler), notification), notificationProgress: newClientMethodInfo(clientSessionMethod((*ClientSession).callProgressNotificationHandler), notification),
} }
func (cs *ClientSession) sendingMethodInfos() map[string]methodInfo { func (cs *ClientSession) sendingMethodInfos() map[string]methodInfo {
@ -342,25 +344,29 @@ func (*ClientSession) ping(context.Context, *PingParams) (*emptyResult, error) {
return &emptyResult{}, nil 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. // Ping makes an MCP "ping" request to the server.
func (cs *ClientSession) Ping(ctx context.Context, params *PingParams) error { 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 return err
} }
// ListPrompts lists prompts that are currently available on the server. // ListPrompts lists prompts that are currently available on the server.
func (cs *ClientSession) ListPrompts(ctx context.Context, params *ListPromptsParams) (*ListPromptsResult, error) { 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. // GetPrompt gets a prompt from the server.
func (cs *ClientSession) GetPrompt(ctx context.Context, params *GetPromptParams) (*GetPromptResult, error) { 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. // ListTools lists tools that are currently available on the server.
func (cs *ClientSession) ListTools(ctx context.Context, params *ListToolsParams) (*ListToolsResult, error) { 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. // 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. // Avoid sending nil over the wire.
params.Arguments = map[string]any{} 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 { 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 return err
} }
// ListResources lists the resources that are currently available on the server. // ListResources lists the resources that are currently available on the server.
func (cs *ClientSession) ListResources(ctx context.Context, params *ListResourcesParams) (*ListResourcesResult, error) { 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. // ListResourceTemplates lists the resource templates that are currently available on the server.
func (cs *ClientSession) ListResourceTemplates(ctx context.Context, params *ListResourceTemplatesParams) (*ListResourceTemplatesResult, error) { 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. // ReadResource asks the server to read a resource and return its contents.
func (cs *ClientSession) ReadResource(ctx context.Context, params *ReadResourceParams) (*ReadResourceResult, error) { 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) { 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 // Subscribe sends a "resources/subscribe" request to the server, asking for
// notifications when the specified resource changes. // notifications when the specified resource changes.
func (cs *ClientSession) Subscribe(ctx context.Context, params *SubscribeParams) error { 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 return err
} }
// Unsubscribe sends a "resources/unsubscribe" request to the server, cancelling // Unsubscribe sends a "resources/unsubscribe" request to the server, cancelling
// a previous subscription. // a previous subscription.
func (cs *ClientSession) Unsubscribe(ctx context.Context, params *UnsubscribeParams) error { 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 return err
} }
func (c *Client) callToolChangedHandler(ctx context.Context, s *ClientSession, params *ToolListChangedParams) (Result, error) { func (c *Client) callToolChangedHandler(ctx context.Context, req *ClientRequest[*ToolListChangedParams]) (Result, error) {
return callNotificationHandler(ctx, c.opts.ToolListChangedHandler, s, params) 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) { func (c *Client) callPromptChangedHandler(ctx context.Context, req *ClientRequest[*PromptListChangedParams]) (Result, error) {
return callNotificationHandler(ctx, c.opts.PromptListChangedHandler, s, params) 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) { func (c *Client) callResourceChangedHandler(ctx context.Context, req *ClientRequest[*ResourceListChangedParams]) (Result, error) {
return callNotificationHandler(ctx, c.opts.ResourceListChangedHandler, s, params) 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) { func (c *Client) callResourceUpdatedHandler(ctx context.Context, req *ClientRequest[*ResourceUpdatedNotificationParams]) (Result, error) {
return callNotificationHandler(ctx, c.opts.ResourceUpdatedHandler, s, params) 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 { if h := c.opts.LoggingMessageHandler; h != nil {
h(ctx, cs, params) h(ctx, req)
} }
return nil, nil return nil, nil
} }
func (cs *ClientSession) callProgressNotificationHandler(ctx context.Context, params *ProgressNotificationParams) (Result, error) { 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 // 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 // This can be used if the client is performing a long-running task that was
// initiated by the server // initiated by the server
func (cs *ClientSession) NotifyProgress(ctx context.Context, params *ProgressNotificationParams) error { 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, // 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( return func(
ctx context.Context, ctx context.Context,
session *mcp.ServerSession,
method string, method string,
params mcp.Params, req mcp.Request,
) (mcp.Result, error) { ) (mcp.Result, error) {
logger.Info("MCP method started", logger.Info("MCP method started",
"method", method, "method", method,
"session_id", session.ID(), "session_id", req.GetSession().ID(),
"has_params", params != nil, "has_params", req.GetParams() != nil,
) )
start := time.Now() start := time.Now()
result, err := next(ctx, session, method, params) result, err := next(ctx, method, req)
duration := time.Since(start) duration := time.Since(start)
if err != nil { if err != nil {
logger.Error("MCP method failed", logger.Error("MCP method failed",
"method", method, "method", method,
"session_id", session.ID(), "session_id", req.GetSession().ID(),
"duration_ms", duration.Milliseconds(), "duration_ms", duration.Milliseconds(),
"err", err, "err", err,
) )
} else { } else {
logger.Info("MCP method completed", logger.Info("MCP method completed",
"method", method, "method", method,
"session_id", session.ID(), "session_id", req.GetSession().ID(),
"duration_ms", duration.Milliseconds(), "duration_ms", duration.Milliseconds(),
"has_result", result != nil, "has_result", result != nil,
) )
@ -90,10 +89,9 @@ func Example_loggingMiddleware() {
}, },
func( func(
ctx context.Context, ctx context.Context,
ss *mcp.ServerSession, req *mcp.ServerRequest[*mcp.CallToolParamsFor[map[string]any]],
params *mcp.CallToolParamsFor[map[string]any],
) (*mcp.CallToolResultFor[any], error) { ) (*mcp.CallToolResultFor[any], error) {
name, ok := params.Arguments["name"].(string) name, ok := req.Params.Arguments["name"].(string)
if !ok { if !ok {
return nil, fmt.Errorf("name parameter is required and must be a string") 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 _ = c
} }
func addProgressToken[S mcp.Session](h mcp.MethodHandler[S]) mcp.MethodHandler[S] { func addProgressToken[S mcp.Session](h mcp.MethodHandler) mcp.MethodHandler {
return func(ctx context.Context, s S, method string, params mcp.Params) (result mcp.Result, err error) { return func(ctx context.Context, method string, req mcp.Request) (result mcp.Result, err error) {
if rp, ok := params.(mcp.RequestParams); ok { if rp, ok := req.GetParams().(mcp.RequestParams); ok {
rp.SetProgressToken(nextProgressToken.Add(1)) 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. // TODO(jba): after schemas are stateless (WIP), this can be a variable.
func greetTool() *Tool { return &Tool{Name: "greet", Description: "say hi"} } func greetTool() *Tool { return &Tool{Name: "greet", Description: "say hi"} }
func sayHi(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[hiParams]) (*CallToolResultFor[any], error) { func sayHi(ctx context.Context, req *ServerRequest[*CallToolParamsFor[hiParams]]) (*CallToolResultFor[any], error) {
if err := ss.Ping(ctx, nil); err != nil { if err := req.Session.Ping(ctx, nil); err != nil {
return nil, fmt.Errorf("ping failed: %v", err) 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{ var codeReviewPrompt = &Prompt{
@ -73,16 +73,20 @@ func TestEndToEnd(t *testing.T) {
} }
sopts := &ServerOptions{ sopts := &ServerOptions{
InitializedHandler: func(context.Context, *ServerSession, *InitializedParams) { notificationChans["initialized"] <- 0 }, InitializedHandler: func(context.Context, *ServerRequest[*InitializedParams]) {
RootsListChangedHandler: func(context.Context, *ServerSession, *RootsListChangedParams) { notificationChans["roots"] <- 0 }, notificationChans["initialized"] <- 0
ProgressNotificationHandler: func(context.Context, *ServerSession, *ProgressNotificationParams) { },
RootsListChangedHandler: func(context.Context, *ServerRequest[*RootsListChangedParams]) {
notificationChans["roots"] <- 0
},
ProgressNotificationHandler: func(context.Context, *ServerRequest[*ProgressNotificationParams]) {
notificationChans["progress_server"] <- 0 notificationChans["progress_server"] <- 0
}, },
SubscribeHandler: func(context.Context, *ServerSession, *SubscribeParams) error { SubscribeHandler: func(context.Context, *ServerRequest[*SubscribeParams]) error {
notificationChans["subscribe"] <- 0 notificationChans["subscribe"] <- 0
return nil return nil
}, },
UnsubscribeHandler: func(context.Context, *ServerSession, *UnsubscribeParams) error { UnsubscribeHandler: func(context.Context, *ServerRequest[*UnsubscribeParams]) error {
notificationChans["unsubscribe"] <- 0 notificationChans["unsubscribe"] <- 0
return nil return nil
}, },
@ -93,7 +97,7 @@ func TestEndToEnd(t *testing.T) {
Description: "say hi", Description: "say hi",
}, sayHi) }, sayHi)
s.AddTool(&Tool{Name: "fail", InputSchema: &jsonschema.Schema{}}, 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 return nil, errTestFailure
}) })
s.AddPrompt(codeReviewPrompt, codReviewPromptHandler) s.AddPrompt(codeReviewPrompt, codReviewPromptHandler)
@ -124,19 +128,25 @@ func TestEndToEnd(t *testing.T) {
loggingMessages := make(chan *LoggingMessageParams, 100) // big enough for all logging loggingMessages := make(chan *LoggingMessageParams, 100) // big enough for all logging
opts := &ClientOptions{ 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 return &CreateMessageResult{Model: "aModel", Content: &TextContent{}}, nil
}, },
ToolListChangedHandler: func(context.Context, *ClientSession, *ToolListChangedParams) { notificationChans["tools"] <- 0 }, ToolListChangedHandler: func(context.Context, *ClientRequest[*ToolListChangedParams]) {
PromptListChangedHandler: func(context.Context, *ClientSession, *PromptListChangedParams) { notificationChans["prompts"] <- 0 }, notificationChans["tools"] <- 0
ResourceListChangedHandler: func(context.Context, *ClientSession, *ResourceListChangedParams) { notificationChans["resources"] <- 0 },
LoggingMessageHandler: func(_ context.Context, _ *ClientSession, lm *LoggingMessageParams) {
loggingMessages <- lm
}, },
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 notificationChans["progress_client"] <- 0
}, },
ResourceUpdatedHandler: func(context.Context, *ClientSession, *ResourceUpdatedNotificationParams) { ResourceUpdatedHandler: func(context.Context, *ClientRequest[*ResourceUpdatedNotificationParams]) {
notificationChans["resource_updated"] <- 0 notificationChans["resource_updated"] <- 0
}, },
} }
@ -500,8 +510,8 @@ var embeddedResources = map[string]string{
"info": "This is the MCP test server.", "info": "This is the MCP test server.",
} }
func handleEmbeddedResource(_ context.Context, _ *ServerSession, params *ReadResourceParams) (*ReadResourceResult, error) { func handleEmbeddedResource(_ context.Context, req *ServerRequest[*ReadResourceParams]) (*ReadResourceResult, error) {
u, err := url.Parse(params.URI) u, err := url.Parse(req.Params.URI)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -515,7 +525,7 @@ func handleEmbeddedResource(_ context.Context, _ *ServerSession, params *ReadRes
} }
return &ReadResourceResult{ return &ReadResourceResult{
Contents: []*ResourceContents{ Contents: []*ResourceContents{
{URI: params.URI, MIMEType: "text/plain", Text: text}, {URI: req.Params.URI, MIMEType: "text/plain", Text: text},
}, },
}, nil }, nil
} }
@ -637,7 +647,7 @@ func TestCancellation(t *testing.T) {
cancelled = make(chan struct{}, 1) // don't block the request 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{}{} start <- struct{}{}
select { select {
case <-ctx.Done(): 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 // traceCalls creates a middleware function that prints the method before and after each call
// with the given prefix. // with the given prefix.
func traceCalls[S Session](w io.Writer, prefix string) Middleware[S] { func traceCalls[S Session](w io.Writer, prefix string) Middleware {
return func(h MethodHandler[S]) MethodHandler[S] { return func(h MethodHandler) MethodHandler {
return func(ctx context.Context, sess S, method string, params Params) (Result, error) { return func(ctx context.Context, method string, req Request) (Result, error) {
fmt.Fprintf(w, "%s >%s\n", prefix, method) fmt.Fprintf(w, "%s >%s\n", prefix, method)
defer 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 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. // A ResourceHandler is a function that reads a resource.
// It will be called when the client calls [ClientSession.ReadResource]. // It will be called when the client calls [ClientSession.ReadResource].
// If it cannot find the resource, it should return the result of calling [ResourceNotFoundError]. // 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 // ResourceNotFoundError returns an error indicating that a resource being read could
// not be found. // not be found.

194
mcp/server.go
View file

@ -41,8 +41,8 @@ type Server struct {
resources *featureSet[*serverResource] resources *featureSet[*serverResource]
resourceTemplates *featureSet[*serverResourceTemplate] resourceTemplates *featureSet[*serverResourceTemplate]
sessions []*ServerSession sessions []*ServerSession
sendingMethodHandler_ MethodHandler[*ServerSession] sendingMethodHandler_ MethodHandler
receivingMethodHandler_ MethodHandler[*ServerSession] receivingMethodHandler_ MethodHandler
resourceSubscriptions map[string]map[*ServerSession]bool // uri -> session -> bool resourceSubscriptions map[string]map[*ServerSession]bool // uri -> session -> bool
} }
@ -51,24 +51,24 @@ type ServerOptions struct {
// Optional instructions for connected clients. // Optional instructions for connected clients.
Instructions string Instructions string
// If non-nil, called when "notifications/initialized" is received. // 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 // PageSize is the maximum number of items to return in a single page for
// list methods (e.g. ListTools). // list methods (e.g. ListTools).
PageSize int PageSize int
// If non-nil, called when "notifications/roots/list_changed" is received. // 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. // 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. // 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 non-zero, defines an interval for regular "ping" requests.
// If the peer fails to respond to pings originating from the keepalive check, // If the peer fails to respond to pings originating from the keepalive check,
// the session is automatically closed. // the session is automatically closed.
KeepAlive time.Duration KeepAlive time.Duration
// Function called when a client session subscribes to a resource. // 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. // 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, // If true, advertises the prompts capability during initialization,
// even if no prompts have been registered. // even if no prompts have been registered.
HasPrompts bool HasPrompts bool
@ -258,11 +258,11 @@ func (s *Server) capabilities() *serverCapabilities {
return caps 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 { if s.opts.CompletionHandler == nil {
return nil, jsonrpc2.ErrMethodNotFound 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. // 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) 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() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
if params == nil { if req.Params == nil {
params = &ListPromptsParams{} 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 res.Prompts = []*Prompt{} // avoid JSON null
for _, p := range prompts { for _, p := range prompts {
res.Prompts = append(res.Prompts, p.prompt) 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() s.mu.Lock()
prompt, ok := s.prompts.get(params.Name) prompt, ok := s.prompts.get(req.Params.Name)
s.mu.Unlock() s.mu.Unlock()
if !ok { if !ok {
// TODO: surface the error code over the wire, instead of flattening it into the string. // 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() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
if params == nil { if req.Params == nil {
params = &ListToolsParams{} 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 res.Tools = []*Tool{} // avoid JSON null
for _, t := range tools { for _, t := range tools {
res.Tools = append(res.Tools, t.tool) 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() s.mu.Lock()
st, ok := s.tools.get(params.Name) st, ok := s.tools.get(req.Params.Name)
s.mu.Unlock() s.mu.Unlock()
if !ok { 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() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
if params == nil { if req.Params == nil {
params = &ListResourcesParams{} 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 res.Resources = []*Resource{} // avoid JSON null
for _, r := range resources { for _, r := range resources {
res.Resources = append(res.Resources, r.resource) 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() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
if params == nil { if req.Params == nil {
params = &ListResourceTemplatesParams{} 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) { func(res *ListResourceTemplatesResult, rts []*serverResourceTemplate) {
res.ResourceTemplates = []*ResourceTemplate{} // avoid JSON null res.ResourceTemplates = []*ResourceTemplate{} // avoid JSON null
for _, rt := range rts { 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) { func (s *Server) readResource(ctx context.Context, req *ServerRequest[*ReadResourceParams]) (*ReadResourceResult, error) {
uri := params.URI uri := req.Params.URI
// Look up the resource URI in the lists of resources and resource templates. // Look up the resource URI in the lists of resources and resource templates.
// This is a security check as well as an information lookup. // This is a security check as well as an information lookup.
handler, mimeType, ok := s.lookupResourceHandler(uri) 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. // Treat an unregistered resource the same as a registered one that couldn't be found.
return nil, ResourceNotFoundError(uri) return nil, ResourceNotFoundError(uri)
} }
res, err := handler(ctx, ss, params) res, err := handler(ctx, req)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -430,11 +430,11 @@ func fileResourceHandler(dir string) ResourceHandler {
if err != nil { if err != nil {
panic(err) panic(err)
} }
return func(ctx context.Context, ss *ServerSession, params *ReadResourceParams) (_ *ReadResourceResult, err error) { return func(ctx context.Context, req *ServerRequest[*ReadResourceParams]) (_ *ReadResourceResult, err error) {
defer util.Wrapf(&err, "reading resource %s", params.URI) defer util.Wrapf(&err, "reading resource %s", req.Params.URI)
// TODO: use a memoizing API here. // TODO: use a memoizing API here.
rootRes, err := ss.ListRoots(ctx, nil) rootRes, err := req.Session.ListRoots(ctx, nil)
if err != nil { if err != nil {
return nil, fmt.Errorf("listing roots: %w", err) return nil, fmt.Errorf("listing roots: %w", err)
} }
@ -442,13 +442,13 @@ func fileResourceHandler(dir string) ResourceHandler {
if err != nil { if err != nil {
return nil, err return nil, err
} }
data, err := readFileResource(params.URI, dirFilepath, roots) data, err := readFileResource(req.Params.URI, dirFilepath, roots)
if err != nil { if err != nil {
return nil, err return nil, err
} }
// TODO(jba): figure out mime type. Omit for now: Server.readResource will fill it in. // TODO(jba): figure out mime type. Omit for now: Server.readResource will fill it in.
return &ReadResourceResult{Contents: []*ResourceContents{ return &ReadResourceResult{Contents: []*ResourceContents{
{URI: params.URI, Blob: data}, {URI: req.Params.URI, Blob: data},
}}, nil }}, nil
} }
} }
@ -465,39 +465,39 @@ func (s *Server) ResourceUpdated(ctx context.Context, params *ResourceUpdatedNot
return nil 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 { if s.opts.SubscribeHandler == nil {
return nil, fmt.Errorf("%w: server does not support resource subscriptions", jsonrpc2.ErrMethodNotFound) 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 return nil, err
} }
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
if s.resourceSubscriptions[params.URI] == nil { if s.resourceSubscriptions[req.Params.URI] == nil {
s.resourceSubscriptions[params.URI] = make(map[*ServerSession]bool) 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 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 { if s.opts.UnsubscribeHandler == nil {
return nil, jsonrpc2.ErrMethodNotFound 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 return nil, err
} }
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
if subscribedSessions, ok := s.resourceSubscriptions[params.URI]; ok { if subscribedSessions, ok := s.resourceSubscriptions[req.Params.URI]; ok {
delete(subscribedSessions, ss) delete(subscribedSessions, req.Session)
if len(subscribedSessions) == 0 { 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 { if wasInitd {
return nil, fmt.Errorf("duplicate %q received", notificationInitialized) 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) { func (s *Server) callRootsListChangedHandler(ctx context.Context, req *ServerRequest[*RootsListChangedParams]) (Result, error) {
return callNotificationHandler(ctx, s.opts.RootsListChangedHandler, ss, params) if h := s.opts.RootsListChangedHandler; h != nil {
h(ctx, req)
}
return nil, nil
} }
func (ss *ServerSession) callProgressNotificationHandler(ctx context.Context, params *ProgressNotificationParams) (Result, error) { func (ss *ServerSession) callProgressNotificationHandler(ctx context.Context, p *ProgressNotificationParams) (Result, error) {
return callNotificationHandler(ctx, ss.server.opts.ProgressNotificationHandler, ss, params) 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 // 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 // This is typically used to report on the status of a long-running request
// that was initiated by the client. // that was initiated by the client.
func (ss *ServerSession) NotifyProgress(ctx context.Context, params *ProgressNotificationParams) error { 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 // 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. // Ping pings the client.
func (ss *ServerSession) Ping(ctx context.Context, params *PingParams) error { 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 return err
} }
// ListRoots lists the client roots. // ListRoots lists the client roots.
func (ss *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) { 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. // CreateMessage sends a sampling request to the client.
func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessageParams) (*CreateMessageResult, error) { 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. // 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 { if compareLevels(params.Level, logLevel) < 0 {
return nil 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 // 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 // Sending middleware is called when a request is sent. It is useful for tasks
// such as tracing, metrics, and adding progress tokens. // such as tracing, metrics, and adding progress tokens.
func (s *Server) AddSendingMiddleware(middleware ...Middleware[*ServerSession]) { func (s *Server) AddSendingMiddleware(middleware ...Middleware) {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
addMiddleware(&s.sendingMethodHandler_, middleware) 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 // Receiving middleware is called when a request is received. It is useful for tasks
// such as authentication, request logging and metrics. // such as authentication, request logging and metrics.
func (s *Server) AddReceivingMiddleware(middleware ...Middleware[*ServerSession]) { func (s *Server) AddReceivingMiddleware(middleware ...Middleware) {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
addMiddleware(&s.receivingMethodHandler_, middleware) 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 // TODO(rfindley): actually load and validate the protocol schema, rather than
// curating these method flags. // curating these method flags.
var serverMethodInfos = map[string]methodInfo{ var serverMethodInfos = map[string]methodInfo{
methodComplete: newMethodInfo(serverMethod((*Server).complete), 0), methodComplete: newServerMethodInfo(serverMethod((*Server).complete), 0),
methodInitialize: newMethodInfo(sessionMethod((*ServerSession).initialize), 0), methodInitialize: newServerMethodInfo(serverSessionMethod((*ServerSession).initialize), 0),
methodPing: newMethodInfo(sessionMethod((*ServerSession).ping), missingParamsOK), methodPing: newServerMethodInfo(serverSessionMethod((*ServerSession).ping), missingParamsOK),
methodListPrompts: newMethodInfo(serverMethod((*Server).listPrompts), missingParamsOK), methodListPrompts: newServerMethodInfo(serverMethod((*Server).listPrompts), missingParamsOK),
methodGetPrompt: newMethodInfo(serverMethod((*Server).getPrompt), 0), methodGetPrompt: newServerMethodInfo(serverMethod((*Server).getPrompt), 0),
methodListTools: newMethodInfo(serverMethod((*Server).listTools), missingParamsOK), methodListTools: newServerMethodInfo(serverMethod((*Server).listTools), missingParamsOK),
methodCallTool: newMethodInfo(serverMethod((*Server).callTool), 0), methodCallTool: newServerMethodInfo(serverMethod((*Server).callTool), 0),
methodListResources: newMethodInfo(serverMethod((*Server).listResources), missingParamsOK), methodListResources: newServerMethodInfo(serverMethod((*Server).listResources), missingParamsOK),
methodListResourceTemplates: newMethodInfo(serverMethod((*Server).listResourceTemplates), missingParamsOK), methodListResourceTemplates: newServerMethodInfo(serverMethod((*Server).listResourceTemplates), missingParamsOK),
methodReadResource: newMethodInfo(serverMethod((*Server).readResource), 0), methodReadResource: newServerMethodInfo(serverMethod((*Server).readResource), 0),
methodSetLevel: newMethodInfo(sessionMethod((*ServerSession).setLevel), 0), methodSetLevel: newServerMethodInfo(serverSessionMethod((*ServerSession).setLevel), 0),
methodSubscribe: newMethodInfo(serverMethod((*Server).subscribe), 0), methodSubscribe: newServerMethodInfo(serverMethod((*Server).subscribe), 0),
methodUnsubscribe: newMethodInfo(serverMethod((*Server).unsubscribe), 0), methodUnsubscribe: newServerMethodInfo(serverMethod((*Server).unsubscribe), 0),
notificationInitialized: newMethodInfo(sessionMethod((*ServerSession).initialized), notification|missingParamsOK), notificationInitialized: newServerMethodInfo(serverSessionMethod((*ServerSession).initialized), notification|missingParamsOK),
notificationRootsListChanged: newMethodInfo(serverMethod((*Server).callRootsListChangedHandler), notification|missingParamsOK), notificationRootsListChanged: newServerMethodInfo(serverMethod((*Server).callRootsListChangedHandler), notification|missingParamsOK),
notificationProgress: newMethodInfo(sessionMethod((*ServerSession).callProgressNotificationHandler), notification), notificationProgress: newServerMethodInfo(serverSessionMethod((*ServerSession).callProgressNotificationHandler), notification),
} }
func (ss *ServerSession) sendingMethodInfos() map[string]methodInfo { return clientMethodInfos } 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) receivingMethodInfos() map[string]methodInfo { return serverMethodInfos }
func (ss *ServerSession) sendingMethodHandler() methodHandler { func (ss *ServerSession) sendingMethodHandler() methodHandler {
ss.server.mu.Lock() s := ss.server
defer ss.server.mu.Unlock() s.mu.Lock()
return ss.server.sendingMethodHandler_ defer s.mu.Unlock()
return s.sendingMethodHandler_
} }
func (ss *ServerSession) receivingMethodHandler() methodHandler { func (ss *ServerSession) receivingMethodHandler() methodHandler {
ss.server.mu.Lock() s := ss.server
defer ss.server.mu.Unlock() s.mu.Lock()
return ss.server.receivingMethodHandler_ defer s.mu.Unlock()
return s.receivingMethodHandler_
} }
// getConn implements [session.getConn]. // getConn implements [session.getConn].
@ -808,13 +823,14 @@ func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParam
version = latestProtocolVersion version = latestProtocolVersion
} }
s := ss.server
return &InitializeResult{ return &InitializeResult{
// TODO(rfindley): alter behavior when falling back to an older version: // TODO(rfindley): alter behavior when falling back to an older version:
// reject unsupported features. // reject unsupported features.
ProtocolVersion: version, ProtocolVersion: version,
Capabilities: ss.server.capabilities(), Capabilities: s.capabilities(),
Instructions: ss.server.opts.Instructions, Instructions: s.opts.Instructions,
ServerInfo: ss.server.impl, ServerInfo: s.impl,
}, nil }, nil
} }

4
mcp/server_example_test.go
View file

@ -16,10 +16,10 @@ type SayHiParams struct {
Name string `json:"name"` 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]{ return &mcp.CallToolResultFor[any]{
Content: []mcp.Content{ Content: []mcp.Content{
&mcp.TextContent{Text: "Hi " + params.Arguments.Name}, &mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name},
}, },
}, nil }, nil
} }

12
mcp/server_test.go
View file

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

185
mcp/shared.go
View file

@ -36,7 +36,7 @@ var supportedProtocolVersions = []string{
// A MethodHandler handles MCP messages. // A MethodHandler handles MCP messages.
// For methods, exactly one of the return values must be nil. // For methods, exactly one of the return values must be nil.
// For notifications, both 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. // 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 // 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]. // A Session is either a [ClientSession] or a [ServerSession].
type Session interface { type Session interface {
*ClientSession | *ServerSession
// ID returns the session ID, or the empty string if there is none. // ID returns the session ID, or the empty string if there is none.
ID() string ID() string
@ -58,29 +57,29 @@ type Session interface {
} }
// Middleware is a function from [MethodHandler] to [MethodHandler]. // 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. // 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) { for _, m := range slices.Backward(middleware) {
*handlerp = m(*handlerp) *handlerp = m(*handlerp)
} }
} }
func defaultSendingMethodHandler[S Session](ctx context.Context, session S, method string, params Params) (Result, error) { func defaultSendingMethodHandler[S Session](ctx context.Context, method string, req Request) (Result, error) {
info, ok := session.sendingMethodInfos()[method] info, ok := req.GetSession().sendingMethodInfos()[method]
if !ok { if !ok {
// This can be called from user code, with an arbitrary value for method. // This can be called from user code, with an arbitrary value for method.
return nil, jsonrpc2.ErrNotHandled return nil, jsonrpc2.ErrNotHandled
} }
// Notifications don't have results. // Notifications don't have results.
if strings.HasPrefix(method, "notifications/") { 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. // Create the result to unmarshal into.
// The concrete type of the result is the return type of the receiving function. // The concrete type of the result is the return type of the receiving function.
res := info.newResult() 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 nil, err
} }
return res, nil return res, nil
@ -95,16 +94,16 @@ func orZero[T any, P *U, U any](p P) T {
return any(p).(T) return any(p).(T)
} }
func handleNotify[S Session](ctx context.Context, session S, method string, params Params) error { func handleNotify(ctx context.Context, method string, req Request) error {
mh := session.sendingMethodHandler().(MethodHandler[S]) mh := req.GetSession().sendingMethodHandler().(MethodHandler)
_, err := mh(ctx, session, method, params) _, err := mh(ctx, method, req)
return err return err
} }
func handleSend[R Result, S Session](ctx context.Context, s S, method string, params Params) (R, error) { func handleSend[R Result](ctx context.Context, method string, req Request) (R, error) {
mh := s.sendingMethodHandler().(MethodHandler[S]) mh := req.GetSession().sendingMethodHandler().(MethodHandler)
// mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol. // 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 { if err != nil {
var z R var z R
return z, err 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. // 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) { func defaultReceivingMethodHandler[S Session](ctx context.Context, method string, req Request) (Result, error) {
info, ok := session.receivingMethodInfos()[method] info, ok := req.GetSession().receivingMethodInfos()[method]
if !ok { if !ok {
// This can be called from user code, with an arbitrary value for method. // This can be called from user code, with an arbitrary value for method.
return nil, jsonrpc2.ErrNotHandled 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) { func handleReceive[S Session](ctx context.Context, session S, jreq *jsonrpc.Request) (Result, error) {
info, err := checkRequest(req, session.receivingMethodInfos()) info, err := checkRequest(jreq, session.receivingMethodInfos())
if err != nil { if err != nil {
return nil, err return nil, err
} }
params, err := info.unmarshalParams(req.Params) params, err := info.unmarshalParams(jreq.Params)
if err != nil { 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. // 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 { if err != nil {
return nil, err return nil, err
} }
@ -179,6 +179,7 @@ type methodInfo struct {
// Unmarshal params from the wire into a Params struct. // Unmarshal params from the wire into a Params struct.
// Used on the receive side. // Used on the receive side.
unmarshalParams func(json.RawMessage) (Params, error) unmarshalParams func(json.RawMessage) (Params, error)
newRequest func(Session, Params) Request
// Run the code when a call to the method is received. // Run the code when a call to the method is received.
// Used on the receive side. // Used on the receive side.
handleMethod methodHandler handleMethod methodHandler
@ -194,7 +195,10 @@ type methodInfo struct {
// - R: results // - R: results
// A typedMethodHandler is like a MethodHandler, but with type information. // 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 { type paramsPtr[T any] interface {
*T *T
@ -208,11 +212,45 @@ const (
missingParamsOK // params may be missing or null 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. // newMethodInfo creates a methodInfo from a typedMethodHandler.
// //
// If isRequest is set, the method is treated as a request rather than a // If isRequest is set, the method is treated as a request rather than a
// notification. // 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{ return methodInfo{
flags: flags, flags: flags,
unmarshalParams: func(m json.RawMessage) (Params, error) { 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 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. // 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. // 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 // 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. // serverMethod is glue for creating a typedMethodHandler from a method on Server.
func serverMethod[P Params, R Result]( func serverMethod[P Params, R Result](
f func(*Server, context.Context, *ServerSession, P) (R, error), f func(*Server, context.Context, *ServerRequest[P]) (R, error),
) typedMethodHandler[*ServerSession, P, R] { ) typedServerMethodHandler[P, R] {
return func(ctx context.Context, ss *ServerSession, p P) (R, error) { return func(ctx context.Context, req *ServerRequest[P]) (R, error) {
return f(ss.server, ctx, ss, p) return f(req.Session.server, ctx, req)
} }
} }
// clientMethod is glue for creating a typedMethodHandler from a method on Client. // clientMethod is glue for creating a typedMethodHandler from a method on Client.
func clientMethod[P Params, R Result]( func clientMethod[P Params, R Result](
f func(*Client, context.Context, *ClientSession, P) (R, error), f func(*Client, context.Context, *ClientRequest[P]) (R, error),
) typedMethodHandler[*ClientSession, P, R] { ) typedClientMethodHandler[P, R] {
return func(ctx context.Context, cs *ClientSession, p P) (R, error) { return func(ctx context.Context, req *ClientRequest[P]) (R, error) {
return f(cs.client, ctx, cs, p) return f(req.Session.client, ctx, req)
} }
} }
// sessionMethod is glue for creating a typedMethodHandler from a method on ServerSession. // serverSessionMethod is glue for creating a typedServerMethodHandler 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] { func serverSessionMethod[P Params, R Result](f func(*ServerSession, context.Context, P) (R, error)) typedServerMethodHandler[P, R] {
return func(ctx context.Context, sess S, p P) (R, error) { return func(ctx context.Context, req *ServerRequest[P]) (R, error) {
return f(sess, ctx, p) 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 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. // notifySessions calls Notify on all the sessions.
// Should be called on a copy of the peer 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 { if sessions == nil {
return return
} }
@ -299,13 +331,25 @@ func notifySessions[S Session](sessions []S, method string, params Params) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel() defer cancel()
for _, s := range sessions { 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 // TODO(jba): surface this error better
log.Printf("calling %s: %v", method, err) 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. // Meta is additional metadata for requests, responses and other types.
type Meta map[string]any type Meta map[string]any
@ -335,6 +379,43 @@ func setProgressToken(p Params, pt any) {
m[progressTokenKey] = pt 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. // Params is a parameter (input) type for an MCP call or notification.
type Params interface { type Params interface {
// isParams discourages implementation of Params outside of this package. // isParams discourages implementation of Params outside of this package.

Some files were not shown because too many files have changed in this diff Show more