diff --git a/mcp/client.go b/mcp/client.go index 3a41c53..341142e 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -319,27 +319,72 @@ func (c *Client) callLoggingHandler(ctx context.Context, cs *ClientSession, para // Tools provides an iterator for all tools available on the server, // automatically fetching pages and managing cursors. // The `params` argument can set the initial cursor. +// Iteration stops at the first encountered error, which will be yielded. func (cs *ClientSession) Tools(ctx context.Context, params *ListToolsParams) iter.Seq2[Tool, error] { - currentParams := &ListToolsParams{} - if params != nil { - *currentParams = *params + if params == nil { + params = &ListToolsParams{} } - return func(yield func(Tool, error) bool) { + return paginate(ctx, params, cs.ListTools, func(res *ListToolsResult) []*Tool { + return res.Tools + }) +} + +// Resources provides an iterator for all resources available on the server, +// automatically fetching pages and managing cursors. +// The `params` argument can set the initial cursor. +// Iteration stops at the first encountered error, which will be yielded. +func (cs *ClientSession) Resources(ctx context.Context, params *ListResourcesParams) iter.Seq2[Resource, error] { + if params == nil { + params = &ListResourcesParams{} + } + return paginate(ctx, params, cs.ListResources, func(res *ListResourcesResult) []*Resource { + return res.Resources + }) +} + +// Prompts provides an iterator for all prompts available on the server, +// automatically fetching pages and managing cursors. +// The `params` argument can set the initial cursor. +// Iteration stops at the first encountered error, which will be yielded. +func (cs *ClientSession) Prompts(ctx context.Context, params *ListPromptsParams) iter.Seq2[Prompt, error] { + if params == nil { + params = &ListPromptsParams{} + } + return paginate(ctx, params, cs.ListPrompts, func(res *ListPromptsResult) []*Prompt { + return res.Prompts + }) +} + +type ListParams interface { + // Returns a pointer to the param's Cursor field. + cursorPtr() *string +} + +type ListResult[T any] interface { + // Returns a pointer to the param's NextCursor field. + nextCursorPtr() *string +} + +// paginate is a generic helper function to provide a paginated iterator. +func paginate[P ListParams, R ListResult[E], E any](ctx context.Context, params P, listFunc func(context.Context, P) (R, error), items func(R) []*E) iter.Seq2[E, error] { + return func(yield func(E, error) bool) { for { - res, err := cs.ListTools(ctx, currentParams) + res, err := listFunc(ctx, params) if err != nil { - yield(Tool{}, err) + var zero E + yield(zero, err) return } - for _, t := range res.Tools { - if !yield(*t, nil) { + for _, r := range items(res) { + if !yield(*r, nil) { return } } - if res.NextCursor == "" { + nextCursorVal := res.nextCursorPtr() + if nextCursorVal == nil || *nextCursorVal == "" { return } - currentParams.Cursor = res.NextCursor + *params.cursorPtr() = *nextCursorVal } } } diff --git a/mcp/client_test.go b/mcp/client_test.go new file mode 100644 index 0000000..9d24a92 --- /dev/null +++ b/mcp/client_test.go @@ -0,0 +1,201 @@ +// Copyright 2025 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by an MIT-style +// license that can be found in the LICENSE file. + +package mcp + +import ( + "context" + "fmt" + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/google/go-cmp/cmp/cmpopts" + "github.com/modelcontextprotocol/go-sdk/jsonschema" +) + +type Item struct { + Name string + Value string +} + +type ListTestParams struct { + Cursor string +} + +func (p *ListTestParams) cursorPtr() *string { + return &p.Cursor +} + +type ListTestResult struct { + Items []*Item + NextCursor string +} + +func (r *ListTestResult) nextCursorPtr() *string { + return &r.NextCursor +} + +var allItems = []*Item{ + {"alpha", "val-A"}, + {"bravo", "val-B"}, + {"charlie", "val-C"}, + {"delta", "val-D"}, + {"echo", "val-E"}, + {"foxtrot", "val-F"}, + {"golf", "val-G"}, + {"hotel", "val-H"}, + {"india", "val-I"}, + {"juliet", "val-J"}, + {"kilo", "val-K"}, +} + +func toItemValueSlice(ptrSlice []*Item) []Item { + var valSlice []Item + for _, ptr := range ptrSlice { + valSlice = append(valSlice, *ptr) + } + return valSlice +} + +// generatePaginatedResults is a helper to create a sequence of mock responses for pagination. +// It simulates a server returning items in pages based on a given page size. +func generatePaginatedResults(all []*Item, pageSize int) []*ListTestResult { + if len(all) == 0 { + return []*ListTestResult{{Items: []*Item{}, NextCursor: ""}} + } + if pageSize <= 0 { + panic("pageSize must be greater than 0") + } + numPages := (len(all) + pageSize - 1) / pageSize // Ceiling division + var results []*ListTestResult + for i := range numPages { + startIndex := i * pageSize + endIndex := min(startIndex+pageSize, len(all)) // Use min to prevent out of bounds + nextCursor := "" + if endIndex < len(all) { // If there are more items after this page + nextCursor = fmt.Sprintf("cursor_%d", endIndex) + } + results = append(results, &ListTestResult{Items: all[startIndex:endIndex], NextCursor: nextCursor}) + } + return results +} + +func TestClientPaginateBasic(t *testing.T) { + ctx := context.Background() + testCases := []struct { + name string + results []*ListTestResult + mockError error + initialParams *ListTestParams + expected []Item + expectError bool + }{ + { + name: "SinglePageAllItems", + results: generatePaginatedResults(allItems, len(allItems)), + expected: toItemValueSlice(allItems), + }, + { + name: "MultiplePages", + results: generatePaginatedResults(allItems, 3), + expected: toItemValueSlice(allItems), + }, + { + name: "EmptyResults", + results: generatePaginatedResults([]*Item{}, 10), + expected: nil, + }, + { + name: "ListFuncReturnsErrorImmediately", + results: []*ListTestResult{{}}, + mockError: fmt.Errorf("API error on first call"), + expected: nil, + expectError: true, + }, + { + name: "InitialCursorProvided", + initialParams: &ListTestParams{Cursor: "cursor_2"}, + results: generatePaginatedResults(allItems[2:], 3), + expected: toItemValueSlice(allItems[2:]), + }, + { + name: "CursorBeyondAllItems", + initialParams: &ListTestParams{Cursor: "cursor_999"}, + results: []*ListTestResult{{Items: []*Item{}, NextCursor: ""}}, + expected: nil, + }, + } + + for _, tc := range testCases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + listFunc := func(ctx context.Context, params *ListTestParams) (*ListTestResult, error) { + if len(tc.results) == 0 { + t.Fatalf("listFunc called but no more results defined for test case %q", tc.name) + } + res := tc.results[0] + tc.results = tc.results[1:] + var err error + if tc.mockError != nil { + err = tc.mockError + } + return res, err + } + + params := tc.initialParams + if tc.initialParams == nil { + params = &ListTestParams{} + } + + var gotItems []Item + var iterationErr error + seq := paginate(ctx, params, listFunc, func(r *ListTestResult) []*Item { return r.Items }) + for item, err := range seq { + if err != nil { + iterationErr = err + break + } + gotItems = append(gotItems, item) + } + if tc.expectError { + if iterationErr == nil { + t.Errorf("paginate() expected an error during iteration, but got none") + } + } else { + if iterationErr != nil { + t.Errorf("paginate() got: %v, want: nil", iterationErr) + } + } + if diff := cmp.Diff(tc.expected, gotItems, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("paginate() mismatch (-want +got):\n%s", diff) + } + }) + } +} + +func TestClientPaginateVariousPageSizes(t *testing.T) { + ctx := context.Background() + for i := 1; i < len(allItems)+1; i++ { + testname := fmt.Sprintf("PageSize=%d", i) + t.Run(testname, func(t *testing.T) { + results := generatePaginatedResults(allItems, i) + listFunc := func(ctx context.Context, params *ListTestParams) (*ListTestResult, error) { + res := results[0] + results = results[1:] + return res, nil + } + var gotItems []Item + seq := paginate(ctx, &ListTestParams{}, listFunc, func(r *ListTestResult) []*Item { return r.Items }) + for item, err := range seq { + if err != nil { + t.Fatalf("paginate() unexpected error during iteration: %v", err) + } + gotItems = append(gotItems, item) + } + if diff := cmp.Diff(toItemValueSlice(allItems), gotItems, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("paginate() mismatch (-want +got):\n%s", diff) + } + }) + } +} diff --git a/mcp/generate.go b/mcp/generate.go index 7db7bf4..809b2ec 100644 --- a/mcp/generate.go +++ b/mcp/generate.go @@ -332,6 +332,13 @@ func writeDecl(configName string, config typeConfig, def *jsonschema.Schema, nam fmt.Fprintf(w, "\nfunc (x *%s) GetMeta() *Meta { return &x.Meta }", typeName) } + if _, ok := def.Properties["cursor"]; ok { + fmt.Fprintf(w, "\nfunc (x *%s) cursorPtr() *string { return &x.Cursor }", typeName) + } + if _, ok := def.Properties["nextCursor"]; ok { + fmt.Fprintf(w, "\nfunc (x *%s) nextCursorPtr() *string { return &x.NextCursor }", typeName) + } + return nil } diff --git a/mcp/protocol.go b/mcp/protocol.go index 7c9d39f..79965ab 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -214,7 +214,8 @@ type ListPromptsParams struct { Cursor string `json:"cursor,omitempty"` } -func (x *ListPromptsParams) GetMeta() *Meta { return &x.Meta } +func (x *ListPromptsParams) GetMeta() *Meta { return &x.Meta } +func (x *ListPromptsParams) cursorPtr() *string { return &x.Cursor } // The server's response to a prompts/list request from the client. type ListPromptsResult struct { @@ -227,7 +228,8 @@ type ListPromptsResult struct { Prompts []*Prompt `json:"prompts"` } -func (x *ListPromptsResult) GetMeta() *Meta { return &x.Meta } +func (x *ListPromptsResult) GetMeta() *Meta { return &x.Meta } +func (x *ListPromptsResult) nextCursorPtr() *string { return &x.NextCursor } type ListResourcesParams struct { // This property is reserved by the protocol to allow clients and servers to @@ -238,7 +240,8 @@ type ListResourcesParams struct { Cursor string `json:"cursor,omitempty"` } -func (x *ListResourcesParams) GetMeta() *Meta { return &x.Meta } +func (x *ListResourcesParams) GetMeta() *Meta { return &x.Meta } +func (x *ListResourcesParams) cursorPtr() *string { return &x.Cursor } // The server's response to a resources/list request from the client. type ListResourcesResult struct { @@ -251,7 +254,8 @@ type ListResourcesResult struct { Resources []*Resource `json:"resources"` } -func (x *ListResourcesResult) GetMeta() *Meta { return &x.Meta } +func (x *ListResourcesResult) GetMeta() *Meta { return &x.Meta } +func (x *ListResourcesResult) nextCursorPtr() *string { return &x.NextCursor } type ListRootsParams struct { // This property is reserved by the protocol to allow clients and servers to @@ -282,7 +286,8 @@ type ListToolsParams struct { Cursor string `json:"cursor,omitempty"` } -func (x *ListToolsParams) GetMeta() *Meta { return &x.Meta } +func (x *ListToolsParams) GetMeta() *Meta { return &x.Meta } +func (x *ListToolsParams) cursorPtr() *string { return &x.Cursor } // The server's response to a tools/list request from the client. type ListToolsResult struct { @@ -295,7 +300,8 @@ type ListToolsResult struct { Tools []*Tool `json:"tools"` } -func (x *ListToolsResult) GetMeta() *Meta { return &x.Meta } +func (x *ListToolsResult) GetMeta() *Meta { return &x.Meta } +func (x *ListToolsResult) nextCursorPtr() *string { return &x.NextCursor } // The severity of a log message. // diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index ded355a..91f7f1b 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -57,8 +57,8 @@ func ExampleServer() { } // createSessions creates and connects an in-memory client and server session for testing purposes. -func createSessions(ctx context.Context, opts *mcp.ServerOptions) (*mcp.ClientSession, *mcp.ServerSession, *mcp.Server) { - server := mcp.NewServer("server", "v0.0.1", opts) +func createSessions(ctx context.Context) (*mcp.ClientSession, *mcp.ServerSession, *mcp.Server) { + server := mcp.NewServer("server", "v0.0.1", nil) client := mcp.NewClient("client", "v0.0.1", nil) serverTransport, clientTransport := mcp.NewInMemoryTransports() serverSession, err := server.Connect(ctx, serverTransport) @@ -77,42 +77,33 @@ func TestListTools(t *testing.T) { toolB := mcp.NewTool("banana", "banana tool", SayHi) toolC := mcp.NewTool("cherry", "cherry tool", SayHi) tools := []*mcp.ServerTool{toolA, toolB, toolC} - wantListTools := []*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool} - wantIteratorTools := []mcp.Tool{*toolA.Tool, *toolB.Tool, *toolC.Tool} ctx := context.Background() + clientSession, serverSession, server := createSessions(ctx) + defer clientSession.Close() + defer serverSession.Close() + server.AddTools(tools...) t.Run("ListTools", func(t *testing.T) { - clientSession, serverSession, server := createSessions(ctx, nil) - defer clientSession.Close() - defer serverSession.Close() - server.AddTools(tools...) + wantTools := []*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool} res, err := clientSession.ListTools(ctx, nil) if err != nil { t.Fatal("ListTools() failed:", err) } - if diff := cmp.Diff(wantListTools, res.Tools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + if diff := cmp.Diff(wantTools, res.Tools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { t.Fatalf("ListTools() mismatch (-want +got):\n%s", diff) } }) t.Run("ToolsIterator", func(t *testing.T) { - for pageSize := range len(tools) + 1 { - testName := fmt.Sprintf("PageSize=%v", pageSize) - t.Run(testName, func(t *testing.T) { - clientSession, serverSession, server := createSessions(ctx, &mcp.ServerOptions{PageSize: pageSize}) - defer clientSession.Close() - defer serverSession.Close() - server.AddTools(tools...) - var gotTools []mcp.Tool - seq := clientSession.Tools(ctx, nil) - for tool, err := range seq { - if err != nil { - t.Fatalf("Tools(%s) failed: %v", testName, err) - } - gotTools = append(gotTools, tool) - } - if diff := cmp.Diff(wantIteratorTools, gotTools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { - t.Fatalf("Tools(%s) mismatch (-want +got):\n%s", testName, diff) - } - }) + wantTools := []mcp.Tool{*toolA.Tool, *toolB.Tool, *toolC.Tool} + var gotTools []mcp.Tool + seq := clientSession.Tools(ctx, nil) + for tool, err := range seq { + if err != nil { + t.Fatalf("Tools() failed: %v", err) + } + gotTools = append(gotTools, tool) + } + if diff := cmp.Diff(wantTools, gotTools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("Tools() mismatch (-want +got):\n%s", diff) } }) } @@ -122,19 +113,36 @@ func TestListResources(t *testing.T) { resourceB := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://banana"}} resourceC := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://cherry"}} resources := []*mcp.ServerResource{resourceA, resourceB, resourceC} - wantResource := []*mcp.Resource{resourceA.Resource, resourceB.Resource, resourceC.Resource} ctx := context.Background() - clientSession, serverSession, server := createSessions(ctx, nil) + clientSession, serverSession, server := createSessions(ctx) defer clientSession.Close() defer serverSession.Close() server.AddResources(resources...) - res, err := clientSession.ListResources(ctx, nil) - if err != nil { - t.Fatal("ListResources() failed:", err) - } - if diff := cmp.Diff(wantResource, res.Resources, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { - t.Fatalf("ListResources() mismatch (-want +got):\n%s", diff) - } + t.Run("ListResources", func(t *testing.T) { + wantResources := []*mcp.Resource{resourceA.Resource, resourceB.Resource, resourceC.Resource} + res, err := clientSession.ListResources(ctx, nil) + if err != nil { + t.Fatal("ListResources() failed:", err) + } + if diff := cmp.Diff(wantResources, res.Resources, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("ListResources() mismatch (-want +got):\n%s", diff) + } + }) + t.Run("ResourcesIterator", func(t *testing.T) { + wantResources := []mcp.Resource{*resourceA.Resource, *resourceB.Resource, *resourceC.Resource} + var gotResources []mcp.Resource + seq := clientSession.Resources(ctx, nil) + for resource, err := range seq { + if err != nil { + t.Fatalf("Resources() failed: %v", err) + } + gotResources = append(gotResources, resource) + } + if diff := cmp.Diff(wantResources, gotResources, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("Resources() mismatch (-want +got):\n%s", diff) + } + }) + } func TestListPrompts(t *testing.T) { @@ -142,17 +150,33 @@ func TestListPrompts(t *testing.T) { promptB := mcp.NewPrompt("banana", "banana prompt", testPromptHandler[struct{}]) promptC := mcp.NewPrompt("cherry", "cherry prompt", testPromptHandler[struct{}]) prompts := []*mcp.ServerPrompt{promptA, promptB, promptC} - wantPrompts := []*mcp.Prompt{promptA.Prompt, promptB.Prompt, promptC.Prompt} ctx := context.Background() - clientSession, serverSession, server := createSessions(ctx, nil) + clientSession, serverSession, server := createSessions(ctx) defer clientSession.Close() defer serverSession.Close() server.AddPrompts(prompts...) - res, err := clientSession.ListPrompts(ctx, nil) - if err != nil { - t.Fatal("ListPrompts() failed:", err) - } - if diff := cmp.Diff(wantPrompts, res.Prompts, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { - t.Fatalf("ListPrompts() mismatch (-want +got):\n%s", diff) - } + t.Run("ListPrompts", func(t *testing.T) { + wantPrompts := []*mcp.Prompt{promptA.Prompt, promptB.Prompt, promptC.Prompt} + res, err := clientSession.ListPrompts(ctx, nil) + if err != nil { + t.Fatal("ListPrompts() failed:", err) + } + if diff := cmp.Diff(wantPrompts, res.Prompts, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("ListPrompts() mismatch (-want +got):\n%s", diff) + } + }) + t.Run("PromptsIterator", func(t *testing.T) { + wantPrompts := []mcp.Prompt{*promptA.Prompt, *promptB.Prompt, *promptC.Prompt} + var gotPrompts []mcp.Prompt + seq := clientSession.Prompts(ctx, nil) + for prompt, err := range seq { + if err != nil { + t.Fatalf("Prompts() failed: %v", err) + } + gotPrompts = append(gotPrompts, prompt) + } + if diff := cmp.Diff(wantPrompts, gotPrompts, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("Prompts() mismatch (-want +got):\n%s", diff) + } + }) } diff --git a/mcp/server_test.go b/mcp/server_test.go index 11223c6..15250ca 100644 --- a/mcp/server_test.go +++ b/mcp/server_test.go @@ -41,7 +41,7 @@ func getCursor(input string) string { return cursor } -func TestPaginateBasic(t *testing.T) { +func TestServerPaginateBasic(t *testing.T) { testCases := []struct { name string initialItems []*TestItem @@ -170,7 +170,7 @@ func TestPaginateBasic(t *testing.T) { } } -func TestPaginateVariousPageSizes(t *testing.T) { +func TestServerPaginateVariousPageSizes(t *testing.T) { fs := newFeatureSet(func(t *TestItem) string { return t.Name }) fs.add(allTestItems...) // Try all possible page sizes, ensuring we get the correct list of items.