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 <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>

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

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)

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, &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)
+}

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