diff --git a/mcp/server.go b/mcp/server.go index 297f82e..7e1a089 100644 --- a/mcp/server.go +++ b/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 diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index 99353e9..0887d47 100644 --- a/mcp/server_example_test.go +++ b/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) { - 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) } } diff --git a/mcp/server_test.go b/mcp/server_test.go new file mode 100644 index 0000000..11223c6 --- /dev/null +++ b/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)) + } + } +}