fix(knowledge): recover vector candidates an HNSW post-filter discards (#7971)

* fix(knowledge): recover vector candidates an HNSW post-filter discards

pgvector's HNSW index post-filters by construction, so a visibility predicate
can only discard neighbours the graph walk already committed to. When the
documents a caller may read are a small share of the index, the traversal
returns a handful of candidates instead of its limit.

Rank the permitted set exactly when the traversal comes back underfilled and a
bounded probe says that set is small enough to afford. A traversal that fills
its limit is returned untouched, so a scope the graph serves well pays nothing.

Bound the probe by documents examined rather than chunks accumulated, and give
it its own sub-budget so deciding against a rescue can never cost the leg its
results.

Retune the scan settings so the tuple budget sits an order of magnitude above
the beam, which is what lets the iterative scan iterate at all, and so the
narrower beam lowers the leg's uninterruptible floor.

* test(knowledge): cover the capped step budget

* fix(knowledge): require an enabled chunk when a tag filter decides a document

A document whose only tagged chunk is disabled could be admitted by the probe
and then discarded by ranking, spending the probe's document bound on a
document that can contribute no candidate.

Update the KB block fan-out integration test to the strategy it now exercises:
a scope too small to fill the traversal probes once and rescues once.

Author: waleedlatif1

apps/sim/app/api/knowledge/search/utils.test.ts

@@ -212,11 +212,14 @@ describe('Knowledge Search Utils', () => {
   describe('handleTagAndVectorSearch', () => {
     it('returns only bounded ranked rows without first materializing every matching tag ID', async () => {
       resetDbChainMock()
-      queueTableRows(
-        schemaMock.embedding,
-        Array.from({ length: 201 }, (_, index) => ({ id: `candidate-${index}` }))
-      )
-      queueTableRows(schemaMock.embedding, [makeResult('second', 0.2), makeResult('first', 0.1)])
+      dbChainMockFns.execute.mockImplementation(async (query) => {
+        const statement = (query as { toSQL: () => { sql: string } }).toSQL().sql
+        if (statement.includes('AS visible')) return []
+        if (statement.includes(') + 0 LIMIT')) return [{ id: 'first' }, { id: 'second' }]
+        if (statement.includes('WITH scored_search_candidates'))
+          return [makeResult('second', 0.2), makeResult('first', 0.1)]
+        return [{ id: 'doc-first' }, { id: 'doc-second' }]
+      })
       queueTableRows(schemaMock.embedding, [makeResult('second', 0.2), makeResult('first', 0.1)])
 
       const results = await handleTagAndVectorSearch({
@@ -231,11 +234,13 @@ describe('Knowledge Search Utils', () => {
       })
 
       expect(results.map((row) => row.id)).toEqual(['first', 'second'])
-      expect(dbChainMockFns.select).toHaveBeenCalledTimes(3)
-      expect(Object.keys(dbChainMockFns.select.mock.calls[0][0])).toEqual(['id'])
-      expect(dbChainMockFns.limit).toHaveBeenNthCalledWith(1, 400)
-      expect(dbChainMockFns.select.mock.calls[1][0]).toHaveProperty('distance')
-      expect(dbChainMockFns.limit).toHaveBeenCalledWith(20)
+      /** Only hydration reads through the query builder; ranking never materializes tag IDs. */
+      expect(dbChainMockFns.select).toHaveBeenCalledTimes(1)
+      expect(dbChainMockFns.select.mock.calls[0][0]).toHaveProperty('distance')
+      const exact = dbChainMockFns.execute.mock.calls
+        .map(([query]) => (query as { toSQL: () => { sql: string; params: unknown[] } }).toSQL())
+        .find((statement) => statement.sql.includes(') + 0 LIMIT'))!
+      expect(exact.params).toContain(400)
     })
 
     it('should throw error when no filters provided', async () => {
@@ -552,13 +557,13 @@ describe('Knowledge Search Utils', () => {
       })
 
       expect(results.map((r) => r.id)).toEqual(['vector-hit'])
-      expect(dbChainMockFns.select).toHaveBeenCalledTimes(2)
+      expect(dbChainMockFns.select).toHaveBeenCalledTimes(1)
     })
 
     it('runs both legs and fuses them in hybrid mode', async () => {
       /**
-       * The raw vector probe does not consume a table chain. Keyword ranking and
-       * hydration complete before vector exact ranking and content hydration.
+       * Vector ranking is raw SQL throughout and consumes no table chain. Keyword ranking and
+       * hydration complete before vector content hydration.
        */
       dbChainMockFns.execute.mockResolvedValue([{ id: 'vector-hit' }])
       queueTableRows(schemaMock.embedding, [{ id: 'keyword-hit', keywordRank: 0.9 }])
@@ -576,7 +581,7 @@ describe('Knowledge Search Utils', () => {
       })
 
       expect(results.map((r) => r.id).sort()).toEqual(['keyword-hit', 'vector-hit'])
-      expect(dbChainMockFns.select).toHaveBeenCalledTimes(4)
+      expect(dbChainMockFns.select).toHaveBeenCalledTimes(3)
     })
 
     it('propagates unexpected keyword errors after the vector leg finishes', async () => {
@@ -601,7 +606,7 @@ describe('Knowledge Search Utils', () => {
           queryVector: { vector: JSON.stringify(TEST_EMBEDDING), dimensions: 1536 },
         })
       ).rejects.toBe(failure)
-      expect(dbChainMockFns.select).toHaveBeenCalledTimes(3)
+      expect(dbChainMockFns.select).toHaveBeenCalledTimes(2)
     })
 
     it('skips both query legs when only tag filters are provided', async () => {

apps/sim/lib/knowledge/__integration__/kb-block-search.integration.ts

@@ -91,7 +91,7 @@ describe('API-key KB block fan-out', () => {
       const previousDebug = db.$client.options.debug
       const statements: string[] = []
       db.$client.options.debug = (_connection, query) => {
-        if (statements.length < 250) statements.push(query)
+        if (statements.length < 1000) statements.push(query)
       }
       try {
         const results = await Promise.all(
@@ -124,9 +124,26 @@ describe('API-key KB block fan-out', () => {
           expect(result.rows[0].knowledgeBaseId).toBe(bases[index].id)
           expect(result.rows[0].distance).toBeCloseTo(0)
         }
-        expect(statements.filter((query) => query.includes('statement_timeout'))).toHaveLength(54)
-        expect(statements.filter((query) => query.includes('+ 0'))).toHaveLength(18)
-        expect(statements.some((query) => query.includes('hnsw.iterative_scan'))).toBe(false)
+        const matching = (fragment: string) =>
+          statements.filter((query) => query.includes(fragment))
+        /**
+         * Every statement runs under the leg's deadline: the candidate search reinstates it after
+         * tuning the scan, and the probe, the exact ranking, the rerank and hydration each open
+         * with one of their own.
+         */
+        expect(matching('statement_timeout')).toHaveLength(bases.length * 6)
+        /**
+         * A scope this small leaves the bounded traversal short of its candidate limit, so every
+         * search probes once and rescues once — never a widening retry loop.
+         */
+        expect(matching('hnsw.iterative_scan')).toHaveLength(bases.length)
+        expect(matching('AS visible')).toHaveLength(bases.length)
+        expect(matching(') + 0 LIMIT')).toHaveLength(bases.length)
+        expect(matching('scored_search_candidates')).toHaveLength(bases.length)
+        /** The probe enumerates visible documents; it never ranks them. */
+        expect(
+          statements.filter((query) => query.includes('AS id FROM') && !query.includes('ORDER BY'))
+        ).toHaveLength(bases.length)
       } finally {
         db.$client.options.debug = previousDebug
       }

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

@@ -44,6 +44,23 @@ describe('search SQL deadline', () => {
     }
   })
 
+  it('caps a step without letting it shorten or outlive the leg it came from', () => {
+    vi.spyOn(performance, 'now').mockReturnValue(0)
+    const controller = new AbortController()
+    const leg = new SearchBudget('vector', 1000, controller.signal)
+    const step = leg.capped(600)
+    expect(step.deadline).toBe(600)
+    expect(step.leg).toBe('vector')
+    expect(step.signal).toBe(controller.signal)
+    /** Spending the step is the step's own business; the leg keeps its deadline and its verdict. */
+    step.isTimeout(new SearchDeadlineError())
+    expect(step.timedOut).toBe(true)
+    expect(leg.timedOut).toBe(false)
+    expect(leg.remaining()).toBe(1000)
+    /** A step longer than what the leg has left cannot extend it. */
+    expect(new SearchBudget('vector', 200).capped(600).deadline).toBe(200)
+  })
+
   it('preserves cancellation and unexpected errors instead of labeling them incomplete evidence', async () => {
     const controller = new AbortController()
     const budget = new SearchBudget('keyword', performance.now() - 1, controller.signal)

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

@@ -39,6 +39,18 @@ export class SearchBudget {
     return remaining
   }
 
+  /**
+   * A shorter deadline for one step of the leg, carrying the same leg and cancellation signal.
+   * Spending it is the step's own business: the leg's deadline is untouched and still governs.
+   */
+  capped(milliseconds: number): SearchBudget {
+    return new SearchBudget(
+      this.leg,
+      Math.min(this.deadline, performance.now() + milliseconds),
+      this.signal
+    )
+  }
+
   isTimeout(error: unknown): boolean {
     this.signal?.throwIfAborted()
     if (error instanceof SearchDeadlineError || getPostgresErrorCode(error) === '57014') {

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

@@ -47,6 +47,7 @@ export type SearchStage =
   | 'vector.settings'
   | 'vector.probe'
   | 'vector.rerank'
+  | 'vector.exact_candidates'
   | 'vector.exact'
   | 'vector.candidate_search'
   | 'source_overview'
@@ -76,7 +77,7 @@ export interface SearchDiagnosticMetadata {
   searchMode?: 'hybrid' | 'vector'
   boostRecency?: boolean
   embeddingDimensions?: number
-  vectorRanking?: 'exact' | 'candidate-rerank'
+  vectorRanking?: 'exact' | 'exact-candidates' | 'candidate-rerank'
   vectorCandidateStorage?: 'stored-halfvec'
   /**
    * Whether the bounded traversal filled its candidate limit. `underfilled` means visibility
@@ -86,6 +87,8 @@ export interface SearchDiagnosticMetadata {
   vectorCandidateScan?: 'planned' | 'underfilled'
   vectorBudgetMs?: number
   vectorCandidateLimit?: number
+  /** Visible documents the tractability probe enumerated, capped at its own document limit. */
+  vectorProbeDocumentCount?: number
   vectorCandidateCount?: number
   vectorCandidateDimensions?: number
   resultCount?: number