fix(webapp): stop the billing limits page timing out under enforcement

.server-changes/billing-limits-page-timeout.md

@@ -0,0 +1,6 @@
+---
+area: webapp
+type: fix
+---
+
+Fixed the billing limits page timing out for organizations with many preview branches, especially while a spend limit was being enforced. The page now loads quickly, so you can raise or resolve your limit without delay.

apps/webapp/app/v3/services/billingLimit/billingLimitConstants.ts

@@ -16,6 +16,9 @@ export const BILLING_LIMIT_RECONCILE_LOOKUP_CONCURRENCY = 10;
 /** Inline bulk-cancel budget for billing limit resolve (worker visibility is 10 min). */
 export const BILLING_LIMIT_RESOLVE_BULK_CANCEL_BUDGET_MS = 8 * 60_000;
 
+/** Server-side cap (seconds) on the org-level queued-run count; it renders in a loader. */
+export const BILLING_LIMIT_QUEUED_COUNT_MAX_EXECUTION_S = 10;
+
 export type BillingLimitConvergeTargetState = "grace" | "rejected" | "ok";
 
 export function isBillableEnvironmentType(type: RuntimeEnvironmentType): boolean {

apps/webapp/app/v3/services/billingLimit/billingLimitQueuedRuns.server.ts

@@ -1,16 +1,25 @@
+import type { ClickHouse } from "@internal/clickhouse";
 import type { PrismaClient, TaskRunStatus } from "@trigger.dev/database";
 import { QUEUED_STATUSES, RUNNING_STATUSES } from "~/components/runs/v3/TaskRunStatus";
 import { prisma } from "~/db.server";
 import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactoryInstance.server";
 import { RunsRepository } from "~/services/runsRepository/runsRepository.server";
-import { BILLABLE_ENVIRONMENT_TYPES } from "./billingLimitConstants";
+import {
+  BILLABLE_ENVIRONMENT_TYPES,
+  BILLING_LIMIT_QUEUED_COUNT_MAX_EXECUTION_S,
+} from "./billingLimitConstants";
 
 import { boundedIn } from "@trigger.dev/database";
 export type BillableEnvironmentRef = {
   id: string;
   projectId: string;
 };
 
+/**
+ * Environments whose runs a billing limit must act on. Archived branches stay included:
+ * archiving is a soft update that cancels nothing, so an archived branch can still hold
+ * executing runs that enforcement has to cancel.
+ */
 export async function getBillableEnvironmentsForBillingLimit(
   organizationId: string,
   prismaClient: PrismaClient = prisma
@@ -73,27 +82,37 @@ async function countRunsForBillableEnvironment(
   });
 }
 
