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 <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>

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
+		}
+	}
+}

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) {
+func createSessions(ctx context.Context, opts *mcp.ServerOptions) (*mcp.ClientSession, *mcp.ServerSession, *mcp.Server) {
-	server := mcp.NewServer("server", "v0.0.1", nil)
+	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)
+	t.Run("ListTools", func(t *testing.T) {
-	defer clientSession.Close()
+		clientSession, serverSession, server := createSessions(ctx, nil)
-	defer serverSession.Close()
+		defer clientSession.Close()
-	server.AddTools(tools...)
+		defer serverSession.Close()
-	res, err := clientSession.ListTools(ctx, nil)
+		server.AddTools(tools...)
-	if err != nil {
+		res, err := clientSession.ListTools(ctx, nil)
-		t.Fatal("ListTools() failed:", err)
+		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 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...)