Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 22 additions & 1 deletion apps/sim/background/knowledge-processing.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ describe('knowledge processing worker', () => {
}
return value
})
mockProcessDocumentAsync.mockResolvedValue(undefined)
mockProcessDocumentAsync.mockResolvedValue({ outcome: 'indexed' })
mockResolveTriggerRegion.mockResolvedValue('us-east-1')
mockTrigger.mockResolvedValue({ id: 'quota-continuation-run' })
})
Expand All @@ -128,6 +128,27 @@ describe('knowledge processing worker', () => {
vi.restoreAllMocks()
})

it('reports indexed only when the document service committed the index', async () => {
expect(await runDocumentProcessing(WORKSPACE_PAYLOAD)).toMatchObject({
success: true,
outcome: 'indexed',
documentId: WORKSPACE_PAYLOAD.documentId,
})
})

it.each(['unavailable', 'not_claimed', 'superseded'] as const)(
'reports a harmless %s skip without turning it into a task failure or an indexed success',
async (reason) => {
mockProcessDocumentAsync.mockResolvedValue({ outcome: 'skipped', reason })
expect(await runDocumentProcessing(WORKSPACE_PAYLOAD)).toMatchObject({
success: false,
outcome: 'skipped',
reason,
})
expect(mockTrigger).not.toHaveBeenCalled()
}
)

