mcp: preserve number precision during schema validation

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.

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)