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 0000000000..b1191d2746 --- /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 986172775c..f4bf4b940a 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 0000000000..3375ba8aa6 --- /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 583bd78a92..575a6f40ca 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,19 @@ 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) { + // 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 }; + } + 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 0000000000..a02212feae --- /dev/null +++ b/internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts @@ -0,0 +1,90 @@ +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 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", () => { + 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); + }); + + 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"); + }); +});