internal/mcp: change paginateList to a generic helper

samthanawalla · gopherbot · commit 82ee0fd1228b · 2025-06-05T07:47:33.000-07:00
This CL simplifies the paginateList function in server.go to use a generic helper for tools, resources, and prompts. Change-Id: Ide0d2a90d715374280067e094d8870882e6bddfa Reviewed-on: https://go-review.googlesource.com/c/tools/+/678055 Auto-Submit: Sam Thanawalla <samthanawalla@google.com> Reviewed-by: Robert Findley <rfindley@google.com> LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
diff --git a/internal/mcp/client.go b/internal/mcp/client.go
@@ -391,23 +391,13 @@ func (cs *ClientSession) Prompts(ctx context.Context, params *ListPromptsParams)
 	})
 }
 
-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) {
+func paginate[P listParams, R listResult[T], T any](ctx context.Context, params P, listFunc func(context.Context, P) (R, error), items func(R) []*T) iter.Seq2[T, error] {
+	return func(yield func(T, error) bool) {
 		for {
 			res, err := listFunc(ctx, params)
 			if err != nil {
-				var zero E
+				var zero T
 				yield(zero, err)
 				return
 			}
diff --git a/internal/mcp/server.go b/internal/mcp/server.go
@@ -182,21 +182,15 @@ func (s *Server) Sessions() iter.Seq[*ServerSession] {
 func (s *Server) listPrompts(_ context.Context, _ *ServerSession, params *ListPromptsParams) (*ListPromptsResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	var cursor string
-	if params != nil {
-		cursor = params.Cursor
+	if params == nil {
+		params = &ListPromptsParams{}
 	}
-	prompts, nextCursor, err := paginateList(s.prompts, cursor, s.opts.PageSize)
-	if err != nil {
-		return nil, err
-	}
-	res := new(ListPromptsResult)
-	res.NextCursor = nextCursor
-	res.Prompts = []*Prompt{} // avoid JSON null
-	for _, p := range prompts {
-		res.Prompts = append(res.Prompts, p.Prompt)
-	}
-	return res, nil
+	return paginateList(s.prompts, s.opts.PageSize, params, &ListPromptsResult{}, func(res *ListPromptsResult, prompts []*ServerPrompt) {
+		res.Prompts = []*Prompt{} // avoid JSON null
+		for _, p := range prompts {
+			res.Prompts = append(res.Prompts, p.Prompt)
+		}
+	})
 }
 
 func (s *Server) getPrompt(ctx context.Context, cc *ServerSession, params *GetPromptParams) (*GetPromptResult, error) {
@@ -213,21 +207,15 @@ func (s *Server) getPrompt(ctx context.Context, cc *ServerSession, params *GetPr
 func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListToolsParams) (*ListToolsResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	var cursor string
-	if params != nil {
-		cursor = params.Cursor
-	}
-	tools, nextCursor, err := paginateList(s.tools, cursor, s.opts.PageSize)
-	if err != nil {
-		return nil, err
-	}
-	res := new(ListToolsResult)
-	res.NextCursor = nextCursor
-	res.Tools = []*Tool{} // avoid JSON null
-	for _, t := range tools {
-		res.Tools = append(res.Tools, t.Tool)
+	if params == nil {
+		params = &ListToolsParams{}
 	}
-	return res, nil
+	return paginateList(s.tools, s.opts.PageSize, params, &ListToolsResult{}, func(res *ListToolsResult, tools []*ServerTool) {
+		res.Tools = []*Tool{} // avoid JSON null
+		for _, t := range tools {
+			res.Tools = append(res.Tools, t.Tool)
+		}
+	})
 }
 
 func (s *Server) callTool(ctx context.Context, cc *ServerSession, params *CallToolParams[json.RawMessage]) (*CallToolResult, error) {
@@ -243,21 +231,15 @@ func (s *Server) callTool(ctx context.Context, cc *ServerSession, params *CallTo
 func (s *Server) listResources(_ context.Context, _ *ServerSession, params *ListResourcesParams) (*ListResourcesResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	var cursor string
-	if params != nil {
-		cursor = params.Cursor
-	}
-	resources, nextCursor, err := paginateList(s.resources, cursor, s.opts.PageSize)
-	if err != nil {
-		return nil, err
+	if params == nil {
+		params = &ListResourcesParams{}
 	}
-	res := new(ListResourcesResult)
-	res.NextCursor = nextCursor
-	res.Resources = []*Resource{} // avoid JSON null
-	for _, r := range resources {
-		res.Resources = append(res.Resources, r.Resource)
-	}
-	return res, nil
+	return paginateList(s.resources, s.opts.PageSize, params, &ListResourcesResult{}, func(res *ListResourcesResult, resources []*ServerResource) {
+		res.Resources = []*Resource{} // avoid JSON null
+		for _, r := range resources {
+			res.Resources = append(res.Resources, r.Resource)
+		}
+	})
 }
 
 func (s *Server) readResource(ctx context.Context, ss *ServerSession, params *ReadResourceParams) (*ReadResourceResult, error) {
@@ -618,22 +600,25 @@ func decodeCursor(cursor string) (*pageToken, error) {
 	return &token, nil
 }
 
-// paginateList returns a slice of features from the given featureSet, based on
-// the provided cursor and page size. It also returns a new cursor for the next
-// page, or an empty string if there are no more pages.
-func paginateList[T any](fs *featureSet[T], cursor string, pageSize int) (features []T, nextCursor string, err error) {
+// paginateList is a generic helper that returns a paginated slice of items
+// from a featureSet. It populates the provided result res with the items
+// and sets its next cursor for subsequent pages.
+// If there are no more pages, the next cursor within the result will be an empty string.
+func paginateList[P listParams, R listResult[T], T any](fs *featureSet[T], pageSize int, params P, res R, setFunc func(R, []T)) (R, error) {
 	var seq iter.Seq[T]
-	if cursor == "" {
+	if params.cursorPtr() == nil || *params.cursorPtr() == "" {
 		seq = fs.all()
 	} else {
-		pageToken, err := decodeCursor(cursor)
+		pageToken, err := decodeCursor(*params.cursorPtr())
 		// According to the spec, invalid cursors should return Invalid params.
 		if err != nil {
-			return nil, "", jsonrpc2.ErrInvalidParams
+			var zero R
+			return zero, jsonrpc2.ErrInvalidParams
 		}
 		seq = fs.above(pageToken.LastUID)
 	}
 	var count int
+	var features []T
 	for f := range seq {
 		count++
 		// If we've seen pageSize + 1 elements, we've gathered enough info to determine
@@ -643,13 +628,16 @@ func paginateList[T any](fs *featureSet[T], cursor string, pageSize int) (featur
 		}
 		features = append(features, f)
 	}
+	setFunc(res, features)
 	// No remaining pages.
 	if count < pageSize+1 {
-		return features, "", nil
+		return res, nil
 	}
-	nextCursor, err = encodeCursor(fs.uniqueID(features[len(features)-1]))
+	nextCursor, err := encodeCursor(fs.uniqueID(features[len(features)-1]))
 	if err != nil {
-		return nil, "", err
+		var zero R
+		return zero, err
 	}
-	return features, nextCursor, nil
+	*res.nextCursorPtr() = nextCursor
+	return res, nil
 }
diff --git a/internal/mcp/server_test.go b/internal/mcp/server_test.go
@@ -12,12 +12,29 @@ import (
 	"github.com/google/go-cmp/cmp"
 )
 
-type TestItem struct {
+type testItem struct {
 	Name  string
 	Value string
 }
 
-var allTestItems = []*TestItem{
+type testListParams struct {
+	Cursor string
+}
+
+func (p *testListParams) cursorPtr() *string {
+	return &p.Cursor
+}
+
+type testListResult struct {
+	Items      []*testItem
+	NextCursor string
+}
+
+func (r *testListResult) nextCursorPtr() *string {
+	return &r.NextCursor
+}
+
+var allTestItems = []*testItem{
 	{"alpha", "val-A"},
 	{"bravo", "val-B"},
 	{"charlie", "val-C"},
@@ -44,10 +61,10 @@ func getCursor(input string) string {
 func TestServerPaginateBasic(t *testing.T) {
 	testCases := []struct {
 		name           string
-		initialItems   []*TestItem
+		initialItems   []*testItem
 		inputCursor    string
 		inputPageSize  int
-		wantFeatures   []*TestItem
+		wantFeatures   []*testItem
 		wantNextCursor string
 		wantErr        bool
 	}{
@@ -154,41 +171,51 @@ func TestServerPaginateBasic(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			fs := newFeatureSet(func(t *TestItem) string { return t.Name })
+			fs := newFeatureSet(func(t *testItem) string { return t.Name })
 			fs.add(tc.initialItems...)
-			gotFeatures, gotNextCursor, err := paginateList(fs, tc.inputCursor, tc.inputPageSize)
+			params := &testListParams{Cursor: tc.inputCursor}
+			gotResult, err := paginateList(fs, tc.inputPageSize, params, &testListResult{}, func(res *testListResult, items []*testItem) {
+				res.Items = items
+			})
 			if (err != nil) != tc.wantErr {
 				t.Errorf("paginateList(%s) error, got %v, wantErr %v", tc.name, err, tc.wantErr)
 			}
-			if diff := cmp.Diff(tc.wantFeatures, gotFeatures); diff != "" {
+			if tc.wantErr {
+				return
+			}
+			if diff := cmp.Diff(tc.wantFeatures, gotResult.Items); diff != "" {
 				t.Errorf("paginateList(%s) mismatch (-want +got):\n%s", tc.name, diff)
 			}
-			if tc.wantNextCursor != gotNextCursor {
-				t.Errorf("paginateList(%s) nextCursor, got %v, want %v", tc.name, gotNextCursor, tc.wantNextCursor)
+			if tc.wantNextCursor != gotResult.NextCursor {
+				t.Errorf("paginateList(%s) nextCursor, got %v, want %v", tc.name, gotResult.NextCursor, tc.wantNextCursor)
 			}
 		})
 	}
 }
 
 func TestServerPaginateVariousPageSizes(t *testing.T) {
-	fs := newFeatureSet(func(t *TestItem) string { return t.Name })
+	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.
 	for pageSize := 1; pageSize < len(allTestItems)+1; pageSize++ {
-		var gotItems []*TestItem
+		var gotItems []*testItem
 		var nextCursor string
 		wantChunks := slices.Collect(slices.Chunk(allTestItems, pageSize))
 		index := 0
 		// Iterate through all pages, comparing sub-slices to the paginated list.
 		for {
-			gotFeatures, gotNextCursor, err := paginateList(fs, nextCursor, pageSize)
+			params := &testListParams{Cursor: nextCursor}
+			gotResult, err := paginateList(fs, pageSize, params, &testListResult{}, func(res *testListResult, items []*testItem) {
+				res.Items = items
+			})
 			if err != nil {
+				t.Fatalf("paginateList() unexpected error for pageSize %d, cursor %q: %v", pageSize, nextCursor, err)
 			}
-			if diff := cmp.Diff(wantChunks[index], gotFeatures); diff != "" {
+			if diff := cmp.Diff(wantChunks[index], gotResult.Items); diff != "" {
 				t.Errorf("paginateList mismatch (-want +got):\n%s", diff)
 			}
-			gotItems = append(gotItems, gotFeatures...)
-			nextCursor = gotNextCursor
+			gotItems = append(gotItems, gotResult.Items...)
+			nextCursor = gotResult.NextCursor
 			if nextCursor == "" {
 				break
 			}
diff --git a/internal/mcp/shared.go b/internal/mcp/shared.go
@@ -271,3 +271,13 @@ type Result interface {
 type emptyResult struct{}
 
 func (*emptyResult) GetMeta() *Meta { panic("should never be called") }
+
+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
+}
