fix(billing): bound the enterprise preflight's workflow run count to the preview budget

apps/sim/lib/admin/dashboard-usage.ts

@@ -1,14 +1,16 @@
 import { db } from '@sim/db'
 import { member, permissions, usageLog, workspace } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
-import { getPostgresCancellationReason } from '@sim/utils/errors'
 import { and, countDistinct, eq, inArray, notInArray, type SQL, sql } from 'drizzle-orm'
 import { unionAll } from 'drizzle-orm/pg-core'
-import { readLedgerBounded } from '@/lib/billing/core/ledger-read'
+import {
+  isLedgerReadOverBudget,
+  readLedgerBounded,
+  WORKFLOW_RUN_COUNT_BUDGET_MS,
+} from '@/lib/billing/core/ledger-read'
 import type { ResolvedUsagePeriod } from '@/lib/billing/core/reporting-period'
 import { buildUsageAnalyticsScope } from '@/lib/billing/core/usage-analytics'
 import type { BillingEntity } from '@/lib/billing/core/usage-log'
-import { DatabaseReadDeadlineError } from '@/lib/db/read-retry'
 
 /**
  * Ledger reads behind the admin dashboard: each billing entity's spend and workflow runs
@@ -26,7 +28,7 @@ import { DatabaseReadDeadlineError } from '@/lib/db/read-retry'
  * Widening both entity indexes with it would grow the hot billing sums' working set to
  * speed up a staff-only page. So only one organization's detail counts runs, in a single
  * scan that yields the total and every actor's share, under
- * {@link WORKFLOW_RUNS_BUDGET_MS}; past it the count is unknown (`null`) and the page
+ * {@link WORKFLOW_RUN_COUNT_BUDGET_MS}; past it the count is unknown (`null`) and the page
  * still loads.
  *
  * They stay on the primary, where the hot billing sums keep those index pages cached. On
@@ -54,12 +56,6 @@ interface DashboardWorkflowRuns {
   formerActors: number
 }
 
-/**
- * How long one detail page may spend counting workflow runs. Well inside the admin
- * proxy's request timeout, so a payer too large to count still gets its page.
- */
-const WORKFLOW_RUNS_BUDGET_MS = 5_000
-
 const WORKFLOW_SOURCE = sql`${usageLog.source} = 'workflow'`
 
 export function organizationWindow(
@@ -164,7 +160,7 @@ function isFormerActor(organizationId: string): SQL {
 /**
  * An organization's workflow runs over its period, in total, per acting user and across its
  * former actors, from one scan. `null` when the count does not finish within
- * {@link WORKFLOW_RUNS_BUDGET_MS}.
+ * {@link WORKFLOW_RUN_COUNT_BUDGET_MS}.
  */
 export async function readDashboardOrganizationWorkflowRuns(
   organizationId: string,
@@ -193,16 +189,13 @@ export async function readDashboardOrganizationWorkflowRuns(
           .from(workflowRows)
           .groupBy(sql`grouping sets ((${workflowRows.userId}), (${workflowRows.formerActor}), ())`)
       },
-      WORKFLOW_RUNS_BUDGET_MS
+      WORKFLOW_RUN_COUNT_BUDGET_MS
     )
   } catch (error) {
-    const overBudget =
-      error instanceof DatabaseReadDeadlineError ||
-      getPostgresCancellationReason(error) === 'statement_timeout'
-    if (!overBudget) throw error
+    if (!isLedgerReadOverBudget(error)) throw error
     logger.warn('Workflow run count exceeded its budget', {
       organizationId,
-      budgetMs: WORKFLOW_RUNS_BUDGET_MS,
+      budgetMs: WORKFLOW_RUN_COUNT_BUDGET_MS,
     })
     return null
   }

apps/sim/lib/billing/core/ledger-read.ts

@@ -1,3 +1,4 @@
+import { getPostgresCancellationReason } from '@sim/utils/errors'
 import { sql } from 'drizzle-orm'
 import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants'
 import { DatabaseReadDeadlineError, withDatabaseReadRetry } from '@/lib/db/read-retry'
@@ -36,3 +37,17 @@ export function readLedgerBounded<T>(
     { label: 'usage ledger', maxAttempts: 2, maxElapsedMs: budgetMs }
   )
 }
+
+/**
+ * How long a staff-only workflow run count may take before it is reported as unknown. Well inside
+ * the admin proxy's request timeout, so a payer too large to count still gets its page.
+ */
+export const WORKFLOW_RUN_COUNT_BUDGET_MS = 5_000
+
+/** Whether a {@link readLedgerBounded} read failed because it ran out of its budget. */
+export function isLedgerReadOverBudget(error: unknown): boolean {
+  return (
+    error instanceof DatabaseReadDeadlineError ||
+    getPostgresCancellationReason(error) === 'statement_timeout'
+  )
+}

apps/sim/lib/billing/core/usage-log.integration.ts

