diff --git a/mcp/tool.go b/mcp/tool.go index f0287782..6db4686f 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -8,6 +8,7 @@ import ( "context" "encoding/json" "fmt" + "strconv" "strings" "github.com/google/jsonschema-go/jsonschema" @@ -95,7 +96,7 @@ func applySchema(data json.RawMessage, resolved *jsonschema.Resolved, forOutput if !forOutput { v := make(map[string]any) if len(data) > 0 { - if err := internaljson.Unmarshal(data, &v); err != nil { + if err := internaljson.UnmarshalUseNumber(data, &v); err != nil { return nil, fmt.Errorf("unmarshaling arguments: %w", err) } } @@ -107,7 +108,7 @@ func applySchema(data json.RawMessage, resolved *jsonschema.Resolved, forOutput unmarshaled = v } else { if len(data) > 0 { - if err := internaljson.Unmarshal(data, &unmarshaled); err != nil { + if err := internaljson.UnmarshalUseNumber(data, &unmarshaled); err != nil { return nil, fmt.Errorf("unmarshaling output: %w", err) } } @@ -131,7 +132,11 @@ func applySchema(data json.RawMessage, resolved *jsonschema.Resolved, forOutput appliedDefaults = true } - if err := resolved.Validate(&unmarshaled); err != nil { + instance, err := validationInstance(unmarshaled) + if err != nil { + return nil, err + } + if err := resolved.Validate(&instance); err != nil { return nil, err } @@ -146,6 +151,49 @@ func applySchema(data json.RawMessage, resolved *jsonschema.Resolved, forOutput return out, nil } +// validationInstance returns a copy of v, a value decoded with UseNumber, that +// jsonschema can validate. jsonschema treats a json.Number as a string, so each +// one is replaced by an int64 or uint64 if it is an integer in range, keeping +// validation exact, and by a float64 otherwise. +func validationInstance(v any) (any, error) { + switch v := v.(type) { + case json.Number: + if i, err := v.Int64(); err == nil { + return i, nil + } + if u, err := strconv.ParseUint(v.String(), 10, 64); err == nil { + return u, nil + } + f, err := v.Float64() + if err != nil { + return nil, fmt.Errorf("number %s cannot be represented as a float64", v) + } + return f, nil + case map[string]any: + m := make(map[string]any, len(v)) + for k, e := range v { + c, err := validationInstance(e) + if err != nil { + return nil, err + } + m[k] = c + } + return m, nil + case []any: + s := make([]any, len(v)) + for i, e := range v { + c, err := validationInstance(e) + if err != nil { + return nil, err + } + s[i] = c + } + return s, nil + default: + return v, nil + } +} + // isObjectJSON reports whether data is a JSON object (i.e., starts with '{' // after any leading whitespace). Returns false for arrays, primitives, null, // or empty input. diff --git a/mcp/tool_test.go b/mcp/tool_test.go index 1e779f40..a31193bf 100644 --- a/mcp/tool_test.go +++ b/mcp/tool_test.go @@ -138,6 +138,64 @@ func TestApplySchemaOutput(t *testing.T) { } } +func TestApplySchemaNumberPrecision(t *testing.T) { + // Numbers must come out of applySchema exactly as they went in, + // and be validated exactly. + maxSafe := float64(1 << 53) + schema := &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "i": {Type: "integer"}, + "n": {Type: "number"}, + "max": {Type: "integer", Maximum: &maxSafe}, + "enum": {Enum: []any{int64(1<<53 + 1)}}, + "list": {Type: "array", Items: &jsonschema.Schema{Type: "integer"}}, + "d": {Type: "integer", Default: json.RawMessage("3")}, + }, + } + resolved, err := schema.Resolve(&jsonschema.ResolveOptions{ValidateDefaults: true}) + if err != nil { + t.Fatal(err) + } + for _, tc := range []struct { + name string + data string + want string // the "d" default is always added + wantErr string + }{ + {name: "above 2^53", data: `{"i":9007199254740993}`, want: `{"d":3,"i":9007199254740993}`}, + {name: "min int64", data: `{"i":-9223372036854775808}`, want: `{"d":3,"i":-9223372036854775808}`}, + {name: "max uint64", data: `{"i":18446744073709551615}`, want: `{"d":3,"i":18446744073709551615}`}, + {name: "long decimal", data: `{"n":0.1000000000000000000001}`, want: `{"d":3,"n":0.1000000000000000000001}`}, + {name: "array elements", data: `{"list":[9007199254740993,1]}`, want: `{"d":3,"list":[9007199254740993,1]}`}, + {name: "maximum is exact", data: `{"max":9007199254740993}`, wantErr: "maximum"}, + {name: "maximum boundary", data: `{"max":9007199254740992}`, want: `{"d":3,"max":9007199254740992}`}, + {name: "enum is exact", data: `{"enum":9007199254740993}`, want: `{"d":3,"enum":9007199254740993}`}, + {name: "enum rejects neighbor", data: `{"enum":9007199254740992}`, wantErr: "enum"}, + {name: "integer rejects fraction", data: `{"i":1.5}`, wantErr: `"integer"`}, + {name: "integer rejects string", data: `{"i":"9007199254740993"}`, wantErr: `"integer"`}, + {name: "float64 overflow", data: `{"n":1e400}`, wantErr: "1e400"}, + } { + for _, forOutput := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/forOutput=%t", tc.name, forOutput), func(t *testing.T) { + got, err := applySchema(json.RawMessage(tc.data), resolved, forOutput) + if tc.wantErr != "" { + if err == nil || !strings.Contains(err.Error(), tc.wantErr) { + t.Fatalf("got %s, %v; want error containing %q", got, err, tc.wantErr) + } + return + } + if err != nil { + t.Fatal(err) + } + if string(got) != tc.want { + t.Errorf("got %s, want %s", got, tc.want) + } + }) + } + } +} + func TestToolErrorHandling(t *testing.T) { // Construct server and add both tools at the top level server := NewServer(testImpl, nil)