From 6a6c44171c149c82b6e1c711a0eab49433da40f8 Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Tue, 10 Jun 2025 08:03:00 -0400 Subject: [PATCH] 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 LUCI-TryBot-Result: Go LUCI --- mcp/server.go | 5 ++-- mcp/tool.go | 4 ++- mcp/tool_test.go | 66 ++++++++++++++++++++++++++++++++++++++---------- 3 files changed, 59 insertions(+), 16 deletions(-) diff --git a/mcp/server.go b/mcp/server.go index 7e9de0e..b5e6d5a 100644 --- a/mcp/server.go +++ b/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 // } diff --git a/mcp/tool.go b/mcp/tool.go index 380fb5d..ced7896 100644 --- a/mcp/tool.go +++ b/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) } diff --git a/mcp/tool_test.go b/mcp/tool_test.go index ad36bab..128147e 100644 --- a/mcp/tool_test.go +++ b/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) + } + + } +}