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

internal/mcp: add iterator methods resources and prompts

Browse source 
This CL streamlines the client side pagination logic by adding a generic
helper that handles pagination. This also simplifies the existing tools
iterator method.

Change-Id: If4304deda1ccd22d0c4580d76deaefc64e5dfeb9
Reviewed-on: https://go-review.googlesource.com/c/tools/+/676975
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>
Auto-Submit: Sam Thanawalla <samthanawalla@google.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
This commit is contained in:
Sam Thanawalla 2025-05-28 17:48:35 +00:00 • committed by Gopher Robot
commit bb2347385e
6 changed files with 347 additions and 64 deletions
Download patch file Download diff file Expand all files Collapse all files

65
mcp/client.go
View file

@ -319,27 +319,72 @@ 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.
// Iteration stops at the first encountered error, which will be yielded.
func (cs *ClientSession) Tools(ctx context.Context, params *ListToolsParams) iter.Seq2[Tool, error] {
currentParams := &ListToolsParams{}
if params != nil {
*currentParams = *params
if params == nil {
params = &ListToolsParams{}
}
return func(yield func(Tool, error) bool) {
return paginate(ctx, params, cs.ListTools, func(res *ListToolsResult) []*Tool {
return res.Tools
})
}
// 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.
// Iteration stops at the first encountered error, which will be yielded.
func (cs *ClientSession) Resources(ctx context.Context, params *ListResourcesParams) iter.Seq2[Resource, error] {
if params == nil {
params = &ListResourcesParams{}
}
return paginate(ctx, params, cs.ListResources, func(res *ListResourcesResult) []*Resource {
return res.Resources
})
}
// 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.
// Iteration stops at the first encountered error, which will be yielded.
func (cs *ClientSession) Prompts(ctx context.Context, params *ListPromptsParams) iter.Seq2[Prompt, error] {
if params == nil {
params = &ListPromptsParams{}
}
return paginate(ctx, params, cs.ListPrompts, func(res *ListPromptsResult) []*Prompt {
return res.Prompts
})
}
type ListParams interface {
// Returns a pointer to the param's Cursor field.
cursorPtr() *string
}
type ListResult[T any] interface {
// Returns a pointer to the param's NextCursor field.
nextCursorPtr() *string
}
// paginate is a generic helper function to provide a paginated iterator.
func paginate[P ListParams, R ListResult[E], E any](ctx context.Context, params P, listFunc func(context.Context, P) (R, error), items func(R) []*E) iter.Seq2[E, error] {
return func(yield func(E, error) bool) {
for {
res, err := cs.ListTools(ctx, currentParams)
res, err := listFunc(ctx, params)
if err != nil {
yield(Tool{}, err)
var zero E
yield(zero, err)
return
}
for _, t := range res.Tools {
if !yield(*t, nil) {
for _, r := range items(res) {
if !yield(*r, nil) {
return
}
}
if res.NextCursor == "" {
nextCursorVal := res.nextCursorPtr()
if nextCursorVal == nil || *nextCursorVal == "" {
return
}
currentParams.Cursor = res.NextCursor
*params.cursorPtr() = *nextCursorVal
}
}
}

201
mcp/client_test.go Normal file
View file

@ -0,0 +1,201 @@
// 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
import (
"context"
"fmt"
"testing"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
"github.com/modelcontextprotocol/go-sdk/jsonschema"
)
type Item struct {
Name string
Value string
}
type ListTestParams struct {
Cursor string
}
func (p *ListTestParams) cursorPtr() *string {
return &p.Cursor
}
type ListTestResult struct {
Items []*Item
NextCursor string
}
func (r *ListTestResult) nextCursorPtr() *string {
return &r.NextCursor
}
var allItems = []*Item{
{"alpha", "val-A"},
{"bravo", "val-B"},
{"charlie", "val-C"},
{"delta", "val-D"},
{"echo", "val-E"},
{"foxtrot", "val-F"},
{"golf", "val-G"},
{"hotel", "val-H"},
{"india", "val-I"},
{"juliet", "val-J"},
{"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 {
if len(all) == 0 {
return []*ListTestResult{{Items: []*Item{}, NextCursor: ""}}
}
if pageSize <= 0 {
panic("pageSize must be greater than 0")
}
numPages := (len(all) + pageSize - 1) / pageSize // Ceiling division
var results []*ListTestResult
for i := range numPages {
startIndex := i * pageSize
endIndex := min(startIndex+pageSize, len(all)) // Use min to prevent out of bounds
nextCursor := ""
if endIndex < len(all) { // If there are more items after this page
nextCursor = fmt.Sprintf("cursor_%d", endIndex)
}
results = append(results, &ListTestResult{Items: all[startIndex:endIndex], NextCursor: nextCursor})
}
return results
}
func TestClientPaginateBasic(t *testing.T) {
ctx := context.Background()
testCases := []struct {
name string
results []*ListTestResult
mockError error
initialParams *ListTestParams
expected []Item
expectError bool
}{
{
name: "SinglePageAllItems",
results: generatePaginatedResults(allItems, len(allItems)),
expected: toItemValueSlice(allItems),
},
{
name: "MultiplePages",
results: generatePaginatedResults(allItems, 3),
expected: toItemValueSlice(allItems),
},
{
name: "EmptyResults",
results: generatePaginatedResults([]*Item{}, 10),
expected: nil,
},
{
name: "ListFuncReturnsErrorImmediately",
results: []*ListTestResult{{}},
mockError: fmt.Errorf("API error on first call"),
expected: nil,
expectError: true,
},
{
name: "InitialCursorProvided",
initialParams: &ListTestParams{Cursor: "cursor_2"},
results: generatePaginatedResults(allItems[2:], 3),
expected: toItemValueSlice(allItems[2:]),
},
{
name: "CursorBeyondAllItems",
initialParams: &ListTestParams{Cursor: "cursor_999"},
results: []*ListTestResult{{Items: []*Item{}, NextCursor: ""}},
expected: nil,
},
}
for _, tc := range testCases {
tc := tc
t.Run(tc.name, func(t *testing.T) {
listFunc := func(ctx context.Context, params *ListTestParams) (*ListTestResult, error) {
if len(tc.results) == 0 {
t.Fatalf("listFunc called but no more results defined for test case %q", tc.name)
}
res := tc.results[0]
tc.results = tc.results[1:]
var err error
if tc.mockError != nil {
err = tc.mockError
}
return res, err
}
params := tc.initialParams
if tc.initialParams == nil {
params = &ListTestParams{}
}
var gotItems []Item
var iterationErr error
seq := paginate(ctx, params, listFunc, func(r *ListTestResult) []*Item { return r.Items })
for item, err := range seq {
if err != nil {
iterationErr = err
break
}
gotItems = append(gotItems, item)
}
if tc.expectError {
if iterationErr == nil {
t.Errorf("paginate() expected an error during iteration, but got none")
}
} else {
if iterationErr != nil {
t.Errorf("paginate() got: %v, want: nil", iterationErr)
}
}
if diff := cmp.Diff(tc.expected, gotItems, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
t.Fatalf("paginate() mismatch (-want +got):\n%s", diff)
}
})
}
}
func TestClientPaginateVariousPageSizes(t *testing.T) {
ctx := context.Background()
for i := 1; i < len(allItems)+1; i++ {
testname := fmt.Sprintf("PageSize=%d", i)
t.Run(testname, func(t *testing.T) {
results := generatePaginatedResults(allItems, i)
listFunc := func(ctx context.Context, params *ListTestParams) (*ListTestResult, error) {
res := results[0]
results = results[1:]
return res, nil
}
var gotItems []Item
seq := paginate(ctx, &ListTestParams{}, listFunc, func(r *ListTestResult) []*Item { return r.Items })
for item, err := range seq {
if err != nil {
t.Fatalf("paginate() unexpected error during iteration: %v", err)
}
gotItems = append(gotItems, item)
}
if diff := cmp.Diff(toItemValueSlice(allItems), gotItems, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
t.Fatalf("paginate() mismatch (-want +got):\n%s", diff)
}
})
}
}

7
mcp/generate.go
View file

@ -332,6 +332,13 @@ func writeDecl(configName string, config typeConfig, def *jsonschema.Schema, nam
fmt.Fprintf(w, "\nfunc (x *%s) GetMeta() *Meta { return &x.Meta }", typeName)
}
if _, ok := def.Properties["cursor"]; ok {
fmt.Fprintf(w, "\nfunc (x *%s) cursorPtr() *string { return &x.Cursor }", typeName)
}
if _, ok := def.Properties["nextCursor"]; ok {
fmt.Fprintf(w, "\nfunc (x *%s) nextCursorPtr() *string { return &x.NextCursor }", typeName)
}
return nil
}

18
mcp/protocol.go
View file

@ -214,7 +214,8 @@ type ListPromptsParams struct {
Cursor string `json:"cursor,omitempty"`
}
func (x *ListPromptsParams) GetMeta() *Meta { return &x.Meta }
func (x *ListPromptsParams) GetMeta() *Meta { return &x.Meta }
func (x *ListPromptsParams) cursorPtr() *string { return &x.Cursor }
// The server's response to a prompts/list request from the client.
type ListPromptsResult struct {
@ -227,7 +228,8 @@ type ListPromptsResult struct {
Prompts []*Prompt `json:"prompts"`
}
func (x *ListPromptsResult) GetMeta() *Meta { return &x.Meta }
func (x *ListPromptsResult) GetMeta() *Meta { return &x.Meta }
func (x *ListPromptsResult) nextCursorPtr() *string { return &x.NextCursor }
type ListResourcesParams struct {
// This property is reserved by the protocol to allow clients and servers to
@ -238,7 +240,8 @@ type ListResourcesParams struct {
Cursor string `json:"cursor,omitempty"`
}
func (x *ListResourcesParams) GetMeta() *Meta { return &x.Meta }
func (x *ListResourcesParams) GetMeta() *Meta { return &x.Meta }
func (x *ListResourcesParams) cursorPtr() *string { return &x.Cursor }
// The server's response to a resources/list request from the client.
type ListResourcesResult struct {
@ -251,7 +254,8 @@ type ListResourcesResult struct {
Resources []*Resource `json:"resources"`
}
func (x *ListResourcesResult) GetMeta() *Meta { return &x.Meta }
func (x *ListResourcesResult) GetMeta() *Meta { return &x.Meta }
func (x *ListResourcesResult) nextCursorPtr() *string { return &x.NextCursor }
type ListRootsParams struct {
// This property is reserved by the protocol to allow clients and servers to
@ -282,7 +286,8 @@ type ListToolsParams struct {
Cursor string `json:"cursor,omitempty"`
}
func (x *ListToolsParams) GetMeta() *Meta { return &x.Meta }
func (x *ListToolsParams) GetMeta() *Meta { return &x.Meta }
func (x *ListToolsParams) cursorPtr() *string { return &x.Cursor }
// The server's response to a tools/list request from the client.
type ListToolsResult struct {
@ -295,7 +300,8 @@ type ListToolsResult struct {
Tools []*Tool `json:"tools"`
}
func (x *ListToolsResult) GetMeta() *Meta { return &x.Meta }
func (x *ListToolsResult) GetMeta() *Meta { return &x.Meta }
func (x *ListToolsResult) nextCursorPtr() *string { return &x.NextCursor }
// The severity of a log message.
//

116
mcp/server_example_test.go
View file

@ -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, opts *mcp.ServerOptions) (*mcp.ClientSession, *mcp.ServerSession, *mcp.Server) {
server := mcp.NewServer("server", "v0.0.1", opts)
func createSessions(ctx context.Context) (*mcp.ClientSession, *mcp.ServerSession, *mcp.Server) {
server := mcp.NewServer("server", "v0.0.1", nil)
client := mcp.NewClient("client", "v0.0.1", nil)
serverTransport, clientTransport := mcp.NewInMemoryTransports()
serverSession, err := server.Connect(ctx, serverTransport)
@ -77,42 +77,33 @@ 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}
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...)
t.Run("ListTools", func(t *testing.T) {
clientSession, serverSession, server := createSessions(ctx, nil)
defer clientSession.Close()
defer serverSession.Close()
server.AddTools(tools...)
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(wantListTools, res.Tools, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
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) {
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)
}
})
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)
}
})
}
@ -122,19 +113,36 @@ func TestListResources(t *testing.T) {
resourceB := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://banana"}}
resourceC := &mcp.ServerResource{Resource: &mcp.Resource{URI: "http://cherry"}}
resources := []*mcp.ServerResource{resourceA, resourceB, resourceC}
wantResource := []*mcp.Resource{resourceA.Resource, resourceB.Resource, resourceC.Resource}
ctx := context.Background()
clientSession, serverSession, server := createSessions(ctx, nil)
clientSession, serverSession, server := createSessions(ctx)
defer clientSession.Close()
defer serverSession.Close()
server.AddResources(resources...)
res, err := clientSession.ListResources(ctx, nil)
if err != nil {
t.Fatal("ListResources() failed:", err)
}
if diff := cmp.Diff(wantResource, res.Resources, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
t.Fatalf("ListResources() mismatch (-want +got):\n%s", diff)
}
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) {
@ -142,17 +150,33 @@ func TestListPrompts(t *testing.T) {
promptB := mcp.NewPrompt("banana", "banana prompt", testPromptHandler[struct{}])
promptC := mcp.NewPrompt("cherry", "cherry prompt", testPromptHandler[struct{}])
prompts := []*mcp.ServerPrompt{promptA, promptB, promptC}
wantPrompts := []*mcp.Prompt{promptA.Prompt, promptB.Prompt, promptC.Prompt}
ctx := context.Background()
clientSession, serverSession, server := createSessions(ctx, nil)
clientSession, serverSession, server := createSessions(ctx)
defer clientSession.Close()
defer serverSession.Close()
server.AddPrompts(prompts...)
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("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)
}
})
}

4
mcp/server_test.go
View file

@ -41,7 +41,7 @@ func getCursor(input string) string {
return cursor
}
func TestPaginateBasic(t *testing.T) {
func TestServerPaginateBasic(t *testing.T) {
testCases := []struct {
name string
initialItems []*TestItem
@ -170,7 +170,7 @@ func TestPaginateBasic(t *testing.T) {
}
}
func TestPaginateVariousPageSizes(t *testing.T) {
func TestServerPaginateVariousPageSizes(t *testing.T) {
fs := newFeatureSet(func(t *TestItem) string { return t.Name })
fs.add(allTestItems...)
// Try all possible page sizes, ensuring we get the correct list of items.