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: add iterator method for tools

Browse source 
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>
This commit is contained in:
Sam Thanawalla 2025-05-28 16:50:51 +00:00
commit d45330dd8f
2 changed files with 70 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

29
mcp/client.go
View file

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

57
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) (*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...)