fix(knowledge): serialize document processing attempts

Author: TheodoreSpeaks

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

@@ -138,6 +138,7 @@ describe('knowledge document processing source', () => {
       .mockResolvedValueOnce([PERSISTED_CONTEXT])
       .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW])
       .mockResolvedValueOnce([{ id: 'document-1' }])
+    dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'document-1' }])
     mockCheckActorUsageLimits.mockResolvedValue({ isExceeded: false })
     mockGetFileMetadataByKeys.mockImplementation(async (_keys: string[], context: string) =>
       context === 'workspace' ? [SOURCE_BINDING] : []
@@ -228,4 +229,18 @@ describe('knowledge document processing source', () => {
     expect(mockProcessDocument).not.toHaveBeenCalled()
     expect(mockGenerateEmbeddings).not.toHaveBeenCalled()
   })
+
+  it('does not start work when another processing attempt already owns the document', async () => {
+    dbChainMockFns.returning.mockReset().mockResolvedValueOnce([])
+
+    await processDocumentAsync('knowledge-base-1', 'document-1', {
+      filename: 'stale.pdf',
+      fileUrl: 'https://example.com/stale.pdf',
+      fileSize: 1,
+      mimeType: 'text/plain',
+    })
+
+    expect(mockProcessDocument).not.toHaveBeenCalled()
+    expect(mockGenerateEmbeddings).not.toHaveBeenCalled()
+  })
 })

apps/sim/lib/knowledge/documents/processing-outbox-handler.test.ts

@@ -105,6 +105,19 @@ describe('knowledge document processing outbox handler', () => {
     }
   )
 
+  it('keeps the event retryable while an earlier processing attempt is active', async () => {
+    mocks.getKnowledgeDocument.mockResolvedValueOnce({
+      ...DOCUMENT,
+      processingStatus: 'processing',
+    })
+
+    await expect(handler()(PAYLOAD, createContext())).rejects.toThrow(
+      'Knowledge document document-1 is already being processed'
+    )
+
+    expect(mocks.processDocumentsWithQueue).not.toHaveBeenCalled()
+  })
+
   it('propagates dispatch failures so the outbox schedules a retry', async () => {
     const failure = new Error('queue unavailable')
     mocks.processDocumentsWithQueue.mockRejectedValueOnce(failure)

apps/sim/lib/knowledge/documents/processing-outbox-handler.ts

@@ -59,6 +59,9 @@ const processKnowledgeDocument: OutboxHandler<unknown> = async (rawPayload, cont
   context.signal.throwIfAborted()
   const document = await getKnowledgeDocument(payload.knowledgeBaseId, payload.documentId)
   if (!document || document.processingStatus === 'completed') return
+  if (document.processingStatus === 'processing') {
+    throw new Error(`Knowledge document ${document.id} is already being processed`)
+  }
 
   context.signal.throwIfAborted()
   await processDocumentsWithQueue(

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

@@ -778,6 +778,7 @@ export async function processDocumentAsync(
   providedBillingContext?: BillingAttributionSnapshot | DocumentProcessingBillingContext
 ): Promise<void> {
   const startTime = Date.now()
+  const processingStartedAt = new Date()
   try {
     logger.info(`[${documentId}] Starting document processing`, {
       knowledgeBaseId,
@@ -815,6 +816,7 @@ export async function processDocumentAsync(
         boolean1: document.boolean1,
         boolean2: document.boolean2,
         boolean3: document.boolean3,
+        processingStatus: document.processingStatus,
       })
       .from(document)
       .innerJoin(knowledgeBase, eq(knowledgeBase.id, document.knowledgeBaseId))
@@ -856,17 +858,30 @@ export async function processDocumentAsync(
       mimeType: ctx.mimeType,
     }
 
-    await db
+    const [claimedDocument] = await db
       .update(document)
       .set({
         processingStatus: 'processing',
-        processingStartedAt: new Date(),
+        processingStartedAt,
         processingCompletedAt: null,
         processingError: null,
       })
       .where(
-        and(eq(document.id, documentId), isNull(document.archivedAt), isNull(document.deletedAt))
+        and(
+          eq(document.id, documentId),
+          inArray(document.processingStatus, ['pending', 'failed']),
+          isNull(document.archivedAt),
+          isNull(document.deletedAt)
+        )
       )
+      .returning({ id: document.id })
+
+    if (!claimedDocument) {
+      logger.info(`[${documentId}] Skipping document processing because another attempt owns it`, {
+        processingStatus: ctx.processingStatus,
+      })
+      return
+    }
 
     logger.info(`[${documentId}] Status updated to 'processing', starting document processor`)
 
@@ -935,7 +950,13 @@ export async function processDocumentAsync(
             usageGate.message ?? 'Usage limit exceeded. Please upgrade your plan to continue.',
           processingCompletedAt: new Date(),
         })
-        .where(eq(document.id, documentId))
+        .where(
+          and(
+            eq(document.id, documentId),
+            eq(document.processingStatus, 'processing'),
+            eq(document.processingStartedAt, processingStartedAt)
+          )
+        )
       return
     }
     let billableEmbeddingTokens = 0
@@ -954,6 +975,7 @@ export async function processDocumentAsync(
       currentSourceFileProvenance
     )
 
+    let processingCommitted = false
     await withTimeout(
       runWithKnowledgeModelInputProvenance(
         documentSecretContext.registry,
@@ -1069,23 +1091,26 @@ export async function processDocumentAsync(
             updatedAt: now,
           }))
 
-          await db.transaction(async (tx) => {
+          processingCommitted = await db.transaction(async (tx) => {
             const activeDocument = await tx
               .select({ id: document.id })
               .from(document)
               .innerJoin(knowledgeBase, eq(document.knowledgeBaseId, knowledgeBase.id))
               .where(
                 and(
                   eq(document.id, documentId),
+                  eq(document.processingStatus, 'processing'),
+                  eq(document.processingStartedAt, processingStartedAt),
                   isNull(document.archivedAt),
                   isNull(document.deletedAt),
                   isNull(knowledgeBase.deletedAt)
                 )
               )
+              .for('update', { of: document })
               .limit(1)
 
             if (activeDocument.length === 0) {
-              return
+              return false
             }
 
             if (embeddingRecords.length > 0) {
@@ -1131,7 +1156,14 @@ export async function processDocumentAsync(
                 processingCompletedAt: now,
                 processingError: null,
               })
-              .where(eq(document.id, documentId))
+              .where(
+                and(
+                  eq(document.id, documentId),
+                  eq(document.processingStatus, 'processing'),
+                  eq(document.processingStartedAt, processingStartedAt)
+                )
+              )
+            return true
           })
         },
         {
@@ -1144,6 +1176,11 @@ export async function processDocumentAsync(
       'Document processing'
     )
 
+    if (!processingCommitted) {
+      logger.info(`[${documentId}] Discarded output from an obsolete processing attempt`)
+      return
+    }
+
     const processingTime = Date.now() - startTime
     logger.info(`[${documentId}] Successfully processed document in ${processingTime}ms`)
 
@@ -1205,7 +1242,13 @@ export async function processDocumentAsync(
         processingError: errorMessage,
         processingCompletedAt: new Date(),
       })
-      .where(eq(document.id, documentId))
+      .where(
+        and(
+          eq(document.id, documentId),
+          eq(document.processingStatus, 'processing'),
+          eq(document.processingStartedAt, processingStartedAt)
+        )
+      )
 
     throw error
   }