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

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 }
+}