fix(file-search): bound dispatcher database work

Author: icecrasher321

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

@@ -42,6 +42,10 @@ export const FILE_SEARCH_INDEX_MAX_OUTSTANDING = 100
 export const FILE_SEARCH_INDEX_DISPATCH_WORKSPACES = 100
 export const FILE_SEARCH_DISPATCH_INTERVAL_MS = 60 * 1000
 export const FILE_SEARCH_DISPATCH_MAX_DURATION_SECONDS = 60
+/** Leave room for connection setup, rollback, and task failure reporting before the hard cutoff. */
+export const FILE_SEARCH_DISPATCH_STATEMENT_TIMEOUT_MS = 10 * 1000
+export const FILE_SEARCH_DISPATCH_LOCK_TIMEOUT_MS = 2 * 1000
+export const FILE_SEARCH_DISPATCH_TRANSACTION_TIMEOUT_MS = 20 * 1000
 export const FILE_SEARCH_INDEX_MAX_DURATION_SECONDS = 15 * 60
 export const FILE_SEARCH_INDEX_STALE_DISPATCH_MS = 6 * 60 * 60 * 1000
 export const FILE_SEARCH_INDEX_STALE_REAP_LIMIT = 100

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

@@ -0,0 +1,111 @@
+/** Real PostgreSQL cancellation must roll back preparation and release its advisory lock. */
+import { withUtcTimestamps } from '@sim/db/timestamps'
+import { getPostgresErrorCode } from '@sim/utils/errors'
+import { generateId } from '@sim/utils/id'
+import { drizzle, type PostgresJsDatabase } from 'drizzle-orm/postgres-js'
+import postgres from 'postgres'
+import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'
+
+const database = vi.hoisted(() => ({ current: undefined as PostgresJsDatabase | undefined }))
+vi.mock('@sim/db', () => ({
+  get db() {
+    if (!database.current) throw new Error('Dispatcher test database is not initialized')
+    return database.current
+  },
+}))
+vi.mock('@/lib/workspace-files/search/indexing', () => ({
+  indexWorkspaceFileForSearch: vi.fn(),
+  markWorkspaceFileSearchIndexFailed: vi.fn(),
+}))
+
+import { prepareWorkspaceFileSearchDispatch } from '@/lib/workspace-files/search/dispatcher'
+
+describe('workspace file search dispatch PostgreSQL deadlines', () => {
+  const schemaName = `dispatch_test_${generateId().replaceAll('-', '')}`
+  const databaseUrl = process.env.KNOWLEDGE_ACL_TEST_DATABASE_URL
+  if (!databaseUrl) throw new Error('Dispatcher tests require a disposable local database')
+  const connection = postgres(
+    databaseUrl,
+    withUtcTimestamps({
+      max: 3,
+      prepare: false,
+      fetch_types: false,
+      connection: { search_path: schemaName },
+      onnotice: () => {},
+    })
+  )
+
+  beforeAll(async () => {
+    await connection`CREATE SCHEMA ${connection(schemaName)}`
+    await connection`CREATE TABLE workspace_file_search_backfill (
+      id text PRIMARY KEY, after_workspace_id text, after_file_id text,
+      completed_at timestamp, updated_at timestamp NOT NULL
+    )`
+    await connection`INSERT INTO workspace_file_search_backfill (id, updated_at)
+      VALUES ('workspace-file-search-v1', '2026-09-16 00:00:00')`
+    database.current = drizzle(connection)
+  })
+
+  afterAll(async () => {
+    try {
+      await connection`DROP SCHEMA ${connection(schemaName)} CASCADE`
+    } finally {
+      await connection.end()
+      database.current = undefined
+    }
+  })
+
+  async function expectAdvisoryLockReleased() {
+    await connection.begin(async (tx) => {
+      const [row] = await tx`SELECT pg_try_advisory_xact_lock(
+        hashtextextended('workspace-file-search-dispatch', 0)
+      ) AS acquired`
+      expect(row.acquired).toBe(true)
+    })
+  }
+
+  it('fails on a locked backfill row and releases the dispatcher lock', async () => {
+    let release = () => {}
+    let locked = () => {}
+    const releaseLock = new Promise<void>((resolve) => {
+      release = resolve
+    })
+    const lockReady = new Promise<void>((resolve) => {
+      locked = resolve
+    })
+    const blocker = connection.begin(async (tx) => {
+      await tx`SELECT id FROM workspace_file_search_backfill FOR UPDATE`
+      locked()
+      await releaseLock
+    })
+    await lockReady
+    try {
+      const failure = await prepareWorkspaceFileSearchDispatch().catch((error: unknown) => error)
+      expect(getPostgresErrorCode(failure)).toBe('55P03')
+      await expectAdvisoryLockReleased()
+    } finally {
+      release()
+      await blocker
+    }
+  })
+
+  it('cancels a slow statement and rolls back its earlier writes', async () => {
+    await connection`CREATE FUNCTION slow_backfill() RETURNS trigger LANGUAGE plpgsql AS $$
+      BEGIN
+        UPDATE workspace_file_search_backfill SET updated_at = '2099-01-01';
+        PERFORM pg_sleep(15);
+        RETURN NEW;
+      END
+    $$`
+    await connection`CREATE TRIGGER slow_backfill BEFORE INSERT ON workspace_file_search_backfill
+      FOR EACH ROW EXECUTE FUNCTION slow_backfill()`
+
+    const failure = await prepareWorkspaceFileSearchDispatch().catch((error: unknown) => error)
+
+    expect(getPostgresErrorCode(failure)).toBe('57014')
+    const [row] =
+      await connection`SELECT updated_at::text AS updated_at FROM workspace_file_search_backfill`
+    expect(row.updated_at).toBe('2026-09-16 00:00:00')
+    await expectAdvisoryLockReleased()
+  }, 20_000)
+})

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

