fix(file-search): stop cleanup starting a batch it cannot fund (#7965)

* fix(file-search): stop cleanup starting a batch it cannot fund

The batch loop admitted another batch whenever any budget remained, then passed
that remainder through as the batch's statement timeout, clamped to 1ms. A batch
admitted with a sliver left either runs past the budget it was given or aborts on
its own statement timeout, which the caller reports as a cleanup failure rather
than as work still to do.

Stop once less than one batch's nominal share of the budget remains. The floor is
derived from the budget and the batch cap rather than fixed, so retuning either
cannot leave it admitting batches at more than their share again.

* fix(file-search): measure the cleanup budget after the connection is held

Hoisting the remaining-budget read out of the transaction callback made it
describe the moment the batch was admitted rather than the moment its statements
begin. Time spent waiting for a pooled connection then went unaccounted, and the
batch installed a timeout larger than the budget actually left.

Keep the cheap check before opening a transaction, and re-read once the
connection is in hand so the installed timeout is the budget that remains.

Author: waleedlatif1

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

@@ -25,6 +25,7 @@ vi.mock('@/lib/file-parsers', () => ({ parseBuffer: vi.fn(), isSupportedFileType
 
 import {
   FILE_SEARCH_CLEANUP_BATCH_ROWS,
+  FILE_SEARCH_CLEANUP_BUDGET_MS,
   FILE_SEARCH_CLEANUP_MAX_BATCHES,
 } from '@/lib/workspace-files/search/constants'
 import { prepareWorkspaceFileSearchDispatch } from '@/lib/workspace-files/search/dispatcher'
@@ -373,6 +374,51 @@ describe('chunked workspace file search on PostgreSQL', () => {
       (await connection`SELECT count(*)::int AS count FROM workspace_file_search_chunk`)[0].count
     ).toBe(0)
   })
+  it('ends a cleanup run cleanly when its time budget runs out', async () => {
+    const build = (await beginFileSearchBuild(revision))!
+    await connection`INSERT INTO workspace_file_search_chunk (build_id, workspace_id, ordinal, line_start, fragment, content)
+      SELECT ${build.id}, 'workspace-1', n, n + 1, false, 'x' FROM generate_series(0, 999) n`
+    await connection`UPDATE workspace_file_search_build SET expires_at = now() WHERE id = ${build.id}`
+
+    /** The deadline is read first; every read after it reports a budget all but consumed. */
+    const startedAt = Date.now()
+    const clock = vi
+      .spyOn(Date, 'now')
+      .mockReturnValueOnce(startedAt)
+      .mockReturnValue(startedAt + FILE_SEARCH_CLEANUP_BUDGET_MS - 1)
+    try {
+      await expect(cleanupFileSearchBuilds()).resolves.toBe(0)
+    } finally {
+      clock.mockRestore()
+    }
+
+    expect(
+      (await connection`SELECT count(*)::int AS count FROM workspace_file_search_chunk`)[0].count
+    ).toBe(1000)
+  })
+  it('abandons a batch whose budget was spent acquiring its connection', async () => {
+    const build = (await beginFileSearchBuild(revision))!
+    await connection`INSERT INTO workspace_file_search_chunk (build_id, workspace_id, ordinal, line_start, fragment, content)
+      SELECT ${build.id}, 'workspace-1', n, n + 1, false, 'x' FROM generate_series(0, 999) n`
+    await connection`UPDATE workspace_file_search_build SET expires_at = now() WHERE id = ${build.id}`
+
+    /** Full budget when the batch is admitted, none left once its connection is in hand. */
+    const startedAt = Date.now()
+    const clock = vi
+      .spyOn(Date, 'now')
+      .mockReturnValueOnce(startedAt)
+      .mockReturnValueOnce(startedAt)
+      .mockReturnValue(startedAt + FILE_SEARCH_CLEANUP_BUDGET_MS - 1)
+    try {
+      await expect(cleanupFileSearchBuilds()).resolves.toBe(0)
+    } finally {
+      clock.mockRestore()
+    }
+
+    expect(
+      (await connection`SELECT count(*)::int AS count FROM workspace_file_search_chunk`)[0].count
+    ).toBe(1000)
+  })
   it('retires many small builds within one cleanup run', async () => {
     await connection`INSERT INTO workspace_file_search_build (id, file_id, workspace_id, source_content_updated_at, expires_at)
       SELECT 'retired-' || n, 'file-1', 'workspace-1', now(), now() FROM generate_series(1, 100) n`

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

@@ -32,6 +32,16 @@ export const FILE_SEARCH_CLEANUP_BATCH_BUILDS = 100
 export const FILE_SEARCH_CLEANUP_BACKLOG_ROWS = 10000
 export const FILE_SEARCH_CLEANUP_MAX_BATCHES = 10
 export const FILE_SEARCH_CLEANUP_BUDGET_MS = 5000
+/**
+ * One batch's nominal share of the run budget, and so the smallest slice worth starting another
+ * with.
+ *
+ * Running out of budget is how cleanup normally ends. A batch admitted with less than its share
+ * either runs past the budget it was given or aborts on its own statement timeout, which the caller
+ * reports as a cleanup failure rather than as work still to do.
+ */
+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
 export const FILE_SEARCH_INSERT_BATCH_BYTES = 1024 * 1024

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

@@ -14,6 +14,7 @@ import {
   FILE_SEARCH_CLEANUP_BATCH_ROWS,
   FILE_SEARCH_CLEANUP_BUDGET_MS,
   FILE_SEARCH_CLEANUP_MAX_BATCHES,
+  FILE_SEARCH_CLEANUP_MIN_BATCH_MS,
   FILE_SEARCH_INSERT_BATCH_BYTES,
   FILE_SEARCH_INSERT_BATCH_ROWS,
 } from '@/lib/workspace-files/search/constants'
@@ -257,11 +258,13 @@ export async function failFileSearchRevision(
 export async function cleanupFileSearchBuilds(): Promise<number> {
   const deadline = Date.now() + FILE_SEARCH_CLEANUP_BUDGET_MS
   let deleted = 0
-  for (let batch = 0; batch < FILE_SEARCH_CLEANUP_MAX_BATCHES && Date.now() < deadline; batch++) {
+  for (let batch = 0; batch < FILE_SEARCH_CLEANUP_MAX_BATCHES; batch++) {
+    if (deadline - Date.now() < FILE_SEARCH_CLEANUP_MIN_BATCH_MS) break
     const result = await db.transaction(async (tx) => {
-      await configureFileSearchTransaction(tx, {
-        statementTimeout: Math.max(1, deadline - Date.now()),
-      })
+      /** Re-read: acquiring the connection can itself have spent the rest of the budget. */
+      const remainingBudget = deadline - Date.now()
+      if (remainingBudget < FILE_SEARCH_CLEANUP_MIN_BATCH_MS) return null
+      await configureFileSearchTransaction(tx, { statementTimeout: remainingBudget })
       const builds = await tx.execute<{
         id: string
       }>(sql`SELECT id FROM workspace_file_search_build