From 254f8c2e22666f63707a99e16b7aec97a056f819 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 18 Sep 2026 16:40:35 -0700 Subject: [PATCH 1/3] fix(knowledge): preserve live document processing during recovery --- apps/sim/lib/internal/mistral/client.test.ts | 49 ++++- apps/sim/lib/internal/mistral/client.ts | 12 +- .../mistral/error-diagnostics.test.ts | 49 +++++ .../lib/internal/mistral/error-diagnostics.ts | 57 +++++ .../stored-document-recovery.integration.ts | 204 +++++++++++++++++- .../knowledge/connectors/sync-primitives.ts | 28 ++- .../knowledge/documents/document-processor.ts | 7 + .../documents/pdf-ocr-triage.test.ts | 39 ++++ .../documents/processing-liveness.test.ts | 137 ++++++++++++ .../documents/processing-liveness.ts | 188 ++++++++++++++++ .../documents/processing-queue.test.ts | 182 ++-------------- .../documents/processing-recovery.ts | 17 +- .../documents/retry-processing-grace.test.ts | 13 ++ apps/sim/lib/knowledge/documents/service.ts | 163 +++++++------- .../google-service-account-transport.test.ts | 47 ++++ .../oauth/google-service-account-transport.ts | 32 ++- 16 files changed, 937 insertions(+), 287 deletions(-) create mode 100644 apps/sim/lib/internal/mistral/error-diagnostics.test.ts create mode 100644 apps/sim/lib/internal/mistral/error-diagnostics.ts create mode 100644 apps/sim/lib/knowledge/documents/processing-liveness.test.ts create mode 100644 apps/sim/lib/knowledge/documents/processing-liveness.ts diff --git a/apps/sim/lib/internal/mistral/client.test.ts b/apps/sim/lib/internal/mistral/client.test.ts index 774882e107d..62bf449afd5 100644 --- a/apps/sim/lib/internal/mistral/client.test.ts +++ b/apps/sim/lib/internal/mistral/client.test.ts @@ -3,6 +3,11 @@ */ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +const { errorLog } = vi.hoisted(() => ({ errorLog: vi.fn() })) +vi.mock('@sim/logger', () => ({ + createLogger: () => ({ error: errorLog, info: vi.fn(), warn: vi.fn(), debug: vi.fn() }), +})) + const { fetchPinned, admit, settle, validate } = vi.hoisted(() => ({ fetchPinned: vi.fn(), admit: vi.fn(), @@ -63,9 +68,41 @@ describe('Mistral provider transport', () => { expect(settle).toHaveBeenCalledWith('success', undefined) }) + it('records cooldown without waiting for a stalled 429 error body', async () => { + const cancel = vi.fn() + fetchPinned.mockResolvedValue( + new Response(new ReadableStream({ cancel }), { + status: 429, + headers: { 'retry-after': '60' }, + }) + ) + await expect(submitMistralOcr('private-key', {})).rejects.toMatchObject({ + reason: 'rate_limit', + retryAfterMs: 60_000, + }) + expect(settle).toHaveBeenCalledWith('rate_limit', 60_000) + expect(cancel).toHaveBeenCalledOnce() + expect(fetchPinned).toHaveBeenCalledOnce() + }) + + it('preserves provider rejection when its diagnostic body exceeds the byte limit', async () => { + fetchPinned.mockResolvedValue(new Response('x'.repeat(70_000), { status: 400 })) + await expect(submitMistralOcr('private-key', {})).rejects.toMatchObject({ + status: 400, + body: { success: false, error: 'Mistral API error: HTTP 400' }, + }) + expect(errorLog).toHaveBeenCalledWith( + 'Mistral API error', + expect.objectContaining({ status: 400, bodyFormat: 'unavailable' }) + ) + }) + it('identifies provider request rejection without retaining echoed document contents', async () => { fetchPinned.mockResolvedValue( - Response.json({ message: 'Sensitive fixture document text' }, { status: 400 }) + Response.json( + { type: 'invalid_request_error', code: 400, message: 'Sensitive fixture document text' }, + { status: 400, headers: { 'x-request-id': 'ocr-request-123' } } + ) ) await expect(submitMistralOcr('key', {})).rejects.toMatchObject({ source: 'provider', @@ -74,6 +111,16 @@ describe('Mistral provider transport', () => { }) expect(fetchPinned).toHaveBeenCalledOnce() expect(settle).toHaveBeenCalledWith('failure', undefined) + expect(errorLog).toHaveBeenCalledWith( + 'Mistral API error', + expect.objectContaining({ + status: 400, + providerRequestId: 'ocr-request-123', + providerErrorCode: '400', + providerErrorType: 'invalid_request_error', + }) + ) + expect(JSON.stringify(errorLog.mock.calls)).not.toContain('Sensitive fixture document text') }) it.each([ diff --git a/apps/sim/lib/internal/mistral/client.ts b/apps/sim/lib/internal/mistral/client.ts index 31ecc0f2ef1..010cf8ade43 100644 --- a/apps/sim/lib/internal/mistral/client.ts +++ b/apps/sim/lib/internal/mistral/client.ts @@ -11,9 +11,10 @@ import { validateUrlWithDNS, } from '@/lib/core/security/input-validation.server' import { getMistralCapacityConfig, getMistralCapacityScope } from '@/lib/internal/mistral/capacity' +import { getOcrResponseDiagnostic } from '@/lib/internal/mistral/error-diagnostics' import { MistralOperationError } from '@/lib/internal/mistral/errors' import { MISTRAL_OCR_REQUEST_POLICY } from '@/lib/knowledge/documents/ocr-request-policy' -import { readBoundedHttpErrorBody, resolveRetryDelayMs } from '@/lib/knowledge/documents/utils' +import { readBoundedHttpErrorPayload, resolveRetryDelayMs } from '@/lib/knowledge/documents/utils' const logger = createLogger('MistralClient') const MISTRAL_ENDPOINT = 'https://api.mistral.ai/v1/ocr' @@ -138,8 +139,13 @@ export async function submitMistralOcr( retryAfterMs, }) } - await readBoundedHttpErrorBody(response) - logger.error('Mistral API error', { status: response.status }) + const payload = await readBoundedHttpErrorPayload(response) + logger.error('Mistral API error', { + provider: 'mistral', + operation: 'ocr', + status: response.status, + ...getOcrResponseDiagnostic(response.headers, payload.ok ? payload.body : ''), + }) throw new MistralOperationError( response.status, { success: false, error: `Mistral API error: HTTP ${response.status}` }, diff --git a/apps/sim/lib/internal/mistral/error-diagnostics.test.ts b/apps/sim/lib/internal/mistral/error-diagnostics.test.ts new file mode 100644 index 00000000000..46001741011 --- /dev/null +++ b/apps/sim/lib/internal/mistral/error-diagnostics.test.ts @@ -0,0 +1,49 @@ +/** @vitest-environment node */ +import { describe, expect, it } from 'vitest' +import { getOcrResponseDiagnostic } from '@/lib/internal/mistral/error-diagnostics' + +describe('OCR error diagnostics', () => { + it.each([ + { code: 400, type: 'invalid_request_error', message: 'private document' }, + { error: { code: 400, type: 'invalid_request_error', message: 'private document' } }, + ])('projects only safe fields from a provider envelope', (body) => { + expect( + getOcrResponseDiagnostic(new Headers({ 'x-request-id': 'request_123' }), JSON.stringify(body)) + ).toEqual({ + bodyFormat: 'json', + providerRequestId: 'request_123', + providerErrorCode: '400', + providerErrorType: 'invalid_request_error', + }) + }) + + it.each(['not json', 'private input', '', 'null', '[]', '42'])( + 'tolerates an unavailable or non-object error: %s', + (body) => { + expect(getOcrResponseDiagnostic(new Headers(), body)).toMatchObject({ + providerRequestId: null, + providerErrorCode: null, + providerErrorType: null, + }) + } + ) + + it('does not trust arbitrary codes, error types, echoed input, or request IDs', () => { + const diagnostic = getOcrResponseDiagnostic( + new Headers({ 'x-request-id': 'secret '.repeat(30), 'apim-request-id': 'safe-123' }), + JSON.stringify({ + code: 'private_document', + type: { input: 'private_document' }, + message: 'private_document', + param: 'secret-key', + detail: { input: 'private_document' }, + }) + ) + expect(diagnostic).toMatchObject({ + providerRequestId: 'safe-123', + providerErrorCode: 'unrecognized', + providerErrorType: 'unrecognized', + }) + expect(JSON.stringify(diagnostic)).not.toMatch(/private_document|secret-key/) + }) +}) diff --git a/apps/sim/lib/internal/mistral/error-diagnostics.ts b/apps/sim/lib/internal/mistral/error-diagnostics.ts new file mode 100644 index 00000000000..5f2df64d35c --- /dev/null +++ b/apps/sim/lib/internal/mistral/error-diagnostics.ts @@ -0,0 +1,57 @@ +const SAFE_ERROR_CODES = new Set([ + 'invalid_request_error', + 'authentication_error', + 'permission_error', + 'rate_limit_error', + 'server_error', + 'unknown_model', + 'BadRequest', + 'InvalidRequest', + 'DeploymentNotFound', + 'ResourceNotFound', + 'OperationNotSupported', + 'Unauthorized', + 'Forbidden', + 'TooManyRequests', + 'InternalServerError', + 'ServiceUnavailable', +]) + +function safeCode(value: unknown): string | null { + if (value === undefined || value === null) return null + if (typeof value === 'string' && SAFE_ERROR_CODES.has(value)) return value + if (typeof value === 'number' && Number.isInteger(value) && value >= 400 && value <= 599) + return String(value) + return 'unrecognized' +} + +function safeRequestId(value: string | null): string | null { + return value && /^[a-zA-Z0-9_-]{1,128}$/.test(value) ? value : null +} + +/** Projects bounded OCR error responses without retaining messages, document data or URLs. */ +export function getOcrResponseDiagnostic(headers: Pick, body: string) { + const diagnostic = { + providerRequestId: + safeRequestId(headers.get('x-request-id')) ?? + safeRequestId(headers.get('apim-request-id')) ?? + safeRequestId(headers.get('x-ms-request-id')), + bodyFormat: body ? 'non_json' : 'unavailable', + providerErrorCode: null as string | null, + providerErrorType: null as string | null, + } + if (!body) return diagnostic + let parsed: unknown + try { + parsed = JSON.parse(body) + } catch { + return diagnostic + } + diagnostic.bodyFormat = 'json' + if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) return diagnostic + const error = 'error' in parsed ? parsed.error : parsed + if (!error || typeof error !== 'object' || Array.isArray(error)) return diagnostic + diagnostic.providerErrorCode = safeCode('code' in error ? error.code : undefined) + diagnostic.providerErrorType = safeCode('type' in error ? error.type : undefined) + return diagnostic +} 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..034a9319a6c 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, + useTrigger: false, + listRuns: vi.fn(), +})) +vi.mock('@/lib/core/config/trigger-runtime', () => ({ + isInsideTriggerRun: () => fixture.useTrigger, +})) +vi.mock('@trigger.dev/sdk', async (importOriginal) => ({ + ...(await importOriginal()), + runs: { list: fixture.listRuns }, +})) vi.mock('@/lib/uploads/core/setup.server', () => ({ get UPLOAD_DIR_SERVER() { return fixture.root @@ -52,13 +64,14 @@ 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, KNOWLEDGE_DOCUMENT_RECOVERY_OUTBOX_EVENT, recoverKnowledgeDocumentProcessing, } from '@/lib/knowledge/documents/processing-recovery' -import { processDocumentAsync } from '@/lib/knowledge/documents/service' +import { processDocumentAsync, retryDocumentProcessing } from '@/lib/knowledge/documents/service' import { MAX_PROCESSING_ATTEMPTS, QUEUED_DISPATCH_GRACE_MS } from '@/lib/knowledge/documents/types' const fixtures: ReturnType[] = [] @@ -119,6 +132,11 @@ async function failedFile( return file } +afterEach(() => { + fixture.useTrigger = false + fixture.listRuns.mockReset() +}) + beforeAll(() => { fixture.root = mkdtempSync(path.join(tmpdir(), 'sim-stored-recovery-')) }) @@ -138,6 +156,186 @@ afterAll(async () => { }) describe('independent recovery of retained connector documents', () => { + it.each(['QUEUED', 'DELAYED', 'WAITING'])( + 'preserves a seven-hour %s run and its existing attempt', + async (status) => { + const ids = await seed() + const file = await failedFile(ids) + const queuedAt = new Date(Date.now() - 7 * 60 * 60_000) + await db + .update(document) + .set({ + processingStatus: 'pending', + processingQueuedAt: queuedAt, + processingCompletedAt: null, + }) + .where(eq(document.id, file.documentId)) + fixture.useTrigger = true + fixture.listRuns.mockResolvedValue({ data: [{ status }], hasNextPage: () => false }) + expect(await recoverKnowledgeDocumentProcessing()).toBe(0) + const [row] = await db.select().from(document).where(eq(document.id, file.documentId)) + expect(row).toMatchObject({ + processingStatus: 'pending', + processingAttempts: 1, + processingQueueToken: 'old-fixture-generation', + processingQueuedAt: queuedAt, + }) + expect(row.processingRecoveryAfter!.getTime()).toBeGreaterThan(Date.now()) + expect(await eventsFor(ids)).toHaveLength(0) + await db + .update(knowledgeConnector) + .set({ status: 'paused' }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + } + ) + + it.each(['pending', 'processing'])( + 'preserves a %s outbox carrier without replacing its generation', + async (status) => { + const ids = await seed() + const file = await failedFile(ids) + const token = generateId() + await db + .update(document) + .set({ processingStatus: 'pending', processingQueueToken: token }) + .where(eq(document.id, file.documentId)) + await db.insert(outboxEvent).values({ + id: token, + status, + eventType: 'knowledge.document.processing', + payload: { knowledgeBaseId: ids.knowledgeBaseId }, + }) + expect(await recoverKnowledgeDocumentProcessing()).toBe(0) + const [row] = await db.select().from(document).where(eq(document.id, file.documentId)) + expect(row.processingQueueToken).toBe(token) + expect(row.processingAttempts).toBe(1) + expect(await eventsFor(ids)).toHaveLength(0) + await db + .update(knowledgeConnector) + .set({ status: 'paused' }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + } + ) + + it.each(['processing', 'completed', 'pending'])( + 'does not replace a concurrent %s transition after inspecting work', + async (status) => { + const ids = await seed() + const file = await failedFile(ids) + fixture.useTrigger = true + fixture.listRuns.mockImplementation(async () => { + await db + .update(document) + .set({ + processingStatus: status, + processingQueueToken: 'winning-generation', + processingStartedAt: new Date(), + }) + .where(eq(document.id, file.documentId)) + return { data: [], hasNextPage: () => false } + }) + expect(await recoverKnowledgeDocumentProcessing()).toBe(0) + const [row] = await db.select().from(document).where(eq(document.id, file.documentId)) + expect(row.processingQueueToken).toBe('winning-generation') + expect(row.processingAttempts).toBe(1) + expect(await eventsFor(ids)).toHaveLength(0) + await db + .update(knowledgeConnector) + .set({ status: 'paused' }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + } + ) + + it.each(['manual', 'connector'] as const)( + 'preserves live work through the %s retry path', + async (path) => { + const ids = await seed() + const file = await failedFile(ids) + await db + .update(document) + .set({ processingStatus: 'pending' }) + .where(eq(document.id, file.documentId)) + const [original] = await db.select().from(document).where(eq(document.id, file.documentId)) + fixture.useTrigger = true + fixture.listRuns.mockResolvedValue({ data: [{ status: 'QUEUED' }], hasNextPage: () => false }) + const billing = await resolveSystemBillingAttribution(ids.workspaceId) + if (path === 'manual') { + const result = await retryDocumentProcessing( + ids.knowledgeBaseId, + file.documentId, + { + filename: original.filename, + fileUrl: original.fileUrl, + fileSize: original.fileSize, + mimeType: original.mimeType, + }, + generateId(), + billing + ) + expect(result.message).toContain('already queued') + } else { + const result = { + 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: billing, + result, + lease: createContentSyncLease(ids.connectorId, ids.lockId), + }) + expect(result.processingDispatch.requested).toBe(0) + } + const [row] = await db.select().from(document).where(eq(document.id, file.documentId)) + expect(row.processingAttempts).toBe(original.processingAttempts) + expect(row.processingQueueToken).toBe(original.processingQueueToken) + expect(await eventsFor(ids)).toHaveLength(0) + await db + .update(knowledgeConnector) + .set({ status: 'paused' }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + } + ) + + it('continues past a protected batch while its cooldown is active', async () => { + const ids = await seed() + const file = await failedFile(ids) + const [original] = await db.select().from(document).where(eq(document.id, file.documentId)) + await db.insert(document).values( + Array.from({ length: DOCUMENT_RECOVERY_BATCH_SIZE - 1 }, () => ({ + ...original, + id: generateId(), + externalId: generateId(), + secretProvenanceVersion: null, + })) + ) + const abandoned = await failedFile(ids) + await db + .update(document) + .set({ uploadedAt: new Date(original.uploadedAt.getTime() + 1) }) + .where(eq(document.id, abandoned.documentId)) + fixture.useTrigger = true + fixture.listRuns.mockImplementation(async ({ tag }: { tag: string }) => ({ + data: tag === `documentId:${abandoned.documentId}` ? [] : [{ status: 'QUEUED' }], + hasNextPage: () => false, + })) + expect(await recoverKnowledgeDocumentProcessing()).toBe(0) + expect(await recoverKnowledgeDocumentProcessing()).toBe(1) + expect((await eventsFor(ids))[0].payload).toMatchObject({ documentId: abandoned.documentId }) + await db + .update(knowledgeConnector) + .set({ status: 'paused' }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + }) + 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..230b3802a60 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' @@ -27,6 +27,12 @@ import { persistSkippedRetryHashes, updateDocument, } from '@/lib/knowledge/connectors/sync-persistence' +import { + DOCUMENT_LIVENESS_BATCH_SIZE, + documentProcessingSnapshotCondition, + findAbandonedDocumentProcessing, + processingSnapshotColumns, +} from '@/lib/knowledge/documents/processing-liveness' import { documentProcessingRecoveryCondition } from '@/lib/knowledge/documents/processing-recovery-policy' import { DOCUMENT_PROCESSING_STALE_THRESHOLD_MS } from '@/lib/knowledge/documents/processing-timeouts.server' import type { DocumentData } from '@/lib/knowledge/documents/service' @@ -1330,16 +1336,11 @@ export async function sweepStuckDocuments(input: SweepStuckDocumentsInput): Prom const sweepEvaluatedAt = new Date() const sweepCandidates = await db .select({ - id: document.id, + ...processingSnapshotColumns, fileUrl: document.fileUrl, filename: document.filename, fileSize: document.fileSize, mimeType: document.mimeType, - processingStatus: document.processingStatus, - processingQueuedAt: document.processingQueuedAt, - processingStartedAt: document.processingStartedAt, - processingDeferredUntil: document.processingDeferredUntil, - processingCompletedAt: document.processingCompletedAt, uploadedAt: document.uploadedAt, }) .from(document) @@ -1360,8 +1361,9 @@ export async function sweepStuckDocuments(input: SweepStuckDocumentsInput): Prom END`), asc(document.id) ) - .limit(STUCK_RETRY_MAX_CANDIDATES_PER_SYNC) - const stuckDocs = sweepCandidates.filter( + .limit(Math.min(STUCK_RETRY_MAX_CANDIDATES_PER_SYNC, DOCUMENT_LIVENESS_BATCH_SIZE)) + const abandonedCandidates = await findAbandonedDocumentProcessing(sweepCandidates) + const stuckDocs = abandonedCandidates.filter( (row): row is typeof row & { processingStatus: DocumentProcessingStatus } => isDocumentProcessingStatus(row.processingStatus) ) @@ -1396,22 +1398,18 @@ export async function sweepStuckDocuments(input: SweepStuckDocumentsInput): Prom const lockedCandidates = await tx .select({ - id: document.id, + ...processingSnapshotColumns, fileUrl: document.fileUrl, filename: document.filename, fileSize: document.fileSize, mimeType: document.mimeType, - processingStatus: document.processingStatus, - processingQueuedAt: document.processingQueuedAt, - processingStartedAt: document.processingStartedAt, - processingDeferredUntil: document.processingDeferredUntil, - processingCompletedAt: document.processingCompletedAt, uploadedAt: document.uploadedAt, }) .from(document) .where( and( inArray(document.id, stuckDocIds), + or(...stuckDocs.map(documentProcessingSnapshotCondition)), eq(document.connectorId, connectorId), documentProcessingRecoveryCondition(sweepEvaluatedAt, retryCutoff), lt(document.uploadedAt, syncStartedAt) diff --git a/apps/sim/lib/knowledge/documents/document-processor.ts b/apps/sim/lib/knowledge/documents/document-processor.ts index b285dd68d09..eb00b2807f3 100644 --- a/apps/sim/lib/knowledge/documents/document-processor.ts +++ b/apps/sim/lib/knowledge/documents/document-processor.ts @@ -38,6 +38,7 @@ import { FileParserError, isFileParserError } from '@/lib/file-parsers/errors' import { openPdfDocument } from '@/lib/file-parsers/pdfjs-server' import type { FileParseMetadata, FileParseResult } from '@/lib/file-parsers/types' import { getMistralOcrPagesPerRequest } from '@/lib/internal/mistral/capacity' +import { getOcrResponseDiagnostic } from '@/lib/internal/mistral/error-diagnostics' import { MistralOperationError } from '@/lib/internal/mistral/errors' import { mistralParseInputSchema } from '@/lib/internal/mistral/input' import { executeMistralParse } from '@/lib/internal/mistral/operations' @@ -651,6 +652,12 @@ async function makeOCRRequest( } if (!response.ok) { + logger.warn('OCR provider request failed', { + provider: 'azure-mistral', + operation: 'ocr', + status: response.status, + ...getOcrResponseDiagnostic(response.headers, responseText), + }) if ([400, 415, 422].includes(response.status)) { throw new OcrRequestRejectedError(response.status) } diff --git a/apps/sim/lib/knowledge/documents/pdf-ocr-triage.test.ts b/apps/sim/lib/knowledge/documents/pdf-ocr-triage.test.ts index 11fa8e213b8..9af849f6f21 100644 --- a/apps/sim/lib/knowledge/documents/pdf-ocr-triage.test.ts +++ b/apps/sim/lib/knowledge/documents/pdf-ocr-triage.test.ts @@ -8,6 +8,11 @@ */ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +const { warnLog } = vi.hoisted(() => ({ warnLog: vi.fn() })) +vi.mock('@sim/logger', () => ({ + createLogger: () => ({ warn: warnLog, error: vi.fn(), info: vi.fn(), debug: vi.fn() }), +})) + const { mockParseBuffer, mockDownload, mockToken, mockBaseUrl, mockExecuteMistralParse } = vi.hoisted(() => ({ mockParseBuffer: vi.fn(), @@ -528,6 +533,40 @@ describe('PDF OCR triage', () => { } ) + it.each([400, 415, 422])( + 'records safe Azure HTTP %i diagnostics without retrying a rejection', + async (status) => { + Object.assign(env, { + OCR_PROVIDER: 'azure-mistral', + OCR_AZURE_API_KEY: 'test-key', + OCR_AZURE_ENDPOINT: 'https://example.openai.azure.com', + OCR_AZURE_MODEL_NAME: 'mistral-document-ai-2512', + }) + mockParseBuffer.mockResolvedValue({ content: '', metadata: {} }) + const fetch = vi + .fn() + .mockResolvedValue( + Response.json( + { error: { code: 'BadRequest', message: 'private fixture text' } }, + { status, headers: { 'apim-request-id': 'azure-request-123' } } + ) + ) + vi.stubGlobal('fetch', fetch) + await expect(parse()).rejects.toBeInstanceOf(OcrRequestRejectedError) + expect(fetch).toHaveBeenCalledOnce() + expect(warnLog).toHaveBeenCalledWith( + 'OCR provider request failed', + expect.objectContaining({ + provider: 'azure-mistral', + status, + providerErrorCode: 'BadRequest', + providerRequestId: 'azure-request-123', + }) + ) + expect(JSON.stringify(warnLog.mock.calls)).not.toContain('private fixture text') + } + ) + it('keeps an internal request-building failure distinct from provider rejection', async () => { mockParseBuffer.mockResolvedValue({ content: '', metadata: {} }) mockExecuteMistralParse.mockRejectedValue(new MistralOperationError(400, {})) diff --git a/apps/sim/lib/knowledge/documents/processing-liveness.test.ts b/apps/sim/lib/knowledge/documents/processing-liveness.test.ts new file mode 100644 index 00000000000..a9f46cac6ae --- /dev/null +++ b/apps/sim/lib/knowledge/documents/processing-liveness.test.ts @@ -0,0 +1,137 @@ +/** @vitest-environment node */ +import { dbChainMockFns, resetDbChainMock, resetEnvFlagsMock, setEnvFlags } from '@sim/testing' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const { listRuns } = vi.hoisted(() => ({ listRuns: vi.fn() })) +vi.mock('@trigger.dev/sdk', () => ({ runs: { list: listRuns } })) + +import { env } from '@/lib/core/config/env' +import { resetInsideTriggerRunForTests } from '@/lib/core/config/trigger-runtime' +import { + DOCUMENT_LIVENESS_BATCH_SIZE, + type DocumentProcessingSnapshot, + findAbandonedDocumentProcessing, +} from '@/lib/knowledge/documents/processing-liveness' + +const snapshot: DocumentProcessingSnapshot = { + id: 'doc-1', + processingStatus: 'pending', + processingQueueToken: 'generation-1', + processingQueuedAt: new Date('2026-09-01T00:00:00Z'), + processingStartedAt: null, + processingDeferredUntil: null, + processingCompletedAt: null, +} +const originalSecret = env.TRIGGER_SECRET_KEY +beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + resetInsideTriggerRunForTests() + setEnvFlags({ isTriggerDevEnabled: true }) + env.TRIGGER_SECRET_KEY = 'test-secret' + listRuns.mockResolvedValue({ data: [], hasNextPage: () => false }) +}) +afterEach(() => { + env.TRIGGER_SECRET_KEY = originalSecret + resetEnvFlagsMock() + vi.useRealTimers() +}) + +describe('document processing liveness', () => { + it.each(['QUEUED', 'DELAYED', 'WAITING', 'EXECUTING', 'PENDING_VERSION', 'DEQUEUED'])( + 'preserves an old document with a %s job without spending an attempt', + async (status) => { + listRuns.mockResolvedValue({ data: [{ status }], hasNextPage: () => false }) + expect(await findAbandonedDocumentProcessing([snapshot])).toEqual([]) + expect(dbChainMockFns.set).toHaveBeenCalledWith({ processingRecoveryAfter: expect.any(Date) }) + expect(dbChainMockFns.delete).not.toHaveBeenCalled() + expect(listRuns).toHaveBeenCalledWith( + expect.objectContaining({ + tag: 'documentId:doc-1', + limit: 1, + status: expect.arrayContaining([ + 'QUEUED', + 'WAITING', + 'DELAYED', + 'EXECUTING', + 'DEQUEUED', + 'PENDING_VERSION', + ]), + }), + { retry: { maxAttempts: 1 } } + ) + } + ) + + it('allows recovery only when no live run or outbox carrier remains', async () => { + expect(await findAbandonedDocumentProcessing([snapshot])).toEqual([snapshot]) + expect(dbChainMockFns.set).not.toHaveBeenCalled() + }) + + it('protects a pending outbox delivery without contacting Trigger', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([{ id: snapshot.processingQueueToken }]) + expect(await findAbandonedDocumentProcessing([snapshot])).toEqual([]) + expect(listRuns).not.toHaveBeenCalled() + }) + + it('protects legacy jobs without a generation token', async () => { + listRuns.mockResolvedValue({ data: [{ status: 'QUEUED' }], hasNextPage: () => false }) + expect( + await findAbandonedDocumentProcessing([{ ...snapshot, processingQueueToken: null }]) + ).toEqual([]) + }) + + it.each(['rejected', 'incomplete'])('fails closed on %s job evidence', async (mode) => { + if (mode === 'rejected') listRuns.mockRejectedValue(new Error('provider unavailable')) + else listRuns.mockResolvedValue({ data: [], hasNextPage: () => true }) + expect(await findAbandonedDocumentProcessing([snapshot])).toEqual([]) + expect(dbChainMockFns.set).toHaveBeenCalledWith({ processingRecoveryAfter: expect.any(Date) }) + }) + + it('does not substitute missing outbox evidence for abandonment', async () => { + dbChainMockFns.limit.mockRejectedValueOnce(new Error('database unavailable')) + expect(await findAbandonedDocumentProcessing([snapshot])).toEqual([]) + expect(listRuns).not.toHaveBeenCalled() + }) + + it('bounds stalled lookups and starts no more than four requests after the deadline', async () => { + vi.useFakeTimers() + listRuns.mockImplementation(() => new Promise(() => {})) + const candidates = Array.from({ length: 20 }, (_, i) => ({ ...snapshot, id: `doc-${i}` })) + const result = findAbandonedDocumentProcessing(candidates) + await vi.advanceTimersByTimeAsync(8_000) + expect(await result).toEqual([]) + expect(listRuns).toHaveBeenCalledTimes(4) + }) + + it('retains recovery on installations using the in-process fallback', async () => { + env.TRIGGER_SECRET_KEY = undefined + expect(await findAbandonedDocumentProcessing([snapshot])).toEqual([snapshot]) + expect(listRuns).not.toHaveBeenCalled() + }) + + it('refuses an unbounded candidate batch before reading external state', async () => { + await expect( + findAbandonedDocumentProcessing( + Array.from({ length: DOCUMENT_LIVENESS_BATCH_SIZE + 1 }, () => snapshot) + ) + ).rejects.toThrow('exceeds its limit') + expect(listRuns).not.toHaveBeenCalled() + expect(dbChainMockFns.select).not.toHaveBeenCalled() + }) + + it('keeps live work protected even when persisting its cooldown fails', async () => { + listRuns.mockResolvedValue({ data: [{ status: 'QUEUED' }], hasNextPage: () => false }) + dbChainMockFns.update.mockImplementationOnce(() => { + throw new Error('database lock timeout') + }) + expect(await findAbandonedDocumentProcessing([snapshot])).toEqual([]) + }) + + it('preserves caller cancellation without resetting a generation', async () => { + const signal = AbortSignal.abort(new Error('cancelled')) + await expect(findAbandonedDocumentProcessing([snapshot], signal)).rejects.toBe(signal.reason) + expect(listRuns).not.toHaveBeenCalled() + expect(dbChainMockFns.set).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/lib/knowledge/documents/processing-liveness.ts b/apps/sim/lib/knowledge/documents/processing-liveness.ts new file mode 100644 index 00000000000..98d6612ec3d --- /dev/null +++ b/apps/sim/lib/knowledge/documents/processing-liveness.ts @@ -0,0 +1,188 @@ +import { db } from '@sim/db' +import { document, outboxEvent } from '@sim/db/schema' +import { createLogger } from '@sim/logger' +import type { RunStatus } from '@trigger.dev/core/v3' +import { runs } from '@trigger.dev/sdk' +import { and, eq, inArray, or, sql } from 'drizzle-orm' +import { env } from '@/lib/core/config/env' +import { isTriggerDevEnabled } from '@/lib/core/config/env-flags' +import { isInsideTriggerRun } from '@/lib/core/config/trigger-runtime' +import { withinDeadline } from '@/lib/core/utils/deadline' + +const logger = createLogger('DocumentProcessingLiveness') +export const DOCUMENT_LIVENESS_BATCH_SIZE = 200 +const LOOKUP_CONCURRENCY = 4 +const LOOKUP_BUDGET_MS = 8_000 +const LIVE_RECHECK_MS = 15 * 60_000 +const UNKNOWN_RECHECK_MS = 60_000 +type TriggerRunStatus = Extract + +/** Exhaustive against the SDK: adding a lifecycle state requires classifying it here. */ +const RUN_STATUS_LIVENESS = { + PENDING_VERSION: true, + QUEUED: true, + DEQUEUED: true, + EXECUTING: true, + WAITING: true, + DELAYED: true, + COMPLETED: false, + CANCELED: false, + FAILED: false, + CRASHED: false, + SYSTEM_FAILURE: false, + EXPIRED: false, + TIMED_OUT: false, +} satisfies Record +const ACTIVE_RUN_STATUSES = (Object.keys(RUN_STATUS_LIVENESS) as TriggerRunStatus[]).filter( + (status) => RUN_STATUS_LIVENESS[status] +) + +export const processingSnapshotColumns = { + id: document.id, + processingStatus: document.processingStatus, + processingQueueToken: document.processingQueueToken, + processingQueuedAt: document.processingQueuedAt, + processingStartedAt: document.processingStartedAt, + processingDeferredUntil: document.processingDeferredUntil, + processingCompletedAt: document.processingCompletedAt, +} + +export type DocumentProcessingSnapshot = Pick< + typeof document.$inferSelect, + keyof typeof processingSnapshotColumns +> + +/** A claim or continuation installed during the external lookup must win over recovery. */ +export function documentProcessingSnapshotCondition(snapshot: DocumentProcessingSnapshot) { + return and( + eq(document.id, snapshot.id), + eq(document.processingStatus, snapshot.processingStatus), + sql`${document.processingQueueToken} IS NOT DISTINCT FROM ${snapshot.processingQueueToken}`, + sql`${document.processingQueuedAt} IS NOT DISTINCT FROM ${sql.param(snapshot.processingQueuedAt, document.processingQueuedAt)}`, + sql`${document.processingStartedAt} IS NOT DISTINCT FROM ${sql.param(snapshot.processingStartedAt, document.processingStartedAt)}`, + sql`${document.processingDeferredUntil} IS NOT DISTINCT FROM ${sql.param(snapshot.processingDeferredUntil, document.processingDeferredUntil)}`, + sql`${document.processingCompletedAt} IS NOT DISTINCT FROM ${sql.param(snapshot.processingCompletedAt, document.processingCompletedAt)}` + ) +} + +type ProcessingLiveness = 'live' | 'abandoned' | 'unknown' + +async function inspectTriggerWork(documentId: string, deadlineAt: number, signal?: AbortSignal) { + return withinDeadline( + async () => { + const page = await runs.list( + { + taskIdentifier: ['knowledge-process-document'], + tag: `documentId:${documentId}`, + status: ACTIVE_RUN_STATUSES, + limit: 1, + }, + { retry: { maxAttempts: 1 } } + ) + if (page.data.length > 0) return 'live' as const + return page.hasNextPage() ? ('unknown' as const) : ('abandoned' as const) + }, + deadlineAt, + signal + ) +} + +/** + * Age only nominates candidates. Inspect durable work before replacing its generation, + * outside row locks. Any live document run protects continuation handoffs and legacy jobs. + * Failed or incomplete lookups defer recovery; they never authorize another admission. + */ +export async function findAbandonedDocumentProcessing( + candidates: readonly T[], + signal?: AbortSignal +): Promise { + if (candidates.length === 0) return [] + if (candidates.length > DOCUMENT_LIVENESS_BATCH_SIZE) { + throw new Error('Document liveness batch exceeds its limit') + } + signal?.throwIfAborted() + const deadlineAt = Date.now() + LOOKUP_BUDGET_MS + const states = new Map() + try { + const tokens = candidates.flatMap((row) => + row.processingQueueToken ? [row.processingQueueToken] : [] + ) + const carriers = + tokens.length === 0 + ? [] + : await db.transaction(async (tx) => { + signal?.throwIfAborted() + await tx.execute( + sql`SELECT set_config('statement_timeout', '2000', true), set_config('lock_timeout', '500', true)` + ) + signal?.throwIfAborted() + return tx + .select({ id: outboxEvent.id }) + .from(outboxEvent) + .where( + and( + inArray(outboxEvent.id, tokens), + inArray(outboxEvent.status, ['pending', 'processing']) + ) + ) + .limit(DOCUMENT_LIVENESS_BATCH_SIZE) + }) + const liveTokens = new Set(carriers.map((row) => row.id)) + let next = 0 + await Promise.all( + Array.from({ length: Math.min(LOOKUP_CONCURRENCY, candidates.length) }, async () => { + while (next < candidates.length && Date.now() < deadlineAt && !signal?.aborted) { + const candidate = candidates[next++]! + if (candidate.processingQueueToken && liveTokens.has(candidate.processingQueueToken)) { + states.set(candidate.id, 'live') + continue + } + if (!(isInsideTriggerRun() || (isTriggerDevEnabled && env.TRIGGER_SECRET_KEY))) { + states.set(candidate.id, 'abandoned') + continue + } + try { + states.set(candidate.id, await inspectTriggerWork(candidate.id, deadlineAt, signal)) + } catch { + states.set(candidate.id, 'unknown') + } + } + }) + ) + } catch { + /** Missing outbox evidence cannot establish abandonment either. */ + } + signal?.throwIfAborted() + + for (const state of ['live', 'unknown'] as const) { + const protectedRows = candidates.filter((row) => (states.get(row.id) ?? 'unknown') === state) + if (protectedRows.length === 0) continue + try { + await db.transaction(async (tx) => { + signal?.throwIfAborted() + await tx.execute( + sql`SELECT set_config('statement_timeout', '2000', true), set_config('lock_timeout', '500', true)` + ) + signal?.throwIfAborted() + await tx + .update(document) + .set({ + processingRecoveryAfter: new Date( + Date.now() + (state === 'live' ? LIVE_RECHECK_MS : UNKNOWN_RECHECK_MS) + ), + }) + .where(or(...protectedRows.map(documentProcessingSnapshotCondition))) + }) + } catch { + signal?.throwIfAborted() + logger.warn('Document recovery cooldown could not be persisted', { + count: protectedRows.length, + }) + } + if (state === 'unknown') + logger.warn('Document recovery deferred: work status could not be established', { + count: protectedRows.length, + }) + } + return candidates.filter((row) => states.get(row.id) === 'abandoned') +} diff --git a/apps/sim/lib/knowledge/documents/processing-queue.test.ts b/apps/sim/lib/knowledge/documents/processing-queue.test.ts index 64debf56f43..14554f79e67 100644 --- a/apps/sim/lib/knowledge/documents/processing-queue.test.ts +++ b/apps/sim/lib/knowledge/documents/processing-queue.test.ts @@ -22,7 +22,6 @@ import { } from '@/lib/core/config/trigger-runtime' import { SyncLockLostException } from '@/lib/knowledge/connectors/sync-lock' import { DOCUMENT_PROCESSING_STALE_THRESHOLD_MS } from '@/lib/knowledge/documents/processing-timeouts.server' -import { QUEUED_DISPATCH_GRACE_MS } from '@/lib/knowledge/documents/types' const { mockBatchTrigger, mockResolveTriggerRegion } = vi.hoisted(() => ({ mockBatchTrigger: vi.fn(), @@ -261,53 +260,6 @@ describe('processDocumentsWithQueue dispatch backend', () => { return dbChainMockFns.where.mock.calls[whereIndex]?.[0] } - function resumeAlternatives(guard: unknown): MockCondition[] { - const alternatives = flattenMockConditions(guard).find( - (node) => - node.type === 'or' && - (node.conditions as MockCondition[]).some((condition) => - hasMockCondition( - condition, - (nested) => - nested.type === 'eq' && - nested.left === schemaMock.document.processingQueueToken && - nested.right === 'request-1' - ) - ) && - (node.conditions as MockCondition[]).some((condition) => - hasMockCondition( - condition, - (nested) => - nested.type === 'isNull' && nested.column === schemaMock.document.processingQueueToken - ) - ) - )?.conditions - expect(alternatives).toBeDefined() - expect(alternatives).toHaveLength(2) - const conditions = alternatives as MockCondition[] - expect( - conditions.filter((condition) => - hasMockCondition( - condition, - (node) => - node.type === 'eq' && - node.left === schemaMock.document.processingQueueToken && - node.right === 'request-1' - ) - ) - ).toHaveLength(1) - expect( - conditions.filter((condition) => - hasMockCondition( - condition, - (node) => - node.type === 'isNull' && node.column === schemaMock.document.processingQueueToken - ) - ) - ).toHaveLength(1) - return conditions - } - beforeEach(() => { vi.clearAllMocks() resetDbChainMock() @@ -573,12 +525,7 @@ describe('processDocumentsWithQueue dispatch backend', () => { (node: MockCondition) => node.type === 'gte' && node.left === schemaMock.document.processingQueuedAt ) - ).toBe(true) - const queuedFreshness = flattenMockConditions(pendingWithQueueState).find( - (node: MockCondition) => - node.type === 'gte' && node.left === schemaMock.document.processingQueuedAt - ) - expect(queuedFreshness?.right).toEqual(new Date(now.getTime() - QUEUED_DISPATCH_GRACE_MS)) + ).toBe(false) const liveProcessingState = acceptedStatuses.find( (condition) => condition.type === 'and' && @@ -652,52 +599,30 @@ describe('processDocumentsWithQueue dispatch backend', () => { node.type === 'isNull' && node.column === schemaMock.document.processingDeferredUntil ) ).toBe(true) - const sameTokenBranch = resumeAlternatives(resumeGuard).find((condition) => + expect( hasMockCondition( - condition, + resumeGuard, (node: MockCondition) => node.type === 'eq' && node.left === schemaMock.document.processingQueueToken && node.right === 'request-1' ) - ) - expect(sameTokenBranch).toBeDefined() + ).toBe(true) expect( hasMockCondition( resumeGuard, (node: MockCondition) => - node.type === 'isNotNull' && node.column === schemaMock.document.processingQueuedAt + node.type === 'inArray' && + node.column === schemaMock.document.processingStatus && + JSON.stringify(node.values) === JSON.stringify(['pending', 'failed']) ) ).toBe(true) - const statusGuard = flattenMockConditions(sameTokenBranch).find( - (node: MockCondition) => node.type === 'or' - ) - expect(statusGuard).toBeDefined() - const statusConditions = statusGuard?.conditions as MockCondition[] - expect(statusConditions).toEqual( - expect.arrayContaining([ - expect.objectContaining({ - type: 'eq', - left: schemaMock.document.processingStatus, - right: 'pending', - }), - expect.objectContaining({ - type: 'eq', - left: schemaMock.document.processingStatus, - right: 'failed', - }), - ]) - ) }) - it('treats a recent legacy queued-at-only row as live without redispatching it', async () => { - vi.useFakeTimers() - const now = new Date('2026-08-24T22:00:00.000Z') - vi.setSystemTime(now) + it('preserves an old legacy generation for the liveness-aware recovery path', async () => { markInsideTriggerRun() dbChainMockFns.returning.mockResolvedValueOnce([]).mockResolvedValueOnce([]) queueTableRows(schemaMock.document, [{ id: 'document-1' }]) - const result = await processDocumentsWithQueue( [DOCUMENT], 'knowledge-base-1', @@ -706,98 +631,15 @@ describe('processDocumentsWithQueue dispatch backend', () => { BILLING_ATTRIBUTION, 'interactive' ) - expect(result).toEqual({ requested: 1, accepted: 1, failed: 0, failedDocumentIds: [] }) expect(mockBatchTrigger).not.toHaveBeenCalled() - const resumeWrite = dbChainMockFns.set.mock.calls.find( - (call) => - (call[0] as Record | undefined)?.processingQueueToken === 'request-1' && - !('processingQueuedAt' in ((call[0] as Record | undefined) ?? {})) - ) - expect(resumeWrite).toBeDefined() - const resumeGuard = guardForResumeWrite() - const legacyBranch = resumeAlternatives(resumeGuard).find( - (condition) => - hasMockCondition( - condition, - (node: MockCondition) => - node.type === 'isNull' && node.column === schemaMock.document.processingQueueToken - ) && - hasMockCondition( - condition, - (node: MockCondition) => - node.type === 'lt' && node.left === schemaMock.document.processingQueuedAt - ) - ) - expect(legacyBranch).toBeDefined() - const cutoff = flattenMockConditions(legacyBranch).find( - (node: MockCondition) => - node.type === 'lt' && node.left === schemaMock.document.processingQueuedAt - ) - expect(cutoff?.right).toEqual(new Date(now.getTime() - QUEUED_DISPATCH_GRACE_MS)) - }) - - it('CAS-adopts a stale legacy queued-at-only row without charging again', async () => { - markInsideTriggerRun() - const legacyQueuedAt = new Date('2020-01-01T00:00:00.000Z') - dbChainMockFns.returning - .mockResolvedValueOnce([]) - .mockResolvedValueOnce([{ id: 'document-1', processingQueuedAt: legacyQueuedAt }]) - - const result = await processDocumentsWithQueue( - [DOCUMENT], - 'knowledge-base-1', - {}, - 'request-1', - BILLING_ATTRIBUTION, - 'interactive' - ) - - expect(result).toEqual({ requested: 1, accepted: 1, failed: 0, failedDocumentIds: [] }) - expect(mockBatchTrigger.mock.calls[0][1][0].payload).toMatchObject({ - processingQueueToken: 'request-1', - processingQueuedAt: legacyQueuedAt.toISOString(), - chargedAtDispatch: false, - }) - - const legacyAdoptionGuard = guardForResumeWrite() - const legacyBranch = resumeAlternatives(legacyAdoptionGuard).find( - (condition) => - hasMockCondition( - condition, - (node: MockCondition) => - node.type === 'isNull' && node.column === schemaMock.document.processingQueueToken - ) && - hasMockCondition( - condition, - (node: MockCondition) => - node.type === 'lt' && node.left === schemaMock.document.processingQueuedAt - ) - ) - expect(legacyBranch).toBeDefined() expect( hasMockCondition( - legacyAdoptionGuard, + guardForResumeWrite(), (node: MockCondition) => - node.type === 'eq' && - node.left === schemaMock.document.knowledgeBaseId && - node.right === 'knowledge-base-1' + node.type === 'isNull' && node.column === schemaMock.document.processingQueueToken ) - ).toBe(true) - for (const column of [schemaMock.document.archivedAt, schemaMock.document.deletedAt]) { - expect( - hasMockCondition( - legacyAdoptionGuard, - (node: MockCondition) => node.type === 'isNull' && node.column === column - ) - ).toBe(true) - } - const adoptionWrite = dbChainMockFns.set.mock.calls.find( - (call) => - (call[0] as Record | undefined)?.processingQueueToken === 'request-1' && - !('processingQueuedAt' in ((call[0] as Record | undefined) ?? {})) - ) - expect(adoptionWrite?.[0]).not.toHaveProperty('processingAttempts') + ).toBe(false) }) it('keeps a pre-claim same-request fallback failure retryable without clearing its stamp', async () => { @@ -1059,7 +901,7 @@ describe('processDocumentsWithQueue dispatch backend', () => { (node: MockCondition) => node.type === 'gte' && node.left === schemaMock.document.processingQueuedAt ) - expect(queuedFreshness?.right).toEqual(new Date(now.getTime() - QUEUED_DISPATCH_GRACE_MS)) + expect(queuedFreshness).toBeUndefined() const processingState = acceptedStatuses.find( (condition) => condition.type === 'and' && diff --git a/apps/sim/lib/knowledge/documents/processing-recovery.ts b/apps/sim/lib/knowledge/documents/processing-recovery.ts index 4a0a457d2ac..4a6df1bdce6 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, @@ -11,6 +11,12 @@ import { import { enqueueOutboxEvent } from '@/lib/core/outbox/service' import { withinDeadline } from '@/lib/core/utils/deadline' import { getConnectorFailureDiagnostic } from '@/lib/knowledge/connectors/connector-error' +import { + DOCUMENT_LIVENESS_BATCH_SIZE, + documentProcessingSnapshotCondition, + findAbandonedDocumentProcessing, + processingSnapshotColumns, +} from '@/lib/knowledge/documents/processing-liveness' import { createDocumentProcessingPayload, createOrganizationDocumentProcessingBillingContext, @@ -69,7 +75,7 @@ async function recoverStoredDocumentBatch( limit: number ): Promise { /** Discovery does not claim work. Ownership is rechecked under lifecycle locks below. */ - const candidates = await db.transaction(async (tx) => { + const observedCandidates = await db.transaction(async (tx) => { signal.throwIfAborted() await tx.execute( sql`SELECT set_config('statement_timeout', '5000', true), set_config('lock_timeout', '1000', true)` @@ -77,7 +83,7 @@ async function recoverStoredDocumentBatch( signal.throwIfAborted() return tx .select({ - id: document.id, + ...processingSnapshotColumns, knowledgeBaseId: document.knowledgeBaseId, connectorId: knowledgeConnector.id, workspaceId: knowledgeBase.workspaceId, @@ -102,9 +108,11 @@ async function recoverStoredDocumentBatch( ) ) .orderBy(asc(document.uploadedAt), asc(document.id)) - .limit(limit) + .limit(Math.min(limit, DOCUMENT_LIVENESS_BATCH_SIZE)) }) signal.throwIfAborted() + for (const candidate of observedCandidates) attemptedConnectors.add(candidate.connectorId) + const candidates = await findAbandonedDocumentProcessing(observedCandidates, signal) if (candidates.length === 0) return 0 let recovered = 0 @@ -183,6 +191,7 @@ async function recoverStoredDocumentBatch( document.id, group.map((row) => row.id) ), + or(...group.map(documentProcessingSnapshotCondition)), eq(document.knowledgeBaseId, knowledgeBaseId), inArray( document.connectorId, diff --git a/apps/sim/lib/knowledge/documents/retry-processing-grace.test.ts b/apps/sim/lib/knowledge/documents/retry-processing-grace.test.ts index 4c9db45ba91..bcee1680b23 100644 --- a/apps/sim/lib/knowledge/documents/retry-processing-grace.test.ts +++ b/apps/sim/lib/knowledge/documents/retry-processing-grace.test.ts @@ -26,6 +26,16 @@ import { } from '@/lib/knowledge/documents/service' import { QUEUED_DISPATCH_GRACE_MS } from '@/lib/knowledge/documents/types' +const OBSERVED_DOCUMENT = { + id: 'doc-1', + processingStatus: 'completed', + processingQueueToken: 'old-token', + processingQueuedAt: new Date(0), + processingStartedAt: null, + processingDeferredUntil: null, + processingCompletedAt: new Date(0), +} + const DOC_DATA = { filename: 'report.pdf', fileUrl: 'https://example.com/report.pdf', @@ -66,6 +76,7 @@ describe('retryDocumentProcessing requeue stamp', () => { beforeEach(() => { vi.clearAllMocks() resetDbChainMock() + dbChainMockFns.limit.mockResolvedValueOnce([OBSERVED_DOCUMENT]) }) it('clears the previous attempt terminal state', async () => { @@ -171,6 +182,7 @@ describe('retryDocumentProcessing requeue guard', () => { beforeEach(() => { vi.clearAllMocks() resetDbChainMock() + dbChainMockFns.limit.mockResolvedValueOnce([OBSERVED_DOCUMENT]) }) /** @@ -393,6 +405,7 @@ describe('retryDocumentProcessing dispatch unwind', () => { beforeEach(() => { vi.clearAllMocks() resetDbChainMock() + dbChainMockFns.limit.mockResolvedValueOnce([OBSERVED_DOCUMENT]) }) /** diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts index 7b1bbef023b..87c4ac7fa78 100644 --- a/apps/sim/lib/knowledge/documents/service.ts +++ b/apps/sim/lib/knowledge/documents/service.ts @@ -104,6 +104,11 @@ import { } from '@/lib/knowledge/documents/processing-claim' import type { DocumentProcessingContinuation } from '@/lib/knowledge/documents/processing-continuation-dispatch' import { documentProcessingQueueOptions } from '@/lib/knowledge/documents/processing-lane' +import { + documentProcessingSnapshotCondition, + findAbandonedDocumentProcessing, + processingSnapshotColumns, +} from '@/lib/knowledge/documents/processing-liveness' import { enqueueKnowledgeDocumentProcessing } from '@/lib/knowledge/documents/processing-outbox-event' import { assertDocumentProcessingBillingContext, @@ -824,15 +829,10 @@ interface MarkDocumentsQueuedResult { } function acceptedDocumentStateCondition(observedAt: Date): SQL | undefined { - const queuedCutoff = new Date(observedAt.getTime() - QUEUED_DISPATCH_GRACE_MS) const processingCutoff = new Date(observedAt.getTime() - DOCUMENT_PROCESSING_STALE_THRESHOLD_MS) return or( eq(document.processingStatus, 'completed'), - and( - eq(document.processingStatus, 'pending'), - isNotNull(document.processingQueuedAt), - gte(document.processingQueuedAt, queuedCutoff) - ), + and(eq(document.processingStatus, 'pending'), isNotNull(document.processingQueuedAt)), and( eq(document.processingStatus, 'processing'), isNotNull(document.processingStartedAt), @@ -881,7 +881,6 @@ async function markDocumentsQueued( queuedAt: Date, lease: ProcessingDispatchLease | undefined ): Promise { - const legacyAdoptionCutoff = new Date(queuedAt.getTime() - QUEUED_DISPATCH_GRACE_MS) return db.transaction(async (tx) => { if (lease) await assertSyncLeaseHeldInTx(tx, lease.connectorId, lease) const claimed = await tx @@ -925,20 +924,8 @@ async function markDocumentsQueued( and( inArray(document.id, unclaimedIds), eq(document.knowledgeBaseId, knowledgeBaseId), - or( - and( - or( - eq(document.processingStatus, 'pending'), - eq(document.processingStatus, 'failed') - ), - eq(document.processingQueueToken, queueToken) - ), - and( - eq(document.processingStatus, 'pending'), - isNull(document.processingQueueToken), - lt(document.processingQueuedAt, legacyAdoptionCutoff) - ) - ), + inArray(document.processingStatus, ['pending', 'failed']), + eq(document.processingQueueToken, queueToken), isNotNull(document.processingQueuedAt), isNull(document.processingDeferredUntil), eq(document.userExcluded, false), @@ -3436,78 +3423,76 @@ export async function retryDocumentProcessing( requestId: string, billingAttribution: BillingAttributionSnapshot | undefined ): Promise<{ success: boolean; status: string; message: string }> { - /** - * A document may be retried from a terminal state, or from a `pending` state - * old enough that its dispatch is certainly lost. - * - * Unguarded, a double-click issued two full passes: the second reset a - * document that the first had already queued, so both dispatches ran, both - * indexed, and both billed. A terminal-only guard closes that, but it also - * strands a document that never left `pending` — a worker killed before its - * claim UPDATE burns an attempt without changing status, and once the - * processing-attempt budget is spent the connector sweep drops it too. The row - * then matches nothing anywhere. - * - * The `pending` arm is admitted only past {@link QUEUED_DISPATCH_GRACE_MS}, - * which is the same grace the connector sweep waits out, so a second click - * still lands inside a live dispatch's window and still matches no rows. - * - * Age is measured from `COALESCE(processingQueuedAt, uploadedAt)`, exactly as - * `isStuckDocumentSweepEligible` measures it. `processingQueuedAt` is NULL - * only for a document no dispatch has ever stamped, and falling back to - * `uploadedAt` — rather than treating NULL as retryable — keeps the grace - * window closed for a document created moments ago whose first dispatch is - * still in flight. - */ + const [observed] = await db + .select(processingSnapshotColumns) + .from(document) + .where( + and( + eq(document.id, documentId), + eq(document.knowledgeBaseId, knowledgeBaseId), + isNull(document.deletedAt), + isNull(document.archivedAt) + ) + ) + .limit(1) + const mayReplace = + observed && + (observed.processingStatus === 'completed' || + (['pending', 'failed'].includes(observed.processingStatus) && + (await findAbandonedDocumentProcessing([observed])).length === 1)) + /** Age alone does not prove a queued generation was lost. */ const queuedGraceCutoff = new Date(Date.now() - QUEUED_DISPATCH_GRACE_MS) - const requeued = await db.transaction(async (tx) => { - const reset = await tx - .update(document) - .set({ - processingStatus: 'pending', - /** - * Invalidates the prior dispatch generation in the same write that - * reopens the row. The dispatch below installs its fresh generation. - */ - processingQueuedAt: null, - processingQueueToken: null, - processingStartedAt: null, - processingDeferredUntil: null, - processingCompletedAt: null, - processingError: null, - chunkCount: 0, - tokenCount: 0, - characterCount: 0, - }) - .where( - and( - eq(document.id, documentId), - or(isNull(document.connectorId), isNotNull(document.contentHash)), - not(skippedDocumentCondition()), - or( - inArray(document.processingStatus, ['completed', 'failed']), - and( - eq(document.processingStatus, 'pending'), - sql`COALESCE(${document.processingQueuedAt}, ${document.uploadedAt}) < ${sql.param(queuedGraceCutoff, document.processingQueuedAt)}`, - or( - isNull(document.processingDeferredUntil), - lt(document.processingDeferredUntil, queuedGraceCutoff) + const requeued = + mayReplace && + (await db.transaction(async (tx) => { + const reset = await tx + .update(document) + .set({ + processingStatus: 'pending', + /** + * Invalidates the prior dispatch generation in the same write that + * reopens the row. The dispatch below installs its fresh generation. + */ + processingQueuedAt: null, + processingQueueToken: null, + processingStartedAt: null, + processingDeferredUntil: null, + processingCompletedAt: null, + processingError: null, + chunkCount: 0, + tokenCount: 0, + characterCount: 0, + }) + .where( + and( + eq(document.id, documentId), + eq(document.knowledgeBaseId, knowledgeBaseId), + documentProcessingSnapshotCondition(observed), + or(isNull(document.connectorId), isNotNull(document.contentHash)), + not(skippedDocumentCondition()), + or( + inArray(document.processingStatus, ['completed', 'failed']), + and( + eq(document.processingStatus, 'pending'), + sql`COALESCE(${document.processingQueuedAt}, ${document.uploadedAt}) < ${sql.param(queuedGraceCutoff, document.processingQueuedAt)}`, + or( + isNull(document.processingDeferredUntil), + lt(document.processingDeferredUntil, queuedGraceCutoff) + ) ) - ) - ), - isNull(document.archivedAt), - isNull(document.deletedAt) + ), + isNull(document.archivedAt), + isNull(document.deletedAt) + ) ) - ) - .returning({ id: document.id }) + .returning({ id: document.id }) - // Embeddings are dropped only for a document this call actually claimed, - // so a losing double-click cannot wipe the winner's in-flight work. - if (reset.length > 0) { - await tx.delete(embedding).where(eq(embedding.documentId, documentId)) - } - return reset.length > 0 - }) + /** Only the winning reset may remove embeddings. */ + if (reset.length > 0) { + await tx.delete(embedding).where(eq(embedding.documentId, documentId)) + } + return reset.length > 0 + })) if (!requeued) { const [skipped] = await db diff --git a/apps/sim/lib/oauth/google-service-account-transport.test.ts b/apps/sim/lib/oauth/google-service-account-transport.test.ts index b62c7adc713..36cac725adf 100644 --- a/apps/sim/lib/oauth/google-service-account-transport.test.ts +++ b/apps/sim/lib/oauth/google-service-account-transport.test.ts @@ -2,6 +2,11 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { exchangeGoogleServiceAccountJwt } from '@/lib/oauth/google-service-account-transport' +const { warn } = vi.hoisted(() => ({ warn: vi.fn() })) +vi.mock('@sim/logger', () => ({ + createLogger: () => ({ warn, info: vi.fn(), error: vi.fn(), debug: vi.fn() }), +})) + const TOKEN_URI = 'https://oauth2.googleapis.com/token' const ASSERTION = 'private-jwt-assertion' const fetchMock = vi.fn() @@ -9,6 +14,7 @@ const fetchMock = vi.fn() beforeEach(() => { vi.useFakeTimers() fetchMock.mockReset() + warn.mockClear() vi.stubGlobal('fetch', fetchMock) }) @@ -158,6 +164,18 @@ describe('Google service-account token transport', () => { const checked = expect(request).rejects.toMatchObject({ name: 'TimeoutError' }) await vi.advanceTimersByTimeAsync(30_000) await checked + expect(warn).toHaveBeenCalledWith( + 'Google service account token transport failed', + expect.objectContaining({ + operation: 'google.oauth.token_exchange', + stage: 'reading_response', + currentStatus: status, + lastHttpStatus: status, + attempts: 1, + elapsedMs: 30_000, + timedOut: true, + }) + ) expect(requestSignal?.aborted).toBe(true) expect(fetchMock).toHaveBeenCalledTimes(1) } @@ -187,5 +205,34 @@ describe('Google service-account token transport', () => { caller.abort(reason) await checked expect(fetchMock).toHaveBeenCalledTimes(1) + expect(warn).not.toHaveBeenCalled() }) }) + +it('distinguishes a prior HTTP response from a stalled next request without logging secrets', async () => { + fetchMock + .mockResolvedValueOnce(Response.json({ error: 'private-provider-detail' }, { status: 503 })) + .mockImplementationOnce( + (_url, init) => + new Promise((_resolve, reject) => { + init?.signal?.addEventListener('abort', () => reject(init.signal?.reason), { once: true }) + }) + ) + const result = exchangeGoogleServiceAccountJwt(TOKEN_URI, ASSERTION) + const checked = expect(result).rejects.toMatchObject({ name: 'TimeoutError' }) + await vi.advanceTimersByTimeAsync(30_000) + await checked + expect(warn).toHaveBeenCalledWith( + 'Google service account token transport failed', + expect.objectContaining({ + stage: 'awaiting_response', + attempts: 2, + lastHttpStatus: 503, + timedOut: true, + }) + ) + expect(warn.mock.calls[0][1]).not.toHaveProperty('currentStatus') + expect(JSON.stringify(warn.mock.calls)).not.toMatch( + /private-provider-detail|private-jwt-assertion|oauth2.googleapis.com/ + ) +}) diff --git a/apps/sim/lib/oauth/google-service-account-transport.ts b/apps/sim/lib/oauth/google-service-account-transport.ts index 8cf93a77844..f3e3b58df0d 100644 --- a/apps/sim/lib/oauth/google-service-account-transport.ts +++ b/apps/sim/lib/oauth/google-service-account-transport.ts @@ -1,3 +1,4 @@ +import { createLogger } from '@sim/logger' import { parseRetryAfter } from '@sim/utils/retry' import { isRetryableError, @@ -7,6 +8,7 @@ import { const RETRYABLE_STATUSES = new Set([429, 500, 502, 503, 504]) const TOKEN_EXCHANGE_BUDGET_MS = 30_000 +const logger = createLogger('GoogleServiceAccountTransport') interface GoogleTokenExchangeResponse { ok: boolean @@ -33,10 +35,18 @@ export async function exchangeGoogleServiceAccountJwt( jwt: string, signal?: AbortSignal ): Promise { + const startedAt = Date.now() + let attempts = 0 + let stage: 'awaiting_response' | 'reading_response' | 'between_attempts' = 'awaiting_response' + let currentStatus: number | undefined + let lastStatus: number | undefined let lastResponse: GoogleTokenExchangeResponse | undefined try { return await retryWithExponentialBackoff( async (attemptSignal) => { + attempts++ + stage = 'awaiting_response' + currentStatus = undefined lastResponse = undefined const response = await fetch(tokenUri, { method: 'POST', @@ -48,6 +58,9 @@ export async function exchangeGoogleServiceAccountJwt( }), signal: attemptSignal, }) + stage = 'reading_response' + currentStatus = response.status + lastStatus = response.status const payload = await readBoundedHttpErrorPayload(response) attemptSignal.throwIfAborted() if (response.ok && !payload.ok) @@ -68,12 +81,27 @@ export async function exchangeGoogleServiceAccountJwt( maxDelayMs: 2000, retryBudgetMs: TOKEN_EXCHANGE_BUDGET_MS, signal, - retryCondition: (error) => - error instanceof TokenExchangeRetryError || isRetryableError(error), + retryCondition: (error) => { + const retryable = error instanceof TokenExchangeRetryError || isRetryableError(error) + if (retryable) stage = 'between_attempts' + return retryable + }, } ) } catch (error) { if (error instanceof TokenExchangeRetryError && lastResponse) return lastResponse + if (!signal?.aborted) { + logger.warn('Google service account token transport failed', { + operation: 'google.oauth.token_exchange', + stage, + attempts, + elapsedMs: Date.now() - startedAt, + budgetMs: TOKEN_EXCHANGE_BUDGET_MS, + ...(currentStatus !== undefined ? { currentStatus } : {}), + ...(lastStatus !== undefined ? { lastHttpStatus: lastStatus } : {}), + timedOut: error instanceof Error && error.name === 'TimeoutError', + }) + } throw error } } From 655313546ab175a79d2f2a589ec9777b90c5a7d9 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 18 Sep 2026 16:57:19 -0700 Subject: [PATCH 2/3] fix(knowledge): recover abandoned redelivery and cancel liveness requests --- .../stored-document-recovery.integration.ts | 115 +++++++++++++++++- .../documents/processing-liveness.test.ts | 97 +++++++++++++-- .../documents/processing-liveness.ts | 40 ++++-- .../documents/processing-queue.test.ts | 7 +- apps/sim/lib/knowledge/documents/service.ts | 85 ++++++++++++- 5 files changed, 312 insertions(+), 32 deletions(-) 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 034a9319a6c..eac7af3652e 100644 --- a/apps/sim/lib/knowledge/__integration__/stored-document-recovery.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/stored-document-recovery.integration.ts @@ -24,13 +24,31 @@ const fixture = vi.hoisted(() => ({ embeddingCalls: 0, useTrigger: false, listRuns: vi.fn(), + batchTrigger: vi.fn(), })) vi.mock('@/lib/core/config/trigger-runtime', () => ({ isInsideTriggerRun: () => fixture.useTrigger, })) -vi.mock('@trigger.dev/sdk', async (importOriginal) => ({ - ...(await importOriginal()), - runs: { list: fixture.listRuns }, +vi.mock('@trigger.dev/core/v3', async (importOriginal) => ({ + ...(await importOriginal()), + apiClientManager: { + clientOrThrow: () => ({ baseUrl: 'https://api.trigger.dev', getHeaders: () => ({}) }), + }, +})) +vi.mock('@trigger.dev/core/v3/zodfetch', () => ({ + zodfetchCursorPage: (_schema: unknown, _url: string, params: { query: URLSearchParams }) => + fixture.listRuns({ tag: params.query.get('filter[tag]') }), +})) +vi.mock('@trigger.dev/sdk', async (importOriginal) => { + const original = await importOriginal() + return { + ...original, + runs: { list: fixture.listRuns }, + tasks: { ...original.tasks, batchTrigger: fixture.batchTrigger }, + } +}) +vi.mock('@/lib/core/async-jobs/region', () => ({ + resolveTriggerRegion: async () => 'us-east-1', })) vi.mock('@/lib/uploads/core/setup.server', () => ({ get UPLOAD_DIR_SERVER() { @@ -65,6 +83,7 @@ import { searchScopedKnowledge } from '@/lib/knowledge/application/workspace-sea 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 { enqueueKnowledgeDocumentProcessing } from '@/lib/knowledge/documents/processing-outbox-event' import { knowledgeDocumentProcessingOutboxHandlers } from '@/lib/knowledge/documents/processing-outbox-handler' import { DOCUMENT_RECOVERY_BATCH_SIZE, @@ -135,6 +154,7 @@ async function failedFile( afterEach(() => { fixture.useTrigger = false fixture.listRuns.mockReset() + fixture.batchTrigger.mockReset() }) beforeAll(() => { @@ -156,6 +176,95 @@ afterAll(async () => { }) describe('independent recovery of retained connector documents', () => { + it.each([null, 'abandoned-generation'])( + 'redelivers an abandoned upload with token %s without spending another admission', + async (processingQueueToken) => { + const ids = await seed() + const file = await failedFile(ids) + const queuedAt = old() + await db + .update(document) + .set({ + connectorId: null, + processingStatus: 'pending', + processingQueueToken, + processingQueuedAt: queuedAt, + processingCompletedAt: null, + }) + .where(eq(document.id, file.documentId)) + const eventId = await enqueueKnowledgeDocumentProcessing(db, { + knowledgeBaseId: ids.knowledgeBaseId, + documentId: file.documentId, + processingOptions: {}, + billingAttribution: await resolveSystemBillingAttribution(ids.workspaceId), + processingLane: 'interactive', + }) + fixture.useTrigger = true + fixture.listRuns.mockResolvedValue({ data: [], hasNextPage: () => false }) + fixture.batchTrigger.mockResolvedValue({ batchId: 'fixture-batch' }) + expect( + await outbox.processOutboxEventById(eventId, knowledgeDocumentProcessingOutboxHandlers) + ).toBe('completed') + expect(fixture.batchTrigger).toHaveBeenCalledOnce() + expect(fixture.batchTrigger.mock.calls[0][1][0].payload).toMatchObject({ + documentId: file.documentId, + processingQueueToken: eventId, + processingQueuedAt: queuedAt.toISOString(), + chargedAtDispatch: false, + }) + const [row] = await db.select().from(document).where(eq(document.id, file.documentId)) + expect(row.processingQueueToken).toBe(eventId) + expect(row.processingAttempts).toBe(1) + } + ) + + it.each(['live', 'unknown', 'race'])( + 'does not adopt a legacy upload when its processing is %s', + async (state) => { + const ids = await seed() + const file = await failedFile(ids) + await db + .update(document) + .set({ + connectorId: null, + processingStatus: 'pending', + processingQueueToken: null, + processingCompletedAt: null, + }) + .where(eq(document.id, file.documentId)) + const eventId = await enqueueKnowledgeDocumentProcessing(db, { + knowledgeBaseId: ids.knowledgeBaseId, + documentId: file.documentId, + processingOptions: {}, + billingAttribution: await resolveSystemBillingAttribution(ids.workspaceId), + processingLane: 'interactive', + }) + fixture.useTrigger = true + if (state === 'unknown') + fixture.listRuns.mockRejectedValue(new Error('Synthetic lookup failure')) + else if (state === 'live') + fixture.listRuns.mockResolvedValue({ + data: [{ status: 'QUEUED' }], + hasNextPage: () => false, + }) + else + fixture.listRuns.mockImplementation(async () => { + await db + .update(document) + .set({ processingQueueToken: 'winning-generation', processingQueuedAt: new Date() }) + .where(eq(document.id, file.documentId)) + return { data: [], hasNextPage: () => false } + }) + expect( + await outbox.processOutboxEventById(eventId, knowledgeDocumentProcessingOutboxHandlers) + ).toBe(state === 'live' ? 'completed' : 'pending') + expect(fixture.batchTrigger).not.toHaveBeenCalled() + const [row] = await db.select().from(document).where(eq(document.id, file.documentId)) + expect(row.processingQueueToken).toBe(state === 'race' ? 'winning-generation' : null) + expect(row.processingAttempts).toBe(1) + } + ) + it.each(['QUEUED', 'DELAYED', 'WAITING'])( 'preserves a seven-hour %s run and its existing attempt', async (status) => { diff --git a/apps/sim/lib/knowledge/documents/processing-liveness.test.ts b/apps/sim/lib/knowledge/documents/processing-liveness.test.ts index a9f46cac6ae..cc8b3b019d2 100644 --- a/apps/sim/lib/knowledge/documents/processing-liveness.test.ts +++ b/apps/sim/lib/knowledge/documents/processing-liveness.test.ts @@ -3,7 +3,19 @@ import { dbChainMockFns, resetDbChainMock, resetEnvFlagsMock, setEnvFlags } from import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' const { listRuns } = vi.hoisted(() => ({ listRuns: vi.fn() })) -vi.mock('@trigger.dev/sdk', () => ({ runs: { list: listRuns } })) +vi.mock('@trigger.dev/core/v3', async (importOriginal) => ({ + ...(await importOriginal()), + apiClientManager: { + clientOrThrow: () => ({ + baseUrl: 'https://api.trigger.dev', + getHeaders: () => ({ + Authorization: 'Bearer fixture-key', + 'x-trigger-branch': 'fixture-branch', + }), + }), + }, +})) +vi.mock('@trigger.dev/core/v3/zodfetch', () => ({ zodfetchCursorPage: listRuns })) import { env } from '@/lib/core/config/env' import { resetInsideTriggerRunForTests } from '@/lib/core/config/trigger-runtime' @@ -29,12 +41,13 @@ beforeEach(() => { resetInsideTriggerRunForTests() setEnvFlags({ isTriggerDevEnabled: true }) env.TRIGGER_SECRET_KEY = 'test-secret' - listRuns.mockResolvedValue({ data: [], hasNextPage: () => false }) + listRuns.mockReset().mockResolvedValue({ data: [], hasNextPage: () => false }) }) afterEach(() => { env.TRIGGER_SECRET_KEY = originalSecret resetEnvFlagsMock() vi.useRealTimers() + vi.unstubAllGlobals() }) describe('document processing liveness', () => { @@ -46,20 +59,28 @@ describe('document processing liveness', () => { expect(dbChainMockFns.set).toHaveBeenCalledWith({ processingRecoveryAfter: expect.any(Date) }) expect(dbChainMockFns.delete).not.toHaveBeenCalled() expect(listRuns).toHaveBeenCalledWith( + expect.anything(), + 'https://api.trigger.dev/api/v1/runs', expect.objectContaining({ - tag: 'documentId:doc-1', + query: expect.any(URLSearchParams), limit: 1, - status: expect.arrayContaining([ - 'QUEUED', - 'WAITING', - 'DELAYED', - 'EXECUTING', - 'DEQUEUED', - 'PENDING_VERSION', - ]), }), + expect.objectContaining({ signal: expect.any(AbortSignal) }), { retry: { maxAttempts: 1 } } ) + const params = listRuns.mock.calls[0][2].query as URLSearchParams + expect(params.get('filter[tag]')).toBe('documentId:doc-1') + expect(params.get('filter[taskIdentifier]')).toBe('knowledge-process-document') + expect(params.get('filter[status]')?.split(',')).toEqual( + expect.arrayContaining([ + 'QUEUED', + 'WAITING', + 'DELAYED', + 'EXECUTING', + 'DEQUEUED', + 'PENDING_VERSION', + ]) + ) } ) @@ -104,6 +125,60 @@ describe('document processing liveness', () => { expect(listRuns).toHaveBeenCalledTimes(4) }) + it('cancels the actual SDK HTTP requests at the deadline without accumulating requests', async () => { + const { zodfetchCursorPage } = await vi.importActual< + typeof import('@trigger.dev/core/v3/zodfetch') + >('@trigger.dev/core/v3/zodfetch') + listRuns.mockImplementation(zodfetchCursorPage) + vi.useFakeTimers() + let active = 0 + const fetch = vi.fn( + (_url: string, init: RequestInit) => + new Promise((_resolve, reject) => { + active++ + init.signal!.addEventListener( + 'abort', + () => { + active-- + reject(init.signal!.reason) + }, + { once: true } + ) + }) + ) + vi.stubGlobal('fetch', fetch) + const candidates = Array.from({ length: 20 }, (_, i) => ({ ...snapshot, id: `doc-${i}` })) + for (let attempt = 0; attempt < 2; attempt++) { + const result = findAbandonedDocumentProcessing(candidates) + await vi.advanceTimersByTimeAsync(8_000) + expect(await result).toEqual([]) + expect(active).toBe(0) + expect(fetch).toHaveBeenCalledTimes((attempt + 1) * 4) + } + const headers = new Headers(fetch.mock.calls[0][1].headers) + expect(headers.get('Authorization')).toBe('Bearer fixture-key') + expect(headers.get('x-trigger-branch')).toBe('fixture-branch') + }) + + it.each([true, false])( + 'validates the SDK response before declaring abandonment: valid=%s', + async (valid) => { + const { zodfetchCursorPage } = await vi.importActual< + typeof import('@trigger.dev/core/v3/zodfetch') + >('@trigger.dev/core/v3/zodfetch') + listRuns.mockImplementation(zodfetchCursorPage) + const fetch = vi + .fn() + .mockResolvedValue(Response.json(valid ? { data: [], pagination: {} } : { data: [] })) + vi.stubGlobal('fetch', fetch) + expect(await findAbandonedDocumentProcessing([snapshot])).toEqual(valid ? [snapshot] : []) + expect(fetch).toHaveBeenCalledOnce() + const url = new URL(fetch.mock.calls[0][0]) + expect(url.searchParams.get('page[size]')).toBe('1') + expect(url.searchParams.get('filter[tag]')).toBe('documentId:doc-1') + } + ) + it('retains recovery on installations using the in-process fallback', async () => { env.TRIGGER_SECRET_KEY = undefined expect(await findAbandonedDocumentProcessing([snapshot])).toEqual([snapshot]) diff --git a/apps/sim/lib/knowledge/documents/processing-liveness.ts b/apps/sim/lib/knowledge/documents/processing-liveness.ts index 98d6612ec3d..18dbf47a56c 100644 --- a/apps/sim/lib/knowledge/documents/processing-liveness.ts +++ b/apps/sim/lib/knowledge/documents/processing-liveness.ts @@ -1,8 +1,8 @@ import { db } from '@sim/db' import { document, outboxEvent } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import type { RunStatus } from '@trigger.dev/core/v3' -import { runs } from '@trigger.dev/sdk' +import { apiClientManager, ListRunResponseItem, type RunStatus } from '@trigger.dev/core/v3' +import { zodfetchCursorPage } from '@trigger.dev/core/v3/zodfetch' import { and, eq, inArray, or, sql } from 'drizzle-orm' import { env } from '@/lib/core/config/env' import { isTriggerDevEnabled } from '@/lib/core/config/env-flags' @@ -69,16 +69,24 @@ type ProcessingLiveness = 'live' | 'abandoned' | 'unknown' async function inspectTriggerWork(documentId: string, deadlineAt: number, signal?: AbortSignal) { return withinDeadline( - async () => { - const page = await runs.list( + async (requestSignal) => { + const client = apiClientManager.clientOrThrow() + /** The SDK's runs.list wrapper omits RequestInit.signal; its transport supports it. */ + const page = await zodfetchCursorPage( + ListRunResponseItem, + `${client.baseUrl}/api/v1/runs`, { - taskIdentifier: ['knowledge-process-document'], - tag: `documentId:${documentId}`, - status: ACTIVE_RUN_STATUSES, + query: new URLSearchParams({ + 'filter[taskIdentifier]': 'knowledge-process-document', + 'filter[tag]': `documentId:${documentId}`, + 'filter[status]': ACTIVE_RUN_STATUSES.join(','), + }), limit: 1, }, + { method: 'GET', headers: client.getHeaders(), signal: requestSignal }, { retry: { maxAttempts: 1 } } ) + requestSignal.throwIfAborted() if (page.data.length > 0) return 'live' as const return page.hasNextPage() ? ('unknown' as const) : ('abandoned' as const) }, @@ -92,11 +100,11 @@ async function inspectTriggerWork(documentId: string, deadlineAt: number, signal * outside row locks. Any live document run protects continuation handoffs and legacy jobs. * Failed or incomplete lookups defer recovery; they never authorize another admission. */ -export async function findAbandonedDocumentProcessing( +export async function inspectDocumentProcessingLiveness( candidates: readonly T[], signal?: AbortSignal -): Promise { - if (candidates.length === 0) return [] +): Promise<{ abandoned: T[]; live: T[] }> { + if (candidates.length === 0) return { abandoned: [], live: [] } if (candidates.length > DOCUMENT_LIVENESS_BATCH_SIZE) { throw new Error('Document liveness batch exceeds its limit') } @@ -184,5 +192,15 @@ export async function findAbandonedDocumentProcessing states.get(row.id) === 'abandoned') + return { + abandoned: candidates.filter((row) => states.get(row.id) === 'abandoned'), + live: candidates.filter((row) => states.get(row.id) === 'live'), + } +} + +export async function findAbandonedDocumentProcessing( + candidates: readonly T[], + signal?: AbortSignal +): Promise { + return (await inspectDocumentProcessingLiveness(candidates, signal)).abandoned } diff --git a/apps/sim/lib/knowledge/documents/processing-queue.test.ts b/apps/sim/lib/knowledge/documents/processing-queue.test.ts index 14554f79e67..85d46ee6a51 100644 --- a/apps/sim/lib/knowledge/documents/processing-queue.test.ts +++ b/apps/sim/lib/knowledge/documents/processing-queue.test.ts @@ -22,6 +22,7 @@ import { } from '@/lib/core/config/trigger-runtime' import { SyncLockLostException } from '@/lib/knowledge/connectors/sync-lock' import { DOCUMENT_PROCESSING_STALE_THRESHOLD_MS } from '@/lib/knowledge/documents/processing-timeouts.server' +import { QUEUED_DISPATCH_GRACE_MS } from '@/lib/knowledge/documents/types' const { mockBatchTrigger, mockResolveTriggerRegion } = vi.hoisted(() => ({ mockBatchTrigger: vi.fn(), @@ -525,7 +526,7 @@ describe('processDocumentsWithQueue dispatch backend', () => { (node: MockCondition) => node.type === 'gte' && node.left === schemaMock.document.processingQueuedAt ) - ).toBe(false) + ).toBe(true) const liveProcessingState = acceptedStatuses.find( (condition) => condition.type === 'and' && @@ -619,7 +620,7 @@ describe('processDocumentsWithQueue dispatch backend', () => { ).toBe(true) }) - it('preserves an old legacy generation for the liveness-aware recovery path', async () => { + it('preserves a recently queued generation without dispatching duplicate work', async () => { markInsideTriggerRun() dbChainMockFns.returning.mockResolvedValueOnce([]).mockResolvedValueOnce([]) queueTableRows(schemaMock.document, [{ id: 'document-1' }]) @@ -901,7 +902,7 @@ describe('processDocumentsWithQueue dispatch backend', () => { (node: MockCondition) => node.type === 'gte' && node.left === schemaMock.document.processingQueuedAt ) - expect(queuedFreshness).toBeUndefined() + expect(queuedFreshness?.right).toEqual(new Date(now.getTime() - QUEUED_DISPATCH_GRACE_MS)) const processingState = acceptedStatuses.find( (condition) => condition.type === 'and' && diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts index 87c4ac7fa78..714b3ddde8d 100644 --- a/apps/sim/lib/knowledge/documents/service.ts +++ b/apps/sim/lib/knowledge/documents/service.ts @@ -105,8 +105,10 @@ import { import type { DocumentProcessingContinuation } from '@/lib/knowledge/documents/processing-continuation-dispatch' import { documentProcessingQueueOptions } from '@/lib/knowledge/documents/processing-lane' import { + DOCUMENT_LIVENESS_BATCH_SIZE, documentProcessingSnapshotCondition, findAbandonedDocumentProcessing, + inspectDocumentProcessingLiveness, processingSnapshotColumns, } from '@/lib/knowledge/documents/processing-liveness' import { enqueueKnowledgeDocumentProcessing } from '@/lib/knowledge/documents/processing-outbox-event' @@ -829,10 +831,15 @@ interface MarkDocumentsQueuedResult { } function acceptedDocumentStateCondition(observedAt: Date): SQL | undefined { + const queuedCutoff = new Date(observedAt.getTime() - QUEUED_DISPATCH_GRACE_MS) const processingCutoff = new Date(observedAt.getTime() - DOCUMENT_PROCESSING_STALE_THRESHOLD_MS) return or( eq(document.processingStatus, 'completed'), - and(eq(document.processingStatus, 'pending'), isNotNull(document.processingQueuedAt)), + and( + eq(document.processingStatus, 'pending'), + isNotNull(document.processingQueuedAt), + gte(document.processingQueuedAt, queuedCutoff) + ), and( eq(document.processingStatus, 'processing'), isNotNull(document.processingStartedAt), @@ -879,9 +886,10 @@ async function markDocumentsQueued( knowledgeBaseId: string, queueToken: string, queuedAt: Date, - lease: ProcessingDispatchLease | undefined + lease: ProcessingDispatchLease | undefined, + signal?: AbortSignal ): Promise { - return db.transaction(async (tx) => { + const result = await db.transaction(async (tx) => { if (lease) await assertSyncLeaseHeldInTx(tx, lease.connectorId, lease) const claimed = await tx .update(document) @@ -983,6 +991,68 @@ async function markDocumentsQueued( ), } }) + + if (result.unresolvedIds.length === 0) return result + const candidates = await db + .select(processingSnapshotColumns) + .from(document) + .where( + and( + inArray(document.id, result.unresolvedIds), + eq(document.knowledgeBaseId, knowledgeBaseId), + eq(document.processingStatus, 'pending'), + lt(document.processingQueuedAt, new Date(queuedAt.getTime() - QUEUED_DISPATCH_GRACE_MS)), + isNull(document.processingDeferredUntil), + eq(document.userExcluded, false), + isNull(document.archivedAt), + isNull(document.deletedAt) + ) + ) + .limit(DOCUMENT_LIVENESS_BATCH_SIZE) + const { abandoned, live } = await inspectDocumentProcessingLiveness(candidates, signal) + /** Redelivery resumes an abandoned admission; it does not spend another attempt. */ + const adopted = + abandoned.length === 0 + ? [] + : await db.transaction(async (tx) => { + signal?.throwIfAborted() + if (lease) await assertSyncLeaseHeldInTx(tx, lease.connectorId, lease) + return tx + .update(document) + .set({ processingQueueToken: queueToken }) + .where( + and( + eq(document.knowledgeBaseId, knowledgeBaseId), + or(...abandoned.map(documentProcessingSnapshotCondition)), + eq(document.userExcluded, false), + isNull(document.archivedAt), + isNull(document.deletedAt) + ) + ) + .returning({ id: document.id, processingQueuedAt: document.processingQueuedAt }) + }) + const acceptedIds = new Set([...live, ...adopted].map((row) => row.id)) + return { + generations: [ + ...result.generations, + ...adopted.flatMap((row) => + row.processingQueuedAt + ? [ + { + documentId: row.id, + processingQueuedAt: row.processingQueuedAt, + chargedAtDispatch: false, + }, + ] + : [] + ), + ], + acceptedWithoutDispatchIds: [ + ...result.acceptedWithoutDispatchIds, + ...live.map((row) => row.id), + ], + unresolvedIds: result.unresolvedIds.filter((id) => !acceptedIds.has(id)), + } } /** @@ -1095,7 +1165,14 @@ export async function processDocumentsWithQueue( generations: queuedGenerations, acceptedWithoutDispatchIds, unresolvedIds, - } = await markDocumentsQueued(documentIds, knowledgeBaseId, requestId, queuedAt, lease) + } = await markDocumentsQueued( + documentIds, + knowledgeBaseId, + requestId, + queuedAt, + lease, + executionContext?.signal + ) const generationByDocumentId = new Map( queuedGenerations.map((generation) => [generation.documentId, generation]) ) From e27c6b52d8c1369a952720d52aa951fe70a4be25 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Fri, 18 Sep 2026 17:08:31 -0700 Subject: [PATCH 3/3] fix(knowledge): fence recovery against concurrent protection --- .../stored-document-recovery.integration.ts | 66 ++++++++++++- .../processing-liveness-transport.test.ts | 98 +++++++++++++++++++ .../documents/processing-liveness.test.ts | 60 +----------- .../documents/processing-liveness.ts | 6 +- .../documents/retry-processing-grace.test.ts | 1 + 5 files changed, 172 insertions(+), 59 deletions(-) create mode 100644 apps/sim/lib/knowledge/documents/processing-liveness-transport.test.ts 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 eac7af3652e..7f115a85d18 100644 --- a/apps/sim/lib/knowledge/__integration__/stored-document-recovery.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/stored-document-recovery.integration.ts @@ -90,7 +90,11 @@ import { KNOWLEDGE_DOCUMENT_RECOVERY_OUTBOX_EVENT, recoverKnowledgeDocumentProcessing, } from '@/lib/knowledge/documents/processing-recovery' -import { processDocumentAsync, retryDocumentProcessing } from '@/lib/knowledge/documents/service' +import { + processDocumentAsync, + processDocumentsWithQueue, + retryDocumentProcessing, +} from '@/lib/knowledge/documents/service' import { MAX_PROCESSING_ATTEMPTS, QUEUED_DISPATCH_GRACE_MS } from '@/lib/knowledge/documents/types' const fixtures: ReturnType[] = [] @@ -176,6 +180,66 @@ afterAll(async () => { }) describe('independent recovery of retained connector documents', () => { + it.each(['manual', 'redelivery'])( + 'respects a concurrent liveness cooldown before %s replacement', + async (path) => { + const ids = await seed() + const file = await failedFile(ids) + await db + .update(document) + .set({ processingStatus: 'pending' }) + .where(eq(document.id, file.documentId)) + const [original] = await db.select().from(document).where(eq(document.id, file.documentId)) + const protectedUntil = new Date(Date.now() + 60_000) + fixture.useTrigger = true + fixture.batchTrigger.mockResolvedValue({ batchId: 'fixture-batch' }) + fixture.listRuns.mockImplementation(async () => { + await db + .update(document) + .set({ processingRecoveryAfter: protectedUntil }) + .where(eq(document.id, file.documentId)) + return { data: [], hasNextPage: () => false } + }) + const billing = await resolveSystemBillingAttribution(ids.workspaceId) + const docData = { + documentId: file.documentId, + filename: original.filename, + fileUrl: original.fileUrl, + fileSize: original.fileSize, + mimeType: original.mimeType, + } + if (path === 'manual') { + const result = await retryDocumentProcessing( + ids.knowledgeBaseId, + file.documentId, + docData, + generateId(), + billing + ) + expect(result.message).toContain('already queued') + } else { + const result = await processDocumentsWithQueue( + [docData], + ids.knowledgeBaseId, + {}, + generateId(), + billing, + 'interactive' + ) + expect(result).toMatchObject({ accepted: 0, failed: 1 }) + } + expect(fixture.batchTrigger).not.toHaveBeenCalled() + const [row] = await db.select().from(document).where(eq(document.id, file.documentId)) + expect(row.processingQueueToken).toBe(original.processingQueueToken) + expect(row.processingAttempts).toBe(1) + expect(row.processingRecoveryAfter).toEqual(protectedUntil) + await db + .update(knowledgeConnector) + .set({ status: 'paused' }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + } + ) + it.each([null, 'abandoned-generation'])( 'redelivers an abandoned upload with token %s without spending another admission', async (processingQueueToken) => { diff --git a/apps/sim/lib/knowledge/documents/processing-liveness-transport.test.ts b/apps/sim/lib/knowledge/documents/processing-liveness-transport.test.ts new file mode 100644 index 00000000000..dcfd31ab920 --- /dev/null +++ b/apps/sim/lib/knowledge/documents/processing-liveness-transport.test.ts @@ -0,0 +1,98 @@ +/** @vitest-environment node */ +import { resetDbChainMock, resetEnvFlagsMock, setEnvFlags } from '@sim/testing' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +/** Keep the SDK transport real to verify signal forwarding and response validation. */ +vi.unmock('@trigger.dev/core/v3') + +import { apiClientManager } from '@trigger.dev/core/v3' +import { env } from '@/lib/core/config/env' +import { resetInsideTriggerRunForTests } from '@/lib/core/config/trigger-runtime' +import { + type DocumentProcessingSnapshot, + findAbandonedDocumentProcessing, +} from '@/lib/knowledge/documents/processing-liveness' + +const snapshot: DocumentProcessingSnapshot = { + id: 'doc-1', + processingStatus: 'pending', + processingQueueToken: 'generation-1', + processingQueuedAt: new Date('2026-09-01T00:00:00Z'), + processingStartedAt: null, + processingDeferredUntil: null, + processingCompletedAt: null, + processingRecoveryAfter: null, +} +const originalSecret = env.TRIGGER_SECRET_KEY + +function inspect(candidates: DocumentProcessingSnapshot[]) { + return apiClientManager.runWithConfig( + { + baseURL: 'https://api.trigger.dev', + accessToken: 'fixture-key', + previewBranch: 'fixture-branch', + }, + () => findAbandonedDocumentProcessing(candidates) + ) +} + +beforeEach(() => { + resetDbChainMock() + resetInsideTriggerRunForTests() + setEnvFlags({ isTriggerDevEnabled: true }) + env.TRIGGER_SECRET_KEY = 'fixture-key' +}) +afterEach(() => { + env.TRIGGER_SECRET_KEY = originalSecret + resetEnvFlagsMock() + vi.useRealTimers() + vi.unstubAllGlobals() +}) + +describe('document liveness SDK transport', () => { + it('cancels the actual SDK HTTP requests at the deadline without accumulating requests', async () => { + vi.useFakeTimers() + let active = 0 + const fetch = vi.fn( + (_url: string, init: RequestInit) => + new Promise((_resolve, reject) => { + active++ + init.signal!.addEventListener( + 'abort', + () => { + active-- + reject(init.signal!.reason) + }, + { once: true } + ) + }) + ) + vi.stubGlobal('fetch', fetch) + const candidates = Array.from({ length: 20 }, (_, i) => ({ ...snapshot, id: `doc-${i}` })) + for (let attempt = 0; attempt < 2; attempt++) { + const result = inspect(candidates) + await vi.advanceTimersByTimeAsync(8_000) + expect(await result).toEqual([]) + expect(active).toBe(0) + expect(fetch).toHaveBeenCalledTimes((attempt + 1) * 4) + } + const headers = new Headers(fetch.mock.calls[0][1].headers) + expect(headers.get('Authorization')).toBe('Bearer fixture-key') + expect(headers.get('x-trigger-branch')).toBe('fixture-branch') + }) + + it.each([true, false])( + 'validates the SDK response before declaring abandonment: valid=%s', + async (valid) => { + const fetch = vi + .fn() + .mockResolvedValue(Response.json(valid ? { data: [], pagination: {} } : { data: [] })) + vi.stubGlobal('fetch', fetch) + expect(await inspect([snapshot])).toEqual(valid ? [snapshot] : []) + expect(fetch).toHaveBeenCalledOnce() + const url = new URL(fetch.mock.calls[0][0]) + expect(url.searchParams.get('page[size]')).toBe('1') + expect(url.searchParams.get('filter[tag]')).toBe('documentId:doc-1') + } + ) +}) diff --git a/apps/sim/lib/knowledge/documents/processing-liveness.test.ts b/apps/sim/lib/knowledge/documents/processing-liveness.test.ts index cc8b3b019d2..6a634cf1093 100644 --- a/apps/sim/lib/knowledge/documents/processing-liveness.test.ts +++ b/apps/sim/lib/knowledge/documents/processing-liveness.test.ts @@ -3,8 +3,9 @@ import { dbChainMockFns, resetDbChainMock, resetEnvFlagsMock, setEnvFlags } from import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' const { listRuns } = vi.hoisted(() => ({ listRuns: vi.fn() })) -vi.mock('@trigger.dev/core/v3', async (importOriginal) => ({ - ...(await importOriginal()), +vi.mock('@trigger.dev/core/v3', () => ({ + taskContext: { isInsideTask: false }, + ListRunResponseItem: {}, apiClientManager: { clientOrThrow: () => ({ baseUrl: 'https://api.trigger.dev', @@ -33,6 +34,7 @@ const snapshot: DocumentProcessingSnapshot = { processingStartedAt: null, processingDeferredUntil: null, processingCompletedAt: null, + processingRecoveryAfter: null, } const originalSecret = env.TRIGGER_SECRET_KEY beforeEach(() => { @@ -125,60 +127,6 @@ describe('document processing liveness', () => { expect(listRuns).toHaveBeenCalledTimes(4) }) - it('cancels the actual SDK HTTP requests at the deadline without accumulating requests', async () => { - const { zodfetchCursorPage } = await vi.importActual< - typeof import('@trigger.dev/core/v3/zodfetch') - >('@trigger.dev/core/v3/zodfetch') - listRuns.mockImplementation(zodfetchCursorPage) - vi.useFakeTimers() - let active = 0 - const fetch = vi.fn( - (_url: string, init: RequestInit) => - new Promise((_resolve, reject) => { - active++ - init.signal!.addEventListener( - 'abort', - () => { - active-- - reject(init.signal!.reason) - }, - { once: true } - ) - }) - ) - vi.stubGlobal('fetch', fetch) - const candidates = Array.from({ length: 20 }, (_, i) => ({ ...snapshot, id: `doc-${i}` })) - for (let attempt = 0; attempt < 2; attempt++) { - const result = findAbandonedDocumentProcessing(candidates) - await vi.advanceTimersByTimeAsync(8_000) - expect(await result).toEqual([]) - expect(active).toBe(0) - expect(fetch).toHaveBeenCalledTimes((attempt + 1) * 4) - } - const headers = new Headers(fetch.mock.calls[0][1].headers) - expect(headers.get('Authorization')).toBe('Bearer fixture-key') - expect(headers.get('x-trigger-branch')).toBe('fixture-branch') - }) - - it.each([true, false])( - 'validates the SDK response before declaring abandonment: valid=%s', - async (valid) => { - const { zodfetchCursorPage } = await vi.importActual< - typeof import('@trigger.dev/core/v3/zodfetch') - >('@trigger.dev/core/v3/zodfetch') - listRuns.mockImplementation(zodfetchCursorPage) - const fetch = vi - .fn() - .mockResolvedValue(Response.json(valid ? { data: [], pagination: {} } : { data: [] })) - vi.stubGlobal('fetch', fetch) - expect(await findAbandonedDocumentProcessing([snapshot])).toEqual(valid ? [snapshot] : []) - expect(fetch).toHaveBeenCalledOnce() - const url = new URL(fetch.mock.calls[0][0]) - expect(url.searchParams.get('page[size]')).toBe('1') - expect(url.searchParams.get('filter[tag]')).toBe('documentId:doc-1') - } - ) - it('retains recovery on installations using the in-process fallback', async () => { env.TRIGGER_SECRET_KEY = undefined expect(await findAbandonedDocumentProcessing([snapshot])).toEqual([snapshot]) diff --git a/apps/sim/lib/knowledge/documents/processing-liveness.ts b/apps/sim/lib/knowledge/documents/processing-liveness.ts index 18dbf47a56c..d809bd009a2 100644 --- a/apps/sim/lib/knowledge/documents/processing-liveness.ts +++ b/apps/sim/lib/knowledge/documents/processing-liveness.ts @@ -45,6 +45,7 @@ export const processingSnapshotColumns = { processingStartedAt: document.processingStartedAt, processingDeferredUntil: document.processingDeferredUntil, processingCompletedAt: document.processingCompletedAt, + processingRecoveryAfter: document.processingRecoveryAfter, } export type DocumentProcessingSnapshot = Pick< @@ -52,7 +53,7 @@ export type DocumentProcessingSnapshot = Pick< keyof typeof processingSnapshotColumns > -/** A claim or continuation installed during the external lookup must win over recovery. */ +/** A claim, continuation or protection cooldown installed during lookup must win over recovery. */ export function documentProcessingSnapshotCondition(snapshot: DocumentProcessingSnapshot) { return and( eq(document.id, snapshot.id), @@ -61,7 +62,8 @@ export function documentProcessingSnapshotCondition(snapshot: DocumentProcessing sql`${document.processingQueuedAt} IS NOT DISTINCT FROM ${sql.param(snapshot.processingQueuedAt, document.processingQueuedAt)}`, sql`${document.processingStartedAt} IS NOT DISTINCT FROM ${sql.param(snapshot.processingStartedAt, document.processingStartedAt)}`, sql`${document.processingDeferredUntil} IS NOT DISTINCT FROM ${sql.param(snapshot.processingDeferredUntil, document.processingDeferredUntil)}`, - sql`${document.processingCompletedAt} IS NOT DISTINCT FROM ${sql.param(snapshot.processingCompletedAt, document.processingCompletedAt)}` + sql`${document.processingCompletedAt} IS NOT DISTINCT FROM ${sql.param(snapshot.processingCompletedAt, document.processingCompletedAt)}`, + sql`${document.processingRecoveryAfter} IS NOT DISTINCT FROM ${sql.param(snapshot.processingRecoveryAfter, document.processingRecoveryAfter)}` ) } diff --git a/apps/sim/lib/knowledge/documents/retry-processing-grace.test.ts b/apps/sim/lib/knowledge/documents/retry-processing-grace.test.ts index bcee1680b23..bfa49404a9c 100644 --- a/apps/sim/lib/knowledge/documents/retry-processing-grace.test.ts +++ b/apps/sim/lib/knowledge/documents/retry-processing-grace.test.ts @@ -34,6 +34,7 @@ const OBSERVED_DOCUMENT = { processingStartedAt: null, processingDeferredUntil: null, processingCompletedAt: new Date(0), + processingRecoveryAfter: null, } const DOC_DATA = {