improvement(knowledge): cache admitted usage checks on the search path

Author: waleedlatif1

apps/sim/app/api/v2/knowledge/search/route.provenance.test.ts

@@ -31,7 +31,10 @@ vi.mock('@sim/platform-authz/workspace', () => ({
 vi.mock('@/lib/billing/core/billing-attribution', () => ({
   resolveBillingAttribution: mocks.resolveBilling,
   resolveSystemBillingAttribution: mocks.resolveBilling,
-  checkAttributedUsageLimits: mocks.checkUsage,
+}))
+
+vi.mock('@/lib/billing/core/usage-gate-cache', () => ({
+  checkSearchUsageLimits: mocks.checkUsage,
 }))
 
 /** Retrieval defaults are the flag's concern; here the flag is off so the search stays as configured. */

apps/sim/lib/billing/core/ingestion-usage-gate.ts

@@ -13,9 +13,10 @@ vi.mock('@/lib/billing/core/billing-attribution', () => ({
 import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution'
 import {
   checkIngestionUsageLimits,
-  INGESTION_USAGE_GATE_TTL_MS,
-  resetIngestionUsageGateCache,
-} from '@/lib/billing/core/ingestion-usage-gate'
+  checkSearchUsageLimits,
+  resetUsageGateCache,
+  USAGE_GATE_TTL_MS,
+} from '@/lib/billing/core/usage-gate-cache'
 
 const ATTRIBUTION: BillingAttributionSnapshot = {
   actorUserId: 'member-1',
@@ -31,9 +32,19 @@ const ATTRIBUTION: BillingAttributionSnapshot = {
   payerSubscription: null,
 }
 
+const SUBSCRIPTION: BillingAttributionSnapshot['payerSubscription'] = {
+  id: 'sub-1',
+  referenceId: 'org-1',
+  plan: 'team',
+  status: 'active',
+  seats: 5,
+  periodStart: '2026-09-01T00:00:00.000Z',
+  periodEnd: '2026-10-01T00:00:00.000Z',
+}
+
 describe('checkIngestionUsageLimits', () => {
   beforeEach(() => {
-    resetIngestionUsageGateCache()
+    resetUsageGateCache()
     mockCheck.mockReset().mockResolvedValue({ isExceeded: false })
   })
   afterEach(() => vi.restoreAllMocks())
@@ -68,13 +79,13 @@ describe('checkIngestionUsageLimits', () => {
 
     /** `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)
+    vi.spyOn(performance, 'now').mockReturnValue(start + 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 () => {
+  it('separates answers by actor, period, payer and plan', async () => {
     await checkIngestionUsageLimits(ATTRIBUTION)
     await checkIngestionUsageLimits({ ...ATTRIBUTION, actorUserId: 'member-2' })
     await checkIngestionUsageLimits({
@@ -86,7 +97,12 @@ describe('checkIngestionUsageLimits', () => {
       billedAccountUserId: 'owner-2',
       billingEntity: { type: 'user', id: 'owner-2' },
     })
-    expect(mockCheck).toHaveBeenCalledTimes(4)
+    await checkIngestionUsageLimits({ ...ATTRIBUTION, payerSubscription: SUBSCRIPTION })
+    await checkIngestionUsageLimits({
+      ...ATTRIBUTION,
+      payerSubscription: { ...SUBSCRIPTION, plan: 'enterprise' },
+    })
+    expect(mockCheck).toHaveBeenCalledTimes(6)
   })
 
   it('does not cache a failed read', async () => {
@@ -96,3 +112,38 @@ describe('checkIngestionUsageLimits', () => {
     expect(mockCheck).toHaveBeenCalledTimes(2)
   })
 })
+
+describe('checkSearchUsageLimits', () => {
+  beforeEach(() => {
+    resetUsageGateCache()
+    mockCheck.mockReset().mockResolvedValue({ isExceeded: false })
+  })
+
+  it('reuses an admission across workspaces of the same payer', async () => {
+    await checkSearchUsageLimits(ATTRIBUTION)
+    await checkSearchUsageLimits({ ...ATTRIBUTION, workspaceId: 'workspace-9' })
+    expect(mockCheck).toHaveBeenCalledTimes(1)
+  })
+
+  it('re-reads a refusal, so a raised limit applies on the next search', async () => {
+    mockCheck.mockResolvedValueOnce({ isExceeded: true, scope: 'payer', message: 'over' })
+    expect((await checkSearchUsageLimits(ATTRIBUTION)).isExceeded).toBe(true)
+    expect((await checkSearchUsageLimits(ATTRIBUTION)).isExceeded).toBe(false)
+    expect(mockCheck).toHaveBeenCalledTimes(2)
+  })
+
+  it('does not serve a refusal cached by ingestion', async () => {
+    mockCheck.mockResolvedValueOnce({ isExceeded: true, scope: 'payer', message: 'over' })
+    expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true)
+    expect((await checkSearchUsageLimits(ATTRIBUTION)).isExceeded).toBe(false)
+    expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false)
+    expect(mockCheck).toHaveBeenCalledTimes(2)
+  })
+
+  it('does not cache a failed read', async () => {
+    mockCheck.mockRejectedValueOnce(new Error('ledger unavailable'))
+    await expect(checkSearchUsageLimits(ATTRIBUTION)).rejects.toThrow('ledger unavailable')
+    await checkSearchUsageLimits(ATTRIBUTION)
+    expect(mockCheck).toHaveBeenCalledTimes(2)
+  })
+})

…illing/core/ingestion-usage-gate.test.ts …ib/billing/core/usage-gate-cache.test.tsapps/sim/lib/billing/core/ingestion-usage-gate.test.ts renamed to apps/sim/lib/billing/core/usage-gate-cache.test.ts apps/sim/lib/billing/core/ingestion-usage-gate.test.ts renamed to apps/sim/lib/billing/core/usage-gate-cache.test.ts

@@ -0,0 +1,106 @@
+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 high-frequency paths.
+ *
+ * The gate sums the payer's usage ledger for the billing period, which grows
+ * with the payer's activity, so a large organization scans its whole period on
+ * every uncached call. Bulk ingestion re-checks per document and knowledge
+ * search checks per query. Staleness is bounded by this TTL and fails in the
+ * harmless direction: a payer who crosses their limit keeps going for at most
+ * this long, which charges nobody wrongly.
+ */
+export const USAGE_GATE_TTL_MS = 60 * 1000
+
+/**
+ * Recent gate answers, admitted and refused, with `LRUCache` supplying the TTL
+ * and the size bound. Each entry point decides which of them it may serve.
+ */
+const gateCache = new LRUCache<string, AttributedUsageLimitsResult>({
+  max: 10_000,
+  ttl: USAGE_GATE_TTL_MS,
+})
+
+/**
+ * The gate depends on who pays, for which period, under which plan, and which
+ * member acts: the payer pool, its limit and the per-member cap are all part of
+ * the answer. The workspace is not, so every workspace of one payer shares an
+ * entry.
+ */
+function gateKey(attribution: BillingAttributionSnapshot): string {
+  const subscription = attribution.payerSubscription
+  return [
+    attribution.billingEntity.type,
+    attribution.billingEntity.id,
+    attribution.billingPeriod.start,
+    attribution.billingPeriod.end,
+    attribution.billedAccountUserId,
+    attribution.actorUserId,
+    subscription?.id ?? '',
+    subscription?.plan ?? '',
+    subscription?.status ?? '',
+    subscription?.seats ?? '',
+  ].join(':')
+}
+
+/**
+ * Serves a cached answer the caller accepts, otherwise reads the gate.
+ *
+ * `coalesceLocally` collapses concurrent misses onto one ledger read and bounds
+ * a hung read at its settle deadline. A failed read throws without writing, so
+ * an outage is never recorded as an answer. The write stays on the value this
+ * caller received, so a producer that timed out and later resolved cannot
+ * overwrite a fresher answer.
+ *
+ * There is deliberately no invalidator: usage and limit changes land in other
+ * processes (execution workers, Stripe webhooks), so the TTL is the real bound.
+ */
+async function checkUsageLimitsThroughCache(
+  attribution: BillingAttributionSnapshot,
+  serveCachedRefusal: boolean
+): Promise<AttributedUsageLimitsResult> {
+  const key = gateKey(attribution)
+  const cached = gateCache.get(key)
+  if (cached !== undefined && (serveCachedRefusal || !cached.isExceeded)) return cached
+
+  const result = await coalesceLocally(`usage-gate:${key}`, () =>
+    checkAttributedUsageLimits(attribution)
+  )
+  gateCache.set(key, result)
+  return result
+}
+
+/**
+ * {@link checkAttributedUsageLimits} for background ingestion. Serves admitted
+ * and refused answers alike: nothing on this path has a person waiting for a
+ * raised limit to apply, so a refused payer waits at most the TTL.
+ */
+export function checkIngestionUsageLimits(
+  attribution: BillingAttributionSnapshot
+): Promise<AttributedUsageLimitsResult> {
+  return checkUsageLimitsThroughCache(attribution, true)
+}
+
+/**
+ * {@link checkAttributedUsageLimits} for knowledge search. Serves only a cached
+ * admission: a refusal is always re-read, so a payer who just raised their limit
+ * or upgraded is never held behind a cached block while they wait on a search.
+ * Every other interactive caller (uploads, execution admission, the settings
+ * surfaces) keeps reading the gate fresh.
+ */
+export function checkSearchUsageLimits(
+  attribution: BillingAttributionSnapshot
+): Promise<AttributedUsageLimitsResult> {
+  return checkUsageLimitsThroughCache(attribution, false)
+}
+
+/** Drops every cached gate answer. Test seam; never called in production code. */
+export function resetUsageGateCache(): void {
+  gateCache.clear()
+}

apps/sim/lib/billing/core/usage-gate-cache.ts

@@ -56,7 +56,10 @@ vi.mock('@/lib/billing/core/billing-attribution', () => ({
   resolveBillingAttribution: mocks.resolveBilling,
   resolveSystemBillingAttribution: mocks.resolveBilling,
   resolveOrganizationBillingAttribution: mocks.resolveBilling,
-  checkAttributedUsageLimits: mocks.checkUsage,
+}))
+
+vi.mock('@/lib/billing/core/usage-gate-cache', () => ({
+  checkSearchUsageLimits: mocks.checkUsage,
 }))
 
 /** Retrieval defaults are the flag's concern; here the flag is off so the search stays as configured. */

apps/sim/lib/knowledge/application/search.test.ts

@@ -3,9 +3,9 @@ import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
 import {
   type BillingAttributionSnapshot,
-  checkAttributedUsageLimits,
   toBillingContext,
 } from '@/lib/billing/core/billing-attribution'
+import { checkSearchUsageLimits } from '@/lib/billing/core/usage-gate-cache'
 import { recordUsage } from '@/lib/billing/core/usage-log'
 import { checkAndBillPayerOverageThreshold } from '@/lib/billing/threshold-billing'
 import { OrchestrationError } from '@/lib/core/orchestration/types'
@@ -297,7 +297,7 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({
       : undefined
     if (shouldMeter && billingAttribution) {
       const usage = await measureSearchStage('usage_admission', () =>
-        checkAttributedUsageLimits(billingAttribution)
+        checkSearchUsageLimits(billingAttribution)
       )
       if (usage.isExceeded) {
         throw new KnowledgeUsageLimitExceededError(

apps/sim/lib/knowledge/application/search.ts

@@ -65,13 +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 { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache'
 import * as embeddingClient from '@/lib/embeddings/client'
 import { processDocumentAsync } from '@/lib/knowledge/documents/service'
 
 const mockEmbeddingCapacity = vi.fn<typeof embeddingClient.assertKnowledgeEmbeddingCapacity>()
 beforeEach(() => {
-  resetIngestionUsageGateCache()
+  resetUsageGateCache()
   vi.spyOn(billingAttribution, 'checkAttributedUsageLimits').mockImplementation(
     mockCheckAttributedUsageLimits
   )

apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts

@@ -74,7 +74,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 { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache'
 import { env } from '@/lib/core/config/env'
 import {
   markInsideTriggerRun,
@@ -102,7 +102,7 @@ import { MAX_PROCESSING_ATTEMPTS } from '@/lib/knowledge/documents/types'
 
 const mockEmbeddingCapacity = vi.fn<typeof embeddingClient.assertKnowledgeEmbeddingCapacity>()
 beforeEach(() => {
-  resetIngestionUsageGateCache()
+  resetUsageGateCache()
   vi.spyOn(billingAttribution, 'checkAttributedUsageLimits').mockImplementation(
     mockCheckAttributedUsageLimits
   )

apps/sim/lib/knowledge/documents/document-processing-source.test.ts

@@ -36,7 +36,7 @@ import {
   type BillingAttributionSnapshot,
   toBillingContext,
 } from '@/lib/billing/core/billing-attribution'
-import { checkIngestionUsageLimits } from '@/lib/billing/core/ingestion-usage-gate'
+import { checkIngestionUsageLimits } from '@/lib/billing/core/usage-gate-cache'
 import { recordUsage } from '@/lib/billing/core/usage-log'
 import {
   applyStorageUsageDeltasInTx,