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: reorganize the tool API, keeping generics

Browse source 
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>
This commit is contained in:
Jonathan Amsterdam 2025-06-05 16:17:03 -04:00
commit f211cc8e4a
19 changed files with 328 additions and 228 deletions
Download patch file Download diff file Expand all files Collapse all files

8
README.md
View file

@ -50,11 +50,11 @@ func main() {
}
defer session.Close()
// Call a tool on the server.
params := &mcp.CallToolParams[map[string]any]{
params := &mcp.CallToolParams{
Name: "greet",
Arguments: map[string]any{"name": "you"},
}
if res, err := mcp.CallTool(ctx, session, params); err != nil {
if res, err := session.CallTool(ctx, params); err != nil {
log.Printf("CallTool failed: %v", err)
} else {
if res.IsError {
@ -82,8 +82,8 @@ type HiParams struct {
Name string `json:"name"`
}
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParams[HiParams]) (*mcp.CallToolResult, error) {
return &mcp.CallToolResult{
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[HiParams]) (*mcp.CallToolResultFor[string], error) {
return &mcp.CallToolResultFor[string]{
Content: []*mcp.Content{mcp.NewTextContent("Hi " + params.Name)},
}, nil
}

4
examples/hello/main.go
View file

@ -21,8 +21,8 @@ type HiArgs struct {
Name string `json:"name"`
}
func SayHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParams[HiArgs]) (*mcp.CallToolResult, error) {
return &mcp.CallToolResult{
func SayHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[HiArgs]) (*mcp.CallToolResultFor[any], error) {
return &mcp.CallToolResultFor[any]{
Content: []*mcp.Content{
mcp.NewTextContent("Hi " + params.Name),
},

4
examples/sse/main.go
View file

@ -19,8 +19,8 @@ type SayHiParams struct {
Name string `json:"name"`
}
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParams[SayHiParams]) (*mcp.CallToolResult, error) {
return &mcp.CallToolResult{
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[SayHiParams]) (*mcp.CallToolResultFor[any], error) {
return &mcp.CallToolResultFor[any]{
Content: []*mcp.Content{
mcp.NewTextContent("Hi " + params.Name),
},

4
internal/readme/client/client.go
View file

@ -25,11 +25,11 @@ func main() {
}
defer session.Close()
// Call a tool on the server.
params := &mcp.CallToolParams[map[string]any]{
params := &mcp.CallToolParams{
Name: "greet",
Arguments: map[string]any{"name": "you"},
}
if res, err := mcp.CallTool(ctx, session, params); err != nil {
if res, err := session.CallTool(ctx, params); err != nil {
log.Printf("CallTool failed: %v", err)
} else {
if res.IsError {

4
internal/readme/server/server.go
View file

@ -15,8 +15,8 @@ type HiParams struct {
Name string `json:"name"`
}
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParams[HiParams]) (*mcp.CallToolResult, error) {
return &mcp.CallToolResult{
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[HiParams]) (*mcp.CallToolResultFor[any], error) {
return &mcp.CallToolResultFor[any]{
Content: []*mcp.Content{mcp.NewTextContent("Hi " + params.Name)},
}, nil
}

35
mcp/client.go
View file

@ -6,8 +6,6 @@ package mcp
import (
"context"
"encoding/json"
"fmt"
"iter"
"slices"
"sync"
@ -288,36 +286,15 @@ func (cs *ClientSession) ListTools(ctx context.Context, params *ListToolsParams)
}
// CallTool calls the tool with the given name and arguments.
// Pass a [CallToolOptions] to provide additional request fields.
func (cs *ClientSession) CallTool(ctx context.Context, params *CallToolParams[json.RawMessage]) (*CallToolResult, error) {
// The arguments can be any value that marshals into a JSON object.
func (cs *ClientSession) CallTool(ctx context.Context, params *CallToolParams) (*CallToolResult, error) {
if params.Arguments == nil {
// Avoid sending nil over the wire.
params.Arguments = map[string]any{}
}
return handleSend[*CallToolResult](ctx, cs, methodCallTool, params)
}
// CallTool is a helper to call a tool with any argument type. It returns an
// error if params.Arguments fails to marshal to JSON.
func CallTool[TArgs any](ctx context.Context, cs *ClientSession, params *CallToolParams[TArgs]) (*CallToolResult, error) {
wireParams, err := toWireParams(params)
if err != nil {
return nil, err
}
return cs.CallTool(ctx, wireParams)
}
func toWireParams[TArgs any](params *CallToolParams[TArgs]) (*CallToolParams[json.RawMessage], error) {
data, err := json.Marshal(params.Arguments)
if err != nil {
return nil, fmt.Errorf("failed to marshal arguments: %v", err)
}
// The field mapping here must be kept up to date with the CallToolParams.
// This is partially enforced by TestToWireParams, which verifies that all
// comparable fields are mapped.
return &CallToolParams[json.RawMessage]{
Meta: params.Meta,
Name: params.Name,
Arguments: data,
}, nil
}
func (cs *ClientSession) SetLevel(ctx context.Context, params *SetLevelParams) error {
_, err := handleSend[*emptyResult](ctx, cs, methodSetLevel, params)
return err

26
mcp/client_test.go
View file

@ -7,7 +7,6 @@ package mcp
import (
"context"
"fmt"
"reflect"
"testing"
"github.com/google/go-cmp/cmp"
@ -200,28 +199,3 @@ func TestClientPaginateVariousPageSizes(t *testing.T) {
})
}
}
func TestToWireParams(t *testing.T) {
// This test verifies that toWireParams maps all fields.
// The Meta and Arguments fields are not comparable, so can't be checked by
// this simple test. However, this test will fail if new fields are added,
// and not handled by toWireParams.
params := &CallToolParams[struct{}]{
Name: "tool",
}
wireParams, err := toWireParams(params)
if err != nil {
t.Fatal(err)
}
v := reflect.ValueOf(wireParams).Elem()
for i := range v.Type().NumField() {
f := v.Type().Field(i)
if f.Name == "Meta" || f.Name == "Arguments" {
continue // not comparable
}
fv := v.Field(i)
if fv.Interface() == reflect.Zero(f.Type).Interface() {
t.Fatalf("toWireParams: unmapped field %q", f.Name)
}
}
}

2
mcp/cmd_test.go
View file

@ -56,7 +56,7 @@ func TestCmdTransport(t *testing.T) {
if err != nil {
log.Fatal(err)
}
got, err := mcp.CallTool(ctx, session, &mcp.CallToolParams[map[string]any]{
got, err := session.CallTool(ctx, &mcp.CallToolParams{
Name: "greet",
Arguments: map[string]any{"name": "user"},
})

5
mcp/features_test.go
View file

@ -18,8 +18,8 @@ type SayHiParams struct {
Name string `json:"name"`
}
func SayHi(ctx context.Context, cc *ServerSession, params *CallToolParams[SayHiParams]) (*CallToolResult, error) {
return &CallToolResult{
func SayHi(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[SayHiParams]) (*CallToolResultFor[any], error) {
return &CallToolResultFor[any]{
Content: []*Content{
NewTextContent("Hi " + params.Name),
},
@ -71,7 +71,6 @@ func TestFeatureSetAbove(t *testing.T) {
got := slices.Collect(fs.above(tc.above))
if diff := cmp.Diff(got, tc.want, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
t.Errorf("expected %v, got %v, (-want +got):\n%s", tc.want, got, diff)
}
}
}

37
mcp/generate.go
View file

@ -58,10 +58,9 @@ var declarations = config{
Name: "-",
Fields: config{
"Params": {
Name: "CallToolParams",
TypeParams: [][2]string{{"TArgs", "any"}},
Name: "CallToolParams",
Fields: config{
"Arguments": {Substitute: "TArgs"},
"Arguments": {Substitute: "any"},
},
},
},
@ -307,6 +306,9 @@ func loadSchema(schemaFile string) (data []byte, err error) {
return data, nil
}
// Suffixes of type names with a meta field.
var metaSuffixes = []string{"Params", "ParamsFor", "Result", "ResultFor"}
func writeDecl(configName string, config typeConfig, def *jsonschema.Schema, named map[string]*bytes.Buffer) error {
var w io.Writer = io.Discard
var typeName string
@ -321,8 +323,11 @@ func writeDecl(configName string, config typeConfig, def *jsonschema.Schema, nam
// Every Params and Result type should have a _meta property.
// Also, those with a progress token will turn into a struct; we want the progress token to
// be a map item. So replace all metas.
if strings.HasSuffix(typeName, "Params") || strings.HasSuffix(typeName, "Result") {
def.Properties["_meta"] = metaSchema
for _, s := range metaSuffixes {
if strings.HasSuffix(typeName, s) {
def.Properties["_meta"] = metaSchema
break
}
}
buf := new(bytes.Buffer)
w = buf
@ -371,6 +376,28 @@ func writeDecl(configName string, config typeConfig, def *jsonschema.Schema, nam
fmt.Fprintf(w, "\nfunc (x *%s) nextCursorPtr() *string { return &x.NextCursor }", typeName)
}
// Some types have generic counterparts.
var typeArg, substField string
switch typeName {
case "CallToolParams":
typeName = "CallToolParamsFor"
typeArg = "In"
substField = "Arguments"
case "CallToolResult":
typeName = "CallToolResultFor"
typeArg = "Out"
substField = "StructuredContent"
}
if typeArg != "" {
fmt.Fprintf(w, "\n\ntype %s[%s any]", typeName, typeArg)
if config.Fields != nil {
config.Fields[substField].Substitute = typeArg
}
if err := writeType(w, &config, def, named); err != nil {
return err
}
fmt.Fprintf(w, "\nfunc (x *%s[%s]) GetMeta() *Meta {return &x.Meta}", typeName, typeArg)
}
return nil
}

43
mcp/mcp_test.go
View file

@ -30,11 +30,11 @@ type hiParams struct {
Name string
}
func sayHi(ctx context.Context, ss *ServerSession, params *CallToolParams[hiParams]) (*CallToolResult, error) {
func sayHi(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[hiParams]) (*CallToolResultFor[any], error) {
if err := ss.Ping(ctx, nil); err != nil {
return nil, fmt.Errorf("ping failed: %v", err)
}
return &CallToolResult{Content: []*Content{NewTextContent("hi " + params.Arguments.Name)}}, nil
return &CallToolResultFor[any]{Content: []*Content{NewTextContent("hi " + params.Arguments.Name)}}, nil
}
func TestEndToEnd(t *testing.T) {
@ -164,11 +164,7 @@ func TestEndToEnd(t *testing.T) {
wantTools := []*Tool{
{
Name: "fail",
Description: "just fail",
InputSchema: &jsonschema.Schema{
Type: "object",
AdditionalProperties: falseSchema(),
},
InputSchema: nil,
},
{
Name: "greet",
@ -187,7 +183,7 @@ func TestEndToEnd(t *testing.T) {
t.Fatalf("tools/list mismatch (-want +got):\n%s", diff)
}
gotHi, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{
gotHi, err := cs.CallTool(ctx, &CallToolParams{
Name: "greet",
Arguments: map[string]any{"name": "user"},
})
@ -201,7 +197,7 @@ func TestEndToEnd(t *testing.T) {
t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff)
}
gotFail, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{
gotFail, err := cs.CallTool(ctx, &CallToolParams{
Name: "fail",
Arguments: map[string]any{},
})
@ -218,7 +214,7 @@ func TestEndToEnd(t *testing.T) {
t.Errorf("tools/call 'fail' mismatch (-want +got):\n%s", diff)
}
s.AddTools(&ServerTool{Tool: &Tool{Name: "T"}})
s.AddTools(&ServerTool{Tool: &Tool{Name: "T"}, Handler: nopHandler})
waitForNotification(t, "tools")
s.RemoveTools("T")
waitForNotification(t, "tools")
@ -416,9 +412,12 @@ var (
tools = map[string]*ServerTool{
"greet": NewTool("greet", "say hi", sayHi),
"fail": NewTool("fail", "just fail", func(context.Context, *ServerSession, *CallToolParams[struct{}]) (*CallToolResult, error) {
return nil, errTestFailure
}),
"fail": {
Tool: &Tool{Name: "fail"},
Handler: func(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
return nil, errTestFailure
},
},
}
prompts = map[string]*ServerPrompt{
@ -547,7 +546,7 @@ func TestServerClosing(t *testing.T) {
}
wg.Done()
}()
if _, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{
if _, err := cs.CallTool(ctx, &CallToolParams{
Name: "greet",
Arguments: map[string]any{"name": "user"},
}); err != nil {
@ -555,7 +554,7 @@ func TestServerClosing(t *testing.T) {
}
cc.Close()
wg.Wait()
if _, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{
if _, err := cs.CallTool(ctx, &CallToolParams{
Name: "greet",
Arguments: map[string]any{"name": "user"},
}); !errors.Is(err, ErrConnectionClosed) {
@ -607,7 +606,7 @@ func TestCancellation(t *testing.T) {
cancelled = make(chan struct{}, 1) // don't block the request
)
slowRequest := func(ctx context.Context, cc *ServerSession, params *CallToolParams[struct{}]) (*CallToolResult, error) {
slowRequest := func(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
start <- struct{}{}
select {
case <-ctx.Done():
@ -617,11 +616,15 @@ func TestCancellation(t *testing.T) {
}
return nil, nil
}
_, cs := basicConnection(t, NewTool("slow", "a slow request", slowRequest))
st := &ServerTool{
Tool: &Tool{Name: "slow"},
Handler: slowRequest,
}
_, cs := basicConnection(t, st)
defer cs.Close()
ctx, cancel := context.WithCancel(context.Background())
go CallTool(ctx, cs, &CallToolParams[struct{}]{Name: "slow"})
go cs.CallTool(ctx, &CallToolParams{Name: "slow"})
<-start
cancel()
select {
@ -732,3 +735,7 @@ func traceCalls[S Session](w io.Writer, prefix string) Middleware[S] {
// A function, because schemas must form a tree (they have hidden state).
func falseSchema() *jsonschema.Schema { return &jsonschema.Schema{Not: &jsonschema.Schema{}} }
func nopHandler(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error) {
return nil, nil
}

29
mcp/protocol.go
View file

@ -26,15 +26,25 @@ type Annotations struct {
Priority float64 `json:"priority,omitempty"`
}
type CallToolParams[TArgs any] struct {
type CallToolParams struct {
// This property is reserved by the protocol to allow clients and servers to
// attach additional metadata to their responses.
Meta Meta `json:"_meta,omitempty"`
Arguments TArgs `json:"arguments,omitempty"`
Arguments any `json:"arguments,omitempty"`
Name string `json:"name"`
}
func (x *CallToolParams[TArgs]) GetMeta() *Meta { return &x.Meta }
func (x *CallToolParams) GetMeta() *Meta { return &x.Meta }
type CallToolParamsFor[In any] struct {
// This property is reserved by the protocol to allow clients and servers to
// attach additional metadata to their responses.
Meta Meta `json:"_meta,omitempty"`
Arguments In `json:"arguments,omitempty"`
Name string `json:"name"`
}
func (x *CallToolParamsFor[In]) GetMeta() *Meta { return &x.Meta }
// The server's response to a tool call.
//
@ -59,6 +69,19 @@ type CallToolResult struct {
func (x *CallToolResult) GetMeta() *Meta { return &x.Meta }
type CallToolResultFor[Out any] struct {
// This property is reserved by the protocol to allow clients and servers to
// attach additional metadata to their responses.
Meta Meta `json:"_meta,omitempty"`
Content []*Content `json:"content"`
// Whether the tool call ended in an error.
//
// If not set, this is assumed to be false (the call was successful).
IsError bool `json:"isError,omitempty"`
}
func (x *CallToolResultFor[Out]) GetMeta() *Meta { return &x.Meta }
type CancelledParams struct {
// This property is reserved by the protocol to allow clients and servers to
// attach additional metadata to their responses.

47
mcp/server.go
View file

@ -110,15 +110,56 @@ func (s *Server) RemovePrompts(names ...string) {
// AddTools adds the given tools to the server,
// replacing any with the same names.
// The arguments must not be modified after this call.
//
// AddTools panics if errors are detected.
// Call [AddToolsErr] to obtain an error instead.
func (s *Server) AddTools(tools ...*ServerTool) {
if err := s.addToolsErr(tools...); err != nil {
panic(err)
}
}
// addToolsErr is like [AddTools], but returns an error instead of panicking.
func (s *Server) addToolsErr(tools ...*ServerTool) error {
// Only notify if something could change.
if len(tools) == 0 {
return
return nil
}
// Wrap the user's Handlers with rawHandlers that take a json.RawMessage.
for _, st := range tools {
if st.rawHandler == nil {
// This ServerTool was not created with NewTool.
if st.Handler == nil {
return fmt.Errorf("AddTools: tool %q has no handler", st.Tool.Name)
}
st.rawHandler = newRawHandler(st)
// Resolve the schemas, with no base URI. We don't expect tool schemas to
// refer outside of themselves.
if st.Tool.InputSchema != nil {
r, err := st.Tool.InputSchema.Resolve(nil)
if err != nil {
return err
}
st.inputResolved = r
}
// TODO: uncomment when output schemas drop.
// if st.Tool.OutputSchema != nil {
// st.outputResolved, err := st.Tool.OutputSchema.Resolve(nil)
// if err != nil {
// return err
// }
// }
}
}
// Assume there was a change, since add replaces existing tools.
// (It's possible a tool was replaced with an identical one, but not worth checking.)
// TODO: surface notify error here?
s.changeAndNotify(notificationToolListChanged, &ToolListChangedParams{},
func() bool { s.tools.add(tools...); return true })
return nil
}
// RemoveTools removes the tools with the given names.
@ -245,14 +286,14 @@ func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListTool
})
}
func (s *Server) callTool(ctx context.Context, cc *ServerSession, params *CallToolParams[json.RawMessage]) (*CallToolResult, error) {
func (s *Server) callTool(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) {
s.mu.Lock()
tool, ok := s.tools.get(params.Name)
s.mu.Unlock()
if !ok {
return nil, fmt.Errorf("%s: unknown tool %q", jsonrpc2.ErrInvalidParams, params.Name)
}
return tool.Handler(ctx, cc, params)
return tool.rawHandler(ctx, cc, params)
}
func (s *Server) listResources(_ context.Context, _ *ServerSession, params *ListResourcesParams) (*ListResourcesResult, error) {

7
mcp/server_example_test.go
View file

@ -20,8 +20,8 @@ type SayHiParams struct {
Name string `json:"name"`
}
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParams[SayHiParams]) (*mcp.CallToolResult, error) {
return &mcp.CallToolResult{
func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[SayHiParams]) (*mcp.CallToolResultFor[any], error) {
return &mcp.CallToolResultFor[any]{
Content: []*mcp.Content{
mcp.NewTextContent("Hi " + params.Arguments.Name),
},
@ -46,7 +46,7 @@ func ExampleServer() {
log.Fatal(err)
}
res, err := mcp.CallTool(ctx, clientSession, &mcp.CallToolParams[map[string]any]{
res, err := clientSession.CallTool(ctx, &mcp.CallToolParams{
Name: "greet",
Arguments: map[string]any{"name": "user"},
})
@ -147,7 +147,6 @@ func TestListResources(t *testing.T) {
t.Fatalf("Resources() mismatch (-want +got):\n%s", diff)
}
})
}
func TestListPrompts(t *testing.T) {

78
mcp/shared_test.go
View file

@ -5,6 +5,7 @@
package mcp
import (
"context"
"encoding/json"
"strings"
"testing"
@ -68,3 +69,80 @@ func roundTrip[T any](t *testing.T, v T) T {
}
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)
}
}
})
}
}

6
mcp/sse_example_test.go
View file

@ -18,8 +18,8 @@ type AddParams struct {
X, Y int
}
func Add(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParams[AddParams]) (*mcp.CallToolResult, error) {
return &mcp.CallToolResult{
func Add(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[AddParams]) (*mcp.CallToolResultFor[any], error) {
return &mcp.CallToolResultFor[any]{
Content: []*mcp.Content{mcp.NewTextContent(fmt.Sprintf("%d", params.Arguments.X+params.Arguments.Y))},
}, nil
}
@ -41,7 +41,7 @@ func ExampleSSEHandler() {
}
defer cs.Close()
res, err := mcp.CallTool(ctx, cs, &mcp.CallToolParams[map[string]any]{
res, err := cs.CallTool(ctx, &mcp.CallToolParams{
Name: "add",
Arguments: map[string]any{"x": 1, "y": 2},
})

2
mcp/sse_test.go
View file

@ -44,7 +44,7 @@ func TestSSEServer(t *testing.T) {
t.Fatal(err)
}
ss := <-conns
gotHi, err := CallTool(ctx, cs, &CallToolParams[map[string]any]{
gotHi, err := cs.CallTool(ctx, &CallToolParams{
Name: "greet",
Arguments: map[string]any{"name": "user"},
})

144
mcp/tool.go
View file

@ -15,48 +15,116 @@ import (
)
// A ToolHandler handles a call to tools/call.
type ToolHandler[TArgs any] func(context.Context, *ServerSession, *CallToolParams[TArgs]) (*CallToolResult, error)
// [CallToolParams.Arguments] will contain a map[string]any that has been validated
// against the input schema.
// Perhaps this should be an alias for ToolHandlerFor[map[string]any, map[string]any].
type ToolHandler func(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error)
// A ToolHandlerFor[Treq, Tres] handles a call to tools/call with typed arguments and results.
type ToolHandlerFor[Treq, Tres any] func(context.Context, *ServerSession, *CallToolParamsFor[Treq]) (*CallToolResultFor[Tres], error)
// A rawToolHandler is like a ToolHandler, but takes the arguments as as json.RawMessage.
type rawToolHandler = func(context.Context, *ServerSession, *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error)
// A Tool is a tool definition that is bound to a tool handler.
type ServerTool struct {
Tool *Tool
Handler ToolHandler[json.RawMessage]
Handler ToolHandler
// Set in NewTool or Server.AddToolsErr.
rawHandler rawToolHandler
// Resolved tool schemas. Set in Server.AddToolsErr.
inputResolved, outputResolved *jsonschema.Resolved
}
// NewTool is a helper to make a tool using reflection on the given handler.
// NewTool is a helper to make a tool using reflection on the given type parameters.
// When the tool is called, CallToolParams.Arguments will be of type TReq.
//
// If provided, variadic [ToolOption] values may be used to customize the tool.
//
// The input schema for the tool is extracted from the request type for the
// handler, and used to unmmarshal and validate requests to the handler. This
// schema may be customized using the [Input] option.
//
// The handler request type must translate to a valid schema, as documented by
// [jsonschema.ForType]; otherwise, NewTool panics.
//
// TODO: just have the handler return a CallToolResult: returning []Content is
// going to be inconsistent with other server features.
func NewTool[TReq any](name, description string, handler ToolHandler[TReq], opts ...ToolOption) *ServerTool {
schema, err := jsonschema.For[TReq]()
func NewTool[TReq, TRes any](name, description string, handler ToolHandlerFor[TReq, TRes], opts ...ToolOption) *ServerTool {
st, err := newToolErr[TReq, TRes](name, description, handler, opts...)
if err != nil {
panic(err)
panic(fmt.Errorf("NewTool(%q): %w", name, err))
}
// We must resolve the schema after the ToolOptions have had a chance to update it.
// But the handler needs access to the resolved schema, and the options may change
// the handler too.
// The best we can do is use the resolved schema in our own wrapped handler,
// and hope that no ToolOption replaces it.
// TODO(jba): at a minimum, document this.
var resolved *jsonschema.Resolved
wrapped := func(ctx context.Context, cc *ServerSession, params *CallToolParams[json.RawMessage]) (*CallToolResult, error) {
var params2 CallToolParams[TReq]
if params.Arguments != nil {
if err := unmarshalSchema(params.Arguments, resolved, &params2.Arguments); err != nil {
return nil, err
}
return st
}
func newToolErr[TReq, TRes any](name, description string, handler ToolHandlerFor[TReq, TRes], opts ...ToolOption) (*ServerTool, error) {
// TODO: check that TReq is a struct.
ischema, err := jsonschema.For[TReq]()
if err != nil {
return nil, err
}
// TODO: uncomment when output schemas drop.
// oschema, err := jsonschema.For[TRes]()
// if err != nil {
// return nil, err
// }
t := &ServerTool{
Tool: &Tool{
Name: name,
Description: description,
InputSchema: ischema,
// OutputSchema: oschema,
},
}
for _, opt := range opts {
opt.set(t)
}
t.rawHandler = func(ctx context.Context, ss *ServerSession, rparams *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) {
var args TReq
if err := unmarshalSchema(rparams.Arguments, t.inputResolved, &args); err != nil {
return nil, err
}
res, err := handler(ctx, cc, &params2)
// TODO: investigate why server errors are embedded in this strange way,
// TODO(jba): future-proof this copy.
params := &CallToolParamsFor[TReq]{
Meta: rparams.Meta,
Name: rparams.Name,
Arguments: args,
}
res, err := handler(ctx, ss, params)
if err != nil {
return nil, err
}
var ctr CallToolResult
if res != nil {
// TODO(jba): future-proof this copy.
ctr.Meta = res.Meta
ctr.Content = res.Content
ctr.IsError = res.IsError
}
return &ctr, nil
}
return t, nil
}
// newRawHandler creates a rawToolHandler for tools not created through NewTool.
// It unmarshals the arguments into a map[string]any and validates them against the
// schema, then calls the ServerTool's handler.
func newRawHandler(st *ServerTool) rawToolHandler {
if st.Handler == nil {
panic("st.Handler is nil")
}
return func(ctx context.Context, ss *ServerSession, rparams *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) {
// Unmarshal the args into what should be a map.
var args map[string]any
if err := unmarshalSchema(rparams.Arguments, st.inputResolved, &args); err != nil {
return nil, err
}
// TODO: generate copy
params := &CallToolParamsFor[map[string]any]{
Meta: rparams.Meta,
Name: rparams.Name,
Arguments: args,
}
res, err := st.Handler(ctx, ss, params)
// TODO(rfindley): investigate why server errors are embedded in this strange way,
// rather than returned as jsonrpc2 server errors.
if err != nil {
return &CallToolResult{
@ -66,26 +134,6 @@ func NewTool[TReq any](name, description string, handler ToolHandler[TReq], opts
}
return res, nil
}
t := &ServerTool{
Tool: &Tool{
Name: name,
Description: description,
InputSchema: schema,
},
Handler: wrapped,
}
for _, opt := range opts {
opt.set(t)
}
if schema := t.Tool.InputSchema; schema != nil {
// Resolve the schema, with no base URI. We don't expect tool schemas to
// refer outside of themselves.
resolved, err = schema.Resolve(nil)
if err != nil {
panic(fmt.Errorf("resolving input schema %s: %w", schemaJSON(schema), err))
}
}
return t
}
// unmarshalSchema unmarshals data into v and validates the result according to
@ -103,6 +151,8 @@ func unmarshalSchema(data json.RawMessage, resolved *jsonschema.Resolved, v any)
if err := dec.Decode(v); err != nil {
return fmt.Errorf("unmarshaling: %w", err)
}
// TODO(jba): apply defaults.
// TODO: test with nil args.
if resolved != nil {
if err := resolved.Validate(v); err != nil {
return fmt.Errorf("validating\n\t%s\nagainst\n\t %s:\n %w", data, schemaJSON(resolved.Schema()), err)

77
mcp/tool_test.go
View file

@ -6,8 +6,6 @@ package mcp_test
import (
"context"
"encoding/json"
"strings"
"testing"
"github.com/google/go-cmp/cmp"
@ -17,7 +15,7 @@ import (
)
// testToolHandler is used for type inference in TestNewTool.
func testToolHandler[T any](context.Context, *mcp.ServerSession, *mcp.CallToolParams[T]) (*mcp.CallToolResult, error) {
func testToolHandler[T any](context.Context, *mcp.ServerSession, *mcp.CallToolParamsFor[T]) (*mcp.CallToolResultFor[any], error) {
panic("not implemented")
}
@ -90,76 +88,3 @@ func TestNewTool(t *testing.T) {
}
}
}
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, *mcp.ServerSession, *mcp.CallToolParams[req]) (*mcp.CallToolResult, error) {
return nil, nil
}
tool := mcp.NewTool("test", "test", dummyHandler)
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.Handler(context.Background(), nil,
&mcp.CallToolParams[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)
}
}
})
}
}