fix(billing): bound every ledger aggregate at the database and size the gate deadline for two

Author: waleedlatif1

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

@@ -0,0 +1,35 @@
+/**
+ * @vitest-environment node
+ */
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants'
+import { readLedgerBounded } from '@/lib/billing/core/ledger-read'
+import type { DbClient } from '@/lib/db/types'
+
+const renderedSql = (statement: unknown) =>
+  (statement as { toSQL: () => { sql: string } }).toSQL().sql
+
+describe('readLedgerBounded', () => {
+  const execute = vi.fn().mockResolvedValue([])
+  const tx = { execute }
+  const transaction = vi.fn((callback: (client: typeof tx) => Promise<unknown>) => callback(tx))
+  const executor = { transaction } as unknown as DbClient
+
+  beforeEach(() => vi.clearAllMocks())
+
+  it('bounds the statement inside one transaction on the given client, before the read', async () => {
+    const read = vi.fn().mockResolvedValue([{ cost: '12.5' }])
+    await expect(readLedgerBounded(executor, read)).resolves.toEqual([{ cost: '12.5' }])
+    expect(transaction).toHaveBeenCalledTimes(1)
+    expect(execute.mock.calls.map(([statement]) => renderedSql(statement))).toEqual([
+      `SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'`,
+    ])
+    expect(read).toHaveBeenCalledWith(tx)
+    expect(execute.mock.invocationCallOrder[0]).toBeLessThan(read.mock.invocationCallOrder[0])
+  })
+
+  it('surfaces the read failure to the caller', async () => {
+    const failure = new Error('canceling statement due to statement timeout')
+    await expect(readLedgerBounded(executor, () => Promise.reject(failure))).rejects.toBe(failure)
+  })
+})

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

@@ -0,0 +1,23 @@
+import { sql } from 'drizzle-orm'
+import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants'
+import type { DbClient, DbTransaction } from '@/lib/db/types'
+
+/**
+ * Runs one aggregate over the usage ledger in a transaction of its own, bounded by
+ * {@link USAGE_LEDGER_STATEMENT_TIMEOUT_MS}. `SET LOCAL` scopes the bound to that transaction,
+ * so it ends with the read and never reaches the pool. Every ledger aggregate reads through
+ * here: a payer whose period has grown past what one statement can sum within the bound fails
+ * at the database instead of holding a connection without limit, and a caller that admits on
+ * the answer can size its own deadline from the bound.
+ */
+export function readLedgerBounded<T>(
+  executor: DbClient,
+  read: (tx: DbTransaction) => Promise<T>
+): Promise<T> {
+  return executor.transaction(async (tx) => {
+    await tx.execute(
+      sql.raw(`SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'`)
+    )
+    return read(tx)
+  })
+}

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

@@ -23,14 +23,16 @@ export const USAGE_GATE_TTL_MS = 5 * 60 * 1000
 
 /**
  * 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.
+ * its ledger aggregates, each of which the database ends at
+ * {@link USAGE_LEDGER_STATEMENT_TIMEOUT_MS}; at most two run in sequence (the payer's usage,
+ * then a member's cap), and 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
+ * aggregate 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 one alongside it. Sized from the statement
+ * bound so the database always ends the aggregates first, and the gate only gives up on a
+ * connection that never answers.
  */
-export const USAGE_GATE_SETTLE_TIMEOUT_MS = USAGE_LEDGER_STATEMENT_TIMEOUT_MS + 15_000
+export const USAGE_GATE_SETTLE_TIMEOUT_MS = 2 * 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

@@ -38,6 +38,9 @@ import {
   CUMULATIVE_COST_EPSILON,
   CumulativeUsageContextMismatchError,
   getBillingPeriodUsageCost,
+  getBillingPeriodUsageCostByUser,
+  getBillingPeriodUsageCostWithSourceSubset,
+  getBillingPeriodWorkflowRunCount,
   getUserUsageLogs,
   getWorkspaceUsageLogs,
   recordCumulativeUsage,
@@ -557,34 +560,73 @@ describe('usage-log query scopes', () => {
   })
 })
 
