fix(knowledge): preserve processing claim ownership

Author: TheodoreSpeaks

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

@@ -230,17 +230,22 @@ describe('knowledge document processing source', () => {
     expect(mockGenerateEmbeddings).not.toHaveBeenCalled()
   })
 
-  it('does not start work when another processing attempt already owns the document', async () => {
+  it('keeps the carrier retryable when another processing attempt 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',
-    })
+    await expect(
+      processDocumentAsync('knowledge-base-1', 'document-1', {
+        filename: 'stale.pdf',
+        fileUrl: 'https://example.com/stale.pdf',
+        fileSize: 1,
+        mimeType: 'text/plain',
+      })
+    ).rejects.toThrow('processing claim is owned by another attempt')
 
     expect(mockProcessDocument).not.toHaveBeenCalled()
     expect(mockGenerateEmbeddings).not.toHaveBeenCalled()
+    expect(dbChainMockFns.set).not.toHaveBeenCalledWith(
+      expect.objectContaining({ processingStatus: 'failed' })
+    )
   })
 })

apps/sim/lib/knowledge/documents/processing-claim.test.ts

@@ -5,6 +5,7 @@
 import { dbChainMockFns, resetDbChainMock } from '@sim/testing'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 import {
+  failStaleDocumentProcessingClaim,
   KNOWLEDGE_DOCUMENT_PROCESSING_STALE_THRESHOLD_MS,
   reclaimStaleDocumentProcessingClaim,
 } from '@/lib/knowledge/documents/processing-claim'
@@ -68,3 +69,61 @@ describe('reclaimStaleDocumentProcessingClaim', () => {
     expect(reclaimed).toBe(false)
   })
 })
+
+describe('failStaleDocumentProcessingClaim', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    resetDbChainMock()
+  })
+
+  it('rejects an active processing claim', async () => {
+    await expect(
+      failStaleDocumentProcessingClaim({
+        documentId: 'document-1',
+        processingStartedAt: new Date(
+          NOW.getTime() - KNOWLEDGE_DOCUMENT_PROCESSING_STALE_THRESHOLD_MS
+        ),
+        now: NOW,
+      })
+    ).rejects.toThrow('Document has not been processing long enough to be considered dead')
+
+    expect(dbChainMockFns.set).not.toHaveBeenCalled()
+  })
+
+  it('fails the exact abandoned processing claim', async () => {
+    const processingStartedAt = new Date(
+      NOW.getTime() - KNOWLEDGE_DOCUMENT_PROCESSING_STALE_THRESHOLD_MS - 1
+    )
+    dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'document-1' }])
+
+    const result = await failStaleDocumentProcessingClaim({
+      documentId: 'document-1',
+      processingStartedAt,
+      now: NOW,
+    })
+
+    expect(result).toEqual({
+      success: true,
+      processingDuration: KNOWLEDGE_DOCUMENT_PROCESSING_STALE_THRESHOLD_MS + 1,
+    })
+    expect(dbChainMockFns.set).toHaveBeenCalledWith({
+      processingStatus: 'failed',
+      processingError: 'Processing timed out. Please retry or re-sync the connector.',
+      processingCompletedAt: NOW,
+    })
+  })
+
+  it('does not fail a replacement processing claim', async () => {
+    dbChainMockFns.returning.mockResolvedValueOnce([])
+
+    const result = await failStaleDocumentProcessingClaim({
+      documentId: 'document-1',
+      processingStartedAt: new Date(
+        NOW.getTime() - KNOWLEDGE_DOCUMENT_PROCESSING_STALE_THRESHOLD_MS - 1
+      ),
+      now: NOW,
+    })
+
+    expect(result.success).toBe(false)
+  })
+})

apps/sim/lib/knowledge/documents/processing-claim.ts

@@ -11,6 +11,12 @@ interface ReclaimStaleDocumentProcessingClaimParams {
   now?: Date
 }
 
