fix(knowledge): refill after a denied gated source, reserve the index build's session, lock the backfill's documents

- once a live proof shows the caller does not hold a gated source, rebuild the candidate pages without it, so its candidates no longer hold the slots of sources the caller can read
- run the source index build on one reserved connection, so its memory setting and its reset apply to the session that builds
- share-lock a page's documents before copying their source, so a detachment in flight waits for the page and then fans its own change out through the trigger

Author: waleedlatif1

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

@@ -1660,6 +1660,48 @@ describe('permitted-document planner', () => {
     expect(getForConnectors).not.toHaveBeenCalled()
   })
 
+  it('rebuilds the pages without a gated source the caller turns out not to hold', async () => {
+    queueTableRows(schemaMock.knowledgeConnector, [
+      {
+        id: 'gated-src',
+        accessMode: 'admin',
+        connectorType: 'confluence',
+        githubRepository: false,
+      },
+    ])
+    /**
+     * The first pool is filled by the gated source alone; only a pool built without it — the
+     * exclusion carries the source id into the statement — reaches the accessible candidate.
+     */
+    dbChainMockFns.execute.mockImplementation(async (query) => {
+      const statement = render(query).sql
+      /** The exclusion is the only clause that negates a connector membership. */
+      const rebuilt = JSON.stringify(query).includes('OR NOT (')
+      if (statement.includes('WITH scored_search_candidates'))
+        return rebuilt ? [hit('b', 'other-src')] : [hit('a', 'gated-src')]
+      if (isExactRanking(statement)) return [{ id: rebuilt ? 'b' : 'a' }]
+      if (isProbeStatement(statement))
+        return [{ id: 'doc-a', connectorId: 'gated-src', saturated: false }]
+      return []
+    })
+    queueTableRows(schemaMock.embedding, [])
+    queueTableRows(schemaMock.embedding, [hit('b', 'other-src')])
+    /** No grants come back, so the gated source is denied. */
+    const getForConnectors = vi.fn<KnowledgeAccessProvider['getForConnectors']>(async () => reader)
+    const result = await retrieveKnowledgeSearch({
+      ...liveSearch,
+      searchMode: 'vector',
+      access: reader,
+      accessProvider: { ...provider, getForConnectors },
+    })
+    expect(getForConnectors).toHaveBeenCalledOnce()
+    expect(result.rows.map((row) => row.id)).toEqual(['b'])
+    const reranks = statements().filter((query) => query.sql.includes('scored_search_candidates'))
+    expect(reranks).toHaveLength(2)
+    expect(JSON.stringify(reranks[0])).not.toContain('OR NOT (')
+    expect(JSON.stringify(reranks[1])).toContain('OR NOT (')
+  })
+
   it('asks a live source for its grants once, when a candidate of its own is read', async () => {
     queueTableRows(schemaMock.knowledgeConnector, [
       {

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

@@ -596,7 +596,8 @@ const SEARCH_READ_CANDIDATE_FIELDS = {
  */
 export interface LiveSourceAccess {
   gates: (connectorId: string) => boolean
-  resolve: () => Promise<KnowledgeAccessScope>
+  /** The caller's scope with its grants, and the gated sources those grants do not cover. */
+  resolve: () => Promise<{ access: KnowledgeAccessScope; denied: ReadonlySet<string> }>
 }
 
 /** Binds a search's gated sources to one memoized resolution of the caller's grants. */
@@ -608,7 +609,7 @@ export function liveSourceAccessFor(
 ): LiveSourceAccess | undefined {
   const gated = new Set(plan?.connectors.liveProofRequired ?? [])
   if (!accessProvider || gated.size === 0) return undefined
-  let pending: Promise<KnowledgeAccessScope> | undefined
+  let pending: Promise<{ access: KnowledgeAccessScope; denied: ReadonlySet<string> }> | undefined
   /** The provider authorizes connectors in bounded pages, so a wide scope resolves page by page. */
   const pages: string[][] = []
   for (const id of gated) {
@@ -624,7 +625,6 @@ export function liveSourceAccessFor(
           accessProvider.getForConnectors(page, signal)
         )
         const [first] = scopes
-        if (scopes.length === 1 || first.kind !== 'user') return first
         const githubInstallationGrants: GitHubInstallationReadGrant[] = []
         const confluenceSiteGrants: ConfluenceSiteReadGrant[] = []
         for (const scope of scopes) {
@@ -633,7 +633,16 @@ export function liveSourceAccessFor(
             githubInstallationGrants.push(...scope.githubInstallationGrants)
           if (scope.confluenceSiteGrants) confluenceSiteGrants.push(...scope.confluenceSiteGrants)
         }
-        return { ...first, githubInstallationGrants, confluenceSiteGrants }
+        const granted = new Set([
+          ...githubInstallationGrants.map((grant) => grant.connectorId),
+          ...confluenceSiteGrants.map((grant) => grant.connectorId),
+        ])
+        const denied = new Set([...gated].filter((id) => !granted.has(id)))
+        const merged =
+          first.kind === 'user'
+            ? { ...first, githubInstallationGrants, confluenceSiteGrants }
+            : first
+        return { access: merged, denied }
       })
       return pending
     },
@@ -655,7 +664,11 @@ async function selectAuthorizedSearchResults(input: {
   signal?: AbortSignal
   budget?: SearchBudget
   topK: number
-  selectPage: (limit: number, offset: number) => Promise<SearchReadCandidatePage>
+  selectPage: (
+    limit: number,
+    offset: number,
+    excludedSources: readonly string[]
+  ) => Promise<SearchReadCandidatePage>
   compareResults?: (a: SearchResult, b: SearchResult) => number
   hydrate: (ids: string[], access: KnowledgeAccessScope) => Promise<SearchResult[]>
   liveSourceAccess?: LiveSourceAccess
@@ -664,6 +677,8 @@ async function selectAuthorizedSearchResults(input: {
   const pageSize = Math.min(AUTHORIZED_SEARCH_PAGE_SIZE, Math.max(input.topK, 20))
   const results = new Map<string, SearchResult>()
   const considered = new Set<string>()
+  /** Gated sources the caller turned out not to hold: left out of every page once that is known. */
+  let excluded: ReadonlySet<string> = new Set()
   let scanned = 0
   let offset = 0
   try {
@@ -675,7 +690,7 @@ async function selectAuthorizedSearchResults(input: {
       input.signal?.throwIfAborted()
       input.budget?.remaining()
       const page = await measureSearchStage(`${input.leg}.candidates`, () =>
-        input.selectPage(pageSize, offset)
+        input.selectPage(pageSize, offset, [...excluded])
       )
       if (!page.candidates.length) break
       scanned += page.candidates.length
@@ -691,11 +706,25 @@ async function selectAuthorizedSearchResults(input: {
        * its own reaches this page, and then once for the whole search: a scope that ranks none
        * of them — most scopes — never asks, and one that ranks many asks once.
        */
-      const access = candidates.some(
-        (candidate) => candidate.connectorId && input.liveSourceAccess?.gates(candidate.connectorId)
-      )
-        ? await input.liveSourceAccess!.resolve()
-        : input.access
+      const proof = input.liveSourceAccess
+      const gatedOnPage =
+        proof !== undefined &&
+        candidates.some((candidate) => candidate.connectorId && proof.gates(candidate.connectorId))
+      let access = input.access
+      let refill = false
+      if (gatedOnPage && proof) {
+        const resolved = await proof.resolve()
+        access = resolved.access
+        /**
+         * A denied source's candidates cannot hydrate, yet they took the slots of sources the
+         * caller does hold. Once the denial is known the pages are rebuilt without that source,
+         * from the start; the ids already seen are not read twice.
+         */
+        if (resolved.denied.size > excluded.size) {
+          excluded = resolved.denied
+          refill = true
+        }
+      }
       const hydrated = await measureSearchStage(`${input.leg}.hydration`, () =>
         input.hydrate(
           candidates.map((candidate) => candidate.id),
@@ -713,6 +742,10 @@ async function selectAuthorizedSearchResults(input: {
         results.clear()
         for (const row of ranked) results.set(row.id, row)
       }
+      if (refill) {
+        offset = 0
+        continue
+      }
       /** A short page is the end of the candidates, whether or not they were reordered. */
       if (page.candidates.length < pageSize) break
     }
@@ -723,6 +756,13 @@ async function selectAuthorizedSearchResults(input: {
   return [...results.values()]
 }
 
+/** Keeps the candidates of sources the caller turned out not to hold out of a page. */
+function excludeSearchSources(sourceIds: readonly string[]): SQL | undefined {
+  return sourceIds.length
+    ? sql`(${document.connectorId} IS NULL OR NOT (${inArray(document.connectorId, [...sourceIds])}))`
+    : undefined
+}
+
 /**
  * Loads the content of candidates that survived ranking, under the read predicate.
  *
@@ -799,7 +839,7 @@ export async function handleTagOnlySearch(params: SearchParams): Promise<SearchR
       signal: params.signal,
       budget: params.budget,
       topK,
-      selectPage: async (limit, offset) => {
+      selectPage: async (limit, offset, excludedSources) => {
         const candidates = await runSearchQuery(params.budget, 'tags.sql', (executor) =>
           executor
             .select(SEARCH_READ_CANDIDATE_FIELDS)
@@ -812,7 +852,8 @@ export async function handleTagOnlySearch(params: SearchParams): Promise<SearchR
                   access,
                   params.filters,
                   candidateAccessCondition(access, params.accessPlan)
-                )
+                ),
+                excludeSearchSources(excludedSources)
               )
             )
             .orderBy(embedding.id)
@@ -1277,7 +1318,7 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
    * refill reuses the pool it already has. Excluding another source is the only thing that
    * changes which candidates belong in it, and that resets the offset to zero anyway.
    */
-  let candidatePool: Array<{ id: string }> | undefined
+  let candidatePool: { excludedKey: string; ids: Array<{ id: string }> } | undefined
   return selectAuthorizedSearchResults({
     leg: 'vector',
     access: params.access,
@@ -1287,9 +1328,11 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
     budget: params.budget,
     topK: params.topK,
     compareResults: (a, b) => a.distance - b.distance,
-    selectPage: async (limit, offset) => {
+    selectPage: async (limit, offset, excludedSources) => {
+      const excludedKey = [...excludedSources].sort().join(',')
       const visibility = [
         ...getVisibilityConditions(params.access, params.filters, candidateAccess),
+        excludeSearchSources(excludedSources),
       ]
       const candidateDocumentVisibility = [
         ...candidateDocumentConditions(
@@ -1298,6 +1341,7 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
           params.filters,
           candidateAccess
         ),
+        excludeSearchSources(excludedSources),
       ]
       /** Explicit document IDs are already a bounded scope, and retain exhaustive ordering. */
       const exactPage = async () => {
@@ -1315,7 +1359,7 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
         return { candidates, nextOffset: offset + candidates.length }
       }
       if (params.filters?.documentIds?.length) return exactPage()
-      if (!candidatePool) {
+      if (candidatePool?.excludedKey !== excludedKey) {
         annotateSearchDiagnostics({
           vectorRanking: 'candidate-rerank',
           vectorCandidateStorage: 'stored-halfvec',
@@ -1421,13 +1465,13 @@ async function selectVectorResults(params: SearchParams): Promise<SearchResult[]
             }
           }
         }
-        candidatePool = selected
+        candidatePool = { excludedKey, ids: selected }
         annotateSearchDiagnostics({
           vectorCandidateCount: selected.length,
           vectorCandidateScan: selected.length < candidateLimit ? 'underfilled' : 'planned',
         })
       }
-      const identities = candidatePool
+      const identities = candidatePool.ids
       if (!identities.length) return { candidates: [], nextOffset: offset }
       /** Score each bounded candidate once; sorting the materialized scalar cannot invoke HNSW again. */
       const page = await runSearchQuery(params.budget, 'vector.rerank', (executor) =>
@@ -1555,14 +1599,15 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
         ? { keywordRanking: tinQuery ? 'tin' : 'gin' }
         : {}),
     })
-    const documentConditions = () =>
+    const documentConditions = (excludedSources: readonly string[]) =>
       and(
         ...candidateDocumentConditions(
           knowledgeBaseIds,
           access,
           params.filters,
           candidateAccessCondition(access, params.accessPlan)
-        )
+        ),
+        excludeSearchSources(excludedSources)
       )
     /**
      * One page from the top of Tin's ranking. The window of ranked chunks widens while too few of
@@ -1572,7 +1617,8 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
     const selectTinPage = async (
       scopedQuery: SQL,
       limit: number,
-      offset: number
+      offset: number,
+      excludedSources: readonly string[]
     ): Promise<SearchReadCandidatePage | null> => {
       for (const window of TIN_KEYWORD_WINDOWS) {
         if (window < offset + limit) continue
@@ -1590,7 +1636,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
               SELECT ${document.id} AS id FROM ${document}
               WHERE ${and(
                 sql`${document.id} = ANY (ARRAY(SELECT document_id FROM ranked_tin_chunks))`,
-                documentConditions()
+                documentConditions(excludedSources)
               )}
             ), page AS (
               SELECT ranked_tin_chunks.id, ${document.id} AS "documentId",
@@ -1635,7 +1681,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
       signal: params.signal,
       budget: params.budget,
       topK,
-      selectPage: async (limit, offset) => {
+      selectPage: async (limit, offset, excludedSources) => {
         /**
          * A bounded permitted set confines matching to the chunks the caller may read, so a term
          * common across the index is ranked only where it can surface. The visibility CTE below
@@ -1647,7 +1693,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
             : undefined
         if (permittedIds?.length === 0) return { candidates: [], nextOffset: offset }
         if (tinScope && !permittedIds) {
-          const tinPage = await selectTinPage(tinScope, limit, offset)
+          const tinPage = await selectTinPage(tinScope, limit, offset, excludedSources)
           if (tinPage) return tinPage
           annotateSearchDiagnostics({ keywordRanking: 'gin' })
         }
@@ -1694,7 +1740,7 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
               SELECT ${document.id} AS id FROM ${document}
               WHERE ${and(
                 sql`${document.id} = ANY (ARRAY(SELECT document_id FROM matched_keyword_chunks))`,
-                documentConditions()
+                documentConditions(excludedSources)
               )}
             ), ranked_keyword_candidates AS MATERIALIZED (
               SELECT matched_keyword_chunks.id, matched_keyword_chunks.document_id,

apps/sim/lib/knowledge/search/source-vector-indexes.test.ts

@@ -1,6 +1,7 @@
 /**
  * @vitest-environment node
  */
+import { db } from '@sim/db'
 import { dbChainMockFns, resetDbChainMock } from '@sim/testing'
 import { beforeEach, describe, expect, it } from 'vitest'
 import {
@@ -27,6 +28,18 @@ describe('source vector indexes', () => {
       if (text.includes('sampled')) return [{ column: 'vector_512' }]
       return []
     })
+    /** The build reserves one connection; its statements are recorded with the rest. */
+    Object.assign(db, {
+      $client: {
+        reserve: async () => ({
+          unsafe: async (text: string) => {
+            statements.push(text)
+            return []
+          },
+          release: () => undefined,
+        }),
+      },
+    })
     /** Warms the catalog cache with this case's state, so a build decision is not a stale read. */
     expect((await indexedVectorSources()).size).toBe(indexed.length)
     statements = []

apps/sim/lib/knowledge/search/source-vector-indexes.ts

@@ -103,14 +103,18 @@ export async function ensureSourceVectorIndex(connectorId: string): Promise<bool
   if (!column) return false
   const name = indexName(connectorId)
   const startedAt = Date.now()
+  /**
+   * The memory setting and the build must share a session, and `CONCURRENTLY` forbids a
+   * transaction, so one connection is reserved from the pool for the whole build and its setting
+   * is reset before the connection goes back.
+   */
+  const session = await db.$client.reserve()
   try {
     await dropInvalidIndex(name)
-    await db.execute(sql`SET maintenance_work_mem = '2GB'`)
-    await db.execute(
-      sql.raw(`CREATE INDEX CONCURRENTLY "${name}" ON embedding_search
+    await session.unsafe("SET maintenance_work_mem = '2GB'")
+    await session.unsafe(`CREATE INDEX CONCURRENTLY "${name}" ON embedding_search
         USING hnsw (${column} halfvec_cosine_ops) WITH (m = 16, ef_construction = 64)
         WHERE connector_id = '${connectorId}' AND enabled`)
-    )
     logger.info('Built a source vector index', {
       connectorId,
       documents,
@@ -123,7 +127,8 @@ export async function ensureSourceVectorIndex(connectorId: string): Promise<bool
     await dropInvalidIndex(name)
     return false
   } finally {
-    await db.execute(sql`RESET maintenance_work_mem`)
+    await session.unsafe('RESET maintenance_work_mem').catch(() => undefined)
+    session.release()
   }
 }