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,