From a97389d69f03ba3db078e4936642621ba8204533 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 4 Nov 2026 02:53:42 -0800 Subject: [PATCH] fix(billing): bound the enterprise preflight's workflow run count to the preview budget 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. --- apps/sim/lib/admin/dashboard-usage.ts | 27 ++--- apps/sim/lib/billing/core/ledger-read.ts | 15 +++ .../lib/billing/core/usage-log.integration.ts | 33 ++++++ apps/sim/lib/billing/core/usage-log.ts | 100 ++++++++++++------ .../lib/billing/enterprise-provisioning.ts | 7 +- .../src/mocks/billing-usage-log.mock.ts | 3 + 6 files changed, 132 insertions(+), 53 deletions(-) diff --git a/apps/sim/lib/admin/dashboard-usage.ts b/apps/sim/lib/admin/dashboard-usage.ts index f76a77a062d..95d02d1491f 100644 --- a/apps/sim/lib/admin/dashboard-usage.ts +++ b/apps/sim/lib/admin/dashboard-usage.ts @@ -1,14 +1,16 @@ import { db } from '@sim/db' import { member, permissions, usageLog, workspace } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { getPostgresCancellationReason } from '@sim/utils/errors' import { and, countDistinct, eq, inArray, notInArray, type SQL, sql } from 'drizzle-orm' import { unionAll } from 'drizzle-orm/pg-core' -import { readLedgerBounded } from '@/lib/billing/core/ledger-read' +import { + isLedgerReadOverBudget, + readLedgerBounded, + WORKFLOW_RUN_COUNT_BUDGET_MS, +} from '@/lib/billing/core/ledger-read' import type { ResolvedUsagePeriod } from '@/lib/billing/core/reporting-period' import { buildUsageAnalyticsScope } from '@/lib/billing/core/usage-analytics' import type { BillingEntity } from '@/lib/billing/core/usage-log' -import { DatabaseReadDeadlineError } from '@/lib/db/read-retry' /** * 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' * Widening both entity indexes with it would grow the hot billing sums' working set to * speed up a staff-only page. So only one organization's detail counts runs, in a single * scan that yields the total and every actor's share, under - * {@link WORKFLOW_RUNS_BUDGET_MS}; past it the count is unknown (`null`) and the page + * {@link WORKFLOW_RUN_COUNT_BUDGET_MS}; past it the count is unknown (`null`) and the page * still loads. * * They stay on the primary, where the hot billing sums keep those index pages cached. On @@ -54,12 +56,6 @@ interface DashboardWorkflowRuns { formerActors: number } -/** - * How long one detail page may spend counting workflow runs. Well inside the admin - * proxy's request timeout, so a payer too large to count still gets its page. - */ -const WORKFLOW_RUNS_BUDGET_MS = 5_000 - const WORKFLOW_SOURCE = sql`${usageLog.source} = 'workflow'` export function organizationWindow( @@ -164,7 +160,7 @@ function isFormerActor(organizationId: string): SQL { /** * An organization's workflow runs over its period, in total, per acting user and across its * former actors, from one scan. `null` when the count does not finish within - * {@link WORKFLOW_RUNS_BUDGET_MS}. + * {@link WORKFLOW_RUN_COUNT_BUDGET_MS}. */ export async function readDashboardOrganizationWorkflowRuns( organizationId: string, @@ -193,16 +189,13 @@ export async function readDashboardOrganizationWorkflowRuns( .from(workflowRows) .groupBy(sql`grouping sets ((${workflowRows.userId}), (${workflowRows.formerActor}), ())`) }, - WORKFLOW_RUNS_BUDGET_MS + WORKFLOW_RUN_COUNT_BUDGET_MS ) } catch (error) { - const overBudget = - error instanceof DatabaseReadDeadlineError || - getPostgresCancellationReason(error) === 'statement_timeout' - if (!overBudget) throw error + if (!isLedgerReadOverBudget(error)) throw error logger.warn('Workflow run count exceeded its budget', { organizationId, - budgetMs: WORKFLOW_RUNS_BUDGET_MS, + budgetMs: WORKFLOW_RUN_COUNT_BUDGET_MS, }) return null } diff --git a/apps/sim/lib/billing/core/ledger-read.ts b/apps/sim/lib/billing/core/ledger-read.ts index 78ed1cd5a73..49e28ffd7bb 100644 --- a/apps/sim/lib/billing/core/ledger-read.ts +++ b/apps/sim/lib/billing/core/ledger-read.ts @@ -1,3 +1,4 @@ +import { getPostgresCancellationReason } from '@sim/utils/errors' import { sql } from 'drizzle-orm' import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants' import { DatabaseReadDeadlineError, withDatabaseReadRetry } from '@/lib/db/read-retry' @@ -36,3 +37,17 @@ export function readLedgerBounded( { label: 'usage ledger', maxAttempts: 2, maxElapsedMs: budgetMs } ) } + +/** + * How long a staff-only workflow run count may take before it is reported as unknown. Well inside + * the admin proxy's request timeout, so a payer too large to count still gets its page. + */ +export const WORKFLOW_RUN_COUNT_BUDGET_MS = 5_000 + +/** Whether a {@link readLedgerBounded} read failed because it ran out of its budget. */ +export function isLedgerReadOverBudget(error: unknown): boolean { + return ( + error instanceof DatabaseReadDeadlineError || + getPostgresCancellationReason(error) === 'statement_timeout' + ) +} diff --git a/apps/sim/lib/billing/core/usage-log.integration.ts b/apps/sim/lib/billing/core/usage-log.integration.ts index 4f91dad0d5b..b3bf624f002 100644 --- a/apps/sim/lib/billing/core/usage-log.integration.ts +++ b/apps/sim/lib/billing/core/usage-log.integration.ts @@ -32,6 +32,7 @@ import { CumulativeUsagePeriodClosedError, getBillingPeriodUsageCost, getBillingPeriodUsageCostByUser, + getBillingPeriodWorkflowRunCountWithinBudget, getStampedPeriodRangeUsageCostByUser, type RecordCumulativeUsageParams, recordCumulativeUsage, @@ -375,6 +376,38 @@ describe('Cumulative billing with PostgreSQL', () => { ) }) + it('reports a workflow run count as unknown once it outlasts the preview budget', async () => { + const entity = { type: 'organization' as const, id: 'org-runs' } + const period = { + start: new Date('2026-08-01T00:00:00.000Z'), + end: new Date('2026-09-01T00:00:00.000Z'), + source: 'reporting' as const, + } + await connection` + insert into usage_log (id, user_id, category, source, description, cost, + billing_entity_type, billing_entity_id, execution_id, created_at) + values ('r1', 'u1', 'model', 'workflow', 'run', 1, 'organization', 'org-runs', 'e1', '2026-08-02'), + ('r2', 'u1', 'tool', 'workflow', 'run', 1, 'organization', 'org-runs', 'e1', '2026-08-02'), + ('r3', 'u1', 'model', 'workflow', 'run', 1, 'organization', 'org-runs', 'e2', '2026-08-03') + ` + expect(await getBillingPeriodWorkflowRunCountWithinBudget(entity, period)).toBe(2) + + const release = deferred() + const locked = deferred() + const holder = connection.begin(async (tx) => { + await tx.unsafe('LOCK TABLE usage_log IN ACCESS EXCLUSIVE MODE') + locked.resolve() + await release.promise + }) + try { + await locked.promise + expect(await getBillingPeriodWorkflowRunCountWithinBudget(entity, period)).toBeNull() + } finally { + release.resolve() + await holder + } + }, 20_000) + describe('a request that outlives its billing period', () => { // Past periods: the old period's row is written under the subscription lock only once // that period has ended. diff --git a/apps/sim/lib/billing/core/usage-log.ts b/apps/sim/lib/billing/core/usage-log.ts index 837ab470c00..57dae19ad2d 100644 --- a/apps/sim/lib/billing/core/usage-log.ts +++ b/apps/sim/lib/billing/core/usage-log.ts @@ -18,7 +18,11 @@ import { timestampKey, } from '@/lib/api/list-query' import { defaultBillingPeriod } from '@/lib/billing/core/billing-period' -import { readLedgerBounded } from '@/lib/billing/core/ledger-read' +import { + isLedgerReadOverBudget, + readLedgerBounded, + WORKFLOW_RUN_COUNT_BUDGET_MS, +} from '@/lib/billing/core/ledger-read' import { getHighestPrioritySubscription } from '@/lib/billing/core/plan' import { resolveSubscriptionUsagePeriod, @@ -269,44 +273,74 @@ export async function getBillingPeriodUsageCost( export async function getBillingPeriodWorkflowRunCount( billingEntity: BillingEntity, billingPeriod: UsageQueryPeriod, - executor: DbClient = db + executor: DbClient = db, + budgetMs?: number ): Promise { - 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`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`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), + ]) + ) + ), + budgetMs ) return row?.workflowRuns ?? 0 } +/** + * {@link getBillingPeriodWorkflowRunCount} for a staff-only preview that can do without it: + * `null` when the count does not finish within {@link WORKFLOW_RUN_COUNT_BUDGET_MS}, so a payer + * too large to count still gets its preview inside the admin proxy's request timeout. + */ +export async function getBillingPeriodWorkflowRunCountWithinBudget( + billingEntity: BillingEntity, + billingPeriod: UsageQueryPeriod +): Promise { + try { + return await getBillingPeriodWorkflowRunCount( + billingEntity, + billingPeriod, + db, + WORKFLOW_RUN_COUNT_BUDGET_MS + ) + } catch (error) { + if (!isLedgerReadOverBudget(error)) throw error + logger.warn('Billing period workflow run count exceeded its budget', { + billingEntityType: billingEntity.type, + budgetMs: WORKFLOW_RUN_COUNT_BUDGET_MS, + }) + return null + } +} + /** * Period total plus the portion attributable to `source`, in a single scan. * diff --git a/apps/sim/lib/billing/enterprise-provisioning.ts b/apps/sim/lib/billing/enterprise-provisioning.ts index 93620a470d9..439a2890de8 100644 --- a/apps/sim/lib/billing/enterprise-provisioning.ts +++ b/apps/sim/lib/billing/enterprise-provisioning.ts @@ -37,7 +37,7 @@ import { getBillingConcurrencyLimit } from '@/lib/billing/concurrency-limits' import { resolveEnterpriseReportingPeriod } from '@/lib/billing/core/reporting-period' import { getBillingPeriodUsageCost, - getBillingPeriodWorkflowRunCount, + getBillingPeriodWorkflowRunCountWithinBudget, } from '@/lib/billing/core/usage-log' import { creditsToDollars, dollarsToCredits } from '@/lib/billing/credits/conversion' import { @@ -462,7 +462,8 @@ export interface EnterpriseIssuancePreflight { limitDollars: number usedCredits: number limitCredits: number - workflowRuns: number + /** `null` when the count did not finish within its budget. */ + workflowRuns: number | null } invoiceAmountUsd: number configuredUsageLimitDollars: number @@ -865,7 +866,7 @@ export async function getEnterpriseIssuancePreflight({ { type: 'organization', id: membership.organizationId }, reportingPeriod ), - getBillingPeriodWorkflowRunCount( + getBillingPeriodWorkflowRunCountWithinBudget( { type: 'organization', id: membership.organizationId }, reportingPeriod ), diff --git a/packages/testing/src/mocks/billing-usage-log.mock.ts b/packages/testing/src/mocks/billing-usage-log.mock.ts index ff7def0f720..f3beace60bd 100644 --- a/packages/testing/src/mocks/billing-usage-log.mock.ts +++ b/packages/testing/src/mocks/billing-usage-log.mock.ts @@ -100,6 +100,7 @@ export const billingUsageLogMockFns = { mockDeriveBillingContext: vi.fn(), mockGetBillingPeriodUsageCost: vi.fn(), mockGetBillingPeriodWorkflowRunCount: vi.fn(), + mockGetBillingPeriodWorkflowRunCountWithinBudget: vi.fn(), mockGetBillingPeriodUsageCostWithSourceSubset: vi.fn(), mockGetBillingPeriodUsageCostByUser: vi.fn(), mockGetStampedPeriodRangeUsageCostByUser: vi.fn(), @@ -138,6 +139,8 @@ export const billingUsageLogMock = { deriveBillingContext: billingUsageLogMockFns.mockDeriveBillingContext, getBillingPeriodUsageCost: billingUsageLogMockFns.mockGetBillingPeriodUsageCost, getBillingPeriodWorkflowRunCount: billingUsageLogMockFns.mockGetBillingPeriodWorkflowRunCount, + getBillingPeriodWorkflowRunCountWithinBudget: + billingUsageLogMockFns.mockGetBillingPeriodWorkflowRunCountWithinBudget, getBillingPeriodUsageCostWithSourceSubset: billingUsageLogMockFns.mockGetBillingPeriodUsageCostWithSourceSubset, getBillingPeriodUsageCostByUser: billingUsageLogMockFns.mockGetBillingPeriodUsageCostByUser,