From 2b4488edc04a93e6b23def5f074f1ee3b69c3898 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 18 Sep 2026 17:04:06 -0700 Subject: [PATCH] fix(knowledge): preserve live document jobs during recovery --- .../background/knowledge-processing.test.ts | 23 +++- apps/sim/background/knowledge-processing.ts | 7 +- .../stored-document-recovery.integration.ts | 123 ++++++++++++++++- .../knowledge/connectors/sync-primitives.ts | 12 +- .../document-processing-source.test.ts | 38 +++++- .../documents/processing-recovery-policy.ts | 7 +- .../processing-recovery-queue.test.ts | 118 ++++++++++++++++ .../documents/processing-recovery-queue.ts | 129 ++++++++++++++++++ .../documents/processing-recovery.ts | 20 ++- apps/sim/lib/knowledge/documents/service.ts | 39 +++--- 10 files changed, 484 insertions(+), 32 deletions(-) create mode 100644 apps/sim/lib/knowledge/documents/processing-recovery-queue.test.ts create mode 100644 apps/sim/lib/knowledge/documents/processing-recovery-queue.ts diff --git a/apps/sim/background/knowledge-processing.test.ts b/apps/sim/background/knowledge-processing.test.ts index b200f9a0dbd..e9c327d988f 100644 --- a/apps/sim/background/knowledge-processing.test.ts +++ b/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({ diff --git a/apps/sim/background/knowledge-processing.ts b/apps/sim/background/knowledge-processing.ts index 3e8a6c2f81e..1964e315958 100644 --- a/apps/sim/background/knowledge-processing.ts +++ b/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, diff --git a/apps/sim/lib/knowledge/__integration__/stored-document-recovery.integration.ts b/apps/sim/lib/knowledge/__integration__/stored-document-recovery.integration.ts index 8d3a3f45a09..36720c4ff32 100644 --- a/apps/sim/lib/knowledge/__integration__/stored-document-recovery.integration.ts +++ b/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()), + 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[] = [] 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, + 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({ diff --git a/apps/sim/lib/knowledge/connectors/sync-primitives.ts b/apps/sim/lib/knowledge/connectors/sync-primitives.ts index 562378d26b0..13a1dd25139 100644 --- a/apps/sim/lib/knowledge/connectors/sync-primitives.ts +++ b/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) diff --git a/apps/sim/lib/knowledge/documents/document-processing-source.test.ts b/apps/sim/lib/knowledge/documents/document-processing-source.test.ts index 92b64b9424d..32e7eed699e 100644 --- a/apps/sim/lib/knowledge/documents/document-processing-source.test.ts +++ b/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() }) diff --git a/apps/sim/lib/knowledge/documents/processing-recovery-policy.ts b/apps/sim/lib/knowledge/documents/processing-recovery-policy.ts index 4c2bb0e9053..f30246e20b6 100644 --- a/apps/sim/lib/knowledge/documents/processing-recovery-policy.ts +++ b/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), diff --git a/apps/sim/lib/knowledge/documents/processing-recovery-queue.test.ts b/apps/sim/lib/knowledge/documents/processing-recovery-queue.test.ts new file mode 100644 index 00000000000..2a99b144b71 --- /dev/null +++ b/apps/sim/lib/knowledge/documents/processing-recovery-queue.test.ts @@ -0,0 +1,118 @@ +/** @vitest-environment node */ +import { dbChainMock, dbChainMockFns, resetDbChainMock } from '@sim/testing' +import { afterEach, beforeEach, expect, it, vi } from 'vitest' + +const mocks = vi.hoisted(() => ({ list: vi.fn(), enabled: true, insideRun: false })) +vi.mock('@sim/db', () => dbChainMock) +vi.mock('@trigger.dev/sdk', () => ({ runs: { list: mocks.list } })) +vi.mock('@/lib/core/config/env-flags', () => ({ + get isTriggerDevEnabled() { + return mocks.enabled + }, +})) +vi.mock('@/lib/core/config/trigger-runtime', () => ({ + isInsideTriggerRun: () => mocks.insideRun, +})) + +import { filterAbandonedDocumentProcessing } from '@/lib/knowledge/documents/processing-recovery-queue' + +const candidate = { + id: 'document-1', + processingQueueToken: 'generation-1', + processingQueuedAt: new Date('2026-01-01T00:00:00Z'), + processingStartedAt: null, + uploadedAt: new Date('2026-01-01T00:00:00Z'), +} + +beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + mocks.enabled = true + mocks.insideRun = false + mocks.list.mockResolvedValue({ data: [] }) +}) +afterEach(() => vi.useRealTimers()) + +it.each(['PENDING_VERSION', 'DELAYED', 'QUEUED', 'DEQUEUED', 'EXECUTING', 'WAITING'])( + 'preserves a %s run regardless of queue age, without spending another attempt', + async (status) => { + mocks.list.mockResolvedValue({ data: [{ id: 'run-1', status }] }) + expect(await filterAbandonedDocumentProcessing([candidate])).toEqual([]) + expect(mocks.list).toHaveBeenCalledWith( + expect.objectContaining({ + taskIdentifier: 'knowledge-process-document', + tag: 'documentId:document-1', + from: new Date('2025-12-31T20:00:00Z'), + status: expect.arrayContaining([status]), + limit: 1, + }), + { retry: { maxAttempts: 1 } } + ) + expect(dbChainMockFns.set).toHaveBeenCalledExactlyOnceWith({ + processingRecoveryAfter: expect.any(Date), + }) + } +) + +it('allows the existing recovery policy after confirming no live job remains', async () => { + expect(await filterAbandonedDocumentProcessing([candidate])).toEqual([candidate]) + expect(dbChainMockFns.set).not.toHaveBeenCalled() +}) + +it('fails closed and backs off when the queue cannot be inspected', async () => { + mocks.list.mockRejectedValue(new Error('unavailable')) + expect(await filterAbandonedDocumentProcessing([candidate])).toEqual([]) + expect(dbChainMockFns.set).toHaveBeenCalledExactlyOnceWith({ + processingRecoveryAfter: expect.any(Date), + }) +}) + +it('bounds concurrency and stops scheduling further lookups after the deadline', async () => { + vi.useFakeTimers() + const candidates = Array.from({ length: 200 }, (_, index) => ({ + ...candidate, + id: `doc-${index}`, + })) + let finishLookup: (value: { data: [] }) => void = () => undefined + mocks.list.mockReturnValue( + new Promise((resolve) => { + finishLookup = resolve + }) + ) + const pending = filterAbandonedDocumentProcessing(candidates) + await vi.advanceTimersByTimeAsync(10_001) + expect(await pending).toEqual([]) + expect(mocks.list).toHaveBeenCalledTimes(4) + finishLookup({ data: [] }) + await vi.runAllTimersAsync() + expect(mocks.list).toHaveBeenCalledTimes(4) +}) + +it('keeps completed checks when another lookup fails, without recovering unknown jobs', async () => { + mocks.list.mockResolvedValueOnce({ data: [] }).mockRejectedValueOnce(new Error('unavailable')) + expect( + await filterAbandonedDocumentProcessing([candidate, { ...candidate, id: 'doc-2' }]) + ).toEqual([candidate]) +}) + +it('does not look up Trigger runs on a deployment without Trigger', async () => { + mocks.enabled = false + expect(await filterAbandonedDocumentProcessing([candidate])).toEqual([candidate]) + expect(mocks.list).not.toHaveBeenCalled() +}) + +it('still inspects the queue inside a Trigger worker with a disabled environment flag', async () => { + mocks.enabled = false + mocks.insideRun = true + mocks.list.mockResolvedValue({ data: [{ id: 'run-1' }] }) + expect(await filterAbandonedDocumentProcessing([candidate])).toEqual([]) +}) + +it('does not mutate any recovery state when the caller is canceled', async () => { + const controller = new AbortController() + controller.abort(new Error('canceled')) + await expect(filterAbandonedDocumentProcessing([candidate], controller.signal)).rejects.toThrow( + 'canceled' + ) + expect(dbChainMockFns.set).not.toHaveBeenCalled() +}) diff --git a/apps/sim/lib/knowledge/documents/processing-recovery-queue.ts b/apps/sim/lib/knowledge/documents/processing-recovery-queue.ts new file mode 100644 index 00000000000..3d8f0aa241a --- /dev/null +++ b/apps/sim/lib/knowledge/documents/processing-recovery-queue.ts @@ -0,0 +1,129 @@ +import { db } from '@sim/db' +import { document } from '@sim/db/schema' +import { createLogger } from '@sim/logger' +import { getErrorMessage } from '@sim/utils/errors' +import { runs } from '@trigger.dev/sdk' +import { and, eq, isNull, or, sql } from 'drizzle-orm' +import { isTriggerDevEnabled } from '@/lib/core/config/env-flags' +import { isInsideTriggerRun } from '@/lib/core/config/trigger-runtime' +import { mapWithConcurrency } from '@/lib/core/utils/concurrency' +import { withinDeadline } from '@/lib/core/utils/deadline' +import { documentProcessingRecoveryCondition } from '@/lib/knowledge/documents/processing-recovery-policy' +import { QUEUED_DISPATCH_GRACE_MS } from '@/lib/knowledge/documents/types' + +const logger = createLogger('DocumentRecoveryQueue') +const LOOKUP_CONCURRENCY = 4 +const LOOKUP_BUDGET_MS = 10_000 +const RECHECK_DELAY_MS = 15 * 60_000 + +export interface DocumentRecoveryGeneration { + id: string + processingQueueToken: string | null + processingQueuedAt: Date | null + processingStartedAt: Date | null + uploadedAt: Date +} + +/** Recovery may replace only the generation whose queue state was inspected outside the transaction. */ +export function documentRecoveryGenerationCondition(candidate: DocumentRecoveryGeneration) { + return and( + eq(document.id, candidate.id), + candidate.processingQueueToken === null + ? isNull(document.processingQueueToken) + : eq(document.processingQueueToken, candidate.processingQueueToken), + candidate.processingQueuedAt === null + ? isNull(document.processingQueuedAt) + : eq(document.processingQueuedAt, candidate.processingQueuedAt), + candidate.processingStartedAt === null + ? isNull(document.processingStartedAt) + : eq(document.processingStartedAt, candidate.processingStartedAt) + ) +} + +/** + * Queue age is not evidence of abandonment. Any live run for the document protects + * it, including legacy dispatches and continuations. Lookup failures fail closed. + * Callers supply bounded candidate pages; only one metadata row is read per lookup. + */ +export async function filterAbandonedDocumentProcessing( + candidates: T[], + signal?: AbortSignal +): Promise { + if (!candidates.length || (!isTriggerDevEnabled && !isInsideTriggerRun())) return candidates + + const abandoned: T[] = [] + let lookupError: unknown + try { + await withinDeadline( + async (lookupSignal) => { + await mapWithConcurrency(candidates, LOOKUP_CONCURRENCY, async (candidate) => { + lookupSignal.throwIfAborted() + if (lookupError) return + const page = await runs + .list( + { + taskIdentifier: 'knowledge-process-document', + tag: `documentId:${candidate.id}`, + /** Include the entire eligible document lifetime, with the existing dispatch grace for clock skew. */ + from: new Date(candidate.uploadedAt.getTime() - QUEUED_DISPATCH_GRACE_MS), + status: [ + 'PENDING_VERSION', + 'DELAYED', + 'QUEUED', + 'DEQUEUED', + 'EXECUTING', + 'WAITING', + ], + limit: 1, + }, + { retry: { maxAttempts: 1 } } + ) + .catch((error: unknown) => { + lookupError = error + return null + }) + lookupSignal.throwIfAborted() + if (page?.data.length === 0) abandoned.push(candidate) + }) + if (lookupError) throw lookupError + }, + Date.now() + LOOKUP_BUDGET_MS, + signal + ) + } catch (error) { + signal?.throwIfAborted() + logger.warn('Could not verify all document jobs; leaving unverified generations unchanged', { + candidates: candidates.length, + abandoned: abandoned.length, + error: getErrorMessage(error), + }) + } + + signal?.throwIfAborted() + const abandonedIds = new Set(abandoned.map((candidate) => candidate.id)) + const retained = candidates.filter((candidate) => !abandonedIds.has(candidate.id)) + if (retained.length) { + try { + await db.transaction(async (tx) => { + await tx.execute( + sql`SELECT set_config('statement_timeout', '5000', true), set_config('lock_timeout', '1000', true)` + ) + signal?.throwIfAborted() + await tx + .update(document) + .set({ processingRecoveryAfter: new Date(Date.now() + RECHECK_DELAY_MS) }) + .where( + and( + documentProcessingRecoveryCondition(new Date()), + or(...retained.map(documentRecoveryGenerationCondition)) + ) + ) + signal?.throwIfAborted() + }) + } catch (error) { + signal?.throwIfAborted() + logger.warn('Could not postpone document queue recheck', { error: getErrorMessage(error) }) + } + } + return [...abandoned] +} diff --git a/apps/sim/lib/knowledge/documents/processing-recovery.ts b/apps/sim/lib/knowledge/documents/processing-recovery.ts index 4a0a457d2ac..369a8252cab 100644 --- a/apps/sim/lib/knowledge/documents/processing-recovery.ts +++ b/apps/sim/lib/knowledge/documents/processing-recovery.ts @@ -2,7 +2,7 @@ import { db } from '@sim/db' import { document, knowledgeBase, knowledgeConnector } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { generateId } from '@sim/utils/id' -import { and, asc, eq, inArray, isNull, notInArray, sql } from 'drizzle-orm' +import { and, asc, eq, inArray, isNull, notInArray, or, sql } from 'drizzle-orm' import { assertBillingAttributionOwner, resolveSystemBillingAttribution, @@ -17,6 +17,10 @@ import { createWorkspaceDocumentProcessingBillingContext, } from '@/lib/knowledge/documents/processing-payload' import { documentProcessingRecoveryCondition } from '@/lib/knowledge/documents/processing-recovery-policy' +import { + documentRecoveryGenerationCondition, + filterAbandonedDocumentProcessing, +} from '@/lib/knowledge/documents/processing-recovery-queue' const logger = createLogger('KnowledgeDocumentRecovery') @@ -29,7 +33,7 @@ const RECOVERABLE_CONNECTOR_STATUSES = ['active', 'error', 'pending', 'syncing'] /** * Re-admits bounded, abandoned connector documents from our retained bytes, independently * of source sync schedules and credentials. The generation, attempt and outbox event commit - * together; no provider call or source lease is needed. Paused/deleted sources stay paused. + * together; no source-provider call or source lease is needed. Paused/deleted sources stay paused. */ export async function recoverKnowledgeDocumentProcessing(now = new Date()): Promise { const deadlineAt = Date.now() + RECOVERY_RUNTIME_MS @@ -78,6 +82,10 @@ async function recoverStoredDocumentBatch( return tx .select({ id: document.id, + processingQueueToken: document.processingQueueToken, + processingQueuedAt: document.processingQueuedAt, + processingStartedAt: document.processingStartedAt, + uploadedAt: document.uploadedAt, knowledgeBaseId: document.knowledgeBaseId, connectorId: knowledgeConnector.id, workspaceId: knowledgeBase.workspaceId, @@ -107,9 +115,11 @@ async function recoverStoredDocumentBatch( signal.throwIfAborted() if (candidates.length === 0) return 0 + const abandoned = await filterAbandonedDocumentProcessing(candidates, signal) + for (const candidate of candidates) attemptedConnectors.add(candidate.connectorId) let recovered = 0 const groups = new Map() - for (const candidate of candidates) { + for (const candidate of abandoned) { const group = groups.get(candidate.knowledgeBaseId) ?? [] group.push(candidate) groups.set(candidate.knowledgeBaseId, group) @@ -117,7 +127,6 @@ async function recoverStoredDocumentBatch( for (const [knowledgeBaseId, group] of groups) { if (Date.now() >= deadlineAt) break const connectorIds = [...new Set(group.map((row) => row.connectorId))] - for (const connectorId of connectorIds) attemptedConnectors.add(connectorId) const owner = group[0] let ownerVerified = false try { @@ -188,7 +197,8 @@ async function recoverStoredDocumentBatch( document.connectorId, connectors.map((row) => row.id) ), - documentProcessingRecoveryCondition(now) + documentProcessingRecoveryCondition(now), + or(...group.map(documentRecoveryGenerationCondition)) ) ) .orderBy(asc(document.id)) diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts index 7b1bbef023b..2a90e1b4038 100644 --- a/apps/sim/lib/knowledge/documents/service.ts +++ b/apps/sim/lib/knowledge/documents/service.ts @@ -1447,6 +1447,22 @@ function queueGenerationConditions( : [isNull(document.processingQueueToken)] } +/** + * Who the processor reads a document's source file as. Always the actor, not + * the payer: authorizing as the KB owner would let a writer ingest an internal + * file only the owner can read. A connector-owned row was written by the sync + * from bytes it fetched, not from a caller-supplied URL, so it is read as the + * system: in members mode the row stays hidden until the sync materializes who + * observed it, and the actor's own scope would deny the read. + */ +function sourceFileAccessFor(connectorId: string | null, actorUserId: string): SourceFileAccess { + return { userId: actorUserId, knowledgeAccess: connectorId ? SYSTEM_ACCESS_SCOPE : undefined } +} + +export type DocumentProcessingResult = + | { outcome: 'indexed' } + | { outcome: 'skipped'; reason: 'unavailable' | 'not_claimed' | 'superseded' } + /** * Parses, embeds, and indexes one document. * @@ -1461,18 +1477,6 @@ function queueGenerationConditions( * invocation against the document's retry budget. Direct callers omit it and * therefore cannot refund an attempt they never charged. */ -/** - * Who the processor reads a document's source file as. Always the actor, not - * the payer: authorizing as the KB owner would let a writer ingest an internal - * file only the owner can read. A connector-owned row was written by the sync - * from bytes it fetched, not from a caller-supplied URL, so it is read as the - * system: in members mode the row stays hidden until the sync materializes who - * observed it, and the actor's own scope would deny the read. - */ -function sourceFileAccessFor(connectorId: string | null, actorUserId: string): SourceFileAccess { - return { userId: actorUserId, knowledgeAccess: connectorId ? SYSTEM_ACCESS_SCOPE : undefined } -} - export async function processDocumentAsync( knowledgeBaseId: string, documentId: string, @@ -1486,7 +1490,7 @@ export async function processDocumentAsync( providedBillingContext?: BillingAttributionSnapshot | DocumentProcessingBillingContext, indexingPassId?: string, attemptContext?: DocumentProcessingAttemptContext -): Promise { +): Promise { const startTime = Date.now() const processingStartedAt = new Date() let processingFilename = docData.filename @@ -1570,12 +1574,12 @@ export async function processDocumentAsync( documentConnectorIsActive() ) ) - return + return { outcome: 'skipped', reason: 'unavailable' } } const ctx = contextRows[0] processingFilename = ctx.filename - await withResourceOutboundScope(ctx, async () => { + return await withResourceOutboundScope(ctx, async (): Promise => { const persistedDocData = { filename: ctx.filename, fileUrl: ctx.fileUrl, @@ -1645,7 +1649,7 @@ export async function processDocumentAsync( logger.info( `[${documentId}] Skipping document processing: superseded, already active, completed, archived, or deleted` ) - return + return { outcome: 'skipped', reason: 'not_claimed' } } attemptContext?.onClaimed?.() @@ -2003,7 +2007,7 @@ export async function processDocumentAsync( if (!processingCommitted) { logger.info(`[${documentId}] Discarded output from an obsolete processing attempt`) - return + return { outcome: 'skipped', reason: 'superseded' } } const processingTime = Date.now() - startTime @@ -2074,6 +2078,7 @@ export async function processDocumentAsync( logger.error(`[${documentId}] Failed to record embedding usage`, { error: billingError }) } } + return { outcome: 'indexed' } }) } catch (error) { const processingTime = Date.now() - startTime