go-sdk/mcp/shared_test.go
Jonathan Amsterdam f211cc8e4a internal/mcp: reorganize the tool API, keeping generics
CallToolParams.Arguments is any.

Clients can send whatever they want.

Tool authors who create their own ToolHandler (bypassing NewTool)
will receive the arguments as a map[string]any which has been
validated against the input schema.

Tool authors who call NewTool can choose the type that the arguments
unmarshal into.

There is no double-unmarshaling and no downcasting.

There is no way for tool authors to avoid the unmarshal, and no way to
avoid the validation unless the omit an input schema, which is probably
bad for LLMs. If tool authors want this optimization, we can provide
it later.

Change-Id: Ieb201e56c7fb8a23bb1a93b27e946630ec9e79ad
Reviewed-on: https://go-review.googlesource.com/c/tools/+/679595
TryBot-Bypass: Jonathan Amsterdam <jba@google.com>
Auto-Submit: Jonathan Amsterdam <jba@google.com>
Reviewed-by: Robert Findley <rfindley@google.com>
2025-06-10 04:57:48 -07:00

148 lines
3.4 KiB
Go

// 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"
"encoding/json"
"strings"
"testing"
"github.com/google/go-cmp/cmp"
)
func TestMetaMarshal(t *testing.T) {
// Verify that Meta values round-trip.
for _, meta := range []Meta{
{Data: nil, ProgressToken: nil},
{Data: nil, ProgressToken: "p"},
{Data: map[string]any{"d": true}, ProgressToken: nil},
{Data: map[string]any{"d": true}, ProgressToken: "p"},
} {
got := roundTrip(t, meta)
if !cmp.Equal(got, meta) {
t.Errorf("\ngot %#v\nwant %#v", got, meta)
}
}
// Check errors.
for _, tt := range []struct {
meta Meta
want string
}{
{
Meta{Data: map[string]any{"progressToken": "p"}, ProgressToken: 1},
"duplicate",
},
{
Meta{ProgressToken: true},
"bad type",
},
} {
_, err := json.Marshal(tt.meta)
if err == nil || !strings.Contains(err.Error(), tt.want) {
t.Errorf("%+v: got %v, want error containing %q", tt.meta, err, tt.want)
}
}
// Accept progressToken in map if the field is nil.
// It will unmarshal by populating ProgressToken.
meta := Meta{Data: map[string]any{"progressToken": "p"}}
got := roundTrip(t, meta)
want := Meta{ProgressToken: "p"}
if !cmp.Equal(got, want) {
t.Errorf("got %+v, want %+v", got, want)
}
}
func roundTrip[T any](t *testing.T, v T) T {
t.Helper()
bytes, err := json.Marshal(v)
if err != nil {
t.Fatal(err)
}
var res T
if err := json.Unmarshal(bytes, &res); err != nil {
t.Fatal(err)
}
return res
}
// TODO(jba): this shouldn't be in this file, but tool_test.go doesn't have access to unexported symbols.
func TestNewToolValidate(t *testing.T) {
// Check that the tool returned from NewTool properly validates its input schema.
type req struct {
I int
B bool
S string `json:",omitempty"`
P *int `json:",omitempty"`
}
dummyHandler := func(context.Context, *ServerSession, *CallToolParamsFor[req]) (*CallToolResultFor[any], error) {
return nil, nil
}
tool := NewTool("test", "test", dummyHandler)
// Need to add the tool to a server to get resolved schemas.
// s := NewServer("", "", nil)
for _, tt := range []struct {
desc string
args map[string]any
want string // error should contain this string; empty for success
}{
{
"both required",
map[string]any{"I": 1, "B": true},
"",
},
{
"optional",
map[string]any{"I": 1, "B": true, "S": "foo"},
"",
},
{
"wrong type",
map[string]any{"I": 1.5, "B": true},
"cannot unmarshal",
},
{
"extra property",
map[string]any{"I": 1, "B": true, "C": 2},
"unknown field",
},
{
"value for pointer",
map[string]any{"I": 1, "B": true, "P": 3},
"",
},
{
"null for pointer",
map[string]any{"I": 1, "B": true, "P": nil},
"",
},
} {
t.Run(tt.desc, func(t *testing.T) {
raw, err := json.Marshal(tt.args)
if err != nil {
t.Fatal(err)
}
_, err = tool.rawHandler(context.Background(), nil,
&CallToolParamsFor[json.RawMessage]{Arguments: json.RawMessage(raw)})
if err == nil && tt.want != "" {
t.Error("got success, wanted failure")
}
if err != nil {
if tt.want == "" {
t.Fatalf("failed with:\n%s\nwanted success", err)
}
if !strings.Contains(err.Error(), tt.want) {
t.Fatalf("got:\n%s\nwanted to contain %q", err, tt.want)
}
}
})
}
}