fix(ai): run_sql reruns misspelt column names and keeps internal SQL out of errors

apps/ai/src/mcp/__evals__/fake-warehouse.ts

@@ -4,12 +4,19 @@ import { __testables } from "@maple/backend/services/warehouse/WarehouseQuerySer
 import { makeLargeTraceSpans, makeTraceLogs } from "./fixtures"
 import { FIXTURES } from "./utils"
 
-export interface FixtureRule {
+export type FixtureRule = {
 	/** Match against the compiled SQL the WarehouseQueryService would execute. */
 	readonly match: (sql: string) => boolean
-	/** Fixed rows, or rows computed from the SQL (to honor a service filter, say). */
-	readonly rows: ReadonlyArray<unknown> | ((sql: string) => ReadonlyArray<unknown>)
-}
+} & (
+	| {
+			/** Fixed rows, or rows computed from the SQL (to honor a service filter, say). */
+			readonly rows: ReadonlyArray<unknown> | ((sql: string) => ReadonlyArray<unknown>)
+	  }
+	| {
+			/** Fail the statement as ClickHouse would, with this message and error type. */
+			readonly failure: { readonly message: string; readonly type: string }
+	  }
+)
 
 /** A read of known traces' spans: `traces` seeked through the `trace_index` subquery. */
 export const readsTraceSpans = (sql: string): boolean =>
@@ -75,6 +82,16 @@ export const installFakeWarehouse = (
 								}),
 							)
 						}
+						if ("failure" in rule) {
+							return Effect.fail(
+								new WarehouseDriverError({
+									reason: "server",
+									status: 400,
+									message: rule.failure.message,
+									type: rule.failure.type,
+								}),
+							)
+						}
 						const rows = typeof rule.rows === "function" ? rule.rows(sql) : rule.rows
 						return Effect.succeed({ data: rows as ReadonlyArray<Record<string, unknown>> })
 					},

apps/ai/src/mcp/__evals__/run-sql-recovery.test.ts