it('rejects workspace work without attribution before document processing starts', async () => {
await expect(
runDocumentProcessing({
Expand Down
7 changes: 4 additions & 3 deletions apps/sim/background/knowledge-processing.ts
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ export async function runDocumentProcessing(
logger.info(`[${requestId}] Starting Trigger.dev processing for document: ${docData.filename}`)

try {
await processDocumentAsync(
const result = await processDocumentAsync(
knowledgeBaseId,
documentId,
docData,
Expand Down Expand Up @@ -90,10 +90,11 @@ export async function runDocumentProcessing(
}
)

logger.info(`[${requestId}] Successfully processed document: ${docData.filename}`)
logger.info(`[${requestId}] Document processing finished`, { documentId, ...result })

return {
success: true,
success: result.outcome === 'indexed',
...result,
documentId,
filename: docData.filename,
processingTime: Date.now() - startedAt,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,21 @@ import {
} from '@sim/db/schema'
import { generateId } from '@sim/utils/id'
import { and, eq, inArray, sql } from 'drizzle-orm'
import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'
import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from 'vitest'

const fixture = vi.hoisted(() => ({ root: '', embeddingCalls: 0 }))
const fixture = vi.hoisted(() => ({
root: '',
embeddingCalls: 0,
queueEnabled: false,
listRuns: vi.fn(),
}))
vi.mock('@/lib/core/config/trigger-runtime', () => ({
isInsideTriggerRun: () => fixture.queueEnabled,
}))
vi.mock('@trigger.dev/sdk', async (original) => ({
...(await original<typeof import('@trigger.dev/sdk')>()),
runs: { list: fixture.listRuns },
}))
vi.mock('@/lib/uploads/core/setup.server', () => ({
get UPLOAD_DIR_SERVER() {
return fixture.root
Expand Down Expand Up @@ -52,6 +64,7 @@ import { searchKnowledge } from '@/lib/knowledge/application/search'
import { searchScopedKnowledge } from '@/lib/knowledge/application/workspace-search'
import { createContentSyncLease } from '@/lib/knowledge/connectors/sync-lock'
import { addDocument } from '@/lib/knowledge/connectors/sync-persistence'
import { sweepStuckDocuments } from '@/lib/knowledge/connectors/sync-primitives'
import { knowledgeDocumentProcessingOutboxHandlers } from '@/lib/knowledge/documents/processing-outbox-handler'
import {
DOCUMENT_RECOVERY_BATCH_SIZE,
Expand All @@ -60,6 +73,7 @@ import {
} from '@/lib/knowledge/documents/processing-recovery'
import { processDocumentAsync } from '@/lib/knowledge/documents/service'
import { MAX_PROCESSING_ATTEMPTS, QUEUED_DISPATCH_GRACE_MS } from '@/lib/knowledge/documents/types'
import type { SyncResult } from '@/connectors/types'

const fixtures: ReturnType<typeof createKnowledgeAclFixtureIds>[] = []
const old = () => new Date(Date.now() - QUEUED_DISPATCH_GRACE_MS - 60_000)
Expand Down Expand Up @@ -119,6 +133,11 @@ async function failedFile(
return file
}

afterEach(() => {
fixture.queueEnabled = false
fixture.listRuns.mockReset()
})

beforeAll(() => {
fixture.root = mkdtempSync(path.join(tmpdir(), 'sim-stored-recovery-'))
})
Expand All @@ -137,7 +156,107 @@ afterAll(async () => {
await db.$client.end()
})

async function recoverFixture(
ids: ReturnType<typeof createKnowledgeAclFixtureIds>,
mode: 'independent' | 'connector'
) {
if (mode === 'independent') return recoverKnowledgeDocumentProcessing()
const result: SyncResult = {
docsAdded: 0,
docsUpdated: 0,
docsDeleted: 0,
docsUnchanged: 0,
docsSkipped: 0,
docsFailed: 0,
processingDispatch: { requested: 0, accepted: 0, failed: 0 },
}
await sweepStuckDocuments({
connectorId: ids.connectorId,
knowledgeBaseId: ids.knowledgeBaseId,
syncStartedAt: new Date(),
retryCutoff: new Date(Date.now() - 7 * 24 * 60 * 60_000),
billingAttribution: await resolveSystemBillingAttribution(ids.workspaceId),
result,
lease: createContentSyncLease(ids.connectorId, ids.lockId),
})
return result.processingDispatch.requested
}

describe('independent recovery of retained connector documents', () => {
it.each(['independent', 'connector'] as const)(
'%s recovery preserves a job queued beyond the grace period',
async (mode) => {
const ids = await seed()
const file = await failedFile(ids)
await db
.update(document)
.set({ processingStatus: 'pending' })
.where(eq(document.id, file.documentId))
fixture.queueEnabled = true
fixture.listRuns.mockResolvedValue({ data: [{ id: 'run-queued', status: 'QUEUED' }] })
expect(await recoverFixture(ids, mode)).toBe(0)
const [row] = await db.select().from(document).where(eq(document.id, file.documentId))
expect(row.processingAttempts).toBe(1)
expect(row.processingQueueToken).toBe('old-fixture-generation')
expect(row.processingRecoveryAfter).not.toBeNull()
expect(await eventsFor(ids)).toHaveLength(0)
await db
.update(knowledgeConnector)
.set({ status: 'paused' })
.where(eq(knowledgeConnector.id, ids.connectorId))
}
)

it.each(['independent', 'connector'] as const)(
'%s recovery rechecks the generation after its remote lookup',
async (mode) => {
const ids = await seed()
const file = await failedFile(ids)
fixture.queueEnabled = true
fixture.listRuns.mockImplementation(async () => {
await db
.update(document)
.set({ processingQueueToken: 'replacement-generation' })
.where(eq(document.id, file.documentId))
return { data: [] }
})
expect(await recoverFixture(ids, mode)).toBe(0)
const [row] = await db.select().from(document).where(eq(document.id, file.documentId))
expect(row.processingAttempts).toBe(1)
expect(row.processingQueueToken).toBe('replacement-generation')
expect(await eventsFor(ids)).toHaveLength(0)
await db
.update(knowledgeConnector)
.set({ status: 'paused' })
.where(eq(knowledgeConnector.id, ids.connectorId))
}
)

it.each(['pending', 'processing'])(
'does not replace an aged %s outbox continuation',
async (status) => {
const ids = await seed()
const file = await failedFile(ids)
const token = generateId()
await db
.update(document)
.set({ processingQueueToken: token })
.where(eq(document.id, file.documentId))
await db.insert(outboxEvent).values({
id: token,
eventType: 'knowledge.document.processing.resume',
payload: { knowledgeBaseId: ids.knowledgeBaseId, documentId: file.documentId },
status,
availableAt: old(),
})
expect(await recoverFixture(ids, 'independent')).toBe(0)
expect(await recoverFixture(ids, 'connector')).toBe(0)
const [row] = await db.select().from(document).where(eq(document.id, file.documentId))
expect(row.processingAttempts).toBe(1)
expect(row.processingQueueToken).toBe(token)
}
)

it('uses the organization owner and preserves Search visibility during source backoff', async () => {
const ids = await seed()
await db.insert(member).values({
Expand Down
12 changes: 10 additions & 2 deletions apps/sim/lib/knowledge/connectors/sync-primitives.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand All @@ -28,6 +28,10 @@ import {
updateDocument,
} from '@/lib/knowledge/connectors/sync-persistence'
import { documentProcessingRecoveryCondition } from '@/lib/knowledge/documents/processing-recovery-policy'
import {
documentRecoveryGenerationCondition,
filterAbandonedDocumentProcessing,
} from '@/lib/knowledge/documents/processing-recovery-queue'
import { DOCUMENT_PROCESSING_STALE_THRESHOLD_MS } from '@/lib/knowledge/documents/processing-timeouts.server'
import type { DocumentData } from '@/lib/knowledge/documents/service'
import { isTriggerAvailable, processDocumentsWithQueue } from '@/lib/knowledge/documents/service'
Expand Down Expand Up @@ -1336,6 +1340,7 @@ export async function sweepStuckDocuments(input: SweepStuckDocumentsInput): Prom
fileSize: document.fileSize,
mimeType: document.mimeType,
processingStatus: document.processingStatus,
processingQueueToken: document.processingQueueToken,
processingQueuedAt: document.processingQueuedAt,
processingStartedAt: document.processingStartedAt,
processingDeferredUntil: document.processingDeferredUntil,
Expand All @@ -1361,7 +1366,8 @@ export async function sweepStuckDocuments(input: SweepStuckDocumentsInput): Prom
asc(document.id)
)
.limit(STUCK_RETRY_MAX_CANDIDATES_PER_SYNC)
const stuckDocs = sweepCandidates.filter(
const abandoned = await filterAbandonedDocumentProcessing(sweepCandidates)
const stuckDocs = abandoned.filter(
(row): row is typeof row & { processingStatus: DocumentProcessingStatus } =>
isDocumentProcessingStatus(row.processingStatus)
)
Expand Down Expand Up @@ -1402,6 +1408,7 @@ export async function sweepStuckDocuments(input: SweepStuckDocumentsInput): Prom
fileSize: document.fileSize,
mimeType: document.mimeType,
processingStatus: document.processingStatus,
processingQueueToken: document.processingQueueToken,
processingQueuedAt: document.processingQueuedAt,
processingStartedAt: document.processingStartedAt,
processingDeferredUntil: document.processingDeferredUntil,
Expand All @@ -1412,6 +1419,7 @@ export async function sweepStuckDocuments(input: SweepStuckDocumentsInput): Prom
.where(
and(
inArray(document.id, stuckDocIds),
or(...stuckDocs.map(documentRecoveryGenerationCondition)),
eq(document.connectorId, connectorId),
documentProcessingRecoveryCondition(sweepEvaluatedAt, retryCutoff),
lt(document.uploadedAt, syncStartedAt)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1460,11 +1460,46 @@ describe('processDocumentAsync write guards', () => {
expect(schedule).not.toHaveBeenCalled()
})

it('reports an unavailable document without parsing or indexing it', async () => {
dbChainMockFns.limit.mockResolvedValueOnce([])
expect(
await processDocumentAsync(
'knowledge-base-1',
'document-1',
PERSISTED_CONTEXT,
{},
BILLING_ATTRIBUTION
)
).toEqual({ outcome: 'skipped', reason: 'unavailable' })
expect(mockProcessDocument).not.toHaveBeenCalled()
})

it('reports discarded output when the generation changes before the index commit', async () => {
armProviderSource()
dbChainMockFns.limit.mockReset()
dbChainMockFns.limit
.mockResolvedValueOnce([PERSISTED_CONTEXT])
.mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW])
.mockResolvedValueOnce([])
expect(
await processDocumentAsync(
'knowledge-base-1',
'document-1',
PERSISTED_CONTEXT,
{},
BILLING_ATTRIBUTION
)
).toEqual({ outcome: 'skipped', reason: 'superseded' })
expect(
dbChainMockFns.set.mock.calls.some(([value]) => value.processingStatus === 'completed')
).toBe(false)
})

it('does not parse or reschedule a superseded provider continuation', async () => {
armProviderSource()
dbChainMockFns.returning.mockResolvedValueOnce([])
const schedule = vi.fn()
await processDocumentAsync(
const result = await processDocumentAsync(
'knowledge-base-1',
'document-1',
PERSISTED_CONTEXT,
Expand All @@ -1477,6 +1512,7 @@ describe('processDocumentAsync write guards', () => {
scheduleProviderContinuation: schedule,
}
)
expect(result).toEqual({ outcome: 'skipped', reason: 'not_claimed' })
expect(mockProcessDocument).not.toHaveBeenCalled()
expect(schedule).not.toHaveBeenCalled()
})
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { document } from '@sim/db/schema'
import { document, outboxEvent } from '@sim/db/schema'
import { and, eq, gt, isNotNull, isNull, lt, lte, or, sql } from 'drizzle-orm'
import { DOCUMENT_PROCESSING_STALE_THRESHOLD_MS } from '@/lib/knowledge/documents/processing-timeouts.server'
import { MAX_PROCESSING_ATTEMPTS, QUEUED_DISPATCH_GRACE_MS } from '@/lib/knowledge/documents/types'
Expand All @@ -15,6 +15,11 @@ export function documentProcessingRecoveryCondition(
return and(
sql`${document.processingStatus} IN ('pending', 'processing', 'failed')`,
isNotNull(document.connectorId),
sql`NOT EXISTS (
SELECT 1 FROM ${outboxEvent}
WHERE ${outboxEvent.id} = ${document.processingQueueToken}
AND ${outboxEvent.status} IN ('pending', 'processing')
)`,
isNotNull(document.contentHash),
isNotNull(document.storageKey),
eq(document.userExcluded, false),
Expand Down
Loading
Loading