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: generic handling of CallToolParams

Browse source 
Make Tool APIs future-proof by pushing generics into CallToolParams, and
using CallToolParams throughout. This has several advantages:

- We no longer need the awkward CallToolOptions type, which is almost
  but not quite like CallToolParams.
- NewTool no longer needs to define its own handler type: it accepts a
  ToolHandler[TArgs]. This also means its signature is future-proof.

Additionally, provide a `CallTool` helper (the analog of NewTool) to
serve as a more strongly typed variant of `ClientSession.CallTool`.

+ design, protocol generation

Change-Id: I80172a25bc50734dae84a541aa4776ed9161694e
Reviewed-on: https://go-review.googlesource.com/c/tools/+/676495
Reviewed-by: Jonathan Amsterdam <jba@google.com>
Auto-Submit: Robert Findley <rfindley@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
This commit is contained in:
Rob Findley 2025-05-27 14:15:09 +00:00 • committed by Gopher Robot
commit 04ca8ac03b
1 changed files with 170 additions and 108 deletions
Download patch file Download diff file Expand all files Collapse all files

44
mcp/mcp_test.go
View file

@ -30,16 +30,18 @@ type hiParams struct {
Name string
}
func sayHi(ctx context.Context, ss *ServerSession, v hiParams) ([]*Content, error) {
func sayHi(ctx context.Context, ss *ServerSession, params *CallToolParams[hiParams]) (*CallToolResult, error) {
if err := ss.Ping(ctx, nil); err != nil {
return nil, fmt.Errorf("ping failed: %v", err)
}
return []*Content{NewTextContent("hi " + v.Name)}, nil
return &CallToolResult{Content: []*Content{NewTextContent("hi " + params.Arguments.Name)}}, nil
}
func TestEndToEnd(t *testing.T) {
ctx := context.Background()
ct, st := NewInMemoryTransports()
var ct, st Transport = NewInMemoryTransports()
// ct = NewLoggingTransport(ct, os.Stderr)
// st = NewLoggingTransport(st, os.Stderr)
// Channels to check if notification callbacks happened.
notificationChans := map[string]chan int{}
@ -185,7 +187,10 @@ func TestEndToEnd(t *testing.T) {
t.Fatalf("tools/list mismatch (-want +got):\n%s", diff)
}
gotHi, err := cs.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil)
gotHi, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{
Name: "greet",
Arguments: map[string]any{"name": "user"},
})
if err != nil {
t.Fatal(err)
}
@ -196,7 +201,10 @@ func TestEndToEnd(t *testing.T) {
t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff)
}
gotFail, err := cs.CallTool(ctx, "fail", map[string]any{}, nil)
gotFail, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{
Name: "fail",
Arguments: map[string]any{},
})
// Counter-intuitively, when a tool fails, we don't expect an RPC error for
// call tool: instead, the failure is embedded in the result.
if err != nil {
@ -387,7 +395,7 @@ var (
tools = map[string]*ServerTool{
"greet": NewTool("greet", "say hi", sayHi),
"fail": NewTool("fail", "just fail", func(context.Context, *ServerSession, struct{}) ([]*Content, error) {
"fail": NewTool("fail", "just fail", func(context.Context, *ServerSession, *CallToolParams[struct{}]) (*CallToolResult, error) {
return nil, errTestFailure
}),
}
@ -506,24 +514,30 @@ func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerSession, *Clien
}
func TestServerClosing(t *testing.T) {
cc, c := basicConnection(t, NewTool("greet", "say hi", sayHi))
defer c.Close()
cc, cs := basicConnection(t, NewTool("greet", "say hi", sayHi))
defer cs.Close()
ctx := context.Background()
var wg sync.WaitGroup
wg.Add(1)
go func() {
if err := c.Wait(); err != nil {
if err := cs.Wait(); err != nil {
t.Errorf("server connection failed: %v", err)
}
wg.Done()
}()
if _, err := c.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil); err != nil {
if _, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{
Name: "greet",
Arguments: map[string]any{"name": "user"},
}); err != nil {
t.Fatalf("after connecting: %v", err)
}
cc.Close()
wg.Wait()
if _, err := c.CallTool(ctx, "greet", map[string]any{"name": "user"}, nil); !errors.Is(err, ErrConnectionClosed) {
if _, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{
Name: "greet",
Arguments: map[string]any{"name": "user"},
}); !errors.Is(err, ErrConnectionClosed) {
t.Errorf("after disconnection, got error %v, want EOF", err)
}
}
@ -572,7 +586,7 @@ func TestCancellation(t *testing.T) {
cancelled = make(chan struct{}, 1) // don't block the request
)
slowRequest := func(ctx context.Context, cc *ServerSession, v struct{}) ([]*Content, error) {
slowRequest := func(ctx context.Context, cc *ServerSession, params *CallToolParams[struct{}]) (*CallToolResult, error) {
start <- struct{}{}
select {
case <-ctx.Done():
@ -582,11 +596,11 @@ func TestCancellation(t *testing.T) {
}
return nil, nil
}
_, sc := basicConnection(t, NewTool("slow", "a slow request", slowRequest))
defer sc.Close()
_, cs := basicConnection(t, NewTool("slow", "a slow request", slowRequest))
defer cs.Close()
ctx, cancel := context.WithCancel(context.Background())
go sc.CallTool(ctx, "slow", map[string]any{}, nil)
go CallTool(ctx, cs, &CallToolParams[struct{}]{Name: "slow"})
<-start
cancel()
select {