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

@@ -396,6 +396,29 @@ describe('chunked workspace file search on PostgreSQL', () => {
       (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/index-state.ts

@@ -259,9 +259,11 @@ 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; batch++) {
-    const remainingBudget = deadline - Date.now()
-    if (remainingBudget < FILE_SEARCH_CLEANUP_MIN_BATCH_MS) break
+    if (deadline - Date.now() < FILE_SEARCH_CLEANUP_MIN_BATCH_MS) break
     const result = await db.transaction(async (tx) => {
+      /** 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