diff --git a/mcp/server.go b/mcp/server.go index caff8dc..febf2e6 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -145,15 +145,17 @@ func (s *Server) RemovePrompts(names ...string) { // AddTool adds a [Tool] to the server, or replaces one with the same name. // The Tool argument must not be modified after this call. // -// The tool's input schema must be non-nil. For a tool that takes no input, -// or one where any input is valid, set [Tool.InputSchema] to the empty schema, -// &jsonschema.Schema{}. +// The tool's input schema must be non-nil and have the type "object". For a tool +// that takes no input, or one where any input is valid, set [Tool.InputSchema] to +// &jsonschema.Schema{Type: "object"}. +// +// If present, the output schema must also have type "object". // // When the handler is invoked as part of a CallTool request, req.Params.Arguments // will be a json.RawMessage. Unmarshaling the arguments and validating them against the // input schema are the handler author's responsibility. // -// Most users will prefer the top-level function [AddTool]. +// Most users should use the top-level function [AddTool]. func (s *Server) AddTool(t *Tool, h ToolHandler) { if t.InputSchema == nil { // This prevents the tool author from forgetting to write a schema where @@ -165,6 +167,9 @@ func (s *Server) AddTool(t *Tool, h ToolHandler) { if t.InputSchema.Type != "object" { panic(fmt.Errorf(`AddTool %q: input schema must have type "object"`, t.Name)) } + if t.OutputSchema != nil && t.OutputSchema.Type != "object" { + panic(fmt.Errorf(`AddTool %q: output schema must have type "object"`, t.Name)) + } st := &serverTool{tool: t, handler: h} // 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.) @@ -176,9 +181,12 @@ func (s *Server) AddTool(t *Tool, h ToolHandler) { // ToolFor returns a shallow copy of t and a [ToolHandler] that wraps h. // If the tool's input schema is nil, it is set to the schema inferred from the In -// type parameter, using [jsonschema.For]. +// type parameter, using [jsonschema.For]. The In type parameter must be a map +// or a struct, so that its inferred JSON Schema has type "object". +// // If the tool's output schema is nil and the Out type parameter is not the empty -// interface, then the output schema is set to the schema inferred from Out. +// interface, then the output schema is set to the schema inferred from Out, which +// must be a map or a struct. // // Most users will call [AddTool]. Use [ToolFor] if you wish to modify the tool's // schemas or wrap the ToolHandler before calling [Server.AddTool]. @@ -305,12 +313,7 @@ func setSchema[T any](sfield **jsonschema.Schema, rfield **jsonschema.Resolved) } // AddTool adds a tool and handler to the server. -// -// A shallow copy of the tool is made first. -// If the tool's input schema is nil, the copy's input schema is set to the schema -// inferred from the In type parameter, using [jsonschema.For]. -// If the tool's output schema is nil and the Out type parameter is not the empty -// interface, then the copy's output schema is set to the schema inferred from Out. +// It is a convenience for s.AddTool(ToolFor(t, h)). func AddTool[In, Out any](s *Server, t *Tool, h ToolHandlerFor[In, Out]) { s.AddTool(ToolFor(t, h)) } diff --git a/mcp/server_test.go b/mcp/server_test.go index 81a5961..cb1c05d 100644 --- a/mcp/server_test.go +++ b/mcp/server_test.go @@ -456,3 +456,36 @@ func TestServerSessionkeepaliveCancelOverwritten(t *testing.T) { t.Fatal("expected ServerSession.keepaliveCancel to be nil after we manually niled it and re-initialized") } } + +// panicks reports whether f() panics. +func panics(f func()) (b bool) { + defer func() { + b = recover() != nil + }() + f() + return false +} + +func TestAddTool(t *testing.T) { + // AddTool should panic if In or Out are not JSON objects. + s := NewServer(testImpl, nil) + if !panics(func() { + AddTool(s, &Tool{Name: "T1"}, func(context.Context, *CallToolRequest, string) (*CallToolResult, any, error) { return nil, nil, nil }) + }) { + t.Error("bad In: expected panic") + } + if panics(func() { + AddTool(s, &Tool{Name: "T2"}, func(context.Context, *CallToolRequest, map[string]any) (*CallToolResult, any, error) { + return nil, nil, nil + }) + }) { + t.Error("good In: expected no panic") + } + if !panics(func() { + AddTool(s, &Tool{Name: "T2"}, func(context.Context, *CallToolRequest, map[string]any) (*CallToolResult, int, error) { + return nil, 0, nil + }) + }) { + t.Error("bad Out: expected panic") + } +}