fix(file-search): bound chunk inserts by trigram work

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

@@ -10,9 +10,9 @@ PostgreSQL stores complete extracted text in bounded chunks. Object storage rema
 
 8 KiB values may still use PostgreSQL TOAST. The bound controls the size of each logical value and detoast operation; avoiding TOAST entirely is not the objective. Tiny lines share rows, so row count scales with bytes instead of newline count. Worst-case line packing can leave roughly half a block unused; long-line overlap adds at most eight bytes per fragment.
 
-Workers download and extract outside database transactions, then insert batches of at most 250 rows / 128 KiB. Each batch checks the build token and lease. Publication locks the canonical file, build, and revision in that order, verifies the stored chunk count, and changes the visible pointer only after every batch succeeds. Old dispatch failure callbacks cannot overwrite newer dispatches or successful builds.
+Workers download and extract outside database transactions, then insert batches of at most 250 rows / 128 KiB and a target of 8,192 estimated trigram keys. Each batch checks the build token and lease. Publication locks the canonical file, build, and revision in that order, verifies the stored chunk count, and changes the visible pointer only after every batch succeeds. Old dispatch failure callbacks cannot overwrite newer dispatches or successful builds.
 
-The chunk GIN index uses `fastupdate = off`. Each bounded insert updates the main index directly instead of appending to a shared pending list. With deferred updates enabled, even a small insert can cross the pending-list threshold and synchronously merge accumulated work from other files. Direct updates trade some bulk-write throughput for avoiding that foreground cleanup cliff. They do not eliminate normal index I/O, vacuum, or storage contention; the row and worker limits still apply. The 128 KiB batch budget bounds direct index work per transaction without changing the 25 MiB file coverage limit. Dense text with many distinct trigrams and an index working set larger than the available cache can still exceed the statement deadline. Capacity validation must include that cache pressure, not only a small corpus or a row count. A batch that takes at least two seconds, including one a statement timeout cancels, is logged with its rows, bytes, and estimated trigram key count (pg_trgm's extraction, reproduced exactly for the database's `en_US.UTF-8` ctype), so a key-based batch budget can be sized from production timings. A failed query is reported to the task runner by error code only; Drizzle's message carries the bound file text.
+The chunk GIN index uses `fastupdate = off`. Each bounded insert updates the main index directly instead of appending to a shared pending list. With deferred updates enabled, even a small insert can cross the pending-list threshold and synchronously merge accumulated work from other files. Direct updates trade some bulk-write throughput for avoiding that foreground cleanup cliff. They do not eliminate normal index I/O, vacuum, or storage contention; the row and worker limits still apply. Bytes alone do not bound GIN posting updates: dense text can generate thousands of distinct keys in each chunk. The batch planner sums each chunk's distinct trigram count, including repeated keys across rows, and flushes before the next chunk exceeds the key target. A single 8 KiB chunk is always allowed to make progress even if word padding pushes its estimate slightly above the target; it is written alone. Storage validates the same budget before opening the transaction. These are write scheduling bounds, not file exclusions: chunk boundaries, complete-file publication, and the 25 MiB file coverage limit are unchanged. The estimate mirrors pg_trgm under `en_US.UTF-8`; it is a work estimate, not a latency guarantee. Capacity validation must still include dense text, concurrent writers, and an index working set larger than available cache. A batch that takes at least two seconds, including a canceled statement, is logged with its rows, bytes, and estimated key count. A failed query is reported to the task runner by error code only; Drizzle's message carries the bound file text.
 
 Indexing transactions have separate limits from search: ten seconds per statement, five seconds waiting for a lock, and thirty seconds total on PostgreSQL 17. The outer limit leaves time for ordinary statement cancellation and rollback instead of terminating the connection at the same ten-second deadline. PostgreSQL 16 uses the compatible idle-transaction guard. A canceled batch remains unpublished; the existing task retry starts a fresh fenced build, and cleanup retires the previous attempt. This does not automatically retry revisions already marked failed.
 

