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 tools

Browse source 
This CL adds paginating functionality for tools. It uses the gob encoder
for creating opaque cursors. More CLs to follow for paginating other features.

Change-Id: I1443c0213ceb6238d844a8eee9a0be52934f5cab
Reviewed-on: https://go-review.googlesource.com/c/tools/+/675055
Reviewed-by: Robert Findley <rfindley@google.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
This commit is contained in:
Sam Thanawalla 2025-05-21 15:31:56 +00:00
commit 8962b38ba5
2 changed files with 253 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

93
mcp/server.go
View file

@ -5,7 +5,10 @@
package mcp
import (
"bytes"
"context"
"encoding/base64"
"encoding/gob"
"fmt"
"iter"
"net/url"
@ -16,6 +19,8 @@ import (
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
)
const DefaultPageSize = 1000
// A Server is an instance of an MCP server.
//
// Servers expose server-side MCP features, which can serve one or more MCP
@ -40,6 +45,9 @@ type ServerOptions struct {
Instructions string
// If non-nil, called when "notifications/intialized" is received.
InitializedHandler func(context.Context, *ServerSession, *InitializedParams)
// PageSize is the maximum number of items to return in a single page for
// list methods (e.g. ListTools).
PageSize int
// If non-nil, called when "notifications/roots/list_changed" is received.
RootsListChangedHandler func(context.Context, *ServerSession, *RootsListChangedParams)
}
@ -55,6 +63,12 @@ func NewServer(name, version string, opts *ServerOptions) *Server {
if opts == nil {
opts = new(ServerOptions)
}
if opts.PageSize < 0 {
panic(fmt.Errorf("invalid page size %d", opts.PageSize))
}
if opts.PageSize == 0 {
opts.PageSize = DefaultPageSize
}
return &Server{
name: name,
version: version,
@ -186,8 +200,17 @@ 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
}
tools, nextCursor, err := paginateList(s.tools, cursor, s.opts.PageSize)
if err != nil {
return nil, err
}
res := new(ListToolsResult)
for t := range s.tools.all() {
res.NextCursor = nextCursor
for _, t := range tools {
res.Tools = append(res.Tools, t.Tool)
}
return res, nil
@ -509,3 +532,71 @@ func (ss *ServerSession) Close() error {
func (ss *ServerSession) Wait() error {
return ss.conn.Wait()
}
// pageToken is the internal structure for the opaque pagination cursor.
// It will be Gob-encoded and then Base64-encoded for use as a string token.
type pageToken struct {
LastUID string // The unique ID of the last resource seen.
}
// 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()
} else {
pageToken, err := decodeCursor(cursor)
// According to the spec, invalid cursors should return Invalid params.
if err != nil {
return nil, "", jsonrpc2.ErrInvalidParams
}
seq = fs.above(pageToken.LastUID)
}
var count int
for f := range seq {
count++
// If we've seen pageSize + 1 elements, we've gathered enough info to determine
// if there's a next page. Stop processing the sequence.
if count == pageSize+1 {
break
}
features = append(features, f)
}
// No remaining pages.
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
}
return features, nextCursor, nil
}

161
mcp/server_example_test.go
View file

@ -8,8 +8,13 @@ import (
"context"
"fmt"
"log"
"slices"
"testing"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
"github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/modelcontextprotocol/go-sdk/jsonschema"
)
type SayHiParams struct {
@ -51,3 +56,159 @@ 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()
server := mcp.NewServer("server", "v0.0.1", nil)
server.AddTools(toolA)
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)
}
_, err = clientSession.ListTools(ctx, &mcp.ListToolsParams{Cursor: "invalid"})
if err == nil {
t.Fatalf("expected error, got none")
}
serverSession.Close()
clientSession.Close()
}
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)
}
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()
if len(gotTools) != len(wantTools) {
t.Fatalf("expected %d tools, got %d", len(wantTools), len(gotTools))
}
}
}