diff --git a/apps/ai/src/mcp/__evals__/run-sql-recovery.test.ts b/apps/ai/src/mcp/__evals__/run-sql-recovery.test.ts index 2bd89349b3..08f6042f00 100644 --- a/apps/ai/src/mcp/__evals__/run-sql-recovery.test.ts +++ b/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 = [] 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") + }) }) diff --git a/apps/ai/src/mcp/lib/run-raw-sql.ts b/apps/ai/src/mcp/lib/run-raw-sql.ts index 2daa4b3913..91654f309d 100644 --- a/apps/ai/src/mcp/lib/run-raw-sql.ts +++ b/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( diff --git a/apps/ai/src/mcp/lib/trace-seek.test.ts b/apps/ai/src/mcp/lib/trace-seek.test.ts new file mode 100644 index 0000000000..c9ce37980e --- /dev/null +++ b/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() + }) +}) diff --git a/apps/ai/src/mcp/lib/trace-seek.ts b/apps/ai/src/mcp/lib/trace-seek.ts new file mode 100644 index 0000000000..1b9684dcf1 --- /dev/null +++ b/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 = ''` or `TraceId IN ('', ...)`, optionally qualified. */ +const TRACE_ID_PREDICATE = new RegExp( + `(?|<|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 +} diff --git a/apps/ai/src/mcp/tools/run-sql.ts b/apps/ai/src/mcp/tools/run-sql.ts index e0a1fe6ec6..8a6ac895af 100644 --- a/apps/ai/src/mcp/tools/run-sql.ts +++ b/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 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.`, diff --git a/packages/domain/src/mcp-outputs/services.ts b/packages/domain/src/mcp-outputs/services.ts index ce52f4dbbb..6bbb9a5dc6 100644 --- a/packages/domain/src/mcp-outputs/services.ts +++ b/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)), }) diff --git a/packages/query-engine/src/runtime/raw-sql.ts b/packages/query-engine/src/runtime/raw-sql.ts index 92e496d70d..8890150c0e 100644 --- a/packages/query-engine/src/runtime/raw-sql.ts +++ b/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 })