Skip to content

Commit b8b7201

Browse files
waleedlatif1waleedlatif1
authored
fix(billing): bound the enterprise preflight's workflow run count to the preview budget (#9673)
The enterprise issuance preflight counted the organization's billable workflow runs under the ledger's default 60s bound, the same COUNT DISTINCT the organization detail page already caps at 5s because it can outlast the admin proxy's request timeout on a large payer. The preflight now reads it through getBillingPeriodWorkflowRunCountWithinBudget, which reports null past the shared WORKFLOW_RUN_COUNT_BUDGET_MS, so the preview still loads; the contract already allows null. The over-budget classification moves next to readLedgerBounded as isLedgerReadOverBudget, shared by both callers.
1 parent 32cb270 commit b8b7201

6 files changed

Lines changed: 132 additions & 53 deletions

File tree

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

Lines changed: 10 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,16 @@
11
import { db } from '@sim/db'
22
import { member, permissions, usageLog, workspace } from '@sim/db/schema'
33
import { createLogger } from '@sim/logger'
4-
import { getPostgresCancellationReason } from '@sim/utils/errors'
54
import { and, countDistinct, eq, inArray, notInArray, type SQL, sql } from 'drizzle-orm'
65
import { unionAll } from 'drizzle-orm/pg-core'
7-
import { readLedgerBounded } from '@/lib/billing/core/ledger-read'
6+
import {
7+
isLedgerReadOverBudget,
8+
readLedgerBounded,
9+
WORKFLOW_RUN_COUNT_BUDGET_MS,
10+
} from '@/lib/billing/core/ledger-read'
811
import type { ResolvedUsagePeriod } from '@/lib/billing/core/reporting-period'
912
import { buildUsageAnalyticsScope } from '@/lib/billing/core/usage-analytics'
1013
import type { BillingEntity } from '@/lib/billing/core/usage-log'
11-
import { DatabaseReadDeadlineError } from '@/lib/db/read-retry'
1214

1315
/**
1416
* 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'
2628
* Widening both entity indexes with it would grow the hot billing sums' working set to
2729
* speed up a staff-only page. So only one organization's detail counts runs, in a single
2830
* scan that yields the total and every actor's share, under
29-
* {@link WORKFLOW_RUNS_BUDGET_MS}; past it the count is unknown (`null`) and the page
31+
* {@link WORKFLOW_RUN_COUNT_BUDGET_MS}; past it the count is unknown (`null`) and the page
3032
* still loads.
3133
*
3234
* They stay on the primary, where the hot billing sums keep those index pages cached. On
@@ -54,12 +56,6 @@ interface DashboardWorkflowRuns {
5456
formerActors: number
5557
}
5658

57-
/**
58-
* How long one detail page may spend counting workflow runs. Well inside the admin
59-
* proxy's request timeout, so a payer too large to count still gets its page.
60-
*/
61-
const WORKFLOW_RUNS_BUDGET_MS = 5_000
62-
6359
const WORKFLOW_SOURCE = sql`${usageLog.source} = 'workflow'`
6460

