fix(knowledge): count a connector's documents before the completion locks and index the live set

Author: waleedlatif1

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

@@ -2252,6 +2252,61 @@ describe('completeSuccessfulSync', () => {
     }
   )
 
+  it('counts the documents before taking the completion locks', async () => {
+    const { completeSuccessfulSync } = await import('@/lib/knowledge/connectors/sync-engine')
+    queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-1' }])
+    queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }])
+    queueTableRows(schemaMock.document, [{ count: 4 }])
+    dbChainMockFns.returning
+      .mockResolvedValueOnce([])
+      .mockResolvedValueOnce([{ id: 'log-1' }])
+      .mockResolvedValueOnce([{ id: 'c-1' }])
+
+    await expect(completeSuccessfulSync('c-1', 'kb-1', 'log-1', 60, RESULT, null)).resolves.toBe(
+      true
+    )
+
+    const countOrder = dbChainMockFns.from.mock.invocationCallOrder[0]
+    const transactionOrder = dbChainMockFns.transaction.mock.invocationCallOrder[0]
+    expect(dbChainMockFns.from.mock.calls[0][0]).toBe(schemaMock.document)
+    expect(countOrder).toBeLessThan(transactionOrder)
+    expect(dbChainMockFns.set).toHaveBeenCalledWith(
+      expect.objectContaining({ status: 'active', lastSyncDocCount: 4 })
+    )
+  })
+
+  it('keeps the previous document count when the count cannot be read', async () => {
+    const { completeSuccessfulSync } = await import('@/lib/knowledge/connectors/sync-engine')
+    queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-1' }])
+    queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }])
+    dbChainMockFns.where.mockImplementationOnce(() =>
+      Promise.reject(
+        new DrizzleQueryError(
+          'select private SQL',
+          [],
+          Object.assign(new Error('canceling statement due to statement timeout'), {
+            code: '57014',
+          })
+        )
+      )
+    )
+    dbChainMockFns.returning
+      .mockResolvedValueOnce([])
+      .mockResolvedValueOnce([{ id: 'log-1' }])
+      .mockResolvedValueOnce([{ id: 'c-1' }])
+
+    await expect(completeSuccessfulSync('c-1', 'kb-1', 'log-1', 60, RESULT, null)).resolves.toBe(
+      true
+    )
+
+    const successWrite = dbChainMockFns.set.mock.calls
+      .map(([value]) => value as Record<string, unknown>)
+      .find((value) => value.status === 'active')
+    expect(successWrite).toBeDefined()
+    expect(successWrite).not.toHaveProperty('lastSyncDocCount')
+    expect(successWrite).toMatchObject({ consecutiveFailures: 0 })
+  })
+
   it('commits the completed log and connector state in one guarded transaction', async () => {
     const { completeSuccessfulSync } = await import('@/lib/knowledge/connectors/sync-engine')
 

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

@@ -393,6 +393,18 @@ export async function completeSuccessfulSync(
     : null
   const completionNotice =
     [directoryNotice, listingNotice, contentNotice].filter(Boolean).join('\n') || null
+  /**
+   * A display snapshot, taken before the completion transaction so a slow count neither holds
+   * the connector lock nor turns a sync whose documents already landed into a failure. When it
+   * cannot be read the previous count stands.
+   */
+  const actualDocCount = await countLiveConnectorDocuments(connectorId).catch((error: unknown) => {
+    logger.warn('Could not count connector documents; keeping the previous count', {
+      connectorId,
+      diagnostic: getConnectorFailureDiagnostic(error),
+    })
+    return null
+  })
   try {
     const completed = await db.transaction(async (tx) => {
       const [lockedKnowledgeBase] = await tx
@@ -424,18 +436,6 @@ export async function completeSuccessfulSync(
         })
       }
 
-      const [{ count: actualDocCount }] = await tx
-        .select({ count: sql<number>`count(*)::int` })
-        .from(document)
-        .where(
-          and(
-            eq(document.connectorId, connectorId),
-            eq(document.userExcluded, false),
-            isNull(document.archivedAt),
-            isNull(document.deletedAt)
-          )
-        )
-
       const now = new Date()
       const [closedLog] = await tx
         .update(knowledgeConnectorSyncLog)
@@ -703,9 +703,25 @@ export function buildSyncCapacityUpdate(
  * `consecutiveFailures` still resets: a held pass is a healthy sync that declined
  * to delete, not a failure, and marking it broken would stop it syncing at all.
  */
+/** Live documents the connector owns: what the connector list shows as its document count. */
+async function countLiveConnectorDocuments(connectorId: string): Promise<number> {
+  const [row] = await db
+    .select({ count: sql<number>`count(*)::int` })
+    .from(document)
+    .where(
+      and(
+        eq(document.connectorId, connectorId),
+        eq(document.userExcluded, false),
+        isNull(document.archivedAt),
+        isNull(document.deletedAt)
+      )
+    )
+  return row?.count ?? 0
+}
+
 export function buildSyncSuccessUpdate(
   now: Date,
-  actualDocCount: number,
+  actualDocCount: number | null,
   nextSyncAt: Date | null,
   holdNotice: string | null,
   advanceLastSyncAt = true
@@ -714,7 +730,7 @@ export function buildSyncSuccessUpdate(
     status: 'active' as const,
     ...(advanceLastSyncAt ? { lastSyncAt: now } : {}),
     lastSyncError: holdNotice,
-    lastSyncDocCount: actualDocCount,
+    ...(actualDocCount === null ? {} : { lastSyncDocCount: actualDocCount }),
     nextSyncAt,
     consecutiveFailures: 0,
     // Releases the lock so a stale token can never match a later run, and closes

packages/db/migrations/0375_document_connector_live_index.sql

@@ -0,0 +1,9 @@
+-- Every sync completion counts the live documents a connector owns. The reconciliation index keeps
+-- tombstones on purpose, so without a live-only index that count walks the connector's heap tuples
+-- and a large connector cannot finish inside the statement timeout.
+COMMIT;--> statement-breakpoint
+SET lock_timeout = 0;--> statement-breakpoint
+-- migration-safe: replay replaces only this new index to recover an interrupted concurrent build; existing indexes remain available.
+DROP INDEX CONCURRENTLY IF EXISTS "doc_connector_live_idx";--> statement-breakpoint
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "doc_connector_live_idx" ON "document" USING btree ("connector_id") WHERE "user_excluded" = false AND "archived_at" IS NULL AND "deleted_at" IS NULL;--> statement-breakpoint
+SET lock_timeout = '5s';