v0.8.46: keyword search optimizations

Author: waleedlatif1

apps/sim/lib/knowledge/__integration__/search-latency.integration.ts

@@ -477,7 +477,7 @@ async function sample(
         item.query.includes('CROSS JOIN LATERAL') ||
         isVectorCandidateQuery(item.query) ||
         item.query.includes('WITH scored_search_candidates') ||
-        item.query.includes('WITH visible_keyword_documents'))
+        item.query.includes('WITH matched_keyword_chunks'))
   )
   const plans: Array<
     CapturedQuery & {
@@ -550,7 +550,7 @@ async function sample(
         assertIndexedCandidates(parsedPlan[0].Plan, diagnostics.vectorCandidateLimit!, width)
       }
     }
-    if (query.query.includes('WITH visible_keyword_documents')) {
+    if (query.query.includes('WITH matched_keyword_chunks')) {
       assertScalarKeywordSorts(parsedPlan[0].Plan)
     }
   }

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

@@ -884,7 +884,7 @@ describe('live repository authorization follows ranked candidates', () => {
       const statement = render(query).sql
       if (statement.includes('AS visible')) return candidatePages.shift() ?? []
       if (statement.includes('WITH scored_search_candidates')) return rerankPages.shift() ?? []
-      if (statement.includes('WITH visible_keyword_documents')) return keywordPages.shift() ?? []
+      if (statement.includes('WITH matched_keyword_chunks')) return keywordPages.shift() ?? []
       if (isExactRanking(statement)) return exactPages.shift() ?? []
       if (statement.includes('AS id FROM')) return probePages.shift() ?? []
       return []
@@ -1102,8 +1102,8 @@ describe('live repository authorization follows ranked candidates', () => {
       expect(getForConnectors).toHaveBeenCalledWith(['allowed-source'], undefined)
       if (mode === 'keyword') {
         const ranking = render(dbChainMockFns.execute.mock.calls[0][0]).sql
-        expect(ranking).toContain('scored_keyword_candidates AS MATERIALIZED')
-        expect(ranking).toContain('ORDER BY keyword_rank DESC, id LIMIT')
+        expect(ranking).toContain('matched_keyword_chunks AS MATERIALIZED')
+        expect(ranking).toContain('ORDER BY keyword_rank DESC, matched_keyword_chunks.id')
         expect(ranking).not.toContain('<=>')
         expect(ranking).not.toContain('"content"')
       } else if (mode === 'tags') {
@@ -1240,6 +1240,22 @@ describe('live repository authorization follows ranked candidates', () => {
     expect(refillPredicate).toContain('revoked-source')
   })
 
+  it('matches keyword chunks before the visibility predicate and ranks only what survives it', async () => {
+    keywordPages.push([candidate('selected', 'allowed-source')])
+    queueTableRows(schemaMock.embedding, [{ id: 'selected', content: 'verified result' }])
+    await executeKeywordSearch({ ...params, query: 'release', queryVector: params.queryVector! })
+    const ranking = render(dbChainMockFns.execute.mock.calls[0][0]).sql
+    const matched = ranking.indexOf('matched_keyword_chunks AS MATERIALIZED')
+    const visible = ranking.indexOf('visible_keyword_documents AS MATERIALIZED')
+    expect(matched).toBeGreaterThanOrEqual(0)
+    expect(visible).toBeGreaterThan(matched)
+    expect(ranking.slice(matched, visible)).not.toContain('keyword_rank')
+    expect(ranking.slice(visible)).toContain('FROM matched_keyword_chunks INNER JOIN')
+    /** The predicate fragments are parameterized, so the restriction is read off the query tree. */
+    const fragments = JSON.stringify(dbChainMockFns.execute.mock.calls[0][0])
+    expect(fragments).toContain('= ANY (ARRAY(SELECT document_id FROM matched_keyword_chunks))')
+  })
+
   it('recomputes keyword candidates after excluding a revoked source and rechecks content access', async () => {
     getForConnectors.mockResolvedValueOnce(identity)
     keywordPages.push(

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

@@ -1100,6 +1100,19 @@ export interface KeywordSearchParams {
  * with `topK` (measured at ~59x the buffer reads on a 20k-chunk base for a term
  * matching every row). Ranking therefore touches no vectors, and only the rows
  * that survive the limit are hydrated.
+ *
+ * The live-scope ranking query runs in three stages: match, authorize, rank. The
+ * visibility predicate carries correlated subqueries — one per connector, one per
+ * search-integration decision — so evaluating it across a base ahead of the query costs a table
+ * pass priced by how many documents the base holds rather than by how many the query matched.
+ * Matching first restricts that predicate to the documents the query actually matched.
+ *
+ * Two details keep that ordering from paying the saving back. Restricting the predicate with
+ * `document.id = ANY (...)` rather than a subquery keeps the narrowed lookup on a bitmap scan,
+ * which prefetches, where a plain `IN (SELECT ...)` plans as an index walk that does not. And
+ * the match stage carries identifiers only: ranking every match rather than every *visible*
+ * match would detoast one text-search vector per match, which on a mid-frequency term costs
+ * more than the pass it replaces.
  */
 export async function executeKeywordSearch(params: KeywordSearchParams): Promise<SearchResult[]> {
   const { knowledgeBaseIds, topK, query, queryVector, structuredFilters, access } = params
@@ -1133,38 +1146,41 @@ export async function executeKeywordSearch(params: KeywordSearchParams): Promise
       selectPage: async (limit, offset, excludedSources) => {
         const candidates = await runSearchQuery(params.budget, 'keyword.sql', (executor) =>
           executor.execute<SearchReadCandidate>(sql`
-            WITH visible_keyword_documents AS MATERIALIZED (
-              SELECT ${document.id} AS id FROM ${document}
-              WHERE ${and(
-                inArray(document.knowledgeBaseId, knowledgeBaseIds),
-                ...getDocumentVisibilityConditions(
-                  access,
-                  params.filters,
-                  knowledgeMetadataCandidateAccessCondition(access)
-                ),
-                excludeSearchSources(excludedSources)
-              )}
-            ), scored_keyword_candidates AS MATERIALIZED (
+            WITH matched_keyword_chunks AS MATERIALIZED (
               SELECT ${embeddingKeywordSearch.id} AS id,
-                ${embeddingKeywordSearch.documentId} AS document_id,
-                ${candidateRank} AS keyword_rank
+                ${embeddingKeywordSearch.documentId} AS document_id
               FROM ${embeddingKeywordSearch}
               WHERE ${and(
                 inArray(embeddingKeywordSearch.knowledgeBaseId, knowledgeBaseIds),
                 eq(embeddingKeywordSearch.enabled, true),
                 sql`${embeddingKeywordSearch.contentTsv} @@ ${tsQuery}`,
-                sql`${embeddingKeywordSearch.documentId} IN (SELECT id FROM visible_keyword_documents)`,
-                sql`EXISTS (SELECT 1 FROM visible_keyword_documents)`,
                 tagFilterConditions.length
                   ? sql`EXISTS (
                   SELECT 1 FROM ${embedding} WHERE ${embedding.id} = ${embeddingKeywordSearch.id}
                     AND ${and(...tagFilterConditions)}
                 )`
                   : undefined
               )}
+            ), visible_keyword_documents AS MATERIALIZED (
+              SELECT ${document.id} AS id FROM ${document}
+              WHERE ${and(
+                inArray(document.knowledgeBaseId, knowledgeBaseIds),
+                sql`${document.id} = ANY (ARRAY(SELECT document_id FROM matched_keyword_chunks))`,
+                ...getDocumentVisibilityConditions(
+                  access,
+                  params.filters,
+                  knowledgeMetadataCandidateAccessCondition(access)
+                ),
+                excludeSearchSources(excludedSources)
+              )}
             ), ranked_keyword_candidates AS MATERIALIZED (
-              SELECT * FROM scored_keyword_candidates
-              ORDER BY keyword_rank DESC, id LIMIT ${limit} OFFSET ${offset}
+              SELECT matched_keyword_chunks.id, matched_keyword_chunks.document_id,
+                ${candidateRank} AS keyword_rank
+              FROM matched_keyword_chunks INNER JOIN ${embeddingKeywordSearch}
+                ON ${embeddingKeywordSearch.id} = matched_keyword_chunks.id
+              WHERE matched_keyword_chunks.document_id IN (SELECT id FROM visible_keyword_documents)
+              ORDER BY keyword_rank DESC, matched_keyword_chunks.id
+              LIMIT ${limit} OFFSET ${offset}
             )
             SELECT ranked_keyword_candidates.id, ${document.id} AS "documentId",
               ${document.connectorId} AS "connectorId",

apps/sim/lib/workspace-files/application/search-workspace-file-content.test.ts

@@ -0,0 +1,75 @@
+/** @vitest-environment node */
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const mocks = vi.hoisted(() => ({
+  load: vi.fn(),
+  permission: vi.fn(),
+  search: vi.fn(),
+  folders: vi.fn(),
+}))
+vi.mock('@sim/platform-authz/workspace', () => ({
+  permissionSatisfies: (actual: string | null) => actual !== null,
+  resolveEffectiveWorkspacePermission: mocks.permission,
+}))
+vi.mock('@/lib/uploads/contexts/workspace', () => ({ loadActiveWorkspaceContext: mocks.load }))
+vi.mock('@/lib/workspace-files/search/repository', () => ({
+  searchWorkspaceFileIndex: mocks.search,
+}))
+vi.mock('@/lib/workspace-files/resolve-folder-scope', () => ({
+  resolveWorkspaceFolderScope: mocks.folders,
+}))
+
+import { searchWorkspaceFileContent } from '@/lib/workspace-files/application/search-workspace-file-content'
+
+const principal = { kind: 'session', userId: 'user-1', sessionId: 'session-1' } as const
+const input = {
+  workspaceId: 'workspace-1',
+  query: 'needle',
+  mode: 'exact',
+  maxResults: 10,
+} as const
+
+describe('searchWorkspaceFileContent cancellation', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mocks.load.mockResolvedValue({
+      workspaceId: 'workspace-1',
+      workspaceOrganizationId: null,
+      allowPersonalApiKeys: true,
+      billedAccountUserId: 'user-1',
+    })
+    mocks.permission.mockResolvedValue('read')
+    mocks.search.mockResolvedValue({ results: [] })
+  })
+
+  it.each(['request', 'input'] as const)(
+    'propagates the %s signal through the authorized application operation',
+    async (source) => {
+      const controller = new AbortController()
+      await searchWorkspaceFileContent.execute({
+        principal,
+        input: { ...input, ...(source === 'input' ? { signal: controller.signal } : {}) },
+        request: {
+          headers: new Headers(),
+          ...(source === 'request' ? { signal: controller.signal } : {}),
+        },
+      })
+      expect(mocks.search).toHaveBeenCalledWith(
+        expect.objectContaining({ signal: controller.signal })
+      )
+    }
+  )
+
+  it('does not resolve folders or enqueue database work for a cancelled HTTP request', async () => {
+    const signal = AbortSignal.abort(new Error('cancelled'))
+    await expect(
+      searchWorkspaceFileContent.execute({
+        principal,
+        input: { ...input, folderPaths: ['/notes'] },
+        request: { headers: new Headers(), signal },
+      })
+    ).rejects.toBe(signal.reason)
+    expect(mocks.folders).not.toHaveBeenCalled()
+    expect(mocks.search).not.toHaveBeenCalled()
+  })
+})

apps/sim/lib/workspace-files/application/search-workspace-file-content.ts

@@ -3,15 +3,13 @@ import { loadActiveWorkspaceContext } from '@/lib/uploads/contexts/workspace'
 import { defineAuthorizedWorkspaceFileUseCase } from '@/lib/workspace-files/application/authorized-workspace-file-use-case'
 import { fileOperations } from '@/lib/workspace-files/application/operations'
 import { resolveWorkspaceFolderScope } from '@/lib/workspace-files/resolve-folder-scope'
+import { WorkspaceFileSearchUnavailableError } from '@/lib/workspace-files/search/errors'
 import {
   compileFileSearchPattern,
   type FileSearchMode,
   FileSearchPatternError,
 } from '@/lib/workspace-files/search/pattern'
-import {
-  searchWorkspaceFileIndex,
-  WorkspaceFileSearchUnavailableError,
-} from '@/lib/workspace-files/search/repository'
+import { searchWorkspaceFileIndex } from '@/lib/workspace-files/search/repository'
 
 export interface SearchWorkspaceFileContentInput {
   workspaceId: string
@@ -37,14 +35,16 @@ export const searchWorkspaceFileContent = defineAuthorizedWorkspaceFileUseCase({
   operation: fileOperations.searchContent,
   resolveContext: ({ input }: { input: SearchWorkspaceFileContentInput }) =>
     resolveSearchWorkspaceFileContext(input),
-  execute: async ({ principal, input, context }) => {
-    /*
+  execute: async ({ principal, input, context, request }) => {
+    const signal = input.signal ?? request?.signal
+    signal?.throwIfAborted()
+    /**
      * Resolved here rather than at the surface so every caller (the File
      * block, the v2 route) is confined by the same check. A folder tree
      * holding one subtree per user makes this scope the isolation boundary,
      * not a convenience filter.
      */
-    /*
+    /**
      * `!== undefined`, not a length check: an explicitly empty list is a scope
      * that names no folder, which must match nothing. Treating it as "absent"
      * would answer a request for nothing with the whole workspace.
@@ -58,15 +58,15 @@ export const searchWorkspaceFileContent = defineAuthorizedWorkspaceFileUseCase({
             includeSubfolders: input.includeSubfolders,
           })
         : undefined
-    input.signal?.throwIfAborted()
+    signal?.throwIfAborted()
 
     try {
       return await searchWorkspaceFileIndex({
         workspaceId: context.workspaceId,
         pattern: compileFileSearchPattern(input.query, input.mode),
         maxResults: input.maxResults,
         folderScope,
-        signal: input.signal,
+        signal,
       })
     } catch (error) {
       /**

apps/sim/lib/workspace-files/search/README.md

@@ -22,7 +22,9 @@ Search joins the current file revision and resolved workspace/folder scope. A re
 
 For regular chunks, PostgreSQL checks the pattern with newline-aware semantics, then verifies individual logical lines. Long-line fragments use only necessary three-character literals as a conservative prefilter, including all required alternation branches. Two-code-point overlap preserves those literals at every boundary. PostgreSQL reconstructs the complete candidate line and evaluates the original regex, so anchors, word boundaries, repetitions, and arbitrarily long match spans retain line semantics. Fixed overlap alone is never treated as proof of a match. The supported regex grammar and minimum literal requirement are unchanged.
 
-Regular blocks are verified in batches of at most 16 (128 KiB of indexed text); long lines are reconstructed one at a time. Only bounded match-centered previews leave PostgreSQL: at most 201 rows to detect truncation, and at most 2 KiB per rendered result. A single search has a ten-second application deadline with per-statement guards. PostgreSQL 17 additionally enforces a total transaction timeout; PostgreSQL 16 uses the compatible idle-transaction guard. Transaction advisory locks admit at most two simultaneous searches per workspace and ten globally per database. Busy and timed-out searches fail explicitly; they never report an incomplete scan as an authoritative empty result. The reader uses the normal application database connection, so admission is coordinated on the same database as the index.
+Regular blocks are verified in batches of at most 16 (128 KiB of indexed text); long lines are reconstructed one at a time. Only bounded match-centered previews leave PostgreSQL: at most 201 rows to detect truncation, and at most 2 KiB per rendered result. A single search has a fifteen-second caller deadline covering queueing, connection acquisition, and execution; SQL runs for at most ten seconds within that budget, with per-statement guards. PostgreSQL 17 additionally enforces a total transaction timeout; PostgreSQL 16 uses the compatible idle-transaction guard. Transaction advisory locks admit at most 20 simultaneous searches per workspace and 5,000 globally per database. Search transactions use the dedicated `dbFor('search')` primary pool, with five connections per process, so they cannot occupy the application or execution client pools. Before acquiring a connection, a process admits five active searches and at most 100 waiting requests, capped at 20 waiting requests per workspace so one burst cannot fill the entire queue. Queued workspaces rotate after each grant; waiting requests expire after five seconds or leave immediately on cancellation. The local active budget comes from the same pool profile as the driver. A slot is released only when the transaction settles, including errors. Cancellation reaches this boundary from both HTTP requests and File-tool execution. The existing deadline helper ends the caller’s wait promptly even when connection acquisition stalls. A late transaction checks cancellation before running search SQL; its active slot remains reserved until the driver settles, so a timed-out request cannot start replacement work on top of a still-running query. The driver and upstream pooler still own physical connection cleanup; the application deadline does not cancel a queued PostgreSQL protocol command. Busy and timed-out searches fail explicitly; they never report an incomplete scan as an authoritative empty result. These bounds apply identically to exact and regex search and do not change result or line-number semantics.
+
+The 20/workspace and 5,000/global advisory ceilings bound admitted transactions across processes; they are not promises of simultaneous execution or throughput. Local queues absorb short bursts without holding database connections. They are not durable jobs or a fleet-wide fair scheduler. The dedicated client pool isolates connection ownership, not PostgreSQL CPU, memory, I/O, or an upstream PgBouncer server pool. Its default URL is the process primary URL; any `DATABASE_URL_SEARCH` override must target the same primary database so current revisions and advisory admission remain coherent. Independent PgBouncer server budgets require separate database/user pool configuration. Total client connections can increase by five per participating process. Before raising execution capacity, measure the number of processes, backend pool budget, queue wait/rejection rates, search latency, and database resource headroom under representative exact and broad-regex workloads. More queueing cannot increase sustained throughput.
 
 Arbitrary regex cannot have a fixed latency guarantee. Common terms, broad alternatives, and punctuation-only literals may require scanning significant scoped text. Larger capacity decisions need representative query plans and workload measurements; neither a per-file byte cap nor a PostgreSQL row-count claim establishes a total corpus capacity.