diff --git a/apps/sim/lib/core/rate-limiter/provider-admission.test.ts b/apps/sim/lib/core/rate-limiter/provider-admission.test.ts index 180306504c9..8bd52a15c51 100644 --- a/apps/sim/lib/core/rate-limiter/provider-admission.test.ts +++ b/apps/sim/lib/core/rate-limiter/provider-admission.test.ts @@ -167,8 +167,8 @@ describe('provider admission', () => { maxWaitMs: 1, } let stored: { key: string; tokens: string; lastRefillAt: Date } | undefined - dbChainMockFns.values.mockImplementation((rows) => { - stored ??= rows.find((row: { key: string }) => row.key.endsWith(':requests')) + dbChainMockFns.values.mockImplementation((row) => { + stored ??= { ...row } return { onConflictDoNothing: vi.fn().mockResolvedValue(undefined) } }) dbChainMockFns.limit.mockImplementation(async () => [stored]) @@ -177,9 +177,15 @@ describe('provider admission', () => { return { where: vi.fn().mockResolvedValue(undefined) } }) const bucket = new DbTokenBucket() - consumeTokens.mockImplementation((reservations, options) => - bucket.consumeTokensAtomically(reservations, options) - ) + /** The single-dimension refill arithmetic the atomic statement mirrors in SQL. */ + consumeTokens.mockImplementation(async ([reservation]) => { + const result = await bucket.consumeTokens( + reservation.key, + reservation.cost, + reservation.config + ) + return { allowed: result.allowed, retryAfterMs: result.retryAfterMs ?? 0 } + }) for (let request = 0; request < 16; request++) await waitForProviderAdmission(input) await expect(waitForProviderAdmission(input)).rejects.toMatchObject({ retryAfterMs: 1000 }) diff --git a/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.test.ts b/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.test.ts index 6c1655d7112..443f3dc8100 100644 --- a/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.test.ts +++ b/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.test.ts @@ -44,41 +44,6 @@ describe('PostgreSQL token bucket', () => { expect(dbChainMockFns.set).toHaveBeenCalledWith(expect.objectContaining({ tokens: '0' })) }) - it('initializes request and cooldown buckets in a consistent order without duplicate keys', async () => { - const now = new Date() - dbChainMockFns.limit.mockResolvedValue([ - { key: 'cooldown', tokens: '0', lastRefillAt: now, blockedUntil: null }, - { key: 'requests', tokens: '10', lastRefillAt: now, blockedUntil: null }, - { key: 'tokens', tokens: '10', lastRefillAt: now, blockedUntil: null }, - ]) - - expect( - await new DbTokenBucket().consumeTokensAtomically( - [ - { key: 'tokens', cost: 2, config: CONFIG }, - { key: 'requests', cost: 1, config: CONFIG }, - ], - { cooldownKeys: ['cooldown', 'cooldown'], deadlineAt: now.getTime() + 1000 } - ) - ).toEqual({ allowed: true, retryAfterMs: 0 }) - - expect(dbChainMockFns.values).toHaveBeenCalledExactlyOnceWith([ - { key: 'cooldown', tokens: '0', lastRefillAt: now, updatedAt: now }, - { key: 'requests', tokens: '10', lastRefillAt: now, updatedAt: now }, - { key: 'tokens', tokens: '10', lastRefillAt: now, updatedAt: now }, - ]) - expect(dbChainMockFns.set).toHaveBeenCalledWith({ - tokens: '8', - lastRefillAt: now, - updatedAt: now, - }) - expect(dbChainMockFns.set).toHaveBeenCalledWith({ - tokens: '9', - lastRefillAt: now, - updatedAt: now, - }) - }) - it('accepts an empty reservation without attempting an empty insert', async () => { dbChainMockFns.limit.mockResolvedValue([]) diff --git a/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts b/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts index 6f818980916..097931a60aa 100644 --- a/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts +++ b/apps/sim/lib/core/rate-limiter/storage/db-token-bucket.ts @@ -1,6 +1,6 @@ import { db } from '@sim/db' import { rateLimitBucket } from '@sim/db/schema' -import { eq, inArray, sql } from 'drizzle-orm' +import { eq, type SQL, sql } from 'drizzle-orm' import type { AtomicAdmissionOptions, AtomicAdmissionResult, @@ -11,6 +11,191 @@ import type { TokenStatus, } from '@/lib/core/rate-limiter/storage/adapter' +/** Thrown inside the admission transaction to roll back a spend that finished past its deadline. */ +class AdmissionExpiredError extends Error {} + +interface AdmissionRow extends Record { + kind: 'reservation' | 'cooldown' + cost: number | null + available: number | null + last_refill_ms: number | null + refill_rate: number | null + refill_interval_ms: number | null + blocked_until_ms: number | null + allowed: boolean + spent: number + now_ms: number +} + +/** An array of scalar binds; the pool cannot bind a JavaScript array as one parameter. */ +function sqlArray(values: readonly (string | number)[], type: 'text' | 'float8'): SQL { + return sql`ARRAY[${sql.join( + values.map((value) => sql`${value}`), + sql`, ` + )}]::${sql.raw(type)}[]` +} + +/** + * Locks every key in sorted order, refills each reservation to the moment every lock is held + * with the same float8 arithmetic {@link DbTokenBucket.consumeTokens} uses in JavaScript, and + * spends all reservations only when no cooldown is active and every balance covers its cost. + * Returns one row per reservation and per cooldown key, each carrying the decision, how many + * buckets were spent, and the clock the decision used. + * + * That clock is the caller's `sentAt` advanced by the time spent acquiring the locks, in whole + * milliseconds. `lock_start` is joined into `locked`, so it is read before the first row can + * reach the lock; `clock` counts every `locked` row before reading `clock_timestamp()` again, so + * that read happens only once the last lock is held. An uncontended admission therefore decides + * at exactly `sentAt`, and measuring the wait rather than reading the database's absolute time + * keeps every stored instant on the application clock that wrote it. + */ +function buildAdmissionStatement( + reservations: readonly TokenBucketReservation[], + cooldownKeys: readonly string[], + keys: readonly string[], + sentAt: number +): SQL { + return sql` + WITH reservation AS ( + SELECT * FROM unnest( + ${sqlArray( + reservations.map((item) => item.key), + 'text' + )}, + ${sqlArray( + reservations.map((item) => item.cost), + 'float8' + )}, + ${sqlArray( + reservations.map((item) => item.config.maxTokens), + 'float8' + )}, + ${sqlArray( + reservations.map((item) => item.config.refillRate), + 'float8' + )}, + ${sqlArray( + reservations.map((item) => item.config.refillIntervalMs), + 'float8' + )} + ) AS r(key, cost, max_tokens, refill_rate, refill_interval_ms) + ), + lock_start AS ( + SELECT clock_timestamp() AS started_at + ), + locked AS ( + SELECT bucket.key, bucket.tokens, bucket.last_refill_at, bucket.blocked_until + FROM ${rateLimitBucket} AS bucket + CROSS JOIN lock_start + WHERE bucket.key = ANY(${sqlArray(keys, 'text')}) + ORDER BY bucket.key + FOR UPDATE OF bucket + ), + clock AS ( + SELECT ${sentAt}::float8 + + GREATEST(0, floor(extract(epoch FROM (held.locked_at - lock_start.started_at)) * 1000)) + AS now_ms + FROM (SELECT clock_timestamp() AS locked_at FROM (SELECT count(*) FROM locked) every_lock) held + CROSS JOIN lock_start + ), + balance AS ( + SELECT + r.key, + r.cost, + r.refill_rate, + r.refill_interval_ms, + LEAST(r.max_tokens, GREATEST(0, l.tokens::float8) + e.elapsed * r.refill_rate) AS available, + m.last_ms + e.elapsed * r.refill_interval_ms AS last_refill_ms, + l.last_refill_at + (e.elapsed * r.refill_interval_ms) * interval '1 millisecond' AS last_refill_at + FROM reservation r + JOIN locked l ON l.key = r.key + CROSS JOIN clock + CROSS JOIN LATERAL ( + SELECT (extract(epoch FROM l.last_refill_at) * 1000)::float8 AS last_ms + ) m + CROSS JOIN LATERAL ( + SELECT GREATEST(0, floor((clock.now_ms - m.last_ms) / r.refill_interval_ms)) AS elapsed + ) e + ), + cooldown AS ( + SELECT (extract(epoch FROM blocked_until) * 1000)::float8 AS blocked_until_ms + FROM locked + WHERE key = ANY(${sqlArray(cooldownKeys, 'text')}) + ), + decision AS ( + SELECT + NOT EXISTS (SELECT 1 FROM balance WHERE available < cost) + AND NOT EXISTS ( + SELECT 1 FROM cooldown, clock WHERE cooldown.blocked_until_ms > clock.now_ms + ) AS allowed + ), + spent AS ( + UPDATE ${rateLimitBucket} AS bucket + SET + tokens = (balance.available - balance.cost)::text::numeric, + last_refill_at = balance.last_refill_at, + updated_at = timestamp '1970-01-01' + clock.now_ms * interval '1 millisecond' + FROM balance, decision, clock + WHERE decision.allowed AND bucket.key = balance.key + RETURNING bucket.key + ), + outcome AS ( + SELECT decision.allowed, (SELECT count(*) FROM spent)::int AS spent, clock.now_ms + FROM decision, clock + ) + SELECT 'reservation' AS kind, b.cost, b.available, b.last_refill_ms, b.refill_rate, + b.refill_interval_ms, NULL::float8 AS blocked_until_ms, o.allowed, o.spent, o.now_ms + FROM balance b CROSS JOIN outcome o + UNION ALL + SELECT 'cooldown', NULL, NULL, NULL, NULL, NULL, c.blocked_until_ms, o.allowed, o.spent, + o.now_ms + FROM cooldown c CROSS JOIN outcome o + ` +} + +/** + * Derives the wait from the locked balances and cross-checks it against the statement's own + * decision, so the JavaScript reading and the SQL spend can never disagree silently: a mismatch, + * or a reservation whose bucket vanished, throws and rolls the spend back. + */ +function readAdmissionResult( + rows: readonly AdmissionRow[], + reservations: readonly TokenBucketReservation[], + now: number +): AtomicAdmissionResult { + let retryAfterMs = 0 + let reserved = 0 + for (const row of rows) { + if (row.kind === 'cooldown') { + if (row.blocked_until_ms !== null) { + retryAfterMs = Math.max(retryAfterMs, Number(row.blocked_until_ms) - now) + } + continue + } + reserved++ + const available = Number(row.available) + const cost = Number(row.cost) + if (available < cost) { + retryAfterMs = Math.max( + retryAfterMs, + Number(row.last_refill_ms) + + Math.ceil((cost - available) / Number(row.refill_rate)) * Number(row.refill_interval_ms) - + now + ) + } + } + if (reserved !== reservations.length) { + throw new Error('Rate limit bucket disappeared during admission') + } + const allowed = rows[0]?.allowed ?? true + const spent = rows[0]?.spent ?? 0 + const spentExpected = allowed ? new Set(reservations.map((item) => item.key)).size : 0 + if (allowed !== retryAfterMs <= 0 || spent !== spentExpected) { + throw new Error('Rate limit admission decision is inconsistent') + } + return allowed ? { allowed: true, retryAfterMs: 0 } : { allowed: false, retryAfterMs } +} + export class DbTokenBucket implements RateLimitStorageAdapter { async consumeTokens( key: string, @@ -78,23 +263,31 @@ export class DbTokenBucket implements RateLimitStorageAdapter { }) } - /** Locks all dimensions and provider gates before spending any capacity. */ + /** + * Locks all dimensions and provider gates before spending any capacity. + * + * The lock, refill, decision, and spend run as one statement, so the bucket rows stay locked + * for that statement and the commit rather than for a round trip per dimension: every waiter + * on a hot provider key queues behind that whole hold. Refill and cooldowns are judged at the + * moment the locks are held, as before, so a wait for a contended lock is refilled. + */ async consumeTokensAtomically( reservations: readonly TokenBucketReservation[], options: AtomicAdmissionOptions ): Promise { options.signal?.throwIfAborted() if (Date.now() >= options.deadlineAt) return { allowed: false, retryAfterMs: 0 } - return db.transaction(async (tx) => { - const remainingMs = Math.max(1, options.deadlineAt - Date.now()) - await tx.execute( - sql`SELECT set_config('statement_timeout', ${String(remainingMs)}, true), set_config('lock_timeout', ${String(remainingMs)}, true)` - ) - const keys = [ - ...new Set([...options.cooldownKeys, ...reservations.map((item) => item.key)]), - ].sort() - const createdAt = new Date() - if (keys.length > 0) { + const keys = [ + ...new Set([...options.cooldownKeys, ...reservations.map((item) => item.key)]), + ].sort() + if (keys.length === 0) return { allowed: true, retryAfterMs: 0 } + try { + return await db.transaction(async (tx) => { + const remainingMs = Math.max(1, options.deadlineAt - Date.now()) + await tx.execute( + sql`SELECT set_config('statement_timeout', ${String(remainingMs)}, true), set_config('lock_timeout', ${String(remainingMs)}, true)` + ) + const createdAt = new Date() await tx .insert(rateLimitBucket) .values( @@ -109,62 +302,20 @@ export class DbTokenBucket implements RateLimitStorageAdapter { }) ) .onConflictDoNothing() - } - const rows = await tx - .select() - .from(rateLimitBucket) - .where(inArray(rateLimitBucket.key, keys)) - .orderBy(rateLimitBucket.key) - .for('update') - .limit(keys.length) - options.signal?.throwIfAborted() - const now = Date.now() - if (now >= options.deadlineAt) return { allowed: false, retryAfterMs: 0 } - let retryAfterMs = 0 - for (const key of options.cooldownKeys) { - const row = rows.find((candidate) => candidate.key === key) - retryAfterMs = Math.max(retryAfterMs, (row?.blockedUntil?.getTime() ?? 0) - now) - } - const balances = reservations.map((reservation) => { - const row = rows.find((candidate) => candidate.key === reservation.key) - if (!row) throw new Error('Rate limit bucket disappeared during admission') - const { config, cost } = reservation - const elapsed = Math.max( - 0, - Math.floor((now - row.lastRefillAt.getTime()) / config.refillIntervalMs) - ) - const available = Math.min( - config.maxTokens, - Math.max(0, Number(row.tokens)) + elapsed * config.refillRate + options.signal?.throwIfAborted() + const sentAt = Date.now() + const rows = await tx.execute( + buildAdmissionStatement(reservations, options.cooldownKeys, keys, sentAt) ) - const lastRefillAt = new Date( - row.lastRefillAt.getTime() + elapsed * config.refillIntervalMs - ) - if (available < cost) { - retryAfterMs = Math.max( - retryAfterMs, - lastRefillAt.getTime() + - Math.ceil((cost - available) / config.refillRate) * config.refillIntervalMs - - now - ) - } - return { reservation, available, lastRefillAt } + const now = rows[0]?.now_ms ?? sentAt + options.signal?.throwIfAborted() + if (Math.max(now, Date.now()) >= options.deadlineAt) throw new AdmissionExpiredError() + return readAdmissionResult(rows, reservations, now) }) - if (retryAfterMs > 0) return { allowed: false, retryAfterMs } - options.signal?.throwIfAborted() - for (const { reservation, available, lastRefillAt } of balances) { - await tx - .update(rateLimitBucket) - .set({ - tokens: String(available - reservation.cost), - lastRefillAt, - updatedAt: new Date(now), - }) - .where(eq(rateLimitBucket.key, reservation.key)) - } - options.signal?.throwIfAborted() - return { allowed: true, retryAfterMs: 0 } - }) + } catch (error) { + if (error instanceof AdmissionExpiredError) return { allowed: false, retryAfterMs: 0 } + throw error + } } async getCooldownUntil(key: string): Promise { diff --git a/apps/sim/lib/knowledge/__integration__/provider-admission.integration.ts b/apps/sim/lib/knowledge/__integration__/provider-admission.integration.ts index b2fa4386a74..421650dea1c 100644 --- a/apps/sim/lib/knowledge/__integration__/provider-admission.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/provider-admission.integration.ts @@ -1,8 +1,15 @@ +import { db } from '@sim/db' +import { rateLimitBucket } from '@sim/db/schema' import { readTestRedisUrl } from '@sim/db/testing/test-infrastructure' +import { sleep } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' +import { eq } from 'drizzle-orm' import Redis from 'ioredis' import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' -import type { RateLimitStorageAdapter } from '@/lib/core/rate-limiter/storage/adapter' +import type { + AtomicAdmissionResult, + RateLimitStorageAdapter, +} from '@/lib/core/rate-limiter/storage/adapter' import { DbTokenBucket } from '@/lib/core/rate-limiter/storage/db-token-bucket' import { RedisTokenBucket } from '@/lib/core/rate-limiter/storage/redis-token-bucket' @@ -96,6 +103,37 @@ describe.each(['PostgreSQL', 'Redis'] as const)('%s provider admission', (backen }) }) + it('refills fractional reservations by whole intervals and returns the full wait', async () => { + const reservation = { + key: `${key}:requests`, + cost: 1, + config: { maxTokens: 1, refillRate: 0.25, refillIntervalMs: 1000 }, + } + const options = { cooldownKeys: [`${key}:cooldown`], deadlineAt: now.getTime() + 600_000 } + expect(await first.consumeTokensAtomically([reservation], options)).toEqual({ + allowed: true, + retryAfterMs: 0, + }) + vi.setSystemTime(now.getTime() + 1500) + expect(await second.consumeTokensAtomically([reservation], options)).toEqual({ + allowed: false, + retryAfterMs: 2500, + }) + expect(await first.getTokenStatus(reservation.key, reservation.config)).toMatchObject({ + tokensAvailable: 0.25, + lastRefillAt: new Date(now.getTime() + 1000), + }) + vi.setSystemTime(now.getTime() + 4000) + expect(await second.consumeTokensAtomically([reservation], options)).toEqual({ + allowed: true, + retryAfterMs: 0, + }) + expect(await first.getTokenStatus(reservation.key, reservation.config)).toMatchObject({ + tokensAvailable: 0, + lastRefillAt: new Date(now.getTime() + 4000), + }) + }) + it.each([ { state: 'new', spentRequests: 0, spentTokens: 0, admitted: 8, remainingTokens: 20 }, { state: 'existing', spentRequests: 2, spentTokens: 20, admitted: 6, remainingTokens: 20 }, @@ -218,3 +256,136 @@ describe.each(['PostgreSQL', 'Redis'] as const)('%s provider admission', (backen }) }) }) + +describe('PostgreSQL admission under a held bucket lock', () => { + const bucket = new DbTokenBucket() + let key: string + let now: number + + /** How long the competing transaction holds the row; long enough to cross a refill interval. */ + const LOCK_HOLD_MS = 1200 + + beforeEach(() => { + now = Date.now() + vi.useFakeTimers({ toFake: ['Date'] }) + vi.setSystemTime(now) + key = `provider-integration:${generateId()}` + }) + + afterEach(async () => { + vi.useRealTimers() + await bucket.resetBucket(key) + }) + + /** Locks the bucket row from another connection, runs `whileHeld`, then commits the holder. */ + async function withRowLockHeld(whileHeld: () => Promise): Promise { + let held: () => void = () => {} + let release: () => void = () => {} + const acquired = new Promise((resolve) => { + held = resolve + }) + const released = new Promise((resolve) => { + release = resolve + }) + const holder = db.transaction(async (tx) => { + await tx.select().from(rateLimitBucket).where(eq(rateLimitBucket.key, key)).for('update') + held() + await released + }) + await acquired + try { + await whileHeld() + } finally { + release() + await holder + } + } + + /** + * Starts an admission and resolves once it is queued behind the held row lock. The result is + * wrapped so awaiting this does not also await the blocked admission. + */ + async function queuedBehindLock( + admit: () => Promise + ): Promise<{ result: Promise }> { + const result = admit() + result.catch(() => {}) + await sleep(300) + return { result } + } + + async function storedTokens(): Promise { + const [row] = await db + .select({ tokens: rateLimitBucket.tokens }) + .from(rateLimitBucket) + .where(eq(rateLimitBucket.key, key)) + return Number(row?.tokens) + } + + it('refills the time spent waiting for the lock and measures the wait from then', async () => { + const reservation = { + key, + cost: 2, + config: { maxTokens: 2, refillRate: 1, refillIntervalMs: 1000 }, + } + const options = { cooldownKeys: [], deadlineAt: now + 60_000 } + expect(await bucket.consumeTokensAtomically([reservation], options)).toEqual({ + allowed: true, + retryAfterMs: 0, + }) + let denied: { result: Promise } | undefined + await withRowLockHeld(async () => { + denied = await queuedBehindLock(() => bucket.consumeTokensAtomically([reservation], options)) + await sleep(LOCK_HOLD_MS - 300) + }) + const result = await denied?.result + expect(result?.allowed).toBe(false) + expect(result?.retryAfterMs).toBeGreaterThan(0) + expect(result?.retryAfterMs).toBeLessThan(1000) + + let admitted: { result: Promise } | undefined + await withRowLockHeld(async () => { + admitted = await queuedBehindLock(() => + bucket.consumeTokensAtomically([{ ...reservation, cost: 1 }], options) + ) + await sleep(LOCK_HOLD_MS - 300) + }) + expect(await admitted?.result).toEqual({ allowed: true, retryAfterMs: 0 }) + }) + + it.each(['deadline', 'abort'] as const)( + 'rolls back a spend whose %s arrives while the statement runs', + async (interruption) => { + const reservation = { + key, + cost: 5, + config: { maxTokens: 10, refillRate: 1, refillIntervalMs: 600_000 }, + } + expect( + await bucket.consumeTokensAtomically([{ ...reservation, cost: 0 }], { + cooldownKeys: [], + deadlineAt: now + 60_000, + }) + ).toEqual({ allowed: true, retryAfterMs: 0 }) + const controller = new AbortController() + let pending: { result: Promise } | undefined + await withRowLockHeld(async () => { + pending = await queuedBehindLock(() => + bucket.consumeTokensAtomically([reservation], { + cooldownKeys: [], + deadlineAt: now + 5000, + signal: controller.signal, + }) + ) + if (interruption === 'deadline') vi.setSystemTime(now + 6000) + else controller.abort(new Error('cancelled')) + }) + if (interruption === 'deadline') { + expect(await pending?.result).toEqual({ allowed: false, retryAfterMs: 0 }) + } else { + await expect(pending?.result).rejects.toThrow('cancelled') + } + expect(await storedTokens()).toBe(10) + } + ) +}) diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts index a9e18722987..c356acad943 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts @@ -1591,6 +1591,61 @@ describe('completeSuccessfulSync', () => { expect(successWrite).toMatchObject({ consecutiveFailures: 0 }) }) + it.each([ + { + scenario: 'after a finished pass', + lastSyncAt: new Date(), + lastSyncDocCount: 7, + written: undefined, + }, + { scenario: 'with no count yet', lastSyncAt: null, lastSyncDocCount: null, written: 4 }, + { + scenario: 'on a later slice of the first pass', + lastSyncAt: null, + lastSyncDocCount: 3, + written: 4, + }, + ])( + 'counts an unfinished slice only until a pass has finished: $scenario', + async ({ lastSyncAt, lastSyncDocCount, written }) => { + const { completeSuccessfulSync } = await import('@/lib/knowledge/connectors/sync-engine') + queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-1' }]) + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1', lastSyncDocCount: 7 }]) + queueTableRows(schemaMock.document, [{ count: 4 }]) + dbChainMockFns.returning + .mockResolvedValueOnce([{ id: 'log-1' }]) + .mockResolvedValueOnce([{ id: 'c-1' }]) + + await expect( + completeSuccessfulSync( + 'c-1', + 'kb-1', + 'log-1', + 60, + { ...RESULT, docsFailed: 0 }, + null, + { + complete: false, + checkpoint: { + unsafe: false, + startedAt: '2026-09-04T00:00:00Z', + listedCount: 25, + resumeAt: null, + }, + }, + null, + false, + { lastSyncAt, lastSyncDocCount } + ) + ).resolves.toBe(true) + + const successWrite = dbChainMockFns.set.mock.calls + .map(([value]) => value as Record) + .find((value) => value.status === 'active') + expect(successWrite?.lastSyncDocCount).toBe(written) + } + ) + it('publishes neither terminal state when lock ownership is gone', async () => { const { completeSuccessfulSync } = await import('@/lib/knowledge/connectors/sync-engine') diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.ts b/apps/sim/lib/knowledge/connectors/sync-engine.ts index 10f467de4a1..5102b99df22 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.ts @@ -399,7 +399,9 @@ export async function completeSuccessfulSync( contentPass?: ContentPassOutcome, directoryNotice: string | null = null, /** A pending ACL rewrite this run began but did not finish: the flag stays and the next run resumes it. */ - accessRewriteUnfinished = false + accessRewriteUnfinished = false, + /** The connector as this run locked it: whether a pass has finished, and the count it holds. */ + lockedConnector: { lastSyncAt: Date | null; lastSyncDocCount: number | null } | null = null ): Promise { const processingDispatchFailed = result.processingDispatch.failed > 0 const contentNotice = @@ -431,16 +433,29 @@ export async function completeSuccessfulSync( * owned count the next pass's mass-deletion guard starts from, so a tally of this run's * adds and deletes would drift; when it cannot be read the previous count stands, which the * guard already treats as a floor. + * + * Once a pass has finished, a slice that leaves the next pass unfinished keeps the count that + * finished pass recorded, the same pass `lastSyncAt` still names, so a large connector is not + * counted again on every slice of a pass whose count is replaced when it finishes. Unfinished + * slices never advance `lastSyncAt`, so until a pass finishes every slice is counted and a + * first pass shows its progress. */ - const actualDocCount = await countLiveConnectorDocuments(connectorId).catch((error: unknown) => { - const diagnostic = getConnectorFailureDiagnostic(error) - if (diagnostic?.category !== 'database') throw error - logger.warn('Could not count connector documents; keeping the previous count', { - connectorId, - diagnostic, - }) - return null - }) + const keepRecordedCount = + contentPass !== undefined && + !contentPass.complete && + lockedConnector?.lastSyncAt != null && + lockedConnector.lastSyncDocCount !== null + const actualDocCount = keepRecordedCount + ? null + : await countLiveConnectorDocuments(connectorId).catch((error: unknown) => { + const diagnostic = getConnectorFailureDiagnostic(error) + if (diagnostic?.category !== 'database') throw error + logger.warn('Could not count connector documents; keeping the previous count', { + connectorId, + diagnostic, + }) + return null + }) try { return await db.transaction(async (tx) => { const [lockedKnowledgeBase] = await tx @@ -1417,7 +1432,8 @@ export async function executeSync( reconciliationHoldNotice, contentPass, directoryNotice, - accessRewriteUnfinished + accessRewriteUnfinished, + connector ) if (!completionLanded) {