diff --git a/docs/docs/reference/project-files/rill-yaml.md b/docs/docs/reference/project-files/rill-yaml.md index cb27076dfa3a..1eb97af463b4 100644 --- a/docs/docs/reference/project-files/rill-yaml.md +++ b/docs/docs/reference/project-files/rill-yaml.md @@ -176,6 +176,8 @@ _[object]_ - A map of key-value pairs for setting variables on your project. It - **`rill.ai.max_time_range_days`** - _[integer]_ - Maximum time range allowed for AI tool queries, in days. Set to 0 for no limit. Default: 0. + - **`rill.ai.mcp_tolerant_args`** - _[boolean]_ - Tolerantly decode MCP tool-call arguments where object/array-typed fields arrive as JSON-encoded strings, to work around a serialization bug in some MCP clients. Default: true. + - **`rill.strict_resolver_properties`** - _[boolean]_ - Return an error when a resolver contains properties not recognized by its implementation. Default: false. - **`rill.strict_model_properties`** - _[boolean]_ - Return an error when a model contains unmapped properties. Default: false. diff --git a/runtime/ai/mcp.go b/runtime/ai/mcp.go index d057540551b6..a12d9bf6852c 100644 --- a/runtime/ai/mcp.go +++ b/runtime/ai/mcp.go @@ -1,12 +1,16 @@ package ai import ( + "bytes" "context" + "encoding/json" "errors" "fmt" "time" + "github.com/google/jsonschema-go/jsonschema" "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/rilldata/rill/runtime/pkg/jsonschemautil" "go.uber.org/zap" ) @@ -80,6 +84,27 @@ func (s *Session) MCPServer(ctx context.Context) *mcp.Server { } }) + // Tolerantly decode tool arguments where object/array-typed fields arrive as JSON-encoded strings. + // This works around a serialization bug in some MCP clients (see https://github.com/anthropics/claude-code/issues/25865). + // It runs before the SDK validates the arguments against the tool's input schema. + // It can be disabled with the rill.ai.mcp_tolerant_args instance config option (enabled by default). + // The internal LLM tool-call path (CallToolWithOptions) receives tool inputs as real JSON objects from the provider API, + // so it does not need this; if that ever changes, the coercion should be applied there too. + cfg, err := s.runner.Runtime.InstanceConfig(ctx, s.instanceID) + if err != nil && !errors.Is(err, ctx.Err()) { + s.logger.Warn("failed to resolve instance config; enabling tolerant MCP argument decoding by default", zap.Error(err)) + } + if err != nil || cfg.AIMCPTolerantArgs { + srv.AddReceivingMiddleware(func(next mcp.MethodHandler) mcp.MethodHandler { + return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { + if method == "tools/call" { + s.coerceMCPToolArgs(req) + } + return next(ctx, method, req) + } + }) + } + // Add only the tools that the user has access to ctx = WithSession(ctx, s) for _, t := range s.runner.Tools { @@ -99,6 +124,42 @@ func (s *Session) MCPServer(ctx context.Context) *mcp.Server { return srv } +// coerceMCPToolArgs rewrites a tool call's raw arguments, +// JSON-decoding any string values in positions where the tool's input schema expects an object or array. +// It fails open: on any lookup or decode failure it leaves the arguments untouched, +// so the SDK's schema validation produces its normal error. +func (s *Session) coerceMCPToolArgs(req mcp.Request) { + params, ok := req.GetParams().(*mcp.CallToolParamsRaw) + if !ok || len(params.Arguments) == 0 { + return + } + t, ok := s.runner.Tools[params.Name] + if !ok || t.Spec == nil { + return + } + schema, ok := t.Spec.InputSchema.(*jsonschema.Schema) + if !ok || schema == nil { + return + } + + var args any + dec := json.NewDecoder(bytes.NewReader(params.Arguments)) + dec.UseNumber() // preserve integer fidelity across the re-marshal + if err := dec.Decode(&args); err != nil { + return + } + + coerced, changed := jsonschemautil.CoerceStringifiedJSON(schema, args) + if !changed { + return + } + data, err := json.Marshal(coerced) + if err != nil { + return + } + params.Arguments = data +} + // InternalError represents an internal error in a tool call. // This is needed because by default, downstream logic (such as the MCP middleware) treats errors returned from tool handlers as user errors, not internal errors. type InternalError struct { diff --git a/runtime/drivers/registry.go b/runtime/drivers/registry.go index f81505762dae..917a4431ecd9 100644 --- a/runtime/drivers/registry.go +++ b/runtime/drivers/registry.go @@ -137,6 +137,8 @@ type InstanceConfig struct { AIMaxTimeRangeDays int64 `mapstructure:"rill.ai.max_time_range_days"` // AIMaxMessageSizeBytes is the maximum allowed size of an AI message's contents (tool call args or results). Exceeding it results in an error. AIMaxMessageSizeBytes int64 `mapstructure:"rill.ai.max_message_size_bytes"` + // AIMCPTolerantArgs indicates whether MCP tool-call arguments with object/array fields JSON-encoded as strings are tolerantly decoded. + AIMCPTolerantArgs bool `mapstructure:"rill.ai.mcp_tolerant_args"` // StrictResolverProps indicates whether to return an error when a resolver contains properties that are not recognized by the resolver implementation. StrictResolverProps bool `mapstructure:"rill.strict_resolver_properties"` // StrictModelProps indicates whether to return an error when a model contains unmapped properties. @@ -223,6 +225,7 @@ func (i *Instance) Config() (InstanceConfig, error) { AIMaxQueryLimit: 250, AIRequireTimeRange: true, AIMaxMessageSizeBytes: 200 * 1024, // 200 KB + AIMCPTolerantArgs: true, ModelPartitionsWarnOnFailure: i.Environment == "prod", ModelTestsWarnOnFailure: i.Environment == "prod", } diff --git a/runtime/parser/schema/rillyaml.schema.yaml b/runtime/parser/schema/rillyaml.schema.yaml index ef8fa09226f7..d87c185b64a0 100644 --- a/runtime/parser/schema/rillyaml.schema.yaml +++ b/runtime/parser/schema/rillyaml.schema.yaml @@ -187,6 +187,9 @@ allOf: rill.ai.max_time_range_days: type: integer description: "Maximum time range allowed for AI tool queries, in days. Set to 0 for no limit. Default: 0." + rill.ai.mcp_tolerant_args: + type: boolean + description: "Tolerantly decode MCP tool-call arguments where object/array-typed fields arrive as JSON-encoded strings, to work around a serialization bug in some MCP clients. Default: true." rill.strict_resolver_properties: type: boolean description: "Return an error when a resolver contains properties not recognized by its implementation. Default: false." diff --git a/runtime/pkg/jsonschemautil/coerce.go b/runtime/pkg/jsonschemautil/coerce.go new file mode 100644 index 000000000000..c2e4a7b84f7b --- /dev/null +++ b/runtime/pkg/jsonschemautil/coerce.go @@ -0,0 +1,296 @@ +package jsonschemautil + +import ( + "encoding/json" + "maps" + "regexp" + "slices" + "strings" + + "github.com/google/jsonschema-go/jsonschema" +) + +// maxRefHops caps `$ref` chain resolution to guard against cyclic references. +const maxRefHops = 32 + +// CoerceStringifiedJSON walks value alongside schema and, wherever the schema +// unambiguously expects an object or array but the value is a string containing valid JSON of that kind, +// replaces the string with the parsed value. +// It returns the (possibly mutated) value and whether anything changed. +// It never errors: on any ambiguity or parse failure it leaves the value untouched +// so that downstream schema validation produces its normal error. +// It exists to work around MCP clients that JSON-encode object-typed tool arguments as strings +// (see https://github.com/anthropics/claude-code/issues/25865). +func CoerceStringifiedJSON(schema *jsonschema.Schema, value any) (any, bool) { + return coerceValue(schema, nil, value) +} + +// coerceValue implements CoerceStringifiedJSON for a single value and its subschema. +// defs holds the `$defs` visible at this point in the schema tree; +// resolveSchema extends it with the current schema's own `$defs`, +// which is required because schemas built with jsonschema.ForOptions.TypeSchemas +// carry `$defs` nested inside property subschemas rather than at the root. +func coerceValue(s *jsonschema.Schema, defs map[string]*jsonschema.Schema, value any) (any, bool) { + s, defs = resolveSchema(s, defs) + if s == nil || value == nil { + return value, false + } + + changed := false + + // If the value is a string but the schema only allows objects (or arrays), attempt to JSON-decode it. + // The type check is deliberately conservative: if the schema also allows strings (e.g. anyOf [object, string]) + // or has no type constraint at all, the value is left untouched. + if str, ok := value.(string); ok { + types, known := effectiveTypes(s, defs, 0) + if known { + var want string + if wantsOnly(types, "object") { + want = "{" + } else if wantsOnly(types, "array") { + want = "[" + } + if want != "" { + if parsed, ok := decodeJSONString(str, want); ok { + value = parsed + changed = true + } + } + } + } + + // Recurse into containers. Mutating in place is safe because the caller owns the freshly decoded value. + switch v := value.(type) { + case map[string]any: + for key, val := range v { + sub, subDefs := propertySchema(s, defs, key, 0) + if sub == nil { + continue + } + if nv, ch := coerceValue(sub, subDefs, val); ch { + v[key] = nv + changed = true + } + } + case []any: + items, itemDefs := itemsSchema(s, defs, 0) + if items != nil { + for i, item := range v { + if nv, ch := coerceValue(items, itemDefs, item); ch { + v[i] = nv + changed = true + } + } + } + } + + return value, changed +} + +// resolveSchema merges the schema's `$defs` into the visible scope and follows local `$ref` chains. +// It returns a nil schema if a ref cannot be resolved or points outside `#/$defs/`. +func resolveSchema(s *jsonschema.Schema, defs map[string]*jsonschema.Schema) (*jsonschema.Schema, map[string]*jsonschema.Schema) { + for range maxRefHops { + if s == nil { + return nil, defs + } + if len(s.Defs) > 0 { + merged := make(map[string]*jsonschema.Schema, len(defs)+len(s.Defs)) + maps.Copy(merged, defs) + maps.Copy(merged, s.Defs) + defs = merged + } + if s.Ref == "" { + return s, defs + } + // A $ref composes with its sibling keywords rather than replacing them. + // Following the ref discards the siblings, so if the node itself permits strings, + // coercion could rewrite a string the schema accepts as-is; fail open instead. + if s.Type == "string" || slices.Contains(s.Types, "string") { + return nil, defs + } + name, ok := strings.CutPrefix(s.Ref, "#/$defs/") + if !ok { + return nil, defs + } + s = defs[name] + } + return nil, defs +} + +// effectiveTypes returns the set of JSON types the schema allows, +// and whether that set could be determined. +// A schema without any type constraint (such as a free-form value field) returns known=false, +// which blocks coercion. +func effectiveTypes(s *jsonschema.Schema, defs map[string]*jsonschema.Schema, depth int) (map[string]bool, bool) { + if depth > maxRefHops { + return nil, false + } + s, defs = resolveSchema(s, defs) + if s == nil { + return nil, false + } + if s.Type != "" { + return map[string]bool{s.Type: true}, true + } + if len(s.Types) > 0 { + types := make(map[string]bool, len(s.Types)) + for _, t := range s.Types { + types[t] = true + } + return types, true + } + // allOf is a conjunction: the value must satisfy every branch, + // so the allowed types are the intersection of the branches' known type sets. + // Branches without a type constraint don't restrict the conjunction. + if len(s.AllOf) > 0 { + var types map[string]bool + known := false + for _, b := range s.AllOf { + branchTypes, ok := effectiveTypes(b, defs, depth+1) + if !ok { + continue + } + if !known { + types, known = branchTypes, true + continue + } + for t := range types { + if !branchTypes[t] { + delete(types, t) + } + } + } + if known { + return types, true + } + } + // anyOf/oneOf is a disjunction: the allowed types are the union of the branches' type sets, + // which is only known if every branch's type set is known. + if len(s.AnyOf) > 0 || len(s.OneOf) > 0 { + types := make(map[string]bool) + for _, branches := range [][]*jsonschema.Schema{s.AnyOf, s.OneOf} { + for _, b := range branches { + branchTypes, known := effectiveTypes(b, defs, depth+1) + if !known { + return nil, false + } + maps.Copy(types, branchTypes) + } + } + return types, true + } + return nil, false +} + +// wantsOnly reports whether the type set allows kind and nothing else except null. +func wantsOnly(types map[string]bool, kind string) bool { + if !types[kind] { + return false + } + for t := range types { + if t != kind && t != "null" { + return false + } + } + return true +} + +// decodeJSONString parses a string as a single JSON value starting with the given delimiter ("{" or "["). +// It uses json.Number to preserve integer fidelity across a later re-marshal. +func decodeJSONString(s, delim string) (any, bool) { + trimmed := strings.TrimSpace(s) + if !strings.HasPrefix(trimmed, delim) { + return nil, false + } + dec := json.NewDecoder(strings.NewReader(trimmed)) + dec.UseNumber() + var v any + if err := dec.Decode(&v); err != nil { + return nil, false + } + // Reject trailing content after the decoded value (dec.More() misses trailing "}" or "]"). + if strings.TrimSpace(trimmed[dec.InputOffset():]) != "" { + return nil, false + } + return v, true +} + +// propertySchema returns the subschema for a map key, along with the `$defs` scope it should be evaluated in, +// or nil if the schema does not unambiguously define one. +// depth guards against cyclic combinator branches (e.g. an allOf branch that refs back to its parent), +// whose recursion is driven purely by the schema and would otherwise overflow the stack. +func propertySchema(s *jsonschema.Schema, defs map[string]*jsonschema.Schema, key string, depth int) (*jsonschema.Schema, map[string]*jsonschema.Schema) { + if depth > maxRefHops { + return nil, nil + } + if sub, ok := s.Properties[key]; ok { + return sub, defs + } + // additionalProperties does not govern keys matched by patternProperties; fail open for those keys. + // On an invalid pattern, also fail open and leave the error to schema validation. + for pattern := range s.PatternProperties { + if matched, err := regexp.MatchString(pattern, key); err != nil || matched { + return nil, nil + } + } + if s.AdditionalProperties != nil { + return s.AdditionalProperties, defs + } + // Search allOf branches; only trust the result if exactly one branch defines the key. + // allOf is a conjunction, so a property schema found in one branch is binding. + // anyOf/oneOf are deliberately not searched: a disjunctive branch that omits the key + // still accepts it by default, so a definition found in one branch is not unambiguous. + var found *jsonschema.Schema + var foundDefs map[string]*jsonschema.Schema + for _, b := range s.AllOf { + b, branchDefs := resolveSchema(b, defs) + if b == nil { + continue + } + sub, subDefs := propertySchema(b, branchDefs, key, depth+1) + if sub == nil { + continue + } + if found != nil { + return nil, nil + } + found, foundDefs = sub, subDefs + } + return found, foundDefs +} + +// itemsSchema returns the subschema for array elements, along with the `$defs` scope it should be evaluated in, +// or nil if the schema does not unambiguously define one. +// depth guards against cyclic combinator branches; see propertySchema. +func itemsSchema(s *jsonschema.Schema, defs map[string]*jsonschema.Schema, depth int) (*jsonschema.Schema, map[string]*jsonschema.Schema) { + if depth > maxRefHops { + return nil, nil + } + // prefixItems changes which elements `items` governs (only those after the prefix); + // fail open rather than coercing tuple elements with the wrong schema. + if len(s.PrefixItems) > 0 { + return nil, nil + } + if s.Items != nil { + return s.Items, defs + } + // Search allOf branches only; see propertySchema for why anyOf/oneOf are excluded. + var found *jsonschema.Schema + var foundDefs map[string]*jsonschema.Schema + for _, b := range s.AllOf { + b, branchDefs := resolveSchema(b, defs) + if b == nil { + continue + } + sub, subDefs := itemsSchema(b, branchDefs, depth+1) + if sub == nil { + continue + } + if found != nil { + return nil, nil + } + found, foundDefs = sub, subDefs + } + return found, foundDefs +} diff --git a/runtime/pkg/jsonschemautil/coerce_test.go b/runtime/pkg/jsonschemautil/coerce_test.go new file mode 100644 index 000000000000..8e5e2a6b9cce --- /dev/null +++ b/runtime/pkg/jsonschemautil/coerce_test.go @@ -0,0 +1,455 @@ +package jsonschemautil + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/google/jsonschema-go/jsonschema" + "github.com/stretchr/testify/require" +) + +// coerceTestSchema is a trimmed version of metricsview.QueryJSONSchema. +// It is embedded here instead of imported because metricsview depends on this package. +const coerceTestSchema = `{ + "type": "object", + "properties": { + "metrics_view": {"type": "string"}, + "contents": {"type": "string"}, + "dimensions": { + "type": "array", + "items": {"$ref": "#/$defs/Dimension"} + }, + "pivot_on": { + "type": "array", + "items": {"type": "string"} + }, + "time_range": {"$ref": "#/$defs/TimeRange"}, + "where": {"$ref": "#/$defs/Expression"}, + "color": { + "anyOf": [ + {"$ref": "#/$defs/TimeRange"}, + {"type": "string"} + ] + }, + "limit": {"type": "integer"} + }, + "$defs": { + "Dimension": { + "type": "object", + "properties": { + "name": {"type": "string"} + } + }, + "TimeRange": { + "type": "object", + "properties": { + "start": {"type": "string"}, + "end": {"type": "string"} + } + }, + "Expression": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "val": {}, + "cond": {"$ref": "#/$defs/Condition"} + } + }, + "Condition": { + "type": "object", + "properties": { + "op": {"type": "string"}, + "exprs": { + "type": "array", + "items": {"$ref": "#/$defs/Expression"} + } + } + } + } +}` + +func TestCoerceStringifiedJSON(t *testing.T) { + tests := []struct { + name string + schema string + args string + want string + wantChanged bool + }{ + { + name: "stringified object via ref", + schema: coerceTestSchema, + args: `{"metrics_view": "mv", "time_range": "{\"start\": \"2024-01-01T00:00:00Z\", \"end\": \"2024-02-01T00:00:00Z\"}"}`, + want: `{"metrics_view": "mv", "time_range": {"start": "2024-01-01T00:00:00Z", "end": "2024-02-01T00:00:00Z"}}`, + wantChanged: true, + }, + { + name: "stringified expression with nested condition", + schema: coerceTestSchema, + args: `{"where": "{\"cond\": {\"op\": \"and\", \"exprs\": [{\"name\": \"country\"}]}}"}`, + want: `{"where": {"cond": {"op": "and", "exprs": [{"name": "country"}]}}}`, + wantChanged: true, + }, + { + name: "stringified array nested inside a real object", + schema: coerceTestSchema, + args: `{"where": {"cond": {"op": "and", "exprs": "[{\"name\": \"country\"}]"}}}`, + want: `{"where": {"cond": {"op": "and", "exprs": [{"name": "country"}]}}}`, + wantChanged: true, + }, + { + name: "stringified array of objects", + schema: coerceTestSchema, + args: `{"dimensions": "[{\"name\": \"country\"}]"}`, + want: `{"dimensions": [{"name": "country"}]}`, + wantChanged: true, + }, + { + name: "stringified element inside a real array", + schema: coerceTestSchema, + args: `{"dimensions": [{"name": "country"}, "{\"name\": \"device\"}"]}`, + want: `{"dimensions": [{"name": "country"}, {"name": "device"}]}`, + wantChanged: true, + }, + { + name: "untyped field is untouched", + schema: coerceTestSchema, + args: `{"where": {"name": "country", "val": "{\"looks\": \"like json\"}"}}`, + want: `{"where": {"name": "country", "val": "{\"looks\": \"like json\"}"}}`, + wantChanged: false, + }, + { + name: "string-typed field holding JSON text is untouched", + schema: coerceTestSchema, + args: `{"contents": "{\"type\": \"model\"}"}`, + want: `{"contents": "{\"type\": \"model\"}"}`, + wantChanged: false, + }, + { + name: "anyOf allowing string is untouched", + schema: coerceTestSchema, + args: `{"color": "{\"start\": \"2024-01-01T00:00:00Z\"}"}`, + want: `{"color": "{\"start\": \"2024-01-01T00:00:00Z\"}"}`, + wantChanged: false, + }, + { + name: "invalid JSON string is untouched", + schema: coerceTestSchema, + args: `{"time_range": "last 7 days"}`, + want: `{"time_range": "last 7 days"}`, + wantChanged: false, + }, + { + name: "wrong-kind JSON string is untouched", + schema: coerceTestSchema, + args: `{"time_range": "[1, 2]"}`, + want: `{"time_range": "[1, 2]"}`, + wantChanged: false, + }, + { + name: "JSON string with trailing garbage is untouched", + schema: coerceTestSchema, + args: `{"time_range": "{\"start\": \"2024-01-01T00:00:00Z\"} trailing"}`, + want: `{"time_range": "{\"start\": \"2024-01-01T00:00:00Z\"} trailing"}`, + wantChanged: false, + }, + { + name: "JSON string with trailing close brace is untouched", + schema: coerceTestSchema, + args: `{"time_range": "{\"start\": \"2024-01-01T00:00:00Z\"}}"}`, + want: `{"time_range": "{\"start\": \"2024-01-01T00:00:00Z\"}}"}`, + wantChanged: false, + }, + { + name: "JSON string with trailing close bracket is untouched", + schema: coerceTestSchema, + args: `{"dimensions": "[{\"name\": \"country\"}]]"}`, + want: `{"dimensions": "[{\"name\": \"country\"}]]"}`, + wantChanged: false, + }, + { + name: "well-formed args are unchanged", + schema: coerceTestSchema, + args: `{"metrics_view": "mv", "time_range": {"start": "2024-01-01T00:00:00Z"}, "dimensions": [{"name": "country"}], "limit": 100}`, + want: `{"metrics_view": "mv", "time_range": {"start": "2024-01-01T00:00:00Z"}, "dimensions": [{"name": "country"}], "limit": 100}`, + wantChanged: false, + }, + { + name: "entire arguments as one stringified object", + schema: coerceTestSchema, + args: `"{\"metrics_view\": \"mv\", \"time_range\": \"{\\\"start\\\": \\\"2024-01-01T00:00:00Z\\\"}\"}"`, + want: `{"metrics_view": "mv", "time_range": {"start": "2024-01-01T00:00:00Z"}}`, + wantChanged: true, + }, + { + name: "nullable type union is coerced", + schema: `{ + "type": "object", + "properties": { + "where": {"type": ["null", "object"], "properties": {"name": {"type": "string"}}} + } + }`, + args: `{"where": "{\"name\": \"country\"}"}`, + want: `{"where": {"name": "country"}}`, + wantChanged: true, + }, + { + name: "type union allowing string is untouched", + schema: `{ + "type": "object", + "properties": { + "where": {"type": ["string", "object"]} + } + }`, + args: `{"where": "{\"name\": \"country\"}"}`, + want: `{"where": "{\"name\": \"country\"}"}`, + wantChanged: false, + }, + { + name: "nested defs scope without root defs", + schema: `{ + "type": "object", + "properties": { + "where": { + "$ref": "#/$defs/Expression", + "$defs": { + "Expression": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "cond": {"$ref": "#/$defs/Condition"} + } + }, + "Condition": { + "type": "object", + "properties": { + "op": {"type": "string"} + } + } + } + } + } + }`, + args: `{"where": "{\"cond\": \"{\\\"op\\\": \\\"and\\\"}\"}"}`, + want: `{"where": {"cond": {"op": "and"}}}`, + wantChanged: true, + }, + { + name: "additionalProperties map values are coerced", + schema: `{ + "type": "object", + "properties": { + "where_per_metrics_view": { + "type": "object", + "additionalProperties": { + "type": "object", + "properties": {"name": {"type": "string"}} + } + } + } + }`, + args: `{"where_per_metrics_view": {"mv1": "{\"name\": \"country\"}"}}`, + want: `{"where_per_metrics_view": {"mv1": {"name": "country"}}}`, + wantChanged: true, + }, + { + name: "property defined in an allOf branch is coerced", + schema: `{ + "type": "object", + "allOf": [ + {"properties": {"where": {"type": "object", "properties": {"name": {"type": "string"}}}}}, + {"properties": {"mode": {"type": "string"}}} + ] + }`, + args: `{"where": "{\"name\": \"country\"}"}`, + want: `{"where": {"name": "country"}}`, + wantChanged: true, + }, + { + name: "property defined in an anyOf branch is untouched", + schema: `{ + "type": "object", + "anyOf": [ + {"properties": {"where": {"type": "object", "properties": {"name": {"type": "string"}}}}}, + {"properties": {"mode": {"type": "string"}}} + ] + }`, + args: `{"where": "{\"name\": \"country\"}"}`, + want: `{"where": "{\"name\": \"country\"}"}`, + wantChanged: false, + }, + { + name: "items defined in a oneOf branch are untouched", + schema: `{ + "type": "object", + "properties": { + "dimensions": { + "type": "array", + "oneOf": [ + {"items": {"type": "object", "properties": {"name": {"type": "string"}}}}, + {"maxItems": 10} + ] + } + } + }`, + args: `{"dimensions": ["{\"name\": \"country\"}"]}`, + want: `{"dimensions": ["{\"name\": \"country\"}"]}`, + wantChanged: false, + }, + { + name: "allOf-wrapped ref is coerced", + schema: `{ + "type": "object", + "properties": { + "time_range": {"allOf": [{"$ref": "#/$defs/TimeRange"}]} + }, + "$defs": { + "TimeRange": { + "type": "object", + "properties": {"start": {"type": "string"}} + } + } + }`, + args: `{"time_range": "{\"start\": \"2024-01-01T00:00:00Z\"}"}`, + want: `{"time_range": {"start": "2024-01-01T00:00:00Z"}}`, + wantChanged: true, + }, + { + name: "allOf branches intersecting to object are coerced", + schema: `{ + "type": "object", + "properties": { + "where": {"allOf": [{"type": ["object", "string"]}, {"type": "object"}]} + } + }`, + args: `{"where": "{\"name\": \"country\"}"}`, + want: `{"where": {"name": "country"}}`, + wantChanged: true, + }, + { + name: "allOf branch allowing string is untouched", + schema: `{ + "type": "object", + "properties": { + "where": {"allOf": [{"type": ["object", "string"]}]} + } + }`, + args: `{"where": "{\"name\": \"country\"}"}`, + want: `{"where": "{\"name\": \"country\"}"}`, + wantChanged: false, + }, + { + name: "cyclic allOf ref is untouched instead of overflowing the stack", + schema: `{ + "type": "object", + "properties": { + "where": {"$ref": "#/$defs/A"}, + "dimensions": {"$ref": "#/$defs/A"} + }, + "$defs": { + "A": {"allOf": [{"$ref": "#/$defs/A"}]} + } + }`, + args: `{"where": {"cond": "{\"op\": \"and\"}"}, "dimensions": ["{\"name\": \"country\"}"]}`, + want: `{"where": {"cond": "{\"op\": \"and\"}"}, "dimensions": ["{\"name\": \"country\"}"]}`, + wantChanged: false, + }, + { + name: "ref with a string-permitting sibling type is untouched", + schema: `{ + "type": "object", + "properties": { + "time_range": {"$ref": "#/$defs/TimeRange", "type": "string"} + }, + "$defs": { + "TimeRange": { + "type": "object", + "properties": {"start": {"type": "string"}} + } + } + }`, + args: `{"time_range": "{\"start\": \"2024-01-01T00:00:00Z\"}"}`, + want: `{"time_range": "{\"start\": \"2024-01-01T00:00:00Z\"}"}`, + wantChanged: false, + }, + { + name: "patternProperties-matched key is not coerced via additionalProperties", + schema: `{ + "type": "object", + "properties": { + "where_per_metrics_view": { + "type": "object", + "patternProperties": {"^x_": {"type": "string"}}, + "additionalProperties": {"type": "object", "properties": {"name": {"type": "string"}}} + } + } + }`, + args: `{"where_per_metrics_view": {"x_note": "{\"name\": \"country\"}", "mv1": "{\"name\": \"country\"}"}}`, + want: `{"where_per_metrics_view": {"x_note": "{\"name\": \"country\"}", "mv1": {"name": "country"}}}`, + wantChanged: true, + }, + { + name: "array with prefixItems is untouched", + schema: `{ + "type": "object", + "properties": { + "dimensions": { + "type": "array", + "prefixItems": [{"type": "string"}], + "items": {"type": "object", "properties": {"name": {"type": "string"}}} + } + } + }`, + args: `{"dimensions": ["{\"name\": \"country\"}", "{\"name\": \"state\"}"]}`, + want: `{"dimensions": ["{\"name\": \"country\"}", "{\"name\": \"state\"}"]}`, + wantChanged: false, + }, + { + name: "unresolvable ref is untouched", + schema: `{ + "type": "object", + "properties": { + "where": {"$ref": "#/$defs/Missing"} + } + }`, + args: `{"where": "{\"name\": \"country\"}"}`, + want: `{"where": "{\"name\": \"country\"}"}`, + wantChanged: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var schema jsonschema.Schema + require.NoError(t, json.Unmarshal([]byte(tt.schema), &schema)) + + var args any + require.NoError(t, json.Unmarshal([]byte(tt.args), &args)) + + got, changed := CoerceStringifiedJSON(&schema, args) + require.Equal(t, tt.wantChanged, changed) + + gotJSON, err := json.Marshal(got) + require.NoError(t, err) + require.JSONEq(t, tt.want, string(gotJSON)) + }) + } +} + +func TestCoerceStringifiedJSONPreservesLargeIntegers(t *testing.T) { + var schema jsonschema.Schema + require.NoError(t, json.Unmarshal([]byte(coerceTestSchema), &schema)) + + var args any + dec := json.NewDecoder(strings.NewReader(`{"limit": 9007199254740993, "time_range": "{\"start\": \"2024-01-01T00:00:00Z\"}"}`)) + dec.UseNumber() + require.NoError(t, dec.Decode(&args)) + + got, changed := CoerceStringifiedJSON(&schema, args) + require.True(t, changed) + + gotJSON, err := json.Marshal(got) + require.NoError(t, err) + require.Contains(t, string(gotJSON), "9007199254740993") +} diff --git a/runtime/server/mcp_test.go b/runtime/server/mcp_test.go index 62df20b4947b..d881846c5823 100644 --- a/runtime/server/mcp_test.go +++ b/runtime/server/mcp_test.go @@ -20,16 +20,18 @@ func TestMCP(t *testing.T) { Files: map[string]string{ "rill.yaml": "", "m.sql": ` -SELECT 'US' AS country +SELECT 'US' AS country, TIMESTAMP '2024-01-15 00:00:00' AS event_time `, // Metrics view "mv1.yaml": ` type: metrics_view model: m +timeseries: event_time dimensions: - column: country measures: -- expression: COUNT(*) +- name: row_count + expression: COUNT(*) explore: skip: true `, @@ -103,4 +105,119 @@ explore: // Test that it handles missing parameters _, err = conn.CallTool(t.Context(), &mcp.CallToolParams{Name: ai.GetMetricsViewName}) require.ErrorContains(t, err, "missing properties") + + // Test a query with object/array-typed arguments sent as JSON-encoded strings. + // Some MCP clients stringify nested arguments (see https://github.com/anthropics/claude-code/issues/25865); + // the server should tolerantly decode them. + res, err := conn.CallTool(t.Context(), &mcp.CallToolParams{ + Name: ai.QueryMetricsViewName, + Arguments: map[string]any{ + "metrics_view": "mv1", + "dimensions": `[{"name": "country"}]`, + "measures": []any{map[string]any{"name": "row_count"}}, + "time_range": `{"start": "2024-01-01T00:00:00Z", "end": "2024-02-01T00:00:00Z"}`, + "where": `{"cond": {"op": "eq", "exprs": [{"name": "country"}, {"val": "US"}]}}`, + "limit": 10, + }, + }) + require.NoError(t, err) + require.False(t, res.IsError) + resText := res.Content[0].(*mcp.TextContent).Text + require.Contains(t, resText, "US") + require.Contains(t, resText, "row_count") + + // Test a query where comparison_time_range, having, sort and a nested measure compute are also stringified + res, err = conn.CallTool(t.Context(), &mcp.CallToolParams{ + Name: ai.QueryMetricsViewName, + Arguments: map[string]any{ + "metrics_view": "mv1", + "dimensions": `[{"name": "country"}]`, + "measures": []any{ + map[string]any{"name": "row_count"}, + map[string]any{"name": "row_count_prev", "compute": `{"comparison_value": {"measure": "row_count"}}`}, + }, + "time_range": `{"start": "2024-01-01T00:00:00Z", "end": "2024-02-01T00:00:00Z"}`, + "comparison_time_range": `{"start": "2023-12-01T00:00:00Z", "end": "2024-01-01T00:00:00Z"}`, + "having": `{"cond": {"op": "gt", "exprs": [{"name": "row_count"}, {"val": 0}]}}`, + "sort": `[{"name": "row_count", "desc": true}]`, + "limit": 10, + }, + }) + require.NoError(t, err) + require.False(t, res.IsError) + resText = res.Content[0].(*mcp.TextContent).Text + require.Contains(t, resText, "US") + require.Contains(t, resText, "row_count_prev") + + // Test that a string that is not valid JSON still fails schema validation + _, err = conn.CallTool(t.Context(), &mcp.CallToolParams{ + Name: ai.QueryMetricsViewName, + Arguments: map[string]any{ + "metrics_view": "mv1", + "measures": []any{map[string]any{"name": "row_count"}}, + "time_range": "last 7 days", + }, + }) + require.ErrorContains(t, err, `want "object"`) +} + +func TestMCPTolerantArgsDisabled(t *testing.T) { + rt, instanceID := testruntime.NewInstanceWithOptions(t, testruntime.InstanceOptions{ + Variables: map[string]string{"rill.ai.mcp_tolerant_args": "false"}, + Files: map[string]string{ + "rill.yaml": "", + "m.sql": ` +SELECT 'US' AS country, TIMESTAMP '2024-01-15 00:00:00' AS event_time +`, + "mv1.yaml": ` +type: metrics_view +model: m +timeseries: event_time +dimensions: +- column: country +measures: +- name: row_count + expression: COUNT(*) +explore: + skip: true +`, + }, + }) + testruntime.RequireReconcileState(t, rt, instanceID, 3, 0, 0) + + srv, err := NewServer(context.Background(), &Options{}, rt, zap.NewNop(), ratelimit.NewNoop(), activity.NewNoopClient()) + require.NoError(t, err) + + httpSrv := httptest.NewServer(auth.HTTPMiddleware(srv.aud, srv.mcpHandler())) + defer httpSrv.Close() + + mcpClient := mcp.NewClient(&mcp.Implementation{Name: "mcp/test", Version: "1.0.0"}, nil) + conn, err := mcpClient.Connect(t.Context(), &mcp.StreamableClientTransport{Endpoint: httpSrv.URL}, nil) + require.NoError(t, err) + defer conn.Close() + + // With the flag disabled, stringified object arguments should fail schema validation instead of being decoded + _, err = conn.CallTool(t.Context(), &mcp.CallToolParams{ + Name: ai.QueryMetricsViewName, + Arguments: map[string]any{ + "metrics_view": "mv1", + "measures": []any{map[string]any{"name": "row_count"}}, + "time_range": `{"start": "2024-01-01T00:00:00Z", "end": "2024-02-01T00:00:00Z"}`, + }, + }) + require.ErrorContains(t, err, `want "object"`) + + // Properly-typed arguments still work + res, err := conn.CallTool(t.Context(), &mcp.CallToolParams{ + Name: ai.QueryMetricsViewName, + Arguments: map[string]any{ + "metrics_view": "mv1", + "measures": []any{map[string]any{"name": "row_count"}}, + "time_range": map[string]any{"start": "2024-01-01T00:00:00Z", "end": "2024-02-01T00:00:00Z"}, + "limit": 10, + }, + }) + require.NoError(t, err) + require.False(t, res.IsError) + require.Contains(t, res.Content[0].(*mcp.TextContent).Text, "row_count") }