internal/mcp: iterators return pointers

Iterators return pointers instead of values, to be consistent.

Move tests out of example file.

Factor out iterator tests.

Change-Id: Icb14edf99d738de9d9dc08c1169057cb3dc5894b
Reviewed-on: https://go-review.googlesource.com/c/tools/+/680295
Reviewed-by: Robert Findley <rfindley@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Sam Thanawalla <samthanawalla@google.com>

mcp/client.go

@@ -336,9 +336,9 @@ func (c *Client) callLoggingHandler(ctx context.Context, cs *ClientSession, para
 
 // 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.
+// The params argument can set the initial cursor.
 // Iteration stops at the first encountered error, which will be yielded.
-func (cs *ClientSession) Tools(ctx context.Context, params *ListToolsParams) iter.Seq2[Tool, error] {
+func (cs *ClientSession) Tools(ctx context.Context, params *ListToolsParams) iter.Seq2[*Tool, error] {
 	if params == nil {
 		params = &ListToolsParams{}
 	}
@@ -349,9 +349,9 @@ func (cs *ClientSession) Tools(ctx context.Context, params *ListToolsParams) ite
 
 // Resources provides an iterator for all resources available on the server,
 // automatically fetching pages and managing cursors.
-// The `params` argument can set the initial cursor.
+// The params argument can set the initial cursor.
 // Iteration stops at the first encountered error, which will be yielded.
-func (cs *ClientSession) Resources(ctx context.Context, params *ListResourcesParams) iter.Seq2[Resource, error] {
+func (cs *ClientSession) Resources(ctx context.Context, params *ListResourcesParams) iter.Seq2[*Resource, error] {
 	if params == nil {
 		params = &ListResourcesParams{}
 	}
@@ -364,7 +364,7 @@ func (cs *ClientSession) Resources(ctx context.Context, params *ListResourcesPar
 // automatically fetching pages and managing cursors.
 // The `params` argument can set the initial cursor.
 // Iteration stops at the first encountered error, which will be yielded.
-func (cs *ClientSession) ResourceTemplates(ctx context.Context, params *ListResourceTemplatesParams) iter.Seq2[ResourceTemplate, error] {
+func (cs *ClientSession) ResourceTemplates(ctx context.Context, params *ListResourceTemplatesParams) iter.Seq2[*ResourceTemplate, error] {
 	if params == nil {
 		params = &ListResourceTemplatesParams{}
 	}
@@ -375,9 +375,9 @@ func (cs *ClientSession) ResourceTemplates(ctx context.Context, params *ListReso
 
 // Prompts provides an iterator for all prompts available on the server,
 // automatically fetching pages and managing cursors.
-// The `params` argument can set the initial cursor.
+// The params argument can set the initial cursor.
 // Iteration stops at the first encountered error, which will be yielded.
-func (cs *ClientSession) Prompts(ctx context.Context, params *ListPromptsParams) iter.Seq2[Prompt, error] {
+func (cs *ClientSession) Prompts(ctx context.Context, params *ListPromptsParams) iter.Seq2[*Prompt, error] {
 	if params == nil {
 		params = &ListPromptsParams{}
 	}
@@ -387,17 +387,16 @@ func (cs *ClientSession) Prompts(ctx context.Context, params *ListPromptsParams)
 }
 
 // paginate is a generic helper function to provide a paginated iterator.
-func paginate[P listParams, R listResult[T], T any](ctx context.Context, params P, listFunc func(context.Context, P) (R, error), items func(R) []*T) iter.Seq2[T, error] {
+func paginate[P listParams, R listResult[T], T any](ctx context.Context, params P, listFunc func(context.Context, P) (R, error), items func(R) []*T) iter.Seq2[*T, error] {
-	return func(yield func(T, error) bool) {
+	return func(yield func(*T, error) bool) {
 		for {
 			res, err := listFunc(ctx, params)
 			if err != nil {
-				var zero T
+				yield(nil, err)
-				yield(zero, err)
 				return
 			}
 			for _, r := range items(res) {
-				if !yield(*r, nil) {
+				if !yield(r, nil) {
 					return
 				}
 			}

mcp/client_list_test.go

@@ -0,0 +1,136 @@
+// 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"
+	"iter"
+	"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"
+)
+
+func TestListTools(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()
+	clientSession, serverSession, server := createSessions(ctx)
+	defer clientSession.Close()
+	defer serverSession.Close()
+	server.AddTools(tools...)
+	t.Run("ListTools", func(t *testing.T) {
+		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("ToolsIterator", func(t *testing.T) {
+		testIterator(ctx, t, clientSession.Tools(ctx, nil), wantTools)
+	})
+}
+
+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"}}
+	wantResources := []*mcp.Resource{resourceA.Resource, resourceB.Resource, resourceC.Resource}
+
+	resources := []*mcp.ServerResource{resourceA, resourceB, resourceC}
+	ctx := context.Background()
+	clientSession, serverSession, server := createSessions(ctx)
+	defer clientSession.Close()
+	defer serverSession.Close()
+	server.AddResources(resources...)
+	t.Run("ListResources", func(t *testing.T) {
+		res, err := clientSession.ListResources(ctx, nil)
+		if err != nil {
+			t.Fatal("ListResources() failed:", err)
+		}
+		if diff := cmp.Diff(wantResources, res.Resources, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
+			t.Fatalf("ListResources() mismatch (-want +got):\n%s", diff)
+		}
+	})
+	t.Run("ResourcesIterator", func(t *testing.T) {
+		testIterator(ctx, t, clientSession.Resources(ctx, nil), wantResources)
+	})
+}
+
+func TestListResourceTemplates(t *testing.T) {
+	resourceTmplA := &mcp.ServerResourceTemplate{ResourceTemplate: &mcp.ResourceTemplate{URITemplate: "http://apple/{x}"}}
+	resourceTmplB := &mcp.ServerResourceTemplate{ResourceTemplate: &mcp.ResourceTemplate{URITemplate: "http://banana/{x}"}}
+	resourceTmplC := &mcp.ServerResourceTemplate{ResourceTemplate: &mcp.ResourceTemplate{URITemplate: "http://cherry/{x}"}}
+	wantResourceTemplates := []*mcp.ResourceTemplate{
+		resourceTmplA.ResourceTemplate, resourceTmplB.ResourceTemplate,
+		resourceTmplC.ResourceTemplate,
+	}
+	resourceTemplates := []*mcp.ServerResourceTemplate{resourceTmplA, resourceTmplB, resourceTmplC}
+	ctx := context.Background()
+	clientSession, serverSession, server := createSessions(ctx)
+	defer clientSession.Close()
+	defer serverSession.Close()
+	server.AddResourceTemplates(resourceTemplates...)
+	t.Run("ListResourceTemplates", func(t *testing.T) {
+		res, err := clientSession.ListResourceTemplates(ctx, nil)
+		if err != nil {
+			t.Fatal("ListResourceTemplates() failed:", err)
+		}
+		if diff := cmp.Diff(wantResourceTemplates, res.ResourceTemplates, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
+			t.Fatalf("ListResourceTemplates() mismatch (-want +got):\n%s", diff)
+		}
+	})
+	t.Run("ResourceTemplatesIterator", func(t *testing.T) {
+		testIterator(ctx, t, clientSession.ResourceTemplates(ctx, nil), wantResourceTemplates)
+	})
+}
+
+func TestListPrompts(t *testing.T) {
+	promptA := mcp.NewPrompt("apple", "apple prompt", testPromptHandler[struct{}])
+	promptB := mcp.NewPrompt("banana", "banana prompt", testPromptHandler[struct{}])
+	promptC := mcp.NewPrompt("cherry", "cherry prompt", testPromptHandler[struct{}])
+	wantPrompts := []*mcp.Prompt{promptA.Prompt, promptB.Prompt, promptC.Prompt}
+
+	prompts := []*mcp.ServerPrompt{promptA, promptB, promptC}
+	ctx := context.Background()
+	clientSession, serverSession, server := createSessions(ctx)
+	defer clientSession.Close()
+	defer serverSession.Close()
+	server.AddPrompts(prompts...)
+	t.Run("ListPrompts", func(t *testing.T) {
+		res, err := clientSession.ListPrompts(ctx, nil)
+		if err != nil {
+			t.Fatal("ListPrompts() failed:", err)
+		}
+		if diff := cmp.Diff(wantPrompts, res.Prompts, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
+			t.Fatalf("ListPrompts() mismatch (-want +got):\n%s", diff)
+		}
+	})
+	t.Run("PromptsIterator", func(t *testing.T) {
+		testIterator(ctx, t, clientSession.Prompts(ctx, nil), wantPrompts)
+	})
+}
+
+func testIterator[T any](ctx context.Context, t *testing.T, seq iter.Seq2[*T, error], want []*T) {
+	t.Helper()
+	var got []*T
+	for x, err := range seq {
+		if err != nil {
+			t.Fatalf("iteration failed: %v", err)
+		}
+		got = append(got, x)
+	}
+	if diff := cmp.Diff(want, got, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
+		t.Fatalf("mismatch (-want +got):\n%s", diff)
+	}
+}

mcp/client_test.go

@@ -50,14 +50,6 @@ var allItems = []*Item{
 	{"kilo", "val-K"},
 }
 
-func toItemValueSlice(ptrSlice []*Item) []Item {
-	var valSlice []Item
-	for _, ptr := range ptrSlice {
-		valSlice = append(valSlice, *ptr)
-	}
-	return valSlice
-}
-
 // generatePaginatedResults is a helper to create a sequence of mock responses for pagination.
 // It simulates a server returning items in pages based on a given page size.
 func generatePaginatedResults(all []*Item, pageSize int) []*ListTestResult {
@@ -88,18 +80,18 @@ func TestClientPaginateBasic(t *testing.T) {
 		results       []*ListTestResult
 		mockError     error
 		initialParams *ListTestParams
-		expected      []Item
+		expected      []*Item
 		expectError   bool
 	}{
 		{
 			name:     "SinglePageAllItems",
 			results:  generatePaginatedResults(allItems, len(allItems)),
-			expected: toItemValueSlice(allItems),
+			expected: allItems,
 		},
 		{
 			name:     "MultiplePages",
 			results:  generatePaginatedResults(allItems, 3),
-			expected: toItemValueSlice(allItems),
+			expected: allItems,
 		},
 		{
 			name:     "EmptyResults",
@@ -117,7 +109,7 @@ func TestClientPaginateBasic(t *testing.T) {
 			name:          "InitialCursorProvided",
 			initialParams: &ListTestParams{Cursor: "cursor_2"},
 			results:       generatePaginatedResults(allItems[2:], 3),
-			expected:      toItemValueSlice(allItems[2:]),
+			expected:      allItems[2:],
 		},
 		{
 			name:          "CursorBeyondAllItems",
@@ -148,7 +140,7 @@ func TestClientPaginateBasic(t *testing.T) {
 				params = &ListTestParams{}
 			}
 
-			var gotItems []Item
+			var gotItems []*Item
 			var iterationErr error
 			seq := paginate(ctx, params, listFunc, func(r *ListTestResult) []*Item { return r.Items })
 			for item, err := range seq {
@@ -185,7 +177,7 @@ func TestClientPaginateVariousPageSizes(t *testing.T) {
 				results = results[1:]
 				return res, nil
 			}
-			var gotItems []Item
+			var gotItems []*Item
 			seq := paginate(ctx, &ListTestParams{}, listFunc, func(r *ListTestResult) []*Item { return r.Items })
 			for item, err := range seq {
 				if err != nil {
@@ -193,7 +185,7 @@ func TestClientPaginateVariousPageSizes(t *testing.T) {
 				}
 				gotItems = append(gotItems, item)
 			}
-			if diff := cmp.Diff(toItemValueSlice(allItems), gotItems, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
+			if diff := cmp.Diff(allItems, gotItems, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
 				t.Fatalf("paginate() mismatch (-want +got):\n%s", diff)
 			}
 		})

mcp/server_example_test.go

@@ -8,12 +8,8 @@ import (
 	"context"
 	"fmt"
 	"log"
-	"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 {
@@ -76,111 +72,3 @@ func createSessions(ctx context.Context) (*mcp.ClientSession, *mcp.ServerSession
 	}
 	return clientSession, serverSession, server
 }
-
-func TestListTools(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}
-	ctx := context.Background()
-	clientSession, serverSession, server := createSessions(ctx)
-	defer clientSession.Close()
-	defer serverSession.Close()
-	server.AddTools(tools...)
-	t.Run("ListTools", func(t *testing.T) {
-		wantTools := []*mcp.Tool{toolA.Tool, toolB.Tool, toolC.Tool}
-		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("ToolsIterator", func(t *testing.T) {
-		wantTools := []mcp.Tool{*toolA.Tool, *toolB.Tool, *toolC.Tool}
-		var gotTools []mcp.Tool
-		seq := clientSession.Tools(ctx, nil)
-		for tool, err := range seq {
-			if err != nil {
-				t.Fatalf("Tools() failed: %v", err)
-			}
-			gotTools = append(gotTools, tool)
-		}
-		if diff := cmp.Diff(wantTools, gotTools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
-			t.Fatalf("Tools() mismatch (-want +got):\n%s", diff)
-		}
-	})
-}
-
-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}
-	ctx := context.Background()
-	clientSession, serverSession, server := createSessions(ctx)
-	defer clientSession.Close()
-	defer serverSession.Close()
-	server.AddResources(resources...)
-	t.Run("ListResources", func(t *testing.T) {
-		wantResources := []*mcp.Resource{resourceA.Resource, resourceB.Resource, resourceC.Resource}
-		res, err := clientSession.ListResources(ctx, nil)
-		if err != nil {
-			t.Fatal("ListResources() failed:", err)
-		}
-		if diff := cmp.Diff(wantResources, res.Resources, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
-			t.Fatalf("ListResources() mismatch (-want +got):\n%s", diff)
-		}
-	})
-	t.Run("ResourcesIterator", func(t *testing.T) {
-		wantResources := []mcp.Resource{*resourceA.Resource, *resourceB.Resource, *resourceC.Resource}
-		var gotResources []mcp.Resource
-		seq := clientSession.Resources(ctx, nil)
-		for resource, err := range seq {
-			if err != nil {
-				t.Fatalf("Resources() failed: %v", err)
-			}
-			gotResources = append(gotResources, resource)
-		}
-		if diff := cmp.Diff(wantResources, gotResources, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
-			t.Fatalf("Resources() mismatch (-want +got):\n%s", diff)
-		}
-	})
-}
-
-func TestListPrompts(t *testing.T) {
-	promptA := mcp.NewPrompt("apple", "apple prompt", testPromptHandler[struct{}])
-	promptB := mcp.NewPrompt("banana", "banana prompt", testPromptHandler[struct{}])
-	promptC := mcp.NewPrompt("cherry", "cherry prompt", testPromptHandler[struct{}])
-	prompts := []*mcp.ServerPrompt{promptA, promptB, promptC}
-	ctx := context.Background()
-	clientSession, serverSession, server := createSessions(ctx)
-	defer clientSession.Close()
-	defer serverSession.Close()
-	server.AddPrompts(prompts...)
-	t.Run("ListPrompts", func(t *testing.T) {
-		wantPrompts := []*mcp.Prompt{promptA.Prompt, promptB.Prompt, promptC.Prompt}
-		res, err := clientSession.ListPrompts(ctx, nil)
-		if err != nil {
-			t.Fatal("ListPrompts() failed:", err)
-		}
-		if diff := cmp.Diff(wantPrompts, res.Prompts, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
-			t.Fatalf("ListPrompts() mismatch (-want +got):\n%s", diff)
-		}
-	})
-	t.Run("PromptsIterator", func(t *testing.T) {
-		wantPrompts := []mcp.Prompt{*promptA.Prompt, *promptB.Prompt, *promptC.Prompt}
-		var gotPrompts []mcp.Prompt
-		seq := clientSession.Prompts(ctx, nil)
-		for prompt, err := range seq {
-			if err != nil {
-				t.Fatalf("Prompts() failed: %v", err)
-			}
-			gotPrompts = append(gotPrompts, prompt)
-		}
-		if diff := cmp.Diff(wantPrompts, gotPrompts, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
-			t.Fatalf("Prompts() mismatch (-want +got):\n%s", diff)
-		}
-	})
-}