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