fix(knowledge): keep vector candidate selection bounded (#7960)

* fix(knowledge): keep vector candidate selection bounded

An underfilled ANN traversal rescored the whole compact projection and joined the
visible document set to it. That fallback is O(corpus): on a 132k-chunk index it
took 1.9s, and on a corpus an order of magnitude larger it exceeds the retrieval
budget, so the vector leg returned nothing at all rather than fewer rows.

Widening the search has no affordable form here. Measured at 10% visibility on the
same index, scanning 6.5k tuples instead of 1.5k took 5.1s and 9.7s on consecutive
identical runs, and joining visibility before scoring took 8.4s because it turns a
sequential scan into a random lookup per document. The bounded traversal itself
costs ~115ms whether or not it fills.

The traversal is now the whole candidate set. An underfilled one yields fewer
candidates and is reported as such, which strictly beats a leg that times out.

* fix(knowledge): match the latency plan assertions to the bounded traversal

The plan assertion still required the removed CTEs, and the predicates selecting
which captured query to assert against matched the old CTE name with a lowercase
fallback the rendered statement never produces. The candidate assertions would
have stopped running rather than failing, so they now key on the visibility
lateral through one shared predicate.

The plan check drops the fallback-specific expectations and gains the one that
guards this change: the traversal must never reach the projection by document
lookup or sequential scan, at any candidate count.

Author: waleedlatif1

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. */