improvement(knowledge): let an on-row walk run to its cap, widen a narrow reader's keyword window stepwise, apply scan settings with the deadline

- an on-row walk keeps walking, up to a 100k-tuple cap, until its limit is met; the separate wider walk and its diagnostic are gone
- the on-row predicate tests the mirrored ACL alone — a member's source is no longer admitted whole, so a document re-owned after its chunk was mirrored is refused at the row
- a reach count that ran out of time decides that search only; it is not remembered
- a narrow reader ranks the narrowest keyword window first and widens to the wide one only when the page is short; resolved scopes leave the widest window short instead of ranking every match
- the HNSW scan settings ride in the deadline statement, one round trip fewer per vector statement

Author: waleedlatif1

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

@@ -600,13 +600,28 @@ describe.runIf(Boolean(databaseUrl))('knowledge ACLs in PostgreSQL', () => {
     const onRow = new PgDialect().sqlToQuery(
       projectionCandidateAccessCondition(schema.embeddingSearch, scope, plan)
     )
-    for (const id of [...cases.map(([documentId]) => documentId), 'upload-doc']) {
+    const onRowAdmits = async (id: string) => {
       const rows = await connection.unsafe(
         `SELECT 1 FROM embedding_search WHERE ${onRow.sql} AND document_id = $${onRow.params.length + 1}`,
         [...(onRow.params as string[]), id]
       )
-      if (await admits(perRow, id)) expect([id, rows.length > 0]).toEqual([id, true])
+      return rows.length > 0
+    }
+    for (const id of [...cases.map(([documentId]) => documentId), 'upload-doc']) {
+      if (await admits(perRow, id)) expect([id, await onRowAdmits(id)]).toEqual([id, true])
     }
+    /**
+     * A source the caller is a member of still holds documents that name other members only; the
+     * mirrored ACL keeps those out of the ranking rather than leaving them for hydration to drop.
+     */
+    await connection.unsafe(
+      "INSERT INTO document(id, connector_id, acl, acl_verified_at) VALUES ('members-other', 'members', ARRAY['s:slack:-:bob'], statement_timestamp())"
+    )
+    await connection.unsafe(
+      "INSERT INTO embedding_search SELECT id || '-chunk', id, connector_id, acl FROM document WHERE id = 'members-other'"
+    )
+    expect(await admits(perRow, 'members-other')).toBe(false)
+    expect(await onRowAdmits('members-other')).toBe(false)
     /**
      * Candidate ranking defers the live source proof, as the per-row candidate predicate does: a
      * caller holds those grants only after authorization, so applying the clause during ranking

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

@@ -307,10 +307,11 @@ export function knowledgeCandidateAccessConditionForConnectors(
  * resolved: `connectorId` and `acl` are mirrored there from the document, so a walk or a keyword
  * window decides readability on the row it scores instead of joining `document` per candidate.
  *
- * It admits a superset of the document predicate, never a subset: a member's source is admitted
- * whole (the observation that vouches for each document is checked at hydration), and requirement
- * clauses live on the document. Both are refused there, under the full predicate, before content is
- * returned — this predicate only decides what is worth ranking.
+ * It admits a superset of the document predicate, never a subset: a mirrored ACL names the members
+ * who observe a document, so overlap with the caller's tokens is the per-row test without the
+ * observation's freshness, and requirement clauses live on the document. Both are refused there,
+ * under the full predicate, before content is returned — this predicate only decides what is worth
+ * ranking.
  */
 export function projectionCandidateAccessCondition(
   projection: { connectorId: AnyPgColumn | SQL; acl: AnyPgColumn | SQL },
@@ -329,11 +330,8 @@ export function projectionCandidateAccessCondition(
     ...plan.connectors.admin,
     ...plan.connectors.members,
   ]
-  return sql`(
-    ${inSources(plan.memberSources)}
-    OR (${projection.acl} && ${tokens}
-      AND (${projection.connectorId} IS NULL OR ${inSources(mirrored)}))
-  )`
+  return sql`(${projection.acl} && ${tokens}
+    AND (${projection.connectorId} IS NULL OR ${inSources(mirrored)}))`
 }
 
 /**

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

@@ -1,6 +1,6 @@
 import { db } from '@sim/db'
 import { getPostgresErrorCode } from '@sim/utils/errors'
-import { sql } from 'drizzle-orm'
+import { type SQL, sql } from 'drizzle-orm'
 import type { DbTransaction } from '@/lib/db/types'
 import {
   measureSearchStage,
@@ -67,7 +67,8 @@ export class SearchBudget {
    */
   async query<T>(
     stage: SearchStage,
-    run: (executor: SearchExecutor) => PromiseLike<T>
+    run: (executor: SearchExecutor) => PromiseLike<T>,
+    settings: readonly SQL[] = []
   ): Promise<T> {
     const started = performance.now()
     const remaining = this.remaining()
@@ -97,9 +98,16 @@ export class SearchBudget {
       if (expired) throw new SearchDeadlineError()
       recordSearchStageDuration(`${this.leg}.connection_acquire`, performance.now() - started)
       const timeout = String(this.remaining())
-      /** Interactive retrieval cannot amortize compilation of the access predicates. */
+      /**
+       * Interactive retrieval cannot amortize compilation of the access predicates. A stage's
+       * own session settings ride in the same statement: each is a round trip otherwise.
+       */
       await tx.execute(
-        sql`SELECT set_config('statement_timeout', ${timeout}, true), set_config('jit', 'off', true)`
+        sessionSettingsStatement([
+          sql`set_config('statement_timeout', ${timeout}, true)`,
+          sql`set_config('jit', 'off', true)`,
+          ...settings,
+        ])
       )
       this.remaining()
       return measureSearchStage(stage, () => run(tx))
@@ -120,6 +128,11 @@ export class SearchBudget {
   }
 }
 
+/** One statement applying `set_config` fragments to the transaction they run in. */
+export function sessionSettingsStatement(settings: readonly SQL[]): SQL {
+  return sql.join([sql`SELECT`, sql.join([...settings], sql`, `)], sql` `)
+}
+
 /** Execute SQL with stage diagnostics and, for live search, the shared retrieval deadline. */
 export function runSearchQuery<T>(
   budget: SearchBudget | undefined,

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

@@ -107,8 +107,6 @@ export interface SearchDiagnosticMetadata {
   vectorSourcesSliced?: number
   /** The sliced sources held more readable documents than one exact ranking may enumerate. */
   vectorSlicedSaturated?: boolean
-  /** The whole-graph walk a broad reach earned came back short, so its sources were searched instead. */
-  vectorBroadWalkUnderfilled?: boolean
   vectorSourcesWalked?: number
   /**
    * Which index ranked an unbounded keyword leg: `tin` ranks by BM25 and checks access on the top

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

@@ -40,6 +40,7 @@ import {
   handleVectorOnlySearch,
   type PermittedDocuments,
   resolvePermittedDocuments,
+  resolveReach,
   retrieveKnowledgeSearch,
   type SearchParams,
   VECTOR_PROBE_DOCUMENT_LIMIT,
@@ -678,7 +679,19 @@ describe('workspace-scoped vector retrieval', () => {
       statements()
         .filter((query) => query.sql.includes('statement_timeout'))
         .map((query) => query.params[0])
-    ).toEqual(['100', '100', '40', '20'])
+    ).toEqual(['100', '40', '20'])
+  })
+
+  it('applies the scan settings in the deadline statement rather than one of their own', async () => {
+    queueTableRows(schemaMock.embedding, ranked)
+    await handleVectorOnlySearch({
+      ...params,
+      budget: new SearchBudget('vector', performance.now() + 8000),
+    })
+    const settings = statements().filter((query) => query.sql.includes('hnsw.iterative_scan'))
+    expect(settings.length).toBeGreaterThan(0)
+    /** Each round trip is latency on the critical path; the settings never earn one alone. */
+    for (const statement of settings) expect(statement.sql).toContain('statement_timeout')
   })
 
   it('does not convert an unexpected candidate failure into partial retrieval', async () => {
@@ -1316,73 +1329,6 @@ describe('permitted-document planner', () => {
     expect(statements().some((query) => query.sql.includes('WITH readable_chunks'))).toBe(false)
   })
 
-  it('widens the walk of a broad caller whose whole-graph walk came back short', async () => {
-    const eligibility = { workspace: [], admin: ['other-src'], members: ['member-src'] }
-    indexedSourceRows = [{ name: 'idx', connectorId: 'member-src' }]
-    /** The graph walk finds one readable neighbour where the pool wants hundreds. */
-    traversedRows = [{ id: 'short-hit', distance: 0.4 }]
-    rerankRows = [hit('short-hit', 'member-src')]
-    queueTableRows(schemaMock.embedding, rerankRows)
-    await handleVectorOnlySearch({
-      ...params,
-      topK: 2,
-      permitted: { kind: 'unbounded', broad: true },
-      accessPlan: {
-        connectors: eligibility,
-        observers: { confirmed: [{ id: 'm-1', connectorId: 'member-src' }], observed: [] },
-        memberSources: ['member-src'],
-      },
-    })
-    /** The same walk twice, the second with the wider scan; no source is searched on its own. */
-    const walks = statements().filter((query) => isWalk(query.sql))
-    expect(walks).toHaveLength(2)
-    expect(JSON.stringify(walks[1])).not.toContain('"right":"member-src"')
-    const settings = statements().filter((query) => query.sql.includes('hnsw.max_scan_tuples'))
-    expect(JSON.stringify(settings.at(-1))).toContain('100000')
-    expect(statements().some((query) => query.sql.includes('WITH readable_chunks'))).toBe(false)
-  })
-
-  it('keeps what a short broad walk found when the wider walk runs out of budget', async () => {
-    const eligibility = { workspace: [], admin: ['other-src'], members: ['member-src'] }
-    indexedSourceRows = [{ name: 'idx', connectorId: 'member-src' }]
-    traversedRows = [{ id: 'short-hit', distance: 0.4 }]
-    rerankRows = [hit('short-hit', 'member-src')]
-    queueTableRows(schemaMock.embedding, rerankRows)
-    /** The per-source search is cancelled by the server's statement timeout; the leg has time left. */
-    const budget = new SearchBudget('vector', performance.now() + 10_000)
-    const base = dbChainMockFns.execute.getMockImplementation()!
-    let walksSeen = 0
-    dbChainMockFns.execute.mockImplementation(async (query) => {
-      const statement = render(query).sql
-      if (isWalk(statement) && walksSeen++ > 0) {
-        throw Object.assign(new Error('canceling statement due to statement timeout'), {
-          code: '57014',
-        })
-      }
-      return base(query)
-    })
-    const rows = await handleVectorOnlySearch({
-      ...params,
-      budget,
-      permitted: { kind: 'unbounded', broad: true },
-      accessPlan: {
-        connectors: eligibility,
-        observers: { confirmed: [{ id: 'm-1', connectorId: 'member-src' }], observed: [] },
-        memberSources: ['member-src'],
-      },
-    })
-    expect(rows.map((row) => row.id)).toEqual(['short-hit'])
-    /** The wider walk ran under a share of the leg's budget, not all of it. */
-    const timeouts = statements()
-      .filter((query) => query.sql.includes("'statement_timeout'"))
-      .map((query) =>
-        Number(query.params.find((param) => typeof param === 'string' && /^\d+$/.test(param)))
-      )
-      .filter((value) => Number.isFinite(value))
-    expect(Math.min(...timeouts)).toBeLessThanOrEqual(5000)
-    expect(Math.max(...timeouts)).toBeGreaterThan(5000)
-  })
-
   it('walks an indexed source a bounded caller is a member of instead of ranking it exactly', async () => {
     const eligibility = { workspace: [], admin: ['small-src'], members: ['member-src'] }
     indexedSourceRows = [{ name: 'idx', connectorId: 'member-src' }]
@@ -1583,8 +1529,32 @@ describe('permitted-document planner', () => {
       expect(statement).toContain('ranked_tin_chunks.acl')
     })
 
-    it('takes one wide window for a narrow resolved scope and leaves a short page short', async () => {
-      tinPages = [{ ranked: 20_000, candidates: [hit('a', 'src-a')] }]
+    it('widens the window for a broad resolved scope whose first page came back short', async () => {
+      tinPages = [
+        { ranked: 2000, candidates: [] },
+        { ranked: 4000, candidates: [hit('b', 'src-a')] },
+      ]
+      queueTableRows(schemaMock.embedding, [{ ...hit('b', 'src-a'), content: 'release notes' }])
+      const results = await keyword({
+        permitted: { kind: 'unbounded', broad: true },
+        accessPlan: {
+          connectors: { workspace: [], admin: ['src-a'], members: [] },
+          observers: { confirmed: [], observed: [] },
+          memberSources: [],
+        },
+      })
+      expect(results.map((row) => row.id)).toEqual(['b'])
+      const windows = tinStatements().map((query) => JSON.stringify(query))
+      expect(windows).toHaveLength(2)
+      expect(windows[0]).toContain('2000')
+      expect(windows[1]).toContain('10000')
+    })
+
+    it('widens a narrow resolved scope straight to its wide window and leaves a short page short', async () => {
+      tinPages = [
+        { ranked: 2000, candidates: [] },
+        { ranked: 20_000, candidates: [hit('a', 'src-a')] },
+      ]
       queueTableRows(schemaMock.embedding, [{ ...hit('a', 'src-a'), content: 'release notes' }])
       const results = await keyword({
         permitted: { kind: 'unbounded', broad: false },
@@ -1595,12 +1565,57 @@ describe('permitted-document planner', () => {
         },
       })
       expect(results.map((row) => row.id)).toEqual(['a'])
-      /** One statement, at the widest window; nothing narrower first, and no ranking of every match after. */
+      /** The narrowest window, then the wide one; nothing between, and no ranking of every match after. */
+      const windows = tinStatements().map((query) => JSON.stringify(query))
+      expect(windows).toHaveLength(2)
+      expect(windows[0]).toContain('2000')
+      expect(windows[1]).toContain('20000')
+      expect(ginStatements()).toHaveLength(0)
+      /** Each window is ranked once for several pages' worth of readable rows, not once per page. */
+      expect(windows[1]).toContain('1000')
+    })
+
+    it('stops a narrow resolved scope at the narrowest window that fills its page', async () => {
+      const page = Array.from({ length: 20 }, (_, i) => hit(`k-${i}`, 'src-a'))
+      tinPages = [{ ranked: 2000, candidates: page }]
+      queueTableRows(
+        schemaMock.embedding,
+        page.map((row) => ({ ...row, content: 'release notes' }))
+      )
+      const results = await keyword({
+        topK: 20,
+        permitted: { kind: 'unbounded', broad: false },
+        accessPlan: {
+          connectors: { workspace: [], admin: ['src-a'], members: [] },
+          observers: { confirmed: [], observed: [] },
+          memberSources: [],
+        },
+      })
+      expect(results).toHaveLength(20)
+      /** Ranking costs grow with the window; a page the narrowest fills never pays for the wide one. */
       expect(tinStatements()).toHaveLength(1)
-      expect(JSON.stringify(tinStatements()[0])).toContain('20000')
+      expect(JSON.stringify(tinStatements()[0])).toContain('2000')
+      expect(JSON.stringify(tinStatements()[0])).not.toContain('20000')
+    })
+
+    it('leaves a broad resolved scope short at its widest window instead of ranking every match', async () => {
+      tinPages = [
+        { ranked: 2000, candidates: [] },
+        { ranked: 10_000, candidates: [] },
+        { ranked: 50_000, candidates: [hit('a', 'src-a')] },
+      ]
+      queueTableRows(schemaMock.embedding, [{ ...hit('a', 'src-a'), content: 'release notes' }])
+      const results = await keyword({
+        permitted: { kind: 'unbounded', broad: true },
+        accessPlan: {
+          connectors: { workspace: [], admin: ['src-a'], members: [] },
+          observers: { confirmed: [], observed: [] },
+          memberSources: [],
+        },
+      })
+      expect(results.map((row) => row.id)).toEqual(['a'])
+      expect(tinStatements()).toHaveLength(3)
       expect(ginStatements()).toHaveLength(0)
-      /** The window is ranked once for several pages' worth of readable rows, not once per page. */
-      expect(JSON.stringify(tinStatements()[0])).toContain('1000')
     })
 
     it('hydrates an oversized keyword page in slices and stops at the results it needs', async () => {
@@ -1777,6 +1792,31 @@ describe('permitted-document planner', () => {
       expect(reachCounts()).toHaveLength(2)
     })
 
+    it('does not remember a reach whose count ran out of time', async () => {
+      dbChainMockFns.execute.mockImplementation(async (query) => {
+        const statement = render(query).sql
+        if (statement.includes('EXPLAIN'))
+          return [{ 'QUERY PLAN': [{ Plan: { 'Plan Rows': 1_000_000 } }] }]
+        if (statement.includes(') reached'))
+          throw Object.assign(new Error('canceling statement due to statement timeout'), {
+            code: '57014',
+          })
+        return []
+      })
+      const budget = new SearchBudget('vector', performance.now() + 10_000)
+      const reachCounts = () => statements().filter((query) => query.sql.includes(') reached'))
+      const plan = await resolveReach(['org-index'], scope('timed-reach'), budget)
+      expect(plan).toEqual({ kind: 'unbounded', broad: true })
+      expect(reachCounts()).toHaveLength(1)
+      /** The next search counts again rather than trusting an answer that never came. */
+      await resolveReach(
+        ['org-index'],
+        scope('timed-reach'),
+        new SearchBudget('vector', performance.now() + 10_000)
+      )
+      expect(reachCounts()).toHaveLength(2)
+    })
+
     it('is remembered, so a broad caller skips the probe on the next search', async () => {
       probeRows = [{ id: null, connectorId: null, saturated: true }]
       const broad = scope('broad')