fix(search): prune idle dispatch owners safely

Author: waleedlatif1

apps/sim/lib/knowledge/__integration__/document-dispatch.integration.ts

@@ -202,6 +202,58 @@ describe('durable fair document admission', () => {
     }
   })
 
+  it('prunes idle owners in bounded batches while retaining delayed and outstanding work', async () => {
+    const idleOwners = Array.from({ length: 1001 }, () => `organization:${generateId()}`)
+    fixtureOwners.push(...idleOwners)
+    await db
+      .insert(knowledgeDocumentDispatchOwner)
+      .values(idleOwners.map((ownerKey) => ({ ownerKey })))
+    const delayed = await createSource()
+    const active = await createSource()
+    const [delayedPayload] = await createDocuments(delayed, 1)
+    const [activePayload] = await createDocuments(active, 1)
+    await enqueueDocumentProcessingDispatch([delayedPayload, activePayload])
+    await db
+      .update(knowledgeDocumentDispatch)
+      .set({ availableAt: new Date(Date.now() + 60_000) })
+      .where(eq(knowledgeDocumentDispatch.documentId, delayedPayload.documentId))
+
+    await claimDocumentProcessingDispatches()
+    const retained = await db.select().from(knowledgeDocumentDispatchOwner)
+    expect(retained.filter((owner) => idleOwners.includes(owner.ownerKey))).toHaveLength(1)
+    expect(retained.map((owner) => owner.ownerKey)).toEqual(
+      expect.arrayContaining([delayed.ownerKey, active.ownerKey])
+    )
+    await claimDocumentProcessingDispatches()
+    expect(await db.select().from(knowledgeDocumentDispatchOwner)).toHaveLength(2)
+
+    await completeDocumentProcessingDispatch(activePayload)
+    await claimDocumentProcessingDispatches()
+    expect(
+      (await db.select().from(knowledgeDocumentDispatchOwner)).map((owner) => owner.ownerKey)
+    ).toEqual([delayed.ownerKey])
+  })
+
+  it('keeps every persisted intent discoverable when idle pruning races with enqueues', async () => {
+    const sources = await Promise.all(Array.from({ length: 12 }, () => createSource()))
+    await db
+      .insert(knowledgeDocumentDispatchOwner)
+      .values(sources.map(({ ownerKey }) => ({ ownerKey })))
+    const payloads = await Promise.all(sources.map((source) => createDocuments(source, 1)))
+    const results = await Promise.allSettled([
+      ...payloads.map((batch) => enqueueDocumentProcessingDispatch(batch)),
+      ...Array.from({ length: 6 }, () => claimDocumentProcessingDispatches()),
+    ])
+    expect(results.every((result) => result.status === 'fulfilled')).toBe(true)
+    const owners = new Set(
+      (await db.select().from(knowledgeDocumentDispatchOwner)).map((owner) => owner.ownerKey)
+    )
+    const intents = await db.select().from(knowledgeDocumentDispatch)
+    expect(intents).toHaveLength(sources.length)
+    expect(intents.every((intent) => owners.has(intent.ownerKey))).toBe(true)
+    expect(await claimDocumentProcessingDispatches()).toHaveLength(sources.length)
+  })
+
   it('lets one owner use available workers without queuing its entire corpus ahead of a newcomer', async () => {
     const large = await createSource()
     await enqueueDocumentProcessingDispatch(await createDocuments(large, 10_000))

apps/sim/lib/knowledge/application/operations.test.ts

@@ -57,6 +57,8 @@ describe('knowledge operation registry', () => {
       'knowledge.connectors.update',
       'knowledge.connectors.access.update',
       'knowledge.search.sources.list',
+      'knowledge.search.sources.overview',
+      'knowledge.search.sources.progress',
       'knowledge.search.integrations.list',
       'knowledge.search.integrations.approve',
       'knowledge.connectors.members.list',

apps/sim/lib/knowledge/documents/processing-dispatch-queue.ts

@@ -24,6 +24,7 @@ const EXECUTION_CONCURRENCY = envNumber(env.KB_CONFIG_CONCURRENCY_LIMIT, 20, { m
 export const DOCUMENT_DISPATCH_OWNER_OUTSTANDING = EXECUTION_CONCURRENCY
 export const DOCUMENT_DISPATCH_MAX_OUTSTANDING = EXECUTION_CONCURRENCY * 2
 const ENQUEUE_BATCH_SIZE = 1000
+const OWNER_PRUNE_BATCH_SIZE = 1000
 const RECONCILE_AFTER_MS = 15 * 60 * 1000
 const DISPATCH_LOCK = 'knowledge-document-dispatch'
 
@@ -119,11 +120,17 @@ async function persistDispatchIntents(
       })
     }
     if (values.length === 0) return []
-    await tx.insert(knowledgeDocumentDispatch).values(values).onConflictDoNothing()
+    /** Lock owners before inserting intents so idle-owner pruning cannot orphan accepted work. */
     await tx
       .insert(knowledgeDocumentDispatchOwner)
-      .values([...new Set(values.map((value) => value.ownerKey))].map((ownerKey) => ({ ownerKey })))
-      .onConflictDoNothing()
+      .values(
+        [...new Set(values.map((value) => value.ownerKey))].sort().map((ownerKey) => ({ ownerKey }))
+      )
+      .onConflictDoUpdate({
+        target: knowledgeDocumentDispatchOwner.ownerKey,
+        set: { ownerKey: sql`excluded.owner_key` },
+      })
+    await tx.insert(knowledgeDocumentDispatch).values(values).onConflictDoNothing()
     return tx
       .select()
       .from(knowledgeDocumentDispatch)
@@ -258,8 +265,45 @@ async function reconcileCompletedDispatches(): Promise<void> {
   })
 }
 
+/** Release idle-owner locks before claims can wait on another owner's concurrent enqueue. */
+async function pruneIdleDispatchOwners(): Promise<void> {
+  await db.transaction(async (tx) => {
+    const idleOwner = sql`
+      NOT EXISTS (SELECT 1 FROM knowledge_document_dispatch AS intent
+        WHERE intent.owner_key = ${knowledgeDocumentDispatchOwner.ownerKey}
+          AND intent.dispatched_at IS NULL)
+      AND NOT EXISTS (SELECT 1 FROM knowledge_document_dispatch AS intent
+        WHERE intent.owner_key = ${knowledgeDocumentDispatchOwner.ownerKey}
+          AND intent.dispatched_at IS NOT NULL)
+    `
+    const idleOwners = await tx
+      .select({ ownerKey: knowledgeDocumentDispatchOwner.ownerKey })
+      .from(knowledgeDocumentDispatchOwner)
+      .where(idleOwner)
+      .orderBy(
+        asc(knowledgeDocumentDispatchOwner.enqueuedAt),
+        asc(knowledgeDocumentDispatchOwner.ownerKey)
+      )
+      .limit(OWNER_PRUNE_BATCH_SIZE)
+      .for('update', { skipLocked: true })
+    if (idleOwners.length > 0) {
+      /** Recheck on a fresh statement snapshot after locking, in case an enqueue just committed. */
+      await tx.delete(knowledgeDocumentDispatchOwner).where(
+        and(
+          inArray(
+            knowledgeDocumentDispatchOwner.ownerKey,
+            idleOwners.map(({ ownerKey }) => ownerKey)
+          ),
+          idleOwner
+        )
+      )
+    }
+  })
+}
+
 /** Reserves both budgets under one short lock; provider calls happen after the transaction. */
 export async function claimDocumentProcessingDispatches(): Promise<DispatchRow[]> {
+  await pruneIdleDispatchOwners()
   return db.transaction(async (tx) => {
     const [lock] = await tx.execute<{ acquired: boolean }>(
       sql`SELECT pg_try_advisory_xact_lock(hashtextextended(${DISPATCH_LOCK}, 0)) AS acquired`