fix(knowledge): give the reach count the probe's share of the deadline, not the leg's (#8111)

Author: waleedlatif1

apps/sim/lib/knowledge/search/queries.test.ts

@@ -2051,6 +2051,20 @@ describe('permitted-document planner', () => {
       expect(reachCounts()).toHaveLength(2)
     })
 
+    it('reports a leg whose own deadline passed during the count as short, not failed', async () => {
+      const budget = new SearchBudget('vector', performance.now() - 1)
+      await expect(
+        resolveReach(['org-index'], scope('spent-leg'), budget, {
+          connectors: { workspace: [], admin: [], members: [], liveProofRequired: [] },
+          observers: { confirmed: [], observed: [] },
+          memberSources: [],
+          connectorTypes: new Map(),
+          uploads: true,
+        })
+      ).resolves.toEqual({ kind: 'unbounded', broad: true })
+      expect(budget.timedOut).toBe(true)
+    })
+
     it('counts a resolved reach against a small index instead of assuming it broad', async () => {
       /** A bound inside the probe limit proves nothing without a saturated probe. */
       dbChainMockFns.execute.mockImplementation(async (query) => {
@@ -2075,9 +2089,10 @@ describe('permitted-document planner', () => {
       )
       expect(reach).toEqual({ kind: 'unbounded', broad: false })
       expect(reachCounts()).toHaveLength(1)
-      /** The count is the search's own read: it runs inside the leg's deadline statement. */
+      /** The count is the search's own read, under the probe's share of the deadline, not the leg's. */
       const countAt = statements().findIndex((query) => query.sql.includes(') reached'))
       expect(statements()[countAt - 1].sql).toContain('statement_timeout')
+      expect(Number(statements()[countAt - 1].params[0])).toBeLessThanOrEqual(600)
     })
 
     it('does not remember a reach whose count ran out of time', async () => {
@@ -2102,6 +2117,8 @@ describe('permitted-document planner', () => {
       })
       expect(plan).toEqual({ kind: 'unbounded', broad: true })
       expect(reachCounts()).toHaveLength(1)
+      /** Only the count's share of the deadline was spent; the leg is not the one that timed out. */
+      expect(budget.timedOut).toBe(false)
       /** The next search counts again rather than trusting an answer that never came. */
       await resolveReach(
         ['org-index'],

apps/sim/lib/knowledge/search/queries.ts

@@ -1253,35 +1253,48 @@ async function estimateFilteredDocuments(
  * documents. Counted once against that bound and remembered, so the first search after the
  * window pays for it and the rest do not. A caller whose probe already saturated is known to
  * reach past the probe's limit, so a bound inside that limit is met without counting.
+ *
+ * The count reads as many index entries as the caller reaches, so on a large index it can cost
+ * more than the leg it serves; it gets the probe's share of the deadline, never the whole leg's.
+ * A count that runs out of that share answers `null`: the leg keeps its time and its deadline
+ * intact, and the caller decides this search alone without remembering anything.
  */
 async function reachIsBroad(
   knowledgeBaseIds: string[],
   access: KnowledgeAccessScope,
   budget: SearchBudget | undefined,
   plan: SearchAccessPlan | undefined,
   saturated: boolean
-): Promise<boolean> {
+): Promise<boolean | null> {
   if (access.kind !== 'user') return true
-  const total =
-    (await indexDocumentCounts.fetch([...knowledgeBaseIds].sort().join(','), {
-      context: budget,
-    })) ?? 0
-  const bound = Math.ceil(total * BROAD_REACH_SHARE)
-  if (saturated && bound <= VECTOR_PROBE_DOCUMENT_LIMIT) return true
-  const [row] = await runSearchQuery(budget, 'permitted_documents', (executor) =>
-    executor.execute<{ n: number }>(sql`
-      SELECT count(*) AS n FROM (
-        SELECT 1 FROM ${document}
-        WHERE ${and(
-          isNull(document.deletedAt),
-          knowledgeAclOverlapCondition(access),
-          inArray(document.knowledgeBaseId, knowledgeBaseIds),
-          planSourceCondition(plan)
-        )}
-        LIMIT ${bound}
-      ) reached`)
-  )
-  return Number(row?.n ?? 0) >= bound
+  const countBudget = budget?.capped(VECTOR_PROBE_BUDGET_MS)
+  try {
+    const total =
+      (await indexDocumentCounts.fetch([...knowledgeBaseIds].sort().join(','), {
+        context: countBudget,
+      })) ?? 0
+    const bound = Math.ceil(total * BROAD_REACH_SHARE)
+    if (saturated && bound <= VECTOR_PROBE_DOCUMENT_LIMIT) return true
+    const [row] = await runSearchQuery(countBudget, 'permitted_documents', (executor) =>
+      executor.execute<{ n: number }>(sql`
+        SELECT count(*) AS n FROM (
+          SELECT 1 FROM ${document}
+          WHERE ${and(
+            isNull(document.deletedAt),
+            knowledgeAclOverlapCondition(access),
+            inArray(document.knowledgeBaseId, knowledgeBaseIds),
+            planSourceCondition(plan)
+          )}
+          LIMIT ${bound}
+        ) reached`)
+    )
+    return Number(row?.n ?? 0) >= bound
+  } catch (error) {
+    if (!budget || !countBudget?.isTimeout(error)) throw error
+    /** Only the count's share was spent; the leg's own deadline still governs. */
+    budget.remaining()
+    return null
+  }
 }
 
 /**
@@ -1333,11 +1346,13 @@ export async function resolveReach(
   if (remembered) return { kind: 'unbounded', broad: remembered.broad }
   try {
     const broad = await reachIsBroad(knowledgeBaseIds, access, budget, plan, false)
+    /** A count that ran out of time decides this search only; the next one counts again. */
+    if (broad === null) return { kind: 'unbounded', broad: true }
     if (key) saturatedReach.set(key, { broad })
     return { kind: 'unbounded', broad }
   } catch (error) {
+    /** The leg's own deadline passed during the count: the leg is short, the search is not failed. */
     if (!budget?.isTimeout(error)) throw error
-    /** A count that ran out of time decides this search only; the next one counts again. */
     return { kind: 'unbounded', broad: true }
   }
 }
@@ -1393,17 +1408,21 @@ export async function resolvePermittedDocuments(params: {
     }
     if (probe.kind === 'saturated') {
       try {
-        broad = await reachIsBroad(
+        const counted = await reachIsBroad(
           params.knowledgeBaseIds,
           params.access,
           params.budget,
           params.accessPlan,
           true
         )
-        if (key) saturatedReach.set(key, { broad })
+        /** A count that ran out of time decides this search only; the next one counts again. */
+        if (counted !== null) {
+          broad = counted
+          if (key) saturatedReach.set(key, { broad })
+        }
       } catch (error) {
+        /** The leg's own deadline passed during the count: the leg is short, the search is not failed. */
         if (!params.budget?.isTimeout(error)) throw error
-        /** A count that ran out of time decides this search only; the next one counts again. */
       }
     }
   }
@@ -2658,9 +2677,7 @@ export async function retrieveKnowledgeSearch(
            * readable documents enumerated ahead of ranking: its reach alone chooses between one
            * walk over the whole graph and a search of each source.
            */
-          await measureSearchStage('permitted_documents', () =>
-            resolveReach(knowledgeBaseIds, access, budgets.vector, accessPlan)
-          )
+          await resolveReach(knowledgeBaseIds, access, budgets.vector, accessPlan)
         : await resolvePermittedDocuments({
             knowledgeBaseIds,
             access,