internal/mcp: apply JSON schema defaults

Before validating a tool input, apply the defaults from the input
schema.

Change-Id: I23fae4087f898d51b3d1f0578d3203b954cb5f2d
Reviewed-on: https://go-review.googlesource.com/c/tools/+/680395
Reviewed-by: Robert Findley <rfindley@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

mcp/server.go

@@ -19,6 +19,7 @@ import (
 
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
 	"github.com/modelcontextprotocol/go-sdk/internal/util"
+	"github.com/modelcontextprotocol/go-sdk/jsonschema"
 )
 
 const DefaultPageSize = 1000
@@ -137,7 +138,7 @@ func (s *Server) addToolsErr(tools ...*ServerTool) error {
 			// Resolve the schemas, with no base URI. We don't expect tool schemas to
 			// refer outside of themselves.
 			if st.Tool.InputSchema != nil {
-				r, err := st.Tool.InputSchema.Resolve(nil)
+				r, err := st.Tool.InputSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
 				if err != nil {
 					return err
 				}
@@ -146,7 +147,7 @@ func (s *Server) addToolsErr(tools ...*ServerTool) error {
 
 			// TODO: uncomment when output schemas drop.
 			// if st.Tool.OutputSchema != nil {
-			// 	st.outputResolved, err := st.Tool.OutputSchema.Resolve(nil)
+			// 	st.outputResolved, err := st.Tool.OutputSchema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
 			// 	if err != nil {
 			// 		return err
 			// 	}

mcp/tool.go

@@ -151,9 +151,11 @@ func unmarshalSchema(data json.RawMessage, resolved *jsonschema.Resolved, v any)
 	if err := dec.Decode(v); err != nil {
 		return fmt.Errorf("unmarshaling: %w", err)
 	}
-	// TODO(jba): apply defaults.
 	// TODO: test with nil args.
 	if resolved != nil {
+		if err := resolved.ApplyDefaults(v); err != nil {
+			return fmt.Errorf("applying defaults from \n\t%s\nto\n\t%s:\n%w", schemaJSON(resolved.Schema()), data, err)
+		}
 		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)
 		}

mcp/tool_test.go

@@ -2,30 +2,31 @@
 // Use of this source code is governed by an MIT-style
 // license that can be found in the LICENSE file.
 
-package mcp_test
+package mcp
 
 import (
 	"context"
+	"encoding/json"
+	"reflect"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
-	"github.com/modelcontextprotocol/go-sdk/mcp"
 	"github.com/modelcontextprotocol/go-sdk/jsonschema"
 )
 
 // testToolHandler is used for type inference in TestNewTool.
-func testToolHandler[T any](context.Context, *mcp.ServerSession, *mcp.CallToolParamsFor[T]) (*mcp.CallToolResultFor[any], error) {
+func testToolHandler[T any](context.Context, *ServerSession, *CallToolParamsFor[T]) (*CallToolResultFor[any], error) {
 	panic("not implemented")
 }
 
 func TestNewTool(t *testing.T) {
 	tests := []struct {
-		tool *mcp.ServerTool
+		tool *ServerTool
 		want *jsonschema.Schema
 	}{
 		{
-			mcp.NewTool("basic", "", testToolHandler[struct {
+			NewTool("basic", "", testToolHandler[struct {
 				Name string `json:"name"`
 			}]),
 			&jsonschema.Schema{
@@ -38,8 +39,8 @@ func TestNewTool(t *testing.T) {
 			},
 		},
 		{
-			mcp.NewTool("enum", "", testToolHandler[struct{ Name string }], mcp.Input(
-				mcp.Property("Name", mcp.Enum("x", "y", "z")),
+			NewTool("enum", "", testToolHandler[struct{ Name string }], Input(
+				Property("Name", Enum("x", "y", "z")),
 			)),
 			&jsonschema.Schema{
 				Type:     "object",
@@ -51,13 +52,13 @@ func TestNewTool(t *testing.T) {
 			},
 		},
 		{
-			mcp.NewTool("required", "", testToolHandler[struct {
+			NewTool("required", "", testToolHandler[struct {
 				Name     string `json:"name"`
 				Language string `json:"language"`
 				X        int    `json:"x,omitempty"`
 				Y        int    `json:"y,omitempty"`
-			}], mcp.Input(
-				mcp.Property("x", mcp.Required(true)))),
+			}], Input(
+				Property("x", Required(true)))),
 			&jsonschema.Schema{
 				Type:     "object",
 				Required: []string{"name", "language", "x"},
@@ -71,11 +72,11 @@ func TestNewTool(t *testing.T) {
 			},
 		},
 		{
-			mcp.NewTool("set_schema", "", testToolHandler[struct {
+			NewTool("set_schema", "", testToolHandler[struct {
 				X int `json:"x,omitempty"`
 				Y int `json:"y,omitempty"`
-			}], mcp.Input(
-				mcp.Schema(&jsonschema.Schema{Type: "object"})),
+			}], Input(
+				Schema(&jsonschema.Schema{Type: "object"})),
 			),
 			&jsonschema.Schema{
 				Type: "object",
@@ -88,3 +89,42 @@ func TestNewTool(t *testing.T) {
 		}
 	}
 }
+
+func TestUnmarshalSchema(t *testing.T) {
+	schema := &jsonschema.Schema{
+		Type: "object",
+		Properties: map[string]*jsonschema.Schema{
+			"x": {Type: "integer", Default: json.RawMessage("3")},
+		},
+	}
+	resolved, err := schema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	type S struct {
+		X int `json:"x"`
+	}
+
+	for _, tt := range []struct {
+		data string
+		v    any
+		want any
+	}{
+		{`{"x": 1}`, new(S), &S{X: 1}},
+		{`{}`, new(S), &S{X: 3}},       // default applied
+		{`{"x": 0}`, new(S), &S{X: 3}}, // FAIL: should be 0. (requires double unmarshal)
+		{`{"x": 1}`, new(map[string]any), &map[string]any{"x": 1.0}},
+		{`{}`, new(map[string]any), &map[string]any{"x": 3.0}}, // default applied
+		{`{"x": 0}`, new(map[string]any), &map[string]any{"x": 0.0}},
+	} {
+		raw := json.RawMessage(tt.data)
+		if err := unmarshalSchema(raw, resolved, tt.v); err != nil {
+			t.Fatal(err)
+		}
+		if !reflect.DeepEqual(tt.v, tt.want) {
+			t.Errorf("got %#v, want %#v", tt.v, tt.want)
+		}
+
+	}
+}