@@ -164,9 +164,9 @@ func (s *Server) RemovePrompts(names ...string) {
164164//
165165// The tool's input schema must be non-nil and have the type "object". For a tool
166166// that takes no input, or one where any input is valid, set [Tool.InputSchema] to
167- // &jsonschema.Schema{Type : "object"}.
167+ // `{"type" : "object"}`, using your preferred library or `json.RawMessage` .
168168//
169- // If present, the output schema must also have type "object".
169+ // If present, [Tool.OutputSchema] must also have type "object".
170170//
171171// When the handler is invoked as part of a CallTool request, req.Params.Arguments
172172// will be a json.RawMessage.
@@ -189,11 +189,29 @@ func (s *Server) AddTool(t *Tool, h ToolHandler) {
189189 // discovered until runtime, when the LLM sent bad data.
190190 panic (fmt .Errorf ("AddTool %q: missing input schema" , t .Name ))
191191 }
192- if t .InputSchema .Type != "object" {
192+ if s , ok := t .InputSchema .( * jsonschema. Schema ); ok && s .Type != "object" {
193193 panic (fmt .Errorf (`AddTool %q: input schema must have type "object"` , t .Name ))
194+ } else {
195+ var m map [string ]any
196+ if err := remarshal (t .InputSchema , & m ); err != nil {
197+ panic (fmt .Errorf ("AddTool %q: can't marshal input schema to a JSON object: %v" , t .Name , err ))
198+ }
199+ if typ := m ["type" ]; typ != "object" {
200+ panic (fmt .Errorf (`AddTool %q: input schema must have type "object" (got %v)` , t .Name , typ ))
201+ }
194202 }
195- if t .OutputSchema != nil && t .OutputSchema .Type != "object" {
196- panic (fmt .Errorf (`AddTool %q: output schema must have type "object"` , t .Name ))
203+ if t .OutputSchema != nil {
204+ if s , ok := t .OutputSchema .(* jsonschema.Schema ); ok && s .Type != "object" {
205+ panic (fmt .Errorf (`AddTool %q: output schema must have type "object"` , t .Name ))
206+ } else {
207+ var m map [string ]any
208+ if err := remarshal (t .OutputSchema , & m ); err != nil {
209+ panic (fmt .Errorf ("AddTool %q: can't marshal output schema to a JSON object: %v" , t .Name , err ))
210+ }
211+ if typ := m ["type" ]; typ != "object" {
212+ panic (fmt .Errorf (`AddTool %q: output schema must have type "object" (got %v)` , t .Name , typ ))
213+ }
214+ }
197215 }
198216 st := & serverTool {tool : t , handler : h }
199217 // Assume there was a change, since add replaces existing tools.
@@ -331,36 +349,50 @@ func toolForErr[In, Out any](t *Tool, h ToolHandlerFor[In, Out]) (*Tool, ToolHan
331349//
332350// TODO(rfindley): we really shouldn't ever return 'null' results. Maybe we
333351// should have a jsonschema.Zero(schema) helper?
334- func setSchema [T any ](sfield * * jsonschema.Schema , rfield * * jsonschema.Resolved ) (zero any , err error ) {
352+ func setSchema [T any ](sfield * any , rfield * * jsonschema.Resolved ) (zero any , err error ) {
353+ var internalSchema * jsonschema.Schema
335354 if * sfield == nil {
336355 rt := reflect .TypeFor [T ]()
337356 if rt .Kind () == reflect .Pointer {
338357 rt = rt .Elem ()
339358 zero = reflect .Zero (rt ).Interface ()
340359 }
341360 // TODO: we should be able to pass nil opts here.
342- * sfield , err = jsonschema .ForType (rt , & jsonschema.ForOptions {})
361+ internalSchema , err = jsonschema .ForType (rt , & jsonschema.ForOptions {})
362+ if err == nil {
363+ * sfield = internalSchema
364+ }
365+ } else {
366+ if err := remarshal (* sfield , & internalSchema ); err != nil {
367+ return zero , err
368+ }
343369 }
344370 if err != nil {
345371 return zero , err
346372 }
347- * rfield , err = ( * sfield ) .Resolve (& jsonschema.ResolveOptions {ValidateDefaults : true })
373+ * rfield , err = internalSchema .Resolve (& jsonschema.ResolveOptions {ValidateDefaults : true })
348374 return zero , err
349375}
350376
351377// AddTool adds a tool and typed tool handler to the server.
352378//
353379// If the tool's input schema is nil, it is set to the schema inferred from the
354- // In type parameter, using [jsonschema.For]. The In type argument must be a
355- // map or a struct, so that its inferred JSON Schema has type "object".
380+ // In type parameter. Types are inferred from Go types, and property
381+ // descriptions are read from the 'jsonschema' struct tag. Internally, the SDK
382+ // uses the github.com/google/jsonschema-go package for ineference and
383+ // validation. The In type argument must be a map or a struct, so that its
384+ // inferred JSON Schema has type "object", as required by the spec. As a
385+ // special case, if the In type is 'any', the tool's input schema is set to an
386+ // empty object schema value.
356387//
357388// If the tool's output schema is nil, and the Out type is not 'any', the
358389// output schema is set to the schema inferred from the Out type argument,
359- // which also must be a map or struct.
390+ // which must also be a map or struct. If the Out type is 'any', the output
391+ // schema is omitted.
360392//
361- // Unlike [Server.AddTool], AddTool does a lot automatically, and forces tools
362- // to conform to the MCP spec. See [ToolHandlerFor] for a detailed description
363- // of this automatic behavior.
393+ // Unlike [Server.AddTool], AddTool does a lot automatically, and forces
394+ // tools to conform to the MCP spec. See [ToolHandlerFor] for a detailed
395+ // description of this automatic behavior.
364396func AddTool [In , Out any ](s * Server , t * Tool , h ToolHandlerFor [In , Out ]) {
365397 tt , hh , err := toolForErr (t , h )
366398 if err != nil {
0 commit comments