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.

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,28 @@ 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('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,11 @@ 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++) {
+    const remainingBudget = deadline - Date.now()
+    if (remainingBudget < FILE_SEARCH_CLEANUP_MIN_BATCH_MS) break
     const result = await db.transaction(async (tx) => {
-      await configureFileSearchTransaction(tx, {
-        statementTimeout: Math.max(1, deadline - Date.now()),
-      })
+      await configureFileSearchTransaction(tx, { statementTimeout: remainingBudget })
       const builds = await tx.execute<{
         id: string
       }>(sql`SELECT id FROM workspace_file_search_build