Skip to content

Commit b68731c

Browse files
committed
fix(connectors): heartbeat every unbounded phase, not just the batch loop
The heartbeat was added where we happened to be looking. The pagination loop — where a large source spends most of its wall clock, since the batch loop does not start until every page is fetched — never beat at all, so a long listing on the uncapped in-process path was still reclaimed as a hard failure. That is the exact ratchet the heartbeat exists to prevent. Auditing the remaining phases found something worse in the stuck-document retry: on the in-process path it handed the entire backlog to a single await that fully parses, embeds and indexes every document before returning, so no beat placement could interrupt it. That dispatch is now chunked, with a beat per chunk. All four call sites share one beatIfDue closure; the two pre-existing inline blocks were collapsed onto it rather than left as copies. An await longer than the TTL — one pathological listing page, or a very large hard delete — is still not covered, and no inline beat can cover it. Closing that needs a concurrent interval, which is a second mechanism and a separate decision. Adds the first test that drives executeSync itself, reaching the pagination loop through the real lock acquisition rather than testing helpers in isolation.
1 parent 5cb942f commit b68731c

2 files changed

Lines changed: 135 additions & 16 deletions

File tree

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

Lines changed: 98 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,12 @@ import {
88
flattenMockConditions,
99
hasMockCondition,
1010
type MockCondition,
11+
queueTableRows,
1112
resetDbChainMock,
1213
schemaMock,
1314
} from '@sim/testing'
1415
import { generateShortId } from '@sim/utils/id'
15-
import { beforeEach, describe, expect, it, vi } from 'vitest'
16+
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
1617
import {
1718
classifySuspectListing,
1819
evaluateListingSafety,
@@ -33,7 +34,14 @@ vi.mock('@/background/knowledge-connector-sync', () => ({
3334
knowledgeConnectorSync: { trigger: vi.fn() },
3435
}))
3536

36-
const { mockMapTags } = vi.hoisted(() => ({ mockMapTags: vi.fn() }))
37+
const { mockMapTags, mockListDocuments } = vi.hoisted(() => ({
38+
mockMapTags: vi.fn(),
39+
mockListDocuments: vi.fn(),
40+
}))
41+
42+
vi.mock('@/lib/billing/core/billing-attribution', () => ({
43+
assertBillingAttributionSnapshot: (snapshot: unknown) => snapshot,
44+
}))
3745

3846
vi.mock('@/connectors/registry.server', () => ({
3947
CONNECTOR_REGISTRY: {
@@ -43,6 +51,11 @@ vi.mock('@/connectors/registry.server', () => ({
4351
'no-tags': {
4452
name: 'No Tags',
4553
},
54+
paged: {
55+
name: 'Paged',
56+
auth: { mode: 'apiKey', optional: true },
57+
listDocuments: mockListDocuments,
58+
},
4659
},
4760
}))
4861

@@ -1530,3 +1543,86 @@ describe('heartbeatSyncLock', () => {
15301543
expect(await heartbeatSyncLock('c-1', 'run-a')).toBe(true)
15311544
})
15321545
})
1546+
1547+
describe('executeSync heartbeats during the listing phase', () => {
1548+
const CONNECTOR = {
1549+
id: 'c-1',
1550+
knowledgeBaseId: 'kb-1',
1551+
connectorType: 'paged',
1552+
credentialId: null,
1553+
encryptedApiKey: null,
1554+
sourceConfig: {},
1555+
syncMode: 'full',
1556+
syncIntervalMinutes: 1440,
1557+
status: 'active',
1558+
lastSyncAt: null,
1559+
lastSyncDocCount: null,
1560+
consecutiveFailures: 0,
1561+
syncLockToken: null,
1562+
}
1563+
1564+
beforeEach(() => {
1565+
vi.clearAllMocks()
1566+
resetDbChainMock()
1567+
vi.useFakeTimers()
1568+
vi.setSystemTime(new Date('2026-08-20T00:00:00.000Z'))
1569+
})
1570+
1571+
afterEach(() => {
1572+
vi.useRealTimers()
1573+
})
1574+
1575+
/** Drives executeSync as far as the pagination loop. */
1576+
function primeSyncUpToListing() {
1577+
queueTableRows(schemaMock.knowledgeConnector, [CONNECTOR])
1578+
queueTableRows(schemaMock.knowledgeBase, [{ userId: 'u-1', workspaceId: 'ws-1' }])
1579+
// The lock CAS; every later `.returning()` falls through to the empty default,
1580+
// which is what makes the heartbeat below report a lost lock.
1581+
dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'c-1' }])
1582+
}
1583+
1584+
it('beats between pages and abandons the run when the lock was reclaimed', async () => {
1585+
const { executeSync } = await import('@/lib/knowledge/connectors/sync-engine')
1586+
const { SYNC_LOCK_HEARTBEAT_INTERVAL_MS } = await import(
1587+
'@/lib/knowledge/connectors/sync-limits'
1588+
)
1589+
1590+
primeSyncUpToListing()
1591+
1592+
/**
1593+
* Listing is where a large source spends most of its wall clock, so a page
1594+
* that pushes the run past the heartbeat interval must trigger a beat before
1595+
* the next page — not only once listing has finished.
1596+
*/
1597+
mockListDocuments.mockImplementation(async () => {
1598+
vi.setSystemTime(new Date(Date.now() + SYNC_LOCK_HEARTBEAT_INTERVAL_MS + 1_000))
1599+
return { documents: [], hasMore: true, nextCursor: 'page-2' }
1600+
})
1601+
1602+
const result = await executeSync('c-1', {
1603+
billingAttribution: { workspaceId: 'ws-1' } as never,
1604+
})
1605+
1606+
// Aborted on the beat before page 2 rather than paging on under a lost lock.
1607+
expect(mockListDocuments).toHaveBeenCalledTimes(1)
1608+
expect(result.error).toBe('sync_superseded')
1609+
})
1610+
1611+
it('does not beat when pages return faster than the interval', async () => {
1612+
const { executeSync } = await import('@/lib/knowledge/connectors/sync-engine')
1613+
1614+
primeSyncUpToListing()
1615+
1616+
let pages = 0
1617+
mockListDocuments.mockImplementation(async () => {
1618+
pages += 1
1619+
vi.setSystemTime(new Date(Date.now() + 1_000))
1620+
return { documents: [], hasMore: pages < 3, nextCursor: `page-${pages}` }
1621+
})
1622+
1623+
await executeSync('c-1', { billingAttribution: { workspaceId: 'ws-1' } as never })
1624+
1625+
// All three pages fetched: the time gate keeps a fast listing beat-free.
1626+
expect(mockListDocuments).toHaveBeenCalledTimes(3)
1627+
})
1628+
})

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

