diff --git a/apps/sim/lib/billing/core/ingestion-usage-gate.test.ts b/apps/sim/lib/billing/core/ingestion-usage-gate.test.ts new file mode 100644 index 00000000000..dbecf11abf0 --- /dev/null +++ b/apps/sim/lib/billing/core/ingestion-usage-gate.test.ts @@ -0,0 +1,98 @@ +/** + * @vitest-environment node + */ +import { sleep } from '@sim/utils/helpers' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockCheck } = vi.hoisted(() => ({ mockCheck: vi.fn() })) + +vi.mock('@/lib/billing/core/billing-attribution', () => ({ + checkAttributedUsageLimits: mockCheck, +})) + +import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution' +import { + checkIngestionUsageLimits, + INGESTION_USAGE_GATE_TTL_MS, + resetIngestionUsageGateCache, +} from '@/lib/billing/core/ingestion-usage-gate' + +const ATTRIBUTION: BillingAttributionSnapshot = { + actorUserId: 'member-1', + workspaceId: null, + organizationId: 'org-1', + billedAccountUserId: 'owner-1', + billingEntity: { type: 'organization', id: 'org-1' }, + billingPeriod: { + start: '2026-09-01T00:00:00.000Z', + end: '2026-10-01T00:00:00.000Z', + source: 'stripe', + }, + payerSubscription: null, +} + +describe('checkIngestionUsageLimits', () => { + beforeEach(() => { + resetIngestionUsageGateCache() + mockCheck.mockReset().mockResolvedValue({ isExceeded: false }) + }) + afterEach(() => vi.restoreAllMocks()) + + it('reads the ledger once per payer, period and actor within the TTL', async () => { + await checkIngestionUsageLimits(ATTRIBUTION) + await checkIngestionUsageLimits(ATTRIBUTION) + await checkIngestionUsageLimits({ ...ATTRIBUTION, workspaceId: 'workspace-9' }) + expect(mockCheck).toHaveBeenCalledTimes(1) + }) + + it('collapses concurrent misses onto one ledger read', async () => { + let resolve!: (value: { isExceeded: boolean }) => void + mockCheck.mockReturnValueOnce(new Promise((r) => (resolve = r))) + const pending = Promise.all([ + checkIngestionUsageLimits(ATTRIBUTION), + checkIngestionUsageLimits(ATTRIBUTION), + checkIngestionUsageLimits(ATTRIBUTION), + ]) + await sleep(0) + resolve({ isExceeded: false }) + const results = await pending + expect(results.every((result) => result.isExceeded === false)).toBe(true) + expect(mockCheck).toHaveBeenCalledTimes(1) + }) + + it('keeps a refusal for the same bounded window, then reads fresh', async () => { + mockCheck.mockResolvedValueOnce({ isExceeded: true, scope: 'payer', message: 'over' }) + expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true) + expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true) + expect(mockCheck).toHaveBeenCalledTimes(1) + + /** `lru-cache` reads `performance.now()` and debounces it behind a real 1 ms timer. */ + const start = performance.now() + vi.spyOn(performance, 'now').mockReturnValue(start + INGESTION_USAGE_GATE_TTL_MS + 1) + await sleep(5) + expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false) + expect(mockCheck).toHaveBeenCalledTimes(2) + }) + + it('separates answers by actor, period and payer', async () => { + await checkIngestionUsageLimits(ATTRIBUTION) + await checkIngestionUsageLimits({ ...ATTRIBUTION, actorUserId: 'member-2' }) + await checkIngestionUsageLimits({ + ...ATTRIBUTION, + billingPeriod: { ...ATTRIBUTION.billingPeriod, start: '2026-10-01T00:00:00.000Z' }, + }) + await checkIngestionUsageLimits({ + ...ATTRIBUTION, + billedAccountUserId: 'owner-2', + billingEntity: { type: 'user', id: 'owner-2' }, + }) + expect(mockCheck).toHaveBeenCalledTimes(4) + }) + + it('does not cache a failed read', async () => { + mockCheck.mockRejectedValueOnce(new Error('ledger unavailable')) + await expect(checkIngestionUsageLimits(ATTRIBUTION)).rejects.toThrow('ledger unavailable') + await checkIngestionUsageLimits(ATTRIBUTION) + expect(mockCheck).toHaveBeenCalledTimes(2) + }) +}) diff --git a/apps/sim/lib/billing/core/ingestion-usage-gate.ts b/apps/sim/lib/billing/core/ingestion-usage-gate.ts new file mode 100644 index 00000000000..2f6f1c5f5f8 --- /dev/null +++ b/apps/sim/lib/billing/core/ingestion-usage-gate.ts @@ -0,0 +1,69 @@ +import { LRUCache } from 'lru-cache' +import { + type AttributedUsageLimitsResult, + type BillingAttributionSnapshot, + checkAttributedUsageLimits, +} from '@/lib/billing/core/billing-attribution' +import { coalesceLocally } from '@/lib/concurrency/singleflight' + +/** + * How long a usage-gate answer stays usable on the ingestion path. + * + * The gate sums the payer's usage ledger for the billing period, which grows + * with every indexed document, so a bulk sync that re-checks per document + * reads the whole period's ledger tens of thousands of times. Staleness fails + * in the harmless direction: a payer at their limit keeps indexing for at most + * this long, and a payer whose limit was just raised waits at most this long. + * Nothing on this path has a person waiting for the answer. + */ +export const INGESTION_USAGE_GATE_TTL_MS = 60 * 1000 + +/** Recent gate answers, with `LRUCache` supplying the TTL and the size bound. */ +const gateCache = new LRUCache({ + max: 10_000, + ttl: INGESTION_USAGE_GATE_TTL_MS, +}) + +/** + * The gate depends on who pays, for which period, and which member acts: the + * payer pool and the per-member cap are both part of the answer. + */ +function gateKey(attribution: BillingAttributionSnapshot): string { + return [ + attribution.billingEntity.type, + attribution.billingEntity.id, + attribution.billingPeriod.start, + attribution.billingPeriod.end, + attribution.billedAccountUserId, + attribution.actorUserId, + ].join(':') +} + +/** + * {@link checkAttributedUsageLimits} for background ingestion, with bounded + * staleness. Interactive callers (uploads, search, the settings surfaces) keep + * reading the gate fresh so a limit change is visible at once. + * + * `coalesceLocally` collapses the concurrent misses of a batch onto one ledger + * read and bounds a hung read at its settle deadline. The cache write stays on + * the value this caller received, so a producer that timed out and later + * resolved cannot overwrite a fresher answer. + */ +export async function checkIngestionUsageLimits( + attribution: BillingAttributionSnapshot +): Promise { + const key = gateKey(attribution) + const cached = gateCache.get(key) + if (cached !== undefined) return cached + + const result = await coalesceLocally(`ingestion-usage-gate:${key}`, () => + checkAttributedUsageLimits(attribution) + ) + gateCache.set(key, result) + return result +} + +/** Drops every cached gate answer. Test seam; never called in production code. */ +export function resetIngestionUsageGateCache(): void { + gateCache.clear() +} diff --git a/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts b/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts index e68b00db64f..437de7bf205 100644 --- a/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts +++ b/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts @@ -65,11 +65,13 @@ vi.mock('@/providers/utils', () => ({ })) import * as billingAttribution from '@/lib/billing/core/billing-attribution' +import { resetIngestionUsageGateCache } from '@/lib/billing/core/ingestion-usage-gate' import * as embeddingClient from '@/lib/embeddings/client' import { processDocumentAsync } from '@/lib/knowledge/documents/service' const mockEmbeddingCapacity = vi.fn() beforeEach(() => { + resetIngestionUsageGateCache() vi.spyOn(billingAttribution, 'checkAttributedUsageLimits').mockImplementation( mockCheckAttributedUsageLimits ) diff --git a/apps/sim/lib/knowledge/documents/document-processing-source.test.ts b/apps/sim/lib/knowledge/documents/document-processing-source.test.ts index c3eeff755ca..5b36cbdf3c2 100644 --- a/apps/sim/lib/knowledge/documents/document-processing-source.test.ts +++ b/apps/sim/lib/knowledge/documents/document-processing-source.test.ts @@ -66,6 +66,7 @@ vi.mock('@/lib/uploads/server/metadata', () => ({ })) import * as billingAttribution from '@/lib/billing/core/billing-attribution' +import { resetIngestionUsageGateCache } from '@/lib/billing/core/ingestion-usage-gate' import { env } from '@/lib/core/config/env' import { markInsideTriggerRun, @@ -92,6 +93,7 @@ import { MAX_PROCESSING_ATTEMPTS } from '@/lib/knowledge/documents/types' const mockEmbeddingCapacity = vi.fn() beforeEach(() => { + resetIngestionUsageGateCache() vi.spyOn(billingAttribution, 'checkAttributedUsageLimits').mockImplementation( mockCheckAttributedUsageLimits ) diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts index 8a555fb1834..7b962342a17 100644 --- a/apps/sim/lib/knowledge/documents/service.ts +++ b/apps/sim/lib/knowledge/documents/service.ts @@ -34,9 +34,9 @@ import { assertBillingAttributionOwner, assertBillingAttributionSnapshot, type BillingAttributionSnapshot, - checkAttributedUsageLimits, toBillingContext, } from '@/lib/billing/core/billing-attribution' +import { checkIngestionUsageLimits } from '@/lib/billing/core/ingestion-usage-gate' import { recordUsage } from '@/lib/billing/core/usage-log' import { applyStorageUsageDeltasInTx, @@ -1618,7 +1618,7 @@ export async function processDocumentAsync( assertBillingAttributionOwner(billingAttribution, ctx) const documentActorUserId = billingAttribution.actorUserId - const usageGate = await checkAttributedUsageLimits(billingAttribution) + const usageGate = await checkIngestionUsageLimits(billingAttribution) if (usageGate.isExceeded) { logger.warn(`[${documentId}] Usage limit reached — skipping document indexing`) throw new UsageLimitDocumentProcessingError(