From f211cc8e4a00add087affa4347f2d49d5e5b2cfb Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Thu, 5 Jun 2025 16:17:03 -0400 Subject: [PATCH] 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 Auto-Submit: Jonathan Amsterdam Reviewed-by: Robert Findley --- README.md | 8 +- examples/hello/main.go | 4 +- examples/sse/main.go | 4 +- internal/readme/client/client.go | 4 +- internal/readme/server/server.go | 4 +- mcp/client.go | 35 ++------ mcp/client_test.go | 26 ------ mcp/cmd_test.go | 2 +- mcp/features_test.go | 5 +- mcp/generate.go | 37 ++++++-- mcp/mcp_test.go | 43 +++++---- mcp/protocol.go | 29 ++++++- mcp/server.go | 47 +++++++++- mcp/server_example_test.go | 7 +- mcp/shared_test.go | 78 +++++++++++++++++ mcp/sse_example_test.go | 6 +- mcp/sse_test.go | 2 +- mcp/tool.go | 144 +++++++++++++++++++++---------- mcp/tool_test.go | 77 +---------------- 19 files changed, 331 insertions(+), 231 deletions(-) diff --git a/README.md b/README.md index 4f27200..11b05b4 100644 --- a/README.md +++ b/README.md @@ -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 } diff --git a/examples/hello/main.go b/examples/hello/main.go index afe087c..b8c02b8 100644 --- a/examples/hello/main.go +++ b/examples/hello/main.go @@ -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), }, diff --git a/examples/sse/main.go b/examples/sse/main.go index 129f956..b778d9c 100644 --- a/examples/sse/main.go +++ b/examples/sse/main.go @@ -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), }, diff --git a/internal/readme/client/client.go b/internal/readme/client/client.go index b9e00b9..af2a009 100644 --- a/internal/readme/client/client.go +++ b/internal/readme/client/client.go @@ -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 { diff --git a/internal/readme/server/server.go b/internal/readme/server/server.go index 269297a..d01d170 100644 --- a/internal/readme/server/server.go +++ b/internal/readme/server/server.go @@ -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 } diff --git a/mcp/client.go b/mcp/client.go index e30eec9..15d4332 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -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 diff --git a/mcp/client_test.go b/mcp/client_test.go index 218fafa..9d24a92 100644 --- a/mcp/client_test.go +++ b/mcp/client_test.go @@ -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) - } - } -} diff --git a/mcp/cmd_test.go b/mcp/cmd_test.go index 79e89bf..9a96d25 100644 --- a/mcp/cmd_test.go +++ b/mcp/cmd_test.go @@ -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"}, }) diff --git a/mcp/features_test.go b/mcp/features_test.go index 7ff3536..ba836ce 100644 --- a/mcp/features_test.go +++ b/mcp/features_test.go @@ -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) - } } } diff --git a/mcp/generate.go b/mcp/generate.go index d91cf0b..8ed7678 100644 --- a/mcp/generate.go +++ b/mcp/generate.go @@ -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 } diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 074e241..fa0ddb3 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -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 +} diff --git a/mcp/protocol.go b/mcp/protocol.go index 05c7007..279c3e0 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -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. diff --git a/mcp/server.go b/mcp/server.go index 351143b..7e9de0e 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -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) { diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index 9855210..e233c23 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -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) { diff --git a/mcp/shared_test.go b/mcp/shared_test.go index 82aeae7..7aa7b60 100644 --- a/mcp/shared_test.go +++ b/mcp/shared_test.go @@ -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) + } + } + }) + } +} diff --git a/mcp/sse_example_test.go b/mcp/sse_example_test.go index 4f0e4c3..12cb91e 100644 --- a/mcp/sse_example_test.go +++ b/mcp/sse_example_test.go @@ -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}, }) diff --git a/mcp/sse_test.go b/mcp/sse_test.go index f122cd3..5ec06cd 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -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"}, }) diff --git a/mcp/tool.go b/mcp/tool.go index a1fce96..380fb5d 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -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, ¶ms2.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, ¶ms2) - // 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) diff --git a/mcp/tool_test.go b/mcp/tool_test.go index af53f0c..ad36bab 100644 --- a/mcp/tool_test.go +++ b/mcp/tool_test.go @@ -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) - } - } - }) - } -}