Skip to content

Commit b26e91f

Browse files
icecrasher321icecrasher321
committed
fix(knowledge): update processing tests for scoped billing
1 parent f85f63b commit b26e91f

2 files changed

Lines changed: 116 additions & 31 deletions

File tree

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

Lines changed: 83 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -6,35 +6,30 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
66

77
const {
88
mockCalculateCost,
9-
mockCheckActorUsageLimits,
10-
mockCheckAndBillOverageThreshold,
9+
mockCheckAttributedUsageLimits,
10+
mockCheckAndBillPayerOverageThreshold,
1111
mockGenerateEmbeddings,
1212
mockGetBoundWorkspaceFileSecretProvenanceByMetadata,
1313
mockGetFileMetadataByKeys,
1414
mockProcessDocument,
1515
mockRecordUsage,
1616
} = vi.hoisted(() => ({
1717
mockCalculateCost: vi.fn(),
18-
mockCheckActorUsageLimits: vi.fn(),
19-
mockCheckAndBillOverageThreshold: vi.fn(),
18+
mockCheckAttributedUsageLimits: vi.fn(),
19+
mockCheckAndBillPayerOverageThreshold: vi.fn(),
2020
mockGenerateEmbeddings: vi.fn(),
2121
mockGetBoundWorkspaceFileSecretProvenanceByMetadata: vi.fn(),
2222
mockGetFileMetadataByKeys: vi.fn(),
2323
mockProcessDocument: vi.fn(),
2424
mockRecordUsage: vi.fn(),
2525
}))
2626

27-
vi.mock('@/lib/billing/calculations/usage-monitor', () => ({
28-
checkActorUsageLimits: mockCheckActorUsageLimits,
29-
}))
30-
3127
vi.mock('@/lib/billing/core/usage-log', () => ({
3228
recordUsage: mockRecordUsage,
3329
}))
3430

3531
vi.mock('@/lib/billing/threshold-billing', () => ({
36-
checkAndBillOverageThreshold: mockCheckAndBillOverageThreshold,
37-
checkAndBillPayerOverageThreshold: vi.fn(),
32+
checkAndBillPayerOverageThreshold: mockCheckAndBillPayerOverageThreshold,
3833
}))
3934

