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..74f556ace0a --- /dev/null +++ b/apps/sim/lib/billing/core/reporting-usage-cache.test.ts @@ -0,0 +1,200 @@ +/** + * @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 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(1) + await expect(readSoftGateUsageCost(org, REPORTING)).resolves.toBe(10) + clock.mockReturnValue(start + REPORTING_USAGE_CACHE_TTL_MS + 1) + await sleep(1) + 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