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 <[email protected]>
Reviewed-by: Robert Findley <[email protected]>

Author: jba

internal/mcp/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
 }
 

internal/mcp/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"},

internal/mcp/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)
 

internal/mcp/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

internal/mcp/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)

internal/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{}} }

internal/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)

internal/mcp/tool.go

@@ -5,8 +5,10 @@
 package mcp
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
+	"fmt"
 	"slices"
 
 	"golang.org/x/tools/internal/mcp/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, &params2.Arguments); err != nil {
+			if err := unmarshalSchema(params.Arguments, resolved, &params2.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("<!%s>", err)
+	}
+	return string(m)
+}

internal/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)
+				}
+			}
+		})
+	}
+}