fix(knowledge): preserve organization billing context in document worker (#7596)

Author: waleedlatif1

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({

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}`)