+interface FailStaleDocumentProcessingClaimParams {
+  documentId: string
+  processingStartedAt: Date
+  now?: Date
+}
+
 /**
  * Reopens an abandoned processing attempt using its start time as a compare-and-set token.
  * The former worker's timestamp-guarded writes then cannot commit after the claim is reclaimed.
@@ -55,3 +61,36 @@ export async function reclaimStaleDocumentProcessingClaim({
 
   return Boolean(reclaimed)
 }
+
+/** Marks only the abandoned processing attempt identified by its start-time token as failed. */
+export async function failStaleDocumentProcessingClaim({
+  documentId,
+  processingStartedAt,
+  now = new Date(),
+}: FailStaleDocumentProcessingClaimParams): Promise<{
+  success: boolean
+  processingDuration: number
+}> {
+  const processingDuration = now.getTime() - processingStartedAt.getTime()
+  if (processingDuration <= KNOWLEDGE_DOCUMENT_PROCESSING_STALE_THRESHOLD_MS) {
+    throw new Error('Document has not been processing long enough to be considered dead')
+  }
+
+  const [failed] = await db
+    .update(document)
+    .set({
+      processingStatus: 'failed',
+      processingError: 'Processing timed out. Please retry or re-sync the connector.',
+      processingCompletedAt: now,
+    })
+    .where(
+      and(
+        eq(document.id, documentId),
+        eq(document.processingStatus, 'processing'),
+        eq(document.processingStartedAt, processingStartedAt)
+      )
+    )
+    .returning({ id: document.id })
+
+  return { success: Boolean(failed), processingDuration }
+}

apps/sim/lib/knowledge/documents/processing-queue.test.ts

