fix(knowledge): preserve live document jobs during recovery (#7994)

Author: waleedlatif1

apps/sim/background/knowledge-processing.test.ts

@@ -119,7 +119,7 @@ describe('knowledge processing worker', () => {
       }
       return value
     })
-    mockProcessDocumentAsync.mockResolvedValue(undefined)
+    mockProcessDocumentAsync.mockResolvedValue({ outcome: 'indexed' })
     mockResolveTriggerRegion.mockResolvedValue('us-east-1')
     mockTrigger.mockResolvedValue({ id: 'quota-continuation-run' })
   })
@@ -128,6 +128,27 @@ describe('knowledge processing worker', () => {
     vi.restoreAllMocks()
   })
 
+  it('reports indexed only when the document service committed the index', async () => {
+    expect(await runDocumentProcessing(WORKSPACE_PAYLOAD)).toMatchObject({
+      success: true,
+      outcome: 'indexed',
+      documentId: WORKSPACE_PAYLOAD.documentId,
+    })
+  })
+
+  it.each(['unavailable', 'not_claimed', 'superseded'] as const)(
+    'reports a harmless %s skip without turning it into a task failure or an indexed success',
+    async (reason) => {
+      mockProcessDocumentAsync.mockResolvedValue({ outcome: 'skipped', reason })
+      expect(await runDocumentProcessing(WORKSPACE_PAYLOAD)).toMatchObject({
+        success: false,
+        outcome: 'skipped',
+        reason,
+      })
+      expect(mockTrigger).not.toHaveBeenCalled()
+    }
+  )
+
   it('rejects workspace work without attribution before document processing starts', async () => {
     await expect(
       runDocumentProcessing({

apps/sim/background/knowledge-processing.ts

@@ -58,7 +58,7 @@ export async function runDocumentProcessing(
   logger.info(`[${requestId}] Starting Trigger.dev processing for document: ${docData.filename}`)
 
   try {
-    await processDocumentAsync(
+    const result = await processDocumentAsync(
       knowledgeBaseId,
       documentId,
       docData,
@@ -90,10 +90,11 @@ export async function runDocumentProcessing(
       }
     )
 
-    logger.info(`[${requestId}] Successfully processed document: ${docData.filename}`)
+    logger.info(`[${requestId}] Document processing finished`, { documentId, ...result })
 
     return {
-      success: true,
+      success: result.outcome === 'indexed',
+      ...result,
       documentId,
       filename: docData.filename,
       processingTime: Date.now() - startedAt,

apps/sim/lib/knowledge/__integration__/stored-document-recovery.integration.ts

@@ -17,9 +17,21 @@ import {
 } from '@sim/db/schema'
 import { generateId } from '@sim/utils/id'
 import { and, eq, inArray, sql } from 'drizzle-orm'
-import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'
+import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from 'vitest'
 
-const fixture = vi.hoisted(() => ({ root: '', embeddingCalls: 0 }))
+const fixture = vi.hoisted(() => ({
+  root: '',
+  embeddingCalls: 0,
+  queueEnabled: false,
+  listRuns: vi.fn(),
+}))
+vi.mock('@/lib/core/config/trigger-runtime', () => ({
+  isInsideTriggerRun: () => fixture.queueEnabled,
+}))
+vi.mock('@trigger.dev/sdk', async (original) => ({
+  ...(await original<typeof import('@trigger.dev/sdk')>()),
+  runs: { list: fixture.listRuns },
+}))
 vi.mock('@/lib/uploads/core/setup.server', () => ({
   get UPLOAD_DIR_SERVER() {
     return fixture.root
@@ -52,6 +64,7 @@ import { searchKnowledge } from '@/lib/knowledge/application/search'
 import { searchScopedKnowledge } from '@/lib/knowledge/application/workspace-search'
 import { createContentSyncLease } from '@/lib/knowledge/connectors/sync-lock'
 import { addDocument } from '@/lib/knowledge/connectors/sync-persistence'
+import { sweepStuckDocuments } from '@/lib/knowledge/connectors/sync-primitives'
 import { knowledgeDocumentProcessingOutboxHandlers } from '@/lib/knowledge/documents/processing-outbox-handler'
 import {
   DOCUMENT_RECOVERY_BATCH_SIZE,
@@ -60,6 +73,7 @@ import {
 } from '@/lib/knowledge/documents/processing-recovery'
 import { processDocumentAsync } from '@/lib/knowledge/documents/service'
 import { MAX_PROCESSING_ATTEMPTS, QUEUED_DISPATCH_GRACE_MS } from '@/lib/knowledge/documents/types'
+import type { SyncResult } from '@/connectors/types'
 
 const fixtures: ReturnType<typeof createKnowledgeAclFixtureIds>[] = []
 const old = () => new Date(Date.now() - QUEUED_DISPATCH_GRACE_MS - 60_000)
@@ -119,6 +133,11 @@ async function failedFile(
   return file
 }
 
+afterEach(() => {
+  fixture.queueEnabled = false
+  fixture.listRuns.mockReset()
+})
+
 beforeAll(() => {
   fixture.root = mkdtempSync(path.join(tmpdir(), 'sim-stored-recovery-'))
 })
@@ -137,7 +156,107 @@ afterAll(async () => {
   await db.$client.end()
 })
 
+async function recoverFixture(
+  ids: ReturnType<typeof createKnowledgeAclFixtureIds>,
+  mode: 'independent' | 'connector'
+) {
+  if (mode === 'independent') return recoverKnowledgeDocumentProcessing()
+  const result: SyncResult = {
+    docsAdded: 0,
+    docsUpdated: 0,
+    docsDeleted: 0,
+    docsUnchanged: 0,
+    docsSkipped: 0,
+    docsFailed: 0,
+    processingDispatch: { requested: 0, accepted: 0, failed: 0 },
+  }
+  await sweepStuckDocuments({
+    connectorId: ids.connectorId,
+    knowledgeBaseId: ids.knowledgeBaseId,
+    syncStartedAt: new Date(),
+    retryCutoff: new Date(Date.now() - 7 * 24 * 60 * 60_000),
+    billingAttribution: await resolveSystemBillingAttribution(ids.workspaceId),
+    result,
+    lease: createContentSyncLease(ids.connectorId, ids.lockId),
+  })
+  return result.processingDispatch.requested
+}
+
 describe('independent recovery of retained connector documents', () => {
+  it.each(['independent', 'connector'] as const)(
+    '%s recovery preserves a job queued beyond the grace period',
+    async (mode) => {
+      const ids = await seed()
+      const file = await failedFile(ids)
+      await db
+        .update(document)
+        .set({ processingStatus: 'pending' })
+        .where(eq(document.id, file.documentId))
+      fixture.queueEnabled = true
+      fixture.listRuns.mockResolvedValue({ data: [{ id: 'run-queued', status: 'QUEUED' }] })
+      expect(await recoverFixture(ids, mode)).toBe(0)
+      const [row] = await db.select().from(document).where(eq(document.id, file.documentId))
+      expect(row.processingAttempts).toBe(1)
+      expect(row.processingQueueToken).toBe('old-fixture-generation')
+      expect(row.processingRecoveryAfter).not.toBeNull()
+      expect(await eventsFor(ids)).toHaveLength(0)
+      await db
+        .update(knowledgeConnector)
+        .set({ status: 'paused' })
+        .where(eq(knowledgeConnector.id, ids.connectorId))
+    }
+  )
+
+  it.each(['independent', 'connector'] as const)(
+    '%s recovery rechecks the generation after its remote lookup',
+    async (mode) => {
+      const ids = await seed()
+      const file = await failedFile(ids)
+      fixture.queueEnabled = true
+      fixture.listRuns.mockImplementation(async () => {
+        await db
+          .update(document)
+          .set({ processingQueueToken: 'replacement-generation' })
+          .where(eq(document.id, file.documentId))
+        return { data: [] }
+      })
+      expect(await recoverFixture(ids, mode)).toBe(0)
+      const [row] = await db.select().from(document).where(eq(document.id, file.documentId))
+      expect(row.processingAttempts).toBe(1)
+      expect(row.processingQueueToken).toBe('replacement-generation')
+      expect(await eventsFor(ids)).toHaveLength(0)
+      await db
+        .update(knowledgeConnector)
+        .set({ status: 'paused' })
+        .where(eq(knowledgeConnector.id, ids.connectorId))
+    }
+  )
+
+  it.each(['pending', 'processing'])(
+    'does not replace an aged %s outbox continuation',
+    async (status) => {
+      const ids = await seed()
+      const file = await failedFile(ids)
+      const token = generateId()
+      await db
+        .update(document)
+        .set({ processingQueueToken: token })
+        .where(eq(document.id, file.documentId))
+      await db.insert(outboxEvent).values({
+        id: token,
+        eventType: 'knowledge.document.processing.resume',
+        payload: { knowledgeBaseId: ids.knowledgeBaseId, documentId: file.documentId },
+        status,
+        availableAt: old(),
+      })
+      expect(await recoverFixture(ids, 'independent')).toBe(0)
+      expect(await recoverFixture(ids, 'connector')).toBe(0)
+      const [row] = await db.select().from(document).where(eq(document.id, file.documentId))
+      expect(row.processingAttempts).toBe(1)
+      expect(row.processingQueueToken).toBe(token)
+    }
+  )
+
   it('uses the organization owner and preserves Search visibility during source backoff', async () => {
     const ids = await seed()
     await db.insert(member).values({

apps/sim/lib/knowledge/connectors/sync-primitives.ts

@@ -9,7 +9,7 @@ import {
 import { createLogger } from '@sim/logger'
 import { toError } from '@sim/utils/errors'
 import { generateId } from '@sim/utils/id'
-import { and, asc, desc, eq, inArray, isNotNull, isNull, lt, ne, sql } from 'drizzle-orm'
+import { and, asc, desc, eq, inArray, isNotNull, isNull, lt, ne, or, sql } from 'drizzle-orm'
 import type { BillingAttributionSnapshot } from '@/lib/billing/core/billing-attribution'
 import { ProviderCapacityDeferredError } from '@/lib/core/rate-limiter/provider-capacity-error'
 import { withDatabaseReadRetry } from '@/lib/db/read-retry'
@@ -28,6 +28,10 @@ import {
   updateDocument,
 } from '@/lib/knowledge/connectors/sync-persistence'
 import { documentProcessingRecoveryCondition } from '@/lib/knowledge/documents/processing-recovery-policy'
+import {
+  documentRecoveryGenerationCondition,
+  filterAbandonedDocumentProcessing,
+} from '@/lib/knowledge/documents/processing-recovery-queue'
 import { DOCUMENT_PROCESSING_STALE_THRESHOLD_MS } from '@/lib/knowledge/documents/processing-timeouts.server'
 import type { DocumentData } from '@/lib/knowledge/documents/service'
 import { isTriggerAvailable, processDocumentsWithQueue } from '@/lib/knowledge/documents/service'
@@ -1336,6 +1340,7 @@ export async function sweepStuckDocuments(input: SweepStuckDocumentsInput): Prom
       fileSize: document.fileSize,
       mimeType: document.mimeType,
       processingStatus: document.processingStatus,
+      processingQueueToken: document.processingQueueToken,
       processingQueuedAt: document.processingQueuedAt,
       processingStartedAt: document.processingStartedAt,
       processingDeferredUntil: document.processingDeferredUntil,
@@ -1361,7 +1366,8 @@ export async function sweepStuckDocuments(input: SweepStuckDocumentsInput): Prom
       asc(document.id)
     )
     .limit(STUCK_RETRY_MAX_CANDIDATES_PER_SYNC)
-  const stuckDocs = sweepCandidates.filter(
+  const abandoned = await filterAbandonedDocumentProcessing(sweepCandidates)
+  const stuckDocs = abandoned.filter(
     (row): row is typeof row & { processingStatus: DocumentProcessingStatus } =>
       isDocumentProcessingStatus(row.processingStatus)
   )
@@ -1402,6 +1408,7 @@ export async function sweepStuckDocuments(input: SweepStuckDocumentsInput): Prom
           fileSize: document.fileSize,
           mimeType: document.mimeType,
           processingStatus: document.processingStatus,
+          processingQueueToken: document.processingQueueToken,
           processingQueuedAt: document.processingQueuedAt,
           processingStartedAt: document.processingStartedAt,
           processingDeferredUntil: document.processingDeferredUntil,
@@ -1412,6 +1419,7 @@ export async function sweepStuckDocuments(input: SweepStuckDocumentsInput): Prom
         .where(
           and(
             inArray(document.id, stuckDocIds),
+            or(...stuckDocs.map(documentRecoveryGenerationCondition)),
             eq(document.connectorId, connectorId),
             documentProcessingRecoveryCondition(sweepEvaluatedAt, retryCutoff),
             lt(document.uploadedAt, syncStartedAt)

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

@@ -1460,11 +1460,46 @@ describe('processDocumentAsync write guards', () => {
     expect(schedule).not.toHaveBeenCalled()
   })
 
+  it('reports an unavailable document without parsing or indexing it', async () => {
+    dbChainMockFns.limit.mockResolvedValueOnce([])
+    expect(
+      await processDocumentAsync(
+        'knowledge-base-1',
+        'document-1',
+        PERSISTED_CONTEXT,
+        {},
+        BILLING_ATTRIBUTION
+      )
+    ).toEqual({ outcome: 'skipped', reason: 'unavailable' })
+    expect(mockProcessDocument).not.toHaveBeenCalled()
+  })
+
+  it('reports discarded output when the generation changes before the index commit', async () => {
+    armProviderSource()
+    dbChainMockFns.limit.mockReset()
+    dbChainMockFns.limit
+      .mockResolvedValueOnce([PERSISTED_CONTEXT])
+      .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW])
+      .mockResolvedValueOnce([])
+    expect(
+      await processDocumentAsync(
+        'knowledge-base-1',
+        'document-1',
+        PERSISTED_CONTEXT,
+        {},
+        BILLING_ATTRIBUTION
+      )
+    ).toEqual({ outcome: 'skipped', reason: 'superseded' })
+    expect(
+      dbChainMockFns.set.mock.calls.some(([value]) => value.processingStatus === 'completed')
+    ).toBe(false)
+  })
+
   it('does not parse or reschedule a superseded provider continuation', async () => {
     armProviderSource()
     dbChainMockFns.returning.mockResolvedValueOnce([])
     const schedule = vi.fn()
-    await processDocumentAsync(
+    const result = await processDocumentAsync(
       'knowledge-base-1',
       'document-1',
       PERSISTED_CONTEXT,
@@ -1477,6 +1512,7 @@ describe('processDocumentAsync write guards', () => {
         scheduleProviderContinuation: schedule,
       }
     )
+    expect(result).toEqual({ outcome: 'skipped', reason: 'not_claimed' })
     expect(mockProcessDocument).not.toHaveBeenCalled()
     expect(schedule).not.toHaveBeenCalled()
   })

apps/sim/lib/knowledge/documents/processing-recovery-policy.ts

@@ -1,4 +1,4 @@
-import { document } from '@sim/db/schema'
+import { document, outboxEvent } from '@sim/db/schema'
 import { and, eq, gt, isNotNull, isNull, lt, lte, or, sql } from 'drizzle-orm'
 import { DOCUMENT_PROCESSING_STALE_THRESHOLD_MS } from '@/lib/knowledge/documents/processing-timeouts.server'
 import { MAX_PROCESSING_ATTEMPTS, QUEUED_DISPATCH_GRACE_MS } from '@/lib/knowledge/documents/types'
@@ -15,6 +15,11 @@ export function documentProcessingRecoveryCondition(
   return and(
     sql`${document.processingStatus} IN ('pending', 'processing', 'failed')`,
     isNotNull(document.connectorId),
+    sql`NOT EXISTS (
+      SELECT 1 FROM ${outboxEvent}
+      WHERE ${outboxEvent.id} = ${document.processingQueueToken}
+        AND ${outboxEvent.status} IN ('pending', 'processing')
+    )`,
     isNotNull(document.contentHash),
     isNotNull(document.storageKey),
     eq(document.userExcluded, false),