modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

internal/mcp: change paginateList to a generic helper

Browse source 
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>
This commit is contained in:
Sam Thanawalla 2025-06-02 16:37:01 +00:00 • committed by Gopher Robot
commit 30af200d5e
4 changed files with 96 additions and 81 deletions
Download patch file Download diff file Expand all files Collapse all files

16
mcp/client.go
View file

@ -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
}

94
mcp/server.go
View file

@ -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
}

57
mcp/server_test.go
View file

@ -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
}

10
mcp/shared.go
View file

@ -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
}