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: add pagination for resources

Browse source 
Change-Id: I9f257fa09fe82e53087616859d1ad0c03f57dab6
Reviewed-on: https://go-review.googlesource.com/c/tools/+/675696
Reviewed-by: Jonathan Amsterdam <jba@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>
This commit is contained in:
Sam Thanawalla 2025-05-23 15:21:57 +00:00
commit dbbdf5a181
1 changed files with 278 additions and 166 deletions
Download patch file Download diff file Expand all files Collapse all files

65
mcp/server.go
View file

@ -229,8 +229,17 @@ 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
}
res := new(ListResourcesResult)
for r := range s.resources.all() {
res.NextCursor = nextCursor
for _, r := range resources {
res.Resources = append(res.Resources, r.Resource)
}
return res, nil
@ -539,35 +548,38 @@ type pageToken struct {
LastUID string // The unique ID of the last resource seen.
}
// encodeCursor encodes a unique identifier (UID) into a opaque pagination cursor
// by serializing a pageToken struct.
func encodeCursor(uid string) (string, error) {
var buf bytes.Buffer
token := pageToken{LastUID: uid}
encoder := gob.NewEncoder(&buf)
if err := encoder.Encode(token); err != nil {
return "", fmt.Errorf("failed to encode page token: %w", err)
}
return base64.URLEncoding.EncodeToString(buf.Bytes()), nil
}
// decodeCursor decodes an opaque pagination cursor into the original pageToken struct.
func decodeCursor(cursor string) (*pageToken, error) {
decodedBytes, err := base64.URLEncoding.DecodeString(cursor)
if err != nil {
return nil, fmt.Errorf("failed to decode cursor: %w", err)
}
var token pageToken
buf := bytes.NewBuffer(decodedBytes)
decoder := gob.NewDecoder(buf)
if err := decoder.Decode(&token); err != nil {
return nil, fmt.Errorf("failed to decode page token: %w, cursor: %v", err, cursor)
}
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) {
encodeCursor := func(uid string) (string, error) {
var buf bytes.Buffer
token := pageToken{LastUID: uid}
encoder := gob.NewEncoder(&buf)
if err := encoder.Encode(token); err != nil {
return "", fmt.Errorf("failed to encode page token: %w", err)
}
return base64.URLEncoding.EncodeToString(buf.Bytes()), nil
}
decodeCursor := func(cursor string) (*pageToken, error) {
decodedBytes, err := base64.URLEncoding.DecodeString(cursor)
if err != nil {
return nil, fmt.Errorf("failed to decode cursor: %w", err)
}
var token pageToken
buf := bytes.NewBuffer(decodedBytes)
decoder := gob.NewDecoder(buf)
if err := decoder.Decode(&token); err != nil {
return nil, fmt.Errorf("failed to decode page token: %w, cursor: %v", err, cursor)
}
return &token, nil
}
var seq iter.Seq[T]
if cursor == "" {
seq = fs.all()
@ -593,7 +605,6 @@ func paginateList[T any](fs *featureSet[T], cursor string, pageSize int) (featur
if count < pageSize+1 {
return features, "", nil
}
// Trim the extra element from the result.
nextCursor, err = encodeCursor(fs.uniqueID(features[len(features)-1]))
if err != nil {
return nil, "", err