Skip to content

Commit ea41888

Browse files
icecrasher321icecrasher321
committed
fix(knowledge): protect connector uploads during attachment
1 parent b85b369 commit ea41888

10 files changed

Lines changed: 199 additions & 13 deletions

‎apps/sim/lib/knowledge/__integration__/connector-upload.integration.ts‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { tmpdir } from 'node:os'
55
import path from 'node:path'
66
import { db } from '@sim/db'
77
import {
8+
document,
89
knowledgeBase,
910
organization,
1011
outboxEvent,
@@ -29,9 +30,12 @@ import {
2930
seedKnowledgeAclFixture,
3031
} from '@/lib/knowledge/__integration__/seed-source-access-fixture'
3132
import { uploadConnectorArtifact } from '@/lib/knowledge/connectors/connector-upload'
33+
import { stillHoldsSyncLock } from '@/lib/knowledge/connectors/sync-lock'
34+
import { addDocument, updateDocument } from '@/lib/knowledge/connectors/sync-persistence'
3235
import * as cleanup from '@/lib/knowledge/documents/storage-cleanup'
3336
import * as storage from '@/lib/uploads/core/storage-service'
3437
import { getFileMetadataByKeys } from '@/lib/uploads/server/metadata'
38+
import type { ExternalDocument } from '@/connectors/types'
3539

3640
describe('connector upload crash recovery', () => {
3741
const ids = createKnowledgeAclFixtureIds()
@@ -140,6 +144,134 @@ describe('connector upload crash recovery', () => {
140144
expect(await getFileMetadataByKeys([fixture.key], 'knowledge-base')).toEqual([])
141145
})
142146

147+
it.each(['add', 'update'] as const)(
148+
'protects an uploaded artifact while %s waits for the knowledge-base lock',
149+
async (operation) => {
150+
const documentId = generateId()
151+
const source: ExternalDocument = {
152+
externalId: generateId(),
153+
title: 'Contended source',
154+
content: 'Synthetic updated source content',
155+
mimeType: 'text/plain',
156+
contentHash: 'updated-content',
157+
}
158+
if (operation === 'update') {
159+
await db.insert(document).values({
160+
id: documentId,
161+
knowledgeBaseId: ids.knowledgeBaseId,
162+
connectorId: ids.connectorId,
163+
externalId: source.externalId,
164+
filename: source.title,
165+
fileUrl: 'data:text/plain,Previous%20content',
166+
fileSize: 16,
167+
mimeType: 'text/plain',
168+
processingStatus: 'completed',
169+
})
170+
}
171+
172+
let releaseKb: (() => void) | undefined
173+
let announceKbLock: ((pid: number) => void) | undefined
174+
const kbReleased = new Promise<void>((resolve) => {
175+
releaseKb = resolve
176+
})
177+
const kbLocked = new Promise<number>((resolve) => {
178+
announceKbLock = resolve
179+
})
180+
const blocker = db.transaction(async (tx) => {
181+
const [row] = await tx.execute<{ pid: number }>(sql`SELECT pg_backend_pid() AS pid`)
182+
await tx.execute(
183+
sql`SELECT id FROM knowledge_base WHERE id = ${ids.knowledgeBaseId} FOR UPDATE`
184+
)
185+
announceKbLock?.(row.pid)
186+
await kbReleased
187+
})
188+
const blockerPid = await kbLocked
189+
190+
let releaseUpload: (() => void) | undefined
191+
let announceUpload:
192+
| ((file: Awaited<ReturnType<typeof storage.uploadFile>>) => void)
193+
| undefined
194+
const uploadReleased = new Promise<void>((resolve) => {
195+
releaseUpload = resolve
196+
})
197+
const uploaded = new Promise<Awaited<ReturnType<typeof storage.uploadFile>>>((resolve) => {
198+
announceUpload = resolve
199+
})
200+
const originalUpload = storage.uploadFile
201+
const upload = vi.spyOn(storage, 'uploadFile').mockImplementation(async (options) => {
202+
const file = await originalUpload(options)
203+
announceUpload?.(file)
204+
await uploadReleased
205+
return file
206+
})
207+
const args = [
208+
ids.knowledgeBaseId,
209+
ids.connectorId,
210+
'confluence',
211+
source,
212+
{ workspaceId: ids.workspaceId, userId: ids.aliceId },
213+
undefined,
214+
'workspace',
215+
{ stillHeld: () => stillHoldsSyncLock(ids.connectorId, ids.lockId) },
216+
] as const
217+
const attachment =
218+
operation === 'add' ? addDocument(...args) : updateDocument(documentId, ...args)
219+
const settled = attachment.then(
220+
(value) => ({ value }),
221+
(error: unknown) => ({ error })
222+
)
223+
try {
224+
const file = await uploaded
225+
const [event] = await db
226+
.select()
227+
.from(outboxEvent)
228+
.where(
229+
sql`${outboxEvent.eventType} = ${cleanup.KNOWLEDGE_STORAGE_CLEANUP_EVENT} AND ${outboxEvent.payload}->>'key' = ${file.key}`
230+
)
231+
.limit(1)
232+
events.push(event.id)
233+
await db
234+
.update(outboxEvent)
235+
.set({ availableAt: new Date(0) })
236+
.where(eq(outboxEvent.id, event.id))
237+
releaseUpload?.()
238+
await expect
239+
.poll(
240+
async () => {
241+
const waiting = await db.execute(
242+
sql`SELECT 1 FROM pg_stat_activity WHERE ${blockerPid} = ANY(pg_blocking_pids(pid)) LIMIT 1`
243+
)
244+
return waiting.length > 0
245+
},
246+
{ interval: 1, timeout: 5000 }
247+
)
248+
.toBe(true)
249+
250+
const handlers = {
251+
[cleanup.KNOWLEDGE_STORAGE_CLEANUP_EVENT]: cleanup.cleanupKnowledgeStorage,
252+
}
253+
expect(await processOutboxEventById(event.id, handlers)).toBe('pending')
254+
expect(await readFile(path.join(fixtureStorage.root, file.key), 'utf8')).toBe(
255+
source.content
256+
)
257+
releaseKb?.()
258+
await blocker
259+
const result = await settled
260+
expect(result).toHaveProperty('value')
261+
expect(await processOutboxEventById(event.id, handlers)).toBe('completed')
262+
expect(await getFileMetadataByKeys([file.key], 'knowledge-base')).toHaveLength(1)
263+
expect(await readFile(path.join(fixtureStorage.root, file.key), 'utf8')).toBe(
264+
source.content
265+
)
266+
} finally {
267+
releaseUpload?.()
268+
releaseKb?.()
269+
await Promise.allSettled([blocker, settled])
270+
upload.mockRestore()
271+
}
272+
}
273+
)
274+
143275
it('preserves an older unbound object when a create-only upload encounters a key collision', async () => {
144276
const fixture = input()
145277
await storage.uploadFile({

‎apps/sim/lib/knowledge/connectors/connector-upload.test.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ const mocks = vi.hoisted(() => ({ insert: vi.fn(), enqueue: vi.fn(), upload: vi.
66
vi.mock('@/lib/uploads', () => ({ StorageService: { uploadFile: mocks.upload } }))
77
vi.mock('@/lib/uploads/server/metadata', () => ({ insertImmutableFileMetadata: mocks.insert }))
88
vi.mock('@/lib/knowledge/documents/storage-cleanup', () => ({
9+
KNOWLEDGE_STORAGE_CLEANUP_EVENT: 'knowledge.document.storage.cleanup',
910
isKnowledgeBaseOwnedStorageKey: (key: string) => key.startsWith('kb/'),
1011
enqueueKnowledgeStorageCleanup: mocks.enqueue,
1112
}))
@@ -31,7 +32,7 @@ describe('connector upload reservation', () => {
3132
id: options.id,
3233
contentUpdatedAt: new Date(0),
3334
}))
34-
mocks.enqueue.mockResolvedValue(undefined)
35+
mocks.enqueue.mockResolvedValue(['cleanup-guard'])
3536
mocks.upload.mockResolvedValue({ key: input.key, path: `/api/files/serve/${input.key}` })
3637
})
3738
afterEach(() => vi.useRealTimers())
@@ -52,6 +53,7 @@ describe('connector upload reservation', () => {
5253
})
5354
expect(mocks.enqueue.mock.calls[0][3]).toMatchObject({ uploadId: options.createOnlyUploadId })
5455
expect(uploaded.metadataId).toBe(mocks.insert.mock.calls[0][0].id)
56+
expect(uploaded.cleanupEventId).toBe('cleanup-guard')
5557
})
5658

5759
it('does not write bytes if durable cleanup cannot be enqueued', async () => {

‎apps/sim/lib/knowledge/connectors/connector-upload.ts‎

Lines changed: 29 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
import { db } from '@sim/db'
2+
import { outboxEvent } from '@sim/db/schema'
23
import { generateId } from '@sim/utils/id'
3-
import { sql } from 'drizzle-orm'
4+
import { and, eq, sql } from 'drizzle-orm'
45
import { resourceScopeFromOwner } from '@/lib/core/resource-scope'
6+
import type { DbTransaction } from '@/lib/db/types'
57
import type { KnowledgeBaseOwner } from '@/lib/knowledge/connectors/sync-persistence'
68
import {
79
enqueueKnowledgeStorageCleanup,
810
isKnowledgeBaseOwnedStorageKey,
11+
KNOWLEDGE_STORAGE_CLEANUP_EVENT,
912
} from '@/lib/knowledge/documents/storage-cleanup'
1013
import { StorageService } from '@/lib/uploads'
1114
import { insertImmutableFileMetadata } from '@/lib/uploads/server/metadata'
@@ -31,7 +34,7 @@ export async function uploadConnectorArtifact(input: {
3134
}
3235
const metadataId = generateId()
3336
const uploadId = generateId()
34-
const binding = await db.transaction(async (tx) => {
37+
const { binding, cleanupEventId } = await db.transaction(async (tx) => {
3538
await tx.execute(sql`SET LOCAL lock_timeout = '5s'`)
3639
await tx.execute(sql`SET LOCAL statement_timeout = '15s'`)
3740
const reserved = await insertImmutableFileMetadata(
@@ -49,7 +52,7 @@ export async function uploadConnectorArtifact(input: {
4952
tx
5053
)
5154
if (reserved.id !== metadataId) throw new Error('Connector upload storage key is already bound')
52-
await enqueueKnowledgeStorageCleanup(
55+
const [cleanupEventId] = await enqueueKnowledgeStorageCleanup(
5356
tx,
5457
[{ id: documentId, fileUrl: `/api/files/serve/${encodeURIComponent(key)}`, ...owner }],
5558
documentId,
@@ -59,7 +62,8 @@ export async function uploadConnectorArtifact(input: {
5962
uploadId,
6063
}
6164
)
62-
return reserved
65+
if (!cleanupEventId) throw new Error('Connector upload cleanup guard was not created')
66+
return { binding: reserved, cleanupEventId }
6367
})
6468

6569
const controller = new AbortController()
@@ -87,8 +91,28 @@ export async function uploadConnectorArtifact(input: {
8791
})
8892
controller.signal.throwIfAborted()
8993
if (file.key !== key) throw new Error('Connector upload changed its reserved storage key')
90-
return { ...file, metadataId, contentUpdatedAt: binding.contentUpdatedAt }
94+
return { ...file, metadataId, contentUpdatedAt: binding.contentUpdatedAt, cleanupEventId }
9195
} finally {
9296
clearTimeout(timer)
9397
}
9498
}
99+
100+
/** Holds the upload's pending cleanup guard before taking KB/connector locks, until attachment commits or rolls back. */
101+
export async function claimConnectorUploadForAttachment(
102+
tx: DbTransaction,
103+
cleanupEventId: string
104+
): Promise<void> {
105+
const [guard] = await tx
106+
.select({ id: outboxEvent.id })
107+
.from(outboxEvent)
108+
.where(
109+
and(
110+
eq(outboxEvent.id, cleanupEventId),
111+
eq(outboxEvent.eventType, KNOWLEDGE_STORAGE_CLEANUP_EVENT),
112+
eq(outboxEvent.status, 'pending')
113+
)
114+
)
115+
.for('update')
116+
.limit(1)
117+
if (!guard) throw new Error('Connector upload expired before it could be attached')
118+
}

‎apps/sim/lib/knowledge/connectors/sync-content-pass.test.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,10 @@ const mocks = vi.hoisted(() => ({
1313
upload: vi.fn(),
1414
deleteFile: vi.fn(),
1515
deleteMetadata: vi.fn(),
16-
enqueueCleanup: vi.fn(),
16+
enqueueCleanup: vi.fn(async () => {
17+
queueTableRows(schemaMock.outboxEvent, [{ id: 'cleanup-guard' }])
18+
return ['cleanup-guard']
19+
}),
1720
dispatch: vi.fn(),
1821
onPage: vi.fn(),
1922
}))
@@ -38,6 +41,7 @@ vi.mock('@/lib/uploads/server/metadata', () => ({
3841
}),
3942
}))
4043
vi.mock('@/lib/knowledge/documents/storage-cleanup', () => ({
44+
KNOWLEDGE_STORAGE_CLEANUP_EVENT: 'knowledge.document.storage.cleanup',
4145
enqueueKnowledgeStorageCleanup: mocks.enqueueCleanup,
4246
isKnowledgeBaseOwnedStorageKey: (key: string) => key.startsWith('kb/'),
4347
}))

‎apps/sim/lib/knowledge/connectors/sync-engine.test.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,10 @@ vi.mock('@/lib/uploads', () => ({ StorageService: { uploadFile: mockUploadFile }
4343
const { mockDeleteFile, mockDeleteFileMetadata, mockEnqueueStorageCleanup } = vi.hoisted(() => ({
4444
mockDeleteFile: vi.fn(),
4545
mockDeleteFileMetadata: vi.fn(),
46-
mockEnqueueStorageCleanup: vi.fn(),
46+
mockEnqueueStorageCleanup: vi.fn(async () => {
47+
queueTableRows(schemaMock.outboxEvent, [{ id: 'cleanup-guard' }])
48+
return ['cleanup-guard']
49+
}),
4750
}))
4851
vi.mock('@/lib/uploads/core/storage-service', () => ({ deleteFile: mockDeleteFile }))
4952
const bindings = vi.hoisted(() => new Map<string, { id: string; contentUpdatedAt: Date }>())
@@ -59,6 +62,7 @@ vi.mock('@/lib/uploads/server/metadata', () => ({
5962
}),
6063
}))
6164
vi.mock('@/lib/knowledge/documents/storage-cleanup', () => ({
65+
KNOWLEDGE_STORAGE_CLEANUP_EVENT: 'knowledge.document.storage.cleanup',
6266
enqueueKnowledgeStorageCleanup: mockEnqueueStorageCleanup,
6367
isKnowledgeBaseOwnedStorageKey: (key: string) => key.startsWith('kb/'),
6468
}))

‎apps/sim/lib/knowledge/connectors/sync-persistence.test.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,11 @@ vi.mock('@/lib/uploads/server/metadata', () => ({
2020
}),
2121
}))
2222
vi.mock('@/lib/knowledge/documents/storage-cleanup', () => ({
23-
enqueueKnowledgeStorageCleanup: vi.fn(),
23+
KNOWLEDGE_STORAGE_CLEANUP_EVENT: 'knowledge.document.storage.cleanup',
24+
enqueueKnowledgeStorageCleanup: vi.fn(async () => {
25+
queueTableRows(schemaMock.outboxEvent, [{ id: 'cleanup-guard' }])
26+
return ['cleanup-guard']
27+
}),
2428
isKnowledgeBaseOwnedStorageKey: (key: string) => key.startsWith('kb/'),
2529
}))
2630
vi.mock('@/connectors/registry.server', () => ({ CONNECTOR_REGISTRY: {} }))

‎apps/sim/lib/knowledge/connectors/sync-persistence.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,10 @@ import {
1414
} from '@/lib/knowledge/access/tokens'
1515
import type { MirroredDocumentAcl } from '@/lib/knowledge/access/types'
1616
import { aclIsDerived, type ConnectorAccessMode } from '@/lib/knowledge/connectors/access-modes'
17-
import { uploadConnectorArtifact } from '@/lib/knowledge/connectors/connector-upload'
17+
import {
18+
claimConnectorUploadForAttachment,
19+
uploadConnectorArtifact,
20+
} from '@/lib/knowledge/connectors/connector-upload'
1821
import { resolveSourceModifiedAt } from '@/lib/knowledge/connectors/source-modified-at'
1922
import { SOURCE_CONTENT_ERROR } from '@/lib/knowledge/connectors/sync-limits'
2023
import { assertSyncLeaseHeldInTx, type SyncWriteLease } from '@/lib/knowledge/connectors/sync-lock'
@@ -546,6 +549,7 @@ export async function addDocument(
546549
? resolveTagMapping(connectorType, extDoc.metadata, sourceConfig)
547550
: undefined
548551
await db.transaction(async (tx) => {
552+
await claimConnectorUploadForAttachment(tx, fileInfo.cleanupEventId)
549553
const isActive = await isKnowledgeBaseActiveInTx(tx, knowledgeBaseId)
550554
if (!isActive) {
551555
throw new Error(`Knowledge base ${knowledgeBaseId} is deleted`)
@@ -648,6 +652,7 @@ export async function updateDocument(
648652
? resolveTagMapping(connectorType, extDoc.metadata, sourceConfig)
649653
: undefined
650654
await db.transaction(async (tx) => {
655+
await claimConnectorUploadForAttachment(tx, fileInfo.cleanupEventId)
651656
const isActive = await isKnowledgeBaseActiveInTx(tx, knowledgeBaseId)
652657
if (!isActive) {
653658
throw new Error(`Knowledge base ${knowledgeBaseId} is deleted`)

‎apps/sim/lib/knowledge/documents/storage-cleanup.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ Each releasing mutation receives a fresh cleanup event ID. A previous event may
88

99
Upload guards also bind the provider upload ID. A crash before the write releases the unused metadata reservation; a crash after the write deletes the matching unreferenced object and reservation. A create-only conflict cannot delete an older object with a different upload ID. Metadata is not registered a second time after the write, so a late worker cannot restore a reservation already removed by cleanup.
1010

11+
Attachment locks its pending cleanup event before waiting for the KB or connector locks. The outbox worker skips that locked event even if its grace period expires during attachment. Commit makes the document reference visible before releasing the guard; rollback releases the guard so orphan cleanup can proceed. The existing KB, connector, and metadata lock order stays intact.
12+
1113
Enqueue reads and inserts at most 100 objects per batch. Each event deletes one object, has a 15-second storage deadline, uses a five-second lock timeout, and has 48 bounded outbox attempts. Exhausted jobs remain visible as dead letters with their identity and final error for operator recovery.
1214

1315
Comparing a `Date` against a `date_trunc(...)` SQL expression must explicitly encode the timestamp parameter. The shared metadata function binds an ISO timestamp with a PostgreSQL timestamp cast. Real PostgreSQL tests reproduce the driver encoding failure and prove deletion of a timestamp with microsecond precision.

‎apps/sim/lib/knowledge/documents/storage-cleanup.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ describe('durable knowledge storage cleanup', () => {
9696
})
9797

9898
it('propagates persistence failures before the parent transaction can commit', async () => {
99-
dbChainMockFns.onConflictDoNothing.mockRejectedValueOnce(new Error('Database unavailable'))
99+
dbChainMockFns.returning.mockRejectedValueOnce(new Error('Database unavailable'))
100100
await expect(
101101
enqueueKnowledgeStorageCleanup(
102102
db,

‎apps/sim/lib/knowledge/documents/storage-cleanup.ts‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,8 @@ export async function enqueueKnowledgeStorageCleanup(
8888
documents: readonly KnowledgeStorageCleanupDocument[],
8989
requestId: string,
9090
options?: { availableAt?: Date; reason?: 'uncommitted-upload'; uploadId?: string }
91-
): Promise<void> {
91+
): Promise<string[]> {
92+
const eventIds: string[] = []
9293
for (let offset = 0; offset < documents.length; offset += ENQUEUE_BATCH_SIZE) {
9394
const entries = documents.slice(offset, offset + ENQUEUE_BATCH_SIZE).flatMap((doc) => {
9495
const key = getKnowledgeBaseStorageKey(doc.fileUrl)
@@ -132,8 +133,16 @@ export async function enqueueKnowledgeStorageCleanup(
132133
...(options?.availableAt ? { availableAt: options.availableAt } : {}),
133134
})
134135
}
135-
if (rows.length) await executor.insert(outboxEvent).values(rows).onConflictDoNothing()
136+
if (rows.length) {
137+
const inserted = await executor
138+
.insert(outboxEvent)
139+
.values(rows)
140+
.onConflictDoNothing()
141+
.returning({ id: outboxEvent.id })
142+
eventIds.push(...inserted.map((row) => row.id))
143+
}
136144
}
145+
return eventIds
137146
}
138147

139148
function isMissingObject(error: unknown): boolean {

0 commit comments

Comments
 (0)