From d45330dd8fc2181a7311043d60660e0c28cbebc7 Mon Sep 17 00:00:00 2001 From: Sam Thanawalla Date: Wed, 28 May 2025 16:50:51 +0000 Subject: [PATCH] internal/mcp: add iterator method for tools Change-Id: Ibe799caf6592b2763868bdb30f474bc224aa8b00 Reviewed-on: https://go-review.googlesource.com/c/tools/+/676915 LUCI-TryBot-Result: Go LUCI Reviewed-by: Robert Findley --- mcp/client.go | 29 +++++++++++++++++++ mcp/server_example_test.go | 57 +++++++++++++++++++++++++++----------- 2 files changed, 70 insertions(+), 16 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 2f8c654..aceb10d 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -8,6 +8,7 @@ import ( "context" "encoding/json" "fmt" + "iter" "slices" "sync" @@ -314,3 +315,31 @@ func (c *Client) callLoggingHandler(ctx context.Context, cs *ClientSession, para } return nil, nil } + +// Tools provides an iterator for all tools available on the server, +// automatically fetching pages and managing cursors. +// The `params` argument can set the initial cursor. +func (c *ClientSession) Tools(ctx context.Context, params *ListToolsParams) iter.Seq2[Tool, error] { + currentParams := &ListToolsParams{} + if params != nil { + *currentParams = *params + } + return func(yield func(Tool, error) bool) { + for { + res, err := c.ListTools(ctx, currentParams) + if err != nil { + yield(Tool{}, err) + return + } + for _, t := range res.Tools { + if !yield(*t, nil) { + return + } + } + if res.NextCursor == "" { + return + } + currentParams.Cursor = res.NextCursor + } + } +} diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index 93679ac..ded355a 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -57,8 +57,8 @@ func ExampleServer() { } // 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) +func createSessions(ctx context.Context, opts *mcp.ServerOptions) (*mcp.ClientSession, *mcp.ServerSession, *mcp.Server) { + server := mcp.NewServer("server", "v0.0.1", opts) client := mcp.NewClient("client", "v0.0.1", nil) serverTransport, clientTransport := mcp.NewInMemoryTransports() serverSession, err := server.Connect(ctx, serverTransport) @@ -77,19 +77,44 @@ func TestListTools(t *testing.T) { 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} + wantListTools := []*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool} + wantIteratorTools := []mcp.Tool{*toolA.Tool, *toolB.Tool, *toolC.Tool} ctx := context.Background() - 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) - } + t.Run("ListTools", func(t *testing.T) { + clientSession, serverSession, server := createSessions(ctx, nil) + 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(wantListTools, res.Tools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("ListTools() mismatch (-want +got):\n%s", diff) + } + }) + t.Run("ToolsIterator", func(t *testing.T) { + for pageSize := range len(tools) + 1 { + testName := fmt.Sprintf("PageSize=%v", pageSize) + t.Run(testName, func(t *testing.T) { + clientSession, serverSession, server := createSessions(ctx, &mcp.ServerOptions{PageSize: pageSize}) + defer clientSession.Close() + defer serverSession.Close() + server.AddTools(tools...) + var gotTools []mcp.Tool + seq := clientSession.Tools(ctx, nil) + for tool, err := range seq { + if err != nil { + t.Fatalf("Tools(%s) failed: %v", testName, err) + } + gotTools = append(gotTools, tool) + } + if diff := cmp.Diff(wantIteratorTools, gotTools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" { + t.Fatalf("Tools(%s) mismatch (-want +got):\n%s", testName, diff) + } + }) + } + }) } func TestListResources(t *testing.T) { @@ -99,7 +124,7 @@ func TestListResources(t *testing.T) { resources := []*mcp.ServerResource{resourceA, resourceB, resourceC} wantResource := []*mcp.Resource{resourceA.Resource, resourceB.Resource, resourceC.Resource} ctx := context.Background() - clientSession, serverSession, server := createSessions(ctx) + clientSession, serverSession, server := createSessions(ctx, nil) defer clientSession.Close() defer serverSession.Close() server.AddResources(resources...) @@ -119,7 +144,7 @@ func TestListPrompts(t *testing.T) { prompts := []*mcp.ServerPrompt{promptA, promptB, promptC} wantPrompts := []*mcp.Prompt{promptA.Prompt, promptB.Prompt, promptC.Prompt} ctx := context.Background() - clientSession, serverSession, server := createSessions(ctx) + clientSession, serverSession, server := createSessions(ctx, nil) defer clientSession.Close() defer serverSession.Close() server.AddPrompts(prompts...)