@@ -0,0 +1,94 @@
+/**
+ * `run_sql` failures an agent hit in production, replayed against the real tool: the column
+ * names it fixes by itself and the internal SQL it keeps out of the error.
+ */
+import { afterAll, beforeAll, describe, expect, it } from "vitest"
+import { Schema } from "effect"
+import { RunSqlOutput } from "@maple/domain/mcp-outputs"
+import { makeEvalRuntime, runToolDirect, markdown, type EvalRuntime } from "./eval-runtime"
+import { installFakeWarehouse, restoreWarehouse } from "./fake-warehouse"
+
+const unknown = (name: string) => ({
+	message: `[Error] Unknown expression or function identifier \`${name}\` in scope SELECT ${name} FROM traces`,
+	type: "UNKNOWN_IDENTIFIER",
+})
+
+const GATEWAY_WRAPPED = {
+	message:
+		"[Error] Unknown expression identifier `Level` in scope SELECT Level FROM (SELECT * FROM (SELECT * FROM " +
+		"maple_us.alert_checks WHERE OrgId = 'org_x') AS maple_us.alert_checks WHERE OrgId = 'org_x') AS alert_checks " +
+		"WHERE (OrgId = 'org_x') ) AS maple_raw_sql_limited LIMIT 1001 SETTINGS max_execution_time=10",
+	type: "UNKNOWN_IDENTIFIER",
+}
+
+let rt: EvalRuntime
+
+beforeAll(() => {
+	installFakeWarehouse([
+		{ match: (sql) => sql.includes("service.name"), failure: unknown("service.name") },
+		{ match: (sql) => sql.includes("FROM alert_checks"), failure: GATEWAY_WRAPPED },
+		{ match: (sql) => /\bbogus\b/.test(sql), failure: unknown("bogus") },
+		// Every pass fixes one name and the warehouse names the next one.
+		{ match: (sql) => sql.includes("trace_id"), failure: unknown("trace_id") },
+		{ match: (sql) => sql.includes("span_id"), failure: unknown("span_id") },
+		{ match: (sql) => sql.includes("span_name"), failure: unknown("span_name") },
+		{ match: (sql) => sql.includes("status_code"), failure: unknown("status_code") },
+		{ match: (sql) => sql.includes("ServiceName"), rows: [{ ServiceName: "checkout", c: 3 }] },
+	])
+	rt = makeEvalRuntime()
+})
+
+afterAll(async () => {
+	restoreWarehouse()
+	await rt.dispose()
+})
+
+const WINDOW = { start_time: "2026-09-24 00:00:00", end_time: "2026-09-24 06:00:00" }
+
+describe("run_sql recovery", () => {
+	it("reruns a misspelt column with its one right name and says so", async () => {
+		const result = await runToolDirect(rt, "run_sql", {
+			...WINDOW,
+			sql: "SELECT service.name, count() AS c FROM traces WHERE $__orgFilter AND service.name != '' GROUP BY service.name",
+		})
+		expect(result.isError).not.toBe(true)
+		const output = Schema.decodeUnknownSync(RunSqlOutput)(result.structuredContent)
+		expect(output.rows).toEqual([{ ServiceName: "checkout", c: 3 }])
+		expect(output.rewrites).toEqual([{ from: "service.name", to: "ServiceName" }])
+		expect(markdown(result)).toContain("Ran with `ServiceName` for `service.name`")
+	})
+
+	it("says which rewrite the error belongs to when the rerun fails on another name", async () => {
+		const result = await runToolDirect(rt, "run_sql", {
+			...WINDOW,
+			sql: "SELECT service.name, bogus FROM traces WHERE $__orgFilter",
+		})
+		expect(result.isError).toBe(true)
+		const text = markdown(result)
+		expect(text).toContain("`bogus`")
+		expect(text).toContain("This error is for the query rerun with `ServiceName` for `service.name`.")
+	})
+
+	it("stops rewriting after three passes", async () => {
+		const result = await runToolDirect(rt, "run_sql", {
+			...WINDOW,
+			sql: "SELECT trace_id, span_id, span_name, status_code FROM traces WHERE $__orgFilter",
+		})
+		expect(result.isError).toBe(true)
+		const text = markdown(result)
+		expect(text).toContain("`status_code`")
+		expect(text).toContain("`TraceId` for `trace_id`, `SpanId` for `span_id`, `SpanName` for `span_name`")
+	})
+
+	it("keeps the gateway's internal database names out of the error", async () => {
+		const result = await runToolDirect(rt, "run_sql", {
+			...WINDOW,
+			sql: "SELECT Level FROM alert_checks WHERE $__orgFilter",
+		})
+		const text = markdown(result)
+		expect(text).toContain("FROM alert_checks")
+		expect(text).not.toContain("maple_us")
+		expect(text).not.toContain("maple_raw_sql_limited")
+		expect(text).not.toContain("org_x")
+	})
+})

apps/ai/src/mcp/lib/raw-sql-error-text.test.ts

