fix(search): bound vector retrieval under selective permissions

Author: icecrasher321

apps/sim/app/api/v1/knowledge/search/route.ts

@@ -269,6 +269,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
         queryVector: {
           vector: JSON.stringify(queryEmbeddingResult.embedding),
           dimensions: queryEmbeddingTarget!.dimensions,
+          model: queryEmbeddingTarget!.model,
         },
         structuredFilters: hasFilters ? structuredFilters : undefined,
       })

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

@@ -139,6 +139,7 @@ describe.each([384, 768, 1024, 1536, 3072] as const)(
         queryVector: {
           vector: JSON.stringify([1, ...Array<number>(dimensions - 1).fill(0)]),
           dimensions,
+          model: embeddingModel,
         },
         structuredFilters: [
           { tagSlot: 'tag1', fieldType: 'text', operator: 'eq', value: 'common' },

apps/sim/lib/knowledge/__integration__/scale.integration.ts

@@ -453,7 +453,7 @@ describe.skipIf(!enabled)('knowledge scale: isolated real PostgreSQL, no provide
           access: { kind: 'workspace', tokens: WORKSPACE_ACCESS_TOKENS },
           searchMode: 'hybrid',
           query: 'Orion',
-          queryVector: { vector, dimensions: DIMENSIONS },
+          queryVector: { vector, dimensions: DIMENSIONS, model: 'text-embedding-3-small' },
         })
       )
       expect(workspaceResults).toEqual([])
@@ -485,7 +485,11 @@ describe.skipIf(!enabled)('knowledge scale: isolated real PostgreSQL, no provide
                     access,
                     searchMode: mode,
                     query: 'Orion',
-                    queryVector: { vector, dimensions: DIMENSIONS },
+                    queryVector: {
+                      vector,
+                      dimensions: DIMENSIONS,
+                      model: 'text-embedding-3-small',
+                    },
                     structuredFilters: filters,
                   })
                 )

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

