fix(billing): bound the ledger sum at the database and derive the gate deadline from it

Author: waleedlatif1

apps/sim/lib/billing/core/usage-gate-cache.test.ts

@@ -183,9 +183,7 @@ describe('checkExecutionUsageLimits', () => {
   it('waits for a slow ledger read past the singleflight default instead of blocking', async () => {
     vi.useFakeTimers()
     try {
-      mockCheck.mockReturnValueOnce(
-        new Promise((resolve) => setTimeout(() => resolve({ isExceeded: false }), 45_000))
-      )
+      mockCheck.mockImplementationOnce(() => sleep(45_000).then(() => ({ isExceeded: false })))
       const pending = checkExecutionUsageLimits(ATTRIBUTION)
       await vi.advanceTimersByTimeAsync(45_000)
       await expect(pending).resolves.toEqual({ isExceeded: false })

apps/sim/lib/billing/core/usage-gate-cache.ts

@@ -4,6 +4,7 @@ import {
   type BillingAttributionSnapshot,
   checkAttributedUsageLimits,
 } from '@/lib/billing/core/billing-attribution'
+import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/core/usage-log'
 import { coalesceLocally } from '@/lib/concurrency/singleflight'
 
 /**
@@ -21,17 +22,15 @@ import { coalesceLocally } from '@/lib/concurrency/singleflight'
 export const USAGE_GATE_TTL_MS = 5 * 60 * 1000
 
 /**
- * How long a coalesced ledger read may take before its callers give up on it. The read sums a
- * payer's ledger for the billing period, which for a large organization is millions of rows and,
- * from a cold cache or under heavy I/O, takes longer than the singleflight default of 30 s. That
- * default exists to bound a hung producer, and a slow read is not a hung one: the database bounds
- * every statement with its own timeout, after which the read fails on its own and the failure is
- * reported rather than cached. The deadline therefore sits above any statement ceiling the
- * deployment applies, so only a connection that never answers is given up on. A shorter deadline
- * fails the callers while the read is still running, and the next caller starts a second read
- * of the same ledger alongside it.
+ * How long a coalesced usage read may take before its callers give up on it. The read's cost is
+ * the ledger sum, which the database ends at {@link USAGE_LEDGER_STATEMENT_TIMEOUT_MS}; the
+ * remainder is a few indexed lookups and the connection waits around them. The singleflight
+ * default of 30 s exists to bound a hung producer, and a slow sum is not a hung one: given up on
+ * early, it keeps running detached while every joined caller fails and the next caller starts a
+ * second sum alongside it. Derived from the statement bound so the database always ends the sum
+ * first, and the gate only gives up on a connection that never answers.
  */
-export const USAGE_GATE_SETTLE_TIMEOUT_MS = 120_000
+export const USAGE_GATE_SETTLE_TIMEOUT_MS = USAGE_LEDGER_STATEMENT_TIMEOUT_MS + 15_000
 
 /**
  * Recent gate answers, admitted and refused, with `LRUCache` supplying the TTL

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

@@ -36,13 +36,15 @@ vi.mock('@/lib/billing/subscriptions/utils', () => ({
 import {
   CUMULATIVE_COST_EPSILON,
   CumulativeUsageContextMismatchError,
+  getBillingPeriodUsageCost,
   getUserUsageLogs,
   getWorkspaceUsageLogs,
   recordCumulativeUsage,
   recordUsage,
   resolveCumulativeTopUp,
   UNKNOWN_CURSOR_MESSAGE,
   UnknownUsageCursorError,
+  USAGE_LEDGER_STATEMENT_TIMEOUT_MS,
 } from '@/lib/billing/core/usage-log'
 import { asOrchestrationError } from '@/lib/core/orchestration/types'
 import { HttpError } from '@/lib/core/utils/http-error'
@@ -554,3 +556,35 @@ describe('usage-log query scopes', () => {
     })
   })
 })
+
+describe('getBillingPeriodUsageCost', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    installSharedDbMocks()
+  })
+
+  it('bounds the ledger sum with its own statement timeout inside one transaction', async () => {
+    const execute = vi.fn().mockResolvedValue([])
+    const where = vi.fn().mockResolvedValue([{ cost: '12.5' }])
+    const tx = { execute, select: vi.fn(() => ({ from: vi.fn(() => ({ where })) })) }
+    mockTransaction.mockImplementation((callback: (client: typeof tx) => Promise<unknown>) =>
+      callback(tx)
+    )
+
+    const cost = await getBillingPeriodUsageCost(
+      { type: 'organization', id: 'org-1' },
+      { start: new Date('2026-05-01T00:00:00Z'), end: new Date('2027-05-01T00:00:00Z') }
+    )
+
+    expect(cost).toBe(12.5)
+    expect(mockTransaction).toHaveBeenCalledTimes(1)
+    const executed = execute.mock.calls.map(
+      ([statement]) => (statement as { toSQL: () => { sql: string } }).toSQL().sql
+    )
+    expect(executed).toContain(
+      `SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'`
+    )
+    /** The bound is set before the sum runs, not after. */
+    expect(execute.mock.invocationCallOrder[0]).toBeLessThan(where.mock.invocationCallOrder[0])
+  })
+})

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

@@ -212,9 +212,22 @@ async function resolveBillingContext(
   }
 }
 
+/**
+ * Bound on one ledger sum. A large payer's period covers millions of rows, and from a cold
+ * cache or under heavy I/O the sum can run for tens of seconds; past this the database ends it
+ * and the read fails, so a caller that admits on the sum fails closed rather than waiting
+ * without limit. The usage gate derives its coalescing deadline from this bound, so the sum
+ * always ends at the database before the gate gives up on it.
+ */
+export const USAGE_LEDGER_STATEMENT_TIMEOUT_MS = 60_000
+
 /**
  * Returns attributed ledger usage for a billing entity/period. The ledger is
  * the sole source of truth for usage — there is no userStats baseline.
+ *
+ * The sum runs in a transaction of its own on the given client so that it can
+ * be bounded by {@link USAGE_LEDGER_STATEMENT_TIMEOUT_MS} for that statement
+ * alone: `SET LOCAL` ends with the transaction and never reaches the pool.
  */
 export async function getBillingPeriodUsageCost(
   billingEntity: BillingEntity,
@@ -238,12 +251,17 @@ export async function getBillingPeriodUsageCost(
     )
   }
 
-  const [row] = await executor
-    .select({
-      cost: sql<string>`COALESCE(SUM(${usageLog.cost}), 0)`,
-    })
-    .from(usageLog)
-    .where(and(...conditions))
+  const [row] = await executor.transaction(async (tx) => {
+    await tx.execute(
+      sql.raw(`SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'`)
+    )
+    return tx
+      .select({
+        cost: sql<string>`COALESCE(SUM(${usageLog.cost}), 0)`,
+      })
+      .from(usageLog)
+      .where(and(...conditions))
+  })
 
   return Number.parseFloat(row?.cost ?? '0')
 }

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

@@ -199,6 +199,8 @@ describe('Enterprise issuance preflight', () => {
     queueTableRows(schemaMock.workspace, [])
     queueTableRows(schemaMock.workspace, [])
     queueTableRows(schemaMock.subscription, [])
+    /** The run count resolves first; the ledger sum opens its bounded transaction before it reads. */
+    queueTableRows(schemaMock.usageLog, [{ workflowRuns: 0 }])
     queueTableRows(schemaMock.usageLog, [{ cost: '150' }])
 
     const result = await getEnterpriseIssuancePreflight({