perf(ai): run_sql seeks trace-id lookups through trace_index and stops at the row cap

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

@@ -21,10 +21,29 @@ const GATEWAY_WRAPPED = {
 	type: "UNKNOWN_IDENTIFIER",
 }
 
+const FALLBACK_ID = "aaaaaaaaaaaaaaaabbbbbbbbbbbbbbbb"
+const issued: Array<string> = []
 let rt: EvalRuntime
 
 beforeAll(() => {
 	installFakeWarehouse([
+		// Records every statement and matches none of them.
+		{
+			match: (sql) => {
+				issued.push(sql)
+				return false
+			},
+			rows: [],
+		},
+		{
+			match: (sql) => sql.includes(FALLBACK_ID) && sql.includes("FROM trace_index"),
+			failure: unknown("ServiceName"),
+		},
+		{ match: (sql) => sql.includes(FALLBACK_ID), rows: [{ n: 1 }] },
+		{
+			match: (sql) => sql.includes("FROM product_events"),
+			rows: Array.from({ length: 1000 }, (_, i) => ({ n: i })),
+		},
 		{ 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") },
@@ -91,4 +110,50 @@ describe("run_sql recovery", () => {
 		expect(text).not.toContain("maple_raw_sql_limited")
 		expect(text).not.toContain("org_x")
 	})
+
+	it("stops at the row cap instead of failing on a longer result", async () => {
+		const result = await runToolDirect(rt, "run_sql", {
+			...WINDOW,
+			sql: "SELECT number AS n FROM product_events WHERE $__orgFilter",
+		})
+		expect(result.isError).not.toBe(true)
+		expect(issued.at(-1)).toMatch(/LIMIT 1000\b/)
+		const output = Schema.decodeUnknownSync(RunSqlOutput)(result.structuredContent)
+		expect(output.rowCount).toBe(1000)
+		expect(output.atRowCap).toBe(true)
+		expect(markdown(result)).toContain("Stopped at the 1000-row cap")
+	})
+
+	it("trusts a full page the query's own LIMIT asked for", async () => {
+		const result = await runToolDirect(rt, "run_sql", {
+			...WINDOW,
+			sql: "SELECT number AS n FROM product_events WHERE $__orgFilter LIMIT 1000",
+		})
+		const output = Schema.decodeUnknownSync(RunSqlOutput)(result.structuredContent)
+		expect(output.rowCount).toBe(1000)
+		expect(output.atRowCap).toBeUndefined()
+		expect(markdown(result)).not.toContain("row cap")
+	})
+
+	it("reads the count, not the offset, of LIMIT offset, count", async () => {
+		const result = await runToolDirect(rt, "run_sql", {
+			...WINDOW,
+			sql: "SELECT number AS n FROM product_events WHERE $__orgFilter LIMIT 10, 5000",
+		})
+		expect(Schema.decodeUnknownSync(RunSqlOutput)(result.structuredContent).atRowCap).toBe(true)
+	})
+
+	// The seek names `ServiceName`/`Timestamp`, which an outer query over a subquery may lack.
+	it("runs a trace-id lookup as written when the warehouse cannot analyse the seeked form", async () => {
+		const result = await runToolDirect(rt, "run_sql", {
+			...WINDOW,
+			sql: `SELECT count() AS n FROM (SELECT TraceId FROM traces WHERE $__orgFilter) WHERE TraceId = '${FALLBACK_ID}'`,
+		})
+		expect(result.isError).not.toBe(true)
+		expect(Schema.decodeUnknownSync(RunSqlOutput)(result.structuredContent).rows).toEqual([{ n: 1 }])
+		const lookups = issued.filter((sql) => sql.includes(FALLBACK_ID))
+		expect(lookups).toHaveLength(2)
+		expect(lookups[0]).toContain("FROM trace_index")
+		expect(lookups[1]).not.toContain("FROM trace_index")
+	})
 })

apps/ai/src/mcp/lib/run-raw-sql.ts

@@ -27,6 +27,8 @@ export interface RunRawSqlInput {
 	readonly startTime: string
 	readonly endTime: string
 	readonly granularitySeconds: number
+	/** Stop at the 1,000-row cap instead of failing on a longer result. */
+	readonly truncateAtCap?: boolean
 }
 
 /**
@@ -58,6 +60,7 @@ export const runRawSql = Effect.fn("McpRawSql.runRawSql")(function* (input: RunR
 		startTime: input.startTime,
 		endTime: input.endTime,
 		granularitySeconds: input.granularitySeconds,
+		...(input.truncateAtCap === true ? { truncateAtCap: true } : undefined),
 		workload: "interactive",
 		context: "mcp.run_sql",
 	}).pipe(

apps/ai/src/mcp/lib/trace-seek.test.ts

@@ -0,0 +1,64 @@
+import { describe, expect, it } from "vitest"
+import { MAX_RAW_SQL_LENGTH } from "@maple/domain/raw-sql"
+import { withTraceIndexSeek } from "./trace-seek"
+
+const ID = "2ea98b74a6330960b2828c8e2edfee38"
+const SEEK = `(ServiceName, toStartOfFiveMinutes(Timestamp), TraceId) IN (SELECT ServiceName, Bucket, TraceId FROM trace_index WHERE $__orgFilter AND TraceId IN ('${ID}'))`
+
+describe("withTraceIndexSeek", () => {
+	it("ANDs the trace_index seek onto an unbounded trace-id predicate", () => {
+		expect(
+			withTraceIndexSeek(`SELECT SpanName FROM traces WHERE $__orgFilter AND TraceId = '${ID}'`),
+		).toBe(`SELECT SpanName FROM traces WHERE $__orgFilter AND (TraceId = '${ID}' AND ${SEEK})`)
+	})
+
+	it("keeps the qualifier and every id of an IN list", () => {
+		const other = "7ac1ca09957fd49ac57461016a5a7819"
+		const seeked = withTraceIndexSeek(
+			`SELECT t.SpanName FROM traces t WHERE $__orgFilter(t) AND t.TraceId IN ('${ID}', '${other}')`,
+		)
+		expect(seeked).toContain(
+			`(t.ServiceName, toStartOfFiveMinutes(t.Timestamp), t.TraceId) IN (SELECT ServiceName, Bucket, TraceId FROM trace_index WHERE $__orgFilter AND TraceId IN ('${ID}', '${other}'))`,
+		)
+	})
+
+	it.each([
+		[
+			"a time bound",
+			`SELECT 1 FROM traces WHERE $__orgFilter AND $__timeFilter(Timestamp) AND TraceId = '${ID}'`,
+		],
+		[
+			"a literal bound",
+			`SELECT 1 FROM traces WHERE $__orgFilter AND Timestamp >= '2026-10-01' AND TraceId = '${ID}'`,
+		],
+		["another table", `SELECT 1 FROM logs WHERE $__orgFilter AND TraceId = '${ID}'`],
+		[
+			"a join",
+			`SELECT 1 FROM traces t JOIN logs l USING TraceId WHERE $__orgFilter(t) AND t.TraceId = '${ID}'`,
+		],
+		[
+			"the id only in a string",
+			`SELECT 1 FROM traces WHERE $__orgFilter AND Body = 'TraceId = ''${ID}'''`,
+		],
+		["no org filter of its own", `SELECT 1 FROM traces WHERE TraceId = '${ID}'`],
+		[
+			"an alias shadowing Timestamp",
+			`SELECT toStartOfHour(Timestamp) AS Timestamp FROM traces WHERE $__orgFilter AND TraceId = '${ID}'`,
+		],
+		["the id inside if()", `SELECT countIf(TraceId = '${ID}') FROM traces WHERE $__orgFilter`],
+		["a negated id", `SELECT 1 FROM traces WHERE $__orgFilter AND NOT TraceId = '${ID}'`],
+		["an OR-ed id", `SELECT 1 FROM traces WHERE $__orgFilter AND (SpanName = 'x' OR TraceId = '${ID}')`],
+		[
+			"a subquery id",
+			"SELECT 1 FROM traces WHERE $__orgFilter AND TraceId IN (SELECT TraceId FROM trace_list_mv)",
+		],
+	])("leaves a query with %s alone", (_, sql) => {
+		expect(withTraceIndexSeek(sql)).toBeUndefined()
+	})
+
+	it("leaves a statement alone when the seek would push it past the length limit", () => {
+		const base = `SELECT SpanName FROM traces WHERE $__orgFilter AND TraceId = '${ID}' AND SpanName != ''`
+		const sql = `${base}${" ".repeat(MAX_RAW_SQL_LENGTH - base.length - 10)}`
+		expect(withTraceIndexSeek(sql)).toBeUndefined()
+	})
+})

apps/ai/src/mcp/lib/trace-seek.ts

@@ -0,0 +1,60 @@
+import { TRACE_BUCKET_FN } from "@maple/domain/tinybird/datasources"
+import { MAX_RAW_SQL_LENGTH, maskLiteralsAndComments, rawSqlIssue } from "@maple/domain/raw-sql"
+import { definedNames, referencedTableNames } from "./sql-identifiers"
+
+const ID = "'[0-9a-fA-F]{16,32}'"
+
+/** `TraceId = '<id>'` or `TraceId IN ('<id>', ...)`, optionally qualified. */
+const TRACE_ID_PREDICATE = new RegExp(
+	`(?<![\\w.])(?:(\\w+)\\.)?TraceId\\s*(?:=\\s*${ID}|IN\\s*\\(\\s*${ID}(?:\\s*,\\s*${ID})*\\s*\\))`,
+	"g",
+)
+
+/** A time bound already lets ClickHouse prune partitions, so the query is left as written. */
+const TIME_BOUND = /\$__timeFilter\(|\$__startTime|\$__endTime|\bTimestamp\s*(?:>|<|BETWEEN\b)/i
+
+// Only a plain conjunct of the WHERE clause: not inside `if()`, `NOT`, or an OR.
+const BEFORE_CONJUNCT = /\b(?:WHERE|PREWHERE|AND)\s*$/i
+const AFTER_CONJUNCT = /^\s*(?:$|\)|AND\b|GROUP\b|ORDER\b|LIMIT\b|HAVING\b|SETTINGS\b|FORMAT\b|UNION\b)/i
+
+/**
+ * A `traces` read by literal trace id with no time bound, seeked through `trace_index`
+ * like `inspect_trace` does. A bare `TraceId = ...` probes every daily partition and ran
+ * past the time limit for large orgs. The seek is ANDed onto the trace-id conjunct, so it
+ * can only select rows the predicate already selects; the index holds every
+ * (trace, service, bucket) its spans occupy.
+ */
+export const withTraceIndexSeek = (sql: string): string | undefined => {
+	const masked = maskLiteralsAndComments(sql)
+	const tables = referencedTableNames(sql)
+	const defined = definedNames(sql)
+	if (
+		tables.length !== 1 ||
+		tables[0] !== "traces" ||
+		TIME_BOUND.test(masked) ||
+		// The seek brings its own `$__orgFilter`; the query must pass the gate without it.
+		rawSqlIssue(sql) !== null ||
+		["ServiceName", "Timestamp", "TraceId"].some((name) => defined.has(name))
+	)
+		return undefined
+	// The literal is blanked in `masked`, so match the original and keep conjuncts outside literals.
+	const hits = [...sql.matchAll(TRACE_ID_PREDICATE)].filter((m) => {
+		const name = m[0].slice(0, m[0].indexOf("TraceId") + "TraceId".length)
+		return (
+			masked.startsWith(name, m.index) &&
+			BEFORE_CONJUNCT.test(masked.slice(0, m.index)) &&
+			AFTER_CONJUNCT.test(masked.slice(m.index + m[0].length))
+		)
+	})
+	if (hits.length === 0) return undefined
+	const seeked = hits.reduceRight((text, m) => {
+		const q = m[1] === undefined ? "" : `${m[1]}.`
+		const ids = m[0].match(new RegExp(ID, "g"))?.join(", ") ?? ""
+		const seek =
+			`(${q}ServiceName, ${TRACE_BUCKET_FN}(${q}Timestamp), ${q}TraceId) IN ` +
+			`(SELECT ServiceName, Bucket, TraceId FROM trace_index WHERE $__orgFilter AND TraceId IN (${ids}))`
+		return `${text.slice(0, m.index)}(${m[0]} AND ${seek})${text.slice(m.index + m[0].length)}`
+	}, sql)
+	// The seek adds text; a statement it pushes past the length limit runs as written.
+	return seeked.length > MAX_RAW_SQL_LENGTH ? undefined : seeked
+}

apps/ai/src/mcp/tools/run-sql.ts

@@ -6,8 +6,9 @@ import { autoBucketSeconds, runRawSql } from "../lib/run-raw-sql"
 import { truncate } from "../lib/format"
 import { toMcpQueryError } from "../lib/map-warehouse-error"
 import * as P from "../lib/params"
-import { MAX_RAW_SQL_RESULT_ROWS } from "@maple/domain/raw-sql"
+import { MAX_RAW_SQL_RESULT_ROWS, maskLiteralsAndComments } from "@maple/domain/raw-sql"
 import { rawSqlErrorText } from "../lib/raw-sql-error-text"
+import { withTraceIndexSeek } from "../lib/trace-seek"
 import { doc } from "../lib/tool-doc"
 import {
 	listWarehouseTables,
@@ -181,7 +182,7 @@ const runSqlSchema = Schema.Struct({
 
 const runSqlDescription =
 	"Run one read-only ClickHouse SELECT against the org's warehouse and return its rows: " +
-	`at most ${MAX_RAW_SQL_RESULT_ROWS} are fetched and the first ${MAX_RENDERED_ROWS} rendered, with the true count reported. ` +
+	`at most ${MAX_RAW_SQL_RESULT_ROWS} are fetched and the first ${MAX_RENDERED_ROWS} rendered, with the count reported. ` +
 	"Use it to answer what query_data cannot express, to spot-check data, or to test SQL before saving it " +
 	"as a raw_sql widget with add_dashboard_widget. For trends and top-N prefer query_data. " +
 	"Spans are in `traces` (there is no otel_spans, otel_traces or spans table) and log lines in `logs`. " +
@@ -291,6 +292,16 @@ const withCallerErrorText =
 				})
 			: error
 
+// The cap leaves no overflow row, so only a statement's own LIMIT within the cap proves it whole.
+const OWN_LIMIT = /\bLIMIT\s+(\d+)(?:\s*,\s*(\d+)|\s+OFFSET\s+\d+)?\s*;?\s*$/i
+
+/** The statement ends in its own LIMIT of at most the cap, so a full page is its answer, not ours. */
+const limitedWithinCap = (sql: string): boolean => {
+	const own = OWN_LIMIT.exec(maskLiteralsAndComments(sql).trimEnd())
+	// `LIMIT offset, count` puts the count second.
+	return own !== null && Number(own[2] ?? own[1]) <= MAX_RAW_SQL_RESULT_ROWS
+}
+
 type RawSqlRun = ReturnType<typeof runRawSql>
 
 interface FixFailure {
@@ -315,7 +326,14 @@ export function registerRunSqlTool(server: McpToolRegistrar) {
 			const granularitySeconds = params.granularity_seconds ?? autoBucketSeconds(st, et)
 
 			const execute = (sql: string) =>
-				runRawSql({ tenant, sql, startTime: st, endTime: et, granularitySeconds })
+				runRawSql({
+					tenant,
+					sql,
+					startTime: st,
+					endTime: et,
+					granularitySeconds,
+					truncateAtCap: true,
+				})
 
 			// A misspelt column with exactly one right answer is fixed and rerun rather than
 			// costing the agent a round trip; each pass fixes the name the warehouse reported.
@@ -341,7 +359,22 @@ export function registerRunSqlTool(server: McpToolRegistrar) {
 					}),
 				)
 
-			const { rows, columns, rowCount, expandedSql, rewrites } = yield* runFixing(params.sql, []).pipe(
+			// An unbounded trace-id lookup reads through `trace_index`. If the warehouse cannot
+			// analyse that form (the predicate sits where `ServiceName` is not in scope), the
+			// query runs as written; any other failure is the seek's to report.
+			const seeked = withTraceIndexSeek(params.sql)
+			const run =
+				seeked === undefined
+					? runFixing(params.sql, [])
+					: runFixing(seeked, []).pipe(
+							Effect.catch((failure) =>
+								failure.error._tag === "@maple/http/errors/WarehouseInvalidSqlError"
+									? runFixing(params.sql, [])
+									: Effect.fail(failure),
+							),
+						)
+
+			const { rows, columns, rowCount, expandedSql, rewrites } = yield* run.pipe(
 				Effect.mapError(({ error, sql, rewrites }) =>
 					// Macro/safety failures (a missing `$__orgFilter` above all) are caller-fixable.
 					error._tag === "@maple/http/errors/RawSqlValidationError"
@@ -385,12 +418,20 @@ export function registerRunSqlTool(server: McpToolRegistrar) {
 				timeRange: { start: st, end: et },
 				maxCellChars,
 				windowApplied: usesWindow(expandedSql, st, et),
+				...(rowCount >= MAX_RAW_SQL_RESULT_ROWS && !limitedWithinCap(params.sql)
+					? { atRowCap: true }
+					: undefined),
 				...(rewrites.length === 0 ? undefined : { rewrites: [...rewrites] }),
 				...(params.show_sql === true ? { expandedSql } : undefined),
 			}
 		}),
 		render: (output) => {
 			const notices = [
+				...(output.atRowCap === true
+					? [
+							`Stopped at the ${MAX_RAW_SQL_RESULT_ROWS}-row cap, so more rows may exist; aggregate or add a LIMIT to see all of them.`,
+						]
+					: []),
 				...(output.rewrites ?? []).map(
 					(r) =>
 						`Ran with \`${r.to}\` for \`${r.from}\`, which is not a column here; use \`${r.to}\` directly.`,

packages/domain/src/mcp-outputs/services.ts

@@ -291,6 +291,8 @@ export const RunSqlOutput = Schema.Struct({
 	windowApplied: Schema.optionalKey(Schema.Boolean),
 	/** The fully macro-expanded SQL that was executed, only when asked for (`show_sql`). */
 	expandedSql: Schema.optionalKey(Schema.String),
+	/** The result stopped at the 1,000-row cap, so more rows may exist. */
+	atRowCap: Schema.optionalKey(Schema.Boolean),
 	/** Column names the warehouse rejected and the server replaced before rerunning. */
 	rewrites: Schema.optionalKey(Schema.Array(RunSqlRewrite)),
 })

packages/query-engine/src/runtime/raw-sql.ts

@@ -33,6 +33,8 @@ export interface PrepareRawSqlInput {
 	readonly workload: RawSqlWorkload
 	/** `$__timeGroup` keys every row at the range start: one bucket per group. Alert evaluation sets it. */
 	readonly wholeRange?: boolean
+	/** Return the first 1,000 rows of a longer result instead of failing on the cap. */
+	readonly truncateAtCap?: boolean
 }
 
 export interface PreparedRawSql {
@@ -153,7 +155,7 @@ export const prepareRawSql = Effect.fn("RawSql.prepare")(function* (input: Prepa
 	// modifiers attached to the inner result (`WITH TOTALS`, `LIMIT BY`,
 	// `WITH FILL`) are discarded by the outer SELECT.
 	return {
-		sql: `SELECT * FROM (\n${splitTerminalClauses(sql).body.trim()}\n) AS maple_raw_sql_limited\nLIMIT ${RAW_SQL_FETCH_ROW_LIMIT}`,
+		sql: `SELECT * FROM (\n${splitTerminalClauses(sql).body.trim()}\n) AS maple_raw_sql_limited\nLIMIT ${input.truncateAtCap === true ? MAX_RAW_SQL_RESULT_ROWS : RAW_SQL_FETCH_ROW_LIMIT}`,
 		granularitySeconds: granularity,
 	} satisfies PreparedRawSql
 })