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

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
3 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

181
mcp/server_example_test.go
View file

@ -8,7 +8,6 @@ import (
"context"
"fmt"
"log"
"slices"
"testing"
"github.com/google/go-cmp/cmp"
@ -57,158 +56,58 @@ func ExampleServer() {
// Output: Hi user
}
func TestListTool(t *testing.T) {
toolA := mcp.NewTool("apple", "apple tool", SayHi)
toolB := mcp.NewTool("banana", "banana tool", SayHi)
toolC := mcp.NewTool("cherry", "cherry tool", SayHi)
testCases := []struct {
tools []*mcp.ServerTool
want []*mcp.Tool
pageSize int
}{
{
// Simple test.
[]*mcp.ServerTool{toolA, toolB, toolC},
[]*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool},
mcp.DefaultPageSize,
},
{
// Tools should be ordered by tool name.
[]*mcp.ServerTool{toolC, toolA, toolB},
[]*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool},
mcp.DefaultPageSize,
},
{
// Page size of 1 should yield the first tool only.
[]*mcp.ServerTool{toolC, toolA, toolB},
[]*mcp.Tool{toolA.Tool},
1,
},
{
// Page size of 2 should yield the first 2 tools only.
[]*mcp.ServerTool{toolC, toolA, toolB},
[]*mcp.Tool{toolA.Tool, toolB.Tool},
2,
},
{
// Page size of 3 should yield all tools.
[]*mcp.ServerTool{toolC, toolA, toolB},
[]*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool},
3,
},
{
[]*mcp.ServerTool{},
nil,
1,
},
}
ctx := context.Background()
for _, tc := range testCases {
server := mcp.NewServer("server", "v0.0.1", &mcp.ServerOptions{PageSize: tc.pageSize})
server.AddTools(tc.tools...)
clientTransport, serverTransport := mcp.NewInMemoryTransports()
serverSession, err := server.Connect(ctx, serverTransport)
if err != nil {
log.Fatal(err)
}
client := mcp.NewClient("client", "v0.0.1", nil)
clientSession, err := client.Connect(ctx, clientTransport)
if err != nil {
log.Fatal(err)
}
res, err := clientSession.ListTools(ctx, nil)
serverSession.Close()
clientSession.Close()
if err != nil {
log.Fatal(err)
}
if len(res.Tools) != len(tc.want) {
t.Fatalf("expected %d tools, got %d", len(tc.want), len(res.Tools))
}
if diff := cmp.Diff(res.Tools, tc.want, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
t.Fatalf("expected tools %+v, got %+v", tc.want, res.Tools)
}
if tc.pageSize < len(tc.tools) && res.NextCursor == "" {
t.Fatalf("expected next cursor, got none")
}
}
}
func TestListToolPaginateInvalidCursor(t *testing.T) {
toolA := mcp.NewTool("apple", "apple tool", SayHi)
ctx := context.Background()
// createSessions creates and connects an in-memory client and server session for testing purposes.
func createSessions(ctx context.Context) (*mcp.ClientSession, *mcp.ServerSession, *mcp.Server) {
server := mcp.NewServer("server", "v0.0.1", nil)
server.AddTools(toolA)
clientTransport, serverTransport := mcp.NewInMemoryTransports()
client := mcp.NewClient("client", "v0.0.1", nil)
serverTransport, clientTransport := mcp.NewInMemoryTransports()
serverSession, err := server.Connect(ctx, serverTransport)
if err != nil {
log.Fatal(err)
}
client := mcp.NewClient("client", "v0.0.1", nil)
clientSession, err := client.Connect(ctx, clientTransport)
if err != nil {
log.Fatal(err)
}
_, err = clientSession.ListTools(ctx, &mcp.ListToolsParams{Cursor: "invalid"})
if err == nil {
t.Fatalf("expected error, got none")
}
serverSession.Close()
clientSession.Close()
return clientSession, serverSession, server
}
func TestListToolPaginate(t *testing.T) {
serverTools := []*mcp.ServerTool{
mcp.NewTool("apple", "apple tool", SayHi),
mcp.NewTool("banana", "banana tool", SayHi),
mcp.NewTool("cherry", "cherry tool", SayHi),
mcp.NewTool("durian", "durian tool", SayHi),
mcp.NewTool("elderberry", "elderberry tool", SayHi),
}
var wantTools []*mcp.Tool
for _, tool := range serverTools {
wantTools = append(wantTools, tool.Tool)
}
func TestListTool(t *testing.T) {
toolA := mcp.NewTool("apple", "apple tool", SayHi)
toolB := mcp.NewTool("banana", "banana tool", SayHi)
toolC := mcp.NewTool("cherry", "cherry tool", SayHi)
tools := []*mcp.ServerTool{toolA, toolB, toolC}
wantTools := []*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool}
ctx := context.Background()
// Try all possible page sizes, ensuring we get the correct list of tools.
for pageSize := 1; pageSize < len(serverTools)+1; pageSize++ {
server := mcp.NewServer("server", "v0.0.1", &mcp.ServerOptions{PageSize: pageSize})
server.AddTools(serverTools...)
clientTransport, serverTransport := mcp.NewInMemoryTransports()
serverSession, err := server.Connect(ctx, serverTransport)
if err != nil {
log.Fatal(err)
}
client := mcp.NewClient("client", "v0.0.1", nil)
clientSession, err := client.Connect(ctx, clientTransport)
if err != nil {
log.Fatal(err)
}
var gotTools []*mcp.Tool
var nextCursor string
wantChunks := slices.Collect(slices.Chunk(wantTools, pageSize))
index := 0
// Iterate through all pages, comparing sub-slices to the paginated list.
for {
res, err := clientSession.ListTools(ctx, &mcp.ListToolsParams{Cursor: nextCursor})
if err != nil {
log.Fatal(err)
}
gotTools = append(gotTools, res.Tools...)
nextCursor = res.NextCursor
if diff := cmp.Diff(res.Tools, wantChunks[index], cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
t.Errorf("expected %v, got %v, (-want +got):\n%s", wantChunks[index], res.Tools, diff)
}
if res.NextCursor == "" {
break
}
index++
}
serverSession.Close()
clientSession.Close()
clientSession, serverSession, server := createSessions(ctx)
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(wantTools, res.Tools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
t.Fatalf("ListTools() mismatch (-want +got):\n%s", diff)
}
}
if len(gotTools) != len(wantTools) {
t.Fatalf("expected %d tools, got %d", len(wantTools), len(gotTools))
}
func TestListResources(t *testing.T) {
resourceA := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://apple"}}
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)
defer clientSession.Close()
defer serverSession.Close()
server.AddResources(resources...)
res, err := clientSession.ListResources(ctx, nil)
if err != nil {
t.Fatal("ListResources() failed:", err)
}
if diff := cmp.Diff(wantResource, res.Resources, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
t.Fatalf("ListResources() mismatch (-want +got):\n%s", diff)
}
}

202
mcp/server_test.go Normal file
View file

@ -0,0 +1,202 @@
// 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 (
"log"
"slices"
"testing"
"github.com/google/go-cmp/cmp"
)
type TestItem struct {
Name string
Value string
}
var allTestItems = []*TestItem{
{"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"},
}
// getCursor encodes a string input into a URL-safe base64 cursor,
// fatally logging any encoding errors.
func getCursor(input string) string {
cursor, err := encodeCursor(input)
if err != nil {
log.Fatalf("encodeCursor(%s) error = %v", input, err)
}
return cursor
}
func TestPaginateBasic(t *testing.T) {
testCases := []struct {
name string
initialItems []*TestItem
inputCursor string
inputPageSize int
wantFeatures []*TestItem
wantNextCursor string
wantErr bool
}{
{
name: "FirstPage_DefaultSize_Full",
initialItems: allTestItems,
inputCursor: "",
inputPageSize: 5,
wantFeatures: allTestItems[0:5],
wantNextCursor: getCursor("echo"), // Based on last item of first page
wantErr: false,
},
{
name: "SecondPage_DefaultSize_Full",
initialItems: allTestItems,
inputCursor: getCursor("echo"),
inputPageSize: 5,
wantFeatures: allTestItems[5:10],
wantNextCursor: getCursor("juliet"), // Based on last item of second page
wantErr: false,
},
{
name: "SecondPage_DefaultSize_Full_OutOfOrder",
initialItems: append(allTestItems[5:], allTestItems[0:5]...),
inputCursor: getCursor("echo"),
inputPageSize: 5,
wantFeatures: allTestItems[5:10],
wantNextCursor: getCursor("juliet"), // Based on last item of second page
wantErr: false,
},
{
name: "SecondPage_DefaultSize_Full_Duplicates",
initialItems: append(allTestItems, allTestItems[0:5]...),
inputCursor: getCursor("echo"),
inputPageSize: 5,
wantFeatures: allTestItems[5:10],
wantNextCursor: getCursor("juliet"), // Based on last item of second page
wantErr: false,
},
{
name: "LastPage_Remaining",
initialItems: allTestItems,
inputCursor: getCursor("juliet"),
inputPageSize: 5,
wantFeatures: allTestItems[10:11], // Only 1 item left
wantNextCursor: "", // No more pages
wantErr: false,
},
{
name: "PageSize_1",
initialItems: allTestItems,
inputCursor: "",
inputPageSize: 1,
wantFeatures: allTestItems[0:1],
wantNextCursor: getCursor("alpha"),
wantErr: false,
},
{
name: "PageSize_All",
initialItems: allTestItems,
inputCursor: "",
inputPageSize: len(allTestItems), // Page size equals total
wantFeatures: allTestItems,
wantNextCursor: "", // No more pages
wantErr: false,
},
{
name: "PageSize_LargerThanAll",
initialItems: allTestItems,
inputCursor: "",
inputPageSize: len(allTestItems) + 5, // Page size larger than total
wantFeatures: allTestItems,
wantNextCursor: "",
wantErr: false,
},
{
name: "EmptySet",
initialItems: nil,
inputCursor: "",
inputPageSize: 5,
wantFeatures: nil,
wantNextCursor: "",
wantErr: false,
},
{
name: "InvalidCursor",
initialItems: allTestItems,
inputCursor: "not-a-valid-gob-base64-cursor",
inputPageSize: 5,
wantFeatures: nil, // Should be nil for error cases
wantNextCursor: "",
wantErr: true,
},
{
name: "AboveNonExistentID",
initialItems: allTestItems,
inputCursor: getCursor("dne"), // A UID that doesn't exist
inputPageSize: 5,
wantFeatures: allTestItems[4:9], // Should return elements above UID.
wantNextCursor: getCursor("india"),
wantErr: false,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
fs := newFeatureSet(func(t *TestItem) string { return t.Name })
fs.add(tc.initialItems...)
gotFeatures, gotNextCursor, err := paginateList(fs, tc.inputCursor, tc.inputPageSize)
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 != "" {
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)
}
})
}
}
func TestPaginateVariousPageSizes(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.
for pageSize := 1; pageSize < len(allTestItems)+1; pageSize++ {
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)
if err != nil {
}
if diff := cmp.Diff(wantChunks[index], gotFeatures); diff != "" {
t.Errorf("paginateList mismatch (-want +got):\n%s", diff)
}
gotItems = append(gotItems, gotFeatures...)
nextCursor = gotNextCursor
if nextCursor == "" {
break
}
index++
}
if len(gotItems) != len(allTestItems) {
t.Fatalf("paginateList() returned %d items, want %d", len(allTestItems), len(gotItems))
}
}
}