Skip to content

Commit 3f245d7

Browse files
waleedlatif1waleedlatif1
authored
improvement(knowledge): cache admitted usage checks on the search path (#7989)
* improvement(knowledge): cache admitted usage checks on the search path * fix(knowledge): keep search-path refusals out of the shared usage gate cache * fix(knowledge): key the usage gate cache by billing period source
1 parent 2d87824 commit 3f245d7

9 files changed

Lines changed: 196 additions & 85 deletions

File tree

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,10 @@ vi.mock('@sim/platform-authz/workspace', () => ({
3131
vi.mock('@/lib/billing/core/billing-attribution', () => ({
3232
resolveBillingAttribution: mocks.resolveBilling,
3333
resolveSystemBillingAttribution: mocks.resolveBilling,
34-
checkAttributedUsageLimits: mocks.checkUsage,
34+
}))
35+
36+
vi.mock('@/lib/billing/core/usage-gate-cache', () => ({
37+
checkSearchUsageLimits: mocks.checkUsage,
3538
}))
3639

3740
/** 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‎

Lines changed: 0 additions & 69 deletions
This file was deleted.

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

Lines changed: 69 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,10 @@ vi.mock('@/lib/billing/core/billing-attribution', () => ({
1313
import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution'
1414
import {
1515
checkIngestionUsageLimits,
16-
INGESTION_USAGE_GATE_TTL_MS,
17-
resetIngestionUsageGateCache,
18-
} from '@/lib/billing/core/ingestion-usage-gate'
16+
checkSearchUsageLimits,
17+
resetUsageGateCache,
18+
USAGE_GATE_TTL_MS,
19+
} from '@/lib/billing/core/usage-gate-cache'
1920

2021
const ATTRIBUTION: BillingAttributionSnapshot = {
2122
actorUserId: 'member-1',
@@ -31,9 +32,19 @@ const ATTRIBUTION: BillingAttributionSnapshot = {
3132
payerSubscription: null,
3233
}
3334

35+
const SUBSCRIPTION: BillingAttributionSnapshot['payerSubscription'] = {
36+
id: 'sub-1',
37+
referenceId: 'org-1',
38+
plan: 'team',
39+
status: 'active',
40+
seats: 5,
41+
periodStart: '2026-09-01T00:00:00.000Z',
42+
periodEnd: '2026-10-01T00:00:00.000Z',
43+
}
44+
3445
describe('checkIngestionUsageLimits', () => {
3546
beforeEach(() => {
36-
resetIngestionUsageGateCache()
47+
resetUsageGateCache()
3748
mockCheck.mockReset().mockResolvedValue({ isExceeded: false })
3849
})
3950
afterEach(() => vi.restoreAllMocks())
@@ -68,13 +79,13 @@ describe('checkIngestionUsageLimits', () => {
6879

6980
/** `lru-cache` reads `performance.now()` and debounces it behind a real 1 ms timer. */
7081
const start = performance.now()
71-
vi.spyOn(performance, 'now').mockReturnValue(start + INGESTION_USAGE_GATE_TTL_MS + 1)
82+
vi.spyOn(performance, 'now').mockReturnValue(start + USAGE_GATE_TTL_MS + 1)
7283
await sleep(5)
7384
expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false)
7485
expect(mockCheck).toHaveBeenCalledTimes(2)
7586
})
7687

77-
it('separates answers by actor, period and payer', async () => {
88+
it('separates answers by actor, period, period source, payer and plan', async () => {
7889
await checkIngestionUsageLimits(ATTRIBUTION)
7990
await checkIngestionUsageLimits({ ...ATTRIBUTION, actorUserId: 'member-2' })
8091
await checkIngestionUsageLimits({
@@ -86,7 +97,16 @@ describe('checkIngestionUsageLimits', () => {
8697
billedAccountUserId: 'owner-2',
8798
billingEntity: { type: 'user', id: 'owner-2' },
8899
})
89-
expect(mockCheck).toHaveBeenCalledTimes(4)
100+
await checkIngestionUsageLimits({
101+
...ATTRIBUTION,
102+
billingPeriod: { ...ATTRIBUTION.billingPeriod, source: 'reporting' },
103+
})
104+
await checkIngestionUsageLimits({ ...ATTRIBUTION, payerSubscription: SUBSCRIPTION })
105+
await checkIngestionUsageLimits({
106+
...ATTRIBUTION,
107+
payerSubscription: { ...SUBSCRIPTION, plan: 'enterprise' },
108+
})
109+
expect(mockCheck).toHaveBeenCalledTimes(7)
90110
})
91111

92112
it('does not cache a failed read', async () => {
@@ -96,3 +116,45 @@ describe('checkIngestionUsageLimits', () => {
96116
expect(mockCheck).toHaveBeenCalledTimes(2)
97117
})
98118
})
119+
120+
describe('checkSearchUsageLimits', () => {
121+
beforeEach(() => {
122+
resetUsageGateCache()
123+
mockCheck.mockReset().mockResolvedValue({ isExceeded: false })
124+
})
125+
126+
it('reuses an admission across workspaces of the same payer', async () => {
127+
await checkSearchUsageLimits(ATTRIBUTION)
128+
await checkSearchUsageLimits({ ...ATTRIBUTION, workspaceId: 'workspace-9' })
129+
expect(mockCheck).toHaveBeenCalledTimes(1)
130+
})
131+
132+
it('re-reads a refusal, so a raised limit applies on the next search', async () => {
133+
mockCheck.mockResolvedValueOnce({ isExceeded: true, scope: 'payer', message: 'over' })
134+
expect((await checkSearchUsageLimits(ATTRIBUTION)).isExceeded).toBe(true)
135+
expect((await checkSearchUsageLimits(ATTRIBUTION)).isExceeded).toBe(false)
136+
expect(mockCheck).toHaveBeenCalledTimes(2)
137+
})
138+
139+
it('does not serve a refusal cached by ingestion', async () => {
140+
mockCheck.mockResolvedValueOnce({ isExceeded: true, scope: 'payer', message: 'over' })
141+
expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true)
142+
expect((await checkSearchUsageLimits(ATTRIBUTION)).isExceeded).toBe(false)
143+
expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false)
144+
expect(mockCheck).toHaveBeenCalledTimes(2)
145+
})
146+
147+
it('never stores a refusal for ingestion to serve', async () => {
148+
mockCheck.mockResolvedValueOnce({ isExceeded: true, scope: 'payer', message: 'over' })
149+
expect((await checkSearchUsageLimits(ATTRIBUTION)).isExceeded).toBe(true)
150+
expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false)
151+
expect(mockCheck).toHaveBeenCalledTimes(2)
152+
})
153+
154+
it('does not cache a failed read', async () => {
155+
mockCheck.mockRejectedValueOnce(new Error('ledger unavailable'))
156+
await expect(checkSearchUsageLimits(ATTRIBUTION)).rejects.toThrow('ledger unavailable')
157+
await checkSearchUsageLimits(ATTRIBUTION)
158+
expect(mockCheck).toHaveBeenCalledTimes(2)
159+
})
160+
})

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

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
import { LRUCache } from 'lru-cache'
2+
import {
3+
type AttributedUsageLimitsResult,
4+
type BillingAttributionSnapshot,
5+
checkAttributedUsageLimits,
6+
} from '@/lib/billing/core/billing-attribution'
7+
import { coalesceLocally } from '@/lib/concurrency/singleflight'
8+
9+
/**
10+
* How long a usage-gate answer stays usable on the high-frequency paths.
11+
*
12+
* The gate sums the payer's usage ledger for the billing period, which grows
13+
* with the payer's activity, so a large organization scans its whole period on
14+
* every uncached call. Bulk ingestion re-checks per document and knowledge
15+
* search checks per query. Staleness is bounded by this TTL and fails in the
16+
* harmless direction: a payer who crosses their limit keeps going for at most
17+
* this long, which charges nobody wrongly.
18+
*/
19+
export const USAGE_GATE_TTL_MS = 60 * 1000
20+
21+
/**
22+
* Recent gate answers, admitted and refused, with `LRUCache` supplying the TTL
23+
* and the size bound. Each entry point decides which of them it may serve.
24+
*/
25+
const gateCache = new LRUCache<string, AttributedUsageLimitsResult>({
26+
max: 10_000,
27+
ttl: USAGE_GATE_TTL_MS,
28+
})
29+
30+
/**
31+
* The gate depends on who pays, for which period (and how that period was
32+
* derived), under which plan, and which member acts: the payer pool, its limit
33+
* and the per-member cap are all part of the answer. The workspace is not, so
34+
* every workspace of one payer shares an entry.
35+
*/
36+
function gateKey(attribution: BillingAttributionSnapshot): string {
37+
const subscription = attribution.payerSubscription
38+
return [
39+
attribution.billingEntity.type,
40+
attribution.billingEntity.id,
41+
attribution.billingPeriod.start,
42+
attribution.billingPeriod.end,
43+
attribution.billingPeriod.source ?? '',
44+
attribution.billedAccountUserId,
45+
attribution.actorUserId,
46+
subscription?.id ?? '',
47+
subscription?.plan ?? '',
48+
subscription?.status ?? '',
49+
subscription?.seats ?? '',
50+
].join(':')
51+
}
52+
53+
/**
54+
* Serves a cached answer the caller accepts, otherwise reads the gate.
55+
*
56+
* `cacheRefusals` governs both directions: a caller that must re-read refusals
57+
* also never stores one, so a refusal read on the search path never reaches
58+
* ingestion. The usage read fails closed (a ledger error comes back as
59+
* exceeded), which makes that the only way a search-path outage stays out of
60+
* the cache. A read that throws writes nothing.
61+
*
62+
* `coalesceLocally` collapses concurrent misses onto one ledger read and bounds
63+
* a hung read at its settle deadline. The write stays on the value this caller
64+
* received, so a producer that timed out and later resolved cannot overwrite a
65+
* fresher answer.
66+
*
67+
* There is deliberately no invalidator: usage and limit changes land in other
68+
* processes (execution workers, Stripe webhooks), so the TTL is the real bound.
69+
*/
70+
async function checkUsageLimitsThroughCache(
71+
attribution: BillingAttributionSnapshot,
72+
cacheRefusals: boolean
73+
): Promise<AttributedUsageLimitsResult> {
74+
const key = gateKey(attribution)
75+
const cached = gateCache.get(key)
76+
if (cached !== undefined && (cacheRefusals || !cached.isExceeded)) return cached
77+
78+
const result = await coalesceLocally(`usage-gate:${key}`, () =>
79+
checkAttributedUsageLimits(attribution)
80+
)
81+
if (cacheRefusals || !result.isExceeded) gateCache.set(key, result)
82+
return result
83+
}
84+
85+
/**
86+
* {@link checkAttributedUsageLimits} for background ingestion. Serves admitted
87+
* and refused answers alike: nothing on this path has a person waiting for a
88+
* raised limit to apply, so a refused payer waits at most the TTL.
89+
*/
90+
export function checkIngestionUsageLimits(
91+
attribution: BillingAttributionSnapshot
92+
): Promise<AttributedUsageLimitsResult> {
93+
return checkUsageLimitsThroughCache(attribution, true)
94+
}
95+
96+
/**
97+
* {@link checkAttributedUsageLimits} for knowledge search. Serves only a cached
98+
* admission: a refusal is always re-read, so a payer who just raised their limit
99+
* or upgraded is never held behind a cached block while they wait on a search.
100+
* Every other interactive caller (uploads, execution admission, the settings
101+
* surfaces) keeps reading the gate fresh.
102+
*/
103+
export function checkSearchUsageLimits(
104+
attribution: BillingAttributionSnapshot
105+
): Promise<AttributedUsageLimitsResult> {
106+
return checkUsageLimitsThroughCache(attribution, false)
107+
}
108+
109+
/** Drops every cached gate answer. Test seam; never called in production code. */
110+
export function resetUsageGateCache(): void {
111+
gateCache.clear()
112+
}

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,10 @@ vi.mock('@/lib/billing/core/billing-attribution', () => ({
5656
resolveBillingAttribution: mocks.resolveBilling,
5757
resolveSystemBillingAttribution: mocks.resolveBilling,
5858
resolveOrganizationBillingAttribution: mocks.resolveBilling,
59-
checkAttributedUsageLimits: mocks.checkUsage,
59+
}))
60+
61+
vi.mock('@/lib/billing/core/usage-gate-cache', () => ({
62+
checkSearchUsageLimits: mocks.checkUsage,
6063
}))
6164

6265
/** Retrieval defaults are the flag's concern; here the flag is off so the search stays as configured. */

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,9 @@ import { createLogger } from '@sim/logger'
33
import { getErrorMessage } from '@sim/utils/errors'
44
import {
55
type BillingAttributionSnapshot,
6-
checkAttributedUsageLimits,
76
toBillingContext,
87
} from '@/lib/billing/core/billing-attribution'
8+
import { checkSearchUsageLimits } from '@/lib/billing/core/usage-gate-cache'
99
import { recordUsage } from '@/lib/billing/core/usage-log'
1010
import { checkAndBillPayerOverageThreshold } from '@/lib/billing/threshold-billing'
1111
import { OrchestrationError } from '@/lib/core/orchestration/types'
@@ -297,7 +297,7 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({
297297
: undefined
298298
if (shouldMeter && billingAttribution) {
299299
const usage = await measureSearchStage('usage_admission', () =>
300-
checkAttributedUsageLimits(billingAttribution)
300+
checkSearchUsageLimits(billingAttribution)
301301
)
302302
if (usage.isExceeded) {
303303
throw new KnowledgeUsageLimitExceededError(

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,13 +65,13 @@ vi.mock('@/providers/utils', () => ({
6565
}))
6666

6767
import * as billingAttribution from '@/lib/billing/core/billing-attribution'
68-
import { resetIngestionUsageGateCache } from '@/lib/billing/core/ingestion-usage-gate'
68+
import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache'
6969
import * as embeddingClient from '@/lib/embeddings/client'
7070
import { processDocumentAsync } from '@/lib/knowledge/documents/service'
7171

7272
const mockEmbeddingCapacity = vi.fn<typeof embeddingClient.assertKnowledgeEmbeddingCapacity>()
7373
beforeEach(() => {
74-
resetIngestionUsageGateCache()
74+
resetUsageGateCache()
7575
vi.spyOn(billingAttribution, 'checkAttributedUsageLimits').mockImplementation(
7676
mockCheckAttributedUsageLimits
7777
)

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ vi.mock('@/lib/uploads/server/metadata', () => ({
7474
}))
7575

7676
import * as billingAttribution from '@/lib/billing/core/billing-attribution'
77-
import { resetIngestionUsageGateCache } from '@/lib/billing/core/ingestion-usage-gate'
77+
import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache'
7878
import { env } from '@/lib/core/config/env'
7979
import {
8080
markInsideTriggerRun,
@@ -102,7 +102,7 @@ import { MAX_PROCESSING_ATTEMPTS } from '@/lib/knowledge/documents/types'
102102

103103
const mockEmbeddingCapacity = vi.fn<typeof embeddingClient.assertKnowledgeEmbeddingCapacity>()
104104
beforeEach(() => {
105-
resetIngestionUsageGateCache()
105+
resetUsageGateCache()
106106
vi.spyOn(billingAttribution, 'checkAttributedUsageLimits').mockImplementation(
107107
mockCheckAttributedUsageLimits
108108
)

‎apps/sim/lib/knowledge/documents/service.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ import {
3636
type BillingAttributionSnapshot,
3737
toBillingContext,
3838
} from '@/lib/billing/core/billing-attribution'
39-
import { checkIngestionUsageLimits } from '@/lib/billing/core/ingestion-usage-gate'
39+
import { checkIngestionUsageLimits } from '@/lib/billing/core/usage-gate-cache'
4040
import { recordUsage } from '@/lib/billing/core/usage-log'
4141
import {
4242
applyStorageUsageDeltasInTx,

0 commit comments

Comments
 (0)