go-sdk/mcp/server_example_test.go
Sam Thanawalla 8962b38ba5 internal/mcp: add pagination for tools
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>
2025-05-23 07:17:52 -07:00

214 lines
6 KiB
Go

// 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_test
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 {
Name string `json:"name" mcp:"the name to say hi to"`
}
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *SayHiParams) ([]*mcp.Content, error) {
return []*mcp.Content{
mcp.NewTextContent("Hi " + params.Name),
}, nil
}
func ExampleServer() {
ctx := context.Background()
clientTransport, serverTransport := mcp.NewInMemoryTransports()
server := mcp.NewServer("greeter", "v0.0.1", nil)
server.AddTools(mcp.NewTool("greet", "say hi", SayHi))
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.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil)
if err != nil {
log.Fatal(err)
}
fmt.Println(res.Content[0].Text)
clientSession.Close()
serverSession.Wait()
// 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))
}
}
}