Lines changed: 37 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,15 @@ const DEFAULT_OP_SIZE_BYTES = 4 * 1024 * 1024
7676
const CONTENT_INFLIGHT_BUDGET_BYTES = 64 * 1024 * 1024
7777
const MAX_PAGES = 500
7878
const MAX_SAFE_TITLE_LENGTH = 200
79+
/**
80+
* How many stuck documents are re-dispatched per call.
81+
*
82+
* The retry backlog is unbounded, and on the in-process fallback path
83+
* `processDocumentsWithQueue` parses, embeds, and indexes every document it is
84+
* given before returning. Handing it the whole backlog made the retry a single
85+
* await no heartbeat could interrupt; chunking gives the beat somewhere to run.
86+
*/
87+
const STUCK_RETRY_DISPATCH_CHUNK_SIZE = 25
7988
const STALE_PROCESSING_MINUTES = 45
8089
const RETRY_WINDOW_DAYS = 7
8190

@@ -1196,6 +1205,20 @@ export async function executeSync(
11961205
const syncStartedAt = new Date()
11971206
/** Seeded at lock acquisition, which wrote `updatedAt` itself. */
11981207
let lastHeartbeatAtMs = Date.now()
1208+
1209+
/**
1210+
* Refreshes the lock if the interval has elapsed, and aborts the run if it has
1211+
* been reclaimed. Called at the top of every unbounded loop in this sync — the
1212+
* time gate makes each call nearly free, so placement only has to guarantee
1213+
* that no unbounded phase runs without reaching one.
1214+
*/
1215+
const beatIfDue = async (): Promise<void> => {
1216+
if (!shouldHeartbeatSyncLock(Date.now(), lastHeartbeatAtMs)) return
1217+
if (!(await heartbeatSyncLock(connectorId, syncLogId))) {
1218+
throw new SyncLockLostException(connectorId)
1219+
}
1220+
lastHeartbeatAtMs = Date.now()
1221+
}
11991222
await db.insert(knowledgeConnectorSyncLog).values({
12001223
id: syncLogId,
12011224
connectorId,
@@ -1303,6 +1326,14 @@ export async function executeSync(
13031326
)
13041327

13051328
for (let pageNum = 0; hasMore && pageNum < MAX_PAGES; pageNum++) {
1329+
/**
1330+
* Listing is where a large source spends most of its wall clock — the
1331+
* batch loop below does not start until every page has been fetched — so
1332+
* without this a big listing outran the TTL and was reclaimed as a hard
1333+
* failure, which is the exact ratchet the heartbeat exists to prevent.
1334+
*/
1335+
await beatIfDue()
1336+
13061337
if (pageNum > 0 && connectorConfig.auth.mode === 'oauth') {
13071338
accessToken = await resolveAccessToken(connector, connectorConfig, credentialUserId)
13081339
}
@@ -1480,12 +1511,7 @@ export async function executeSync(
14801511
// per-file cap never hydrate/upload together and exhaust the worker heap.
14811512
const batches = chunkOpsByByteBudget(pendingOps, CONTENT_INFLIGHT_BUDGET_BYTES, SYNC_BATCH_SIZE)
14821513
for (const rawBatch of batches) {
1483-
if (shouldHeartbeatSyncLock(Date.now(), lastHeartbeatAtMs)) {
1484-
if (!(await heartbeatSyncLock(connectorId, syncLogId))) {
1485-
throw new SyncLockLostException(connectorId)
1486-
}
1487-
lastHeartbeatAtMs = Date.now()
1488-
}
1514+
await beatIfDue()
14891515

14901516
const liveness = await checkSyncLiveness(connectorId, connector.knowledgeBaseId)
14911517
if (liveness.connectorDeleted) {
@@ -1867,12 +1893,7 @@ export async function executeSync(
18671893
result.docsDeleted += await hardDeleteDocuments(safeHardDeleteIds, syncLogId, connectorId)
18681894
}
18691895

1870-
if (shouldHeartbeatSyncLock(Date.now(), lastHeartbeatAtMs)) {
1871-
if (!(await heartbeatSyncLock(connectorId, syncLogId))) {
1872-
throw new SyncLockLostException(connectorId)
1873-
}
1874-
lastHeartbeatAtMs = Date.now()
1875-
}
1896+
await beatIfDue()
18761897

18771898
const postBatchLiveness = await checkSyncLiveness(connectorId, connector.knowledgeBaseId)
18781899
if (postBatchLiveness.connectorDeleted) {
@@ -1976,9 +1997,11 @@ export async function executeSync(
19761997
}
19771998
})
19781999

1979-
if (retryDocs.length > 0) {
2000+
for (let i = 0; i < retryDocs.length; i += STUCK_RETRY_DISPATCH_CHUNK_SIZE) {
2001+
await beatIfDue()
2002+
19802003
await processDocumentsWithQueue(
1981-
retryDocs.map((doc) => ({
2004+
retryDocs.slice(i, i + STUCK_RETRY_DISPATCH_CHUNK_SIZE).map((doc) => ({
19822005
documentId: doc.id,
19832006
filename: doc.filename ?? 'document.txt',
19842007
fileUrl: doc.fileUrl ?? '',

0 commit comments

Comments
 (0)