chore: merge feat/query-safety-tri-11165 (busy-rejection cap fix)

2026-08-12 08:20:01 +00:00
parent 378e00efcb 683bd87b4d
commit c7bbd1db88
8 changed files with 127 additions and 6 deletions
@@ -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.
@@ -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.
@@ -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: {} }]);
  });

@@ -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:
@@ -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;
    }
@@ -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",
@@ -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) {
@@ -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");
+  });
 });