From fbc87b519988aa0a70e2fb8f217f67038c80d8b0 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 22 Sep 2026 11:15:18 -0700 Subject: [PATCH 1/3] fix(billing): let a slow ledger read finish instead of blocking execution at the singleflight default --- .../lib/billing/core/usage-gate-cache.test.ts | 35 +++++++++++++++++++ apps/sim/lib/billing/core/usage-gate-cache.ts | 25 ++++++++++--- 2 files changed, 55 insertions(+), 5 deletions(-) diff --git a/apps/sim/lib/billing/core/usage-gate-cache.test.ts b/apps/sim/lib/billing/core/usage-gate-cache.test.ts index 752ae97e762..343f2d29db2 100644 --- a/apps/sim/lib/billing/core/usage-gate-cache.test.ts +++ b/apps/sim/lib/billing/core/usage-gate-cache.test.ts @@ -16,6 +16,7 @@ import { checkIngestionUsageLimits, checkSearchUsageLimits, resetUsageGateCache, + USAGE_GATE_SETTLE_TIMEOUT_MS, USAGE_GATE_TTL_MS, } from '@/lib/billing/core/usage-gate-cache' @@ -178,4 +179,38 @@ describe('checkExecutionUsageLimits', () => { await checkExecutionUsageLimits(ATTRIBUTION) expect(mockCheck).toHaveBeenCalledTimes(2) }) + + 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)) + ) + const pending = checkExecutionUsageLimits(ATTRIBUTION) + await vi.advanceTimersByTimeAsync(45_000) + await expect(pending).resolves.toEqual({ isExceeded: false }) + expect(mockCheck).toHaveBeenCalledTimes(1) + } finally { + vi.useRealTimers() + } + }) + + it('gives up on a read that never answers at the gate deadline, then reads fresh', async () => { + vi.useFakeTimers() + try { + mockCheck.mockReturnValueOnce(new Promise(() => {})) + const hung = checkExecutionUsageLimits(ATTRIBUTION) + const rejection = expect(hung).rejects.toThrow( + `did not settle within ${USAGE_GATE_SETTLE_TIMEOUT_MS}ms` + ) + await vi.advanceTimersByTimeAsync(USAGE_GATE_SETTLE_TIMEOUT_MS) + await rejection + await expect(checkExecutionUsageLimits(ATTRIBUTION)).resolves.toEqual({ + isExceeded: false, + }) + expect(mockCheck).toHaveBeenCalledTimes(2) + } finally { + vi.useRealTimers() + } + }) }) diff --git a/apps/sim/lib/billing/core/usage-gate-cache.ts b/apps/sim/lib/billing/core/usage-gate-cache.ts index 4e3a1284226..7a4fa799cce 100644 --- a/apps/sim/lib/billing/core/usage-gate-cache.ts +++ b/apps/sim/lib/billing/core/usage-gate-cache.ts @@ -20,6 +20,19 @@ 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. + */ +export const USAGE_GATE_SETTLE_TIMEOUT_MS = 120_000 + /** * Recent gate answers, admitted and refused, with `LRUCache` supplying the TTL * and the size bound. Each entry point decides which of them it may serve. @@ -62,9 +75,9 @@ function gateKey(attribution: BillingAttributionSnapshot): string { * the cache. A read that throws writes nothing. * * `coalesceLocally` collapses concurrent misses onto one ledger read and bounds - * a hung read at its settle deadline. The write stays on the value this caller - * received, so a producer that timed out and later resolved cannot overwrite a - * fresher answer. + * a hung read at {@link USAGE_GATE_SETTLE_TIMEOUT_MS}. The write stays on the + * value this caller received, so a producer that timed out and later resolved + * cannot overwrite a fresher answer. * * There is deliberately no invalidator: usage and limit changes land in other * processes (execution workers, Stripe webhooks), so the TTL is the real bound. @@ -77,8 +90,10 @@ async function checkUsageLimitsThroughCache( const cached = gateCache.get(key) if (cached !== undefined && (cacheRefusals || !cached.isExceeded)) return cached - const result = await coalesceLocally(`usage-gate:${key}`, () => - checkAttributedUsageLimits(attribution) + const result = await coalesceLocally( + `usage-gate:${key}`, + () => checkAttributedUsageLimits(attribution), + USAGE_GATE_SETTLE_TIMEOUT_MS ) if (cacheRefusals || !result.isExceeded) gateCache.set(key, result) return result From 237c22bdd31bcac8434407e0a1b3767ccc2b21ed Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 22 Sep 2026 11:27:01 -0700 Subject: [PATCH 2/3] fix(billing): bound the ledger sum at the database and derive the gate deadline from it --- .../lib/billing/core/usage-gate-cache.test.ts | 4 +-- apps/sim/lib/billing/core/usage-gate-cache.ts | 19 +++++------ apps/sim/lib/billing/core/usage-log.test.ts | 34 +++++++++++++++++++ apps/sim/lib/billing/core/usage-log.ts | 30 ++++++++++++---- .../billing/enterprise-provisioning.test.ts | 2 ++ 5 files changed, 70 insertions(+), 19 deletions(-) diff --git a/apps/sim/lib/billing/core/usage-gate-cache.test.ts b/apps/sim/lib/billing/core/usage-gate-cache.test.ts index 343f2d29db2..3a073fc0640 100644 --- a/apps/sim/lib/billing/core/usage-gate-cache.test.ts +++ b/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 }) diff --git a/apps/sim/lib/billing/core/usage-gate-cache.ts b/apps/sim/lib/billing/core/usage-gate-cache.ts index 7a4fa799cce..f882c80589e 100644 --- a/apps/sim/lib/billing/core/usage-gate-cache.ts +++ b/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 diff --git a/apps/sim/lib/billing/core/usage-log.test.ts b/apps/sim/lib/billing/core/usage-log.test.ts index cd612889af2..dafcf58ab2e 100644 --- a/apps/sim/lib/billing/core/usage-log.test.ts +++ b/apps/sim/lib/billing/core/usage-log.test.ts @@ -36,6 +36,7 @@ vi.mock('@/lib/billing/subscriptions/utils', () => ({ import { CUMULATIVE_COST_EPSILON, CumulativeUsageContextMismatchError, + getBillingPeriodUsageCost, getUserUsageLogs, getWorkspaceUsageLogs, recordCumulativeUsage, @@ -43,6 +44,7 @@ import { 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) => + 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]) + }) +}) diff --git a/apps/sim/lib/billing/core/usage-log.ts b/apps/sim/lib/billing/core/usage-log.ts index b79653f0e49..5559edc7529 100644 --- a/apps/sim/lib/billing/core/usage-log.ts +++ b/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`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`COALESCE(SUM(${usageLog.cost}), 0)`, + }) + .from(usageLog) + .where(and(...conditions)) + }) return Number.parseFloat(row?.cost ?? '0') } diff --git a/apps/sim/lib/billing/enterprise-provisioning.test.ts b/apps/sim/lib/billing/enterprise-provisioning.test.ts index 1368760c0ee..b69eb60f452 100644 --- a/apps/sim/lib/billing/enterprise-provisioning.test.ts +++ b/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({ From 5b6a4a28c8fe3d5526f9ca081777bc08977d4e97 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 22 Sep 2026 11:36:07 -0700 Subject: [PATCH 3/3] fix(billing): keep the ledger statement bound in the billing constants module --- apps/sim/lib/billing/constants.ts | 9 +++++++++ apps/sim/lib/billing/core/usage-gate-cache.ts | 2 +- apps/sim/lib/billing/core/usage-log.test.ts | 2 +- apps/sim/lib/billing/core/usage-log.ts | 10 +--------- 4 files changed, 12 insertions(+), 11 deletions(-) diff --git a/apps/sim/lib/billing/constants.ts b/apps/sim/lib/billing/constants.ts index ebe5e7c17b5..c1fd884a293 100644 --- a/apps/sim/lib/billing/constants.ts +++ b/apps/sim/lib/billing/constants.ts @@ -39,6 +39,15 @@ export const DEFAULT_OVERAGE_THRESHOLD = 100 */ export const BILLING_LOCK_TIMEOUT_MS = 5_000 +/** + * 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 + /** * Available credit tiers. Each tier maps a credit amount to the underlying dollar * cost and carries that tier's fixed weekly refresh allowance. diff --git a/apps/sim/lib/billing/core/usage-gate-cache.ts b/apps/sim/lib/billing/core/usage-gate-cache.ts index f882c80589e..dbcf326f3d4 100644 --- a/apps/sim/lib/billing/core/usage-gate-cache.ts +++ b/apps/sim/lib/billing/core/usage-gate-cache.ts @@ -1,10 +1,10 @@ import { LRUCache } from 'lru-cache' +import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants' import { type AttributedUsageLimitsResult, 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' /** diff --git a/apps/sim/lib/billing/core/usage-log.test.ts b/apps/sim/lib/billing/core/usage-log.test.ts index dafcf58ab2e..ccb88ae4c0b 100644 --- a/apps/sim/lib/billing/core/usage-log.test.ts +++ b/apps/sim/lib/billing/core/usage-log.test.ts @@ -33,6 +33,7 @@ vi.mock('@/lib/billing/subscriptions/utils', () => ({ isOrgScopedSubscription: mockIsOrgScopedSubscription, })) +import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants' import { CUMULATIVE_COST_EPSILON, CumulativeUsageContextMismatchError, @@ -44,7 +45,6 @@ import { 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' diff --git a/apps/sim/lib/billing/core/usage-log.ts b/apps/sim/lib/billing/core/usage-log.ts index 5559edc7529..0849352ef35 100644 --- a/apps/sim/lib/billing/core/usage-log.ts +++ b/apps/sim/lib/billing/core/usage-log.ts @@ -15,6 +15,7 @@ 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 { getHighestPrioritySubscription } from '@/lib/billing/core/plan' import { @@ -212,15 +213,6 @@ 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.