feat(webapp): query boundary pinned end-to-end and a capped query retry

.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.

apps/webapp/app/services/queryService.server.ts

@@ -397,6 +397,7 @@ export async function executeQuery<TOut extends z.ZodSchema>(
         ...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,

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");
+  });
+});

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 };

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");
+  });
+});