@@ -1,9 +1,43 @@
 /**
  * @vitest-environment node
  */
-import { describe, expect, it } from 'vitest'
+import {
+  workspaceFileSearchBackfill,
+  workspaceFileSearchDispatchQueue,
+  workspaceFileSearchIndex,
+} from '@sim/db/schema'
+import { dbChainMock, dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const mocks = vi.hoisted(() => ({
+  batchTrigger: vi.fn(),
+  info: vi.fn(),
+  error: vi.fn(),
+}))
+
+vi.mock('@sim/db/schema', async () => ({
+  ...(await import('@sim/testing/mocks/schema.mock')).schemaMock,
+  workspaceFileSearchBackfill: { id: 'backfill.id' },
+  workspaceFileSearchDispatchQueue: {
+    workspaceId: 'queue.workspaceId',
+    lastDispatchedAt: 'queue.lastDispatchedAt',
+    enqueuedAt: 'queue.enqueuedAt',
+  },
+}))
+
+vi.mock('@sim/logger', () => ({ createLogger: () => ({ info: mocks.info, error: mocks.error }) }))
+vi.mock('@trigger.dev/sdk', () => ({ tasks: { batchTrigger: mocks.batchTrigger } }))
+vi.mock('@/lib/core/config/env-flags', () => ({ isTriggerDevEnabled: true }))
+vi.mock('@/lib/core/async-jobs/region', () => ({ resolveTriggerRegion: async () => 'us-east-1' }))
+vi.mock('@/lib/workspace-files/search/indexing', () => ({
+  indexWorkspaceFileForSearch: vi.fn(),
+  markWorkspaceFileSearchIndexFailed: vi.fn(),
+}))
+
 import {
   buildWorkspaceFileSearchTriggerItems,
+  dispatchWorkspaceFileSearchIndexJobs,
+  prepareWorkspaceFileSearchDispatch,
   shouldUseWorkspaceFileSearchTrigger,
 } from '@/lib/workspace-files/search/dispatcher'
 
@@ -34,3 +68,102 @@ describe('workspace file search dispatch policy', () => {
     ])
   })
 })