@@ -52,6 +52,7 @@ vi.hoisted(() => {
   if (process.env.KNOWLEDGE_SEARCH_PERFORMANCE_TEST === 'true') {
     Object.assign(process.env, {
       OPENAI_API_KEY: 'isolated-embedding-http-fixture',
+      GEMINI_API_KEY: 'isolated-gemini-http-fixture',
       CONFLUENCE_CLIENT_ID: 'isolated-confluence-fixture-client',
       CONFLUENCE_CLIENT_SECRET: 'isolated-confluence-fixture-secret',
     })
@@ -78,7 +79,8 @@ const fixtureSchema = z.object({
 })
 const reuseFile = enabled ? process.env.KNOWLEDGE_SEARCH_PERFORMANCE_REUSE_REPORT_FILE : undefined
 function readFixtureReport(file: string) {
-  if (statSync(file).size > 8 * 1024 * 1024) throw new Error('Fixture report exceeds 8 MiB')
+  /** Captured SQL plans include repeated high-dimensional query parameters. */
+  if (statSync(file).size > 64 * 1024 * 1024) throw new Error('Fixture report exceeds 64 MiB')
   return z
     .object({ fixture: fixtureSchema, unrelatedFixture: fixtureSchema })
     .parse(JSON.parse(readFileSync(file, 'utf8')))
@@ -146,14 +148,6 @@ const explainNodeSchema: z.ZodType<ExplainNode> = z.lazy(() =>
 )
 const explainSchema = z.array(z.object({ Plan: explainNodeSchema }).passthrough()).length(1)
 
-function usesVectorIndex(node: ExplainNode): boolean {
-  return (
-    node['Index Name'] === 'embedding_search_binary_hnsw_idx' ||
-    node['Index Name'] === 'embedding_vector_hnsw_idx' ||
-    (node.Plans?.some(usesVectorIndex) ?? false)
-  )
-}
-
 /** The ANN stage must not fetch full vectors, even for planner-added sort projections. */
 function assertCompactCandidates(node: ExplainNode) {
   expect(node['Relation Name']).not.toBe('embedding')
@@ -253,18 +247,27 @@ async function sample(label: string, run: () => ReturnType<typeof search>) {
   expect(captured.length).toBeLessThan(300)
   const searches = captured.filter(
     (item) =>
-      (item.query.includes('from "embedding"') || item.query.includes('from "embedding_search"')) &&
-      (item.query.includes('order by') || item.query.includes('limit'))
+      (item.query.includes('from "embedding"') ||
+        item.query.includes('FROM "embedding"') ||
+        item.query.includes('from "embedding_search"') ||
+        item.query.includes('FROM "embedding_search"')) &&
+      (item.query.includes('order by') ||
+        item.query.includes('limit') ||
+        item.query.includes('WITH visible_search_documents') ||
+        item.query.includes('WITH scored_search_candidates'))
   )
   const plans = []
   for (const query of searches) {
     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')
+      if (
+        query.query.includes('WITH visible_search_documents') ||
+        (query.query.includes('from "embedding_search"') && query.query.includes('order by'))
+      ) {
+        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')
       }
       return tx.unsafe(
         `EXPLAIN (ANALYZE, BUFFERS, VERBOSE, FORMAT JSON) ${query.query}`,
@@ -274,9 +277,11 @@ async function sample(label: string, run: () => ReturnType<typeof search>) {
     plans.push({
       kind: query.query.includes('keyword_rank')
         ? 'keyword'
-        : query.query.includes('binary_quantize')
+        : query.query.includes('WITH visible_search_documents') ||
+            (query.query.includes('from "embedding_search"') && query.query.includes('order by'))
           ? 'vector'
-          : query.query.includes('order by')
+          : query.query.includes('order by') ||
+              query.query.includes('WITH scored_search_candidates')
             ? 'rerank'
             : 'probe',
       query: query.query,
@@ -322,6 +327,29 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
           ? new Response(null, { status: 403 })
           : Response.json({ type: 'known', accountId: ids.aliceId })
       }
+      if (
+        url ===
+        'https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:batchEmbedContents'
+      ) {
+        const body = z
+          .object({
+            requests: z
+              .array(
+                z.object({
+                  content: z.object({ parts: z.array(z.object({ text: z.string() })).length(1) }),
+                })
+              )
+              .length(1),
+          })
+          .parse(JSON.parse(String(init?.body)))
+        embeddingCalls++
+        const text = body.requests[0].content.parts[0].text
+        const topic = Number(/^Topic (\d+) deployment$/.exec(text)?.[1] ?? 0)
+        return Response.json({
+          embeddings: [{ values: topicVector(topic) }],
+          usageMetadata: { promptTokenCount: 4 },
+        })
+      }
       if (url !== 'https://api.openai.com/v1/embeddings')
         throw new Error(`Unexpected outbound request in search fixture: ${new URL(url).origin}`)
       const body = z
@@ -355,7 +383,11 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
       }
       await db
         .update(knowledgeBase)
-        .set({ workspaceId: ids.workspaceId, organizationId: null })
+        .set({
+          workspaceId: ids.workspaceId,
+          organizationId: null,
+          embeddingModel: 'gemini-embedding-001',
+        })
         .where(eq(knowledgeBase.id, ids.knowledgeBaseId))
       await db
         .update(knowledgeConnector)
@@ -380,6 +412,11 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
     } else {
       await seedKnowledgeAclFixture(ids, { connectorType: 'google_drive' })
       await seedKnowledgeAclFixture(unrelated, { connectorType: 'google_drive' })
+      /** These arbitrary dense vectors are not trained for prefix shortening. */
+      await db
+        .update(knowledgeBase)
+        .set({ embeddingModel: 'gemini-embedding-001' })
+        .where(inArray(knowledgeBase.id, [ids.knowledgeBaseId, unrelated.knowledgeBaseId]))
       await db
         .update(knowledgeBase)
         .set({ isSearchIndex: true })
@@ -646,7 +683,7 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
       expect(plans.length).toBeGreaterThanOrEqual(2)
       const vectorPlans = plans.filter((plan) => plan.kind === 'vector')
       expect(vectorPlans).toHaveLength(1)
-      expect(usesVectorIndex(vectorPlans[0].plan[0].Plan)).toBe(true)
+      expect(vectorPlans[0].plan[0].Plan['Actual Rows']).toBeGreaterThan(0)
       assertCompactCandidates(vectorPlans[0].plan[0].Plan)
       expect(plans.some((plan) => plan.kind === 'rerank')).toBe(true)
       const rerank = plans.find((plan) => plan.kind === 'rerank')!
@@ -711,14 +748,65 @@ describe.skipIf(!enabled)('Assistant search latency on a realistic indexed corpu
     expect(assistantVector).toHaveLength(1)
     expect(dashboardVector[0].query).toBe(assistantVector[0].query)
     expect(dashboardVector[0].parameters).toEqual(assistantVector[0].parameters)
-    expect(usesVectorIndex(dashboardVector[0].plan[0].Plan)).toBe(true)
+    expect(dashboardVector[0].plan[0].Plan['Actual Rows']).toBeGreaterThan(0)
   }, 180_000)
 
   it('keeps inaccessible content out of an otherwise identical search', async () => {
     const { result } = await sample('denied', () => search(ids.bobId))
     expect(result.data.results).toEqual([])
   }, 180_000)
 
+  it('preserves recall when the nearest topic is mostly inaccessible within a broad permission scope', async () => {
+    const reader = `u:${ids.aliceId}@fixture.test`
+    /** Four consecutive topics share each document; hide 99% of the query's topic cluster. */
+    await db.execute(sql`UPDATE document
+      SET acl = ARRAY[${`u:${ids.bobId}@fixture.test`}]
+      WHERE knowledge_base_id = ${ids.knowledgeBaseId}
+        AND external_id::int % 8 = 0 AND external_id::int % 800 <> 0`)
+    try {
+      for (const surface of ['copilot', 'dashboard'] as const) {
+        const { result, plans, diagnostics } = await sample(
+          `filtered-neighborhood.${surface}`,
+          async () => {
+            if (surface === 'copilot') return search()
+            const data = await searchScopedKnowledge.execute({
+              principal: { kind: 'session', userId: ids.aliceId, sessionId: 'fixture-dashboard' },
+              input: {
+                workspaceId: ids.workspaceId,
+                query: 'Orion deployment',
+                topK: 15,
+                surface,
+              },
+            })
+            return resultSchema.parse({ success: true, data })
+          }
+        )
+        expect(diagnostics).toMatchObject({ retrievalStatus: 'complete', timedOutLegs: [] })
+        expect(result.data.results).toHaveLength(15)
+        const rerank = plans.find((plan) => plan.kind === 'rerank')!
+        expect(rerank).toBeDefined()
+        const actual = await db.$client.unsafe(rerank.query, rerank.parameters).values()
+        const expected = await db.execute<{ id: string }>(sql`SELECT e.id FROM embedding e
+          INNER JOIN document d ON d.id = e.document_id
+          WHERE e.knowledge_base_id = ${ids.knowledgeBaseId} AND e.enabled
+            AND d.acl @> ARRAY[${reader}]::text[]
+          ORDER BY (e.embedding <=> ${JSON.stringify(queryVector)}::vector) + 0, e.id
+          LIMIT ${actual.length}`)
+        expect(expected.length).toBeGreaterThan(0)
+        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.filtered-neighborhood.${surface}`] = { neighbors: expected.length, recall }
+        saveReport()
+      }
+    } finally {
+      await db
+        .update(document)
+        .set({ acl: [reader] })
+        .where(eq(document.knowledgeBaseId, ids.knowledgeBaseId))
+    }
+  }, 180_000)
+
   it('ranks a small permission scope by its bounded IDs without a corpus-wide vector probe', async () => {
     const documentIds = [0, 8, 16].map((index) => `${ids.workspaceId}-doc-${index}`)
     await db

apps/sim/lib/knowledge/access/predicate.postgres.test.ts

@@ -55,6 +55,7 @@ describe.runIf(Boolean(databaseUrl))('knowledge ACLs in PostgreSQL', () => {
     await client.unsafe(approvalMigration.replaceAll('"public".', `"${schemaName}".`))
     await client.unsafe(`
       ALTER TABLE knowledge_connector ADD COLUMN credential_id text,
+        ADD COLUMN access_rewrite_pending boolean NOT NULL DEFAULT false,
         ADD COLUMN source_config json NOT NULL DEFAULT '{}',
         ADD COLUMN credential_group_id text, ADD COLUMN credential_group_option_id text;
       ALTER TABLE credential ADD COLUMN revoked_at timestamp, ADD COLUMN credential_group_option_id text;

apps/sim/lib/knowledge/access/predicate.ts

@@ -224,13 +224,11 @@ function storedKnowledgeAccessCondition(
             (${knowledgeConnector.accessMode} = 'workspace' AND ${document.acl} = ARRAY['ws']::text[])
             OR (${document.acl} <> ARRAY['ws']::text[] AND (
             (${knowledgeConnector.accessMode} = 'admin' AND ${document.aclVerifiedAt} > ${cutoff})
-            OR (${knowledgeConnector.accessMode} = 'members' AND EXISTS (
-              SELECT 1 FROM ${knowledgeDocumentObservation}
+            OR (${knowledgeConnector.accessMode} = 'members' AND (${document.id}, ${document.connectorId}) IN (
+              SELECT ${knowledgeDocumentObservation.documentId}, ${knowledgeConnectorMember.connectorId} FROM ${knowledgeDocumentObservation}
               JOIN ${knowledgeConnectorMember}
                 ON ${knowledgeConnectorMember.id} = ${knowledgeDocumentObservation.memberId}
-              WHERE ${knowledgeDocumentObservation.documentId} = ${document.id}
-                AND ${knowledgeConnectorMember.connectorId} = ${document.connectorId}
-                AND ${knowledgeConnectorMember.status} = 'active'
+              WHERE ${knowledgeConnectorMember.status} = 'active'
                 AND ${knowledgeConnectorMember.subjectToken} = ANY(${tokens})
                 AND GREATEST(${knowledgeDocumentObservation.lastSeenAt}, ${knowledgeConnectorMember.memberSyncedThrough}) > ${cutoff}
             ))

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

@@ -420,6 +420,7 @@ const searchKnowledgeUseCase = defineAuthorizedKnowledgeUseCase({
           ? {
               vector: JSON.stringify(queryEmbedding?.embedding ?? null),
               dimensions: embeddingTarget!.dimensions,
+              model: embeddingTarget!.model,
             }
           : undefined,
         structuredFilters: structuredFilters.length > 0 ? structuredFilters : undefined,

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

@@ -49,6 +49,7 @@ export type SearchStage =
   | 'vector.ann'
   | 'vector.rerank'
   | 'vector.exact'
+  | 'vector.candidate_search'
 
 /** Fixed, content-free fields. Never pass queries, filters, document identities, SQL, or errors. */
 export interface SearchDiagnosticMetadata {
@@ -69,11 +70,14 @@ export interface SearchDiagnosticMetadata {
   searchMode?: 'hybrid' | 'vector'
   boostRecency?: boolean
   embeddingDimensions?: number
-  vectorRanking?: 'exact' | 'binary-rerank'
-  vectorCandidateStorage?: 'stored-binary'
+  vectorRanking?: 'exact' | 'candidate-rerank'
+  vectorCandidateStorage?: 'stored-halfvec'
+  vectorCandidateScan?: 'planned' | 'filtered'
   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