-/** Same source as BillingLimitBulkCancelService — ClickHouse countRuns(QUEUED_STATUSES). */
+/**
+ * Same table and statuses as BillingLimitBulkCancelService's per-environment counts, but a
+ * single org-level ClickHouse query filtered on environment_type. Deliberately NOT a
+ * per-environment loop: an org can have thousands of (mostly archived) preview environments,
+ * and sequential per-env counts hold the billing-limits loader open past the edge timeout.
+ * The count is display-only, so a server-side execution cap beats an unbounded query.
+ */
 export async function countBillableQueuedRunsForOrganization(
-  organizationId: string
+  organizationId: string,
+  clickhouse?: ClickHouse
 ): Promise<number> {
-  const environments = await getBillableEnvironmentsForBillingLimit(organizationId);
+  const client =
+    clickhouse ??
+    (await clickhouseFactory.getClickhouseForOrganization(organizationId, "standard"));
 
-  if (environments.length === 0) {
-    return 0;
-  }
+  const queryBuilder = client.taskRuns.countQueryBuilder({
+    settings: { max_execution_time: BILLING_LIMIT_QUEUED_COUNT_MAX_EXECUTION_S },
+  });
 
-  const runsRepository = await createBillingLimitRunsRepository(organizationId);
+  queryBuilder
+    .where("organization_id = {organizationId: String}", { organizationId })
+    .where("environment_type IN {environmentTypes: Array(String)}", {
+      environmentTypes: [...BILLABLE_ENVIRONMENT_TYPES],
+    })
+    .where("status IN {statuses: Array(String)}", { statuses: [...QUEUED_STATUSES] });
 
-  let total = 0;
+  const [queryError, result] = await queryBuilder.execute();
 
-  for (const environment of environments) {
-    total += await countQueuedRunsForBillableEnvironment(
-      runsRepository,
-      organizationId,
-      environment
-    );
+  if (queryError) {
+    throw queryError;
   }
 
-  return total;
+  return result[0]?.count ?? 0;
 }

apps/webapp/app/v3/services/billingLimit/getBillingLimitQueuedRunCount.server.ts

@@ -1,9 +1,26 @@
+import { tryCatch } from "@trigger.dev/core/utils";
 import { EnvironmentPauseSource } from "@trigger.dev/database";
 import { prisma } from "~/db.server";
+import { logger } from "~/services/logger.server";
 import { countBillableQueuedRunsForOrganization } from "./billingLimitQueuedRuns.server";
 
+/**
+ * Display-only count for the billing-limits page. Falls back to 0 on failure (the page hides
+ * the count label at 0) — the recovery panel must stay reachable even when the count errors,
+ * because it is the customer's only self-serve path out of an enforced limit.
+ */
 export async function getBillingLimitQueuedRunCount(organizationId: string): Promise<number> {
-  return countBillableQueuedRunsForOrganization(organizationId);
+  const [error, count] = await tryCatch(countBillableQueuedRunsForOrganization(organizationId));
+
+  if (error) {
+    logger.error("getBillingLimitQueuedRunCount failed, returning 0", {
+      organizationId,
+      error,
+    });
+    return 0;
+  }
+
+  return count ?? 0;
 }
 
 export async function countBillingLimitPausedEnvironments(organizationId: string): Promise<number> {

apps/webapp/test/billingLimitQueuedRuns.test.ts

@@ -2,7 +2,11 @@ import { describe, expect, vi } from "vitest";
 import { setTimeout } from "node:timers/promises";
 import { replicationContainerTest } from "@internal/testcontainers";
 import { RunsRepository } from "~/services/runsRepository/runsRepository.server";
-import { countQueuedRunsForBillableEnvironment } from "~/v3/services/billingLimit/billingLimitQueuedRuns.server";
+import {
+  countBillableQueuedRunsForOrganization,
+  countQueuedRunsForBillableEnvironment,
+  getBillableEnvironmentsForBillingLimit,
+} from "~/v3/services/billingLimit/billingLimitQueuedRuns.server";
 import { setupClickhouseReplication } from "./utils/replicationUtils";
 
 vi.setConfig({ testTimeout: 60_000 });
@@ -109,4 +113,188 @@ describe("billingLimitQueuedRuns", () => {
       expect(developmentCount).toBe(1);
     }
   );
+
+  replicationContainerTest(
+    "counts queued runs org-wide with a single query, spanning billable environment types only",
+    async ({ clickhouseContainer, redisOptions, postgresContainer, prisma }) => {
+      const { clickhouse } = await setupClickhouseReplication({
+        prisma,
+        databaseUrl: postgresContainer.getConnectionUri(),
+        clickhouseUrl: clickhouseContainer.getConnectionUrl(),
+        redisOptions,
+      });
+
+      const organization = await prisma.organization.create({
+        data: { title: "billing-limit-org-count", slug: "billing-limit-org-count" },
+      });
+
+      const project = await prisma.project.create({
+        data: {
+          name: "billing-limit-org-count",
+          slug: "billing-limit-org-count",
+          organizationId: organization.id,
+          externalRef: "billing-limit-org-count",
+        },
+      });
+
+      const productionEnv = await prisma.runtimeEnvironment.create({
+        data: {
+          slug: "prod",
+          type: "PRODUCTION",
+          projectId: project.id,
+          organizationId: organization.id,
+          apiKey: "prod-org-count",
+          pkApiKey: "prod-org-count",
+          shortcode: "prod-org-count",
+        },
+      });
+
+      const developmentEnv = await prisma.runtimeEnvironment.create({
+        data: {
+          slug: "dev",
+          type: "DEVELOPMENT",
+          projectId: project.id,
+          organizationId: organization.id,
+          apiKey: "dev-org-count",
+          pkApiKey: "dev-org-count",
+          shortcode: "dev-org-count",
+        },
+      });
+
+      const previewEnv = await prisma.runtimeEnvironment.create({
+        data: {
+          slug: "preview",
+          type: "PREVIEW",
+          projectId: project.id,
+          organizationId: organization.id,
+          apiKey: "preview-org-count",
+          pkApiKey: "preview-org-count",
+          shortcode: "preview-org-count",
+          branchName: "feature-branch",
+          archivedAt: new Date(),
+        },
+      });
+
+      const runRows = [
+        {
+          friendlyId: "run_org_prod_pending",
+          env: productionEnv,
+          type: "PRODUCTION",
+          status: "PENDING",
+        },
+        {
+          friendlyId: "run_org_prod_delayed",
+          env: productionEnv,
+          type: "PRODUCTION",
+          status: "DELAYED",
+        },
+        {
+          friendlyId: "run_org_prod_done",
+          env: productionEnv,
+          type: "PRODUCTION",
+          status: "COMPLETED_SUCCESSFULLY",
+        },
+        {
+          friendlyId: "run_org_dev_pending",
+          env: developmentEnv,
+          type: "DEVELOPMENT",
+          status: "PENDING",
+        },
+        {
+          friendlyId: "run_org_preview_pending",
+          env: previewEnv,
+          type: "PREVIEW",
+          status: "PENDING",
+        },
+      ] as const;
+
+      for (const row of runRows) {
+        await prisma.taskRun.create({
+          data: {
+            friendlyId: row.friendlyId,
+            taskIdentifier: "queued-task",
+            status: row.status,
+            payload: JSON.stringify({}),
+            traceId: "trace",
+            spanId: "span",
+            queue: "main",
+            runtimeEnvironmentId: row.env.id,
+            projectId: project.id,
+            organizationId: organization.id,
+            environmentType: row.type,
+            engine: "V2",
+          },
+        });
+      }
+
+      await setTimeout(1000);
+
+      const count = await countBillableQueuedRunsForOrganization(organization.id, clickhouse);
+
+      expect(count).toBe(3);
+    }
+  );
+
+  replicationContainerTest(
+    "keeps archived environments in the billable environment list so enforcement can cancel their runs",
+    async ({ prisma }) => {
+      const organization = await prisma.organization.create({
+        data: { title: "billing-limit-archived", slug: "billing-limit-archived" },
+      });
+
+      const project = await prisma.project.create({
+        data: {
+          name: "billing-limit-archived",
+          slug: "billing-limit-archived",
+          organizationId: organization.id,
+          externalRef: "billing-limit-archived",
+        },
+      });
+
+      const activeEnv = await prisma.runtimeEnvironment.create({
+        data: {
+          slug: "preview-active",
+          type: "PREVIEW",
+          projectId: project.id,
+          organizationId: organization.id,
+          apiKey: "preview-active-test",
+          pkApiKey: "preview-active-test",
+          shortcode: "preview-active-test",
+          branchName: "live-branch",
+        },
+      });
+
+      const archivedEnv = await prisma.runtimeEnvironment.create({
+        data: {
+          slug: "preview-archived",
+          type: "PREVIEW",
+          projectId: project.id,
+          organizationId: organization.id,
+          apiKey: "preview-archived-test",
+          pkApiKey: "preview-archived-test",
+          shortcode: "preview-archived-test",
+          branchName: "old-branch",
+          archivedAt: new Date(),
+        },
+      });
+
+      const developmentEnv = await prisma.runtimeEnvironment.create({
+        data: {
+          slug: "dev",
+          type: "DEVELOPMENT",
+          projectId: project.id,
+          organizationId: organization.id,
+          apiKey: "dev-archived-test",
+          pkApiKey: "dev-archived-test",
+          shortcode: "dev-archived-test",
+        },
+      });
+
+      const environments = await getBillableEnvironmentsForBillingLimit(organization.id, prisma);
+      const environmentIds = environments.map((environment) => environment.id).sort();
+
+      expect(environmentIds).toEqual([activeEnv.id, archivedEnv.id].sort());
+      expect(environmentIds).not.toContain(developmentEnv.id);
+    }
+  );
 });