fix(search): use compact candidates for filtered vector retrieval (#7794)

Author: icecrasher321

apps/sim/lib/copilot/tools/server/knowledge/workspace-search.test.ts

@@ -84,6 +84,29 @@ describe('Assistant retrieval tools', () => {
       next: null,
     })
   })
+  it('returns empty incomplete retrieval as a recoverable search outcome and logs coverage', async () => {
+    mocks.search.mockResolvedValue({
+      retrieval: { status: 'partial', timedOutLegs: ['vector', 'keyword'] },
+      knowledgeBases: [{ id: 'index', name: 'Enterprise Search' }],
+      results: [],
+    })
+
+    const result = await searchWorkspaceServerTool.execute({ query: 'canaries' }, context)
+
+    expect(result).toMatchObject({
+      success: true,
+      message: expect.stringContaining('cannot establish absence or completeness'),
+      data: {
+        retrieval: { status: 'partial', timedOutLegs: ['vector', 'keyword'] },
+        results: [],
+      },
+    })
+    expect(result).not.toHaveProperty('error')
+    expect(mocks.info).toHaveBeenCalledWith(
+      'Knowledge search completed',
+      expect.objectContaining({ passageBytes: 0, originalPassageBytes: 0, outcome: 'success' })
+    )
+  })
   it('pins organization and private chat while reusing the canonical search index and citations', async () => {
     const orgContext = {
       ...context,

apps/sim/lib/copilot/tools/server/knowledge/workspace-search.ts

@@ -93,7 +93,7 @@ export const searchWorkspaceServerTool: BaseServerTool = {
             const names = new Map(result.knowledgeBases.map((base) => [base.id, base.name]))
             const output = {
               success: true,
-              message: `${result.retrieval.status === 'partial' ? 'Partial search: a retrieval branch reached its deadline. These results cannot establish absence or completeness. ' : ''}Found ${result.results.length} passage previews. Read a document at its chunkIndex for more context. ${CITATION_INSTRUCTION}`,
+              message: `${result.retrieval.status === 'partial' ? 'Search coverage is incomplete. Continue with a more specific query or source filter; these results cannot establish absence or completeness. ' : ''}Found ${result.results.length} passage previews. Read a document at its chunkIndex for more context. ${CITATION_INSTRUCTION}`,
               data: {
                 query: safeQuery,
                 retrieval: result.retrieval,

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

@@ -83,7 +83,11 @@ const reused = reuseFile ? readFixtureReport(reuseFile) : undefined
 const ids = reused?.fixture ?? createKnowledgeAclFixtureIds()
 const unrelated = reused?.unrelatedFixture ?? createKnowledgeAclFixtureIds()
 const organizationChatId = generateId()
-const queryVector = Array.from({ length: dimensions }, (_, index) => (index === 0 ? 1 : 0))
+const queryVector = Array.from({ length: dimensions }, (_, index) =>
+  Math.sin((index + 1) * 12.9898)
+)
+const queryMagnitude = Math.hypot(...queryVector)
+for (let index = 0; index < queryVector.length; index++) queryVector[index] /= queryMagnitude
 const captured: CapturedQuery[] = []
 const report: Record<string, unknown> = {
   fixture: ids,
@@ -133,6 +137,7 @@ const explainSchema = z.array(z.object({ Plan: explainNodeSchema }).passthrough(
 
 function usesVectorIndex(node: ExplainNode): boolean {
   return (
+    node['Index Name'] === 'embedding_binary_hnsw_idx' ||
     node['Index Name'] === 'embedding_vector_hnsw_idx' ||
     (node.Plans?.some(usesVectorIndex) ?? false)
   )
@@ -237,14 +242,21 @@ async function sample(label: string, run: () => ReturnType<typeof search>) {
     const plan = await db.$client.begin(async (tx) => {
       await tx.unsafe("SET LOCAL hnsw.iterative_scan = 'relaxed_order'")
       await tx.unsafe('SET LOCAL hnsw.max_scan_tuples = 20000')
+      if (query.query.includes('binary_quantize')) {
+        await tx.unsafe('SET LOCAL hnsw.max_scan_tuples = 100000')
+        await tx.unsafe('SET LOCAL hnsw.ef_search = 200')
+        await tx.unsafe('SET LOCAL hnsw.scan_mem_multiplier = 4')
+      }
       return tx.unsafe(`EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) ${query.query}`, query.parameters)
     })
     plans.push({
       kind: query.query.includes('keyword_rank')
         ? 'keyword'
-        : query.query.includes('order by')
+        : query.query.includes('binary_quantize')
           ? 'vector'
-          : 'probe',
+          : query.query.includes('order by')
+            ? 'rerank'
+            : 'probe',
       query: query.query,
       parameters: query.parameters,
       plan: explainSchema.parse(plan[0]['QUERY PLAN']),
@@ -378,8 +390,8 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
             CASE WHEN n % 8 = 0 THEN 'Orion deployment reference. ' ELSE 'Engineering operations reference. ' END ||
               (SELECT string_agg(md5(n::text || ':' || paragraph::text), ' ') FROM generate_series(1, 90) paragraph),
             3000, 750, 0, 3000,
-            l2_normalize(ARRAY(SELECT (CASE WHEN coordinate = n % 32 + 1 THEN 1 ELSE 0 END +
-              0.025 * sin(n::double precision * coordinate * 12.9898 + coordinate * 78.233))::real
+            l2_normalize(ARRAY(SELECT (sin(coordinate * (n % 32 + 1) * 12.9898) +
+              0.25 * sin(n::double precision * coordinate * 12.9898 + coordinate * 78.233))::real
               FROM generate_series(1, ${dimensions}) coordinate)::vector(1536))
           FROM generate_series(${first}::int, ${last}::int) n`)
           })
@@ -494,7 +506,13 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
           }
         )
         if (delayedLegs === 'both') {
-          expect(result).toMatchObject({ success: false, retryable: true })
+          expect(result).toMatchObject({
+            success: true,
+            data: {
+              retrieval: { status: 'partial', timedOutLegs: ['vector', 'keyword'] },
+              results: [],
+            },
+          })
           return
         }
         expect(result).toMatchObject({
@@ -526,6 +544,18 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
       const vectorPlans = plans.filter((plan) => plan.kind === 'vector')
       expect(vectorPlans).toHaveLength(1)
       expect(usesVectorIndex(vectorPlans[0].plan[0].Plan)).toBe(true)
+      expect(plans.some((plan) => plan.kind === 'rerank')).toBe(true)
+      const rerank = plans.find((plan) => plan.kind === 'rerank')!
+      const actual = await db.$client.unsafe(rerank.query, rerank.parameters).values()
+      const expected = await db.execute<{ id: string }>(sql`SELECT id FROM embedding
+        WHERE knowledge_base_id = ${ids.knowledgeBaseId} AND enabled
+        ORDER BY (embedding <=> ${JSON.stringify(queryVector)}::vector) + 0, id
+        LIMIT ${actual.length}`)
+      const expectedIds = new Set(expected.map(({ id }) => id))
+      const recall = actual.filter(([id]) => expectedIds.has(id)).length / expected.length
+      expect(recall).toBeGreaterThanOrEqual(0.95)
+      report[`recall.${iteration}`] = { neighbors: expected.length, recall }
+      saveReport()
     }
     expect(embeddingCalls - before).toBe(2)
   }, 180_000)
@@ -578,7 +608,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
       expect(probe).toHaveLength(1)
       expect(probe[0].query).not.toContain('<=>')
       expect(probe[0].plan[0].Plan['Actual Rows']).toBe(12)
-      const vector = plans.filter((plan) => plan.kind === 'vector')
+      const vector = plans.filter((plan) => plan.kind === 'rerank')
       expect(vector).toHaveLength(1)
       expect(vector[0].query).toContain('"embedding"."id" in')
     } finally {
@@ -672,7 +702,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
   }, 180_000)
   /** Opt in with local Sim and Go URLs; uses the real configured provider, billing adapter, and async resume protocol. */
   it.skipIf(!process.env.KNOWLEDGE_SEARCH_ASSISTANT_URL)(
-    'answers through local Go Assistant with progressive reads and citations',
+    'recovers quietly from incomplete search through local Go Assistant, then reads and cites evidence',
     async () => {
       const assistantUrl = new URL(process.env.KNOWLEDGE_SEARCH_ASSISTANT_URL!)
       const simUrl = new URL(process.env.KNOWLEDGE_SEARCH_SIM_URL!)
@@ -708,6 +738,18 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
         bytes: number
       }> = []
       let answer = ''
+      let incompleteSearch = true
+      const query = SearchBudget.prototype.query
+      const delayed = vi.spyOn(SearchBudget.prototype, 'query').mockImplementation(function <T>(
+        this: SearchBudget,
+        stage: SearchStage,
+        run: (executor: SearchExecutor) => PromiseLike<T>
+      ): Promise<T> {
+        return query.call(this, stage, async (tx) => {
+          if (incompleteSearch) await tx.execute(sql`SELECT pg_sleep(9)`)
+          return run(tx)
+        }) as Promise<T>
+      })
       const started = performance.now()
       try {
         await db
@@ -808,9 +850,20 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
                 bytes: Buffer.byteLength(JSON.stringify(result)),
               })
               const { success } = z.object({ success: z.boolean() }).parse(result)
+              expect(success).toBe(true)
+              if (incompleteSearch) {
+                expect(call.toolName).toBe('search_workspace')
+                expect(result).toMatchObject({
+                  data: {
+                    retrieval: { status: 'partial', timedOutLegs: ['vector', 'keyword'] },
+                    results: [],
+                  },
+                })
+              }
               return { callId, name: call.toolName, success, data: result }
             })
           )
+          incompleteSearch = false
           path = '/api/tools/resume'
           body = {
             checkpointId: checkpoint.checkpointId,
@@ -828,10 +881,12 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
         expect(answer).toContain('SILVER COMET')
         expect(answer).toContain('K7M2-84')
         expect(answer).toContain('<source>')
+        expect(answer).not.toMatch(/timed?\s*out|timeout|internal retr(?:y|ies)/i)
         expect(calls.some((call) => call.name === 'read_document')).toBe(true)
-        expect(calls.some((call) => call.name === 'search_workspace')).toBe(true)
+        expect(calls.filter((call) => call.name === 'search_workspace').length).toBeGreaterThan(1)
         expect(calls.every((call) => call.bytes < 40000)).toBe(true)
       } finally {
+        delayed.mockRestore()
         await db.update(embedding).set(original).where(eq(embedding.id, chunkId))
       }
     },

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

@@ -19,6 +19,7 @@ const mocks = vi.hoisted(() => ({
   checkActorUsage: vi.fn(),
   generateEmbedding: vi.fn(),
   executeSearch: vi.fn(),
+  retrieval: vi.fn(),
   getDocumentMetadata: vi.fn(),
   getTagDefinitions: vi.fn(),
   getTagDefinitionsBatch: vi.fn(),
@@ -91,7 +92,7 @@ vi.mock('@/lib/knowledge/search/queries', () => ({
   generateSearchEmbedding: mocks.generateEmbedding,
   retrieveKnowledgeSearch: async (...args: unknown[]) => ({
     rows: await mocks.executeSearch(...args),
-    retrieval: { status: 'complete', timedOutLegs: [] },
+    retrieval: mocks.retrieval(),
   }),
   getDocumentMetadataByIds: mocks.getDocumentMetadata,
 }))
@@ -129,6 +130,7 @@ const knowledgeBase = {
 
 describe('knowledge search application use case', () => {
   beforeEach(() => {
+    mocks.retrieval.mockReturnValue({ status: 'complete', timedOutLegs: [] })
     vi.clearAllMocks()
     mocks.rerank.mockReset()
     resetDbChainMock()
@@ -204,6 +206,34 @@ describe('knowledge search application use case', () => {
     expect(result.totalResults).toBe(0)
   })
 
+  it.each([false, true])(
+    'requires explicit partial-result support for empty incomplete searches (allowPartialResults=%s)',
+    async (allowPartialResults) => {
+      mocks.retrieval.mockReturnValue({ status: 'partial', timedOutLegs: ['vector', 'keyword'] })
+      mocks.executeSearch.mockResolvedValue([])
+      const result = searchKnowledge.execute({
+        principal: { kind: 'session', userId: 'user-1', sessionId: 'session-1' },
+        input: {
+          workspaceId: 'workspace-1',
+          knowledgeBaseIds: ['knowledge-1'],
+          query: 'canaries',
+          topK: 20,
+          allowPartialResults,
+        },
+      })
+
+      if (!allowPartialResults) {
+        await expect(result).rejects.toThrow('retrieval deadline')
+        return
+      }
+      await expect(result).resolves.toMatchObject({
+        results: [],
+        totalResults: 0,
+        retrieval: { status: 'partial', timedOutLegs: ['vector', 'keyword'] },
+      })
+    }
+  )
+
   describe.each(['workspace', 'organization'] as const)('%s ranking policy', (scope) => {
     beforeEach(() => {
       if (scope === 'organization') {

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

@@ -709,9 +709,6 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({
       }
     }
     annotateSearchDiagnostics({ resultCount: results.length })
-    if (retrieved.retrieval.status === 'partial' && results.length === 0) {
-      throw new SearchDeadlineError()
-    }
     const cost = baseCost
       ? {
           input: baseCost.input,

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

@@ -47,6 +47,7 @@ export type SearchStage =
   | 'vector.settings'
   | 'vector.probe'
   | 'vector.ann'
+  | 'vector.rerank'
   | 'vector.exact'
 
 /** Fixed, content-free fields. Never pass queries, filters, document identities, SQL, or errors. */
@@ -68,6 +69,8 @@ export interface SearchDiagnosticMetadata {
   searchMode?: 'hybrid' | 'vector'
   boostRecency?: boolean
   embeddingDimensions?: number
+  vectorRanking?: 'exact' | 'binary-rerank'
+  vectorCandidateCount?: 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

@@ -528,6 +528,7 @@ describe('live repository authorization follows ranked candidates', () => {
       schemaMock.embedding,
       Array.from({ length: 200 }, (_, index) => candidate(`probe-${index}`, 'allowed-source'))
     )
+    queueTableRows(schemaMock.embedding, [{ id: 'far' }, { id: 'near' }])
     queueTableRows(schemaMock.embedding, [
       { ...candidate('far', 'allowed-source'), distance: 0.3 },
       { ...candidate('near', 'allowed-source'), distance: 0.1 },
@@ -544,13 +545,17 @@ describe('live repository authorization follows ranked candidates', () => {
     })
     expect(rows.map((row) => row.id)).toEqual(['near'])
     expect(dbChainMockFns.orderBy.mock.calls[0]).toHaveLength(1)
-    expect(Object.keys(dbChainMockFns.select.mock.calls[1][0])).toEqual(['id', 'distance'])
+    expect(Object.keys(dbChainMockFns.select.mock.calls[1][0])).toEqual(['id'])
+    expect(render(dbChainMockFns.orderBy.mock.calls[0][0]).sql).toContain('binary_quantize')
+    expect(JSON.stringify(dbChainMockFns.orderBy.mock.calls[1][0])).toContain('<=>')
+    expect(dbChainMockFns.limit.mock.calls[1]).toEqual([4000])
+    expect(JSON.stringify(dbChainMockFns.where.mock.calls[1][0])).not.toContain('<=>')
     expect(dbChainMockFns.limit.mock.invocationCallOrder[1]).toBeLessThan(
       dbChainMockFns.select.mock.invocationCallOrder[2]
     )
     expect(JSON.stringify(dbChainMockFns.where.mock.calls[1][0])).toContain('OFFSET 0')
     expect(getForConnectors).toHaveBeenCalledExactlyOnceWith(['allowed-source'], undefined)
-    expect(JSON.stringify(dbChainMockFns.where.mock.calls[2][0])).toContain('github_read_grant')
+    expect(JSON.stringify(dbChainMockFns.where.mock.calls[3][0])).toContain('github_read_grant')
   })
 
   it('finishes empty scopes after the bounded probe without scanning HNSW or calling providers', async () => {
@@ -592,6 +597,7 @@ describe('live repository authorization follows ranked candidates', () => {
       schemaMock.embedding,
       Array.from({ length: 200 }, (_, index) => candidate(`probe-${index}`, 'allowed-source'))
     )
+    queueTableRows(schemaMock.embedding, [{ id: 'partial' }])
     queueTableRows(schemaMock.embedding, [candidate('partial', 'allowed-source')])
     queueTableRows(schemaMock.embedding, [candidate('selected', 'allowed-source')])
     queueTableRows(schemaMock.embedding, [
@@ -614,6 +620,10 @@ describe('live repository authorization follows ranked candidates', () => {
       candidate(`approximate-${index}`, 'allowed-source')
     )
     queueTableRows(schemaMock.embedding, probe)
+    queueTableRows(
+      schemaMock.embedding,
+      approximate.map(({ id }) => ({ id }))
+    )
     queueTableRows(schemaMock.embedding, approximate)
     queueTableRows(schemaMock.embedding, [])
     queueTableRows(schemaMock.embedding, probe)
@@ -628,7 +638,7 @@ describe('live repository authorization follows ranked candidates', () => {
     const rows = await handleVectorOnlySearch({ ...params, structuredFilters: undefined })
 
     expect(rows.map((row) => row.id)).toEqual(['selected'])
-    expect(dbChainMockFns.offset.mock.calls).toEqual([[0], [20], [0], [20]])
+    expect(dbChainMockFns.offset.mock.calls).toEqual([[0], [0], [20]])
     expect(getForConnectors).toHaveBeenCalledTimes(2)
     expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(2)
   })
@@ -638,6 +648,7 @@ describe('live repository authorization follows ranked candidates', () => {
       candidate(`probe-${index}`, 'allowed-source')
     )
     queueTableRows(schemaMock.embedding, probe)
+    queueTableRows(schemaMock.embedding, [{ id: 'far' }])
     queueTableRows(schemaMock.embedding, [
       { ...candidate('far', 'allowed-source'), distance: 0.7 },
       ...Array.from({ length: 19 }, (_, index) => candidate(`hidden-${index}`, 'allowed-source')),
@@ -662,7 +673,7 @@ describe('live repository authorization follows ranked candidates', () => {
     })
 
     expect(rows.map((row) => row.id)).toEqual(['nearer', 'near'])
-    expect(dbChainMockFns.offset.mock.calls).toEqual([[0], [20], [0]])
+    expect(dbChainMockFns.offset.mock.calls).toEqual([[0], [0]])
     expect(
       hasMockCondition(
         dbChainMockFns.where.mock.calls.at(-1)![0],