fix(dashboard-agent): don't count busy-rejections toward the query failure cap

Author: kathiekiwi

.server-changes/query-busy-retry.md

@@ -0,0 +1,6 @@
+---
+area: webapp
+type: fix
+---
+
+When queries are queued up and one is turned away, you now get a clear "try again shortly" instead of an error that looks like a problem with the query itself.

apps/webapp/app/routes/api.v1.query.ts

@@ -2,7 +2,11 @@ import { json } from "@remix-run/server-runtime";
 import { QueryError } from "@internal/clickhouse";
 import { z } from "zod";
 import { createActionApiRoute, everyResource } from "~/services/routeBuilders/apiBuilder.server";
-import { executeQuery, type QueryScope } from "~/services/queryService.server";
+import {
+  executeQuery,
+  isQueryConcurrencyRejection,
+  type QueryScope,
+} from "~/services/queryService.server";
 import { logger } from "~/services/logger.server";
 import { rowsToCSV } from "~/utils/dataExport";
 import { detectQueryTables } from "~/v3/detectQueryTables";
@@ -78,6 +82,12 @@ const { action, loader } = createActionApiRoute(
     });
 
     if (!queryResult.success) {
+      // A concurrency rejection is "too busy", not a bad query: 429 so callers retry it
+      // instead of rewriting a query that was fine.
+      if (isQueryConcurrencyRejection(queryResult.error)) {
+        return json({ error: queryResult.error.message }, { status: 429 });
+      }
+
       // QueryError surfaces customer SQL problems (invalid syntax,
       // unsupported construct). Returned to the caller as 400; system
       // handles it gracefully, no alert needed.

apps/webapp/test/queryRouteReadOnly.test.ts

@@ -32,6 +32,7 @@ const mocks = vi.hoisted(() => ({
   runtimeEnvironmentFindFirst: vi.fn(),
   queryWithStats: vi.fn(),
   customerQueryCreate: vi.fn(),
+  concurrencyAcquire: vi.fn(),
 }));
 
 vi.mock("~/db.server", () => {
@@ -67,7 +68,7 @@ vi.mock("~/services/clickhouse/clickhouseFactoryInstance.server", () => ({
 vi.mock("~/services/platform.v3.server", () => ({ getLimit: async () => 30 }));
 vi.mock("~/services/queryConcurrencyLimiter.server", () => ({
   queryConcurrencyLimiter: {
-    acquire: async () => ({ success: true }),
+    acquire: mocks.concurrencyAcquire,
     release: async () => {},
   },
   DEFAULT_ORG_CONCURRENCY_LIMIT: 10,
@@ -118,6 +119,7 @@ describe("the query API route", () => {
     vi.clearAllMocks();
     mocks.runtimeEnvironmentFindFirst.mockResolvedValue(environment);
     mocks.customerQueryCreate.mockResolvedValue({ id: "cq_1" });
+    mocks.concurrencyAcquire.mockResolvedValue({ success: true });
     mocks.queryWithStats.mockReturnValue(async () => [null, { rows: [], stats: {} }]);
   });
 
@@ -143,11 +145,22 @@ describe("the query API route", () => {
     expect(result.status).toBe(400);
     expect(mocks.queryWithStats).not.toHaveBeenCalled();
   });
+
+  // A busy service is not a bad query: 400 would tell a caller to rewrite a query that was fine.
+  it("answers a concurrency rejection with 429", async () => {
+    mocks.concurrencyAcquire.mockResolvedValue({ success: false, reason: "key_limit" });
+
+    const result = await runQuery("SELECT count() FROM runs");
+
+    expect(result.status).toBe(429);
+    expect(result.body.error).toContain("try again later");
+  });
 });
 
 describe("the query service", () => {
   beforeEach(() => {
     vi.clearAllMocks();
+    mocks.concurrencyAcquire.mockResolvedValue({ success: true });
     mocks.queryWithStats.mockReturnValue(async () => [null, { rows: [], stats: {} }]);
   });
 

docs/v3-openapi.yaml

@@ -1203,6 +1203,16 @@ paths:
                     description: Error message describing the query error
         "401":
           description: Unauthorized - API key is missing or invalid
+        "429":
+          description: Query service is busy or rate limited - retry shortly
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  error:
+                    type: string
+                    description: Error message describing why the query was turned away
         "500":
           description: Internal server error during query execution
       tags:

internal-packages/dashboard-agent/src/tool-api-client.ts

@@ -42,11 +42,12 @@ const GET_TIMEOUT_MS = 10_000;
 const JWT_TIMEOUT_MS = 10_000;
 const QUERY_TIMEOUT_MS = 30_000;
 
-// "query" is the server rejecting the TRQL, "transport" is the request breaking. Chart
+// "query" is the server rejecting the TRQL, "transport" is the request breaking, "busy" is
+// the server too loaded or rate limited to answer — the same query may work shortly. Chart
 // validation only fails a render on "query".
 export type QueryPostResult =
   | { ok: true; rows: Array<Record<string, unknown>> }
-  | { ok: false; kind: "query" | "transport"; error: string };
+  | { ok: false; kind: "query" | "transport" | "busy"; error: string };
 
 export const NO_AUTH = { error: "No delegated access is available for this turn." } as const;
 
@@ -215,6 +216,15 @@ export function createApiClient(ctx: ApiClientContext): DashboardAgentApiClient
     // The route returns 400 with { error } for invalid TRQL.
     const data = (await res.json().catch(() => ({}))) as { results?: unknown; error?: string };
     if (!res.ok) {
+      // 429 is the concurrency rejection and the rate limiter: nothing is wrong with the
+      // query, so it is not a query error.
+      if (res.status === 429) {
+        return {
+          ok: false,
+          kind: "busy",
+          error: `${data.error ?? "The query service is busy right now."} You can retry the same query shortly.`,
+        };
+      }
       return {
         ok: false,
         kind: res.status >= 500 ? "transport" : "query",
@@ -234,7 +244,7 @@ export function createApiClient(ctx: ApiClientContext): DashboardAgentApiClient
   ): Promise<string | null> {
     const result = await postQuery(query, period);
     if (isEnvUnavailable(result) || result.ok) return null;
-    if (result.kind === "transport") {
+    if (result.kind === "transport" || result.kind === "busy") {
       logger.warn("Skipped chart query validation", { error: result.error });
       return null;
     }

internal-packages/dashboard-agent/src/tool-api-transport.test.ts

@@ -88,6 +88,45 @@ describe("a broken request reads as a broken request, never as an answer", () =>
     expect(JSON.stringify(result)).not.toContain("isn't locked to a deployed version");
   });
 
+  it("classifies a busy query route as busy, not as a bad query", async () => {
+    vi.stubGlobal(
+      "fetch",
+      vi.fn(async (url: string) =>
+        url.endsWith("/jwt")
+          ? Response.json({ token: "jwt" })
+          : Response.json(
+              { error: "We're experiencing a lot of queries at the moment." },
+              { status: 429 }
+            )
+      )
+    );
+
+    const result = await createApiClient(CTX).postQuery("SELECT 1", undefined);
+
+    expect(result).toMatchObject({ ok: false, kind: "busy" });
+    expect((result as { error: string }).error).toContain("retry the same query shortly");
+  });
+
+  // The other half of the same invariant: only 429 is busy, so a rejected query still counts.
+  it("classifies a rejected query as a query error", async () => {
+    vi.stubGlobal(
+      "fetch",
+      vi.fn(async (url: string) =>
+        url.endsWith("/jwt")
+          ? Response.json({ token: "jwt" })
+          : Response.json({ error: "Unknown expression identifier 'createdAt'." }, { status: 400 })
+      )
+    );
+
+    const result = await createApiClient(CTX).postQuery("SELECT createdAt FROM runs", undefined);
+
+    expect(result).toMatchObject({
+      ok: false,
+      kind: "query",
+      error: "Unknown expression identifier 'createdAt'.",
+    });
+  });
+
   it("still reports a real 404 as the answer it is", async () => {
     vi.stubGlobal(
       "fetch",

internal-packages/dashboard-agent/src/tool-api.ts

@@ -360,7 +360,8 @@ export function buildApiTools(args: {
         const result = await postQuery(query, period);
         if (isEnvUnavailable(result)) return envUnavailableError(result, "query");
         if (!result.ok) {
-          // Only SQL errors count toward the cap; transport errors are transient.
+          // Only SQL errors count toward the cap; transport and busy errors are transient,
+          // and the same query may work on a retry.
           if (result.kind === "query") {
             consecutiveQueryFailures++;
             if (consecutiveQueryFailures >= MAX_CONSECUTIVE_QUERY_FAILURES) {

internal-packages/dashboard-agent/src/tool-query-retry-cap.test.ts

@@ -35,6 +35,11 @@ const transportFailure = {
   kind: "transport" as const,
   error: "The environment is temporarily unavailable.",
 };
+const busyFailure = {
+  ok: false as const,
+  kind: "busy" as const,
+  error: "We're experiencing a lot of queries at the moment. You can retry the same query shortly.",
+};
 const success = { ok: true as const, rows: [{ n: 1 }] };
 
 describe("run_query's consecutive-failure cap", () => {
@@ -87,4 +92,31 @@ describe("run_query's consecutive-failure cap", () => {
     expect(result.error).toBe(transportFailure.error);
     expect(result.error).not.toContain("answer the user with what you already have");
   });
+
+  // A "too busy" rejection says nothing about the query, so spending the cap on it would
+  // stop the model over a queue that clears in seconds.
+  it("does not count busy rejections toward the cap", async () => {
+    const run = queryTool(async () => busyFailure);
+
+    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(busyFailure.error);
+    expect(result.error).not.toContain("answer the user with what you already have");
+  });
+
+  it("still caps real SQL errors that follow busy rejections", async () => {
+    const postQuery = vi.fn().mockResolvedValueOnce(busyFailure).mockResolvedValue(failure);
+    const run = queryTool(postQuery as any);
+
+    await run("busy");
+    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("answer the user with what you already have");
+  });
 });