From 0edbb2293a8c874a5190028a92b52762ecb69daa Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Mon, 10 Aug 2026 10:27:05 +0000 Subject: [PATCH 1/2] feat(webapp): pin the query boundary end-to-end and cap the query retry The read-only guard was enforced by the TRQL grammar and a parser test, but nothing proved the route itself refuses a write; a route test now drives api.v1.query with a signed environment JWT and asserts nothing reaches ClickHouse. readonly=1 is no longer overridable by a caller's clickhouseSettings. run_query gives up after three consecutive failures so a broken query can't burn a whole agent turn. TRI-11165 --- .../query-boundary-and-retry-cap.md | 6 + .../app/services/queryService.server.ts | 1 + apps/webapp/test/queryRouteReadOnly.test.ts | 168 ++++++++++++++++++ .../dashboard-agent/src/tool-api.ts | 20 ++- .../src/tool-query-retry-cap.test.ts | 73 ++++++++ 5 files changed, 267 insertions(+), 1 deletion(-) create mode 100644 .server-changes/query-boundary-and-retry-cap.md create mode 100644 apps/webapp/test/queryRouteReadOnly.test.ts create mode 100644 internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts diff --git a/.server-changes/query-boundary-and-retry-cap.md b/.server-changes/query-boundary-and-retry-cap.md new file mode 100644 index 00000000000..b1191d2746c --- /dev/null +++ b/.server-changes/query-boundary-and-retry-cap.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: improvement +--- + +Queries stay read-only, and the agent now stops after a few failed queries in a row and answers with what it found instead of spending the whole reply retrying. diff --git a/apps/webapp/app/services/queryService.server.ts b/apps/webapp/app/services/queryService.server.ts index 986172775cb..f4bf4b940a4 100644 --- a/apps/webapp/app/services/queryService.server.ts +++ b/apps/webapp/app/services/queryService.server.ts @@ -397,6 +397,7 @@ export async function executeQuery( ...getDefaultClickhouseSettings(), ...queryCacheSettings, ...baseOptions.clickhouseSettings, // Allow caller overrides if needed + readonly: "1", // Not overridable: every query through here is read-only. }, querySettings: { maxRows: env.QUERY_CLICKHOUSE_MAX_RETURNED_ROWS, diff --git a/apps/webapp/test/queryRouteReadOnly.test.ts b/apps/webapp/test/queryRouteReadOnly.test.ts new file mode 100644 index 00000000000..3375ba8aa66 --- /dev/null +++ b/apps/webapp/test/queryRouteReadOnly.test.ts @@ -0,0 +1,168 @@ +import { generateJWT } from "@trigger.dev/core/v3/jwt"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +/** + * The query API is read-only, and the grammar is what enforces it. A parser test alone would + * stay green if the route ever compiled agent SQL somewhere else, so these drive the real route + * with a real signed environment JWT and stub only the ClickHouse client. A write must be + * refused before anything reaches ClickHouse. + */ + +const ENVIRONMENT_ID = "env_1234"; +const API_KEY = "tr_dev_abcdefghijklmnop"; + +const environment = { + id: ENVIRONMENT_ID, + type: "DEVELOPMENT", + slug: "dev", + branchName: null, + apiKey: API_KEY, + organizationId: "org_1", + projectId: "proj_1", + archivedAt: null, + concurrencyLimitBurstFactor: { toNumber: () => 1 }, + maximumConcurrencyLimit: 10, + project: { id: "proj_1", externalRef: "proj_ref", deletedAt: null }, + organization: { id: "org_1" }, + orgMember: null, + parentEnvironment: null, +}; + +const mocks = vi.hoisted(() => ({ + runtimeEnvironmentFindFirst: vi.fn(), + queryWithStats: vi.fn(), + customerQueryCreate: vi.fn(), +})); + +vi.mock("~/db.server", () => { + const client = { + runtimeEnvironment: { + findFirst: mocks.runtimeEnvironmentFindFirst, + findMany: async () => [], + }, + revokedApiKey: { findMany: async () => [], findFirst: async () => null }, + project: { findMany: async () => [] }, + customerQuery: { findFirst: async () => null, create: mocks.customerQueryCreate }, + }; + return { prisma: client, $replica: client }; +}); +vi.mock("~/env.server", () => ({ + env: { + SESSION_SECRET: "test-session-secret", + QUERY_CLICKHOUSE_MAX_EXECUTION_TIME: "30", + QUERY_CLICKHOUSE_MAX_MEMORY_USAGE: 1000000, + QUERY_CLICKHOUSE_MAX_AST_ELEMENTS: 50000, + QUERY_CLICKHOUSE_MAX_EXPANDED_AST_ELEMENTS: 500000, + QUERY_CLICKHOUSE_MAX_BYTES_BEFORE_EXTERNAL_GROUP_BY: 1000000, + QUERY_CLICKHOUSE_MAX_RETURNED_ROWS: 1000, + }, +})); +vi.mock("~/services/clickhouse/clickhouseFactoryInstance.server", () => ({ + clickhouseFactory: { + getClickhouseForOrganization: async () => ({ + reader: { queryWithStats: mocks.queryWithStats }, + }), + }, +})); +vi.mock("~/services/platform.v3.server", () => ({ getLimit: async () => 30 })); +vi.mock("~/services/queryConcurrencyLimiter.server", () => ({ + queryConcurrencyLimiter: { + acquire: async () => ({ success: true }), + release: async () => {}, + }, + DEFAULT_ORG_CONCURRENCY_LIMIT: 10, + GLOBAL_CONCURRENCY_LIMIT: 100, +})); +vi.mock("~/services/logger.server", () => ({ + logger: { debug: vi.fn(), error: vi.fn(), warn: vi.fn(), info: vi.fn() }, +})); +vi.mock("~/v3/services/worker/workerGroupTokenService.server", () => ({ + WorkerGroupTokenService: class {}, +})); +vi.mock("~/v3/services/common.server", () => ({ ServiceValidationError: class extends Error {} })); +vi.mock("@internal/run-engine", () => ({ EngineServiceValidationError: class extends Error {} })); + +import { action } from "~/routes/api.v1.query"; +import { executeQuery } from "~/services/queryService.server"; + +/** The claims the env-JWT exchange mints (api.v1.projects.$projectRef.$env.jwt.ts). */ +function mintEnvJwt(scopes: string[]) { + return generateJWT({ + secretKey: API_KEY, + payload: { + sub: ENVIRONMENT_ID, + pub: true, + scopes, + act: { sub: "usr_1", client: "dashboard-agent" }, + }, + expirationTime: "1h", + }); +} + +async function runQuery(query: string): Promise<{ status: number; body: any }> { + const jwt = await mintEnvJwt(["read:query"]); + const response = await action({ + request: new Request("https://api.trigger.dev/api/v1/query", { + method: "POST", + headers: { Authorization: `Bearer ${jwt}`, "Content-Type": "application/json" }, + body: JSON.stringify({ query }), + }), + params: {}, + context: {}, + } as any); + return { status: response.status, body: await response.json() }; +} + +describe("the query API route", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.runtimeEnvironmentFindFirst.mockResolvedValue(environment); + mocks.customerQueryCreate.mockResolvedValue({ id: "cq_1" }); + mocks.queryWithStats.mockReturnValue(async () => [null, { rows: [], stats: {} }]); + }); + + // Pins the seam the two refusals assert against: a read really does reach ClickHouse here, + // so `not.toHaveBeenCalled()` below means refused, not unreachable. + it("runs a read against ClickHouse", async () => { + const result = await runQuery("SELECT count() FROM runs"); + + expect(result.status).toBe(200); + expect(mocks.queryWithStats).toHaveBeenCalled(); + }); + + it("refuses a write smuggled in as a second statement", async () => { + const result = await runQuery("SELECT 1 FROM runs; DROP TABLE runs"); + + expect(result.status).toBe(400); + expect(mocks.queryWithStats).not.toHaveBeenCalled(); + }); + + it("refuses a mutating statement", async () => { + const result = await runQuery("INSERT INTO runs (task_identifier) VALUES ('x')"); + + expect(result.status).toBe(400); + expect(mocks.queryWithStats).not.toHaveBeenCalled(); + }); +}); + +describe("the query service", () => { + beforeEach(() => { + vi.clearAllMocks(); + mocks.queryWithStats.mockReturnValue(async () => [null, { rows: [], stats: {} }]); + }); + + it("keeps ClickHouse read-only when a caller overrides the settings", async () => { + await executeQuery({ + name: "test-query", + query: "SELECT count() FROM runs", + scope: "environment", + organizationId: "org_1", + projectId: "proj_1", + environmentId: ENVIRONMENT_ID, + clickhouseSettings: { readonly: "0" }, + } as any); + + expect(mocks.queryWithStats).toHaveBeenCalled(); + expect(mocks.queryWithStats.mock.calls[0][0].settings.readonly).toBe("1"); + }); +}); diff --git a/internal-packages/dashboard-agent/src/tool-api.ts b/internal-packages/dashboard-agent/src/tool-api.ts index 583bd78a925..27e618180ef 100644 --- a/internal-packages/dashboard-agent/src/tool-api.ts +++ b/internal-packages/dashboard-agent/src/tool-api.ts @@ -177,6 +177,9 @@ export function withLiveState(metrics: unknown, queueType: "task" | "custom", li }; } +/** Failed `run_query` calls in a row before the tool tells the model to stop and answer. */ +export const MAX_CONSECUTIVE_QUERY_FAILURES = 3; + export function buildApiTools(args: { ctx: DashboardAgentToolContext; client: DashboardAgentApiClient; @@ -186,6 +189,12 @@ export function buildApiTools(args: { const { userActorToken, projectRef, environmentName, environmentBranch } = ctx; const { origin, hasAuth, envApiGet, postQuery, validateChartQuery } = client; + // A failed query hands the model the database error to fix, and it usually does. When it + // doesn't, the only other limit is the turn's 10 steps, so one broken query can eat the + // whole turn and leave the user with no answer at all. This tool set is built per turn, + // so the counter caps consecutive failures within one turn. + let consecutiveQueryFailures = 0; + return { list_projects: tool({ ...listProjectsSchema, @@ -341,7 +350,16 @@ export function buildApiTools(args: { execute: async ({ query, period }) => { const result = await postQuery(query, period); if (isEnvUnavailable(result)) return envUnavailableError(result, "query"); - if (!result.ok) return { error: result.error }; + if (!result.ok) { + consecutiveQueryFailures++; + if (consecutiveQueryFailures >= MAX_CONSECUTIVE_QUERY_FAILURES) { + return { + error: `${result.error} That is ${consecutiveQueryFailures} queries in a row that failed. Stop querying and answer the user with what you already have.`, + }; + } + return { error: result.error }; + } + consecutiveQueryFailures = 0; const cap = 200; const rows = result.rows; return { rows: rows.slice(0, cap), rowCount: rows.length, truncated: rows.length > cap }; diff --git a/internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts b/internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts new file mode 100644 index 00000000000..ee0a6db8049 --- /dev/null +++ b/internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts @@ -0,0 +1,73 @@ +import { describe, expect, it, vi } from "vitest"; +import { buildApiTools, MAX_CONSECUTIVE_QUERY_FAILURES } from "./tool-api"; +import type { DashboardAgentApiClient } from "./tool-api-client"; + +/** + * A failed query hands the model the database error to fix. Without a cap, the only other + * limit is the turn's step budget, so a model that keeps rewriting the same broken query + * burns the whole turn and the user gets no answer. After three failures in a row the tool + * tells it to stop and answer. + */ + +function queryTool(postQuery: DashboardAgentApiClient["postQuery"]) { + const client = { + origin: "https://api.example.com", + hasAuth: true, + envApiGet: async () => ({ ok: false as const, status: 500 }), + postQuery, + validateChartQuery: async () => null, + } as unknown as DashboardAgentApiClient; + const tools = buildApiTools({ + ctx: { userActorToken: "uat", apiOrigin: client.origin }, + client, + renderInvestigations: (() => []) as any, + }); + return (query: string) => (tools.run_query as any).execute({ query }, {} as any); +} + +const failure = { + ok: false as const, + kind: "query" as const, + error: "Unknown expression identifier 'createdAt'.", +}; +const success = { ok: true as const, rows: [{ n: 1 }] }; + +describe("run_query's consecutive-failure cap", () => { + it("keeps handing back the plain error until the cap", async () => { + const run = queryTool(async () => failure); + + for (let attempt = 1; attempt < MAX_CONSECUTIVE_QUERY_FAILURES; attempt++) { + const result = await run("SELECT createdAt FROM runs"); + expect(result.error).toBe(failure.error); + } + }); + + it("tells the model to stop and answer at the cap", async () => { + const run = queryTool(async () => failure); + + let result: { error: string } = { error: "" }; + for (let attempt = 0; attempt < MAX_CONSECUTIVE_QUERY_FAILURES; attempt++) { + result = await run("SELECT createdAt FROM runs"); + } + + expect(result.error).toContain(failure.error); + expect(result.error).toContain("answer the user with what you already have"); + }); + + it("counts consecutive failures only, so a good query clears the count", async () => { + const postQuery = vi + .fn() + .mockResolvedValueOnce(failure) + .mockResolvedValueOnce(failure) + .mockResolvedValueOnce(success) + .mockResolvedValue(failure); + const run = queryTool(postQuery as any); + + await run("bad"); + await run("bad"); + await run("good"); + const result = await run("bad"); + + expect(result.error).toBe(failure.error); + }); +}); From 614a7a9768c7fe82fceee748f1ada6b694780b19 Mon Sep 17 00:00:00 2001 From: Katia Bulatova Date: Mon, 10 Aug 2026 18:53:31 +0000 Subject: [PATCH 2/2] fix(dashboard-agent): only count SQL errors toward the query-failure cap --- .../dashboard-agent/src/tool-api.ts | 13 ++++++++----- .../src/tool-query-retry-cap.test.ts | 17 +++++++++++++++++ 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/internal-packages/dashboard-agent/src/tool-api.ts b/internal-packages/dashboard-agent/src/tool-api.ts index 27e618180ef..575a6f40ca8 100644 --- a/internal-packages/dashboard-agent/src/tool-api.ts +++ b/internal-packages/dashboard-agent/src/tool-api.ts @@ -351,11 +351,14 @@ export function buildApiTools(args: { const result = await postQuery(query, period); if (isEnvUnavailable(result)) return envUnavailableError(result, "query"); if (!result.ok) { - consecutiveQueryFailures++; - if (consecutiveQueryFailures >= MAX_CONSECUTIVE_QUERY_FAILURES) { - return { - error: `${result.error} That is ${consecutiveQueryFailures} queries in a row that failed. Stop querying and answer the user with what you already have.`, - }; + // Only SQL errors count toward the cap; transport errors are transient. + if (result.kind === "query") { + consecutiveQueryFailures++; + if (consecutiveQueryFailures >= MAX_CONSECUTIVE_QUERY_FAILURES) { + return { + error: `${result.error} That is ${consecutiveQueryFailures} queries in a row that failed. Stop querying and answer the user with what you already have.`, + }; + } } return { error: result.error }; } diff --git a/internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts b/internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts index ee0a6db8049..a02212feaeb 100644 --- a/internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts +++ b/internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts @@ -30,6 +30,11 @@ const failure = { kind: "query" as const, error: "Unknown expression identifier 'createdAt'.", }; +const transportFailure = { + ok: false as const, + kind: "transport" as const, + error: "The environment is temporarily unavailable.", +}; const success = { ok: true as const, rows: [{ n: 1 }] }; describe("run_query's consecutive-failure cap", () => { @@ -70,4 +75,16 @@ describe("run_query's consecutive-failure cap", () => { expect(result.error).toBe(failure.error); }); + + it("does not count transport errors toward the cap", async () => { + const run = queryTool(async () => transportFailure); + + let result: { error: string } = { error: "" }; + for (let attempt = 0; attempt < MAX_CONSECUTIVE_QUERY_FAILURES + 2; attempt++) { + result = await run("SELECT createdAt FROM runs"); + } + + expect(result.error).toBe(transportFailure.error); + expect(result.error).not.toContain("answer the user with what you already have"); + }); });