refactor: apply MCP argument coercion only as a fallback after validation failure

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.

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.

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",
 	}

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."

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")
-}