diff --git a/README.md b/README.md index 22a3fed..4700d08 100644 --- a/README.md +++ b/README.md @@ -115,9 +115,9 @@ type HiParams struct { 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]{ - Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + params.Arguments.Name}}, + Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}}, }, nil } diff --git a/design/design.md b/design/design.md index 65a2b61..fa2270e 100644 --- a/design/design.md +++ b/design/design.md @@ -408,15 +408,12 @@ We provide a mechanism to add MCP-level middleware on the both the client and se ```go // A MethodHandler handles MCP messages. -// The params argument is an XXXParams struct pointer, such as *GetPromptParams. -// For methods, a MethodHandler must return either an XXResult struct pointer and a nil error, or -// nil with a non-nil error. -// For notifications, a MethodHandler must return nil, nil. -type MethodHandler[S Session] func( - ctx context.Context, _ *S, method string, params Params) (result Result, err error) +// For methods, exactly one of the return values must be nil. +// For notifications, both must be nil. +type MethodHandler func(ctx context.Context, method string, req Request) (result Result, err error) // 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 // 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 // m1(m2(m3(handler))). -func (c *Client) AddSendingMiddleware(middleware ...Middleware[*ClientSession]) -func (c *Client) AddReceivingMiddleware(middleware ...Middleware[*ClientSession]) -func (s *Server) AddSendingMiddleware(middleware ...Middleware[*ServerSession]) -func (s *Server) AddReceivingMiddleware(middleware ...Middleware[*ServerSession]) +func (c *Client) AddSendingMiddleware(middleware ...Middleware) +func (c *Client) AddReceivingMiddleware(middleware ...Middleware) +func (s *Server) AddSendingMiddleware(middleware ...Middleware) +func (s *Server) AddReceivingMiddleware(middleware ...Middleware) ``` As an example, this code adds server-side logging: ```go -func withLogging(h mcp.MethodHandler[*mcp.ServerSession]) mcp.MethodHandler[*mcp.ServerSession]{ - return func(ctx context.Context, s *mcp.ServerSession, method string, params mcp.Params) (res mcp.Result, err error) { +func withLogging(h mcp.MethodHandler) mcp.MethodHandler{ + return func(ctx context.Context, method string, req mcp.Request) (res mcp.Result, err error) { log.Printf("request: %s %v", method, params) defer func() { log.Printf("response: %v, %v", res, err) }() return h(ctx, s , method, params) @@ -597,8 +594,9 @@ type Tool struct { Name string `json:"name"` } -type ToolHandlerFor[In, Out any] func(context.Context, *ServerSession, *CallToolParamsFor[In]) (*CallToolResultFor[Out], error) -type ToolHandler = ToolHandlerFor[map[string]any, any] +// A ToolHandlerFor handles a call to tools/call with typed arguments and results. +type ToolHandlerFor[In, Out any] func(context.Context, *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 @@ -648,8 +646,8 @@ type AddParams struct { Y int `json:"y"` } -func addHandler(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[AddParams]) (*mcp.CallToolResultFor[int], error) { - return &mcp.CallToolResultFor[int]{StructuredContent: params.Arguments.X + params.Arguments.Y}, nil +func addHandler(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[AddParams]]) (*mcp.CallToolResultFor[int], error) { + 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 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. @@ -752,10 +748,10 @@ If a server author wants to support resource subscriptions, they must provide ha ```go type ServerOptions struct { ... - // Function called when a client session subscribes to a resource. - SubscribeHandler func(context.Context, ss *ServerSession, *SubscribeParams) error - // Function called when a client session unsubscribes from a resource. - UnsubscribeHandler func(context.Context, ss *ServerSession, *UnsubscribeParams) error + // Function called when a client session subscribes to a resource. + SubscribeHandler func(context.Context, *ServerRequest[*SubscribeParams]) error + // Function called when a client session unsubscribes from a resource. + 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 type ClientOptions struct { ... - ToolListChangedHandler func(context.Context, *ClientSession, *ToolListChangedParams) - PromptListChangedHandler func(context.Context, *ClientSession, *PromptListChangedParams) + ToolListChangedHandler func(context.Context, *ClientRequest[*ToolListChangedParams]) + PromptListChangedHandler func(context.Context, *ClientRequest[*PromptListChangedParams]) // 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. ```go -// A CompletionHandler handles a call to completion/complete. -type CompletionHandler func(context.Context, *ServerSession, *CompleteParams) (*CompleteResult, error) - type ServerOptions struct { ... // If non-nil, called when a client sends a completion request. - CompletionHandler CompletionHandler + CompletionHandler func(context.Context, *ServerRequest[*CompleteParams]) (*CompleteResult, error) } ``` diff --git a/examples/server/completion/main.go b/examples/server/completion/main.go index d6530b5..f05b272 100644 --- a/examples/server/completion/main.go +++ b/examples/server/completion/main.go @@ -16,17 +16,17 @@ import ( // a CompletionHandler to an MCP Server's options. func main() { // 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. // For this example, we return a fixed set of suggestions. var suggestions []string - switch params.Ref.Type { + switch req.Params.Ref.Type { case "ref/prompt": suggestions = []string{"suggestion1", "suggestion2", "suggestion3"} case "ref/resource": suggestions = []string{"suggestion4", "suggestion5", "suggestion6"} 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{ diff --git a/examples/server/custom-transport/main.go b/examples/server/custom-transport/main.go index cc4f15f..bf0306c 100644 --- a/examples/server/custom-transport/main.go +++ b/examples/server/custom-transport/main.go @@ -85,10 +85,10 @@ type HiArgs struct { } // 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{}]{ Content: []mcp.Content{ - &mcp.TextContent{Text: "Hi " + params.Arguments.Name}, + &mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}, }, }, nil } diff --git a/examples/server/hello/main.go b/examples/server/hello/main.go index 00fb37a..8125441 100644 --- a/examples/server/hello/main.go +++ b/examples/server/hello/main.go @@ -22,10 +22,10 @@ type HiArgs struct { 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{}]{ Content: []mcp.Content{ - &mcp.TextContent{Text: "Hi " + params.Arguments.Name}, + &mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}, }, }, nil } @@ -69,8 +69,8 @@ var embeddedResources = map[string]string{ "info": "This is the hello example server.", } -func handleEmbeddedResource(_ context.Context, _ *mcp.ServerSession, params *mcp.ReadResourceParams) (*mcp.ReadResourceResult, error) { - u, err := url.Parse(params.URI) +func handleEmbeddedResource(_ context.Context, req *mcp.ServerRequest[*mcp.ReadResourceParams]) (*mcp.ReadResourceResult, error) { + u, err := url.Parse(req.Params.URI) if err != nil { return nil, err } @@ -84,7 +84,7 @@ func handleEmbeddedResource(_ context.Context, _ *mcp.ServerSession, params *mcp } return &mcp.ReadResourceResult{ Contents: []*mcp.ResourceContents{ - {URI: params.URI, MIMEType: "text/plain", Text: text}, + {URI: req.Params.URI, MIMEType: "text/plain", Text: text}, }, }, nil } diff --git a/examples/server/memory/kb.go b/examples/server/memory/kb.go index a274f05..f053bee 100644 --- a/examples/server/memory/kb.go +++ b/examples/server/memory/kb.go @@ -84,7 +84,7 @@ func (fs *fileStore) Read() ([]byte, error) { // Write saves data to file with 0600 permissions. 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 nil @@ -431,10 +431,10 @@ func (k knowledgeBase) openNodes(names []string) (KnowledgeGraph, error) { }, 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] - entities, err := k.createEntities(params.Arguments.Entities) + entities, err := k.createEntities(req.Params.Arguments.Entities) if err != nil { return nil, err } @@ -450,10 +450,10 @@ func (k knowledgeBase) CreateEntities(ctx context.Context, ss *mcp.ServerSession 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] - relations, err := k.createRelations(params.Arguments.Relations) + relations, err := k.createRelations(req.Params.Arguments.Relations) if err != nil { return nil, err } @@ -469,10 +469,10 @@ func (k knowledgeBase) CreateRelations(ctx context.Context, ss *mcp.ServerSessio 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] - observations, err := k.addObservations(params.Arguments.Observations) + observations, err := k.addObservations(req.Params.Arguments.Observations) if err != nil { return nil, err } @@ -488,10 +488,10 @@ func (k knowledgeBase) AddObservations(ctx context.Context, ss *mcp.ServerSessio 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{}] - err := k.deleteEntities(params.Arguments.EntityNames) + err := k.deleteEntities(req.Params.Arguments.EntityNames) if err != nil { return nil, err } @@ -503,10 +503,10 @@ func (k knowledgeBase) DeleteEntities(ctx context.Context, ss *mcp.ServerSession 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{}] - err := k.deleteObservations(params.Arguments.Deletions) + err := k.deleteObservations(req.Params.Arguments.Deletions) if err != nil { return nil, err } @@ -518,10 +518,10 @@ func (k knowledgeBase) DeleteObservations(ctx context.Context, ss *mcp.ServerSes 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{}] - err := k.deleteRelations(params.Arguments.Relations) + err := k.deleteRelations(req.Params.Arguments.Relations) if err != nil { return nil, err } @@ -533,7 +533,7 @@ func (k knowledgeBase) DeleteRelations(ctx context.Context, ss *mcp.ServerSessio 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] graph, err := k.loadGraph() @@ -549,10 +549,10 @@ func (k knowledgeBase) ReadGraph(ctx context.Context, ss *mcp.ServerSession, par 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] - graph, err := k.searchNodes(params.Arguments.Query) + graph, err := k.searchNodes(req.Params.Arguments.Query) if err != nil { return nil, err } @@ -565,10 +565,10 @@ func (k knowledgeBase) SearchNodes(ctx context.Context, ss *mcp.ServerSession, p 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] - graph, err := k.openNodes(params.Arguments.Names) + graph, err := k.openNodes(req.Params.Arguments.Names) if err != nil { return nil, err } diff --git a/examples/server/memory/kb_test.go b/examples/server/memory/kb_test.go index e4fbacc..6e29d5e 100644 --- a/examples/server/memory/kb_test.go +++ b/examples/server/memory/kb_test.go @@ -230,7 +230,7 @@ func TestSaveAndLoadGraph(t *testing.T) { // Test malformed data handling 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 { 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 { t.Fatalf("MCP CreateEntities failed: %v", err) } @@ -463,7 +463,7 @@ func TestMCPServerIntegration(t *testing.T) { // Test ReadGraph through MCP readParams := &mcp.CallToolParamsFor[struct{}]{} - readResult, err := kb.ReadGraph(ctx, serverSession, readParams) + readResult, err := kb.ReadGraph(ctx, requestFor(serverSession, readParams)) if err != nil { 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 { 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 { 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 { 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 { 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 { 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 { 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 { t.Fatalf("MCP DeleteEntities failed: %v", err) } @@ -614,7 +614,7 @@ func TestMCPServerIntegration(t *testing.T) { } // Verify final state - finalRead, err := kb.ReadGraph(ctx, serverSession, readParams) + finalRead, err := kb.ReadGraph(ctx, requestFor(serverSession, readParams)) if err != nil { 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 { t.Errorf("expected MCP AddObservations to return error for non-existent entity") } 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 { t.Fatalf("CreateEntities failed: %v", err) } @@ -701,3 +701,7 @@ func TestMCPResponseFormat(t *testing.T) { 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} +} diff --git a/examples/server/sequentialthinking/main.go b/examples/server/sequentialthinking/main.go index e9cb594..45a4fa6 100644 --- a/examples/server/sequentialthinking/main.go +++ b/examples/server/sequentialthinking/main.go @@ -231,8 +231,8 @@ func deepCopyThoughts(thoughts []*Thought) []*Thought { } // 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) { - args := params.Arguments +func StartThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[StartThinkingArgs]]) (*mcp.CallToolResultFor[any], error) { + args := req.Params.Arguments sessionID := args.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. -func ContinueThinking(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[ContinueThinkingArgs]) (*mcp.CallToolResultFor[any], error) { - args := params.Arguments +func ContinueThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[ContinueThinkingArgs]]) (*mcp.CallToolResultFor[any], error) { + args := req.Params.Arguments // Handle revision of existing thought 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. -func ReviewThinking(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[ReviewThinkingArgs]) (*mcp.CallToolResultFor[any], error) { - args := params.Arguments +func ReviewThinking(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[ReviewThinkingArgs]]) (*mcp.CallToolResultFor[any], error) { + args := req.Params.Arguments // Get a snapshot of the session to avoid race conditions 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. -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") - u, err := url.Parse(params.URI) + u, err := url.Parse(req.Params.URI) 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" { 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{ Contents: []*mcp.ResourceContents{ { - URI: params.URI, + URI: req.Params.URI, MIMEType: "application/json", Text: string(data), }, @@ -478,7 +478,7 @@ func ThinkingHistory(ctx context.Context, ss *mcp.ServerSession, params *mcp.Rea return &mcp.ReadResourceResult{ Contents: []*mcp.ResourceContents{ { - URI: params.URI, + URI: req.Params.URI, MIMEType: "application/json", Text: string(data), }, diff --git a/examples/server/sequentialthinking/main_test.go b/examples/server/sequentialthinking/main_test.go index 13cd9cc..c5e4a95 100644 --- a/examples/server/sequentialthinking/main_test.go +++ b/examples/server/sequentialthinking/main_test.go @@ -31,7 +31,7 @@ func TestStartThinking(t *testing.T) { Arguments: args, } - result, err := StartThinking(ctx, nil, params) + result, err := StartThinking(ctx, requestFor(params)) if err != nil { t.Fatalf("StartThinking() error = %v", err) } @@ -89,7 +89,7 @@ func TestContinueThinking(t *testing.T) { Arguments: startArgs, } - _, err := StartThinking(ctx, nil, startParams) + _, err := StartThinking(ctx, requestFor(startParams)) if err != nil { t.Fatalf("StartThinking() error = %v", err) } @@ -105,7 +105,7 @@ func TestContinueThinking(t *testing.T) { Arguments: continueArgs, } - result, err := ContinueThinking(ctx, nil, continueParams) + result, err := ContinueThinking(ctx, requestFor(continueParams)) if err != nil { t.Fatalf("ContinueThinking() error = %v", err) } @@ -158,7 +158,7 @@ func TestContinueThinkingWithCompletion(t *testing.T) { Arguments: startArgs, } - _, err := StartThinking(ctx, nil, startParams) + _, err := StartThinking(ctx, requestFor(startParams)) if err != nil { t.Fatalf("StartThinking() error = %v", err) } @@ -176,7 +176,7 @@ func TestContinueThinkingWithCompletion(t *testing.T) { Arguments: continueArgs, } - result, err := ContinueThinking(ctx, nil, continueParams) + result, err := ContinueThinking(ctx, requestFor(continueParams)) if err != nil { t.Fatalf("ContinueThinking() error = %v", err) } @@ -233,7 +233,7 @@ func TestContinueThinkingRevision(t *testing.T) { Arguments: continueArgs, } - result, err := ContinueThinking(ctx, nil, continueParams) + result, err := ContinueThinking(ctx, requestFor(continueParams)) if err != nil { t.Fatalf("ContinueThinking() error = %v", err) } @@ -289,7 +289,7 @@ func TestContinueThinkingBranching(t *testing.T) { Arguments: continueArgs, } - result, err := ContinueThinking(ctx, nil, continueParams) + result, err := ContinueThinking(ctx, requestFor(continueParams)) if err != nil { t.Fatalf("ContinueThinking() error = %v", err) } @@ -356,7 +356,7 @@ func TestReviewThinking(t *testing.T) { Arguments: reviewArgs, } - result, err := ReviewThinking(ctx, nil, reviewParams) + result, err := ReviewThinking(ctx, requestFor(reviewParams)) if err != nil { t.Fatalf("ReviewThinking() error = %v", err) } @@ -431,7 +431,7 @@ func TestThinkingHistory(t *testing.T) { URI: "thinking://sessions", } - result, err := ThinkingHistory(ctx, nil, listParams) + result, err := ThinkingHistory(ctx, requestFor(listParams)) if err != nil { t.Fatalf("ThinkingHistory() error = %v", err) } @@ -461,7 +461,7 @@ func TestThinkingHistory(t *testing.T) { URI: "thinking://session1", } - result, err = ThinkingHistory(ctx, nil, sessionParams) + result, err = ThinkingHistory(ctx, requestFor(sessionParams)) if err != nil { t.Fatalf("ThinkingHistory() error = %v", err) } @@ -496,7 +496,7 @@ func TestInvalidOperations(t *testing.T) { Arguments: continueArgs, } - _, err := ContinueThinking(ctx, nil, continueParams) + _, err := ContinueThinking(ctx, requestFor(continueParams)) if err == nil { t.Error("Expected error for non-existent session") } @@ -511,7 +511,7 @@ func TestInvalidOperations(t *testing.T) { Arguments: reviewArgs, } - _, err = ReviewThinking(ctx, nil, reviewParams) + _, err = ReviewThinking(ctx, requestFor(reviewParams)) if err == nil { t.Error("Expected error for non-existent session in review") } @@ -541,8 +541,12 @@ func TestInvalidOperations(t *testing.T) { Arguments: invalidReviseArgs, } - _, err = ContinueThinking(ctx, nil, invalidReviseParams) + _, err = ContinueThinking(ctx, requestFor(invalidReviseParams)) if err == nil { t.Error("Expected error for invalid revision step") } } + +func requestFor[P mcp.Params](p P) *mcp.ServerRequest[P] { + return &mcp.ServerRequest[P]{Params: p} +} diff --git a/examples/server/sse/main.go b/examples/server/sse/main.go index 59412b1..2fbd695 100644 --- a/examples/server/sse/main.go +++ b/examples/server/sse/main.go @@ -24,10 +24,10 @@ type SayHiParams struct { Name string `json:"name"` } -func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[SayHiParams]) (*mcp.CallToolResultFor[any], error) { +func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[SayHiParams]]) (*mcp.CallToolResultFor[any], error) { return &mcp.CallToolResultFor[any]{ Content: []mcp.Content{ - &mcp.TextContent{Text: "Hi " + params.Arguments.Name}, + &mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}, }, }, nil } diff --git a/internal/readme/server/server.go b/internal/readme/server/server.go index 3746e19..3aa1037 100644 --- a/internal/readme/server/server.go +++ b/internal/readme/server/server.go @@ -16,9 +16,9 @@ type HiParams struct { 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]{ - Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + params.Arguments.Name}}, + Content: []mcp.Content{&mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}}, }, nil } diff --git a/mcp/client.go b/mcp/client.go index 5798dc5..b0db1d6 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -24,8 +24,8 @@ type Client struct { mu sync.Mutex roots *featureSet[*Root] sessions []*ClientSession - sendingMethodHandler_ MethodHandler[*ClientSession] - receivingMethodHandler_ MethodHandler[*ClientSession] + sendingMethodHandler_ MethodHandler + receivingMethodHandler_ MethodHandler } // NewClient creates a new [Client]. @@ -55,14 +55,14 @@ func NewClient(impl *Implementation, opts *ClientOptions) *Client { type ClientOptions struct { // Handler for sampling. // Called when a server calls CreateMessage. - CreateMessageHandler func(context.Context, *ClientSession, *CreateMessageParams) (*CreateMessageResult, error) + CreateMessageHandler func(context.Context, *ClientRequest[*CreateMessageParams]) (*CreateMessageResult, error) // Handlers for notifications from the server. - ToolListChangedHandler func(context.Context, *ClientSession, *ToolListChangedParams) - PromptListChangedHandler func(context.Context, *ClientSession, *PromptListChangedParams) - ResourceListChangedHandler func(context.Context, *ClientSession, *ResourceListChangedParams) - ResourceUpdatedHandler func(context.Context, *ClientSession, *ResourceUpdatedNotificationParams) - LoggingMessageHandler func(context.Context, *ClientSession, *LoggingMessageParams) - ProgressNotificationHandler func(context.Context, *ClientSession, *ProgressNotificationParams) + ToolListChangedHandler func(context.Context, *ClientRequest[*ToolListChangedParams]) + PromptListChangedHandler func(context.Context, *ClientRequest[*PromptListChangedParams]) + ResourceListChangedHandler func(context.Context, *ClientRequest[*ResourceListChangedParams]) + ResourceUpdatedHandler func(context.Context, *ClientRequest[*ResourceUpdatedNotificationParams]) + LoggingMessageHandler func(context.Context, *ClientRequest[*LoggingMessageParams]) + ProgressNotificationHandler func(context.Context, *ClientRequest[*ProgressNotificationParams]) // If non-zero, defines an interval for regular "ping" requests. // If the peer fails to respond to pings originating from the keepalive check, // the session is automatically closed. @@ -130,7 +130,8 @@ func (c *Client) Connect(ctx context.Context, t Transport, _ *ClientSessionOptio ClientInfo: c.impl, Capabilities: caps, } - res, err := handleSend[*InitializeResult](ctx, cs, methodInitialize, params) + req := &ClientRequest[*InitializeParams]{Session: cs, Params: params} + res, err := handleSend[*InitializeResult](ctx, methodInitialize, req) if err != nil { _ = cs.Close() return nil, err @@ -142,7 +143,8 @@ func (c *Client) Connect(ctx context.Context, t Transport, _ *ClientSessionOptio if hc, ok := cs.mcpConn.(clientConnection); ok { hc.sessionUpdated(cs.state) } - if err := handleNotify(ctx, cs, notificationInitialized, &InitializedParams{}); err != nil { + req2 := &ClientRequest[*InitializedParams]{Session: cs, Params: &InitializedParams{}} + if err := handleNotify(ctx, notificationInitialized, req2); err != nil { _ = cs.Close() return nil, err } @@ -216,7 +218,7 @@ func (c *Client) AddRoots(roots ...*Root) { if len(roots) == 0 { return } - c.changeAndNotify(notificationRootsListChanged, &RootsListChangedParams{}, + changeAndNotify(c, notificationRootsListChanged, &RootsListChangedParams{}, func() bool { c.roots.add(roots...); return true }) } @@ -225,14 +227,14 @@ func (c *Client) AddRoots(roots ...*Root) { // It is not an error to remove a nonexistent root. // TODO: notification func (c *Client) RemoveRoots(uris ...string) { - c.changeAndNotify(notificationRootsListChanged, &RootsListChangedParams{}, + changeAndNotify(c, notificationRootsListChanged, &RootsListChangedParams{}, func() bool { return c.roots.remove(uris...) }) } // changeAndNotify is called when a feature is added or removed. // It calls change, which should do the work and report whether a change actually occurred. // If there was a change, it notifies a snapshot of the sessions. -func (c *Client) changeAndNotify(notification string, params Params, change func() bool) { +func changeAndNotify[P Params](c *Client, notification string, params P, change func() bool) { var sessions []*ClientSession // Lock for the change, but not for the notification. c.mu.Lock() @@ -243,7 +245,7 @@ func (c *Client) changeAndNotify(notification string, params Params, change func notifySessions(sessions, notification, params) } -func (c *Client) listRoots(_ context.Context, _ *ClientSession, _ *ListRootsParams) (*ListRootsResult, error) { +func (c *Client) listRoots(_ context.Context, req *ClientRequest[*ListRootsParams]) (*ListRootsResult, error) { c.mu.Lock() defer c.mu.Unlock() roots := slices.Collect(c.roots.all()) @@ -255,12 +257,12 @@ func (c *Client) listRoots(_ context.Context, _ *ClientSession, _ *ListRootsPara }, nil } -func (c *Client) createMessage(ctx context.Context, cs *ClientSession, params *CreateMessageParams) (*CreateMessageResult, error) { +func (c *Client) createMessage(ctx context.Context, req *ClientRequest[*CreateMessageParams]) (*CreateMessageResult, error) { if c.opts.CreateMessageHandler == nil { // TODO: wrap or annotate this error? Pick a standard code? return nil, jsonrpc2.NewError(CodeUnsupportedMethod, "client does not support CreateMessage") } - return c.opts.CreateMessageHandler(ctx, cs, params) + return c.opts.CreateMessageHandler(ctx, req) } // AddSendingMiddleware wraps the current sending method handler using the provided @@ -272,7 +274,7 @@ func (c *Client) createMessage(ctx context.Context, cs *ClientSession, params *C // // Sending middleware is called when a request is sent. It is useful for tasks // such as tracing, metrics, and adding progress tokens. -func (c *Client) AddSendingMiddleware(middleware ...Middleware[*ClientSession]) { +func (c *Client) AddSendingMiddleware(middleware ...Middleware) { c.mu.Lock() defer c.mu.Unlock() addMiddleware(&c.sendingMethodHandler_, middleware) @@ -287,7 +289,7 @@ func (c *Client) AddSendingMiddleware(middleware ...Middleware[*ClientSession]) // // Receiving middleware is called when a request is received. It is useful for tasks // such as authentication, request logging and metrics. -func (c *Client) AddReceivingMiddleware(middleware ...Middleware[*ClientSession]) { +func (c *Client) AddReceivingMiddleware(middleware ...Middleware) { c.mu.Lock() defer c.mu.Unlock() addMiddleware(&c.receivingMethodHandler_, middleware) @@ -299,16 +301,16 @@ func (c *Client) AddReceivingMiddleware(middleware ...Middleware[*ClientSession] // TODO(rfindley): actually load and validate the protocol schema, rather than // curating these method flags. var clientMethodInfos = map[string]methodInfo{ - methodComplete: newMethodInfo(sessionMethod((*ClientSession).Complete), 0), - methodPing: newMethodInfo(sessionMethod((*ClientSession).ping), missingParamsOK), - methodListRoots: newMethodInfo(clientMethod((*Client).listRoots), missingParamsOK), - methodCreateMessage: newMethodInfo(clientMethod((*Client).createMessage), 0), - notificationToolListChanged: newMethodInfo(clientMethod((*Client).callToolChangedHandler), notification|missingParamsOK), - notificationPromptListChanged: newMethodInfo(clientMethod((*Client).callPromptChangedHandler), notification|missingParamsOK), - notificationResourceListChanged: newMethodInfo(clientMethod((*Client).callResourceChangedHandler), notification|missingParamsOK), - notificationResourceUpdated: newMethodInfo(clientMethod((*Client).callResourceUpdatedHandler), notification|missingParamsOK), - notificationLoggingMessage: newMethodInfo(clientMethod((*Client).callLoggingHandler), notification), - notificationProgress: newMethodInfo(sessionMethod((*ClientSession).callProgressNotificationHandler), notification), + methodComplete: newClientMethodInfo(clientSessionMethod((*ClientSession).Complete), 0), + methodPing: newClientMethodInfo(clientSessionMethod((*ClientSession).ping), missingParamsOK), + methodListRoots: newClientMethodInfo(clientMethod((*Client).listRoots), missingParamsOK), + methodCreateMessage: newClientMethodInfo(clientMethod((*Client).createMessage), 0), + notificationToolListChanged: newClientMethodInfo(clientMethod((*Client).callToolChangedHandler), notification|missingParamsOK), + notificationPromptListChanged: newClientMethodInfo(clientMethod((*Client).callPromptChangedHandler), notification|missingParamsOK), + notificationResourceListChanged: newClientMethodInfo(clientMethod((*Client).callResourceChangedHandler), notification|missingParamsOK), + notificationResourceUpdated: newClientMethodInfo(clientMethod((*Client).callResourceUpdatedHandler), notification|missingParamsOK), + notificationLoggingMessage: newClientMethodInfo(clientMethod((*Client).callLoggingHandler), notification), + notificationProgress: newClientMethodInfo(clientSessionMethod((*ClientSession).callProgressNotificationHandler), notification), } func (cs *ClientSession) sendingMethodInfos() map[string]methodInfo { @@ -342,25 +344,29 @@ func (*ClientSession) ping(context.Context, *PingParams) (*emptyResult, error) { return &emptyResult{}, nil } +func newClientRequest[P Params](cs *ClientSession, params P) *ClientRequest[P] { + return &ClientRequest[P]{Session: cs, Params: params} +} + // Ping makes an MCP "ping" request to the server. func (cs *ClientSession) Ping(ctx context.Context, params *PingParams) error { - _, err := handleSend[*emptyResult](ctx, cs, methodPing, orZero[Params](params)) + _, err := handleSend[*emptyResult](ctx, methodPing, newClientRequest(cs, orZero[Params](params))) return err } // ListPrompts lists prompts that are currently available on the server. func (cs *ClientSession) ListPrompts(ctx context.Context, params *ListPromptsParams) (*ListPromptsResult, error) { - return handleSend[*ListPromptsResult](ctx, cs, methodListPrompts, orZero[Params](params)) + return handleSend[*ListPromptsResult](ctx, methodListPrompts, newClientRequest(cs, orZero[Params](params))) } // GetPrompt gets a prompt from the server. func (cs *ClientSession) GetPrompt(ctx context.Context, params *GetPromptParams) (*GetPromptResult, error) { - return handleSend[*GetPromptResult](ctx, cs, methodGetPrompt, orZero[Params](params)) + return handleSend[*GetPromptResult](ctx, methodGetPrompt, newClientRequest(cs, orZero[Params](params))) } // ListTools lists tools that are currently available on the server. func (cs *ClientSession) ListTools(ctx context.Context, params *ListToolsParams) (*ListToolsResult, error) { - return handleSend[*ListToolsResult](ctx, cs, methodListTools, orZero[Params](params)) + return handleSend[*ListToolsResult](ctx, methodListTools, newClientRequest(cs, orZero[Params](params))) } // CallTool calls the tool with the given name and arguments. @@ -373,72 +379,87 @@ func (cs *ClientSession) CallTool(ctx context.Context, params *CallToolParams) ( // Avoid sending nil over the wire. params.Arguments = map[string]any{} } - return handleSend[*CallToolResult](ctx, cs, methodCallTool, params) + return handleSend[*CallToolResult](ctx, methodCallTool, newClientRequest(cs, orZero[Params](params))) } func (cs *ClientSession) SetLevel(ctx context.Context, params *SetLevelParams) error { - _, err := handleSend[*emptyResult](ctx, cs, methodSetLevel, orZero[Params](params)) + _, err := handleSend[*emptyResult](ctx, methodSetLevel, newClientRequest(cs, orZero[Params](params))) return err } // ListResources lists the resources that are currently available on the server. func (cs *ClientSession) ListResources(ctx context.Context, params *ListResourcesParams) (*ListResourcesResult, error) { - return handleSend[*ListResourcesResult](ctx, cs, methodListResources, orZero[Params](params)) + return handleSend[*ListResourcesResult](ctx, methodListResources, newClientRequest(cs, orZero[Params](params))) } // ListResourceTemplates lists the resource templates that are currently available on the server. func (cs *ClientSession) ListResourceTemplates(ctx context.Context, params *ListResourceTemplatesParams) (*ListResourceTemplatesResult, error) { - return handleSend[*ListResourceTemplatesResult](ctx, cs, methodListResourceTemplates, orZero[Params](params)) + return handleSend[*ListResourceTemplatesResult](ctx, methodListResourceTemplates, newClientRequest(cs, orZero[Params](params))) } // ReadResource asks the server to read a resource and return its contents. func (cs *ClientSession) ReadResource(ctx context.Context, params *ReadResourceParams) (*ReadResourceResult, error) { - return handleSend[*ReadResourceResult](ctx, cs, methodReadResource, orZero[Params](params)) + return handleSend[*ReadResourceResult](ctx, methodReadResource, newClientRequest(cs, orZero[Params](params))) } func (cs *ClientSession) Complete(ctx context.Context, params *CompleteParams) (*CompleteResult, error) { - return handleSend[*CompleteResult](ctx, cs, methodComplete, orZero[Params](params)) + return handleSend[*CompleteResult](ctx, methodComplete, newClientRequest(cs, orZero[Params](params))) } // Subscribe sends a "resources/subscribe" request to the server, asking for // notifications when the specified resource changes. func (cs *ClientSession) Subscribe(ctx context.Context, params *SubscribeParams) error { - _, err := handleSend[*emptyResult](ctx, cs, methodSubscribe, orZero[Params](params)) + _, err := handleSend[*emptyResult](ctx, methodSubscribe, newClientRequest(cs, orZero[Params](params))) return err } // Unsubscribe sends a "resources/unsubscribe" request to the server, cancelling // a previous subscription. func (cs *ClientSession) Unsubscribe(ctx context.Context, params *UnsubscribeParams) error { - _, err := handleSend[*emptyResult](ctx, cs, methodUnsubscribe, orZero[Params](params)) + _, err := handleSend[*emptyResult](ctx, methodUnsubscribe, newClientRequest(cs, orZero[Params](params))) return err } -func (c *Client) callToolChangedHandler(ctx context.Context, s *ClientSession, params *ToolListChangedParams) (Result, error) { - return callNotificationHandler(ctx, c.opts.ToolListChangedHandler, s, params) +func (c *Client) callToolChangedHandler(ctx context.Context, req *ClientRequest[*ToolListChangedParams]) (Result, error) { + if h := c.opts.ToolListChangedHandler; h != nil { + h(ctx, req) + } + return nil, nil } -func (c *Client) callPromptChangedHandler(ctx context.Context, s *ClientSession, params *PromptListChangedParams) (Result, error) { - return callNotificationHandler(ctx, c.opts.PromptListChangedHandler, s, params) +func (c *Client) callPromptChangedHandler(ctx context.Context, req *ClientRequest[*PromptListChangedParams]) (Result, error) { + if h := c.opts.PromptListChangedHandler; h != nil { + h(ctx, req) + } + return nil, nil } -func (c *Client) callResourceChangedHandler(ctx context.Context, s *ClientSession, params *ResourceListChangedParams) (Result, error) { - return callNotificationHandler(ctx, c.opts.ResourceListChangedHandler, s, params) +func (c *Client) callResourceChangedHandler(ctx context.Context, req *ClientRequest[*ResourceListChangedParams]) (Result, error) { + if h := c.opts.ResourceListChangedHandler; h != nil { + h(ctx, req) + } + return nil, nil } -func (c *Client) callResourceUpdatedHandler(ctx context.Context, s *ClientSession, params *ResourceUpdatedNotificationParams) (Result, error) { - return callNotificationHandler(ctx, c.opts.ResourceUpdatedHandler, s, params) +func (c *Client) callResourceUpdatedHandler(ctx context.Context, req *ClientRequest[*ResourceUpdatedNotificationParams]) (Result, error) { + if h := c.opts.ResourceUpdatedHandler; h != nil { + h(ctx, req) + } + return nil, nil } -func (c *Client) callLoggingHandler(ctx context.Context, cs *ClientSession, params *LoggingMessageParams) (Result, error) { +func (c *Client) callLoggingHandler(ctx context.Context, req *ClientRequest[*LoggingMessageParams]) (Result, error) { if h := c.opts.LoggingMessageHandler; h != nil { - h(ctx, cs, params) + h(ctx, req) } return nil, nil } func (cs *ClientSession) callProgressNotificationHandler(ctx context.Context, params *ProgressNotificationParams) (Result, error) { - return callNotificationHandler(ctx, cs.client.opts.ProgressNotificationHandler, cs, params) + if h := cs.client.opts.ProgressNotificationHandler; h != nil { + h(ctx, clientRequestFor(cs, params)) + } + return nil, nil } // NotifyProgress sends a progress notification from the client to the server @@ -446,7 +467,7 @@ func (cs *ClientSession) callProgressNotificationHandler(ctx context.Context, pa // This can be used if the client is performing a long-running task that was // initiated by the server func (cs *ClientSession) NotifyProgress(ctx context.Context, params *ProgressNotificationParams) error { - return handleNotify(ctx, cs, notificationProgress, params) + return handleNotify(ctx, notificationProgress, newClientRequest(cs, orZero[Params](params))) } // Tools provides an iterator for all tools available on the server, diff --git a/mcp/example_middleware_test.go b/mcp/example_middleware_test.go index c91250c..56f7428 100644 --- a/mcp/example_middleware_test.go +++ b/mcp/example_middleware_test.go @@ -29,36 +29,35 @@ func Example_loggingMiddleware() { }, })) - loggingMiddleware := func(next mcp.MethodHandler[*mcp.ServerSession]) mcp.MethodHandler[*mcp.ServerSession] { + loggingMiddleware := func(next mcp.MethodHandler) mcp.MethodHandler { return func( ctx context.Context, - session *mcp.ServerSession, method string, - params mcp.Params, + req mcp.Request, ) (mcp.Result, error) { logger.Info("MCP method started", "method", method, - "session_id", session.ID(), - "has_params", params != nil, + "session_id", req.GetSession().ID(), + "has_params", req.GetParams() != nil, ) start := time.Now() - result, err := next(ctx, session, method, params) + result, err := next(ctx, method, req) duration := time.Since(start) if err != nil { logger.Error("MCP method failed", "method", method, - "session_id", session.ID(), + "session_id", req.GetSession().ID(), "duration_ms", duration.Milliseconds(), "err", err, ) } else { logger.Info("MCP method completed", "method", method, - "session_id", session.ID(), + "session_id", req.GetSession().ID(), "duration_ms", duration.Milliseconds(), "has_result", result != nil, ) @@ -90,10 +89,9 @@ func Example_loggingMiddleware() { }, func( ctx context.Context, - ss *mcp.ServerSession, - params *mcp.CallToolParamsFor[map[string]any], + req *mcp.ServerRequest[*mcp.CallToolParamsFor[map[string]any]], ) (*mcp.CallToolResultFor[any], error) { - name, ok := params.Arguments["name"].(string) + name, ok := req.Params.Arguments["name"].(string) if !ok { return nil, fmt.Errorf("name parameter is required and must be a string") } diff --git a/mcp/example_progress_test.go b/mcp/example_progress_test.go index 6c771e2..304c838 100644 --- a/mcp/example_progress_test.go +++ b/mcp/example_progress_test.go @@ -21,11 +21,11 @@ func Example_progressMiddleware() { _ = c } -func addProgressToken[S mcp.Session](h mcp.MethodHandler[S]) mcp.MethodHandler[S] { - return func(ctx context.Context, s S, method string, params mcp.Params) (result mcp.Result, err error) { - if rp, ok := params.(mcp.RequestParams); ok { +func addProgressToken[S mcp.Session](h mcp.MethodHandler) mcp.MethodHandler { + return func(ctx context.Context, method string, req mcp.Request) (result mcp.Result, err error) { + if rp, ok := req.GetParams().(mcp.RequestParams); ok { rp.SetProgressToken(nextProgressToken.Add(1)) } - return h(ctx, s, method, params) + return h(ctx, method, req) } } diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 9e3eccd..58b0377 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -32,11 +32,11 @@ type hiParams struct { // TODO(jba): after schemas are stateless (WIP), this can be a variable. func greetTool() *Tool { return &Tool{Name: "greet", Description: "say hi"} } -func sayHi(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[hiParams]) (*CallToolResultFor[any], error) { - if err := ss.Ping(ctx, nil); err != nil { +func sayHi(ctx context.Context, req *ServerRequest[*CallToolParamsFor[hiParams]]) (*CallToolResultFor[any], error) { + if err := req.Session.Ping(ctx, nil); err != nil { return nil, fmt.Errorf("ping failed: %v", err) } - return &CallToolResultFor[any]{Content: []Content{&TextContent{Text: "hi " + params.Arguments.Name}}}, nil + return &CallToolResultFor[any]{Content: []Content{&TextContent{Text: "hi " + req.Params.Arguments.Name}}}, nil } var codeReviewPrompt = &Prompt{ @@ -73,16 +73,20 @@ func TestEndToEnd(t *testing.T) { } sopts := &ServerOptions{ - InitializedHandler: func(context.Context, *ServerSession, *InitializedParams) { notificationChans["initialized"] <- 0 }, - RootsListChangedHandler: func(context.Context, *ServerSession, *RootsListChangedParams) { notificationChans["roots"] <- 0 }, - ProgressNotificationHandler: func(context.Context, *ServerSession, *ProgressNotificationParams) { + InitializedHandler: func(context.Context, *ServerRequest[*InitializedParams]) { + notificationChans["initialized"] <- 0 + }, + RootsListChangedHandler: func(context.Context, *ServerRequest[*RootsListChangedParams]) { + notificationChans["roots"] <- 0 + }, + ProgressNotificationHandler: func(context.Context, *ServerRequest[*ProgressNotificationParams]) { notificationChans["progress_server"] <- 0 }, - SubscribeHandler: func(context.Context, *ServerSession, *SubscribeParams) error { + SubscribeHandler: func(context.Context, *ServerRequest[*SubscribeParams]) error { notificationChans["subscribe"] <- 0 return nil }, - UnsubscribeHandler: func(context.Context, *ServerSession, *UnsubscribeParams) error { + UnsubscribeHandler: func(context.Context, *ServerRequest[*UnsubscribeParams]) error { notificationChans["unsubscribe"] <- 0 return nil }, @@ -93,7 +97,7 @@ func TestEndToEnd(t *testing.T) { Description: "say hi", }, sayHi) s.AddTool(&Tool{Name: "fail", InputSchema: &jsonschema.Schema{}}, - func(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error) { + func(context.Context, *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResult, error) { return nil, errTestFailure }) s.AddPrompt(codeReviewPrompt, codReviewPromptHandler) @@ -124,19 +128,25 @@ func TestEndToEnd(t *testing.T) { loggingMessages := make(chan *LoggingMessageParams, 100) // big enough for all logging opts := &ClientOptions{ - CreateMessageHandler: func(context.Context, *ClientSession, *CreateMessageParams) (*CreateMessageResult, error) { + CreateMessageHandler: func(context.Context, *ClientRequest[*CreateMessageParams]) (*CreateMessageResult, error) { return &CreateMessageResult{Model: "aModel", Content: &TextContent{}}, nil }, - ToolListChangedHandler: func(context.Context, *ClientSession, *ToolListChangedParams) { notificationChans["tools"] <- 0 }, - PromptListChangedHandler: func(context.Context, *ClientSession, *PromptListChangedParams) { notificationChans["prompts"] <- 0 }, - ResourceListChangedHandler: func(context.Context, *ClientSession, *ResourceListChangedParams) { notificationChans["resources"] <- 0 }, - LoggingMessageHandler: func(_ context.Context, _ *ClientSession, lm *LoggingMessageParams) { - loggingMessages <- lm + ToolListChangedHandler: func(context.Context, *ClientRequest[*ToolListChangedParams]) { + notificationChans["tools"] <- 0 }, - ProgressNotificationHandler: func(context.Context, *ClientSession, *ProgressNotificationParams) { + PromptListChangedHandler: func(context.Context, *ClientRequest[*PromptListChangedParams]) { + notificationChans["prompts"] <- 0 + }, + ResourceListChangedHandler: func(context.Context, *ClientRequest[*ResourceListChangedParams]) { + notificationChans["resources"] <- 0 + }, + LoggingMessageHandler: func(_ context.Context, req *ClientRequest[*LoggingMessageParams]) { + loggingMessages <- req.Params + }, + ProgressNotificationHandler: func(context.Context, *ClientRequest[*ProgressNotificationParams]) { notificationChans["progress_client"] <- 0 }, - ResourceUpdatedHandler: func(context.Context, *ClientSession, *ResourceUpdatedNotificationParams) { + ResourceUpdatedHandler: func(context.Context, *ClientRequest[*ResourceUpdatedNotificationParams]) { notificationChans["resource_updated"] <- 0 }, } @@ -500,8 +510,8 @@ var embeddedResources = map[string]string{ "info": "This is the MCP test server.", } -func handleEmbeddedResource(_ context.Context, _ *ServerSession, params *ReadResourceParams) (*ReadResourceResult, error) { - u, err := url.Parse(params.URI) +func handleEmbeddedResource(_ context.Context, req *ServerRequest[*ReadResourceParams]) (*ReadResourceResult, error) { + u, err := url.Parse(req.Params.URI) if err != nil { return nil, err } @@ -515,7 +525,7 @@ func handleEmbeddedResource(_ context.Context, _ *ServerSession, params *ReadRes } return &ReadResourceResult{ Contents: []*ResourceContents{ - {URI: params.URI, MIMEType: "text/plain", Text: text}, + {URI: req.Params.URI, MIMEType: "text/plain", Text: text}, }, }, nil } @@ -637,7 +647,7 @@ func TestCancellation(t *testing.T) { cancelled = make(chan struct{}, 1) // don't block the request ) - slowRequest := func(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[map[string]any]) (*CallToolResult, error) { + slowRequest := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResult, error) { start <- struct{}{} select { case <-ctx.Done(): @@ -816,17 +826,17 @@ func TestNoJSONNull(t *testing.T) { // traceCalls creates a middleware function that prints the method before and after each call // with the given prefix. -func traceCalls[S Session](w io.Writer, prefix string) Middleware[S] { - return func(h MethodHandler[S]) MethodHandler[S] { - return func(ctx context.Context, sess S, method string, params Params) (Result, error) { +func traceCalls[S Session](w io.Writer, prefix string) Middleware { + return func(h MethodHandler) MethodHandler { + return func(ctx context.Context, method string, req Request) (Result, error) { fmt.Fprintf(w, "%s >%s\n", prefix, method) defer fmt.Fprintf(w, "%s <%s\n", prefix, method) - return h(ctx, sess, method, params) + return h(ctx, method, req) } } } -func nopHandler(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error) { +func nopHandler(context.Context, *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResult, error) { return nil, nil } diff --git a/mcp/resource.go b/mcp/resource.go index 590e067..5445715 100644 --- a/mcp/resource.go +++ b/mcp/resource.go @@ -35,7 +35,7 @@ type serverResourceTemplate struct { // A ResourceHandler is a function that reads a resource. // It will be called when the client calls [ClientSession.ReadResource]. // If it cannot find the resource, it should return the result of calling [ResourceNotFoundError]. -type ResourceHandler func(context.Context, *ServerSession, *ReadResourceParams) (*ReadResourceResult, error) +type ResourceHandler func(context.Context, *ServerRequest[*ReadResourceParams]) (*ReadResourceResult, error) // ResourceNotFoundError returns an error indicating that a resource being read could // not be found. diff --git a/mcp/server.go b/mcp/server.go index 89f3b6c..e39372d 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -41,8 +41,8 @@ type Server struct { resources *featureSet[*serverResource] resourceTemplates *featureSet[*serverResourceTemplate] sessions []*ServerSession - sendingMethodHandler_ MethodHandler[*ServerSession] - receivingMethodHandler_ MethodHandler[*ServerSession] + sendingMethodHandler_ MethodHandler + receivingMethodHandler_ MethodHandler resourceSubscriptions map[string]map[*ServerSession]bool // uri -> session -> bool } @@ -51,24 +51,24 @@ type ServerOptions struct { // Optional instructions for connected clients. Instructions string // If non-nil, called when "notifications/initialized" is received. - InitializedHandler func(context.Context, *ServerSession, *InitializedParams) + InitializedHandler func(context.Context, *ServerRequest[*InitializedParams]) // PageSize is the maximum number of items to return in a single page for // list methods (e.g. ListTools). PageSize int // If non-nil, called when "notifications/roots/list_changed" is received. - RootsListChangedHandler func(context.Context, *ServerSession, *RootsListChangedParams) + RootsListChangedHandler func(context.Context, *ServerRequest[*RootsListChangedParams]) // If non-nil, called when "notifications/progress" is received. - ProgressNotificationHandler func(context.Context, *ServerSession, *ProgressNotificationParams) + ProgressNotificationHandler func(context.Context, *ServerRequest[*ProgressNotificationParams]) // If non-nil, called when "completion/complete" is received. - CompletionHandler func(context.Context, *ServerSession, *CompleteParams) (*CompleteResult, error) + CompletionHandler func(context.Context, *ServerRequest[*CompleteParams]) (*CompleteResult, error) // If non-zero, defines an interval for regular "ping" requests. // If the peer fails to respond to pings originating from the keepalive check, // the session is automatically closed. KeepAlive time.Duration // Function called when a client session subscribes to a resource. - SubscribeHandler func(context.Context, *ServerSession, *SubscribeParams) error + SubscribeHandler func(context.Context, *ServerRequest[*SubscribeParams]) error // Function called when a client session unsubscribes from a resource. - UnsubscribeHandler func(context.Context, *ServerSession, *UnsubscribeParams) error + UnsubscribeHandler func(context.Context, *ServerRequest[*UnsubscribeParams]) error // If true, advertises the prompts capability during initialization, // even if no prompts have been registered. HasPrompts bool @@ -258,11 +258,11 @@ func (s *Server) capabilities() *serverCapabilities { return caps } -func (s *Server) complete(ctx context.Context, ss *ServerSession, params *CompleteParams) (Result, error) { +func (s *Server) complete(ctx context.Context, req *ServerRequest[*CompleteParams]) (Result, error) { if s.opts.CompletionHandler == nil { return nil, jsonrpc2.ErrMethodNotFound } - return s.opts.CompletionHandler(ctx, ss, params) + return s.opts.CompletionHandler(ctx, req) } // changeAndNotify is called when a feature is added or removed. @@ -287,13 +287,13 @@ func (s *Server) Sessions() iter.Seq[*ServerSession] { return slices.Values(clients) } -func (s *Server) listPrompts(_ context.Context, _ *ServerSession, params *ListPromptsParams) (*ListPromptsResult, error) { +func (s *Server) listPrompts(_ context.Context, req *ServerRequest[*ListPromptsParams]) (*ListPromptsResult, error) { s.mu.Lock() defer s.mu.Unlock() - if params == nil { - params = &ListPromptsParams{} + if req.Params == nil { + req.Params = &ListPromptsParams{} } - return paginateList(s.prompts, s.opts.PageSize, params, &ListPromptsResult{}, func(res *ListPromptsResult, prompts []*serverPrompt) { + return paginateList(s.prompts, s.opts.PageSize, req.Params, &ListPromptsResult{}, func(res *ListPromptsResult, prompts []*serverPrompt) { res.Prompts = []*Prompt{} // avoid JSON null for _, p := range prompts { res.Prompts = append(res.Prompts, p.prompt) @@ -301,24 +301,24 @@ func (s *Server) listPrompts(_ context.Context, _ *ServerSession, params *ListPr }) } -func (s *Server) getPrompt(ctx context.Context, ss *ServerSession, params *GetPromptParams) (*GetPromptResult, error) { +func (s *Server) getPrompt(ctx context.Context, req *ServerRequest[*GetPromptParams]) (*GetPromptResult, error) { s.mu.Lock() - prompt, ok := s.prompts.get(params.Name) + prompt, ok := s.prompts.get(req.Params.Name) s.mu.Unlock() if !ok { // TODO: surface the error code over the wire, instead of flattening it into the string. - return nil, fmt.Errorf("%s: unknown prompt %q", jsonrpc2.ErrInvalidParams, params.Name) + return nil, fmt.Errorf("%s: unknown prompt %q", jsonrpc2.ErrInvalidParams, req.Params.Name) } - return prompt.handler(ctx, ss, params) + return prompt.handler(ctx, req.Session, req.Params) } -func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListToolsParams) (*ListToolsResult, error) { +func (s *Server) listTools(_ context.Context, req *ServerRequest[*ListToolsParams]) (*ListToolsResult, error) { s.mu.Lock() defer s.mu.Unlock() - if params == nil { - params = &ListToolsParams{} + if req.Params == nil { + req.Params = &ListToolsParams{} } - return paginateList(s.tools, s.opts.PageSize, params, &ListToolsResult{}, func(res *ListToolsResult, tools []*serverTool) { + return paginateList(s.tools, s.opts.PageSize, req.Params, &ListToolsResult{}, func(res *ListToolsResult, tools []*serverTool) { res.Tools = []*Tool{} // avoid JSON null for _, t := range tools { res.Tools = append(res.Tools, t.tool) @@ -326,23 +326,23 @@ func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListTool }) } -func (s *Server) callTool(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) { +func (s *Server) callTool(ctx context.Context, req *ServerRequest[*CallToolParamsFor[json.RawMessage]]) (*CallToolResult, error) { s.mu.Lock() - st, ok := s.tools.get(params.Name) + st, ok := s.tools.get(req.Params.Name) s.mu.Unlock() if !ok { - return nil, fmt.Errorf("%s: unknown tool %q", jsonrpc2.ErrInvalidParams, params.Name) + return nil, fmt.Errorf("%s: unknown tool %q", jsonrpc2.ErrInvalidParams, req.Params.Name) } - return st.handler(ctx, cc, params) + return st.handler(ctx, req) } -func (s *Server) listResources(_ context.Context, _ *ServerSession, params *ListResourcesParams) (*ListResourcesResult, error) { +func (s *Server) listResources(_ context.Context, req *ServerRequest[*ListResourcesParams]) (*ListResourcesResult, error) { s.mu.Lock() defer s.mu.Unlock() - if params == nil { - params = &ListResourcesParams{} + if req.Params == nil { + req.Params = &ListResourcesParams{} } - return paginateList(s.resources, s.opts.PageSize, params, &ListResourcesResult{}, func(res *ListResourcesResult, resources []*serverResource) { + return paginateList(s.resources, s.opts.PageSize, req.Params, &ListResourcesResult{}, func(res *ListResourcesResult, resources []*serverResource) { res.Resources = []*Resource{} // avoid JSON null for _, r := range resources { res.Resources = append(res.Resources, r.resource) @@ -350,13 +350,13 @@ func (s *Server) listResources(_ context.Context, _ *ServerSession, params *List }) } -func (s *Server) listResourceTemplates(_ context.Context, _ *ServerSession, params *ListResourceTemplatesParams) (*ListResourceTemplatesResult, error) { +func (s *Server) listResourceTemplates(_ context.Context, req *ServerRequest[*ListResourceTemplatesParams]) (*ListResourceTemplatesResult, error) { s.mu.Lock() defer s.mu.Unlock() - if params == nil { - params = &ListResourceTemplatesParams{} + if req.Params == nil { + req.Params = &ListResourceTemplatesParams{} } - return paginateList(s.resourceTemplates, s.opts.PageSize, params, &ListResourceTemplatesResult{}, + return paginateList(s.resourceTemplates, s.opts.PageSize, req.Params, &ListResourceTemplatesResult{}, func(res *ListResourceTemplatesResult, rts []*serverResourceTemplate) { res.ResourceTemplates = []*ResourceTemplate{} // avoid JSON null for _, rt := range rts { @@ -365,8 +365,8 @@ func (s *Server) listResourceTemplates(_ context.Context, _ *ServerSession, para }) } -func (s *Server) readResource(ctx context.Context, ss *ServerSession, params *ReadResourceParams) (*ReadResourceResult, error) { - uri := params.URI +func (s *Server) readResource(ctx context.Context, req *ServerRequest[*ReadResourceParams]) (*ReadResourceResult, error) { + uri := req.Params.URI // Look up the resource URI in the lists of resources and resource templates. // This is a security check as well as an information lookup. handler, mimeType, ok := s.lookupResourceHandler(uri) @@ -375,7 +375,7 @@ func (s *Server) readResource(ctx context.Context, ss *ServerSession, params *Re // Treat an unregistered resource the same as a registered one that couldn't be found. return nil, ResourceNotFoundError(uri) } - res, err := handler(ctx, ss, params) + res, err := handler(ctx, req) if err != nil { return nil, err } @@ -430,11 +430,11 @@ func fileResourceHandler(dir string) ResourceHandler { if err != nil { panic(err) } - return func(ctx context.Context, ss *ServerSession, params *ReadResourceParams) (_ *ReadResourceResult, err error) { - defer util.Wrapf(&err, "reading resource %s", params.URI) + return func(ctx context.Context, req *ServerRequest[*ReadResourceParams]) (_ *ReadResourceResult, err error) { + defer util.Wrapf(&err, "reading resource %s", req.Params.URI) // TODO: use a memoizing API here. - rootRes, err := ss.ListRoots(ctx, nil) + rootRes, err := req.Session.ListRoots(ctx, nil) if err != nil { return nil, fmt.Errorf("listing roots: %w", err) } @@ -442,13 +442,13 @@ func fileResourceHandler(dir string) ResourceHandler { if err != nil { return nil, err } - data, err := readFileResource(params.URI, dirFilepath, roots) + data, err := readFileResource(req.Params.URI, dirFilepath, roots) if err != nil { return nil, err } // TODO(jba): figure out mime type. Omit for now: Server.readResource will fill it in. return &ReadResourceResult{Contents: []*ResourceContents{ - {URI: params.URI, Blob: data}, + {URI: req.Params.URI, Blob: data}, }}, nil } } @@ -465,39 +465,39 @@ func (s *Server) ResourceUpdated(ctx context.Context, params *ResourceUpdatedNot return nil } -func (s *Server) subscribe(ctx context.Context, ss *ServerSession, params *SubscribeParams) (*emptyResult, error) { +func (s *Server) subscribe(ctx context.Context, req *ServerRequest[*SubscribeParams]) (*emptyResult, error) { if s.opts.SubscribeHandler == nil { return nil, fmt.Errorf("%w: server does not support resource subscriptions", jsonrpc2.ErrMethodNotFound) } - if err := s.opts.SubscribeHandler(ctx, ss, params); err != nil { + if err := s.opts.SubscribeHandler(ctx, req); err != nil { return nil, err } s.mu.Lock() defer s.mu.Unlock() - if s.resourceSubscriptions[params.URI] == nil { - s.resourceSubscriptions[params.URI] = make(map[*ServerSession]bool) + if s.resourceSubscriptions[req.Params.URI] == nil { + s.resourceSubscriptions[req.Params.URI] = make(map[*ServerSession]bool) } - s.resourceSubscriptions[params.URI][ss] = true + s.resourceSubscriptions[req.Params.URI][req.Session] = true return &emptyResult{}, nil } -func (s *Server) unsubscribe(ctx context.Context, ss *ServerSession, params *UnsubscribeParams) (*emptyResult, error) { +func (s *Server) unsubscribe(ctx context.Context, req *ServerRequest[*UnsubscribeParams]) (*emptyResult, error) { if s.opts.UnsubscribeHandler == nil { return nil, jsonrpc2.ErrMethodNotFound } - if err := s.opts.UnsubscribeHandler(ctx, ss, params); err != nil { + if err := s.opts.UnsubscribeHandler(ctx, req); err != nil { return nil, err } s.mu.Lock() defer s.mu.Unlock() - if subscribedSessions, ok := s.resourceSubscriptions[params.URI]; ok { - delete(subscribedSessions, ss) + if subscribedSessions, ok := s.resourceSubscriptions[req.Params.URI]; ok { + delete(subscribedSessions, req.Session) if len(subscribedSessions) == 0 { - delete(s.resourceSubscriptions, params.URI) + delete(s.resourceSubscriptions, req.Params.URI) } } @@ -611,15 +611,24 @@ func (ss *ServerSession) initialized(ctx context.Context, params *InitializedPar if wasInitd { return nil, fmt.Errorf("duplicate %q received", notificationInitialized) } - return callNotificationHandler(ctx, ss.server.opts.InitializedHandler, ss, params) + if h := ss.server.opts.InitializedHandler; h != nil { + h(ctx, serverRequestFor(ss, params)) + } + return nil, nil } -func (s *Server) callRootsListChangedHandler(ctx context.Context, ss *ServerSession, params *RootsListChangedParams) (Result, error) { - return callNotificationHandler(ctx, s.opts.RootsListChangedHandler, ss, params) +func (s *Server) callRootsListChangedHandler(ctx context.Context, req *ServerRequest[*RootsListChangedParams]) (Result, error) { + if h := s.opts.RootsListChangedHandler; h != nil { + h(ctx, req) + } + return nil, nil } -func (ss *ServerSession) callProgressNotificationHandler(ctx context.Context, params *ProgressNotificationParams) (Result, error) { - return callNotificationHandler(ctx, ss.server.opts.ProgressNotificationHandler, ss, params) +func (ss *ServerSession) callProgressNotificationHandler(ctx context.Context, p *ProgressNotificationParams) (Result, error) { + if h := ss.server.opts.ProgressNotificationHandler; h != nil { + h(ctx, serverRequestFor(ss, p)) + } + return nil, nil } // NotifyProgress sends a progress notification from the server to the client @@ -627,7 +636,11 @@ func (ss *ServerSession) callProgressNotificationHandler(ctx context.Context, pa // This is typically used to report on the status of a long-running request // that was initiated by the client. func (ss *ServerSession) NotifyProgress(ctx context.Context, params *ProgressNotificationParams) error { - return handleNotify(ctx, ss, notificationProgress, params) + return handleNotify(ctx, notificationProgress, newServerRequest(ss, orZero[Params](params))) +} + +func newServerRequest[P Params](ss *ServerSession, params P) *ServerRequest[P] { + return &ServerRequest[P]{Session: ss, Params: params} } // A ServerSession is a logical connection from a single MCP client. Its @@ -665,18 +678,18 @@ func (ss *ServerSession) ID() string { // Ping pings the client. func (ss *ServerSession) Ping(ctx context.Context, params *PingParams) error { - _, err := handleSend[*emptyResult](ctx, ss, methodPing, orZero[Params](params)) + _, err := handleSend[*emptyResult](ctx, methodPing, newServerRequest(ss, orZero[Params](params))) return err } // ListRoots lists the client roots. func (ss *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) { - return handleSend[*ListRootsResult](ctx, ss, methodListRoots, orZero[Params](params)) + return handleSend[*ListRootsResult](ctx, methodListRoots, newServerRequest(ss, orZero[Params](params))) } // CreateMessage sends a sampling request to the client. func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessageParams) (*CreateMessageResult, error) { - return handleSend[*CreateMessageResult](ctx, ss, methodCreateMessage, orZero[Params](params)) + return handleSend[*CreateMessageResult](ctx, methodCreateMessage, newServerRequest(ss, orZero[Params](params))) } // Log sends a log message to the client. @@ -695,7 +708,7 @@ func (ss *ServerSession) Log(ctx context.Context, params *LoggingMessageParams) if compareLevels(params.Level, logLevel) < 0 { return nil } - return handleNotify(ctx, ss, notificationLoggingMessage, params) + return handleNotify(ctx, notificationLoggingMessage, newServerRequest(ss, orZero[Params](params))) } // AddSendingMiddleware wraps the current sending method handler using the provided @@ -707,7 +720,7 @@ func (ss *ServerSession) Log(ctx context.Context, params *LoggingMessageParams) // // Sending middleware is called when a request is sent. It is useful for tasks // such as tracing, metrics, and adding progress tokens. -func (s *Server) AddSendingMiddleware(middleware ...Middleware[*ServerSession]) { +func (s *Server) AddSendingMiddleware(middleware ...Middleware) { s.mu.Lock() defer s.mu.Unlock() addMiddleware(&s.sendingMethodHandler_, middleware) @@ -722,7 +735,7 @@ func (s *Server) AddSendingMiddleware(middleware ...Middleware[*ServerSession]) // // Receiving middleware is called when a request is received. It is useful for tasks // such as authentication, request logging and metrics. -func (s *Server) AddReceivingMiddleware(middleware ...Middleware[*ServerSession]) { +func (s *Server) AddReceivingMiddleware(middleware ...Middleware) { s.mu.Lock() defer s.mu.Unlock() addMiddleware(&s.receivingMethodHandler_, middleware) @@ -734,22 +747,22 @@ func (s *Server) AddReceivingMiddleware(middleware ...Middleware[*ServerSession] // TODO(rfindley): actually load and validate the protocol schema, rather than // curating these method flags. var serverMethodInfos = map[string]methodInfo{ - methodComplete: newMethodInfo(serverMethod((*Server).complete), 0), - methodInitialize: newMethodInfo(sessionMethod((*ServerSession).initialize), 0), - methodPing: newMethodInfo(sessionMethod((*ServerSession).ping), missingParamsOK), - methodListPrompts: newMethodInfo(serverMethod((*Server).listPrompts), missingParamsOK), - methodGetPrompt: newMethodInfo(serverMethod((*Server).getPrompt), 0), - methodListTools: newMethodInfo(serverMethod((*Server).listTools), missingParamsOK), - methodCallTool: newMethodInfo(serverMethod((*Server).callTool), 0), - methodListResources: newMethodInfo(serverMethod((*Server).listResources), missingParamsOK), - methodListResourceTemplates: newMethodInfo(serverMethod((*Server).listResourceTemplates), missingParamsOK), - methodReadResource: newMethodInfo(serverMethod((*Server).readResource), 0), - methodSetLevel: newMethodInfo(sessionMethod((*ServerSession).setLevel), 0), - methodSubscribe: newMethodInfo(serverMethod((*Server).subscribe), 0), - methodUnsubscribe: newMethodInfo(serverMethod((*Server).unsubscribe), 0), - notificationInitialized: newMethodInfo(sessionMethod((*ServerSession).initialized), notification|missingParamsOK), - notificationRootsListChanged: newMethodInfo(serverMethod((*Server).callRootsListChangedHandler), notification|missingParamsOK), - notificationProgress: newMethodInfo(sessionMethod((*ServerSession).callProgressNotificationHandler), notification), + methodComplete: newServerMethodInfo(serverMethod((*Server).complete), 0), + methodInitialize: newServerMethodInfo(serverSessionMethod((*ServerSession).initialize), 0), + methodPing: newServerMethodInfo(serverSessionMethod((*ServerSession).ping), missingParamsOK), + methodListPrompts: newServerMethodInfo(serverMethod((*Server).listPrompts), missingParamsOK), + methodGetPrompt: newServerMethodInfo(serverMethod((*Server).getPrompt), 0), + methodListTools: newServerMethodInfo(serverMethod((*Server).listTools), missingParamsOK), + methodCallTool: newServerMethodInfo(serverMethod((*Server).callTool), 0), + methodListResources: newServerMethodInfo(serverMethod((*Server).listResources), missingParamsOK), + methodListResourceTemplates: newServerMethodInfo(serverMethod((*Server).listResourceTemplates), missingParamsOK), + methodReadResource: newServerMethodInfo(serverMethod((*Server).readResource), 0), + methodSetLevel: newServerMethodInfo(serverSessionMethod((*ServerSession).setLevel), 0), + methodSubscribe: newServerMethodInfo(serverMethod((*Server).subscribe), 0), + methodUnsubscribe: newServerMethodInfo(serverMethod((*Server).unsubscribe), 0), + notificationInitialized: newServerMethodInfo(serverSessionMethod((*ServerSession).initialized), notification|missingParamsOK), + notificationRootsListChanged: newServerMethodInfo(serverMethod((*Server).callRootsListChangedHandler), notification|missingParamsOK), + notificationProgress: newServerMethodInfo(serverSessionMethod((*ServerSession).callProgressNotificationHandler), notification), } func (ss *ServerSession) sendingMethodInfos() map[string]methodInfo { return clientMethodInfos } @@ -757,15 +770,17 @@ func (ss *ServerSession) sendingMethodInfos() map[string]methodInfo { return cli func (ss *ServerSession) receivingMethodInfos() map[string]methodInfo { return serverMethodInfos } func (ss *ServerSession) sendingMethodHandler() methodHandler { - ss.server.mu.Lock() - defer ss.server.mu.Unlock() - return ss.server.sendingMethodHandler_ + s := ss.server + s.mu.Lock() + defer s.mu.Unlock() + return s.sendingMethodHandler_ } func (ss *ServerSession) receivingMethodHandler() methodHandler { - ss.server.mu.Lock() - defer ss.server.mu.Unlock() - return ss.server.receivingMethodHandler_ + s := ss.server + s.mu.Lock() + defer s.mu.Unlock() + return s.receivingMethodHandler_ } // getConn implements [session.getConn]. @@ -808,13 +823,14 @@ func (ss *ServerSession) initialize(ctx context.Context, params *InitializeParam version = latestProtocolVersion } + s := ss.server return &InitializeResult{ // TODO(rfindley): alter behavior when falling back to an older version: // reject unsupported features. ProtocolVersion: version, - Capabilities: ss.server.capabilities(), - Instructions: ss.server.opts.Instructions, - ServerInfo: ss.server.impl, + Capabilities: s.capabilities(), + Instructions: s.opts.Instructions, + ServerInfo: s.impl, }, nil } diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index 241008e..f735b84 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -16,10 +16,10 @@ type SayHiParams struct { Name string `json:"name"` } -func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[SayHiParams]) (*mcp.CallToolResultFor[any], error) { +func SayHi(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[SayHiParams]]) (*mcp.CallToolResultFor[any], error) { return &mcp.CallToolResultFor[any]{ Content: []mcp.Content{ - &mcp.TextContent{Text: "Hi " + params.Arguments.Name}, + &mcp.TextContent{Text: "Hi " + req.Params.Arguments.Name}, }, }, nil } diff --git a/mcp/server_test.go b/mcp/server_test.go index 6415dec..202ab5d 100644 --- a/mcp/server_test.go +++ b/mcp/server_test.go @@ -281,10 +281,10 @@ func TestServerCapabilities(t *testing.T) { s.AddResourceTemplate(&ResourceTemplate{URITemplate: "file:///rt"}, nil) }, serverOpts: ServerOptions{ - SubscribeHandler: func(ctx context.Context, _ *ServerSession, sp *SubscribeParams) error { + SubscribeHandler: func(context.Context, *ServerRequest[*SubscribeParams]) error { return nil }, - UnsubscribeHandler: func(ctx context.Context, _ *ServerSession, up *UnsubscribeParams) error { + UnsubscribeHandler: func(context.Context, *ServerRequest[*UnsubscribeParams]) error { return nil }, }, @@ -307,7 +307,7 @@ func TestServerCapabilities(t *testing.T) { name: "With completions", configureServer: func(s *Server) {}, serverOpts: ServerOptions{ - CompletionHandler: func(ctx context.Context, ss *ServerSession, params *CompleteParams) (*CompleteResult, error) { + CompletionHandler: func(context.Context, *ServerRequest[*CompleteParams]) (*CompleteResult, error) { return nil, nil }, }, @@ -325,13 +325,13 @@ func TestServerCapabilities(t *testing.T) { s.AddTool(tool, nil) }, serverOpts: ServerOptions{ - SubscribeHandler: func(ctx context.Context, _ *ServerSession, sp *SubscribeParams) error { + SubscribeHandler: func(context.Context, *ServerRequest[*SubscribeParams]) error { return nil }, - UnsubscribeHandler: func(ctx context.Context, _ *ServerSession, up *UnsubscribeParams) error { + UnsubscribeHandler: func(context.Context, *ServerRequest[*UnsubscribeParams]) error { return nil }, - CompletionHandler: func(ctx context.Context, ss *ServerSession, params *CompleteParams) (*CompleteResult, error) { + CompletionHandler: func(context.Context, *ServerRequest[*CompleteParams]) (*CompleteResult, error) { return nil, nil }, }, diff --git a/mcp/shared.go b/mcp/shared.go index e368864..ca06221 100644 --- a/mcp/shared.go +++ b/mcp/shared.go @@ -36,7 +36,7 @@ var supportedProtocolVersions = []string{ // A MethodHandler handles MCP messages. // For methods, exactly one of the return values must be nil. // For notifications, both must be nil. -type MethodHandler[S Session] func(ctx context.Context, _ S, method string, params Params) (result Result, err error) +type MethodHandler func(ctx context.Context, method string, req Request) (result Result, err error) // A methodHandler is a MethodHandler[Session] for some session. // We need to give up type safety here, or we will end up with a type cycle somewhere @@ -46,7 +46,6 @@ type methodHandler any // MethodHandler[*ClientSession] | MethodHandler[*ServerS // A Session is either a [ClientSession] or a [ServerSession]. type Session interface { - *ClientSession | *ServerSession // ID returns the session ID, or the empty string if there is none. ID() string @@ -58,29 +57,29 @@ type Session interface { } // Middleware is a function from [MethodHandler] to [MethodHandler]. -type Middleware[S Session] func(MethodHandler[S]) MethodHandler[S] +type Middleware func(MethodHandler) MethodHandler // addMiddleware wraps the handler in the middleware functions. -func addMiddleware[S Session](handlerp *MethodHandler[S], middleware []Middleware[S]) { +func addMiddleware(handlerp *MethodHandler, middleware []Middleware) { for _, m := range slices.Backward(middleware) { *handlerp = m(*handlerp) } } -func defaultSendingMethodHandler[S Session](ctx context.Context, session S, method string, params Params) (Result, error) { - info, ok := session.sendingMethodInfos()[method] +func defaultSendingMethodHandler[S Session](ctx context.Context, method string, req Request) (Result, error) { + info, ok := req.GetSession().sendingMethodInfos()[method] if !ok { // This can be called from user code, with an arbitrary value for method. return nil, jsonrpc2.ErrNotHandled } // Notifications don't have results. if strings.HasPrefix(method, "notifications/") { - return nil, session.getConn().Notify(ctx, method, params) + return nil, req.GetSession().getConn().Notify(ctx, method, req.GetParams()) } // Create the result to unmarshal into. // The concrete type of the result is the return type of the receiving function. res := info.newResult() - if err := call(ctx, session.getConn(), method, params, res); err != nil { + if err := call(ctx, req.GetSession().getConn(), method, req.GetParams(), res); err != nil { return nil, err } return res, nil @@ -95,16 +94,16 @@ func orZero[T any, P *U, U any](p P) T { return any(p).(T) } -func handleNotify[S Session](ctx context.Context, session S, method string, params Params) error { - mh := session.sendingMethodHandler().(MethodHandler[S]) - _, err := mh(ctx, session, method, params) +func handleNotify(ctx context.Context, method string, req Request) error { + mh := req.GetSession().sendingMethodHandler().(MethodHandler) + _, err := mh(ctx, method, req) return err } -func handleSend[R Result, S Session](ctx context.Context, s S, method string, params Params) (R, error) { - mh := s.sendingMethodHandler().(MethodHandler[S]) +func handleSend[R Result](ctx context.Context, method string, req Request) (R, error) { + mh := req.GetSession().sendingMethodHandler().(MethodHandler) // mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol. - res, err := mh(ctx, s, method, params) + res, err := mh(ctx, method, req) if err != nil { var z R return z, err @@ -113,28 +112,29 @@ func handleSend[R Result, S Session](ctx context.Context, s S, method string, pa } // defaultReceivingMethodHandler is the initial MethodHandler for servers and clients, before being wrapped by middleware. -func defaultReceivingMethodHandler[S Session](ctx context.Context, session S, method string, params Params) (Result, error) { - info, ok := session.receivingMethodInfos()[method] +func defaultReceivingMethodHandler[S Session](ctx context.Context, method string, req Request) (Result, error) { + info, ok := req.GetSession().receivingMethodInfos()[method] if !ok { // This can be called from user code, with an arbitrary value for method. return nil, jsonrpc2.ErrNotHandled } - return info.handleMethod.(MethodHandler[S])(ctx, session, method, params) + return info.handleMethod.(MethodHandler)(ctx, method, req) } -func handleReceive[S Session](ctx context.Context, session S, req *jsonrpc.Request) (Result, error) { - info, err := checkRequest(req, session.receivingMethodInfos()) +func handleReceive[S Session](ctx context.Context, session S, jreq *jsonrpc.Request) (Result, error) { + info, err := checkRequest(jreq, session.receivingMethodInfos()) if err != nil { return nil, err } - params, err := info.unmarshalParams(req.Params) + params, err := info.unmarshalParams(jreq.Params) if err != nil { - return nil, fmt.Errorf("handling '%s': %w", req.Method, err) + return nil, fmt.Errorf("handling '%s': %w", jreq.Method, err) } - mh := session.receivingMethodHandler().(MethodHandler[S]) + mh := session.receivingMethodHandler().(MethodHandler) + req := info.newRequest(session, params) // mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol. - res, err := mh(ctx, session, req.Method, params) + res, err := mh(ctx, jreq.Method, req) if err != nil { return nil, err } @@ -179,6 +179,7 @@ type methodInfo struct { // Unmarshal params from the wire into a Params struct. // Used on the receive side. unmarshalParams func(json.RawMessage) (Params, error) + newRequest func(Session, Params) Request // Run the code when a call to the method is received. // Used on the receive side. handleMethod methodHandler @@ -194,7 +195,10 @@ type methodInfo struct { // - R: results // A typedMethodHandler is like a MethodHandler, but with type information. -type typedMethodHandler[S Session, P Params, R Result] func(context.Context, S, P) (R, error) +type ( + typedClientMethodHandler[P Params, R Result] func(context.Context, *ClientRequest[P]) (R, error) + typedServerMethodHandler[P Params, R Result] func(context.Context, *ServerRequest[P]) (R, error) +) type paramsPtr[T any] interface { *T @@ -208,11 +212,45 @@ const ( missingParamsOK // params may be missing or null ) +func newClientMethodInfo[P paramsPtr[T], R Result, T any](d typedClientMethodHandler[P, R], flags methodFlags) methodInfo { + mi := newMethodInfo[P, R](flags) + mi.newRequest = func(s Session, p Params) Request { + r := &ClientRequest[P]{Session: s.(*ClientSession)} + if p != nil { + r.Params = p.(P) + } + return r + } + mi.handleMethod = MethodHandler(func(ctx context.Context, _ string, req Request) (Result, error) { + return d(ctx, req.(*ClientRequest[P])) + }) + return mi +} + +func newServerMethodInfo[P paramsPtr[T], R Result, T any](d typedServerMethodHandler[P, R], flags methodFlags) methodInfo { + mi := newMethodInfo[P, R](flags) + mi.newRequest = func(s Session, p Params) Request { + r := &ServerRequest[P]{Session: s.(*ServerSession)} + if p != nil { + r.Params = p.(P) + } + return r + } + mi.handleMethod = MethodHandler(func(ctx context.Context, _ string, req Request) (Result, error) { + rf := &ServerRequest[P]{Session: req.GetSession().(*ServerSession)} + if req.GetParams() != nil { + rf.Params = req.GetParams().(P) + } + return d(ctx, rf) + }) + return mi +} + // newMethodInfo creates a methodInfo from a typedMethodHandler. // // If isRequest is set, the method is treated as a request rather than a // notification. -func newMethodInfo[S Session, P paramsPtr[T], R Result, T any](d typedMethodHandler[S, P, R], flags methodFlags) methodInfo { +func newMethodInfo[P paramsPtr[T], R Result, T any](flags methodFlags) methodInfo { return methodInfo{ flags: flags, unmarshalParams: func(m json.RawMessage) (Params, error) { @@ -234,12 +272,6 @@ func newMethodInfo[S Session, P paramsPtr[T], R Result, T any](d typedMethodHand } return orZero[Params](p), nil }, - handleMethod: MethodHandler[S](func(ctx context.Context, session S, _ string, params Params) (Result, error) { - if params == nil { - return d(ctx, session, nil) - } - return d(ctx, session, params.(P)) - }), // newResult is used on the send side, to construct the value to unmarshal the result into. // R is a pointer to a result struct. There is no way to "unpointer" it without reflection. // TODO(jba): explore generic approaches to this, perhaps by treating R in @@ -250,26 +282,33 @@ func newMethodInfo[S Session, P paramsPtr[T], R Result, T any](d typedMethodHand // serverMethod is glue for creating a typedMethodHandler from a method on Server. func serverMethod[P Params, R Result]( - f func(*Server, context.Context, *ServerSession, P) (R, error), -) typedMethodHandler[*ServerSession, P, R] { - return func(ctx context.Context, ss *ServerSession, p P) (R, error) { - return f(ss.server, ctx, ss, p) + f func(*Server, context.Context, *ServerRequest[P]) (R, error), +) typedServerMethodHandler[P, R] { + return func(ctx context.Context, req *ServerRequest[P]) (R, error) { + return f(req.Session.server, ctx, req) } } // clientMethod is glue for creating a typedMethodHandler from a method on Client. func clientMethod[P Params, R Result]( - f func(*Client, context.Context, *ClientSession, P) (R, error), -) typedMethodHandler[*ClientSession, P, R] { - return func(ctx context.Context, cs *ClientSession, p P) (R, error) { - return f(cs.client, ctx, cs, p) + f func(*Client, context.Context, *ClientRequest[P]) (R, error), +) typedClientMethodHandler[P, R] { + return func(ctx context.Context, req *ClientRequest[P]) (R, error) { + return f(req.Session.client, ctx, req) } } -// sessionMethod is glue for creating a typedMethodHandler from a method on ServerSession. -func sessionMethod[S Session, P Params, R Result](f func(S, context.Context, P) (R, error)) typedMethodHandler[S, P, R] { - return func(ctx context.Context, sess S, p P) (R, error) { - return f(sess, ctx, p) +// serverSessionMethod is glue for creating a typedServerMethodHandler from a method on ServerSession. +func serverSessionMethod[P Params, R Result](f func(*ServerSession, context.Context, P) (R, error)) typedServerMethodHandler[P, R] { + return func(ctx context.Context, req *ServerRequest[P]) (R, error) { + return f(req.GetSession().(*ServerSession), ctx, req.Params) + } +} + +// clientSessionMethod is glue for creating a typedMethodHandler from a method on ServerSession. +func clientSessionMethod[P Params, R Result](f func(*ClientSession, context.Context, P) (R, error)) typedClientMethodHandler[P, R] { + return func(ctx context.Context, req *ClientRequest[P]) (R, error) { + return f(req.GetSession().(*ClientSession), ctx, req.Params) } } @@ -282,16 +321,9 @@ const ( CodeUnsupportedMethod = -31001 ) -func callNotificationHandler[S Session, P any](ctx context.Context, h func(context.Context, S, *P), sess S, params *P) (Result, error) { - if h != nil { - h(ctx, sess, params) - } - return nil, nil -} - // notifySessions calls Notify on all the sessions. // Should be called on a copy of the peer sessions. -func notifySessions[S Session](sessions []S, method string, params Params) { +func notifySessions[S Session, P Params](sessions []S, method string, params P) { if sessions == nil { return } @@ -299,13 +331,25 @@ func notifySessions[S Session](sessions []S, method string, params Params) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() for _, s := range sessions { - if err := handleNotify(ctx, s, method, params); err != nil { + req := newRequest(s, params) + if err := handleNotify(ctx, method, req); err != nil { // TODO(jba): surface this error better log.Printf("calling %s: %v", method, err) } } } +func newRequest[S Session, P Params](s S, p P) Request { + switch s := any(s).(type) { + case *ClientSession: + return &ClientRequest[P]{Session: s, Params: p} + case *ServerSession: + return &ServerRequest[P]{Session: s, Params: p} + default: + panic("bad session") + } +} + // Meta is additional metadata for requests, responses and other types. type Meta map[string]any @@ -335,6 +379,43 @@ func setProgressToken(p Params, pt any) { m[progressTokenKey] = pt } +// A Request is a method request with parameters and additional information, such as the session. +// Request is implemented by [*ClientRequest] and [*ServerRequest]. +type Request interface { + isRequest() + GetSession() Session + GetParams() Params +} + +// A ClientRequest is a request to a client. +type ClientRequest[P Params] struct { + Session *ClientSession + Params P +} + +// A ServerRequest is a request to a server. +type ServerRequest[P Params] struct { + Session *ServerSession + Params P +} + +func (*ClientRequest[P]) isRequest() {} +func (*ServerRequest[P]) isRequest() {} + +func (r *ClientRequest[P]) GetSession() Session { return r.Session } +func (r *ServerRequest[P]) GetSession() Session { return r.Session } + +func (r *ClientRequest[P]) GetParams() Params { return r.Params } +func (r *ServerRequest[P]) GetParams() Params { return r.Params } + +func serverRequestFor[P Params](s *ServerSession, p P) *ServerRequest[P] { + return &ServerRequest[P]{Session: s, Params: p} +} + +func clientRequestFor[P Params](s *ClientSession, p P) *ClientRequest[P] { + return &ClientRequest[P]{Session: s, Params: p} +} + // Params is a parameter (input) type for an MCP call or notification. type Params interface { // isParams discourages implementation of Params outside of this package. diff --git a/mcp/shared_test.go b/mcp/shared_test.go index 0aea194..01d1eff 100644 --- a/mcp/shared_test.go +++ b/mcp/shared_test.go @@ -23,7 +23,7 @@ func TestToolValidate(t *testing.T) { P *int `json:",omitempty"` } - dummyHandler := func(context.Context, *ServerSession, *CallToolParamsFor[req]) (*CallToolResultFor[any], error) { + dummyHandler := func(context.Context, *ServerRequest[*CallToolParamsFor[req]]) (*CallToolResultFor[any], error) { return nil, nil } @@ -73,8 +73,9 @@ func TestToolValidate(t *testing.T) { if err != nil { t.Fatal(err) } - _, err = st.handler(context.Background(), nil, - &CallToolParamsFor[json.RawMessage]{Arguments: json.RawMessage(raw)}) + _, err = st.handler(context.Background(), &ServerRequest[*CallToolParamsFor[json.RawMessage]]{ + Params: &CallToolParamsFor[json.RawMessage]{Arguments: json.RawMessage(raw)}, + }) if err == nil && tt.want != "" { t.Error("got success, wanted failure") } @@ -102,12 +103,12 @@ func TestNilParamsHandling(t *testing.T) { type TestResult = *CallToolResultFor[string] // Simple test handler - testHandler := func(ctx context.Context, ss *ServerSession, params TestParams) (TestResult, error) { - result := "processed: " + params.Arguments.Name + testHandler := func(ctx context.Context, req *ServerRequest[TestParams]) (TestResult, error) { + result := "processed: " + req.Params.Arguments.Name return &CallToolResultFor[string]{StructuredContent: result}, nil } - methodInfo := newMethodInfo(testHandler, missingParamsOK) + methodInfo := newServerMethodInfo(testHandler, missingParamsOK) // Helper function to test that unmarshalParams doesn't panic and handles nil gracefully mustNotPanic := func(t *testing.T, rawMsg json.RawMessage, expectNil bool) Params { @@ -183,11 +184,11 @@ func TestNilParamsEdgeCases(t *testing.T) { } type TestParams = *CallToolParamsFor[TestArgs] - testHandler := func(ctx context.Context, ss *ServerSession, params TestParams) (*CallToolResultFor[string], error) { + testHandler := func(context.Context, *ServerRequest[TestParams]) (*CallToolResultFor[string], error) { return &CallToolResultFor[string]{StructuredContent: "test"}, nil } - methodInfo := newMethodInfo(testHandler, missingParamsOK) + methodInfo := newServerMethodInfo(testHandler, missingParamsOK) // These should fail normally, not be treated as nil params invalidCases := []json.RawMessage{ @@ -209,7 +210,7 @@ func TestNilParamsEdgeCases(t *testing.T) { // Test that methods without missingParamsOK flag properly reject nil params t.Run("reject_when_params_required", func(t *testing.T) { - methodInfoStrict := newMethodInfo(testHandler, 0) // No missingParamsOK flag + methodInfoStrict := newServerMethodInfo(testHandler, 0) // No missingParamsOK flag testCases := []struct { name string diff --git a/mcp/sse_example_test.go b/mcp/sse_example_test.go index cf1e75d..b5dfdc5 100644 --- a/mcp/sse_example_test.go +++ b/mcp/sse_example_test.go @@ -18,10 +18,10 @@ type AddParams struct { X, Y int } -func Add(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[AddParams]) (*mcp.CallToolResultFor[any], error) { +func Add(ctx context.Context, req *mcp.ServerRequest[*mcp.CallToolParamsFor[AddParams]]) (*mcp.CallToolResultFor[any], error) { return &mcp.CallToolResultFor[any]{ Content: []mcp.Content{ - &mcp.TextContent{Text: fmt.Sprintf("%d", params.Arguments.X+params.Arguments.Y)}, + &mcp.TextContent{Text: fmt.Sprintf("%d", req.Params.Arguments.X+req.Params.Arguments.Y)}, }, }, nil } diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index ba7af9f..fd1dc3e 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -136,12 +136,12 @@ func TestClientReplay(t *testing.T) { serverReadyToKillProxy := make(chan struct{}) serverClosed := make(chan struct{}) server.AddTool(&Tool{Name: "multiMessageTool", InputSchema: &jsonschema.Schema{}}, - func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[map[string]any]) (*CallToolResult, error) { + func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResult, error) { go func() { bgCtx := context.Background() // Send the first two messages immediately. - ss.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg1"}) - ss.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg2"}) + req.Session.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg1"}) + req.Session.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg2"}) // Signal the test that it can now kill the proxy. close(serverReadyToKillProxy) @@ -149,8 +149,8 @@ func TestClientReplay(t *testing.T) { // These messages should be queued for replay by the server after // the client's connection drops. - ss.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg3"}) - ss.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg4"}) + req.Session.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg3"}) + req.Session.NotifyProgress(bgCtx, &ProgressNotificationParams{Message: "msg4"}) }() return &CallToolResult{}, nil }) @@ -170,8 +170,8 @@ func TestClientReplay(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() client := NewClient(testImpl, &ClientOptions{ - ProgressNotificationHandler: func(ctx context.Context, cc *ClientSession, params *ProgressNotificationParams) { - notifications <- params.Message + ProgressNotificationHandler: func(ctx context.Context, req *ClientRequest[*ProgressNotificationParams]) { + notifications <- req.Params.Message }, }) clientSession, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: proxy.URL}, nil) @@ -236,9 +236,10 @@ func TestServerInitiatedSSE(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() - client := NewClient(testImpl, &ClientOptions{ToolListChangedHandler: func(ctx context.Context, cc *ClientSession, params *ToolListChangedParams) { - notifications <- "toolListChanged" - }, + client := NewClient(testImpl, &ClientOptions{ + ToolListChangedHandler: func(context.Context, *ClientRequest[*ToolListChangedParams]) { + notifications <- "toolListChanged" + }, }) clientSession, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: httpServer.URL}, nil) if err != nil { @@ -246,7 +247,7 @@ func TestServerInitiatedSSE(t *testing.T) { } defer clientSession.Close() server.AddTool(&Tool{Name: "testTool", InputSchema: &jsonschema.Schema{}}, - func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[map[string]any]) (*CallToolResult, error) { + func(context.Context, *ServerRequest[*CallToolParamsFor[map[string]any]]) (*CallToolResult, error) { return &CallToolResult{}, nil }) receivedNotifications := readNotifications(t, ctx, notifications, 1) @@ -509,9 +510,9 @@ func TestStreamableServerTransport(t *testing.T) { // Create a server containing a single tool, which runs the test tool // behavior, if any. server := NewServer(&Implementation{Name: "testServer", Version: "v1.0.0"}, nil) - AddTool(server, &Tool{Name: "tool"}, func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[any]) (*CallToolResultFor[any], error) { + AddTool(server, &Tool{Name: "tool"}, func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[any]]) (*CallToolResultFor[any], error) { if test.tool != nil { - test.tool(t, ctx, ss) + test.tool(t, ctx, req.Session) } return &CallToolResultFor[any]{}, nil }) @@ -822,14 +823,15 @@ func TestEventID(t *testing.T) { }) } } + func TestStreamableStateless(t *testing.T) { // Test stateless mode behavior ctx := context.Background() // This version of sayHi doesn't make a ping request (we can't respond to // that request from our client). - sayHi := func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[hiParams]) (*CallToolResultFor[any], error) { - return &CallToolResultFor[any]{Content: []Content{&TextContent{Text: "hi " + params.Arguments.Name}}}, nil + sayHi := func(ctx context.Context, req *ServerRequest[*CallToolParamsFor[hiParams]]) (*CallToolResultFor[any], error) { + return &CallToolResultFor[any]{Content: []Content{&TextContent{Text: "hi " + req.Params.Arguments.Name}}}, nil } server := NewServer(testImpl, nil) AddTool(server, &Tool{Name: "greet", Description: "say hi"}, sayHi) diff --git a/mcp/tool.go b/mcp/tool.go index 234cd65..15f17e1 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -20,10 +20,12 @@ import ( type ToolHandler = ToolHandlerFor[map[string]any, any] // A ToolHandlerFor handles a call to tools/call with typed arguments and results. -type ToolHandlerFor[In, Out any] func(context.Context, *ServerSession, *CallToolParamsFor[In]) (*CallToolResultFor[Out], error) +type ToolHandlerFor[In, Out any] func(context.Context, *ServerRequest[*CallToolParamsFor[In]]) (*CallToolResultFor[Out], error) // A rawToolHandler is like a ToolHandler, but takes the arguments as as json.RawMessage. -type rawToolHandler = func(context.Context, *ServerSession, *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) +// Second arg is *Request[*ServerSession, *CallToolParamsFor[json.RawMessage]], but that creates +// a cycle. +type rawToolHandler = func(context.Context, any) (*CallToolResult, error) // A serverTool is a tool definition that is bound to a tool handler. type serverTool struct { @@ -48,20 +50,25 @@ func newServerTool[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*serverTool } } - st.handler = func(ctx context.Context, ss *ServerSession, rparams *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) { + st.handler = func(ctx context.Context, areq any) (*CallToolResult, error) { + req := areq.(*ServerRequest[*CallToolParamsFor[json.RawMessage]]) var args In - if rparams.Arguments != nil { - if err := unmarshalSchema(rparams.Arguments, st.inputResolved, &args); err != nil { + if req.Params.Arguments != nil { + if err := unmarshalSchema(req.Params.Arguments, st.inputResolved, &args); err != nil { return nil, err } } // TODO(jba): future-proof this copy. params := &CallToolParamsFor[In]{ - Meta: rparams.Meta, - Name: rparams.Name, + Meta: req.Params.Meta, + Name: req.Params.Name, Arguments: args, } - res, err := h(ctx, ss, params) + // TODO(jba): improve copy + res, err := h(ctx, &ServerRequest[*CallToolParamsFor[In]]{ + Session: req.Session, + Params: params, + }) // TODO(rfindley): investigate why server errors are embedded in this strange way, // rather than returned as jsonrpc2 server errors. if err != nil { diff --git a/mcp/tool_test.go b/mcp/tool_test.go index 52cac9f..609536c 100644 --- a/mcp/tool_test.go +++ b/mcp/tool_test.go @@ -16,7 +16,7 @@ import ( ) // testToolHandler is used for type inference in TestNewServerTool. -func testToolHandler[In, Out any](context.Context, *ServerSession, *CallToolParamsFor[In]) (*CallToolResultFor[Out], error) { +func testToolHandler[In, Out any](context.Context, *ServerRequest[*CallToolParamsFor[In]]) (*CallToolResultFor[Out], error) { panic("not implemented") }