fix(knowledge): estimate a filter's size under a short deadline of its own

Author: waleedlatif1

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

@@ -2199,6 +2199,38 @@ describe('filters on a resolved scope', () => {
     expect(statements().filter((query) => isExactRanking(query.sql))).toHaveLength(1)
   })
 
+  it("estimates a filter under the leg's deadline and walks when the estimate runs out of time", async () => {
+    dbChainMockFns.execute.mockImplementation(async (query) => {
+      const statement = render(query).sql
+      if (statement.includes('EXPLAIN') && JSON.stringify(render(query)).includes('"type":"gte"'))
+        throw Object.assign(new Error('canceling statement due to statement timeout'), {
+          code: '57014',
+        })
+      if (statement.includes('EXPLAIN'))
+        return [{ 'QUERY PLAN': [{ Plan: { 'Plan Rows': 1_000_000 } }] }]
+      if (statement.includes(') reached')) return [{ n: 250_000 }]
+      if (isWalk(statement)) return traversedRows
+      if (statement.includes('WITH scored_search_candidates')) return rerankRows
+      return []
+    })
+    const result = await retrieveKnowledgeSearch({
+      ...params,
+      accessProvider: provider,
+      searchMode: 'vector',
+      query: 'release',
+      filters: { modifiedAfter: '2026-09-13T00:00:00.000Z' },
+    })
+    /** The estimate ran inside a deadline statement; its own timeout chose the walk and cost the leg nothing. */
+    const estimateAt = statements().findIndex((query) =>
+      JSON.stringify(query).includes('"type":"gte"')
+    )
+    expect(estimateAt).toBeGreaterThan(0)
+    expect(statements()[estimateAt - 1].sql).toContain('statement_timeout')
+    expect(statements().filter((query) => query.sql.includes('AS saturated'))).toHaveLength(0)
+    expect(statements().filter((query) => isWalk(query.sql))).toHaveLength(1)
+    expect(result.retrieval.status).toBe('complete')
+  })
+
   it('enumerates a date filter only while the planner estimates its documents few', async () => {
     let estimated = 50_000
     dbChainMockFns.execute.mockImplementation(async (query) => {

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

@@ -1160,9 +1160,11 @@ const indexDocumentCounts = new LRUCache<string, number>({
 async function estimateFilteredDocuments(
   knowledgeBaseIds: string[],
   filters: WorkspaceSearchFilters,
-  plan: SearchAccessPlan
+  plan: SearchAccessPlan,
+  budget: SearchBudget | undefined
 ): Promise<number> {
-  const [row] = await db.execute<{ 'QUERY PLAN': Array<{ Plan: { 'Plan Rows': number } }> }>(sql`
+  const [row] = await runSearchQuery(budget, 'permitted_documents', (executor) =>
+    executor.execute<{ 'QUERY PLAN': Array<{ Plan: { 'Plan Rows': number } }> }>(sql`
     EXPLAIN (FORMAT JSON) SELECT 1 FROM ${document}
     WHERE ${and(
       inArray(document.knowledgeBaseId, knowledgeBaseIds),
@@ -1172,6 +1174,7 @@ async function estimateFilteredDocuments(
         : undefined,
       filters.source ? planSourceCondition(plan) : undefined
     )}`)
+  )
   return Number(row?.['QUERY PLAN']?.[0]?.Plan?.['Plan Rows'] ?? 0)
 }
 
@@ -2421,11 +2424,21 @@ export async function retrieveKnowledgeSearch(
    * source confined on the row, the date tested through the document — since a set that large
    * holds most of the query's neighbours anyway. The planner's estimate decides which.
    */
+  /** Planning only, so a short cap of its own: running past it answers as the wide window it may be. */
+  const estimateBudget = budgets.vector.capped(VECTOR_PROBE_BUDGET_MS)
   const enumerateFiltered =
     accessPlan && (params.filters?.modifiedAfter || params.filters?.source)
-      ? (await measureSearchStage('permitted_documents', () =>
-          estimateFilteredDocuments(knowledgeBaseIds, params.filters!, accessPlan)
-        )) <= VECTOR_PROBE_DOCUMENT_LIMIT
+      ? await estimateFilteredDocuments(
+          knowledgeBaseIds,
+          params.filters,
+          accessPlan,
+          estimateBudget
+        )
+          .then((estimate) => estimate <= VECTOR_PROBE_DOCUMENT_LIMIT)
+          .catch((error) => {
+            if (!estimateBudget.isTimeout(error)) throw error
+            return false
+          })
       : false
   const permitted =
     access.kind === 'user' && params.accessProvider && !params.filters?.documentIds?.length