fix(knowledge): keep vector candidate selection bounded

apps/sim/lib/knowledge/__integration__/search-latency.integration.ts

@@ -208,6 +208,15 @@ function explainNodes(node: ExplainNode): ExplainNode[] {
 }
 
 /** Broad ranking must stop the ordered ANN scan instead of sorting every accessible chunk. */
+/**
+ * The bounded ANN traversal, identified by the visibility lateral it alone carries. The probe
+ * aliases its own lateral `scoped_chunk`, so this cannot match it, and matching on the rendered
+ * clause casing would silently stop these assertions from running at all.
+ */
+function isVectorCandidateQuery(statement: string) {
+  return statement.toLowerCase().includes(') as visible')
+}
+
 function assertIndexedCandidates(
   plan: ExplainNode,
   candidateLimit: number,
@@ -218,31 +227,25 @@ function assertIndexedCandidates(
       ? 'embedding_search_cosine_hnsw_idx'
       : `embedding_search_${width}_cosine_hnsw_idx`
   const nodes = explainNodes(plan)
-  const initial = nodes.find((node) => node['Subplan Name'] === 'CTE initial_candidates')
-  expect(initial).toBeDefined()
-  const candidateNodes = explainNodes(initial!)
-  expect(
-    candidateNodes.some((node) => node['Index Name'] === indexName && node['Actual Loops'] > 0)
-  ).toBe(true)
-  expect(candidateNodes.some((node) => node['Node Type'] === 'Sort')).toBe(false)
+  expect(nodes.some((node) => node['Index Name'] === indexName && node['Actual Loops'] > 0)).toBe(
+    true
+  )
+  /** The graph walk supplies the order, so a Sort here means the index ordering was discarded. */
+  expect(nodes.some((node) => node['Node Type'] === 'Sort')).toBe(false)
+  /**
+   * The traversal is the whole candidate set. Reaching the projection by document lookup or by
+   * sequential scan is the corpus-wide rescan this query exists to avoid, at any candidate count.
+   */
+  expect(nodes.some((node) => node['Index Name'] === 'embedding_search_document_lookup_idx')).toBe(
+    false
+  )
   expect(
-    candidateNodes.some((node) => node['Index Name'] === 'embedding_search_document_lookup_idx')
+    nodes.some(
+      (node) => node['Relation Name'] === 'embedding_search' && node['Node Type'] === 'Seq Scan'
+    )
   ).toBe(false)
-  const filtered = nodes.find((node) => node['Subplan Name'] === 'CTE filtered_scores')
-  expect(filtered).toBeDefined()
-  expect(filtered!.Output).toHaveLength(3)
-  expect(filtered!.Output![2]).toContain('<=>')
-  if (initial!['Actual Rows'] >= candidateLimit) {
-    for (const node of nodes.filter(
-      (item) =>
-        item['Subplan Name'] === 'CTE visible_search_documents' ||
-        item['CTE Name'] === 'visible_search_documents' ||
-        item['Subplan Name'] === 'CTE filtered_scores' ||
-        item['CTE Name'] === 'filtered_scores'
-    )) {
-      expect(node['Actual Loops']).toBe(0)
-    }
-  }
+  const traversed = nodes.find((node) => node['Index Name'] === indexName)!
+  expect(traversed['Actual Rows']).toBeLessThanOrEqual(candidateLimit)
 }
 
 /** Small scopes must seek chunk metadata by document without reading the full vector projection. */
@@ -294,7 +297,7 @@ const diagnosticSchema = z
     vectorBudgetMs: z.number().positive(),
     vectorCandidateDimensions: z.number().optional(),
     vectorCandidateLimit: z.number().optional(),
-    vectorCandidateScan: z.enum(['planned', 'filtered']).optional(),
+    vectorCandidateScan: z.enum(['planned', 'underfilled']).optional(),
     retrievalStatus: z.enum(['complete', 'partial']),
     timedOutLegs: z.array(z.enum(['vector', 'keyword', 'tags'])),
     toolResultBytes: z.number().int().nonnegative().optional(),
@@ -472,7 +475,7 @@ async function sample(
       (item.query.includes('order by') ||
         item.query.includes('limit') ||
         item.query.includes('CROSS JOIN LATERAL') ||
-        item.query.includes('WITH visible_search_documents') ||
+        isVectorCandidateQuery(item.query) ||
         item.query.includes('WITH scored_search_candidates') ||
         item.query.includes('WITH visible_keyword_documents'))
   )
@@ -497,10 +500,7 @@ async function sample(
       await tx.unsafe('SET LOCAL jit = off')
       await tx.unsafe("SET LOCAL hnsw.iterative_scan = 'relaxed_order'")
       await tx.unsafe('SET LOCAL hnsw.max_scan_tuples = 20000')
-      if (
-        query.query.includes('WITH visible_search_documents') ||
-        (query.query.includes('from "embedding_search"') && query.query.includes('order by'))
-      ) {
+      if (isVectorCandidateQuery(query.query)) {
         await tx.unsafe('SET LOCAL hnsw.max_scan_tuples = 1000')
         await tx.unsafe('SET LOCAL hnsw.ef_search = 1000')
         await tx.unsafe('SET LOCAL hnsw.scan_mem_multiplier = 2')
@@ -514,8 +514,7 @@ async function sample(
     plans.push({
       kind: query.query.includes('keyword_rank')
         ? 'keyword'
-        : query.query.includes('WITH visible_search_documents') ||
-            (query.query.includes('from "embedding_search"') && query.query.includes('order by'))
+        : isVectorCandidateQuery(query.query)
           ? 'vector'
           : query.query.includes('order by') ||
               query.query.includes('WITH scored_search_candidates')
@@ -526,7 +525,7 @@ async function sample(
       plan: parsedPlan,
     })
     saveReport()
-    if (query.query.includes('WITH visible_search_documents')) {
+    if (isVectorCandidateQuery(query.query)) {
       const width = diagnostics.vectorCandidateDimensions!
       expect(query.query).toContain(
         `"embedding_search"."${width === 1536 ? 'vector' : `vector_${width}`}"`

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

@@ -78,13 +78,16 @@ export interface SearchDiagnosticMetadata {
   embeddingDimensions?: number
   vectorRanking?: 'exact' | 'candidate-rerank'
   vectorCandidateStorage?: 'stored-halfvec'
-  /** Requested strategy, not an assertion about the physical index selected by PostgreSQL. */
-  vectorCandidateScan?: 'planned' | 'filtered'
+  /**
+   * Whether the bounded traversal filled its candidate limit. `underfilled` means visibility
+   * removed enough neighbours that the rerank pool is smaller than requested, which lowers recall
+   * without widening the scan. Not an assertion about the physical index PostgreSQL selected.
+   */
+  vectorCandidateScan?: 'planned' | 'underfilled'
   vectorBudgetMs?: number
   vectorCandidateLimit?: number
   vectorCandidateCount?: number
   vectorCandidateDimensions?: number
-  vectorInitialCandidateCount?: number
   resultCount?: number
   /** Tool output before the executor's final egress projection; counts only, never content. */
   toolResultBytes?: number

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

@@ -370,7 +370,7 @@ describe('workspace-scoped vector retrieval', () => {
         if (failSettings) throw failSettings
         return []
       }
-      if (statement.includes('WITH visible_search_documents')) {
+      if (statement.includes('AS visible')) {
         if (failCandidates) throw failCandidates
         return candidates
       }
@@ -434,9 +434,7 @@ describe('workspace-scoped vector retrieval', () => {
   it('uses compact candidates for a large KB and applies full workspace access before its limit', async () => {
     queueTableRows(schemaMock.embedding, [...ranked].reverse())
     expect((await handleVectorOnlySearch(params)).map((row) => row.id)).toEqual(['near', 'far'])
-    const candidate = statements().find((query) =>
-      query.sql.includes('WITH visible_search_documents')
-    )!
+    const candidate = statements().find((query) => query.sql.includes('AS visible'))!
     expect(candidate.sql).toContain('CROSS JOIN LATERAL')
     expect(candidate.sql).toContain('LIMIT 1')
     const serialized = JSON.stringify(candidate)
@@ -501,9 +499,7 @@ describe('workspace-scoped vector retrieval', () => {
     const rows = await handleVectorOnlySearch({ ...params, topK: 1 })
 
     expect(rows.map((row) => row.id)).toEqual(['far'])
-    expect(
-      statements().filter((query) => query.sql.includes('WITH visible_search_documents'))
-    ).toHaveLength(1)
+    expect(statements().filter((query) => query.sql.includes('AS visible'))).toHaveLength(1)
     expect(getForConnectors).not.toHaveBeenCalled()
   })
 
@@ -516,9 +512,7 @@ describe('workspace-scoped vector retrieval', () => {
     })
     expect(rows.map((row) => row.id)).toEqual(['near', 'far'])
     expect(Object.keys(dbChainMockFns.select.mock.calls[0][0])).toEqual(['id'])
-    const candidate = statements().find((query) =>
-      query.sql.includes('WITH visible_search_documents')
-    )!
+    const candidate = statements().find((query) => query.sql.includes('AS visible'))!
     expect(JSON.stringify(candidate)).toContain('common')
     expect(JSON.stringify(candidate)).toContain(String(schemaMock.embedding.tag1))
     expect(JSON.stringify(dbChainMockFns.where.mock.calls.at(-1)![0])).toContain('common')
@@ -534,9 +528,7 @@ describe('workspace-scoped vector retrieval', () => {
     const rows = await handleVectorOnlySearch({ ...params, knowledgeBaseIds })
     expect(rows.map((row) => row.id)).toEqual(['near', 'far'])
     expect(rows.every((row) => row.knowledgeBaseId === 'kb-1')).toBe(true)
-    const candidateQueries = statements().filter((query) =>
-      query.sql.includes('WITH visible_search_documents')
-    )
+    const candidateQueries = statements().filter((query) => query.sql.includes('AS visible'))
     expect(candidateQueries).toHaveLength(1)
     for (const id of knowledgeBaseIds) expect(JSON.stringify(candidateQueries[0])).toContain(id)
     expect(dbChainMockFns.transaction).toHaveBeenCalledOnce()
@@ -616,9 +608,7 @@ describe('workspace-scoped vector retrieval', () => {
     expect(statements().filter((query) => query.sql.includes('hnsw.iterative_scan'))).toHaveLength(
       1
     )
-    const queries = statements().filter((query) =>
-      query.sql.includes('WITH visible_search_documents')
-    )
+    const queries = statements().filter((query) => query.sql.includes('AS visible'))
     expect(queries).toHaveLength(2)
     expect(JSON.stringify(queries[0])).toBe(JSON.stringify(queries[1]))
     await vi.advanceTimersByTimeAsync(10 * 60 * 1000 + 1)
@@ -807,7 +797,7 @@ describe('live repository authorization follows ranked candidates', () => {
     dbChainMockFns.execute.mockImplementation(async (query) =>
       render(query).sql.includes('SELECT scoped_chunk.id')
         ? (probePages.shift() ?? [])
-        : render(query).sql.includes('WITH visible_search_documents')
+        : render(query).sql.includes('AS visible')
           ? (candidatePages.shift() ?? [])
           : render(query).sql.includes('WITH scored_search_candidates')
             ? (rerankPages.shift() ?? [])
@@ -841,9 +831,8 @@ describe('live repository authorization follows ranked candidates', () => {
     })
     expect(rows.map((row) => row.id)).toEqual(['near'])
     const candidateQuery = dbChainMockFns.execute.mock.calls.find(([query]) =>
-      render(query).sql.includes('WITH visible_search_documents')
+      render(query).sql.includes('AS visible')
     )![0]
-    expect(render(candidateQuery).sql).toContain('MATERIALIZED')
     expect(render(candidateQuery).sql).toContain('CROSS JOIN LATERAL')
     expect(render(candidateQuery).sql).toContain('LIMIT 1')
     expect(JSON.stringify(candidateQuery)).toContain('required_clause')
@@ -922,7 +911,7 @@ describe('live repository authorization follows ranked candidates', () => {
     }
   )
 
-  it('scans the filtered projection when ANN cannot fill its limit', async () => {
+  it('keeps an underfilled ANN result instead of rescoring the whole projection', async () => {
     probePages.push(Array.from({ length: 400 }, (_, index) => ({ id: `probe-${index}` })))
     queueCandidates([{ id: 'selected' }], 1)
     queueRerank([candidate('selected', 'allowed-source')])
@@ -933,12 +922,12 @@ describe('live repository authorization follows ranked candidates', () => {
       { id: 'selected', content: 'Verified fallback', distance: 0.1 },
     ])
     const candidateQuery = dbChainMockFns.execute.mock.calls.find(([query]) =>
-      render(query).sql.includes('WITH visible_search_documents')
+      render(query).sql.includes('AS visible')
     )![0]
-    expect(render(candidateQuery).sql).toContain('UNION ALL')
-    expect(render(candidateQuery).sql).toContain('+ 0')
-    expect(render(candidateQuery).sql).toContain('filtered_scores AS MATERIALIZED')
-    expect(render(candidateQuery).sql).toContain('ORDER BY filtered_scores.distance + 0')
+    /** Widening the scan on underfill is what made this leg exceed its budget on a large corpus. */
+    expect(render(candidateQuery).sql).not.toContain('UNION ALL')
+    expect(render(candidateQuery).sql).not.toContain('filtered_scores')
+    expect(render(candidateQuery).sql).toContain('CROSS JOIN LATERAL')
     expect(JSON.stringify(dbChainMockFns.where.mock.calls.at(-1)![0])).toContain(
       'github_read_grant'
     )

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

@@ -50,8 +50,12 @@ const UNDEFINED_OBJECT_SQLSTATE = '42704'
 /** Bound candidate pages retained while live permissions are checked. */
 const MAX_AUTHORIZED_SEARCH_CANDIDATES = 20_000
 /**
- * Stop a permission-starved graph walk early enough to scan the filtered projection instead.
+ * Bounds a permission-starved graph walk, which returns fewer candidates rather than widening.
  * This approximate iterative-visit threshold excludes pgvector's initial scan; it is not a row limit.
+ *
+ * Raising it trades recall for latency far more steeply than its size suggests: on a 132k-chunk
+ * index at 10% visibility, visiting 6.5k tuples instead of 1.5k took 5.1s and 9.7s on consecutive
+ * identical runs, against ~115ms for the bounded walk. Re-measure before changing it.
  */
 const CANDIDATE_HNSW_MAX_SCAN_TUPLES = '1000'
 const CANDIDATE_HNSW_EF_SEARCH = '1000'
@@ -902,57 +906,37 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
         ),
       })
       /**
-       * LIMIT keeps document authorization downstream of vector traversal, with a primary-key
-       * lookup per candidate. Only an underfilled ANN scan materializes the visible document set.
-       * Its exact fallback scores the compact projection once, then joins scalar distances to
-       * visible identities; it cannot turn into a random vector lookup for every document.
+       * The bounded ANN traversal is the whole candidate set. LIMIT keeps document authorization
+       * downstream of the traversal, with a primary-key lookup per candidate.
+       *
+       * An underfilled traversal yields fewer candidates rather than widening the search. Widening
+       * it has no affordable form here: rescoring the projection exhaustively is O(corpus) and a
+       * deeper `hnsw.max_scan_tuples` is worse still — measured on a 132k-chunk index at 10%
+       * visibility, the exhaustive rescan took 1.9s while scanning 6.5k tuples instead of 1.5k took
+       * 5.1s and 9.7s on consecutive identical runs. Both exceed the retrieval budget on a corpus
+       * an order of magnitude larger, and a leg that exceeds its budget returns nothing at all, so
+       * fewer candidates strictly beats every widening strategy available.
        */
       const identities = await withVectorScanSettings(
         (executor) =>
-          executor.execute<{ id: string; initial_count: number }>(sql`
-          WITH visible_search_documents AS MATERIALIZED (
-            SELECT ${document.id} AS id FROM ${document}
-            WHERE ${and(...candidateDocumentVisibility)}
-          ), initial_candidates AS MATERIALIZED (
-            SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch}
-            CROSS JOIN LATERAL (
-              SELECT 1 FROM ${document}
-              WHERE ${and(eq(document.id, embeddingSearch.documentId), ...candidateDocumentVisibility, candidateTagCondition)}
-              LIMIT 1
-            ) AS visible
-            WHERE ${and(
-              inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds),
-              eq(embeddingSearch.enabled, true)
-            )}
-            ORDER BY ${candidateDistance} LIMIT ${candidateLimit}
-          ), filtered_scores AS MATERIALIZED (
-            SELECT ${embeddingSearch.id} AS id, ${embeddingSearch.documentId} AS document_id,
-              ${candidateDistance} AS distance FROM ${embeddingSearch}
-            WHERE ${and(
-              inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds),
-              eq(embeddingSearch.enabled, true),
-              candidateTagCondition
-            )}
-              AND (SELECT count(*) FROM initial_candidates) < ${candidateLimit}
-          ), candidates AS (
-            SELECT id FROM initial_candidates
-            WHERE (SELECT count(*) FROM initial_candidates) >= ${candidateLimit}
-            UNION ALL (
-              SELECT filtered_scores.id FROM filtered_scores
-              INNER JOIN visible_search_documents ON visible_search_documents.id = filtered_scores.document_id
-              WHERE (SELECT count(*) FROM initial_candidates) < ${candidateLimit}
-              ORDER BY filtered_scores.distance + 0, filtered_scores.id
-              LIMIT ${candidateLimit}
-            )
-          ) SELECT id, (SELECT count(*)::int FROM initial_candidates) AS initial_count FROM candidates
+          executor.execute<{ id: string }>(sql`
+          SELECT ${embeddingSearch.id} AS id FROM ${embeddingSearch}
+          CROSS JOIN LATERAL (
+            SELECT 1 FROM ${document}
+            WHERE ${and(eq(document.id, embeddingSearch.documentId), ...candidateDocumentVisibility, candidateTagCondition)}
+            LIMIT 1
+          ) AS visible
+          WHERE ${and(
+            inArray(embeddingSearch.knowledgeBaseId, params.knowledgeBaseIds),
+            eq(embeddingSearch.enabled, true)
+          )}
+          ORDER BY ${candidateDistance} LIMIT ${candidateLimit}
         `),
         params.budget
       )
-      const initialCount = identities[0]?.initial_count ?? 0
       annotateSearchDiagnostics({
         vectorCandidateCount: identities.length,
-        vectorInitialCandidateCount: initialCount,
-        vectorCandidateScan: initialCount < candidateLimit ? 'filtered' : 'planned',
+        vectorCandidateScan: identities.length < candidateLimit ? 'underfilled' : 'planned',
       })
       if (!identities.length) return { candidates: [], nextOffset: offset }
       /** Score each bounded candidate once; sorting the materialized scalar cannot invoke HNSW again. */