@@ -32,6 +32,7 @@ import {
   CumulativeUsagePeriodClosedError,
   getBillingPeriodUsageCost,
   getBillingPeriodUsageCostByUser,
+  getBillingPeriodWorkflowRunCountWithinBudget,
   getStampedPeriodRangeUsageCostByUser,
   type RecordCumulativeUsageParams,
   recordCumulativeUsage,
@@ -375,6 +376,38 @@ describe('Cumulative billing with PostgreSQL', () => {
     )
   })
 
+  it('reports a workflow run count as unknown once it outlasts the preview budget', async () => {
+    const entity = { type: 'organization' as const, id: 'org-runs' }
+    const period = {
+      start: new Date('2026-08-01T00:00:00.000Z'),
+      end: new Date('2026-09-01T00:00:00.000Z'),
+      source: 'reporting' as const,
+    }
+    await connection`
+      insert into usage_log (id, user_id, category, source, description, cost,
+        billing_entity_type, billing_entity_id, execution_id, created_at)
+      values ('r1', 'u1', 'model', 'workflow', 'run', 1, 'organization', 'org-runs', 'e1', '2026-08-02'),
+             ('r2', 'u1', 'tool', 'workflow', 'run', 1, 'organization', 'org-runs', 'e1', '2026-08-02'),
+             ('r3', 'u1', 'model', 'workflow', 'run', 1, 'organization', 'org-runs', 'e2', '2026-08-03')
+    `
+    expect(await getBillingPeriodWorkflowRunCountWithinBudget(entity, period)).toBe(2)
+
+    const release = deferred()
+    const locked = deferred()
+    const holder = connection.begin(async (tx) => {
+      await tx.unsafe('LOCK TABLE usage_log IN ACCESS EXCLUSIVE MODE')
+      locked.resolve()
+      await release.promise
+    })
+    try {
+      await locked.promise
+      expect(await getBillingPeriodWorkflowRunCountWithinBudget(entity, period)).toBeNull()
+    } finally {
+      release.resolve()
+      await holder
+    }
+  }, 20_000)
+
   describe('a request that outlives its billing period', () => {
     // Past periods: the old period's row is written under the subscription lock only once
     // that period has ended.

apps/sim/lib/billing/core/usage-log.ts

@@ -18,7 +18,11 @@ import {
   timestampKey,
 } from '@/lib/api/list-query'
 import { defaultBillingPeriod } from '@/lib/billing/core/billing-period'
-import { readLedgerBounded } from '@/lib/billing/core/ledger-read'
+import {
+  isLedgerReadOverBudget,
+  readLedgerBounded,
+  WORKFLOW_RUN_COUNT_BUDGET_MS,
+} from '@/lib/billing/core/ledger-read'
 import { getHighestPrioritySubscription } from '@/lib/billing/core/plan'
 import {
   resolveSubscriptionUsagePeriod,
@@ -269,44 +273,74 @@ export async function getBillingPeriodUsageCost(
 export async function getBillingPeriodWorkflowRunCount(
   billingEntity: BillingEntity,
   billingPeriod: UsageQueryPeriod,
-  executor: DbClient = db
+  executor: DbClient = db,
+  budgetMs?: number
 ): Promise<number> {
-  const [row] = await readLedgerBounded(executor, (tx) =>
-    tx
-      .select({
-        /**
-         * The exclusion goes through `notInArray`, not `<> ALL(${array})`. Interpolating
-         * a JavaScript array into a `sql` template emits parenthesized scalar binds —
-         * `ALL(($1))` — which Postgres rejects outright with "op ANY/ALL (array)
-         * requires array on right side". Unit tests cannot catch it, because `@sim/db`
-         * is mocked and no statement is ever rendered.
-         */
-        workflowRuns:
-          sql<number>`COUNT(DISTINCT ${usageLog.executionId}) FILTER (WHERE ${usageLog.source} = 'workflow' AND ${notInArray(usageLog.category, [...UNBILLED_USAGE_CATEGORIES])})`.mapWith(
-            Number
-          ),
-      })
-      .from(usageLog)
-      .where(
-        and(
-          eq(usageLog.billingEntityType, billingEntity.type),
-          eq(usageLog.billingEntityId, billingEntity.id),
-          ...(billingPeriod.source === 'reporting'
-            ? [
-                gte(usageLog.createdAt, billingPeriod.start),
-                lt(usageLog.createdAt, billingPeriod.end),
-              ]
-            : [
-                eq(usageLog.billingPeriodStart, billingPeriod.start),
-                eq(usageLog.billingPeriodEnd, billingPeriod.end),
-              ])
-        )
-      )
+  const [row] = await readLedgerBounded(
+    executor,
+    (tx) =>
+      tx
+        .select({
+          /**
+           * The exclusion goes through `notInArray`, not `<> ALL(${array})`. Interpolating
+           * a JavaScript array into a `sql` template emits parenthesized scalar binds —
+           * `ALL(($1))` — which Postgres rejects outright with "op ANY/ALL (array)
+           * requires array on right side". Unit tests cannot catch it, because `@sim/db`
+           * is mocked and no statement is ever rendered.
+           */
+          workflowRuns:
+            sql<number>`COUNT(DISTINCT ${usageLog.executionId}) FILTER (WHERE ${usageLog.source} = 'workflow' AND ${notInArray(usageLog.category, [...UNBILLED_USAGE_CATEGORIES])})`.mapWith(
+              Number
+            ),
+        })
+        .from(usageLog)
+        .where(
+          and(
+            eq(usageLog.billingEntityType, billingEntity.type),
+            eq(usageLog.billingEntityId, billingEntity.id),
+            ...(billingPeriod.source === 'reporting'
+              ? [
+                  gte(usageLog.createdAt, billingPeriod.start),
+                  lt(usageLog.createdAt, billingPeriod.end),
+                ]
+              : [
+                  eq(usageLog.billingPeriodStart, billingPeriod.start),
+                  eq(usageLog.billingPeriodEnd, billingPeriod.end),
+                ])
+          )
+        ),
+    budgetMs
   )
 
   return row?.workflowRuns ?? 0
 }
 
+/**
+ * {@link getBillingPeriodWorkflowRunCount} for a staff-only preview that can do without it:
+ * `null` when the count does not finish within {@link WORKFLOW_RUN_COUNT_BUDGET_MS}, so a payer
+ * too large to count still gets its preview inside the admin proxy's request timeout.
+ */
+export async function getBillingPeriodWorkflowRunCountWithinBudget(
+  billingEntity: BillingEntity,
+  billingPeriod: UsageQueryPeriod
+): Promise<number | null> {
+  try {
+    return await getBillingPeriodWorkflowRunCount(
+      billingEntity,
+      billingPeriod,
+      db,
+      WORKFLOW_RUN_COUNT_BUDGET_MS
+    )
+  } catch (error) {
+    if (!isLedgerReadOverBudget(error)) throw error
+    logger.warn('Billing period workflow run count exceeded its budget', {
+      billingEntityType: billingEntity.type,
+      budgetMs: WORKFLOW_RUN_COUNT_BUDGET_MS,
+    })
+    return null
+  }
+}
+
 /**
  * Period total plus the portion attributable to `source`, in a single scan.
  *

apps/sim/lib/billing/enterprise-provisioning.ts

@@ -37,7 +37,7 @@ import { getBillingConcurrencyLimit } from '@/lib/billing/concurrency-limits'
 import { resolveEnterpriseReportingPeriod } from '@/lib/billing/core/reporting-period'
 import {
   getBillingPeriodUsageCost,
-  getBillingPeriodWorkflowRunCount,
+  getBillingPeriodWorkflowRunCountWithinBudget,
 } from '@/lib/billing/core/usage-log'
 import { creditsToDollars, dollarsToCredits } from '@/lib/billing/credits/conversion'
 import {
@@ -462,7 +462,8 @@ export interface EnterpriseIssuancePreflight {
       limitDollars: number
       usedCredits: number
       limitCredits: number
-      workflowRuns: number
+      /** `null` when the count did not finish within its budget. */
+      workflowRuns: number | null
     }
     invoiceAmountUsd: number
     configuredUsageLimitDollars: number
@@ -865,7 +866,7 @@ export async function getEnterpriseIssuancePreflight({
             { type: 'organization', id: membership.organizationId },
             reportingPeriod
           ),
-          getBillingPeriodWorkflowRunCount(
+          getBillingPeriodWorkflowRunCountWithinBudget(
             { type: 'organization', id: membership.organizationId },
             reportingPeriod
           ),

packages/testing/src/mocks/billing-usage-log.mock.ts

@@ -100,6 +100,7 @@ export const billingUsageLogMockFns = {
   mockDeriveBillingContext: vi.fn(),
   mockGetBillingPeriodUsageCost: vi.fn(),
   mockGetBillingPeriodWorkflowRunCount: vi.fn(),
+  mockGetBillingPeriodWorkflowRunCountWithinBudget: vi.fn(),
   mockGetBillingPeriodUsageCostWithSourceSubset: vi.fn(),
   mockGetBillingPeriodUsageCostByUser: vi.fn(),
   mockGetStampedPeriodRangeUsageCostByUser: vi.fn(),
@@ -138,6 +139,8 @@ export const billingUsageLogMock = {
   deriveBillingContext: billingUsageLogMockFns.mockDeriveBillingContext,
   getBillingPeriodUsageCost: billingUsageLogMockFns.mockGetBillingPeriodUsageCost,
   getBillingPeriodWorkflowRunCount: billingUsageLogMockFns.mockGetBillingPeriodWorkflowRunCount,
+  getBillingPeriodWorkflowRunCountWithinBudget:
+    billingUsageLogMockFns.mockGetBillingPeriodWorkflowRunCountWithinBudget,
   getBillingPeriodUsageCostWithSourceSubset:
     billingUsageLogMockFns.mockGetBillingPeriodUsageCostWithSourceSubset,
   getBillingPeriodUsageCostByUser: billingUsageLogMockFns.mockGetBillingPeriodUsageCostByUser,