6561
export function organizationWindow(
@@ -164,7 +160,7 @@ function isFormerActor(organizationId: string): SQL {
164160
/**
165161
* An organization's workflow runs over its period, in total, per acting user and across its
166162
* former actors, from one scan. `null` when the count does not finish within
167-
* {@link WORKFLOW_RUNS_BUDGET_MS}.
163+
* {@link WORKFLOW_RUN_COUNT_BUDGET_MS}.
168164
*/
169165
export async function readDashboardOrganizationWorkflowRuns(
170166
organizationId: string,
@@ -193,16 +189,13 @@ export async function readDashboardOrganizationWorkflowRuns(
193189
.from(workflowRows)
194190
.groupBy(sql`grouping sets ((${workflowRows.userId}), (${workflowRows.formerActor}), ())`)
195191
},
196-
WORKFLOW_RUNS_BUDGET_MS
192+
WORKFLOW_RUN_COUNT_BUDGET_MS
197193
)
198194
} catch (error) {
199-
const overBudget =
200-
error instanceof DatabaseReadDeadlineError ||
201-
getPostgresCancellationReason(error) === 'statement_timeout'
202-
if (!overBudget) throw error
195+
if (!isLedgerReadOverBudget(error)) throw error
203196
logger.warn('Workflow run count exceeded its budget', {
204197
organizationId,
205-
budgetMs: WORKFLOW_RUNS_BUDGET_MS,
198+
budgetMs: WORKFLOW_RUN_COUNT_BUDGET_MS,
206199
})
207200
return null
208201
}

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { getPostgresCancellationReason } from '@sim/utils/errors'
12
import { sql } from 'drizzle-orm'
23
import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants'
34
import { DatabaseReadDeadlineError, withDatabaseReadRetry } from '@/lib/db/read-retry'
@@ -36,3 +37,17 @@ export function readLedgerBounded<T>(
3637
{ label: 'usage ledger', maxAttempts: 2, maxElapsedMs: budgetMs }
3738
)
3839
}
40+
41+
/**
42+
* How long a staff-only workflow run count may take before it is reported as unknown. Well inside
43+
* the admin proxy's request timeout, so a payer too large to count still gets its page.
44+
*/
45+
export const WORKFLOW_RUN_COUNT_BUDGET_MS = 5_000
46+
47+
/** Whether a {@link readLedgerBounded} read failed because it ran out of its budget. */
48+
export function isLedgerReadOverBudget(error: unknown): boolean {
49+
return (
50+
error instanceof DatabaseReadDeadlineError ||
51+
getPostgresCancellationReason(error) === 'statement_timeout'
52+
)
53+
}

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

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ import {
3232
CumulativeUsagePeriodClosedError,
3333
getBillingPeriodUsageCost,
3434
getBillingPeriodUsageCostByUser,
35+
getBillingPeriodWorkflowRunCountWithinBudget,
3536
getStampedPeriodRangeUsageCostByUser,
3637
type RecordCumulativeUsageParams,
3738
recordCumulativeUsage,
@@ -375,6 +376,38 @@ describe('Cumulative billing with PostgreSQL', () => {
375376
)
376377
})
377378

