From a07ce11b370051eb6b3b6a972d1af7bf5d2a97c7 Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Tue, 6 May 2025 06:55:07 -0400 Subject: [PATCH] internal/mcp: validate tool input schemas A tool now validates its input with its InputSchema. Change schema inference to allow explicit nulls for fields of pointer type. We assume the schema has no external references. For example, the following schema cannot be handled: { "$ref": "https://example.com/other.json" } Schemas with internal references, like to a "$defs", are fine. Change-Id: I6ee7c18c2c5cb609df0b22a66da986f7ea64bbe4 Reviewed-on: https://go-review.googlesource.com/c/tools/+/670676 LUCI-TryBot-Result: Go LUCI Reviewed-by: Robert Findley --- jsonschema/infer.go | 12 ++++++- jsonschema/infer_test.go | 2 +- jsonschema/resolve.go | 4 +++ jsonschema/util.go | 2 +- jsonschema/validate.go | 26 +++++++++----- mcp/mcp_test.go | 1 + mcp/prompt.go | 6 +++- mcp/tool.go | 49 +++++++++++++++++++++++--- mcp/tool_test.go | 75 ++++++++++++++++++++++++++++++++++++++++ 9 files changed, 160 insertions(+), 17 deletions(-) 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) + } + } + }) + } +}