-describe('getBillingPeriodUsageCost', () => {
+describe('ledger aggregates', () => {
+  const billingEntity = { type: 'organization' as const, id: 'org-1' }
+  const billingPeriod = {
+    start: new Date('2026-05-01T00:00:00Z'),
+    end: new Date('2027-05-01T00:00:00Z'),
+  }
+  /** Every aggregate over the ledger, with the row the mocked read hands back and the value it yields. */
+  const aggregates: Array<{
+    name: string
+    read: () => Promise<unknown>
+    rows: unknown[]
+    expected: unknown
+  }> = [
+    {
+      name: 'getBillingPeriodUsageCost',
+      read: () => getBillingPeriodUsageCost(billingEntity, billingPeriod),
+      rows: [{ cost: '12.5' }],
+      expected: 12.5,
+    },
+    {
+      name: 'getBillingPeriodWorkflowRunCount',
+      read: () => getBillingPeriodWorkflowRunCount(billingEntity, billingPeriod),
+      rows: [{ workflowRuns: 7 }],
+      expected: 7,
+    },
+    {
+      name: 'getBillingPeriodUsageCostWithSourceSubset',
+      read: () =>
+        getBillingPeriodUsageCostWithSourceSubset(billingEntity, billingPeriod, ['workflow']),
+      rows: [{ total: '20', subset: '5' }],
+      expected: { total: 20, subset: 5 },
+    },
+    {
+      name: 'getBillingPeriodUsageCostByUser',
+      read: () => getBillingPeriodUsageCostByUser(billingEntity, billingPeriod),
+      rows: [{ userId: 'user-1', cost: '3' }],
+      expected: new Map([['user-1', 3]]),
+    },
+  ]
+
   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])
-  })
+  for (const aggregate of aggregates) {
+    it(`${aggregate.name} reads through the bounded ledger transaction`, async () => {
+      const execute = vi.fn().mockResolvedValue([])
+      const terminal = vi.fn().mockResolvedValue(aggregate.rows)
+      const chain: Record<string, unknown> = {}
+      for (const step of ['select', 'from', 'where', 'leftJoin']) chain[step] = vi.fn(() => chain)
+      chain.groupBy = terminal
+      chain.then = (resolve: (rows: unknown[]) => unknown) => terminal().then(resolve)
+      const tx = { execute, select: chain.select }
+      mockTransaction.mockImplementation((callback: (client: typeof tx) => Promise<unknown>) =>
+        callback(tx)
+      )
+
+      await expect(aggregate.read()).resolves.toEqual(aggregate.expected)
+      expect(mockTransaction).toHaveBeenCalledTimes(1)
+      expect(
+        execute.mock.calls.map(
+          ([statement]) => (statement as { toSQL: () => { sql: string } }).toSQL().sql
+        )
+      ).toEqual([`SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'`])
+      /** The bound is set before the aggregate runs, not after. */
+      expect(execute.mock.invocationCallOrder[0]).toBeLessThan(terminal.mock.invocationCallOrder[0])
+    })
+  }
 })

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

@@ -15,8 +15,8 @@ import {
   textKey,
   timestampKey,
 } from '@/lib/api/list-query'
-import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants'
 import { defaultBillingPeriod } from '@/lib/billing/core/billing-period'
+import { readLedgerBounded } from '@/lib/billing/core/ledger-read'
 import { getHighestPrioritySubscription } from '@/lib/billing/core/plan'
 import {
   resolveSubscriptionUsagePeriod,
@@ -216,10 +216,6 @@ async function resolveBillingContext(
 /**
  * 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,
@@ -243,17 +239,14 @@ export async function getBillingPeriodUsageCost(
     )
   }
 
-  const [row] = await executor.transaction(async (tx) => {
-    await tx.execute(
-      sql.raw(`SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'`)
-    )
-    return tx
+  const [row] = await readLedgerBounded(executor, (tx) =>
+    tx
       .select({
         cost: sql<string>`COALESCE(SUM(${usageLog.cost}), 0)`,
       })
       .from(usageLog)
       .where(and(...conditions))
-  })
+  )
 
   return Number.parseFloat(row?.cost ?? '0')
 }
@@ -274,36 +267,38 @@ export async function getBillingPeriodWorkflowRunCount(
   billingPeriod: UsageQueryPeriod,
   executor: DbClient = db
 ): Promise<number> {
-  const [row] = await executor
-    .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),
+              ])
+        )
       )
-    )
+  )
 
   return row?.workflowRuns ?? 0
 }
@@ -321,27 +316,29 @@ export async function getBillingPeriodUsageCostWithSourceSubset(
   source: UsageLogSource[],
   executor: DbClient = db
 ): Promise<{ total: number; subset: number }> {
-  const [row] = await executor
-    .select({
-      total: sql<string>`COALESCE(SUM(${usageLog.cost}), 0)`,
-      subset: sql<string>`COALESCE(SUM(${usageLog.cost}) FILTER (WHERE ${inArray(usageLog.source, source)}), 0)`,
-    })
-    .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({
+        total: sql<string>`COALESCE(SUM(${usageLog.cost}), 0)`,
+        subset: sql<string>`COALESCE(SUM(${usageLog.cost}) FILTER (WHERE ${inArray(usageLog.source, source)}), 0)`,
+      })
+      .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),
+              ])
+        )
       )
-    )
+  )
 
   return {
     total: Number.parseFloat(row?.total ?? '0'),
@@ -377,14 +374,16 @@ export async function getBillingPeriodUsageCostByUser(
   }
   if (userIds) conditions.push(inArray(usageLog.userId, [...userIds]))
 
-  const rows = await executor
-    .select({
-      userId: usageLog.userId,
-      cost: sql<string>`COALESCE(SUM(${usageLog.cost}), 0)`,
-    })
-    .from(usageLog)
-    .where(and(...conditions))
-    .groupBy(usageLog.userId)
+  const rows = await readLedgerBounded(executor, (tx) =>
+    tx
+      .select({
+        userId: usageLog.userId,
+        cost: sql<string>`COALESCE(SUM(${usageLog.cost}), 0)`,
+      })
+      .from(usageLog)
+      .where(and(...conditions))
+      .groupBy(usageLog.userId)
+  )
 
   return new Map(rows.map((row) => [row.userId, Number.parseFloat(row.cost ?? '0')]))
 }