fix(knowledge): bound KB block vector retrieval (#7903)

* fix(knowledge): bound KB block vector retrieval

* fix(knowledge): align search fixtures with scoped probing

Author: icecrasher321

.github/workflows/test-build.yml

@@ -221,6 +221,7 @@ jobs:
           lib/knowledge/__integration__/search-source-progress.integration.ts
           lib/knowledge/__integration__/search-source-pagination.integration.ts
           lib/knowledge/__integration__/search-reference-batching.integration.ts
+          lib/knowledge/__integration__/kb-block-search.integration.ts
           lib/core/outbox/service.integration.ts
           lib/knowledge/__integration__/connector-upload.integration.ts
           lib/uploads/contexts/organization-logo/application.integration.ts

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

@@ -212,6 +212,10 @@ 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)])
 
       const results = await handleTagAndVectorSearch({
@@ -226,9 +230,11 @@ describe('Knowledge Search Utils', () => {
       })
 
       expect(results.map((row) => row.id)).toEqual(['first', 'second'])
-      expect(dbChainMockFns.select).toHaveBeenCalledTimes(2)
+      expect(dbChainMockFns.select).toHaveBeenCalledTimes(3)
       expect(dbChainMockFns.as).toHaveBeenCalledWith('ranked_embeddings')
-      expect(dbChainMockFns.select.mock.calls[0][0]).toHaveProperty('distance')
+      expect(Object.keys(dbChainMockFns.select.mock.calls[0][0])).toEqual(['id'])
+      expect(dbChainMockFns.limit).toHaveBeenNthCalledWith(1, 201)
+      expect(dbChainMockFns.select.mock.calls[1][0]).toHaveProperty('distance')
       expect(dbChainMockFns.limit).toHaveBeenCalledWith(2)
     })
 
