internal/mcp: add pagination for resources

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>

mcp/server.go

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

mcp/server_example_test.go

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

mcp/server_test.go

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