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

Counting a caller's reach reads as many index entries as they reach, so on a large index it took the whole vector budget and left the leg no time to walk; the search then reported the leg timed out with the count's stage measured twice. The count now runs under the probe's share of the deadline, a count that runs out answers nothing for this search without marking the leg, and the reach memo holds for an hour: reach moves slowly and a stale answer costs speed, never access.

Author: waleedlatif1

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

@@ -2075,9 +2075,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 +2103,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

@@ -1180,9 +1180,11 @@ export const BROAD_REACH_SHARE = 0.25
 /**
  * How long a caller's saturated reach is remembered. Reach counts the documents a caller's tokens
  * touch in the bases, which moves slowly, and an unbounded set only means the legs search the
- * index with the full access predicate, so a stale answer costs speed, never access.
+ * index with the full access predicate, so a stale answer costs speed, never access. Counting it
+ * is the one read of a search that scales with the caller's reach rather than the query, so it is
+ * remembered for long.
  */
-const SATURATED_REACH_TTL_MS = 5 * 60 * 1000
+const SATURATED_REACH_TTL_MS = 60 * 60 * 1000
 
 /** A saturated reach, and whether it is broad enough to walk the whole graph for. */
 const saturatedReach = new LRUCache<string, { broad: boolean }>({
@@ -1253,35 +1255,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
+  }
 }
 
 /**
@@ -1331,15 +1346,11 @@ export async function resolveReach(
   const key = reachKey(knowledgeBaseIds, access, plan)
   const remembered = key ? saturatedReach.get(key) : undefined
   if (remembered) return { kind: 'unbounded', broad: remembered.broad }
-  try {
-    const broad = await reachIsBroad(knowledgeBaseIds, access, budget, plan, false)
-    if (key) saturatedReach.set(key, { broad })
-    return { kind: 'unbounded', broad }
-  } catch (error) {
-    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 }
-  }
+  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 }
 }
 
 /**
@@ -1392,18 +1403,17 @@ export async function resolvePermittedDocuments(params: {
       probe = { kind: 'timed_out' }
     }
     if (probe.kind === 'saturated') {
-      try {
-        broad = await reachIsBroad(
-          params.knowledgeBaseIds,
-          params.access,
-          params.budget,
-          params.accessPlan,
-          true
-        )
+      const counted = await reachIsBroad(
+        params.knowledgeBaseIds,
+        params.access,
+        params.budget,
+        params.accessPlan,
+        true
+      )
+      /** 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) {
-        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 +2668,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,