+
+describe('workspace file search dispatch deadlines', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    resetDbChainMock()
+  })
+
+  it('sets local database deadlines before taking the advisory lock', async () => {
+    dbChainMockFns.execute.mockResolvedValueOnce([]).mockResolvedValueOnce([{ acquired: false }])
+
+    await expect(prepareWorkspaceFileSearchDispatch()).resolves.toEqual({
+      payloads: [],
+      backfilledFiles: 0,
+      reapedClaims: 0,
+      lockAcquired: false,
+    })
+
+    const guards = JSON.stringify(dbChainMockFns.execute.mock.calls[0][0])
+    expect(guards).toContain("set_config('statement_timeout', ")
+    expect(guards).toContain('10000ms')
+    expect(guards).toContain("set_config('lock_timeout', ")
+    expect(guards).toContain('2000ms')
+    expect(guards).toContain("set_config('transaction_timeout', ")
+    expect(guards).toContain('20000ms')
+    expect(guards.match(/, true\)/g)).toHaveLength(3)
+    expect(JSON.stringify(dbChainMockFns.execute.mock.calls[1][0])).toContain(
+      'pg_try_advisory_xact_lock'
+    )
+    expect(dbChainMockFns.insert).not.toHaveBeenCalled()
+  })
+
+  it.each(['57014', '55P03', '25P04'])(
+    'propagates SQLSTATE %s without enqueuing an uncommitted claim',
+    async (code) => {
+      const error = new Error('Failed query\nparams: sensitive-value', {
+        cause: Object.assign(new Error('database timeout'), { code }),
+      })
+      dbChainMockFns.execute.mockResolvedValueOnce([]).mockResolvedValueOnce([{ acquired: true }])
+      dbChainMockFns.onConflictDoNothing.mockRejectedValueOnce(error)
+
+      await expect(dispatchWorkspaceFileSearchIndexJobs()).rejects.toBe(error)
+
+      expect(mocks.batchTrigger).not.toHaveBeenCalled()
+      expect(mocks.error).toHaveBeenCalledWith('Workspace file search dispatch phase failed', {
+        phase: 'backfill',
+        durationMs: expect.any(Number),
+        code,
+        error: 'Failed query',
+      })
+      expect(JSON.stringify(mocks.error.mock.calls)).not.toContain('sensitive-value')
+    }
+  )
+
+  it('reports a transaction failure even after the transaction callback finishes', async () => {
+    const error = Object.assign(new Error('commit failed'), { code: '08006' })
+    dbChainMockFns.execute.mockResolvedValueOnce([]).mockResolvedValueOnce([{ acquired: false }])
+    dbChainMockFns.transaction.mockImplementationOnce(async (callback) => {
+      await callback(dbChainMock.db)
+      throw error
+    })
+
+    await expect(prepareWorkspaceFileSearchDispatch()).rejects.toBe(error)
+
+    expect(mocks.error).toHaveBeenCalledWith('Workspace file search dispatch phase failed', {
+      phase: 'prepare-transaction',
+      durationMs: expect.any(Number),
+      code: '08006',
+      error: 'commit failed',
+    })
+  })
+
+  it('also bounds releasing committed claims when batch submission fails', async () => {
+    queueTableRows(workspaceFileSearchBackfill, [{ completedAt: new Date() }])
+    queueTableRows(workspaceFileSearchIndex, [])
+    queueTableRows(workspaceFileSearchIndex, [{ active: 0 }])
+    queueTableRows(workspaceFileSearchDispatchQueue, [{ workspaceId: 'workspace-1' }])
+    dbChainMockFns.execute
+      .mockResolvedValueOnce([])
+      .mockResolvedValueOnce([{ acquired: true }])
+      .mockResolvedValueOnce([
+        {
+          workspaceId: 'workspace-1',
+          fileId: 'file-1',
+          sourceContentUpdatedAt: new Date('2026-09-16T00:00:00Z'),
+        },
+      ])
+    const error = new Error('Trigger unavailable')
+    mocks.batchTrigger.mockRejectedValueOnce(error)
+
+    await expect(dispatchWorkspaceFileSearchIndexJobs()).rejects.toBe(error)
+
+    expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(2)
+    const guards = dbChainMockFns.execute.mock.calls.filter(([query]) =>
+      JSON.stringify(query).includes('statement_timeout')
+    )
+    expect(guards).toHaveLength(2)
+    expect(dbChainMockFns.set).toHaveBeenCalledWith(expect.objectContaining({ dispatchedAt: null }))
+  })
+})