diff --git a/mcp/server.go b/mcp/server.go index 127b8b8..a8543cd 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -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 +} diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index 2be9035..99353e9 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -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)) + } + } +}