improvement(knowledge): exclude a denied source through its documents while the projection is unfilled

Author: waleedlatif1

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

@@ -2076,6 +2076,48 @@ describe('permitted-document planner', () => {
     expect(JSON.stringify(pages[1].params)).toContain('"b"')
   })
 
+  it('excludes a denied source through its documents while the projection is unfilled', async () => {
+    queueTableRows(schemaMock.knowledgeConnector, [
+      {
+        id: 'gated-src',
+        accessMode: 'admin',
+        connectorType: 'confluence',
+        githubRepository: false,
+      },
+    ])
+    dbChainMockFns.execute.mockImplementation(async (query) => {
+      const statement = render(query).sql
+      /** The fill has not reached every row, so a denied source cannot be read off the row. */
+      if (statement.includes('AS unfilled')) return [{ unfilled: true }]
+      const rebuilt = JSON.stringify(query).includes('NOT EXISTS')
+      if (isPageStatement(statement))
+        return JSON.stringify(render(query).params).includes('"b"')
+          ? [hit('b', 'other-src')]
+          : [hit('a', 'gated-src')]
+      if (isWalk(statement))
+        return Array.from({ length: 400 }, (_, i) => ({
+          id: i === 0 ? (rebuilt ? 'b' : 'a') : `w-${i}`,
+          distance: 0.1,
+        }))
+      return []
+    })
+    queueTableRows(schemaMock.embedding, [])
+    queueTableRows(schemaMock.embedding, [hit('b', 'other-src')])
+    const getForConnectors = vi.fn<KnowledgeAccessProvider['getForConnectors']>(async () => reader)
+    const result = await retrieveKnowledgeSearch({
+      ...liveSearch,
+      searchMode: 'vector',
+      access: reader,
+      accessProvider: { ...provider, getForConnectors },
+    })
+    expect(result.rows.map((row) => row.id)).toEqual(['b'])
+    const walks = statements().filter((query) => isWalk(query.sql))
+    expect(walks).toHaveLength(2)
+    expect(JSON.stringify(walks[1])).toContain('NOT EXISTS')
+    expect(JSON.stringify(walks[1])).toContain('gated-src')
+    forgetProjectionFilled()
+  })
+
   it('hands back the unread slices of a page a denied source made it rebuild', async () => {
     queueTableRows(schemaMock.knowledgeConnector, [
       {

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

@@ -1722,12 +1722,22 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
         return { candidates, nextOffset: offset + candidates.length }
       }
       if (params.filters?.documentIds?.length) return exactPage()
+      const plan = params.access.kind === 'user' ? params.accessPlan : undefined
+      const filled = plan
+        ? ((await projectionFilled.fetch('embedding_search', {
+            context: { budget: params.budget, stage: 'vector.projection_filled' },
+          })) ?? false)
+        : false
       /**
        * A source the caller turned out not to hold is left out where the pool is built: the pool
        * is the page's order now, so a denied source's chunks would otherwise keep their slots.
+       * The row's mirrored source decides it once the fill is complete; until then a row the fill
+       * has not reached carries no source, so its document is asked instead.
        */
       const excludedOnRow = excludedSources.length
-        ? sql`(${embeddingSearch.connectorId} IS NULL OR NOT (${embeddingSearch.connectorId} = ANY(${textArrayLiteral([...excludedSources])})))`
+        ? filled
+          ? sql`(${embeddingSearch.connectorId} IS NULL OR NOT (${embeddingSearch.connectorId} = ANY(${textArrayLiteral([...excludedSources])})))`
+          : sql`NOT EXISTS (SELECT 1 FROM ${document} WHERE ${document.id} = ${embeddingSearch.documentId} AND ${document.connectorId} = ANY(${textArrayLiteral([...excludedSources])}))`
         : undefined
       const needed = offset + limit
       if (
@@ -1776,12 +1786,6 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
           )
         }
         let selected: Array<{ id: string }>
-        const plan = params.access.kind === 'user' ? params.accessPlan : undefined
-        const filled = plan
-          ? ((await projectionFilled.fetch('embedding_search', {
-              context: { budget: params.budget, stage: 'vector.projection_filled' },
-            })) ?? false)
-          : false
         /**
          * A source the caller is a member of that has its own index is walked on its own, which
          * beats ranking it exactly once it is large enough to have earned that index.