internal/mcp: add iterator methods resources and prompts

This CL streamlines the client side pagination logic by adding a generic
helper that handles pagination. This also simplifies the existing tools
iterator method.

Change-Id: If4304deda1ccd22d0c4580d76deaefc64e5dfeb9
Reviewed-on: https://go-review.googlesource.com/c/tools/+/676975
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>
Auto-Submit: Sam Thanawalla <samthanawalla@google.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>

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 {
-	if params != nil {
+		params = &ListToolsParams{}
-		*currentParams = *params
 	}
-	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 {
+			for _, r := range items(res) {
-				if !yield(*t, nil) {
+				if !yield(*r, nil) {
 					return
 				}
 			}
-			if res.NextCursor == "" {
+			nextCursorVal := res.nextCursorPtr()
+			if nextCursorVal == nil || *nextCursorVal == "" {
 				return
 			}
-			currentParams.Cursor = res.NextCursor
+			*params.cursorPtr() = *nextCursorVal
 		}
 	}
 }

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)
+			}
+		})
+	}
+}

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
 }
 

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.
 //

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) {
+func createSessions(ctx context.Context) (*mcp.ClientSession, *mcp.ServerSession, *mcp.Server) {
-	server := mcp.NewServer("server", "v0.0.1", opts)
+	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)
+		wantTools := []*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool}
-		defer clientSession.Close()
-		defer serverSession.Close()
-		server.AddTools(tools...)
 		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 {
+		wantTools := []mcp.Tool{*toolA.Tool, *toolB.Tool, *toolC.Tool}
-			testName := fmt.Sprintf("PageSize=%v", pageSize)
+		var gotTools []mcp.Tool
-			t.Run(testName, func(t *testing.T) {
+		seq := clientSession.Tools(ctx, nil)
-				clientSession, serverSession, server := createSessions(ctx, &mcp.ServerOptions{PageSize: pageSize})
+		for tool, err := range seq {
-				defer clientSession.Close()
+			if err != nil {
-				defer serverSession.Close()
+				t.Fatalf("Tools() failed: %v", err)
-				server.AddTools(tools...)
+			}
-				var gotTools []mcp.Tool
+			gotTools = append(gotTools, tool)
-				seq := clientSession.Tools(ctx, nil)
+		}
-				for tool, err := range seq {
+		if diff := cmp.Diff(wantTools, gotTools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
-					if err != nil {
+			t.Fatalf("Tools() mismatch (-want +got):\n%s", diff)
-						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)
-				}
-			})
 		}
 	})
 }
@@ -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)
+	t.Run("ListResources", func(t *testing.T) {
-	if err != nil {
+		wantResources := []*mcp.Resource{resourceA.Resource, resourceB.Resource, resourceC.Resource}
-		t.Fatal("ListResources() failed:", err)
+		res, err := clientSession.ListResources(ctx, nil)
-	}
+		if err != nil {
-	if diff := cmp.Diff(wantResource, res.Resources, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
+			t.Fatal("ListResources() failed:", err)
-		t.Fatalf("ListResources() mismatch (-want +got):\n%s", diff)
+		}
-	}
+		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)
+	t.Run("ListPrompts", func(t *testing.T) {
-	if err != nil {
+		wantPrompts := []*mcp.Prompt{promptA.Prompt, promptB.Prompt, promptC.Prompt}
-		t.Fatal("ListPrompts() failed:", err)
+		res, err := clientSession.ListPrompts(ctx, nil)
-	}
+		if err != nil {
-	if diff := cmp.Diff(wantPrompts, res.Prompts, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
+			t.Fatal("ListPrompts() failed:", err)
-		t.Fatalf("ListPrompts() mismatch (-want +got):\n%s", diff)
+		}
-	}
+		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)
+		}
+	})
 }

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.