From 08c9641b0e6c3d088d6813330ae99af86560a58a Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 24 Sep 2026 14:20:19 -0700 Subject: [PATCH 1/2] fix(billing): cache enterprise reporting-window usage sums for soft usage gates An enterprise usage period is a reporting window up to a year long, so every soft usage check re-summed the payer's whole year of usage_log, once per billable event (every run admission, Chat request, v1 API response, and knowledge document). For a busy enterprise org that scan ran thousands of times an hour and dominated database CPU. readSoftGateUsageCost serves reporting-window sums from a per-process LRUCache (fetchMethod, max 1000, 30 s TTL) keyed by payer, period source, and window bounds; concurrent misses coalesce and a failed sum is never cached. The raw cached reader is private and typed to reporting periods, so every other period is summed exactly. Pooled org admission (computePooledOrgUsage) and the v1 usage report (getEffectiveCurrentPeriodCost) read through it. The ledger only grows within a window, so a cached sum can only trail the true one by at most one TTL of usage: an admission gate may let a payer run on briefly past its limit, never refuse it wrongly. Exact paths are untouched: getBillingPeriodUsageCost itself (read-your-writes), threshold billing, cycle close, invoicing, analytics, and the execution logger's edge-triggered usage emails, which need an exact baseline. --- .../calculations/usage-monitor.test.ts | 52 +++++ .../lib/billing/calculations/usage-monitor.ts | 10 +- .../core/reporting-usage-cache.test.ts | 197 ++++++++++++++++++ .../lib/billing/core/reporting-usage-cache.ts | 90 ++++++++ apps/sim/lib/billing/core/usage.ts | 6 +- 5 files changed, 352 insertions(+), 3 deletions(-) create mode 100644 apps/sim/lib/billing/core/reporting-usage-cache.test.ts create mode 100644 apps/sim/lib/billing/core/reporting-usage-cache.ts diff --git a/apps/sim/lib/billing/calculations/usage-monitor.test.ts b/apps/sim/lib/billing/calculations/usage-monitor.test.ts index 3fed2e53576..d30208f8fe1 100644 --- a/apps/sim/lib/billing/calculations/usage-monitor.test.ts +++ b/apps/sim/lib/billing/calculations/usage-monitor.test.ts @@ -105,6 +105,58 @@ describe('checkUsageStatus', () => { ) }) + it('shares one pooled sum across admissions in an enterprise reporting window', async () => { + const billingPeriod = { + start: new Date('2026-01-01T00:00:00.000Z'), + end: new Date('2027-01-01T00:00:00.000Z'), + source: 'reporting' as const, + anchorDate: '2026-01-01', + interval: 'year' as const, + } + const subscription = { + referenceId: 'org-reporting-shared', + plan: 'enterprise', + status: 'active', + seats: 1, + periodStart: billingPeriod.start, + periodEnd: billingPeriod.end, + } + const billingContext = { + billingEntity: { type: 'organization' as const, id: 'org-reporting-shared' }, + billingPeriod, + } + + await checkUsageStatus('user-1', subscription, billingContext) + await checkUsageStatus('user-2', subscription, billingContext) + + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledTimes(1) + }) + + it('sums a Stripe-period organization pool exactly on every admission', async () => { + const billingPeriod = { + start: new Date('2026-06-01T00:00:00.000Z'), + end: new Date('2026-07-01T00:00:00.000Z'), + source: 'stripe' as const, + } + const subscription = { + referenceId: 'org-stripe', + plan: 'team', + status: 'active', + seats: 1, + periodStart: null, + periodEnd: null, + } + const billingContext = { + billingEntity: { type: 'organization' as const, id: 'org-stripe' }, + billingPeriod, + } + + await checkUsageStatus('user-1', subscription, billingContext) + await checkUsageStatus('user-1', subscription, billingContext) + + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledTimes(2) + }) + it('reads paid personal ledger usage and refresh from one snapshot', async () => { const periodStart = new Date('2026-06-01T00:00:00.000Z') const periodEnd = new Date('2026-07-01T00:00:00.000Z') diff --git a/apps/sim/lib/billing/calculations/usage-monitor.ts b/apps/sim/lib/billing/calculations/usage-monitor.ts index 2fab6468c14..f4ebe5bc66b 100644 --- a/apps/sim/lib/billing/calculations/usage-monitor.ts +++ b/apps/sim/lib/billing/calculations/usage-monitor.ts @@ -7,6 +7,7 @@ import { isOrganizationBillingBlocked } from '@/lib/billing/core/access' import { defaultBillingPeriod } from '@/lib/billing/core/billing-period' import { getHighestPrioritySubscription } from '@/lib/billing/core/plan' import { resolveSubscriptionUsagePeriod } from '@/lib/billing/core/reporting-period' +import { readSoftGateUsageCost } from '@/lib/billing/core/reporting-usage-cache' import { getUserUsageLimit, type UsageLimitSubscription } from '@/lib/billing/core/usage' import { type BillingContext, @@ -44,6 +45,11 @@ interface UsageData { organizationId: string | null } +/** + * The organization's pooled usage for an admission check. An enterprise reporting window is + * served through {@link readSoftGateUsageCost}; its weekly refresh is zero, so it always takes + * a plain-sum branch below. + */ async function computePooledOrgUsage( organizationId: string, sub: UsageLimitSubscription, @@ -58,12 +64,12 @@ async function computePooledOrgUsage( } if (!isPaid(sub.plan) || !sub.periodStart) { - return getBillingPeriodUsageCost({ type: 'organization', id: organizationId }, billingPeriod) + return readSoftGateUsageCost({ type: 'organization', id: organizationId }, billingPeriod) } const weeklyRefreshDollars = getPlanWeeklyRefreshDollars(sub.plan) if (weeklyRefreshDollars <= 0) { - return getBillingPeriodUsageCost({ type: 'organization', id: organizationId }, billingPeriod) + return readSoftGateUsageCost({ type: 'organization', id: organizationId }, billingPeriod) } const { ledgerUsage, refreshConsumed } = await computeBillingPeriodUsageWithWeeklyRefresh({ diff --git a/apps/sim/lib/billing/core/reporting-usage-cache.test.ts b/apps/sim/lib/billing/core/reporting-usage-cache.test.ts new file mode 100644 index 00000000000..c616a58af5b --- /dev/null +++ b/apps/sim/lib/billing/core/reporting-usage-cache.test.ts @@ -0,0 +1,197 @@ +/** + * @vitest-environment node + */ +import { db } from '@sim/db' +import { sleep } from '@sim/utils/helpers' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockGetBillingPeriodUsageCost } = vi.hoisted(() => ({ + mockGetBillingPeriodUsageCost: vi.fn(), +})) + +vi.mock('@/lib/billing/core/usage-log', () => ({ + getBillingPeriodUsageCost: mockGetBillingPeriodUsageCost, +})) + +import * as reportingUsageCache from '@/lib/billing/core/reporting-usage-cache' +import type { UsageQueryPeriod } from '@/lib/billing/core/usage-log' + +const { REPORTING_USAGE_CACHE_TTL_MS, readSoftGateUsageCost } = reportingUsageCache + +const REPORTING: UsageQueryPeriod = { + start: new Date('2026-01-01T00:00:00.000Z'), + end: new Date('2027-01-01T00:00:00.000Z'), + source: 'reporting', +} +const NEXT_REPORTING: UsageQueryPeriod = { + start: new Date('2027-01-01T00:00:00.000Z'), + end: new Date('2028-01-01T00:00:00.000Z'), + source: 'reporting', +} +const STRIPE: UsageQueryPeriod = { + start: new Date('2026-09-01T00:00:00.000Z'), + end: new Date('2026-10-01T00:00:00.000Z'), + source: 'stripe', +} + +let nextOrg = 0 +/** A fresh payer per test, since the cache is module state shared across tests. */ +function freshOrg() { + nextOrg += 1 + return { type: 'organization' as const, id: `org-${nextOrg}` } +} + +describe('readSoftGateUsageCost on a reporting window', () => { + beforeEach(() => { + vi.clearAllMocks() + mockGetBillingPeriodUsageCost.mockReset() + }) + + afterEach(() => { + vi.restoreAllMocks() + }) + + it('sums the ledger again once a cached sum outlives its TTL', async () => { + const org = freshOrg() + const start = performance.now() + const clock = vi.spyOn(performance, 'now').mockReturnValue(start) + mockGetBillingPeriodUsageCost.mockResolvedValueOnce(10).mockResolvedValueOnce(25) + + await expect(readSoftGateUsageCost(org, REPORTING)).resolves.toBe(10) + /** The cache re-reads its clock only after `ttlResolution` (1 ms) of real time. */ + clock.mockReturnValue(start + REPORTING_USAGE_CACHE_TTL_MS - 1) + await sleep(2) + await expect(readSoftGateUsageCost(org, REPORTING)).resolves.toBe(10) + clock.mockReturnValue(start + REPORTING_USAGE_CACHE_TTL_MS + 1) + await sleep(2) + await expect(readSoftGateUsageCost(org, REPORTING)).resolves.toBe(25) + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledTimes(2) + }) + + it('coalesces concurrent and repeated reads of one window into one ledger sum', async () => { + const org = freshOrg() + let resolveSum: (value: number) => void = () => {} + mockGetBillingPeriodUsageCost.mockReturnValueOnce( + new Promise((resolve) => { + resolveSum = resolve + }) + ) + + const concurrent = Promise.all([ + readSoftGateUsageCost(org, REPORTING), + readSoftGateUsageCost(org, REPORTING), + readSoftGateUsageCost(org, REPORTING), + ]) + resolveSum(42) + + await expect(concurrent).resolves.toEqual([42, 42, 42]) + await expect(readSoftGateUsageCost(org, REPORTING)).resolves.toBe(42) + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledTimes(1) + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledWith(org, REPORTING) + }) + + it('serves a zero sum from cache rather than re-reading it', async () => { + const org = freshOrg() + mockGetBillingPeriodUsageCost.mockResolvedValue(0) + + await expect(readSoftGateUsageCost(org, REPORTING)).resolves.toBe(0) + await expect(readSoftGateUsageCost(org, REPORTING)).resolves.toBe(0) + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledTimes(1) + }) + + it('keeps separate sums for different payers and windows', async () => { + const first = freshOrg() + const second = freshOrg() + mockGetBillingPeriodUsageCost + .mockResolvedValueOnce(1) + .mockResolvedValueOnce(2) + .mockResolvedValueOnce(3) + .mockResolvedValueOnce(4) + + await expect(readSoftGateUsageCost(first, REPORTING)).resolves.toBe(1) + await expect(readSoftGateUsageCost(second, REPORTING)).resolves.toBe(2) + await expect(readSoftGateUsageCost(first, NEXT_REPORTING)).resolves.toBe(3) + await expect(readSoftGateUsageCost({ type: 'user', id: first.id }, REPORTING)).resolves.toBe(4) + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledTimes(4) + }) + + it('surfaces a failed sum to every waiting caller and never caches it', async () => { + const org = freshOrg() + const failure = new Error('canceling statement due to statement timeout') + mockGetBillingPeriodUsageCost.mockRejectedValueOnce(failure).mockResolvedValueOnce(17) + + const results = await Promise.allSettled([ + readSoftGateUsageCost(org, REPORTING), + readSoftGateUsageCost(org, REPORTING), + ]) + expect(results).toEqual([ + { status: 'rejected', reason: failure }, + { status: 'rejected', reason: failure }, + ]) + + await expect(readSoftGateUsageCost(org, REPORTING)).resolves.toBe(17) + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledTimes(2) + }) +}) + +describe('readSoftGateUsageCost', () => { + beforeEach(() => { + vi.clearAllMocks() + mockGetBillingPeriodUsageCost.mockReset() + }) + + it('exposes no cached reader that a non-reporting period could reach', () => { + expect(Object.keys(reportingUsageCache).sort()).toEqual([ + 'REPORTING_USAGE_CACHE_TTL_MS', + 'readSoftGateUsageCost', + ]) + }) + + it('never serves a cached reporting sum to another source with the same bounds', async () => { + const org = freshOrg() + const sameBounds = { start: REPORTING.start, end: REPORTING.end } + mockGetBillingPeriodUsageCost.mockResolvedValueOnce(10).mockResolvedValueOnce(99) + + await expect(readSoftGateUsageCost(org, REPORTING)).resolves.toBe(10) + await expect(readSoftGateUsageCost(org, { ...sameBounds, source: 'stripe' })).resolves.toBe(99) + await expect(readSoftGateUsageCost(org, REPORTING)).resolves.toBe(10) + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledTimes(2) + }) + + it('serves reporting windows from the cache', async () => { + const org = freshOrg() + mockGetBillingPeriodUsageCost.mockResolvedValue(10) + + await readSoftGateUsageCost(org, REPORTING) + await readSoftGateUsageCost(org, REPORTING) + + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledTimes(1) + }) + + it.each([ + ['stripe', STRIPE], + ['default', { ...STRIPE, source: 'default' as const }], + ['unlabelled', { start: STRIPE.start, end: STRIPE.end }], + ])('sums %s periods exactly on every call', async (_label, period) => { + const org = freshOrg() + mockGetBillingPeriodUsageCost.mockResolvedValue(10) + + await readSoftGateUsageCost(org, period) + await readSoftGateUsageCost(org, period) + + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledTimes(2) + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledWith(org, period, undefined, db) + }) + + it('reads a reporting window exactly on a caller-supplied executor', async () => { + const org = freshOrg() + const executor = { transaction: vi.fn() } as unknown as typeof db + mockGetBillingPeriodUsageCost.mockResolvedValue(10) + + await readSoftGateUsageCost(org, REPORTING, executor) + await readSoftGateUsageCost(org, REPORTING, executor) + + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledTimes(2) + expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledWith(org, REPORTING, undefined, executor) + }) +}) diff --git a/apps/sim/lib/billing/core/reporting-usage-cache.ts b/apps/sim/lib/billing/core/reporting-usage-cache.ts new file mode 100644 index 00000000000..2278207c117 --- /dev/null +++ b/apps/sim/lib/billing/core/reporting-usage-cache.ts @@ -0,0 +1,90 @@ +import { db } from '@sim/db' +import { LRUCache } from 'lru-cache' +import { + type BillingEntity, + getBillingPeriodUsageCost, + type UsageQueryPeriod, +} from '@/lib/billing/core/usage-log' +import type { DbClient } from '@/lib/db/types' + +/** + * How long a reporting-window usage sum is served before it is summed again. + * + * A reporting window is an enterprise contract period, up to a year long, so its sum scans every + * ledger row the payer wrote in that year, and the soft gates below re-ran it once per billable + * event. The ledger only grows within a window (rows are inserted at a cost above zero, and the + * one update is a monotonic top-up), so a served sum is never above the true one: it omits at + * most the usage written since it was read. That is the safe direction for every reader here — + * an admission gate lets a payer run on for at most this long past their limit, and a sum at or + * above the limit is a refusal the true sum would also give. Thirty seconds keeps that overrun + * small against a year-long allowance while turning a per-event scan into one per window. + */ +export const REPORTING_USAGE_CACHE_TTL_MS = 30_000 + +/** A usage window known to be an enterprise reporting window — the only kind this cache serves. */ +type ReportingQueryPeriod = UsageQueryPeriod & { source: 'reporting' } + +/** + * Sums shared across callers, one per payer and window. Every key is an enterprise payer's + * current window, a few dozen bytes each, so the ceiling sits far above any process's working + * set and only backstops memory; an eviction inside the TTL costs one extra sum. + * + * `fetchMethod` coalesces concurrent misses onto one sum. A rejected sum is evicted rather than + * stored (`noDeleteOnFetchRejection` and `allowStaleOnFetchRejection` stay off), so every caller + * of that read sees the error it would have seen uncached and the next call sums again. There is + * no settle deadline: the sum runs under the ledger's own `statement_timeout`, so the database + * ends a slow one. There is deliberately no invalidator either — usage is written by execution + * workers in other processes, so the TTL is the real bound. + */ +const reportingUsageCache = new LRUCache< + string, + number, + { entity: BillingEntity; period: ReportingQueryPeriod } +>({ + max: 1_000, + ttl: REPORTING_USAGE_CACHE_TTL_MS, + fetchMethod: (_key, _stale, { context }) => + getBillingPeriodUsageCost(context.entity, context.period), +}) + +/** + * The key names the period's source as well as its bounds, so a sum can only ever be shared + * with a read of the same kind of window, even if another source someday reaches this cache. + */ +function reportingUsageKey(entity: BillingEntity, period: ReportingQueryPeriod): string { + return `${entity.type}:${entity.id}:${period.source}:${period.start.toISOString()}:${period.end.toISOString()}` +} + +function isReportingPeriod(period: UsageQueryPeriod): period is ReportingQueryPeriod { + return period.source === 'reporting' +} + +async function readCachedReportingUsageCost( + entity: BillingEntity, + period: ReportingQueryPeriod +): Promise { + const cost = await reportingUsageCache.fetch(reportingUsageKey(entity, period), { + context: { entity, period }, + }) + return cost !== undefined ? cost : getBillingPeriodUsageCost(entity, period) +} + +/** + * Period usage for a soft reader: an admission check, a display, or a level-triggered + * notification that tolerates the cache's bounded under-count. Enterprise reporting windows are + * served from the shared cache for up to {@link REPORTING_USAGE_CACHE_TTL_MS}, since their + * year-long sum is the expensive one; every other period is summed exactly, as before. A read on + * a caller's own executor (a transaction or a replica) keeps its own snapshot and is never shared. + * Never use it for invoicing, cycle close, an edge-triggered decision, or a read that must see its + * own write. + */ +export function readSoftGateUsageCost( + entity: BillingEntity, + period: UsageQueryPeriod, + executor: DbClient = db +): Promise { + if (isReportingPeriod(period) && executor === db) { + return readCachedReportingUsageCost(entity, period) + } + return getBillingPeriodUsageCost(entity, period, undefined, executor) +} diff --git a/apps/sim/lib/billing/core/usage.ts b/apps/sim/lib/billing/core/usage.ts index b62fe0c8bd9..dbb35478073 100644 --- a/apps/sim/lib/billing/core/usage.ts +++ b/apps/sim/lib/billing/core/usage.ts @@ -15,6 +15,7 @@ import { type ResolvedUsagePeriod, resolveSubscriptionUsagePeriod, } from '@/lib/billing/core/reporting-period' +import { readSoftGateUsageCost } from '@/lib/billing/core/reporting-usage-cache' import { type BillingEntity, getBillingPeriodUsageCost } from '@/lib/billing/core/usage-log' import { computeWeeklyRefreshConsumed } from '@/lib/billing/credits/weekly-refresh' import { @@ -639,6 +640,9 @@ export async function syncUsageLimitsFromSubscription(userId: string): Promise Date: Thu, 24 Sep 2026 14:29:36 -0700 Subject: [PATCH 2/2] test(billing): use 1 ms waits in the reporting usage cache expiry test --- apps/sim/lib/billing/core/reporting-usage-cache.test.ts | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/apps/sim/lib/billing/core/reporting-usage-cache.test.ts b/apps/sim/lib/billing/core/reporting-usage-cache.test.ts index c616a58af5b..74f556ace0a 100644 --- a/apps/sim/lib/billing/core/reporting-usage-cache.test.ts +++ b/apps/sim/lib/billing/core/reporting-usage-cache.test.ts @@ -58,12 +58,15 @@ describe('readSoftGateUsageCost on a reporting window', () => { mockGetBillingPeriodUsageCost.mockResolvedValueOnce(10).mockResolvedValueOnce(25) await expect(readSoftGateUsageCost(org, REPORTING)).resolves.toBe(10) - /** The cache re-reads its clock only after `ttlResolution` (1 ms) of real time. */ + /** + * The cache refreshes its clock on a 1 ms `ttlResolution` timer scheduled when it last read + * the time, so it always fires before a 1 ms sleep queued afterwards. + */ clock.mockReturnValue(start + REPORTING_USAGE_CACHE_TTL_MS - 1) - await sleep(2) + await sleep(1) await expect(readSoftGateUsageCost(org, REPORTING)).resolves.toBe(10) clock.mockReturnValue(start + REPORTING_USAGE_CACHE_TTL_MS + 1) - await sleep(2) + await sleep(1) await expect(readSoftGateUsageCost(org, REPORTING)).resolves.toBe(25) expect(mockGetBillingPeriodUsageCost).toHaveBeenCalledTimes(2) })