diff --git a/apps/sim/background/knowledge-processing.test.ts b/apps/sim/background/knowledge-processing.test.ts index 7e4423df541..f4d85c4f2c2 100644 --- a/apps/sim/background/knowledge-processing.test.ts +++ b/apps/sim/background/knowledge-processing.test.ts @@ -73,6 +73,21 @@ const WORKSPACE_PAYLOAD = { billingAttribution: BILLING_ATTRIBUTION, } +const ORGANIZATION_PAYLOAD = { + ...BASE_PAYLOAD, + billingScope: 'organization' as const, + actorUserId: 'organization-member', + workspaceId: null, + organizationId: 'organization-1', + billingAttribution: { + ...BILLING_ATTRIBUTION, + actorUserId: 'organization-member', + workspaceId: null, + organizationId: 'organization-1', + billingEntity: { type: 'organization' as const, id: 'organization-1' }, + }, +} + function mockQuotaExhaustion(error: EmbeddingQuotaExhaustedError): void { mockProcessDocumentAsync.mockImplementation(async (...args: unknown[]) => { const attemptContext = args[6] as { @@ -229,6 +244,36 @@ describe('knowledge processing worker', () => { ) }) + it('preserves organization ownership and billing attribution in the worker', async () => { + await runDocumentProcessing(structuredClone(ORGANIZATION_PAYLOAD)) + + expect(mockProcessDocumentAsync).toHaveBeenCalledWith( + BASE_PAYLOAD.knowledgeBaseId, + BASE_PAYLOAD.documentId, + BASE_PAYLOAD.docData, + BASE_PAYLOAD.processingOptions, + { + billingScope: 'organization', + actorUserId: ORGANIZATION_PAYLOAD.actorUserId, + workspaceId: null, + organizationId: ORGANIZATION_PAYLOAD.organizationId, + billingAttribution: ORGANIZATION_PAYLOAD.billingAttribution, + }, + BASE_PAYLOAD.requestId, + expect.objectContaining({ + chargedAtDispatch: true, + processingQueuedAt: new Date(BASE_PAYLOAD.processingQueuedAt), + }) + ) + }) + + it('rejects an organization mismatch before document processing starts', async () => { + await expect( + runDocumentProcessing({ ...ORGANIZATION_PAYLOAD, organizationId: 'organization-2' }) + ).rejects.toThrow('Document processing organization does not match billing attribution') + expect(mockProcessDocumentAsync).not.toHaveBeenCalled() + }) + it('rejects an actor mismatch before document processing starts', async () => { await expect( runDocumentProcessing({ diff --git a/apps/sim/background/knowledge-processing.ts b/apps/sim/background/knowledge-processing.ts index 41efe6c6674..35d77895376 100644 --- a/apps/sim/background/knowledge-processing.ts +++ b/apps/sim/background/knowledge-processing.ts @@ -12,8 +12,8 @@ import { isUsageLimitDocumentProcessingError, } from '@/lib/knowledge/documents/document-processing-error' import { + assertDocumentProcessingBillingContext, assertDocumentProcessingPayload, - type DocumentProcessingBillingContext, type DocumentProcessingPayload, } from '@/lib/knowledge/documents/processing-payload' import { @@ -34,19 +34,7 @@ export async function runDocumentProcessing( const startedAt = Date.now() const payload = assertDocumentProcessingPayload(rawPayload) const { knowledgeBaseId, documentId, docData, processingOptions, requestId } = payload - const billingContext: DocumentProcessingBillingContext = - payload.billingScope === 'workspace' - ? { - billingScope: 'workspace', - actorUserId: payload.actorUserId, - workspaceId: payload.workspaceId, - billingAttribution: payload.billingAttribution, - } - : { - billingScope: 'non-workspace', - actorUserId: payload.actorUserId, - workspaceId: null, - } + const billingContext = assertDocumentProcessingBillingContext(payload) const canScheduleQuotaContinuation = canScheduleDocumentProcessingQuotaContinuation(payload) logger.info(`[${requestId}] Starting Trigger.dev processing for document: ${docData.filename}`)