@@ -536,6 +542,7 @@ describe('Knowledge Search Utils', () => {
     })
 
     it('runs a single retrieval leg in vector mode', async () => {
+      queueTableRows(schemaMock.embedding, [{ id: 'vector-hit' }])
       queueTableRows(schemaMock.embedding, [makeResult('vector-hit')])
 
       const results = await executeKnowledgeSearch({
@@ -548,20 +555,19 @@ describe('Knowledge Search Utils', () => {
       })
 
       expect(results.map((r) => r.id)).toEqual(['vector-hit'])
-      expect(dbChainMockFns.select).toHaveBeenCalledTimes(2)
+      expect(dbChainMockFns.select).toHaveBeenCalledTimes(3)
       expect(dbChainMockFns.as).toHaveBeenCalledWith('ranked_embeddings')
     })
 
     it('runs both legs and fuses them in hybrid mode', async () => {
       /**
-       * Chains dequeue in creation order. Hybrid legs over-fetch past the
-       * plain scan's candidate pool, so the vector leg opens its transaction
-       * and applies the scan settings before selecting: the keyword ranking
-       * pass is built first, then the vector select, then hydration.
+       * Chains dequeue in creation order: keyword ranking, the budgeted vector
+       * probe, keyword hydration, then vector ranking and hydration in one query.
        */
       queueTableRows(schemaMock.embedding, [{ id: 'keyword-hit', keywordRank: 0.9 }])
-      queueTableRows(schemaMock.embedding, [makeResult('vector-hit')])
+      queueTableRows(schemaMock.embedding, [{ id: 'vector-hit' }])
       queueTableRows(schemaMock.embedding, [makeResult('keyword-hit')])
+      queueTableRows(schemaMock.embedding, [makeResult('vector-hit')])
 
       const results = await executeKnowledgeSearch({
         knowledgeBaseIds: ['kb-123'],
@@ -573,39 +579,31 @@ 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(5)
     })
 
-    it('falls back to vector results when the keyword leg fails', async () => {
+    it('propagates unexpected keyword errors after the vector leg finishes', async () => {
       /** The failing ranking chain is still built first and takes the first queued set. */
       queueTableRows(schemaMock.embedding, [{ id: 'never-ranked', keywordRank: 0 }])
+      queueTableRows(schemaMock.embedding, [{ id: 'vector-hit' }])
       queueTableRows(schemaMock.embedding, [makeResult('vector-hit')])
 
-      /**
-       * Both legs share one `orderBy` spy, so target the keyword leg by its
-       * ranking expression. Calling the untouched spy first captures the
-       * sentinel that tells the mock to build its normal chain, which the
-       * vector leg still needs.
-       */
-      const chainDefault = dbChainMockFns.orderBy()
-      dbChainMockFns.orderBy.mockImplementation((fragment: unknown) => {
-        const text = (fragment as { strings?: string[] })?.strings?.join('') ?? ''
-        if (text.includes('ts_rank_cd')) {
-          throw new Error('tsquery blew up')
-        }
-        return chainDefault
-      })
-
-      const results = await executeKnowledgeSearch({
-        knowledgeBaseIds: ['kb-123'],
-        access: WORKSPACE_ACCESS_SCOPE,
-        topK: 10,
-        searchMode: 'hybrid',
-        query: 'PROJ-1234',
-        queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+      const failure = new Error('tsquery failed')
+      dbChainMockFns.orderBy.mockImplementationOnce(() => {
+        throw failure
       })
 
-      expect(results.map((r) => r.id)).toEqual(['vector-hit'])
+      await expect(
+        executeKnowledgeSearch({
+          knowledgeBaseIds: ['kb-123'],
+          access: WORKSPACE_ACCESS_SCOPE,
+          topK: 10,
+          searchMode: 'hybrid',
+          query: 'PROJ-1234',
+          queryVector: JSON.stringify([0.1, 0.2, 0.3]),
+        })
+      ).rejects.toBe(failure)
+      expect(dbChainMockFns.as).toHaveBeenCalledWith('ranked_embeddings')
     })
 
     it('skips both query legs when only tag filters are provided', async () => {

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

@@ -0,0 +1,135 @@
+/** KB block retrieval against disposable PostgreSQL, using a workspace API-key identity. */
+import type { Principal } from '@sim/auth/principal'
+import { db } from '@sim/db'
+import { document, embedding, knowledgeBase, organization, user, workspace } from '@sim/db/schema'
+import { generateId } from '@sim/utils/id'
+import { eq, inArray } from 'drizzle-orm'
+import { afterAll, beforeAll, describe, expect, it } from 'vitest'
+import {
+  createKnowledgeAclFixtureIds,
+  seedKnowledgeAclFixture,
+} from '@/lib/knowledge/__integration__/seed-source-access-fixture'
+import { createKnowledgeAccessProvider } from '@/lib/knowledge/access/scope'
+import { retrieveKnowledgeSearch } from '@/lib/knowledge/search/queries'
+import { embeddingVectorValues } from '@/lib/knowledge/vector-columns'
+
+describe('API-key KB block fan-out', () => {
+  const ids = createKnowledgeAclFixtureIds()
+  const bases = Array.from({ length: 18 }, () => ({
+    id: generateId(),
+    visible: generateId(),
+    denied: generateId(),
+    excluded: generateId(),
+  }))
+  const principal: Principal = {
+    kind: 'workspace_api_key',
+    workspaceId: ids.workspaceId,
+    keyId: 'fixture-key',
+  }
+  const vector = [1, ...Array<number>(1535).fill(0)]
+  const queryVector = {
+    vector: JSON.stringify(vector),
+    dimensions: 1536 as const,
+    model: 'text-embedding-3-small',
+  }
+
+  beforeAll(async () => {
+    await seedKnowledgeAclFixture(ids, { connectorType: 'google_drive' })
+    await db.insert(knowledgeBase).values(
+      bases.map((base, index) => ({
+        id: base.id,
+        userId: ids.aliceId,
+        workspaceId: ids.workspaceId,
+        name: `KB block ${index}`,
+      }))
+    )
+    await db.insert(document).values(
+      bases.flatMap((base) =>
+        (['visible', 'denied', 'excluded'] as const).map((kind) => ({
+          id: base[kind],
+          knowledgeBaseId: base.id,
+          filename: kind,
+          fileUrl: `https://fixture.invalid/${base[kind]}`,
+          fileSize: 12,
+          mimeType: 'text/plain',
+          processingStatus: 'completed',
+          acl: kind === 'denied' ? [`u:${ids.aliceId}@fixture.test`] : ['ws'],
+          userExcluded: kind === 'excluded',
+        }))
+      )
+    )
+    await db.insert(embedding).values(
+      bases.flatMap((base) =>
+        (['visible', 'denied', 'excluded'] as const).map((kind) => ({
+          id: generateId(),
+          documentId: base[kind],
+          knowledgeBaseId: base.id,
+          chunkIndex: 0,
+          chunkHash: base[kind],
+          content: `Fixture policy ${kind}`,
+          contentLength: 24,
+          tokenCount: 5,
+          startOffset: 0,
+          endOffset: 24,
+          tag1: 'policy',
+          ...embeddingVectorValues(1536, vector),
+        }))
+      )
+    )
+  })
+
+  afterAll(async () => {
+    await db.delete(workspace).where(eq(workspace.id, ids.workspaceId))
+    await db.delete(organization).where(eq(organization.id, ids.organizationId))
+    await db.delete(user).where(inArray(user.id, [ids.aliceId, ids.bobId]))
+    await db.$client.end()
+  })
+
+  it.each([false, true])(
+    'completes 18 concurrent KB searches with access checks intact (tag filter: %s)',
+    async (withTags) => {
+      const previousDebug = db.$client.options.debug
+      const statements: string[] = []
+      db.$client.options.debug = (_connection, query) => {
+        if (statements.length < 250) statements.push(query)
+      }
+      try {
+        const results = await Promise.all(
+          bases.map(async (base) => {
+            const accessProvider = createKnowledgeAccessProvider(principal, {
+              workspaceId: ids.workspaceId,
+              knowledgeBaseIds: [base.id],
+            })
+            const access = await accessProvider.get()
+            expect(access.kind).toBe('workspace')
+            return retrieveKnowledgeSearch({
+              knowledgeBaseIds: [base.id],
+              topK: 2,
+              access,
+              accessProvider,
+              searchMode: 'vector',
+              query: 'Find the fixture policy',
+              queryVector,
+              ...(withTags && {
+                structuredFilters: [
+                  { tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'policy' },
+                ],
+              }),
+            })
+          })
+        )
+        for (const [index, result] of results.entries()) {
+          expect(result.retrieval).toEqual({ status: 'complete', timedOutLegs: [] })
+          expect(result.rows.map((row) => row.documentId)).toEqual([bases[index].visible])
+          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(36)
+        expect(statements.filter((query) => query.includes('+ 0'))).toHaveLength(18)
+        expect(statements.some((query) => query.includes('hnsw.iterative_scan'))).toBe(false)
+      } finally {
+        db.$client.options.debug = previousDebug
+      }
+    }
+  )
+})