379+
it('reports a workflow run count as unknown once it outlasts the preview budget', async () => {
380+
const entity = { type: 'organization' as const, id: 'org-runs' }
381+
const period = {
382+
start: new Date('2026-08-01T00:00:00.000Z'),
383+
end: new Date('2026-09-01T00:00:00.000Z'),
384+
source: 'reporting' as const,
385+
}
386+
await connection`
387+
insert into usage_log (id, user_id, category, source, description, cost,
388+
billing_entity_type, billing_entity_id, execution_id, created_at)
389+
values ('r1', 'u1', 'model', 'workflow', 'run', 1, 'organization', 'org-runs', 'e1', '2026-08-02'),
390+
('r2', 'u1', 'tool', 'workflow', 'run', 1, 'organization', 'org-runs', 'e1', '2026-08-02'),
391+
('r3', 'u1', 'model', 'workflow', 'run', 1, 'organization', 'org-runs', 'e2', '2026-08-03')
392+
`
393+
expect(await getBillingPeriodWorkflowRunCountWithinBudget(entity, period)).toBe(2)
394+
395+
const release = deferred()
396+
const locked = deferred()
397+
const holder = connection.begin(async (tx) => {
398+
await tx.unsafe('LOCK TABLE usage_log IN ACCESS EXCLUSIVE MODE')
399+
locked.resolve()
400+
await release.promise
401+
})
402+
try {
403+
await locked.promise
404+
expect(await getBillingPeriodWorkflowRunCountWithinBudget(entity, period)).toBeNull()
405+
} finally {
406+
release.resolve()
407+
await holder
408+
}
409+
}, 20_000)
410+
378411
describe('a request that outlives its billing period', () => {
379412
// Past periods: the old period's row is written under the subscription lock only once
380413
// that period has ended.

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

Lines changed: 67 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,11 @@ import {
1818
timestampKey,
1919
} from '@/lib/api/list-query'
2020
import { defaultBillingPeriod } from '@/lib/billing/core/billing-period'
21-
import { readLedgerBounded } from '@/lib/billing/core/ledger-read'
21+
import {
22+
isLedgerReadOverBudget,
23+
readLedgerBounded,
24+
WORKFLOW_RUN_COUNT_BUDGET_MS,
25+
} from '@/lib/billing/core/ledger-read'
2226
import { getHighestPrioritySubscription } from '@/lib/billing/core/plan'
2327
import {
2428
resolveSubscriptionUsagePeriod,
@@ -269,44 +273,74 @@ export async function getBillingPeriodUsageCost(
269273
export async function getBillingPeriodWorkflowRunCount(
270274
billingEntity: BillingEntity,
271275
billingPeriod: UsageQueryPeriod,
272-
executor: DbClient = db
276+
executor: DbClient = db,
277+
budgetMs?: number
273278
): Promise<number> {
274-
const [row] = await readLedgerBounded(executor, (tx) =>
275-
tx
276-
.select({
277-
/**
278-
* The exclusion goes through `notInArray`, not `<> ALL(${array})`. Interpolating
279-
* a JavaScript array into a `sql` template emits parenthesized scalar binds —
280-
* `ALL(($1))` — which Postgres rejects outright with "op ANY/ALL (array)
281-
* requires array on right side". Unit tests cannot catch it, because `@sim/db`
282-
* is mocked and no statement is ever rendered.
283-
*/
284-
workflowRuns:
285-
sql<number>`COUNT(DISTINCT ${usageLog.executionId}) FILTER (WHERE ${usageLog.source} = 'workflow' AND ${notInArray(usageLog.category, [...UNBILLED_USAGE_CATEGORIES])})`.mapWith(
286-
Number
287-
),
288-
})
289-
.from(usageLog)
290-
.where(
291-
and(
292-
eq(usageLog.billingEntityType, billingEntity.type),
293-
eq(usageLog.billingEntityId, billingEntity.id),
294-
...(billingPeriod.source === 'reporting'
295-
? [
296-
gte(usageLog.createdAt, billingPeriod.start),
297-
lt(usageLog.createdAt, billingPeriod.end),
298-
]
299-
: [
300-
eq(usageLog.billingPeriodStart, billingPeriod.start),
301-
eq(usageLog.billingPeriodEnd, billingPeriod.end),
302-
])
303-
)
304-
)
279+
const [row] = await readLedgerBounded(
280+
executor,
281+
(tx) =>
282+
tx
283+
.select({
284+
/**
285+
* The exclusion goes through `notInArray`, not `<> ALL(${array})`. Interpolating
286+
* a JavaScript array into a `sql` template emits parenthesized scalar binds —
287+
* `ALL(($1))` — which Postgres rejects outright with "op ANY/ALL (array)
288+
* requires array on right side". Unit tests cannot catch it, because `@sim/db`
289+
* is mocked and no statement is ever rendered.
290+
*/
291+
workflowRuns:
292+
sql<number>`COUNT(DISTINCT ${usageLog.executionId}) FILTER (WHERE ${usageLog.source} = 'workflow' AND ${notInArray(usageLog.category, [...UNBILLED_USAGE_CATEGORIES])})`.mapWith(
293+
Number
294+
),
295+
})
296+
.from(usageLog)
297+
.where(
298+
and(
299+
eq(usageLog.billingEntityType, billingEntity.type),
300+
eq(usageLog.billingEntityId, billingEntity.id),
301+
...(billingPeriod.source === 'reporting'
302+
? [
303+
gte(usageLog.createdAt, billingPeriod.start),
304+
lt(usageLog.createdAt, billingPeriod.end),
305+
]
306+
: [
307+
eq(usageLog.billingPeriodStart, billingPeriod.start),
308+
eq(usageLog.billingPeriodEnd, billingPeriod.end),
309+
])
310+
)
311+
),
312+
budgetMs
305313
)
306314

307315
return row?.workflowRuns ?? 0
308316
}
309317

318+
/**
319+
* {@link getBillingPeriodWorkflowRunCount} for a staff-only preview that can do without it:
320+
* `null` when the count does not finish within {@link WORKFLOW_RUN_COUNT_BUDGET_MS}, so a payer
321+
* too large to count still gets its preview inside the admin proxy's request timeout.
322+
*/
323+
export async function getBillingPeriodWorkflowRunCountWithinBudget(
324+
billingEntity: BillingEntity,
325+
billingPeriod: UsageQueryPeriod
326+
): Promise<number | null> {
327+
try {
328+
return await getBillingPeriodWorkflowRunCount(
329+
billingEntity,
330+
billingPeriod,
331+
db,
332+
WORKFLOW_RUN_COUNT_BUDGET_MS
333+
)
334+
} catch (error) {
335+
if (!isLedgerReadOverBudget(error)) throw error
336+
logger.warn('Billing period workflow run count exceeded its budget', {
337+
billingEntityType: billingEntity.type,
338+
budgetMs: WORKFLOW_RUN_COUNT_BUDGET_MS,
339+
})
340+
return null
341+
}
342+
}
343+
310344
/**
311345
* Period total plus the portion attributable to `source`, in a single scan.
312346
*

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ import { getBillingConcurrencyLimit } from '@/lib/billing/concurrency-limits'
3737
import { resolveEnterpriseReportingPeriod } from '@/lib/billing/core/reporting-period'
3838
import {
3939
getBillingPeriodUsageCost,
40-
getBillingPeriodWorkflowRunCount,
40+
getBillingPeriodWorkflowRunCountWithinBudget,
4141
} from '@/lib/billing/core/usage-log'
4242
import { creditsToDollars, dollarsToCredits } from '@/lib/billing/credits/conversion'
4343
import {
@@ -462,7 +462,8 @@ export interface EnterpriseIssuancePreflight {
462462
limitDollars: number
463463
usedCredits: number
464464
limitCredits: number
465-
workflowRuns: number
465+
/** `null` when the count did not finish within its budget. */
466+
workflowRuns: number | null
466467
}
467468
invoiceAmountUsd: number
468469
configuredUsageLimitDollars: number
@@ -865,7 +866,7 @@ export async function getEnterpriseIssuancePreflight({
865866
{ type: 'organization', id: membership.organizationId },
866867
reportingPeriod
867868
),
868-
getBillingPeriodWorkflowRunCount(
869+
getBillingPeriodWorkflowRunCountWithinBudget(
869870
{ type: 'organization', id: membership.organizationId },
870871
reportingPeriod
871872
),

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ export const billingUsageLogMockFns = {
100100
mockDeriveBillingContext: vi.fn(),
101101
mockGetBillingPeriodUsageCost: vi.fn(),
102102
mockGetBillingPeriodWorkflowRunCount: vi.fn(),
103+
mockGetBillingPeriodWorkflowRunCountWithinBudget: vi.fn(),
103104
mockGetBillingPeriodUsageCostWithSourceSubset: vi.fn(),
104105
mockGetBillingPeriodUsageCostByUser: vi.fn(),
105106
mockGetStampedPeriodRangeUsageCostByUser: vi.fn(),
@@ -138,6 +139,8 @@ export const billingUsageLogMock = {
138139
deriveBillingContext: billingUsageLogMockFns.mockDeriveBillingContext,
139140
getBillingPeriodUsageCost: billingUsageLogMockFns.mockGetBillingPeriodUsageCost,
140141
getBillingPeriodWorkflowRunCount: billingUsageLogMockFns.mockGetBillingPeriodWorkflowRunCount,
142+
getBillingPeriodWorkflowRunCountWithinBudget:
143+
billingUsageLogMockFns.mockGetBillingPeriodWorkflowRunCountWithinBudget,
141144
getBillingPeriodUsageCostWithSourceSubset:
142145
billingUsageLogMockFns.mockGetBillingPeriodUsageCostWithSourceSubset,
143146
getBillingPeriodUsageCostByUser: billingUsageLogMockFns.mockGetBillingPeriodUsageCostByUser,

0 commit comments

Comments
 (0)