4035
vi.mock('@/lib/knowledge/documents/document-processor', () => ({
@@ -69,11 +64,15 @@ vi.mock('@/providers/utils', () => ({
6964
calculateCost: mockCalculateCost,
7065
}))
7166

67+
import * as billingAttribution from '@/lib/billing/core/billing-attribution'
7268
import * as embeddingClient from '@/lib/embeddings/client'
7369
import { processDocumentAsync } from '@/lib/knowledge/documents/service'
7470

7571
const mockEmbeddingCapacity = vi.fn<typeof embeddingClient.assertKnowledgeEmbeddingCapacity>()
7672
beforeEach(() => {
73+
vi.spyOn(billingAttribution, 'checkAttributedUsageLimits').mockImplementation(
74+
mockCheckAttributedUsageLimits
75+
)
7776
mockEmbeddingCapacity.mockReset().mockResolvedValue(undefined)
7877
vi.spyOn(embeddingClient, 'assertKnowledgeEmbeddingCapacity').mockImplementation(
7978
mockEmbeddingCapacity
@@ -86,13 +85,26 @@ const PERSISTED_KEY = 'workspace/workspace-1/persisted.pdf'
8685
const PERSISTED_URL = `/api/files/serve/${encodeURIComponent(PERSISTED_KEY)}?context=workspace`
8786
const CONTENT_UPDATED_AT = new Date('2026-08-05T12:00:00.000Z')
8887

88+
const BILLING_ATTRIBUTION: billingAttribution.BillingAttributionSnapshot = {
89+
actorUserId: 'uploader-1',
90+
workspaceId: 'workspace-1',
91+
organizationId: null,
92+
billedAccountUserId: 'workspace-owner',
93+
billingEntity: { type: 'user', id: 'workspace-owner' },
94+
billingPeriod: {
95+
start: '2026-08-01T00:00:00.000Z',
96+
end: '2026-09-01T00:00:00.000Z',
97+
source: 'default',
98+
},
99+
payerSubscription: null,
100+
}
101+
89102
const PERSISTED_CONTEXT = {
90-
workspaceId: null,
91-
knowledgeBaseUserId: 'knowledge-owner',
103+
workspaceId: 'workspace-1',
104+
organizationId: null,
92105
chunkingConfig: null,
93106
embeddingModel: 'text-embedding-3-small',
94107
embeddingDimension: 1536,
95-
billedAccountUserId: null,
96108
uploadedBy: 'uploader-1',
97109
filename: 'persisted.pdf',
98110
fileUrl: PERSISTED_URL,
@@ -194,7 +206,7 @@ describe('knowledge document indexing usage', () => {
194206
// The processing claim is guarded and returns the row it claimed; without a
195207
// stub every worker would read as 'already completed' and return early.
196208
dbChainMockFns.returning.mockResolvedValue([{ id: 'document-1' }])
197-
mockCheckActorUsageLimits.mockResolvedValue({ isExceeded: false })
209+
mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false })
198210
mockGetFileMetadataByKeys.mockImplementation(async (_keys: string[], context: string) =>
199211
context === 'workspace' ? [SOURCE_BINDING] : []
200212
)
@@ -233,7 +245,7 @@ describe('knowledge document indexing usage', () => {
233245
DOCUMENT_ID,
234246
DOC_DATA,
235247
{},
236-
undefined,
248+
BILLING_ATTRIBUTION,
237249
'timeout-pass'
238250
)
239251
const rejected = expect(pending).rejects.toThrow('Document processing timed out')
@@ -263,7 +275,7 @@ describe('knowledge document indexing usage', () => {
263275
DOCUMENT_ID,
264276
DOC_DATA,
265277
{},
266-
undefined,
278+
BILLING_ATTRIBUTION,
267279
'timeout-pass'
268280
)
269281
const rejected = expect(pending).rejects.toThrow('Document processing timed out')
@@ -278,22 +290,54 @@ describe('knowledge document indexing usage', () => {
278290

279291
it('records one embedding charge per indexing pass', async () => {
280292
armDocumentReads()
281-
await processDocumentAsync(KNOWLEDGE_BASE_ID, DOCUMENT_ID, DOC_DATA, {}, undefined, 'pass-1')
293+
await processDocumentAsync(
294+
KNOWLEDGE_BASE_ID,
295+
DOCUMENT_ID,
296+
DOC_DATA,
297+
{},
298+
BILLING_ATTRIBUTION,
299+
'pass-1'
300+
)
282301

283302
expect(mockRecordUsage).toHaveBeenCalledTimes(1)
284303
expect(recordedSourceReference(0)).toBe(`knowledge-document:${DOCUMENT_ID}:pass-1`)
304+
expect(mockCheckAttributedUsageLimits).toHaveBeenCalledWith(BILLING_ATTRIBUTION)
305+
expect(mockRecordUsage).toHaveBeenCalledWith(
306+
expect.objectContaining({
307+
userId: BILLING_ATTRIBUTION.actorUserId,
308+
workspaceId: BILLING_ATTRIBUTION.workspaceId,
309+
billingEntity: BILLING_ATTRIBUTION.billingEntity,
310+
})
311+
)
312+
expect(mockCheckAndBillPayerOverageThreshold).toHaveBeenCalledWith(
313+
BILLING_ATTRIBUTION.billingEntity
314+
)
285315
})
286316

287317
it('reuses the same usage source reference across attempts of one indexing pass', async () => {
288318
const nowSpy = vi.spyOn(Date, 'now')
289319

290320
nowSpy.mockReturnValue(1_000)
291321
armDocumentReads()
292-
await processDocumentAsync(KNOWLEDGE_BASE_ID, DOCUMENT_ID, DOC_DATA, {}, undefined, 'pass-1')
322+
await processDocumentAsync(
323+
KNOWLEDGE_BASE_ID,
324+
DOCUMENT_ID,
325+
DOC_DATA,
326+
{},
327+
BILLING_ATTRIBUTION,
328+
'pass-1'
329+
)
293330

294331
nowSpy.mockReturnValue(9_000)
295332
armDocumentReads()
296-
await processDocumentAsync(KNOWLEDGE_BASE_ID, DOCUMENT_ID, DOC_DATA, {}, undefined, 'pass-1')
333+
await processDocumentAsync(
334+
KNOWLEDGE_BASE_ID,
335+
DOCUMENT_ID,
336+
DOC_DATA,
337+
{},
338+
BILLING_ATTRIBUTION,
339+
'pass-1'
340+
)
297341

298342
nowSpy.mockRestore()
299343

@@ -303,10 +347,24 @@ describe('knowledge document indexing usage', () => {
303347

304348
it('uses a distinct usage source reference for a genuinely new indexing pass', async () => {
305349
armDocumentReads()
306-
await processDocumentAsync(KNOWLEDGE_BASE_ID, DOCUMENT_ID, DOC_DATA, {}, undefined, 'pass-1')
350+
await processDocumentAsync(
351+
KNOWLEDGE_BASE_ID,
352+
DOCUMENT_ID,
353+
DOC_DATA,
354+
{},
355+
BILLING_ATTRIBUTION,
356+
'pass-1'
357+
)
307358

308359
armDocumentReads()
309-
await processDocumentAsync(KNOWLEDGE_BASE_ID, DOCUMENT_ID, DOC_DATA, {}, undefined, 'pass-2')
360+
await processDocumentAsync(
361+
KNOWLEDGE_BASE_ID,
362+
DOCUMENT_ID,
363+
DOC_DATA,
364+
{},
365+
BILLING_ATTRIBUTION,
366+
'pass-2'
367+
)
310368

311369
expect(mockRecordUsage).toHaveBeenCalledTimes(2)
312370
expect(recordedSourceReference(1)).not.toBe(recordedSourceReference(0))
@@ -317,11 +375,11 @@ describe('knowledge document indexing usage', () => {
317375

318376
nowSpy.mockReturnValue(1_000)
319377
armDocumentReads()
320-
await processDocumentAsync(KNOWLEDGE_BASE_ID, DOCUMENT_ID, DOC_DATA, {})
378+
await processDocumentAsync(KNOWLEDGE_BASE_ID, DOCUMENT_ID, DOC_DATA, {}, BILLING_ATTRIBUTION)
321379

322380
nowSpy.mockReturnValue(9_000)
323381
armDocumentReads()
324-
await processDocumentAsync(KNOWLEDGE_BASE_ID, DOCUMENT_ID, DOC_DATA, {})
382+
await processDocumentAsync(KNOWLEDGE_BASE_ID, DOCUMENT_ID, DOC_DATA, {}, BILLING_ATTRIBUTION)
325383

326384
nowSpy.mockRestore()
327385

@@ -333,7 +391,7 @@ describe('knowledge document indexing usage', () => {
333391

334392
it('re-bills the fallback reference when the embedding model changes', async () => {
335393
armDocumentReads()
336-
await processDocumentAsync(KNOWLEDGE_BASE_ID, DOCUMENT_ID, DOC_DATA, {})
394+
await processDocumentAsync(KNOWLEDGE_BASE_ID, DOCUMENT_ID, DOC_DATA, {}, BILLING_ATTRIBUTION)
337395

338396
mockGenerateEmbeddings.mockResolvedValue({
339397
embeddings: [[0.3, 0.4]],
@@ -342,7 +400,7 @@ describe('knowledge document indexing usage', () => {
342400
pricingId: 'text-embedding-3-large',
343401
})
344402
armDocumentReads()
345-
await processDocumentAsync(KNOWLEDGE_BASE_ID, DOCUMENT_ID, DOC_DATA, {})
403+
await processDocumentAsync(KNOWLEDGE_BASE_ID, DOCUMENT_ID, DOC_DATA, {}, BILLING_ATTRIBUTION)
346404

347405
expect(recordedSourceReference(1)).not.toBe(recordedSourceReference(0))
348406
})

‎apps/sim/lib/knowledge/documents/retry-processing-grace.test.ts‎

Lines changed: 33 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ vi.mock('@/lib/knowledge/documents/processing-outbox-event', () => ({
1818
vi.mock('@/lib/uploads', () => ({ StorageService: {} }))
1919
vi.mock('@/connectors/registry.server', () => ({ CONNECTOR_REGISTRY: {} }))
2020

21+
import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution'
2122
import { isStuckDocumentSweepEligible } from '@/lib/knowledge/connectors/sync-primitives'
2223
import {
2324
processDocumentsWithQueue,
@@ -32,6 +33,20 @@ const DOC_DATA = {
3233
mimeType: 'application/pdf',
3334
}
3435

36+
const BILLING_ATTRIBUTION: BillingAttributionSnapshot = {
37+
actorUserId: 'user-1',
38+
workspaceId: 'workspace-1',
39+
organizationId: null,
40+
billedAccountUserId: 'workspace-owner',
41+
billingEntity: { type: 'user', id: 'workspace-owner' },
42+
billingPeriod: {
43+
start: '2026-08-01T00:00:00.000Z',
44+
end: '2026-09-01T00:00:00.000Z',
45+
source: 'default',
46+
},
47+
payerSubscription: null,
48+
}
49+
3550
/**
3651
* Runs the requeue and returns the values it wrote. Dispatch runs after the
3752
* reset transaction and needs infrastructure this test does not stand up, so a
@@ -73,7 +88,9 @@ describe('processDocumentsWithQueue dispatch stamp', () => {
7388
beforeEach(() => {
7489
vi.clearAllMocks()
7590
resetDbChainMock()
76-
dbChainMockFns.limit.mockResolvedValue([{ userId: 'user-1', workspaceId: null }])
91+
dbChainMockFns.limit.mockResolvedValue([
92+
{ userId: 'user-1', workspaceId: 'workspace-1', organizationId: null },
93+
])
7794
})
7895

7996
/**
@@ -86,7 +103,7 @@ describe('processDocumentsWithQueue dispatch stamp', () => {
86103
'kb-1',
87104
{},
88105
'req-1',
89-
undefined
106+
BILLING_ATTRIBUTION
90107
).catch(() => {})
91108
}
92109

@@ -340,7 +357,9 @@ describe('processing attempt budget', () => {
340357
beforeEach(() => {
341358
vi.clearAllMocks()
342359
resetDbChainMock()
343-
dbChainMockFns.limit.mockResolvedValue([{ userId: 'user-1', workspaceId: null }])
360+
dbChainMockFns.limit.mockResolvedValue([
361+
{ userId: 'user-1', workspaceId: 'workspace-1', organizationId: null },
362+
])
344363
})
345364

346365
it('spends one attempt per dispatch, in the same guarded write', async () => {
@@ -349,7 +368,7 @@ describe('processing attempt budget', () => {
349368
'kb-1',
350369
{},
351370
'req-1',
352-
undefined
371+
BILLING_ATTRIBUTION
353372
).catch(() => {})
354373

355374
const stampCall = dbChainMockFns.set.mock.calls.find(
@@ -415,9 +434,17 @@ describe('retryDocumentProcessing dispatch unwind', () => {
415434
.mockResolvedValueOnce([{ id: 'doc-1' }])
416435
.mockResolvedValueOnce([])
417436
.mockResolvedValueOnce([])
418-
dbChainMockFns.limit.mockResolvedValue([{ userId: 'user-1', workspaceId: null }])
437+
dbChainMockFns.limit.mockResolvedValue([
438+
{ userId: 'user-1', workspaceId: 'workspace-1', organizationId: null },
439+
])
419440

420-
const result = await retryDocumentProcessing('kb-1', 'doc-1', DOC_DATA, 'req-1', undefined)
441+
const result = await retryDocumentProcessing(
442+
'kb-1',
443+
'doc-1',
444+
DOC_DATA,
445+
'req-1',
446+
BILLING_ATTRIBUTION
447+
)
421448

422449
expect(result).toMatchObject({ success: false, status: 'failed' })
423450
expect(result.message).toContain('was not accepted')

0 commit comments

Comments
 (0)