apps/sim/lib/workspace-files/search/chunks.integration.ts

@@ -32,16 +32,15 @@ vi.mock('@/lib/copilot/tools/server/files/doc-compile', () => ({ resolveServable
 vi.mock('@/lib/file-parsers', () => ({ parseBuffer: vi.fn(), isSupportedFileType: vi.fn() }))
 
 import {
-  FILE_SEARCH_CHUNK_BYTES,
   FILE_SEARCH_CLEANUP_BATCH_ROWS,
   FILE_SEARCH_CLEANUP_BUDGET_MS,
   FILE_SEARCH_CLEANUP_MAX_BATCHES,
-  FILE_SEARCH_INSERT_BATCH_BYTES,
-  FILE_SEARCH_INSERT_BATCH_ROWS,
+  FILE_SEARCH_INSERT_BATCH_TRIGRAM_KEYS,
   FILE_SEARCH_QUERY_GLOBAL_CONCURRENCY,
   FILE_SEARCH_QUERY_WORKSPACE_CONCURRENCY,
 } from '@/lib/workspace-files/search/constants'
 import { prepareWorkspaceFileSearchDispatch } from '@/lib/workspace-files/search/dispatcher'
+import { iterateFileSearchBatches } from '@/lib/workspace-files/search/index-batches'
 import {
   iterateFileSearchChunks,
   planFileSearchIndex,
@@ -142,14 +141,8 @@ describe('chunked workspace file search on PostgreSQL', () => {
     expect(build).not.toBeNull()
     const plan = planFileSearchIndex({ text, partial: false }, signal)
     const chunks = [...iterateFileSearchChunks(plan, signal)]
-    const batchRows = Math.min(
-      FILE_SEARCH_INSERT_BATCH_ROWS,
-      Math.floor(FILE_SEARCH_INSERT_BATCH_BYTES / FILE_SEARCH_CHUNK_BYTES)
-    )
-    for (let offset = 0; offset < chunks.length; offset += batchRows)
-      expect(
-        await appendFileSearchChunks(build!, chunks.slice(offset, offset + batchRows), signal)
-      ).toBe(true)
+    for (const batch of iterateFileSearchBatches(chunks, signal))
+      expect(await appendFileSearchChunks(build!, batch, signal)).toBe(true)
     expect(
       await publishFileSearchBuild(
         build!,
@@ -304,6 +297,43 @@ describe('chunked workspace file search on PostgreSQL', () => {
     expect((await search('needle')).results).toMatchObject([{ fileId: 'file-1', lineNumber: 1 }])
   })
 
+  it('bounds native GIN posting work and keeps dense-file exact and regex line results', async () => {
+    const lines = Array.from(
+      { length: 1200 },
+      (_, i) => `dependency-${i}: sha512-${createHash('sha512').update(String(i)).digest('base64')}`
+    )
+    const text = lines.join('\n')
+    const plan = planFileSearchIndex({ text, partial: false }, signal)
+    const chunks = [...iterateFileSearchChunks(plan, signal)]
+    const build = (await beginFileSearchBuild(revision))!
+    await expect(appendFileSearchChunks(build, chunks.slice(0, 16), signal)).rejects.toThrow(
+      'insert batch exceeds its budget'
+    )
+    for (const batch of iterateFileSearchBatches(chunks, signal)) {
+      const [{ keys }] = await connection`SELECT sum(cardinality(show_trgm(content)))::int AS keys
+        FROM (VALUES ${connection(batch.map((c) => [c.content]))}) AS chunk(content)`
+      expect(keys).toBeLessThanOrEqual(FILE_SEARCH_INSERT_BATCH_TRIGRAM_KEYS)
+      expect(await appendFileSearchChunks(build, batch, signal)).toBe(true)
+    }
+    expect((await search('dependency-1199')).results).toEqual([])
+    expect(
+      await publishFileSearchBuild(
+        build,
+        {
+          status: 'ready',
+          chunkCount: chunks.length,
+          lineCount: plan.lineCount,
+          indexedBytes: plan.indexedBytes,
+        },
+        signal
+      )
+    ).toBe(true)
+    expect((await search(lines[1199])).results).toMatchObject([{ lineNumber: 1200 }])
+    expect(
+      (await search('^dependency-1199: sha512-[A-Za-z0-9+/]+=*$', 'regex')).results
+    ).toMatchObject([{ lineNumber: 1200 }])
+  })
+
   it('packs a million short lines without a million rows and bounds every stored value', async () => {
     await index('abc\n'.repeat(1_000_000))
     const [row] =

apps/sim/lib/workspace-files/search/constants.ts

@@ -61,8 +61,13 @@ export const FILE_SEARCH_CLEANUP_MIN_BATCH_MS =
   FILE_SEARCH_CLEANUP_BUDGET_MS / FILE_SEARCH_CLEANUP_MAX_BATCHES
 export const FILE_SEARCH_RECONCILE_INTERVAL_MS = 60 * 60 * 1000
 export const FILE_SEARCH_INSERT_BATCH_ROWS = 250
-/** Direct GIN writes perform index work in each insert, so transactions use smaller byte batches. */
+/** Bounds the text payload independently of its index work. */
 export const FILE_SEARCH_INSERT_BATCH_BYTES = 128 * 1024
+/**
+ * Sum of each chunk's distinct trigram keys, bounding direct GIN posting updates per insert.
+ * A single 8 KiB chunk may exceed this target slightly due to word padding and is written alone.
+ */
+export const FILE_SEARCH_INSERT_BATCH_TRIGRAM_KEYS = 8 * 1024
 /** Batches at least this slow are logged with their estimated trigram key count. */
 export const FILE_SEARCH_SLOW_INSERT_BATCH_MS = 2000
 

apps/sim/lib/workspace-files/search/index-batches.test.ts

@@ -0,0 +1,126 @@
+/** @vitest-environment node */
+import { createHash } from 'node:crypto'
+import { describe, expect, it } from 'vitest'
+import {
+  FILE_SEARCH_CHUNK_BYTES,
+  FILE_SEARCH_INSERT_BATCH_BYTES,
+  FILE_SEARCH_INSERT_BATCH_ROWS,
+  FILE_SEARCH_INSERT_BATCH_TRIGRAM_KEYS,
+} from '@/lib/workspace-files/search/constants'
+import { iterateFileSearchBatches } from '@/lib/workspace-files/search/index-batches'
+import {
+  estimateTrigramKeys,
+  type FileSearchChunk,
+  iterateFileSearchChunks,
+  planFileSearchIndex,
+} from '@/lib/workspace-files/search/index-plan'
+
+const signal = new AbortController().signal
+const chunk = (content: string, ordinal = 0): FileSearchChunk => ({
+  content,
+  ordinal,
+  lineStart: ordinal + 1,
+  fragment: false,
+  overlap: 0,
+})
+
+/** A de Bruijn prefix has no repeated three-letter windows, exercising the padding overhead. */
+function uniqueTrigrams(): string {
+  const alphabet = 'abcdefghijklmnopqrstuvwxyz'
+  const positions = Array<number>(4).fill(0)
+  let text = ''
+  const visit = (depth: number, period: number) => {
+    if (depth > 3) {
+      if (3 % period === 0) for (let i = 1; i <= period; i++) text += alphabet[positions[i]]
+      return
+    }
+    positions[depth] = positions[depth - period]
+    visit(depth + 1, period)
+    for (let i = positions[depth - period] + 1; i < alphabet.length; i++) {
+      positions[depth] = i
+      visit(depth + 1, depth)
+    }
+  }
+  visit(1, 1)
+  return text.slice(0, FILE_SEARCH_CHUNK_BYTES)
+}
+
+describe('file search insert batches', () => {
+  it('counts GIN work per row, preserving every dense lockfile line and ordinal', () => {
+    const text = Array.from(
+      { length: 1200 },
+      (_, i) =>
+        `dependency-${i}: sha512-${createHash('sha512').update(String(i)).digest('base64')}\n`
+    ).join('')
+    const chunks = [
+      ...iterateFileSearchChunks(planFileSearchIndex({ text, partial: false }, signal), signal),
+    ]
+    const batches = [...iterateFileSearchBatches(chunks, signal)]
+    expect(batches.length).toBeGreaterThan(
+      Math.ceil(Buffer.byteLength(text) / FILE_SEARCH_INSERT_BATCH_BYTES)
+    )
+    expect(batches.flat()).toEqual(chunks)
+    expect(
+      batches
+        .flat()
+        .map((c) => c.content)
+        .join('')
+    ).toBe(text)
+    for (const batch of batches) {
+      expect(batch.reduce((sum, c) => sum + estimateTrigramKeys(c.content), 0)).toBeLessThanOrEqual(
+        FILE_SEARCH_INSERT_BATCH_TRIGRAM_KEYS
+      )
+    }
+    const repeated = chunk(chunks[0].content)
+    const sameContent = [...iterateFileSearchBatches(Array(16).fill(repeated), signal)]
+    expect(sameContent.length).toBeGreaterThan(1)
+  })
+
+  it('retains byte batching for low-key text and the row bound for tiny chunks', () => {
+    const full = chunk('a'.repeat(FILE_SEARCH_CHUNK_BYTES))
+    const byBytes = [...iterateFileSearchBatches(Array(17).fill(full), signal)]
+    expect(byBytes.map((batch) => batch.length)).toEqual([16, 1])
+    const byRows = [
+      ...iterateFileSearchBatches(Array(FILE_SEARCH_INSERT_BATCH_ROWS + 1).fill(chunk('')), signal),
+    ]
+    expect(byRows.map((batch) => batch.length)).toEqual([FILE_SEARCH_INSERT_BATCH_ROWS, 1])
+  })
+
+  it('writes a maximum-key chunk alone without splitting or dropping it', () => {
+    const dense = chunk(uniqueTrigrams(), 1)
+    expect(estimateTrigramKeys(dense.content)).toBeGreaterThan(
+      FILE_SEARCH_INSERT_BATCH_TRIGRAM_KEYS
+    )
+    const first = chunk('first')
+    const last = chunk('last', 2)
+    expect([...iterateFileSearchBatches([first, dense, last], signal)]).toEqual([
+      [first],
+      [dense],
+      [last],
+    ])
+  })
+
+  it('rejects oversized chunks and emits no empty batch', () => {
+    expect([...iterateFileSearchBatches([], signal)]).toEqual([])
+    expect(() => [
+      ...iterateFileSearchBatches([chunk('a'.repeat(FILE_SEARCH_CHUNK_BYTES + 1))], signal),
+    ]).toThrow('chunk exceeds its budget')
+  })
+
+  it('consumes only one batch plus lookahead and honors cancellation between writes', () => {
+    let consumed = 0
+    function* source() {
+      for (let i = 0; i < 100; i++) {
+        consumed++
+        yield chunk('a'.repeat(FILE_SEARCH_CHUNK_BYTES), i)
+      }
+    }
+    const controller = new AbortController()
+    const batches = iterateFileSearchBatches(source(), controller.signal)
+    expect(batches.next().value).toHaveLength(16)
+    expect(consumed).toBe(17)
+    controller.abort(new Error('canceled'))
+    expect(() => batches.next()).toThrow('canceled')
+    expect(consumed).toBe(17)
+  })
+})

apps/sim/lib/workspace-files/search/index-batches.ts

@@ -0,0 +1,51 @@
+import { Buffer } from 'node:buffer'
+import {
+  FILE_SEARCH_CHUNK_BYTES,
+  FILE_SEARCH_INSERT_BATCH_BYTES,
+  FILE_SEARCH_INSERT_BATCH_ROWS,
+  FILE_SEARCH_INSERT_BATCH_TRIGRAM_KEYS,
+} from '@/lib/workspace-files/search/constants'
+import { estimateTrigramKeys, type FileSearchChunk } from '@/lib/workspace-files/search/index-plan'
+
+/** A single bounded chunk must make progress even when its estimate exceeds the key target. */
+export function exceedsFileSearchBatchBudget(rows: number, bytes: number, keys: number): boolean {
+  return (
+    rows > FILE_SEARCH_INSERT_BATCH_ROWS ||
+    bytes > FILE_SEARCH_INSERT_BATCH_BYTES ||
+    (rows > 1 && keys > FILE_SEARCH_INSERT_BATCH_TRIGRAM_KEYS)
+  )
+}
+
+/**
+ * Bounds each insert's payload and estimated GIN work without changing stored chunks or coverage.
+ * Keys are counted per row: repeated keys across rows still require separate posting updates.
+ */
+export function* iterateFileSearchBatches(
+  chunks: Iterable<FileSearchChunk>,
+  signal: AbortSignal
+): Generator<FileSearchChunk[]> {
+  let batch: FileSearchChunk[] = []
+  let bytes = 0
+  let keys = 0
+  for (const chunk of chunks) {
+    signal.throwIfAborted()
+    const chunkBytes = Buffer.byteLength(chunk.content)
+    if (chunkBytes > FILE_SEARCH_CHUNK_BYTES)
+      throw new Error('File search chunk exceeds its budget')
+    const chunkKeys = estimateTrigramKeys(chunk.content)
+    if (
+      batch.length &&
+      exceedsFileSearchBatchBudget(batch.length + 1, bytes + chunkBytes, keys + chunkKeys)
+    ) {
+      yield batch
+      signal.throwIfAborted()
+      batch = []
+      bytes = keys = 0
+    }
+    batch.push(chunk)
+    bytes += chunkBytes
+    keys += chunkKeys
+  }
+  signal.throwIfAborted()
+  if (batch.length) yield batch
+}

apps/sim/lib/workspace-files/search/index-state.ts

@@ -10,16 +10,16 @@ import { and, eq, sql } from 'drizzle-orm'
 import type { DbTransaction } from '@/lib/db/types'
 import {
   FILE_SEARCH_BUILD_LEASE_MS,
+  FILE_SEARCH_CHUNK_BYTES,
   FILE_SEARCH_CLEANUP_BATCH_BUILDS,
   FILE_SEARCH_CLEANUP_BATCH_ROWS,
   FILE_SEARCH_CLEANUP_BUDGET_MS,
   FILE_SEARCH_CLEANUP_MAX_BATCHES,
   FILE_SEARCH_CLEANUP_MIN_BATCH_MS,
   FILE_SEARCH_INDEX_TRANSACTION_LIMITS,
-  FILE_SEARCH_INSERT_BATCH_BYTES,
-  FILE_SEARCH_INSERT_BATCH_ROWS,
 } from '@/lib/workspace-files/search/constants'
-import type { FileSearchChunk } from '@/lib/workspace-files/search/index-plan'
+import { exceedsFileSearchBatchBudget } from '@/lib/workspace-files/search/index-batches'
+import { estimateTrigramKeys, type FileSearchChunk } from '@/lib/workspace-files/search/index-plan'
 import { configureFileSearchTransaction } from '@/lib/workspace-files/search/transaction'
 
 export interface FileSearchRevision {
@@ -139,7 +139,7 @@ export async function beginFileSearchBuild(
   })
 }
 
-/** Each batch is fenced and byte-bounded; no file or parser work runs inside this transaction. */
+/** Each batch is fenced and work-bounded; no file or parser work runs inside this transaction. */
 export async function appendFileSearchChunks(
   build: FileSearchBuild,
   chunks: readonly FileSearchChunk[],
@@ -148,9 +148,12 @@ export async function appendFileSearchChunks(
   signal.throwIfAborted()
   if (!chunks.length) return true
   if (
-    chunks.length > FILE_SEARCH_INSERT_BATCH_ROWS ||
-    chunks.reduce((sum, c) => sum + Buffer.byteLength(c.content), 0) >
-      FILE_SEARCH_INSERT_BATCH_BYTES
+    chunks.some((chunk) => Buffer.byteLength(chunk.content) > FILE_SEARCH_CHUNK_BYTES) ||
+    exceedsFileSearchBatchBudget(
+      chunks.length,
+      chunks.reduce((sum, chunk) => sum + Buffer.byteLength(chunk.content), 0),
+      chunks.reduce((sum, chunk) => sum + estimateTrigramKeys(chunk.content), 0)
+    )
   ) {
     throw new Error('File search insert batch exceeds its budget')
   }

apps/sim/lib/workspace-files/search/indexing.ts

@@ -5,12 +5,11 @@ import { redactDatabaseQueryError } from '@/lib/core/errors/database-query-error
 import { isPayloadSizeLimitError } from '@/lib/core/utils/stream-limits'
 import { getWorkspaceFile } from '@/lib/uploads/contexts/workspace'
 import {
-  FILE_SEARCH_INSERT_BATCH_BYTES,
-  FILE_SEARCH_INSERT_BATCH_ROWS,
   FILE_SEARCH_MAX_SOURCE_BYTES,
   FILE_SEARCH_SLOW_INSERT_BATCH_MS,
 } from '@/lib/workspace-files/search/constants'
 import { extractIndexText, loadIndexableBytes } from '@/lib/workspace-files/search/extract'
+import { iterateFileSearchBatches } from '@/lib/workspace-files/search/index-batches'
 import {
   estimateTrigramKeys,
   type FileSearchChunk,
@@ -46,7 +45,7 @@ function parseRevision(payload: WorkspaceFileSearchIndexPayload): FileSearchRevi
 
 /**
  * Appends one batch and records slow ones, including a batch a statement timeout cancels, with the
- * trigram key load that drives direct GIN insert cost. Keys are estimated only for slow batches.
+ * trigram key load that drives direct GIN insert cost. Logging never includes the indexed text.
  */
 async function appendTimedBatch(
   build: FileSearchBuild,
@@ -107,25 +106,12 @@ export async function indexWorkspaceFileForSearch(
       return
     }
     const plan = planFileSearchIndex(extracted, signal)
-    let batch: FileSearchChunk[] = []
-    let batchBytes = 0
     let chunkCount = 0
-    for (const chunk of iterateFileSearchChunks(plan, signal)) {
-      const chunkBytes = Buffer.byteLength(chunk.content, 'utf8')
-      if (
-        batch.length &&
-        (batch.length >= FILE_SEARCH_INSERT_BATCH_ROWS ||
-          batchBytes + chunkBytes > FILE_SEARCH_INSERT_BATCH_BYTES)
-      ) {
-        if (!(await appendTimedBatch(build, batch, batchBytes, signal))) return
-        batch = []
-        batchBytes = 0
-      }
-      batch.push(chunk)
-      batchBytes += chunkBytes
-      chunkCount++
+    for (const batch of iterateFileSearchBatches(iterateFileSearchChunks(plan, signal), signal)) {
+      const batchBytes = batch.reduce((sum, chunk) => sum + Buffer.byteLength(chunk.content), 0)
+      if (!(await appendTimedBatch(build, batch, batchBytes, signal))) return
+      chunkCount += batch.length
     }
-    if (!(await appendTimedBatch(build, batch, batchBytes, signal))) return
     const published = await publishFileSearchBuild(
       build,
       { status: 'ready', chunkCount, lineCount: plan.lineCount, indexedBytes: plan.indexedBytes },