diff --git a/docs/docs/reference/project-files/rill-yaml.md b/docs/docs/reference/project-files/rill-yaml.md index 1eb97af463b..cb27076dfa3 100644 --- a/docs/docs/reference/project-files/rill-yaml.md +++ b/docs/docs/reference/project-files/rill-yaml.md @@ -176,8 +176,6 @@ _[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 a12d9bf6852..cb6610e0ec2 100644 --- a/runtime/ai/mcp.go +++ b/runtime/ai/mcp.go @@ -9,6 +9,7 @@ import ( "time" "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/jsonrpc" "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/rilldata/rill/runtime/pkg/jsonschemautil" "go.uber.org/zap" @@ -86,24 +87,29 @@ 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 coercion runs only as a fallback: if a tool call fails the SDK's input schema validation, + // the arguments are coerced and the call is retried once; calls that validate as-is are never rewritten. // 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) + srv.AddReceivingMiddleware(func(next mcp.MethodHandler) mcp.MethodHandler { + return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { + res, err := next(ctx, method, req) + if method != "tools/call" || err == nil { + return res, err } - }) - } + var jsonrpcErr *jsonrpc.Error + if !errors.As(err, &jsonrpcErr) || jsonrpcErr.Code != jsonrpc.CodeInvalidParams { + return res, err + } + if !s.coerceMCPToolArgs(req) { + return res, err + } + if params, ok := req.GetParams().(*mcp.CallToolParamsRaw); ok { + s.logger.Debug("tolerantly decoded stringified MCP tool call arguments; retrying the call", zap.String("tool", params.Name)) + } + return next(ctx, method, req) + } + }) // Add only the tools that the user has access to ctx = WithSession(ctx, s) @@ -126,38 +132,40 @@ func (s *Session) MCPServer(ctx context.Context) *mcp.Server { // 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 returns whether the arguments were rewritten. // 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) { +func (s *Session) coerceMCPToolArgs(req mcp.Request) bool { params, ok := req.GetParams().(*mcp.CallToolParamsRaw) if !ok || len(params.Arguments) == 0 { - return + return false } t, ok := s.runner.Tools[params.Name] if !ok || t.Spec == nil { - return + return false } schema, ok := t.Spec.InputSchema.(*jsonschema.Schema) if !ok || schema == nil { - return + return false } 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 + return false } coerced, changed := jsonschemautil.CoerceStringifiedJSON(schema, args) if !changed { - return + return false } data, err := json.Marshal(coerced) if err != nil { - return + return false } params.Arguments = data + return true } // InternalError represents an internal error in a tool call. diff --git a/runtime/drivers/registry.go b/runtime/drivers/registry.go index 917a4431ecd..f81505762da 100644 --- a/runtime/drivers/registry.go +++ b/runtime/drivers/registry.go @@ -137,8 +137,6 @@ 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. @@ -225,7 +223,6 @@ 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 d87c185b64a..ef8fa09226f 100644 --- a/runtime/parser/schema/rillyaml.schema.yaml +++ b/runtime/parser/schema/rillyaml.schema.yaml @@ -187,9 +187,6 @@ 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/server/mcp_test.go b/runtime/server/mcp_test.go index d881846c582..5d842ab9a67 100644 --- a/runtime/server/mcp_test.go +++ b/runtime/server/mcp_test.go @@ -160,64 +160,3 @@ explore: }) 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") -}