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>
This commit is contained in:
Sam Thanawalla
parent
b1780319e4
commit
8962b38ba5
1 changed files with 253 additions and 1 deletions
|
|
@ -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))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue