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

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

* fix(billing): bound the cycle-close per-user sum and state the helper's scope precisely

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,25 @@
+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 a payer's 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 sum over a payer's billing period
+ * reads through here, whether it admits a run, closes a cycle or previews a bill: 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. Reads keyed to one execution or one stamped period
+ * boundary, and the platform-wide admin analytics, are not period sums and read directly.
+ */
+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,10 @@ import {
   CUMULATIVE_COST_EPSILON,
   CumulativeUsageContextMismatchError,
   getBillingPeriodUsageCost,
+  getBillingPeriodUsageCostByUser,
+  getBillingPeriodUsageCostWithSourceSubset,
+  getBillingPeriodWorkflowRunCount,
+  getStampedPeriodRangeUsageCostByUser,
   getUserUsageLogs,
   getWorkspaceUsageLogs,
   recordCumulativeUsage,
@@ -557,34 +561,83 @@ 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]]),
+    },
+    {
+      name: 'getStampedPeriodRangeUsageCostByUser',
+      read: () =>
+        getStampedPeriodRangeUsageCostByUser(billingEntity, {
+          from: billingPeriod.start,
+          to: billingPeriod.end,
+        }),
+      rows: [{ userId: 'user-2', cost: '4' }],
+      expected: new Map([['user-2', 4]]),
+    },
+  ]
+
   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])
+    })
+  }
 })