improvement(knowledge): resolve connector and member state once per search (#8040)

* improvement(knowledge): resolve connectors once and rank each source in its own index

Two costs dominated organization search, and both came from deriving per candidate
what is true of a whole source.

Connector state — deletion, archival, a pending access rewrite, the organization's
integration approval, the access mode — is a fact about a connector, so a search
resolves it once and filters candidates by the resulting ids. For a member reading
~25k documents, the candidate predicate's connector lookup drops from one per
document examined to one per document of a source that still needs this request's
live proof: 24,571 evaluations to 610, and the access check from ~390ms to ~160ms.

Ranking then follows sources. pgvector post-filters, so one walk over every source
spends its scan budget on the sources a caller cannot read: measured recall for a
member reading half an index ranged from 0.00 to 1.00, averaging 0.80 over nine
queries, two of which returned none of the exact page. A member of a source reads
essentially all of it, so that source is walked through an index covering it alone,
built after a sync grows it past the threshold and dropped with the connector.
Every other source is sliced — mirrored permissions give a caller their own mail,
their own files — and those slices are ranked exactly in one statement. Recall over
the same queries rises to 0.95 with none below 0.75.

Retrieval never waits on an index existing: a source without one is ranked exactly,
so a build that is skipped, fails, or has not happened yet costs recall nothing.

* improvement(knowledge): resolve the caller's member identities with the rest of the plan

A members-mode document is readable while one of the caller's active members observes
it, freshly, and which members those are is a fact about the caller. Resolving them
with the connectors — one query, one plan — turns each candidate's check into a lookup
on the observation key instead of a join to the member behind it, and lets the vector
planner read the sources the caller belongs to from the same resolution rather than
asking again.

* chore(knowledge): drop the stage the access plan made unnecessary

* improvement(knowledge): resolve live source grants once per search

- ask a live-proof source for the caller's grants once, before either leg, and only when the scope actually reads one, instead of per candidate page
- drop the per-page source exclusion and refill loop that followed from the per-page proof
- read search result metadata under the stored predicate the rows already passed
- end candidate paging on a short page whether or not the leg reorders its results
- batch the embedding_search.connector_id backfill in independently committed keyset pages, matching the other projection backfills

* test(db): register the projection-source script migration in the push and registry expectations

* fix(knowledge): walk a saturated slice and re-read connector state at hydration

- fall back to a graph walk when the sliced sources hold more readable documents than one exact ranking may enumerate, since that enumeration has no order and would otherwise rank an arbitrary subset
- read content under the full predicate, which re-reads each connector's own lifecycle and approval, so a source deleted, archived or unapproved mid-search stops answering at the gate that returns content

* improvement(knowledge): ask a live source for grants only when one of its candidates is read

Resolving every gated source's grants up front charged a search for sources it never ranked. The grants are now resolved on first need — when a page of ranked candidates actually contains one of those sources — and memoized for the rest of the search, so a scope that ranks none never asks and one that ranks many asks once.

* fix(knowledge): rank uploads without a sliced source and bound the grant resolution

- always run the slice that carries connectorless uploads; a caller who is a member of every indexed source had no sliced source beside them, and their uploaded documents were never ranked
- resolve live source grants in the provider's bounded pages, so a scope holding more gated sources than one page no longer fails
- drop a source's index only while it is invalid, so overlapping syncs cannot have one remove the index the other just built
- fan a document's new source out to the chunks a search can reach, which the document lookup index covers, and take the source again when a chunk is enabled

* 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

* fix(knowledge): start the scan budget over when the pages are rebuilt

The candidates already seen before a denied source was known are skipped when the pages are rebuilt, but they still counted against the scan budget, so a refill could exhaust it before reaching the readable candidates ranked after the denied source.

* fix(knowledge): serialize the projection source with its document

- the set trigger reads the document under FOR SHARE, so a projection write and a document's connector change cannot interleave and leave the older value
- the backfill writes a chunk's source only for the document it read that source from; a chunk moved meanwhile is left to its new document's trigger

* improvement(knowledge): choose the vector plan by the caller's reach

- a caller who reaches at least a quarter of the bases' documents walks the whole graph once, as before: pgvector post-filters, so at that share almost every neighbour the graph visits is theirs and one walk is the cheapest exact answer; searching each source on its own only helps a caller whose reach is a small share of the index
- the share is counted once against that bound when the probe saturates, and remembered with the saturation
- a bounded caller who is a member of an indexed source walks that source instead of ranking it exactly with the rest

* fix(knowledge): search a broad caller's sources when their whole-graph walk comes back short

Reach is counted from token overlap, which every readable document has but which connector state, requirement clauses or observations can still refuse. A caller the count called broad whose walk then underfills was not broad: their sources are searched on their own instead, so the misjudgement costs one walk rather than their neighbours.

* fix(knowledge): read result metadata under the scope the results were read under

Retrieval reports the scope its content was read under — with the live grants it resolved, when a gated source's candidate was read — and the metadata lookup applies the full read predicate under that scope. A gated source's result keeps its name and URL; a revoked one loses them here as it did at hydration.

* fix(knowledge): keep what a short broad walk found when searching its sources runs out of budget

The per-source search adds to the walk's candidates rather than replacing them, and if it is cancelled by the leg's budget the walk's candidates still stand.

* fix(api): fail a v1 knowledge search whose retrieval ran out of time

v1 cannot express an incomplete search, so a leg that reached its deadline fails the request, as it did before the route read the retrieval result directly.

* improvement(knowledge): mirror each chunk's source and ACL onto the ranking projections

- embedding_search and embedding_keyword_tin carry connector_id and acl from their document, set by each projection's own trigger under a share lock and fanned out when a document changes hands or its ACL is rewritten
- script migration 0021 installs both, backfills both projections in independently committed keyset pages, and builds a partial GIN index on each projection's acl concurrently
- projectionCandidateAccessCondition decides candidate readability on the projection row alone, admitting a superset of the per-row document predicate; a real-database test asserts it never refuses what that predicate admits

* improvement(knowledge): decide candidate readability on the ranking row

Once a search has resolved its connectors, the vector walks and the Tin keyword window test the caller's access on the row they score — the source and ACL are mirrored there — instead of joining document per visited or ranked chunk. The graph is no longer stalled by a document lookup per candidate, and a keyword window of mostly unreadable chunks costs an array test per row. Content is still read under the full predicate at hydration.

* improvement(knowledge): rank the sliced sources on the projection row

The sliced sources' readable chunks are enumerated by the projection's ACL index and ranked exactly in one statement, with the saturation bound counted in the same statement; the per-document enumeration through document is gone from the per-source path.

* improvement(knowledge): let a resolved scope's reach choose its plan without enumerating documents

- with readability decided on the projection row, a resolved scope never needs its readable documents enumerated ahead of ranking: its reach alone chooses between one walk over the whole graph and a search of each source, and is remembered per bases and tokens
- the reach share's denominator is the planner's estimate of the bases' documents rather than a count of them
- a resolved scope's keyword leg ranks with Tin and decides readability on the ranked row whether or not its reach is bounded

* improvement(knowledge): walk every source a caller reads whole; one keyword window for a resolved scope

- a source earns its own vector index at a thousand documents: once readability is decided on the row a walk costs a few milliseconds whatever the source's size, where ranking it exactly grows with it
- a resolved scope's keyword page takes one wide Tin window and leaves a short page short, rather than widening through narrower windows and then ranking every match

* improvement(knowledge): widen a broad reader's short walk; keyword windows follow reach

- a broad reader whose whole-graph walk comes back short walks again with a wider scan instead of searching each source: their own sources are already in the walk, and enumerating their readable chunks is a bitmap over most of the index
- a broad reader's keyword page widens through the narrow windows as before; only a narrow reader takes the wide window alone

* improvement(knowledge): index the projection by source; rank a narrow reader's keyword window once

- a partial btree on embedding_search.connector_id lets the planner lead with a few sources where the caller's tokens alone would match most of the index, which is what made a small slice cost a bitmap over all of it
- a narrow reader's keyword page returns several pages' worth of readable rows from one wide ranking, instead of ranking the window again for each page

* improvement(knowledge): widen a broad walk only when it is short of the result need; a narrower keyword window for narrow readers

A walk that found hundreds of candidates needs no rescue; widening it spent the leg's budget and lost them. The wider walk now runs only when the first found fewer candidates than there are results to return. A narrow reader's one keyword window is 20k rows: a few percent of it fills their page several times over, at less than half the widest window's cost.

* fix(knowledge): hydrate an oversized ranking page in slices, only as far as the results need

A ranking may hand back more candidates than one hydration should read — a narrow reader's keyword window is ranked once for several pages' worth — so a page is drained in slices of the result need and what is left waits, unread, until the results still need it.

* fix(knowledge): count a candidate as considered only once its slice is read

The slices a refill discards were never read, so the rebuilt pages may hand their readable candidates back.

* fix(knowledge): pair an observer with its connector, settle the per-source fan-out, serialize a source's index build

- an observation vouches for a document only from a member of the document's own connector: a document that changed hands keeps its old observations, which no longer carry it as a candidate
- a source whose search runs out of budget marks the leg partial while the other sources' results stand
- a source's index is built under a session lock on the one reserved connection that also builds it, so two syncs cannot interleave the invalidity check and the drop; 0021 drops an invalid leftover before building, and analyzes both projections after the backfill
- a broad reader's walk widens when it found fewer candidates than the smallest pool worth reranking

* fix(knowledge): give a broad reader's wider walk half of what the leg has left

The rerank and hydration of whatever is found — the first walk's candidates at least — keep the rest, so a wider walk that runs out of its share can no longer take the leg's results with it.

* 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

* test(knowledge): expect the scan settings to share the deadline statement

Author: waleedlatif1

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

@@ -14,6 +14,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
 const {
   mockResolveV1KnowledgeReadAccess,
   mockExecuteKnowledgeSearch,
+  mockRetrievalStatus,
   mockGenerateSearchEmbedding,
   mockGetDocumentMetadataByIds,
   mockGetDocumentTagDefinitions,
@@ -25,6 +26,7 @@ const {
 } = vi.hoisted(() => ({
   mockResolveV1KnowledgeReadAccess: vi.fn(),
   mockExecuteKnowledgeSearch: vi.fn(),
+  mockRetrievalStatus: vi.fn(() => ({ status: 'complete', timedOutLegs: [] })),
   mockGenerateSearchEmbedding: vi.fn(),
   mockGetDocumentMetadataByIds: vi.fn(),
   mockGetDocumentTagDefinitions: vi.fn(),
@@ -54,7 +56,12 @@ vi.mock('@/lib/knowledge/access/availability', () => ({
 }))
 
 vi.mock('@/lib/knowledge/search/queries', () => ({
-  executeKnowledgeSearch: mockExecuteKnowledgeSearch,
+  /** The route reads the retrieval result; the rows come from the same mock the tests drive. */
+  retrieveKnowledgeSearch: async (params: { access: unknown }) => ({
+    rows: await mockExecuteKnowledgeSearch(params),
+    retrieval: mockRetrievalStatus(),
+    readAccess: params.access,
+  }),
   getDocumentMetadataByIds: mockGetDocumentMetadataByIds,
 }))
 
@@ -139,6 +146,27 @@ describe('v1 knowledge search route — per-KB embedding model', () => {
     mockRecordSearchEmbeddingUsage.mockResolvedValue(undefined)
   })
 
+  it('fails a search whose retrieval ran out of time instead of returning partial rows', async () => {
+    const access = { kind: 'user' as const, userId: 'user-1', tokens: ['reader-token'] }
+    mockResolveV1KnowledgeReadAccess.mockResolvedValue({
+      get: vi.fn().mockResolvedValue(access),
+      getForConnectors: vi.fn(),
+      getForDocuments: vi.fn(),
+    })
+    mockCheckKnowledgeBaseAccess.mockResolvedValueOnce({
+      hasAccess: true,
+      knowledgeBase: baseKb('kb-1', 'text-embedding-3-small'),
+    })
+    mockRetrievalStatus.mockReturnValueOnce({ status: 'partial', timedOutLegs: ['vector'] })
+    mockExecuteKnowledgeSearch.mockResolvedValue([])
+    const response = await POST(
+      createMockRequest('POST', { workspaceId: 'ws-1', knowledgeBaseIds: 'kb-1', query: 'hello' })
+    )
+    expect(mockExecuteKnowledgeSearch).toHaveBeenCalledOnce()
+    expect(response.status).toBe(500)
+    expect(mockGetDocumentMetadataByIds).not.toHaveBeenCalled()
+  })
+
   it('retains the reader provider for ranked results and returned document metadata', async () => {
     const access = { kind: 'user' as const, userId: 'user-1', tokens: ['reader-token'] }
     const provider = {
@@ -165,7 +193,7 @@ describe('v1 knowledge search route — per-KB embedding model', () => {
         accessProvider: provider,
       })
     )
-    expect(mockGetDocumentMetadataByIds).toHaveBeenCalledWith([], access, provider)
+    expect(mockGetDocumentMetadataByIds).toHaveBeenCalledWith([], access)
   })
 
   it.each([
@@ -224,8 +252,7 @@ describe('v1 knowledge search route — per-KB embedding model', () => {
       expect(response.status).toBe(200)
       expect(mockGetDocumentMetadataByIds).toHaveBeenCalledWith(
         ['revoked-document', 'allowed-document'],
-        access,
-        provider
+        access
       )
       expect(body.data.results).toEqual(
         allDenied

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

@@ -14,10 +14,12 @@ import {
   type KbEmbeddingTarget,
   recordSearchEmbeddingUsage,
 } from '@/lib/knowledge/embeddings'
+import { SearchDeadlineError } from '@/lib/knowledge/search/budget'
 import { resolveKnowledgeSearchDefaults } from '@/lib/knowledge/search/defaults'
 import {
-  executeKnowledgeSearch,
   getDocumentMetadataByIds,
+  type KnowledgeRetrievalResult,
+  retrieveKnowledgeSearch,
   type SearchResult,
 } from '@/lib/knowledge/search/queries'
 import { getDocumentTagDefinitions } from '@/lib/knowledge/tags/service'
@@ -226,7 +228,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
         }
       : undefined
 
-    let results: SearchResult[]
+    let retrieved: KnowledgeRetrievalResult
     let queryEmbeddingIsBYOK: boolean | null = null
     const [readAccess, { searchMode, boostRecency }] = await Promise.all([
       resolveV1KnowledgeReadAccess(userId, rateLimit, workspaceId),
@@ -242,7 +244,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
     const access = 'get' in readAccess ? await readAccess.get() : readAccess
 
     if (!hasQuery && hasFilters) {
-      results = await executeKnowledgeSearch({
+      retrieved = await retrieveKnowledgeSearch({
         knowledgeBaseIds: accessibleKbIds,
         topK,
         access,
@@ -258,7 +260,7 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
         workspaceId
       )
       queryEmbeddingIsBYOK = queryEmbeddingResult.isBYOK
-      results = await executeKnowledgeSearch({
+      retrieved = await retrieveKnowledgeSearch({
         knowledgeBaseIds: accessibleKbIds,
         topK,
         access,
@@ -311,8 +313,11 @@ export const POST = withRouteHandler(async (request: NextRequest) => {
       tagDefinitionsMap[kbId] = map
     })
 
+    /** v1 cannot express an incomplete search, so a leg that ran out of time fails the request. */
+    if (retrieved.retrieval.status === 'partial') throw new SearchDeadlineError()
+    const results = retrieved.rows
     const documentIds = results.map((r) => r.documentId)
-    const documentMetadataMap = await getDocumentMetadataByIds(documentIds, access, accessProvider)
+    const documentMetadataMap = await getDocumentMetadataByIds(documentIds, retrieved.readAccess)
     const readableResults = results.filter((result) => documentMetadataMap[result.documentId])
 
     return NextResponse.json({

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

@@ -131,11 +131,11 @@ describe('API-key KB block fan-out', () => {
         const matching = (fragment: string) =>
           statements.filter((query) => query.includes(fragment))
         /**
-         * Every statement runs under the leg's deadline: the candidate search reinstates it after
-         * tuning the scan, and the probe, the exact ranking, the rerank and hydration each open
-         * with one of their own.
+         * Every statement runs under the leg's deadline: the candidate search applies it with the
+         * scan settings in one statement, and the probe, the exact ranking, the rerank and
+         * hydration each open with one of their own.
          */
-        expect(matching('statement_timeout')).toHaveLength(bases.length * 6)
+        expect(matching('statement_timeout')).toHaveLength(bases.length * 5)
         /**
          * A scope this small leaves the bounded traversal short of its candidate limit, so every
          * search probes once and rescues once — never a widening retry loop.

apps/sim/lib/knowledge/access/connector-eligibility.ts

@@ -0,0 +1,122 @@
+import { db } from '@sim/db'
+import { knowledgeConnector, knowledgeConnectorMember } from '@sim/db/schema'
+import { and, eq, inArray, isNull, sql } from 'drizzle-orm'
+import { SOURCE_ACL_MAX_AGE_MS } from '@/lib/knowledge/access/freshness'
+import {
+  type KnowledgeConnectorEligibility,
+  type KnowledgeMemberObserver,
+  type KnowledgeMemberObservers,
+  type SearchAccessPlan,
+  textArrayLiteral,
+} from '@/lib/knowledge/access/predicate'
+import type { KnowledgeAccessScope } from '@/lib/knowledge/access/types'
+import { searchIntegrationAccessCondition } from '@/lib/knowledge/search/integration-policy'
+
+/**
+ * The connectors a search may read from, grouped by access mode, with the ones whose reader access
+ * must be proven live marked.
+ *
+ * Deletion, archival, a pending access rewrite and the organization's integration approval are
+ * facts about a connector. Resolving them once per query — there are tens of connectors against
+ * hundreds of thousands of documents — leaves each candidate its own columns to check.
+ */
+async function resolveConnectorEligibility(
+  knowledgeBaseIds: readonly string[]
+): Promise<KnowledgeConnectorEligibility> {
+  const eligibility: {
+    workspace: string[]
+    admin: string[]
+    members: string[]
+    liveProofRequired: string[]
+  } = { workspace: [], admin: [], members: [], liveProofRequired: [] }
+  if (knowledgeBaseIds.length === 0) return eligibility
+  const rows = await db
+    .select({
+      id: knowledgeConnector.id,
+      accessMode: knowledgeConnector.accessMode,
+      connectorType: knowledgeConnector.connectorType,
+      /** A GitHub connector is gated only where it names the immutable repository behind a grant. */
+      githubRepository: sql<boolean>`${knowledgeConnector.sourceConfig}::jsonb ? 'githubRepositoryId'`,
+    })
+    .from(knowledgeConnector)
+    .where(
+      and(
+        inArray(knowledgeConnector.knowledgeBaseId, [...knowledgeBaseIds]),
+        isNull(knowledgeConnector.deletedAt),
+        isNull(knowledgeConnector.archivedAt),
+        eq(knowledgeConnector.accessRewritePending, false),
+        searchIntegrationAccessCondition()
+      )
+    )
+  for (const row of rows) {
+    if (row.accessMode === 'workspace') eligibility.workspace.push(row.id)
+    else if (row.accessMode === 'admin') eligibility.admin.push(row.id)
+    else if (row.accessMode === 'members') eligibility.members.push(row.id)
+    else continue
+    const live =
+      (row.connectorType === 'github' && row.githubRepository) ||
+      (row.connectorType === 'confluence' && row.accessMode === 'admin')
+    if (live) eligibility.liveProofRequired.push(row.id)
+  }
+  return eligibility
+}
+
+/**
+ * The caller's own member identities on these connectors, split by whether the member's change
+ * feed is itself current.
+ *
+ * A members-mode document is readable while one of the caller's active members observes it,
+ * freshly — and which members those are is a fact about the caller, not about any document. A
+ * member whose feed drained recently confirms every observation it holds, so its observations need
+ * no age check at all; the rest are checked against the age of the observation itself. Resolved
+ * once, the per-document check becomes one lookup on the observation key, with no join to the
+ * member behind it.
+ */
+async function resolveMemberObservers(
+  access: KnowledgeAccessScope,
+  connectorIds: readonly string[]
+): Promise<{ observers: KnowledgeMemberObservers; memberSources: string[] }> {
+  if (access.kind !== 'user' || connectorIds.length === 0 || access.tokens.length === 0) {
+    return { observers: { confirmed: [], observed: [] }, memberSources: [] }
+  }
+  const rows = await db
+    .select({
+      id: knowledgeConnectorMember.id,
+      connectorId: knowledgeConnectorMember.connectorId,
+      syncedThrough: knowledgeConnectorMember.memberSyncedThrough,
+    })
+    .from(knowledgeConnectorMember)
+    .where(
+      and(
+        inArray(knowledgeConnectorMember.connectorId, [...connectorIds]),
+        eq(knowledgeConnectorMember.status, 'active'),
+        sql`${knowledgeConnectorMember.subjectToken} = ANY(${textArrayLiteral([...access.tokens])})`
+      )
+    )
+  const cutoff = Date.now() - SOURCE_ACL_MAX_AGE_MS
+  const confirmed: KnowledgeMemberObserver[] = []
+  const observed: KnowledgeMemberObserver[] = []
+  const memberSources = new Set<string>()
+  for (const row of rows) {
+    const member = { id: row.id, connectorId: row.connectorId }
+    if (row.syncedThrough !== null && row.syncedThrough.getTime() > cutoff) confirmed.push(member)
+    else observed.push(member)
+    memberSources.add(row.connectorId)
+  }
+  return { observers: { confirmed, observed }, memberSources: [...memberSources] }
+}
+
+/**
+ * Everything a search needs to know about its sources and the caller's standing in them, resolved
+ * once: which connectors it may read, the caller's member identities there, and the sources they
+ * are a member of. Each is a fact about a connector or a caller, so deriving them per candidate
+ * document is what made retrieval cost grow with the size of what someone may read.
+ */
+export async function resolveSearchAccessPlan(
+  knowledgeBaseIds: readonly string[],
+  access: KnowledgeAccessScope
+): Promise<SearchAccessPlan> {
+  const connectors = await resolveConnectorEligibility(knowledgeBaseIds)
+  const { observers, memberSources } = await resolveMemberObservers(access, connectors.members)
+  return { connectors, observers, memberSources }
+}