fix(file-search): bound admission waits and propagate cancellation

Author: icecrasher321

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

@@ -35,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.
@@ -56,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,7 @@ 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 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. 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. The ten-second execution deadline starts after queueing; connecting to the database additionally uses the driver connection timeout. 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.
+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.
 

apps/sim/lib/workspace-files/search/admission.test.ts

@@ -7,8 +7,13 @@ describe('FileSearchAdmission', () => {
   beforeEach(() => vi.useFakeTimers())
   afterEach(() => vi.useRealTimers())
 
-  function createAdmission(concurrency = 1, maxPending = 3) {
-    return new FileSearchAdmission({ concurrency, maxPending, timeoutMs: 5000 })
+  function createAdmission(concurrency = 1, maxPending = 3, maxPendingPerWorkspace = maxPending) {
+    return new FileSearchAdmission({
+      concurrency,
+      maxPending,
+      maxPendingPerWorkspace,
+      timeoutMs: 5000,
+    })
   }
 
   it('queues a burst without exceeding the active budget and releases each slot once', async () => {
@@ -129,4 +134,79 @@ describe('FileSearchAdmission', () => {
     first()
     await next
   })
+  it('leaves waiting capacity for another workspace when one burst hits its own cap', async () => {
+    const admission = createAdmission(1, 4, 2)
+    const release = await admission.acquire('hot')
+    const hot = [admission.acquire('hot'), admission.acquire('hot')]
+    await expect(admission.acquire('hot')).rejects.toBeInstanceOf(
+      WorkspaceFileSearchUnavailableError
+    )
+    const quiet = admission.acquire('quiet')
+    release()
+    ;(await hot[0])()
+    ;(await quiet)()
+    ;(await hot[1])()
+    ;(await admission.acquire('hot'))()
+    expect(vi.getTimerCount()).toBe(0)
+  })
+
+  it.each(['timeout', 'abort'] as const)(
+    'keeps a stalled operation counted after caller %s until the operation settles',
+    async (reason) => {
+      const admission = createAdmission()
+      const controller = new AbortController()
+      const dependency = Promise.withResolvers<void>()
+      const lateWork = vi.fn()
+      const running = admission.run(
+        'a',
+        async (signal) => {
+          await dependency.promise
+          signal.throwIfAborted()
+          lateWork()
+        },
+        controller.signal
+      )
+      const rejected = expect(running).rejects.toThrow(
+        reason === 'timeout' ? /timed out/ : /cancelled/
+      )
+      await vi.advanceTimersByTimeAsync(1)
+      if (reason === 'timeout') await vi.advanceTimersByTimeAsync(15000)
+      else controller.abort(new Error('cancelled'))
+      await rejected
+      const nextWork = vi.fn().mockResolvedValue('done')
+      const next = admission.run('b', nextWork)
+      await vi.advanceTimersByTimeAsync(1)
+      expect(nextWork).not.toHaveBeenCalled()
+      dependency.resolve()
+      await expect(next).resolves.toBe('done')
+      expect(lateWork).not.toHaveBeenCalled()
+      expect(vi.getTimerCount()).toBe(0)
+    }
+  )
+
+  it('releases a lease after an operation throws synchronously', async () => {
+    const admission = createAdmission()
+    await expect(
+      admission.run('a', () => {
+        throw new Error('operation failed')
+      })
+    ).rejects.toThrow('operation failed')
+    await expect(admission.run('b', async () => 'done')).resolves.toBe('done')
+    expect(vi.getTimerCount()).toBe(0)
+  })
+
+  it('does not start an operation cancelled while waiting', async () => {
+    const admission = createAdmission()
+    const release = await admission.acquire('a')
+    const controller = new AbortController()
+    const operation = vi.fn()
+    const waiting = expect(admission.run('b', operation, controller.signal)).rejects.toThrow(
+      'cancelled'
+    )
+    controller.abort(new Error('cancelled'))
+    await waiting
+    release()
+    expect(operation).not.toHaveBeenCalled()
+    expect(vi.getTimerCount()).toBe(0)
+  })
 })

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

@@ -1,7 +1,10 @@
 import { DB_POOL_PROFILES } from '@sim/db/pool-profiles'
+import { DeadlineExceededError, withinDeadline } from '@/lib/core/utils/deadline'
 import {
+  FILE_SEARCH_QUERY_WORKSPACE_CONCURRENCY,
   FILE_SEARCH_QUEUE_MAX_PENDING,
   FILE_SEARCH_QUEUE_TIMEOUT_MS,
+  FILE_SEARCH_STATEMENT_TIMEOUT_MS,
 } from '@/lib/workspace-files/search/constants'
 import { WorkspaceFileSearchUnavailableError } from '@/lib/workspace-files/search/errors'
 
@@ -19,13 +22,56 @@ export class FileSearchAdmission {
   private readonly workspaces = new Map<string, Set<Waiter>>()
 
   constructor(
-    private readonly options: { concurrency: number; maxPending: number; timeoutMs: number }
+    private readonly options: {
+      concurrency: number
+      maxPending: number
+      maxPendingPerWorkspace: number
+      timeoutMs: number
+    }
   ) {}
 
+  /**
+   * One caller deadline includes queueing, connection acquisition, and execution.
+   * The lease belongs to the underlying operation, even if its caller stops waiting.
+   */
+  async run<T>(
+    workspaceId: string,
+    operation: (signal: AbortSignal, deadlineAt: number) => Promise<T>,
+    signal?: AbortSignal
+  ): Promise<T> {
+    const deadlineAt = Date.now() + this.options.timeoutMs + FILE_SEARCH_STATEMENT_TIMEOUT_MS
+    try {
+      return await withinDeadline(
+        async (operationSignal) => {
+          const release = await this.acquire(workspaceId, operationSignal)
+          try {
+            operationSignal.throwIfAborted()
+            return await operation(operationSignal, deadlineAt)
+          } finally {
+            release()
+          }
+        },
+        deadlineAt,
+        signal
+      )
+    } catch (error) {
+      signal?.throwIfAborted()
+      if (error instanceof DeadlineExceededError) {
+        throw new WorkspaceFileSearchUnavailableError(
+          'Workspace file search timed out. Retry shortly.'
+        )
+      }
+      throw error
+    }
+  }
+
   async acquire(workspaceId: string, signal?: AbortSignal): Promise<() => void> {
     signal?.throwIfAborted()
     if (this.active < this.options.concurrency) return this.claim()
-    if (this.pending >= this.options.maxPending) {
+    if (
+      this.pending >= this.options.maxPending ||
+      (this.workspaces.get(workspaceId)?.size ?? 0) >= this.options.maxPendingPerWorkspace
+    ) {
       throw new WorkspaceFileSearchUnavailableError('Workspace file search is busy. Retry shortly.')
     }
 
@@ -92,5 +138,6 @@ export class FileSearchAdmission {
 export const fileSearchAdmission = new FileSearchAdmission({
   concurrency: DB_POOL_PROFILES.search.primaryMax,
   maxPending: FILE_SEARCH_QUEUE_MAX_PENDING,
+  maxPendingPerWorkspace: FILE_SEARCH_QUERY_WORKSPACE_CONCURRENCY,
   timeoutMs: FILE_SEARCH_QUEUE_TIMEOUT_MS,
 })

apps/sim/lib/workspace-files/search/repository.test.ts

@@ -3,7 +3,7 @@
  */
 import { db } from '@sim/db'
 import { dbChainMockFns, resetDbChainMock } from '@sim/testing'
-import { beforeEach, describe, expect, it } from 'vitest'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import { WorkspaceFileSearchUnavailableError } from '@/lib/workspace-files/search/errors'
 import {
   compileFileSearchPattern,
@@ -22,6 +22,7 @@ function driverError(code: string): Error {
 }
 
 describe('searchWorkspaceFileIndex fault mapping', () => {
+  afterEach(() => vi.useRealTimers())
   beforeEach(() => {
     resetDbChainMock()
   })
@@ -96,4 +97,26 @@ describe('searchWorkspaceFileIndex fault mapping', () => {
     expect(guards).toContain('lock_timeout')
     expect(db.transaction).toHaveBeenCalled()
   })
+  it('bounds the caller while BEGIN is stalled and does not run late search SQL', async () => {
+    vi.useFakeTimers()
+    const acquired = Promise.withResolvers<void>()
+    let transactionFinished!: Promise<unknown>
+    dbChainMockFns.transaction.mockImplementationOnce((callback) => {
+      transactionFinished = acquired.promise.then(() => callback(db))
+      return transactionFinished
+    })
+    const waiting = expect(
+      searchWorkspaceFileIndex({
+        workspaceId: 'workspace-1',
+        pattern: compileFileSearchPattern('needle', 'exact'),
+        maxResults: 50,
+      })
+    ).rejects.toBeInstanceOf(WorkspaceFileSearchUnavailableError)
+    await vi.advanceTimersByTimeAsync(15000)
+    await waiting
+    expect(dbChainMockFns.execute).not.toHaveBeenCalled()
+    acquired.resolve()
+    await expect(transactionFinished).rejects.toThrow('Operation deadline expired')
+    expect(dbChainMockFns.execute).not.toHaveBeenCalled()
+  })
 })