From 30af200d5e39fb02ec7523d6489685ef54e01280 Mon Sep 17 00:00:00 2001 From: Sam Thanawalla Date: Mon, 2 Jun 2025 16:37:01 +0000 Subject: [PATCH] internal/mcp: change paginateList to a generic helper 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 Reviewed-by: Robert Findley LUCI-TryBot-Result: Go LUCI --- mcp/client.go | 16 ++------ mcp/server.go | 94 ++++++++++++++++++++-------------------------- mcp/server_test.go | 57 ++++++++++++++++++++-------- mcp/shared.go | 10 +++++ 4 files changed, 96 insertions(+), 81 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 258dc31..ec21b69 100644 --- a/mcp/client.go +++ b/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/mcp/server.go b/mcp/server.go index afb2575..3521b25 100644 --- a/mcp/server.go +++ b/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 + if params == nil { + params = &ListToolsParams{} } - 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) - } - 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 + if params == nil { + params = &ListResourcesParams{} } - resources, nextCursor, err := paginateList(s.resources, cursor, s.opts.PageSize) - if err != nil { - return nil, err - } - 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/mcp/server_test.go b/mcp/server_test.go index 15250ca..19701f3 100644 --- a/mcp/server_test.go +++ b/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/mcp/shared.go b/mcp/shared.go index 628a045..57454b8 100644 --- a/mcp/shared.go +++ b/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 +}