modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

internal/mcp: iterators return pointers

Browse source 
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>
This commit is contained in:
Jonathan Amsterdam 2025-06-09 15:15:48 -04:00
commit 77fe91dca0
4 changed files with 154 additions and 139 deletions
Download patch file Download diff file Expand all files Collapse all files

23
mcp/client.go
View file

@ -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] {
return func(yield func(T, error) bool) {
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) {
for {
res, err := listFunc(ctx, params)
if err != nil {
var zero T
yield(zero, err)
yield(nil, err)
return
}
for _, r := range items(res) {
if !yield(*r, nil) {
if !yield(r, nil) {
return
}
}

136
mcp/client_list_test.go Normal file
View file

@ -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)
}
}

22
mcp/client_test.go
View file

@ -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)
}
})

112
mcp/server_example_test.go
View file

@ -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)
}
})
}