fix(knowledge): fail upload completion on dispatch errors

Author: TheodoreSpeaks

apps/sim/lib/knowledge/application/upload-sessions.test.ts

@@ -390,9 +390,85 @@ describe('knowledge-document upload application lifecycle', () => {
     expect(result.value.created).toBe(false)
     expect(mocks.resolveBilling).not.toHaveBeenCalled()
     expect(mocks.createDocument).not.toHaveBeenCalled()
+    expect(mocks.processQueue).not.toHaveBeenCalled()
     expect(mocks.recordAudit).not.toHaveBeenCalled()
   })
 
+  it('fails completion when document processing cannot be dispatched', async () => {
+    const failure = new Error('queue unavailable')
+    mocks.processQueue.mockRejectedValue(failure)
+    mocks.completeUpload.mockImplementation(
+      async (params: {
+        session: UploadSessionRecord
+        finalize: (session: UploadSessionRecord) => Promise<unknown>
+      }) => params.finalize(params.session)
+    )
+
+    await expect(
+      completeKnowledgeDocumentUpload.execute({
+        principal: PRINCIPAL,
+        input: {
+          knowledgeBaseId: 'knowledge-1',
+          assertedWorkspaceId: 'workspace-1',
+          uploadId: 'upload-1',
+          uploadToken: 'token',
+          source: 'api',
+        },
+        request: REQUEST,
+      })
+    ).rejects.toMatchObject({
+      name: 'KnowledgeDocumentProcessingDispatchError',
+      message: 'Knowledge document processing dispatch failed',
+      cause: failure,
+    })
+    expect(mocks.createDocument).toHaveBeenCalledTimes(1)
+    expect(mocks.recordAudit).not.toHaveBeenCalled()
+  })
+
+  it('retries a failed processing dispatch before completing a bound registration', async () => {
+    const recoveringSession = {
+      ...SESSION,
+      status: 'finalizing' as const,
+      completedFileId: null,
+      error: 'Knowledge document processing dispatch failed',
+    }
+    mocks.getUpload.mockResolvedValue(recoveringSession)
+    mocks.findBound.mockResolvedValue({
+      status: 'bound',
+      document: { ...DOCUMENT, processingStatus: 'pending' },
+    })
+    mocks.completeUpload.mockImplementation(
+      async (params: {
+        session: UploadSessionRecord
+        finalize: (session: UploadSessionRecord) => Promise<{
+          value: { document: typeof DOCUMENT; created: boolean; knowledgeBaseName: string | null }
+        }>
+      }) => ({
+        session: { ...params.session, status: 'completed' as const },
+        value: (await params.finalize(params.session)).value,
+        alreadyCompleted: true,
+      })
+    )
+
+    const result = await completeKnowledgeDocumentUpload.execute({
+      principal: PRINCIPAL,
+      input: {
+        knowledgeBaseId: 'knowledge-1',
+        assertedWorkspaceId: 'workspace-1',
+        uploadId: 'upload-1',
+        uploadToken: 'token',
+        source: 'api',
+      },
+      request: REQUEST,
+    })
+
+    expect(result.value.created).toBe(true)
+    expect(mocks.resolveBilling).toHaveBeenCalledTimes(1)
+    expect(mocks.processQueue).toHaveBeenCalledTimes(1)
+    expect(mocks.createDocument).not.toHaveBeenCalled()
+    expect(mocks.recordAudit).toHaveBeenCalledTimes(1)
+  })
+
   it('converges a finalization retry after durable bind without duplicate document or audit', async () => {
     const recoveringSession = {
       ...SESSION,
@@ -454,6 +530,7 @@ describe('knowledge-document upload application lifecycle', () => {
     expect(recovered.value.created).toBe(true)
     expect(retry.value.created).toBe(false)
     expect(mocks.createDocument).not.toHaveBeenCalled()
+    expect(mocks.processQueue).not.toHaveBeenCalled()
     expect(mocks.recordAudit).toHaveBeenCalledTimes(1)
   })
 

apps/sim/lib/knowledge/application/upload-sessions.ts

@@ -1,6 +1,5 @@
 import { AuditAction, AuditResourceType } from '@sim/audit'
 import type { Principal } from '@sim/auth/principal'
-import { createLogger } from '@sim/logger'
 import type { V2KnowledgeDocumentUploadMetadata } from '@/lib/api/contracts/v2/knowledge'
 import { v2KnowledgeDocumentUploadMetadataSchema } from '@/lib/api/contracts/v2/knowledge'
 import { checkAttributedUsageLimits } from '@/lib/billing/core/billing-attribution'
@@ -39,7 +38,14 @@ import {
 } from '@/lib/uploads/upload-session/service'
 import { validateFileType } from '@/lib/uploads/utils/validation'
 
-const logger = createLogger('KnowledgeUploadApplication')
+const PROCESSING_DISPATCH_FAILURE_MESSAGE = 'Knowledge document processing dispatch failed'
+
+class KnowledgeDocumentProcessingDispatchError extends Error {
+  constructor(cause: unknown) {
+    super(PROCESSING_DISPATCH_FAILURE_MESSAGE, { cause })
+    this.name = 'KnowledgeDocumentProcessingDispatchError'
+  }
+}
 
 export class KnowledgeDocumentUnsupportedMediaTypeError extends Error {
   constructor(message: string) {
@@ -253,6 +259,22 @@ export const completeKnowledgeDocumentUpload = defineAuthorizedKnowledgeUseCase(
           )
         }
         if (bound.status === 'bound') {
+          if (
+            claimed.error === PROCESSING_DISPATCH_FAILURE_MESSAGE &&
+            bound.document.processingStatus === 'pending'
+          ) {
+            const billingAttribution = await resolveKnowledgeBillingAttribution(
+              principal,
+              freshContext
+            )
+            await dispatchKnowledgeDocumentProcessing(
+              bound.document,
+              freshContext.knowledgeBaseId,
+              processingOptions,
+              requestId,
+              billingAttribution
+            )
+          }
           return {
             value: {
               document: bound.document,
@@ -315,26 +337,13 @@ export const completeKnowledgeDocumentUpload = defineAuthorizedKnowledgeUseCase(
           throw error
         }
 
-        const processingDocument: DocumentData = {
-          documentId: created.id,
-          filename: created.filename,
-          fileUrl: created.fileUrl,
-          fileSize: created.fileSize,
-          mimeType: created.mimeType,
-        }
-        processDocumentsWithQueue(
-          [processingDocument],
+        await dispatchKnowledgeDocumentProcessing(
+          created,
           registrationContext.knowledgeBaseId,
-          processingOptions ?? {},
+          processingOptions,
           requestId,
           billingAttribution
-        ).catch((error: unknown) => {
-          logger.error('Knowledge document processing pipeline failed', {
-            knowledgeBaseId: registrationContext.knowledgeBaseId,
-            documentId: created.id,
-            error,
-          })
-        })
+        )
         return {
           value: {
             document: created,
@@ -372,6 +381,33 @@ export const completeKnowledgeDocumentUpload = defineAuthorizedKnowledgeUseCase(
   },
 })
 
+async function dispatchKnowledgeDocumentProcessing(
+  document: CreatedKnowledgeDocument,
+  knowledgeBaseId: string,
+  processingOptions: V2KnowledgeDocumentUploadMetadata['processingOptions'],
+  requestId: string,
+  billingAttribution: Awaited<ReturnType<typeof resolveKnowledgeBillingAttribution>>
+): Promise<void> {
+  const processingDocument: DocumentData = {
+    documentId: document.id,
+    filename: document.filename,
+    fileUrl: document.fileUrl,
+    fileSize: document.fileSize,
+    mimeType: document.mimeType,
+  }
+  try {
+    await processDocumentsWithQueue(
+      [processingDocument],
+      knowledgeBaseId,
+      processingOptions ?? {},
+      requestId,
+      billingAttribution
+    )
+  } catch (error) {
+    throw new KnowledgeDocumentProcessingDispatchError(error)
+  }
+}
+
 async function loadBoundKnowledgeDocumentUpload(
   principal: Principal,
   input: KnowledgeDocumentUploadControlInput,