@@ -72,6 +72,7 @@ describe('processDocumentsWithQueue billing attribution', () => {
   beforeEach(() => {
     vi.clearAllMocks()
     resetDbChainMock()
+    setEnvFlags({ isTriggerDevEnabled: true })
     mockBatchTrigger.mockResolvedValue({ batchId: 'batch-1' })
     for (const key of Object.keys(env)) {
       delete (env as Record<string, unknown>)[key]
@@ -152,4 +153,29 @@ describe('processDocumentsWithQueue billing attribution', () => {
     })
     expect(jobs[0].payload).not.toHaveProperty('billingAttribution')
   })
+
+  it('keeps direct dispatch retryable when another worker wins the document claim', async () => {
+    setEnvFlags({ isTriggerDevEnabled: false })
+    env.TRIGGER_SECRET_KEY = undefined
+    dbChainMockFns.limit
+      .mockResolvedValueOnce([{ userId: 'legacy-owner', workspaceId: null }])
+      .mockResolvedValueOnce([
+        {
+          workspaceId: null,
+          knowledgeBaseUserId: 'legacy-owner',
+          processingStatus: 'pending',
+          filename: DOCUMENT.filename,
+          fileUrl: DOCUMENT.fileUrl,
+          fileSize: DOCUMENT.fileSize,
+          mimeType: DOCUMENT.mimeType,
+        },
+      ])
+    dbChainMockFns.returning.mockResolvedValueOnce([])
+
+    await expect(
+      processDocumentsWithQueue([DOCUMENT], 'knowledge-base-1', {}, 'request-1', undefined)
+    ).rejects.toThrow('All 1 document processing dispatches failed')
+
+    expect(mockBatchTrigger).not.toHaveBeenCalled()
+  })
 })

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

@@ -64,7 +64,7 @@ import {
   mergeDurableSecretProvenance,
 } from '@/lib/execution/durable-secret-provenance'
 import { processDocument } from '@/lib/knowledge/documents/document-processor'
-import { KNOWLEDGE_DOCUMENT_PROCESSING_STALE_THRESHOLD_MS } from '@/lib/knowledge/documents/processing-claim'
+import { failStaleDocumentProcessingClaim } from '@/lib/knowledge/documents/processing-claim'
 import { enqueueKnowledgeDocumentProcessing } from '@/lib/knowledge/documents/processing-outbox-event'
 import {
   assertDocumentProcessingBillingContext,
@@ -118,6 +118,13 @@ import { calculateCost } from '@/providers/utils'
 
 const logger = createLogger('DocumentService')
 
+class DocumentProcessingClaimUnavailableError extends Error {
+  constructor(documentId: string) {
+    super(`Knowledge document ${documentId} processing claim is owned by another attempt`)
+    this.name = 'DocumentProcessingClaimUnavailableError'
+  }
+}
+
 /**
  * Thrown when a knowledge-base document's `fileUrl` references an internal
  * knowledge-base storage object not owned by the target knowledge base's workspace.
@@ -881,7 +888,7 @@ export async function processDocumentAsync(
       logger.info(`[${documentId}] Skipping document processing because another attempt owns it`, {
         processingStatus: ctx.processingStatus,
       })
-      return
+      throw new DocumentProcessingClaimUnavailableError(documentId)
     }
 
     logger.info(`[${documentId}] Status updated to 'processing', starting document processor`)
@@ -1236,6 +1243,10 @@ export async function processDocumentAsync(
       fileSize: docData.fileSize,
     })
 
+    if (error instanceof DocumentProcessingClaimUnavailableError) {
+      throw error
+    }
+
     await db
       .update(document)
       .set({
@@ -2371,30 +2382,17 @@ export async function markDocumentAsFailedTimeout(
   processingStartedAt: Date,
   requestId: string
 ): Promise<{ success: boolean; processingDuration: number }> {
-  const now = new Date()
-  const processingDuration = now.getTime() - processingStartedAt.getTime()
+  const result = await failStaleDocumentProcessingClaim({ documentId, processingStartedAt })
 
-  if (processingDuration <= KNOWLEDGE_DOCUMENT_PROCESSING_STALE_THRESHOLD_MS) {
-    throw new Error('Document has not been processing long enough to be considered dead')
+  if (result.success) {
+    logger.info(
+      `[${requestId}] Marked document ${documentId} as failed due to dead process (processing time: ${Math.round(result.processingDuration / 1000)}s)`
+    )
+  } else {
+    logger.info(`[${requestId}] Did not time out document ${documentId} because its claim changed`)
   }
 
-  await db
-    .update(document)
-    .set({
-      processingStatus: 'failed',
-      processingError: 'Processing timed out. Please retry or re-sync the connector.',
-      processingCompletedAt: now,
-    })
-    .where(eq(document.id, documentId))
-
-  logger.info(
-    `[${requestId}] Marked document ${documentId} as failed due to dead process (processing time: ${Math.round(processingDuration / 1000)}s)`
-  )
-
-  return {
-    success: true,
-    processingDuration,
-  }
+  return result
 }
 
 export async function retryDocumentProcessing(

apps/sim/lib/knowledge/orchestration/documents.test.ts

@@ -378,7 +378,13 @@ describe('performDeleteKnowledgeDocument', () => {
 })
 
 describe('document processing state changes', () => {
-  beforeEach(() => vi.clearAllMocks())
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockMarkDocumentAsFailedTimeout.mockResolvedValue({
+      success: true,
+      processingDuration: 600_001,
+    })
+  })
 
   it('refuses to time out a document that is not processing', async () => {
     const outcome = await performMarkKnowledgeDocumentTimedOut({
@@ -409,6 +415,19 @@ describe('document processing state changes', () => {
     expect(outcome).toMatchObject({ success: false, errorCode: 'validation' })
   })
 
+  it('reports a conflict when the processing claim changes before the timeout write', async () => {
+    mockMarkDocumentAsFailedTimeout.mockResolvedValue({
+      success: false,
+      processingDuration: 600_001,
+    })
+
+    const outcome = await performMarkKnowledgeDocumentTimedOut({
+      document: { id: 'doc-1', processingStatus: 'processing', processingStartedAt: new Date() },
+    })
+
+    expect(outcome).toMatchObject({ success: false, errorCode: 'conflict' })
+  })
+
   it('refuses to retry a document that has not failed', async () => {
     const outcome = await performRetryKnowledgeDocumentProcessing({
       knowledgeBaseId: 'kb-1',

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

@@ -527,7 +527,17 @@ export async function performMarkKnowledgeDocumentTimedOut(
   }
 
   try {
-    await markDocumentAsFailedTimeout(document.id, document.processingStartedAt, requestId)
+    const result = await markDocumentAsFailedTimeout(
+      document.id,
+      document.processingStartedAt,
+      requestId
+    )
+    if (!result.success) {
+      return fail(
+        'Document processing attempt changed before timeout could be recorded',
+        'conflict'
+      )
+    }
   } catch (error) {
     // The service rejects a document that has not been processing long enough
     // to be presumed dead; that is a caller-fixable "try again later", not a fault.