@@ -0,0 +1,59 @@
+import { describe, expect, it } from "vitest"
+import { rawSqlErrorText } from "./raw-sql-error-text"
+
+describe("rawSqlErrorText", () => {
+	it("collapses the gateway's per-org subquery back to the table name", () => {
+		const message =
+			"[Error] Unknown identifier in scope SELECT x FROM (SELECT * FROM (SELECT * FROM maple_us.traces WHERE " +
+			"OrgId = 'org_a') AS maple_us.traces WHERE OrgId = 'org_a') AS traces WHERE (OrgId = 'org_a')"
+		expect(rawSqlErrorText(message)).toBe(
+			"[Error] Unknown identifier in scope SELECT x FROM traces WHERE $__orgFilter",
+		)
+	})
+
+	// A preview or staging warehouse uses its own database name.
+	it("collapses the wrapper under any database name, aliased or as a bare subquery", () => {
+		const wrap = (db: string, table: string) =>
+			`(SELECT * FROM (SELECT * FROM ${db}.${table} WHERE OrgId = 'org_a') AS ${db}.${table} WHERE OrgId = 'org_a')`
+		expect(
+			rawSqlErrorText(
+				`In scope SELECT l.x FROM ${wrap("pr_1961", "logs")} AS l WHERE (l.OrgId = 'org_a')`,
+			),
+		).toBe("In scope SELECT l.x FROM logs AS l WHERE $__orgFilter(l)")
+		expect(
+			rawSqlErrorText(
+				"Actual SELECT * FROM (SELECT * FROM pr_1961.traces WHERE OrgId = 'org_a') AS pr_1961.traces WHERE OrgId = 'org_a'.",
+			),
+		).toBe("Actual traces.")
+	})
+
+	it("keeps parentheses balanced around a wrapper with no alias", () => {
+		expect(
+			rawSqlErrorText(
+				"x IN (SELECT * FROM (SELECT * FROM pr_1.traces WHERE OrgId = 'org_a') AS pr_1.traces WHERE OrgId = 'org_a')",
+			),
+		).toBe("x IN (traces)")
+	})
+
+	it("never leaves an org id or a database-qualified table behind", () => {
+		expect(rawSqlErrorText("FROM pr_1961.traces WHERE OrgId = 'org_3KkfPhr7'")).toBe(
+			"FROM traces WHERE OrgId = '<your org>'",
+		)
+	})
+
+	it("drops the row-cap wrapper", () => {
+		expect(
+			rawSqlErrorText(
+				"in scope SELECT * FROM ( SELECT n FROM logs ) AS maple_raw_sql_limited LIMIT 1001 SETTINGS max_execution_time=10, max_threads=2",
+			),
+		).toBe("in scope SELECT n FROM logs")
+	})
+
+	it("turns a datasource permission error into a plain refusal", () => {
+		expect(
+			rawSqlErrorText(
+				"Not enough permissions for datasource 'audit_log', token needs DATASOURCES:READ:audit_log scope. curl -X PUT ...",
+			),
+		).toBe("`audit_log` cannot be read from run_sql.")
+	})
+})

apps/ai/src/mcp/lib/raw-sql-error-text.ts

@@ -0,0 +1,33 @@
+import { listWarehouseTables } from "@maple/backend/services/warehouse/warehouse-catalog"
+
+const DB = "[A-Za-z_]\\w*"
+
+const INNER = `SELECT \\* FROM \\(SELECT \\* FROM ${DB}\\.(\\w+) WHERE OrgId = '[^']*'\\) AS ${DB}\\.\\w+ WHERE OrgId = '[^']*'`
+
+/** The managed gateway's per-org subquery around a table, as `(...) AS alias` or bare. */
+const ALIASED_WRAPPER = new RegExp(`\\(${INNER}\\) AS (\\w+)`, "g")
+const BARE_WRAPPER = new RegExp(INNER, "g")
+
+/**
+ * A warehouse error rewritten in the caller's own terms. The managed gateway wraps every
+ * table in a per-org subquery under its internal database name and our row cap wraps the
+ * statement; agents copied both back into their SQL (`FROM maple_us.alert_checks`).
+ */
+export const rawSqlErrorText = (message: string): string => {
+	const denied = message.match(/Not enough permissions for datasource '([^']+)'/)
+	if (denied !== null) return `\`${denied[1]}\` cannot be read from run_sql.`
+	const tables = listWarehouseTables()
+		.map((t) => t.name)
+		.join("|")
+	return message
+		.replace(ALIASED_WRAPPER, (_whole, table: string, alias: string) =>
+			alias === table ? table : `${table} AS ${alias}`,
+		)
+		.replace(BARE_WRAPPER, "$1")
+		.replace(new RegExp(`\\b${DB}\\.(${tables})\\b`, "g"), "$1")
+		.replace(/(scope )SELECT \* FROM \(\s*/g, "$1")
+		.replace(/\s*\) AS maple_raw_sql_limited\s+LIMIT \d+(?:\s+SETTINGS [\w=, ]+)?(?:\s+FORMAT \w+)?/g, "")
+		.replace(/\((\w+)\.OrgId = 'org_\w+'\)/g, "$__orgFilter($1)")
+		.replace(/\(OrgId = 'org_\w+'\)/g, "$__orgFilter")
+		.replace(/'org_[A-Za-z0-9]+'/g, "'<your org>'")
+}