diff --git a/jsonschema/infer.go b/jsonschema/infer.go index 4785ef7..d84a1a8 100644 --- a/jsonschema/infer.go +++ b/jsonschema/infer.go @@ -42,15 +42,21 @@ func For[T any]() (*Schema, error) { // - complex numbers // - unsafe pointers // +// The cannot be any cycles in the types. // TODO(rfindley): we could perhaps just skip these incompatible fields. func ForType(t reflect.Type) (*Schema, error) { return typeSchema(t) } func typeSchema(t reflect.Type) (*Schema, error) { - if t.Kind() == reflect.Pointer { + // Follow pointers: the schema for *T is almost the same as for T, except that + // an explicit JSON "null" is allowed for the pointer. + allowNull := false + for t.Kind() == reflect.Pointer { + allowNull = true t = t.Elem() } + var ( s = new(Schema) err error @@ -121,6 +127,10 @@ func typeSchema(t reflect.Type) (*Schema, error) { default: return nil, fmt.Errorf("type %v is unsupported by jsonschema", t) } + if allowNull && s.Type != "" { + s.Types = []string{"null", s.Type} + s.Type = "" + } return s, nil } diff --git a/jsonschema/infer_test.go b/jsonschema/infer_test.go index d0cf0d1..9325b83 100644 --- a/jsonschema/infer_test.go +++ b/jsonschema/infer_test.go @@ -57,7 +57,7 @@ func TestForType(t *testing.T) { Properties: map[string]*schema{ "f": {Type: "integer"}, "G": {Type: "array", Items: &schema{Type: "number"}}, - "P": {Type: "boolean"}, + "P": {Types: []string{"null", "boolean"}}, "NoSkip": {Type: "string"}, }, Required: []string{"f", "G", "P"}, diff --git a/jsonschema/resolve.go b/jsonschema/resolve.go index 9e775d6..58a44b2 100644 --- a/jsonschema/resolve.go +++ b/jsonschema/resolve.go @@ -27,6 +27,10 @@ type Resolved struct { resolvedURIs map[string]*Schema } +// Schema returns the schema that was resolved. +// It must not be modified. +func (r *Resolved) Schema() *Schema { return r.root } + // A Loader reads and unmarshals the schema at uri, if any. type Loader func(uri *url.URL) (*Schema, error) diff --git a/jsonschema/util.go b/jsonschema/util.go index fcb632f..469ba1a 100644 --- a/jsonschema/util.go +++ b/jsonschema/util.go @@ -270,7 +270,7 @@ func jsonType(v reflect.Value) (string, bool) { return "string", true case reflect.Slice, reflect.Array: return "array", true - case reflect.Map: + case reflect.Map, reflect.Struct: return "object", true default: return "", false diff --git a/jsonschema/validate.go b/jsonschema/validate.go index 9165b47..9c2c1d3 100644 --- a/jsonschema/validate.go +++ b/jsonschema/validate.go @@ -662,8 +662,8 @@ func property(v reflect.Value, name string) reflect.Value { case reflect.Struct: props := structPropertiesOf(v.Type()) // Ignore nonexistent properties. - if index, ok := props[name]; ok { - return v.FieldByIndex(index) + if sf, ok := props[name]; ok { + return v.FieldByIndex(sf.Index) } return reflect.Value{} default: @@ -673,6 +673,8 @@ func property(v reflect.Value, name string) reflect.Value { // properties returns an iterator over the names and values of all properties // in v, which must be a map or a struct. +// If a struct, zero-valued properties that are marked omitempty or omitzero +// are excluded. func properties(v reflect.Value) iter.Seq2[string, reflect.Value] { return func(yield func(string, reflect.Value) bool) { switch v.Kind() { @@ -683,8 +685,14 @@ func properties(v reflect.Value) iter.Seq2[string, reflect.Value] { } } case reflect.Struct: - for name, index := range structPropertiesOf(v.Type()) { - if !yield(name, v.FieldByIndex(index)) { + for name, sf := range structPropertiesOf(v.Type()) { + val := v.FieldByIndex(sf.Index) + if val.IsZero() { + if tag, ok := sf.Tag.Lookup("json"); ok && (strings.Contains(tag, "omitempty") || strings.Contains(tag, "omitzero")) { + continue + } + } + if !yield(name, val) { return } } @@ -707,8 +715,8 @@ func numPropertiesBounds(v reflect.Value, isRequired map[string]bool) (int, int) case reflect.Struct: sp := structPropertiesOf(v.Type()) min := 0 - for prop, index := range sp { - if !v.FieldByIndex(index).IsZero() || isRequired[prop] { + for prop, sf := range sp { + if !v.FieldByIndex(sf.Index).IsZero() || isRequired[prop] { min++ } } @@ -719,7 +727,7 @@ func numPropertiesBounds(v reflect.Value, isRequired map[string]bool) (int, int) } // A propertyMap is a map from property name to struct field index. -type propertyMap = map[string][]int +type propertyMap = map[string]reflect.StructField var structProperties sync.Map // from reflect.Type to propertyMap @@ -730,10 +738,10 @@ func structPropertiesOf(t reflect.Type) propertyMap { if props, ok := structProperties.Load(t); ok { return props.(propertyMap) } - props := map[string][]int{} + props := map[string]reflect.StructField{} for _, sf := range reflect.VisibleFields(t) { if name, ok := jsonName(sf); ok { - props[name] = sf.Index + props[name] = sf } } structProperties.Store(t, props) diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index e4409bb..1231d00 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -709,4 +709,5 @@ 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{}} } diff --git a/mcp/prompt.go b/mcp/prompt.go index 28265d9..5171a04 100644 --- a/mcp/prompt.go +++ b/mcp/prompt.go @@ -42,6 +42,10 @@ func NewPrompt[TReq any](name, description string, handler func(context.Context, if schema.Type != "object" || !reflect.DeepEqual(schema.AdditionalProperties, &jsonschema.Schema{Not: &jsonschema.Schema{}}) { panic(fmt.Sprintf("handler request type must be a struct")) } + resolved, err := schema.Resolve(nil) + if err != nil { + panic(err) + } prompt := &ServerPrompt{ Prompt: &Prompt{ Name: name, @@ -70,7 +74,7 @@ func NewPrompt[TReq any](name, description string, handler func(context.Context, return nil, err } var v TReq - if err := unmarshalSchema(rawArgs, schema, &v); err != nil { + if err := unmarshalSchema(rawArgs, resolved, &v); err != nil { return nil, err } return handler(ctx, ss, v, params) diff --git a/mcp/tool.go b/mcp/tool.go index 8f7673e..a1fce96 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -5,8 +5,10 @@ package mcp import ( + "bytes" "context" "encoding/json" + "fmt" "slices" "github.com/modelcontextprotocol/go-sdk/jsonschema" @@ -39,10 +41,17 @@ func NewTool[TReq any](name, description string, handler ToolHandler[TReq], opts if err != nil { panic(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, schema, ¶ms2.Arguments); err != nil { + if err := unmarshalSchema(params.Arguments, resolved, ¶ms2.Arguments); err != nil { return nil, err } } @@ -68,15 +77,38 @@ func NewTool[TReq any](name, description string, handler ToolHandler[TReq], opts 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 -// the given schema. -func unmarshalSchema(data json.RawMessage, _ *jsonschema.Schema, v any) error { +// the given resolved schema. +func unmarshalSchema(data json.RawMessage, resolved *jsonschema.Resolved, v any) error { // TODO: use reflection to create the struct type to unmarshal into. // Separate validation from assignment. - return json.Unmarshal(data, v) + + // Disallow unknown fields. + // Otherwise, if the tool was built with a struct, the client could send extra + // fields and json.Unmarshal would ignore them, so the schema would never get + // a chance to declare the extra args invalid. + dec := json.NewDecoder(bytes.NewReader(data)) + dec.DisallowUnknownFields() + if err := dec.Decode(v); err != nil { + return fmt.Errorf("unmarshaling: %w", err) + } + 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) + } + } + return nil } // A ToolOption configures the behavior of a Tool. @@ -177,3 +209,12 @@ func Schema(schema *jsonschema.Schema) SchemaOption { *s = *schema }) } + +// schemaJSON returns the JSON value for s as a string, or a string indicating an error. +func schemaJSON(s *jsonschema.Schema) string { + m, err := json.Marshal(s) + if err != nil { + return fmt.Sprintf("", err) + } + return string(m) +} diff --git a/mcp/tool_test.go b/mcp/tool_test.go index f34dfd2..af53f0c 100644 --- a/mcp/tool_test.go +++ b/mcp/tool_test.go @@ -6,6 +6,8 @@ package mcp_test import ( "context" + "encoding/json" + "strings" "testing" "github.com/google/go-cmp/cmp" @@ -88,3 +90,76 @@ 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) + } + } + }) + } +}