Skip to content

Commit 27600ef

Browse files
committed
fix(knowledge): validate bundles before writing, close import cancel and cleanup gaps, and fix CI
- read every chunk entry before the base is created, reject reversed offsets, and length-check vectors before decoding - confirm the import is still running before creating the base, and remove it if a cancel wins - record the import inside upload completion, so a refusal or crash cannot strand the bundle; re-read on a concurrent completion - serialize progress writes, keep the source character count, cap rebuilt text, and recheck uploads when issuing part URLs - janitor fails in-process documents a dead worker left processing, and deletes bundles when pruning settled imports - shrink the unused-exports baseline, count the new janitor passes in its test, and fix the integration test's Blob typing
1 parent 58f6977 commit 27600ef

10 files changed

Lines changed: 375 additions & 150 deletions

File tree

‎apps/sim/background/cleanup-stale-executions.test.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -301,9 +301,9 @@ describe('stale execution cleanup deadline grace', () => {
301301

302302
it('claims every cleanup page without overlapping concurrent workers', async () => {
303303
await runCleanupStaleExecutions()
304-
// Nine batched arms: the connector sync-log retention pass is the newest.
305-
expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(9)
306-
expect(dbChainMockFns.for).toHaveBeenCalledTimes(9)
304+
// Eleven batched arms: the knowledge import stale and retention passes are the newest.
305+
expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(11)
306+
expect(dbChainMockFns.for).toHaveBeenCalledTimes(11)
307307
for (const [strength, options] of dbChainMockFns.for.mock.calls) {
308308
expect(strength).toBe('update')
309309
expect(options).toEqual({ skipLocked: true })
@@ -372,7 +372,7 @@ describe('stale execution cleanup deadline grace', () => {
372372
expect(mockDeleteFile).toHaveBeenCalledTimes(1000)
373373

374374
const limits = dbChainMockFns.limit.mock.calls.map(([limit]) => limit)
375-
expect(limits.filter((limit) => limit === 100)).toHaveLength(20)
375+
expect(limits.filter((limit) => limit === 100)).toHaveLength(22)
376376
expect(limits.filter((limit) => limit === 1000)).toHaveLength(30)
377377
expect(limits.filter((limit) => limit === 2000)).toHaveLength(12)
378378

‎apps/sim/background/cleanup-stale-executions.ts‎

Lines changed: 42 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import { db } from '@sim/db'
22
import {
33
asyncJobs,
4+
document,
45
knowledgeBaseImport,
56
knowledgeConnectorSyncLog,
67
tableJobs,
@@ -160,6 +161,18 @@ async function runBatchedMutation<TCandidate extends { id: string }, TRow>({
160161
return { affected, reachedLimit: affected >= maxRowsPerRun }
161162
}
162163

164+
/** Deletes import bundles one at a time; a bundle already gone is not an error. */
165+
async function deleteKnowledgeImportBundles(imports: { sourceKey: string }[]): Promise<void> {
166+
for (const { sourceKey } of imports) {
167+
await deleteFile({ key: sourceKey, context: KNOWLEDGE_IMPORT_STORAGE_CONTEXT }).catch((err) => {
168+
logger.warn('Failed to delete a knowledge import bundle', {
169+
sourceKey,
170+
error: toError(err).message,
171+
})
172+
})
173+
}
174+
}
175+
163176
/** Settles everything a dead worker left unfinished and prunes expired bookkeeping rows. */
164177
export async function runCleanupStaleExecutions() {
165178
logger.info('Starting stale execution cleanup job')
@@ -466,7 +479,9 @@ export async function runCleanupStaleExecutions() {
466479
/**
467480
* Fail knowledge imports whose worker died, and delete the bundle that worker
468481
* would have deleted in its `finally`. Documents it already imported stay in
469-
* the base, as they would after a cancel. Settled imports are pruned later.
482+
* the base, as they would after a cancel; one it was indexing in-process has
483+
* no queue behind it, so it is failed for the user to retry. Settled imports
484+
* are pruned later, deleting any bundle a canceled or crashed worker left.
470485
*/
471486
let staleKnowledgeImportsMarkedFailed = 0
472487
try {
@@ -497,18 +512,32 @@ export async function runCleanupStaleExecutions() {
497512
updatedAt: now,
498513
})
499514
.where(and(staleKnowledgeImportPredicate, inArray(knowledgeBaseImport.id, candidateIds)))
500-
.returning({ sourceKey: knowledgeBaseImport.sourceKey }),
515+
.returning({
516+
sourceKey: knowledgeBaseImport.sourceKey,
517+
knowledgeBaseId: knowledgeBaseImport.knowledgeBaseId,
518+
}),
501519
onBatch: async (imports) => {
502-
for (const { sourceKey } of imports) {
503-
await deleteFile({ key: sourceKey, context: KNOWLEDGE_IMPORT_STORAGE_CONTEXT }).catch(
504-
(err) => {
505-
logger.warn('Failed to delete a stale knowledge import bundle', {
506-
sourceKey,
507-
error: toError(err).message,
508-
})
509-
}
510-
)
520+
const knowledgeBaseIds = imports.flatMap(({ knowledgeBaseId }) =>
521+
knowledgeBaseId ? [knowledgeBaseId] : []
522+
)
523+
if (knowledgeBaseIds.length > 0) {
524+
await db
525+
.update(document)
526+
.set({
527+
processingStatus: 'failed',
528+
processingError: 'The import stopped before this document finished indexing',
529+
processingCompletedAt: now,
530+
})
531+
.where(
532+
and(
533+
inArray(document.knowledgeBaseId, knowledgeBaseIds),
534+
eq(document.processingStatus, 'processing'),
535+
isNull(document.processingQueueToken),
536+
lt(document.processingStartedAt, staleKnowledgeImportThreshold)
537+
)
538+
)
511539
}
540+
await deleteKnowledgeImportBundles(imports)
512541
},
513542
})
514543
staleKnowledgeImportsMarkedFailed = staleKnowledgeImportResult.affected
@@ -539,7 +568,8 @@ export async function runCleanupStaleExecutions() {
539568
.where(
540569
and(settledKnowledgeImportPredicate, inArray(knowledgeBaseImport.id, candidateIds))
541570
)
542-
.returning({ id: knowledgeBaseImport.id }),
571+
.returning({ sourceKey: knowledgeBaseImport.sourceKey }),
572+
onBatch: deleteKnowledgeImportBundles,
543573
})
544574
} catch (error) {
545575
logger.error('Failed to clean up stale knowledge imports:', {

‎apps/sim/lib/knowledge/__integration__/bundle-import.integration.ts‎

Lines changed: 123 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import { deleteWorkspaceFixture } from '@sim/db/testing/workspace-fixtures'
2121
import { sha256Hex } from '@sim/security/hash'
2222
import { generateId } from '@sim/utils/id'
2323
import { and, asc, eq, inArray, sql } from 'drizzle-orm'
24+
import JSZip from 'jszip'
2425
import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from 'vitest'
2526

2627
const fixture = vi.hoisted(() => ({
@@ -71,6 +72,7 @@ import {
7172
findKnowledgeImport,
7273
insertKnowledgeImport,
7374
KNOWLEDGE_IMPORT_STORAGE_CONTEXT,
75+
stopKnowledgeImport,
7476
} from '@/lib/knowledge/transfer/import-jobs'
7577
import { runKnowledgeImport } from '@/lib/knowledge/transfer/import-runner'
7678
import { embeddingVectorColumn, embeddingVectorValues } from '@/lib/knowledge/vector-columns'
@@ -243,10 +245,11 @@ afterEach(() => {
243245
fixture.target = { model: 'text-embedding-3-small', dimensions: 1536 }
244246
})
245247
afterAll(async () => {
248+
const workspaceBases = sql`SELECT ${knowledgeBase.id} FROM ${knowledgeBase} WHERE ${knowledgeBase.workspaceId} = ${ids.workspaceId}`
246249
await db
247250
.delete(outboxEvent)
248251
.where(
249-
sql`${outboxEvent.payload}::jsonb ->> 'knowledgeBaseId' IN (SELECT ${knowledgeBase.id} FROM ${knowledgeBase} WHERE ${knowledgeBase.workspaceId} = ${ids.workspaceId})`
252+
sql`${outboxEvent.payload}::jsonb ->> 'knowledgeBaseId' IN (${workspaceBases}) OR ${outboxEvent.payload}::jsonb ->> 'documentId' IN (SELECT ${document.id} FROM ${document} WHERE ${document.knowledgeBaseId} IN (${workspaceBases}))`
250253
)
251254
await db.delete(auditLog).where(eq(auditLog.workspaceId, ids.workspaceId))
252255
await db.delete(knowledgeBase).where(eq(knowledgeBase.workspaceId, ids.workspaceId))
@@ -408,6 +411,77 @@ describe('knowledge base bundle import', () => {
408411
).rejects.toThrow()
409412
})
410413

414+
it('fails a bundle with a damaged chunk entry before creating anything', async () => {
415+
const zip = new JSZip()
416+
zip.file('chunks/doc-1.ndjson', '{"index": 0, "content": "half a line"\n')
417+
zip.file(
418+
'manifest.json',
419+
JSON.stringify({
420+
version: 1,
421+
exportedAt: '2026-10-01T00:00:00.000Z',
422+
embedding: { model: 'text-embedding-3-small', dimension: 1536, vectorsIncluded: false },
423+
knowledgeBase: {
424+
name: 'Damaged fixture',
425+
description: null,
426+
chunkingConfig: { maxSize: 1024, minSize: 1, overlap: 20 },
427+
},
428+
tags: [],
429+
documents: [
430+
{
431+
id: 'doc-1',
432+
filename: 'notes.txt',
433+
mimeType: 'text/plain',
434+
fileSize: 0,
435+
enabled: true,
436+
tags: {},
437+
file: null,
438+
chunks: 'chunks/doc-1.ndjson',
439+
chunkCount: 1,
440+
tokenCount: 1,
441+
characterCount: 11,
442+
},
443+
],
444+
})
445+
)
446+
const importId = generateId()
447+
const sourceKey = await stageBundleBytes(
448+
importId,
449+
await zip.generateAsync({ type: 'nodebuffer' })
450+
)
451+
452+
const record = await runImport(importId, sourceKey, 'Damaged fixture')
453+
454+
expect(record).toMatchObject({ status: 'failed', knowledgeBaseId: null, documentsTotal: 0 })
455+
expect(record.error).toContain('chunks/doc-1.ndjson line 1 is not valid JSON')
456+
expect(
457+
await db
458+
.select({ id: knowledgeBase.id })
459+
.from(knowledgeBase)
460+
.where(eq(knowledgeBase.name, 'Damaged fixture'))
461+
).toEqual([])
462+
})
463+
464+
it('creates no base for an import canceled before its worker reached it', async () => {
465+
const importId = generateId()
466+
const sourceKey = await stageExportedBundle(importId)
467+
const staged = await findKnowledgeImport({ importId })
468+
if (!staged) throw new Error('Import was not staged')
469+
await stopKnowledgeImport(staged)
470+
471+
const record = await runImport(importId, sourceKey, 'Canceled fixture')
472+
473+
expect(record).toMatchObject({ status: 'canceled', knowledgeBaseId: null })
474+
expect(
475+
await db
476+
.select({ id: knowledgeBase.id })
477+
.from(knowledgeBase)
478+
.where(eq(knowledgeBase.name, 'Canceled fixture'))
479+
).toEqual([])
480+
await expect(
481+
downloadFile({ key: sourceKey, context: KNOWLEDGE_IMPORT_STORAGE_CONTEXT })
482+
).rejects.toThrow()
483+
})
484+
411485
it('fails without a partial base when the name is already taken', async () => {
412486
const importId = generateId()
413487
const sourceKey = await stageExportedBundle(importId)
@@ -459,7 +533,7 @@ describe('knowledge base import requests', () => {
459533
await writeLocalPutObject({
460534
uploadId: created.upload.id,
461535
key: created.upload.finalKey,
462-
body: new Blob([bytes]).stream(),
536+
body: new Blob([new Uint8Array(bytes)]).stream(),
463537
expectedSize: bytes.length,
464538
contentType: created.upload.contentType,
465539
metadata: uploadSessionObjectMetadata(created.upload),
@@ -499,7 +573,12 @@ describe('knowledge base import requests', () => {
499573
db
500574
.select({ resourceId: auditLog.resourceId })
501575
.from(auditLog)
502-
.where(eq(auditLog.action, 'knowledge_base.imported'))
576+
.where(
577+
and(
578+
eq(auditLog.workspaceId, ids.workspaceId),
579+
eq(auditLog.action, 'knowledge_base.imported')
580+
)
581+
)
503582
)
504583
.toEqual([{ resourceId: finished.knowledgeBaseId }])
505584
})
@@ -538,34 +617,62 @@ describe('knowledge base import requests', () => {
538617
})
539618

540619
describe('stale knowledge import cleanup', () => {
541-
it('fails an import whose worker went silent, deletes its bundle, and prunes long-settled imports', async () => {
620+
async function seedProcessingDocument(processingQueueToken: string | null) {
621+
const id = generateId()
622+
await db.insert(document).values({
623+
id,
624+
knowledgeBaseId: ids.knowledgeBaseId,
625+
filename: 'stranded.txt',
626+
fileUrl: '',
627+
fileSize: 0,
628+
mimeType: 'text/plain',
629+
processingStatus: 'processing',
630+
processingStartedAt: new Date(Date.now() - 2 * 60 * 60 * 1000),
631+
processingQueueToken,
632+
})
633+
return id
634+
}
635+
636+
async function processingStatus(documentId: string) {
637+
const [row] = await db
638+
.select({ status: document.processingStatus })
639+
.from(document)
640+
.where(eq(document.id, documentId))
641+
return row?.status
642+
}
643+
644+
it('fails an import whose worker went silent, its stranded documents, and its bundle, and prunes settled imports with theirs', async () => {
542645
const staleId = generateId()
543-
const sourceKey = await stageBundleBytes(staleId, Buffer.from('bundle bytes'))
646+
const staleBundle = await stageBundleBytes(staleId, Buffer.from('bundle bytes'))
544647
const settledId = generateId()
545-
await insertKnowledgeImport({
546-
id: settledId,
547-
workspaceId: ids.workspaceId,
548-
userId: ids.aliceId,
549-
sourceKey: `knowledge-import/${ids.workspaceId}/${settledId}/bundle.simkb.zip`,
550-
})
648+
const settledBundle = await stageBundleBytes(settledId, Buffer.from('bundle bytes'))
649+
const stranded = await seedProcessingDocument(null)
650+
const queued = await seedProcessingDocument('queue-token')
551651
const longAgo = new Date(Date.now() - 30 * 24 * 60 * 60 * 1000)
552652
await db
553653
.update(knowledgeBaseImport)
554-
.set({ updatedAt: new Date(Date.now() - 2 * 60 * 60 * 1000) })
654+
.set({
655+
knowledgeBaseId: ids.knowledgeBaseId,
656+
updatedAt: new Date(Date.now() - 2 * 60 * 60 * 1000),
657+
})
555658
.where(eq(knowledgeBaseImport.id, staleId))
556659
await db
557660
.update(knowledgeBaseImport)
558-
.set({ status: 'completed', updatedAt: longAgo, completedAt: longAgo })
661+
.set({ status: 'canceled', updatedAt: longAgo, completedAt: longAgo })
559662
.where(eq(knowledgeBaseImport.id, settledId))
560663

561664
await runCleanupStaleExecutions()
562665

563666
const stale = await findKnowledgeImport({ importId: staleId })
564667
expect(stale).toMatchObject({ status: 'failed' })
565668
expect(stale?.error).toContain('no progress')
566-
await expect(
567-
downloadFile({ key: sourceKey, context: KNOWLEDGE_IMPORT_STORAGE_CONTEXT })
568-
).rejects.toThrow()
669+
expect(await processingStatus(stranded)).toBe('failed')
670+
expect(await processingStatus(queued)).toBe('processing')
569671
expect(await findKnowledgeImport({ importId: settledId })).toBeNull()
672+
for (const key of [staleBundle, settledBundle]) {
673+
await expect(
674+
downloadFile({ key, context: KNOWLEDGE_IMPORT_STORAGE_CONTEXT })
675+
).rejects.toThrow()
676+
}
570677
})
571678
})

0 commit comments

Comments
 (0)