From a4a903cdfb0f3c77ad87117416f976fbf3ab1f9f Mon Sep 17 00:00:00 2001 From: Waleed Date: Tue, 22 Sep 2026 22:23:19 -0700 Subject: [PATCH 01/13] fix(file-search): back off indexing retries after database timeouts (#8189) A statement, lock, or transaction timeout while indexing a workspace file means the database had no capacity for the build, not that the file is bad. The task retried it within seconds, so all three attempts landed in the same slow window and the revision was marked failed for good. Capacity timeouts now retry after about 2, 4, 8, 16 and 30 minutes (six attempts); other failures keep three attempts with the default delays. --- .../workspace-file-search-index.test.ts | 16 +++- .../background/workspace-file-search-index.ts | 6 +- apps/sim/lib/workspace-files/search/README.md | 2 +- .../lib/workspace-files/search/constants.ts | 11 +++ .../workspace-files/search/indexing.test.ts | 85 +++++++++++++++++-- .../lib/workspace-files/search/indexing.ts | 37 +++++++- 6 files changed, 148 insertions(+), 9 deletions(-) diff --git a/apps/sim/background/workspace-file-search-index.test.ts b/apps/sim/background/workspace-file-search-index.test.ts index b261a0e3b7f..1da3dc564b6 100644 --- a/apps/sim/background/workspace-file-search-index.test.ts +++ b/apps/sim/background/workspace-file-search-index.test.ts @@ -5,6 +5,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' const mocks = vi.hoisted(() => ({ indexWorkspaceFile: vi.fn(), + retry: vi.fn(), markFailed: vi.fn(), task: vi.fn((config: unknown) => config), })) @@ -12,10 +13,12 @@ const mocks = vi.hoisted(() => ({ vi.mock('@trigger.dev/sdk', () => ({ task: mocks.task })) vi.mock('@/lib/workspace-files/search/indexing', () => ({ indexWorkspaceFileForSearch: mocks.indexWorkspaceFile, + getWorkspaceFileSearchRetry: mocks.retry, markWorkspaceFileSearchIndexFailed: mocks.markFailed, })) import { + FILE_SEARCH_INDEX_CAPACITY_MAX_ATTEMPTS, FILE_SEARCH_INDEX_GLOBAL_CONCURRENCY, FILE_SEARCH_INDEX_MAX_DURATION_SECONDS, } from '@/lib/workspace-files/search/constants' @@ -37,7 +40,7 @@ describe('workspace file search index task', () => { id: 'workspace-file-search-index', machine: 'medium-2x', maxDuration: FILE_SEARCH_INDEX_MAX_DURATION_SECONDS, - retry: { maxAttempts: 3 }, + retry: { maxAttempts: FILE_SEARCH_INDEX_CAPACITY_MAX_ATTEMPTS }, queue: { name: 'workspace-file-search-index', concurrencyLimit: FILE_SEARCH_INDEX_GLOBAL_CONCURRENCY, @@ -65,4 +68,15 @@ describe('workspace file search index task', () => { await workspaceFileSearchIndexTask.onFailure({ payload }) expect(mocks.markFailed).toHaveBeenCalledWith(payload) }) + + it('lets the indexing retry policy schedule each failed attempt', async () => { + const error = new Error('statement timeout') + const decision = { retryAt: new Date('2026-08-29T12:02:00.000Z') } + mocks.retry.mockReturnValue(decision) + + await expect( + workspaceFileSearchIndexTask.catchError({ error, ctx: { attempt: { number: 2 } } }) + ).resolves.toBe(decision) + expect(mocks.retry).toHaveBeenCalledWith(error, 2) + }) }) diff --git a/apps/sim/background/workspace-file-search-index.ts b/apps/sim/background/workspace-file-search-index.ts index c44bd47acf3..86340f6ae01 100644 --- a/apps/sim/background/workspace-file-search-index.ts +++ b/apps/sim/background/workspace-file-search-index.ts @@ -1,9 +1,11 @@ import { task } from '@trigger.dev/sdk' import { + FILE_SEARCH_INDEX_CAPACITY_MAX_ATTEMPTS, FILE_SEARCH_INDEX_GLOBAL_CONCURRENCY, FILE_SEARCH_INDEX_MAX_DURATION_SECONDS, } from '@/lib/workspace-files/search/constants' import { + getWorkspaceFileSearchRetry, indexWorkspaceFileForSearch, markWorkspaceFileSearchIndexFailed, type WorkspaceFileSearchIndexPayload, @@ -17,13 +19,15 @@ export const workspaceFileSearchIndexTask = task({ id: 'workspace-file-search-index', machine: 'medium-2x', maxDuration: FILE_SEARCH_INDEX_MAX_DURATION_SECONDS, - retry: { maxAttempts: 3 }, + /** The ceiling for capacity retries; `catchError` stops other failures sooner. */ + retry: { maxAttempts: FILE_SEARCH_INDEX_CAPACITY_MAX_ATTEMPTS }, queue: { name: 'workspace-file-search-index', concurrencyLimit: FILE_SEARCH_INDEX_GLOBAL_CONCURRENCY, }, run: (payload: WorkspaceFileSearchIndexPayload, { signal }) => indexWorkspaceFileForSearch(payload, signal), + catchError: async ({ error, ctx }) => getWorkspaceFileSearchRetry(error, ctx.attempt.number), onFailure: async ({ payload }) => { await markWorkspaceFileSearchIndexFailed(payload) }, diff --git a/apps/sim/lib/workspace-files/search/README.md b/apps/sim/lib/workspace-files/search/README.md index a7fd2201888..0ad529f2c2a 100644 --- a/apps/sim/lib/workspace-files/search/README.md +++ b/apps/sim/lib/workspace-files/search/README.md @@ -14,7 +14,7 @@ Workers download and extract outside database transactions, then insert batches The chunk GIN index uses `fastupdate = off`. Each bounded insert updates the main index directly instead of appending to a shared pending list. With deferred updates enabled, even a small insert can cross the pending-list threshold and synchronously merge accumulated work from other files. Direct updates trade some bulk-write throughput for avoiding that foreground cleanup cliff. They do not eliminate normal index I/O, vacuum, or storage contention; the row and worker limits still apply. Bytes alone do not bound GIN posting updates: dense text can generate thousands of distinct keys in each chunk. The batch planner sums each chunk's distinct trigram count, including repeated keys across rows, and flushes before the next chunk exceeds the key target. A single 8 KiB chunk is always allowed to make progress even if word padding pushes its estimate slightly above the target; it is written alone. Storage validates the same budget before opening the transaction. These are write scheduling bounds, not file exclusions: chunk boundaries, complete-file publication, and the 25 MiB file coverage limit are unchanged. The estimate mirrors pg_trgm under `en_US.UTF-8`; it is a work estimate, not a latency guarantee. Capacity validation must still include dense text, concurrent writers, and an index working set larger than available cache. A batch that takes at least two seconds, including a canceled statement, is logged with its rows, bytes, and estimated key count. A failed query is reported to the task runner by error code only; Drizzle's message carries the bound file text. -Indexing transactions have separate limits from search: ten seconds per statement, five seconds waiting for a lock, and thirty seconds total on PostgreSQL 17. The outer limit leaves time for ordinary statement cancellation and rollback instead of terminating the connection at the same ten-second deadline. PostgreSQL 16 uses the compatible idle-transaction guard. A canceled batch remains unpublished; the existing task retry starts a fresh fenced build, and cleanup retires the previous attempt. This does not automatically retry revisions already marked failed. +Indexing transactions have separate limits from search: ten seconds per statement, five seconds waiting for a lock, and thirty seconds total on PostgreSQL 17. The outer limit leaves time for ordinary statement cancellation and rollback instead of terminating the connection at the same ten-second deadline. PostgreSQL 16 uses the compatible idle-transaction guard. A canceled batch remains unpublished; the task retry starts a fresh fenced build, and cleanup retires the previous attempt. One row's direct GIN insert is not interruptible, so when storage is saturated a single ordinary chunk can run well past the statement deadline and the cancellation lands only after it; smaller batches cannot prevent that. A statement, lock, or transaction timeout is therefore treated as missing database capacity rather than a bad file: the task retries it after about 2, 4, 8, 16, and 30 minutes (with jitter), six attempts in all, so the retries outlast a slow window instead of landing inside it. Other failures keep three attempts with the runner's short default delays. A run waiting to retry still holds one of its workspace's two outstanding dispatch slots. Only a revision that exhausts its attempts is marked failed; this does not automatically retry revisions already marked failed. The indexing task uses an isolated `medium-2x` Trigger worker (4 GB RAM). Document parsers can materialize expanded content before chunking, so source and extracted-text byte limits do not bound parser memory. Parser complexity guards and the worker's memory budget remain separate protections. diff --git a/apps/sim/lib/workspace-files/search/constants.ts b/apps/sim/lib/workspace-files/search/constants.ts index 22bf12bb14e..823795f93a7 100644 --- a/apps/sim/lib/workspace-files/search/constants.ts +++ b/apps/sim/lib/workspace-files/search/constants.ts @@ -78,6 +78,17 @@ export const FILE_SEARCH_INDEX_TRANSACTION_LIMITS = { transactionTimeout: 30 * 1000, } as const +/** Attempts for failures other than database capacity cancellations. */ +export const FILE_SEARCH_INDEX_MAX_ATTEMPTS = 3 +/** + * Attempts when PostgreSQL cancels an indexing statement on a timeout. One row's direct GIN insert + * is not interruptible, so under storage saturation even a single ordinary chunk can outlive the + * statement deadline; smaller batches cannot help, only waiting out the slow window can. + */ +export const FILE_SEARCH_INDEX_CAPACITY_MAX_ATTEMPTS = 6 +/** First capacity retry delay; later ones double up to the ceiling, about an hour in total. */ +export const FILE_SEARCH_INDEX_CAPACITY_RETRY_BASE_MS = 2 * 60 * 1000 +export const FILE_SEARCH_INDEX_CAPACITY_RETRY_MAX_MS = 30 * 60 * 1000 export const FILE_SEARCH_INDEX_GLOBAL_CONCURRENCY = 10 export const FILE_SEARCH_INDEX_WORKSPACE_OUTSTANDING = 2 export const FILE_SEARCH_INDEX_MAX_OUTSTANDING = 100 diff --git a/apps/sim/lib/workspace-files/search/indexing.test.ts b/apps/sim/lib/workspace-files/search/indexing.test.ts index 66761c36886..ed6b6412551 100644 --- a/apps/sim/lib/workspace-files/search/indexing.test.ts +++ b/apps/sim/lib/workspace-files/search/indexing.test.ts @@ -24,13 +24,20 @@ vi.mock('@/lib/workspace-files/search/extract', () => ({ })) import { + FILE_SEARCH_INDEX_CAPACITY_MAX_ATTEMPTS, + FILE_SEARCH_INDEX_CAPACITY_RETRY_BASE_MS, + FILE_SEARCH_INDEX_CAPACITY_RETRY_MAX_MS, + FILE_SEARCH_INDEX_MAX_ATTEMPTS, FILE_SEARCH_INSERT_BATCH_BYTES, FILE_SEARCH_INSERT_BATCH_ROWS, FILE_SEARCH_MAX_SOURCE_BYTES, FILE_SEARCH_SLOW_INSERT_BATCH_MS, } from '@/lib/workspace-files/search/constants' import type { FileSearchChunk } from '@/lib/workspace-files/search/index-plan' -import { indexWorkspaceFileForSearch } from '@/lib/workspace-files/search/indexing' +import { + getWorkspaceFileSearchRetry, + indexWorkspaceFileForSearch, +} from '@/lib/workspace-files/search/indexing' const logger = vi.mocked(createLogger).mock.results[ vi.mocked(createLogger).mock.calls.findIndex(([name]) => name === 'WorkspaceFileSearchIndexer') @@ -38,10 +45,11 @@ const logger = vi.mocked(createLogger).mock.results[ const FILE_TEXT = 'confidential customer text' -function statementTimeout(): DrizzleQueryError { - const driverError = Object.assign(new Error('canceling statement due to statement timeout'), { - code: '57014', - }) +function statementTimeout( + message = 'canceling statement due to statement timeout', + code = '57014' +): DrizzleQueryError { + const driverError = Object.assign(new Error(message), { code }) return new DrizzleQueryError( 'insert into "workspace_file_search_chunk" values ($1)', [FILE_TEXT], @@ -182,3 +190,70 @@ describe('complete-file indexing worker', () => { expect(mocks.publish).not.toHaveBeenCalled() }) }) + +describe('indexing retry policy', () => { + const now = Date.parse('2026-01-01T00:00:00.000Z') + + /** The error the task runner receives: the redacted wrapper the worker throws. */ + async function thrownBy(error: unknown): Promise { + vi.clearAllMocks() + mocks.begin.mockResolvedValue({ id: 'build', ...payload }) + mocks.file.mockResolvedValue({ + name: 'notes.txt', + size: 100, + contentUpdatedAt: new Date(payload.sourceContentUpdatedAt), + }) + mocks.load.mockResolvedValue({ buffer: Buffer.from('a\n') }) + mocks.extract.mockResolvedValue({ text: 'a\n', lineCount: 1 }) + mocks.append.mockRejectedValue(error) + return indexWorkspaceFileForSearch(payload, signal).catch((thrown) => thrown) + } + + function delayOf(decision: ReturnType): number { + if (!decision || !('retryAt' in decision)) throw new Error('expected a scheduled retry') + return decision.retryAt.getTime() - now + } + + it.each([ + ['statement timeout', 'canceling statement due to statement timeout', '57014'], + ['lock timeout', 'canceling statement due to lock timeout', '55P03'], + ])('waits minutes, not seconds, after a %s', async (_label, message, code) => { + const thrown = await thrownBy(statementTimeout(message, code)) + const first = delayOf(getWorkspaceFileSearchRetry(thrown, 1, now)) + expect(first).toBeGreaterThanOrEqual(FILE_SEARCH_INDEX_CAPACITY_RETRY_BASE_MS * 0.8) + expect(first).toBeLessThanOrEqual(FILE_SEARCH_INDEX_CAPACITY_RETRY_BASE_MS * 1.2) + }) + + it('backs capacity retries off to a ceiling and spans a slow window', async () => { + const thrown = await thrownBy(statementTimeout()) + let total = 0 + for (let attempt = 1; attempt < FILE_SEARCH_INDEX_CAPACITY_MAX_ATTEMPTS; attempt++) { + const delay = delayOf(getWorkspaceFileSearchRetry(thrown, attempt, now)) + expect(delay).toBeLessThanOrEqual(FILE_SEARCH_INDEX_CAPACITY_RETRY_MAX_MS * 1.2) + total += delay + } + expect(total).toBeGreaterThanOrEqual(45 * 60 * 1000) + }) + + it('stops capacity retries at their attempt ceiling', async () => { + const thrown = await thrownBy(statementTimeout()) + expect( + getWorkspaceFileSearchRetry(thrown, FILE_SEARCH_INDEX_CAPACITY_MAX_ATTEMPTS, now) + ).toEqual({ skipRetrying: true }) + }) + + it('keeps the short default retries and attempt count for other failures', () => { + const parserFailure = new Error('parser failed') + for (let attempt = 1; attempt < FILE_SEARCH_INDEX_MAX_ATTEMPTS; attempt++) { + expect(getWorkspaceFileSearchRetry(parserFailure, attempt, now)).toBeUndefined() + } + expect(getWorkspaceFileSearchRetry(parserFailure, FILE_SEARCH_INDEX_MAX_ATTEMPTS, now)).toEqual( + { skipRetrying: true } + ) + }) + + it('treats a user cancellation as an ordinary failure', async () => { + const thrown = await thrownBy(statementTimeout('canceling statement due to user request')) + expect(getWorkspaceFileSearchRetry(thrown, 1, now)).toBeUndefined() + }) +}) diff --git a/apps/sim/lib/workspace-files/search/indexing.ts b/apps/sim/lib/workspace-files/search/indexing.ts index 4930595bd88..1c04091bbdd 100644 --- a/apps/sim/lib/workspace-files/search/indexing.ts +++ b/apps/sim/lib/workspace-files/search/indexing.ts @@ -1,10 +1,15 @@ import { Buffer } from 'node:buffer' import { createLogger } from '@sim/logger' -import { describeError } from '@sim/utils/errors' +import { describeError, getPostgresCancellationReason } from '@sim/utils/errors' +import { backoffWithJitter } from '@sim/utils/retry' import { redactDatabaseQueryError } from '@/lib/core/errors/database-query-error' import { isPayloadSizeLimitError } from '@/lib/core/utils/stream-limits' import { getWorkspaceFile } from '@/lib/uploads/contexts/workspace' import { + FILE_SEARCH_INDEX_CAPACITY_MAX_ATTEMPTS, + FILE_SEARCH_INDEX_CAPACITY_RETRY_BASE_MS, + FILE_SEARCH_INDEX_CAPACITY_RETRY_MAX_MS, + FILE_SEARCH_INDEX_MAX_ATTEMPTS, FILE_SEARCH_MAX_SOURCE_BYTES, FILE_SEARCH_SLOW_INSERT_BATCH_MS, } from '@/lib/workspace-files/search/constants' @@ -154,3 +159,33 @@ export async function markWorkspaceFileSearchIndexFailed( return await failFileSearchRevision(parseRevision(payload), payload.dispatchToken) } + +const CAPACITY_CANCELLATIONS = new Set(['statement_timeout', 'lock_timeout', 'transaction_timeout']) + +export type WorkspaceFileSearchRetryDecision = + | { retryAt: Date } + | { skipRetrying: true } + | undefined + +/** + * Chooses the next attempt after `attempt` (1-based) failed. A statement, lock, or transaction + * timeout means the database had no capacity for this build right now, not that the file is bad: + * those back off for minutes so the retries outlast a slow window instead of all landing inside + * it. Anything else keeps the ordinary short retries. `undefined` keeps the runner's default delay. + */ +export function getWorkspaceFileSearchRetry( + error: unknown, + attempt: number, + now = Date.now() +): WorkspaceFileSearchRetryDecision { + const reason = getPostgresCancellationReason(error) + if (reason && CAPACITY_CANCELLATIONS.has(reason)) { + if (attempt >= FILE_SEARCH_INDEX_CAPACITY_MAX_ATTEMPTS) return { skipRetrying: true } + const delayMs = backoffWithJitter(attempt, null, { + baseMs: FILE_SEARCH_INDEX_CAPACITY_RETRY_BASE_MS, + maxMs: FILE_SEARCH_INDEX_CAPACITY_RETRY_MAX_MS, + }) + return { retryAt: new Date(now + delayMs) } + } + return attempt >= FILE_SEARCH_INDEX_MAX_ATTEMPTS ? { skipRetrying: true } : undefined +} From 1c597fec1af4f869c55af46442121b312ec31d51 Mon Sep 17 00:00:00 2001 From: Waleed Date: Tue, 22 Sep 2026 22:42:41 -0700 Subject: [PATCH 02/13] fix(db): retry migration lock timeouts within a time budget instead of eight attempts (#8190) * fix(db): retry migration lock timeouts within a time budget instead of eight attempts * fix(db): measure the migration lock budget on a monotonic clock and never start an attempt past it --- .../db/scripts/lock-timeout-retry.test.ts | 139 ++++++++++++++++++ packages/db/scripts/lock-timeout-retry.ts | 70 +++++++++ packages/db/scripts/migrate.ts | 59 ++++---- 3 files changed, 243 insertions(+), 25 deletions(-) create mode 100644 packages/db/scripts/lock-timeout-retry.test.ts create mode 100644 packages/db/scripts/lock-timeout-retry.ts diff --git a/packages/db/scripts/lock-timeout-retry.test.ts b/packages/db/scripts/lock-timeout-retry.test.ts new file mode 100644 index 00000000000..0da4d4977f2 --- /dev/null +++ b/packages/db/scripts/lock-timeout-retry.test.ts @@ -0,0 +1,139 @@ +/** + * @vitest-environment node + */ + +import { retryOnLockTimeout } from '@sim/db/scripts/lock-timeout-retry' +import { describe, expect, it, vi } from 'vitest' + +const BACKOFF = { baseMs: 2_000, maxMs: 30_000 } as const + +function pgError(code: string): Error { + return Object.assign(new Error(`postgres error ${code}`), { code }) +} + +/** A fake clock that only advances when the retry loop sleeps. */ +function fakeClock() { + let nowMs = 0 + return { + now: () => nowMs, + sleep: vi.fn(async (ms: number) => { + nowMs += ms + }), + advance: (ms: number) => { + nowMs += ms + }, + } +} + +describe('retryOnLockTimeout', () => { + it('keeps retrying lock timeouts well past eight attempts while the budget lasts', async () => { + const clock = fakeClock() + let calls = 0 + const result = await retryOnLockTimeout( + async () => { + calls++ + clock.advance(5_000) + if (calls < 20) throw pgError('55P03') + return 'applied' + }, + { budgetMs: 20 * 60_000, backoff: BACKOFF, now: clock.now, sleep: clock.sleep } + ) + + expect(result).toBe('applied') + expect(calls).toBe(20) + expect(clock.sleep).toHaveBeenCalledTimes(19) + }) + + it('starts no attempt after the budget and throws the last lock timeout', async () => { + const clock = fakeClock() + const onRetry = vi.fn() + const startedAt: number[] = [] + const attempt = vi.fn(async () => { + startedAt.push(clock.now()) + clock.advance(5_000) + throw pgError('55P03') + }) + + await expect( + retryOnLockTimeout(attempt, { + budgetMs: 2 * 60_000, + backoff: BACKOFF, + now: clock.now, + sleep: clock.sleep, + onRetry, + }) + ).rejects.toMatchObject({ code: '55P03' }) + + for (const start of startedAt) expect(start).toBeLessThan(2 * 60_000) + /** The last attempt may run one lock timeout past the budget, never more. */ + expect(clock.now()).toBeLessThan(2 * 60_000 + 5_000) + expect(attempt).toHaveBeenCalledTimes(onRetry.mock.calls.length + 1) + }) + + it('does not start an attempt when a timer resolves after the budget', async () => { + let nowMs = 0 + const attempt = vi.fn(async () => { + nowMs += 1_000 + throw pgError('55P03') + }) + + await expect( + retryOnLockTimeout(attempt, { + budgetMs: 60_000, + backoff: BACKOFF, + now: () => nowMs, + /** The process stalls: the timer fires long after its delay. */ + sleep: async () => { + nowMs += 120_000 + }, + }) + ).rejects.toMatchObject({ code: '55P03' }) + expect(attempt).toHaveBeenCalledOnce() + }) + + it('finds a lock timeout wrapped in a cause chain', async () => { + const clock = fakeClock() + let calls = 0 + await retryOnLockTimeout( + async () => { + calls++ + if (calls === 1) throw new Error('Failed query', { cause: pgError('55P03') }) + }, + { budgetMs: 60_000, backoff: BACKOFF, now: clock.now, sleep: clock.sleep } + ) + + expect(calls).toBe(2) + }) + + it('does not retry any other error', async () => { + const clock = fakeClock() + const attempt = vi.fn(async () => { + throw pgError('42P07') + }) + + await expect( + retryOnLockTimeout(attempt, { + budgetMs: 60_000, + backoff: BACKOFF, + now: clock.now, + sleep: clock.sleep, + }) + ).rejects.toMatchObject({ code: '42P07' }) + expect(attempt).toHaveBeenCalledTimes(1) + expect(clock.sleep).not.toHaveBeenCalled() + }) + + it('passes the attempt number to each attempt', async () => { + const clock = fakeClock() + const seen: number[] = [] + await retryOnLockTimeout( + async (attemptNumber) => { + seen.push(attemptNumber) + if (attemptNumber < 3) throw pgError('55P03') + }, + { budgetMs: 60_000, backoff: BACKOFF, now: clock.now, sleep: clock.sleep } + ) + + expect(seen).toEqual([1, 2, 3]) + }) +}) diff --git a/packages/db/scripts/lock-timeout-retry.ts b/packages/db/scripts/lock-timeout-retry.ts new file mode 100644 index 00000000000..dd283e05c55 --- /dev/null +++ b/packages/db/scripts/lock-timeout-retry.ts @@ -0,0 +1,70 @@ +import { getPostgresErrorCode } from '@sim/utils/errors' +import { sleep as defaultSleep } from '@sim/utils/helpers' +import { backoffWithJitter } from '@sim/utils/retry' + +/** SQLSTATE `lock_not_available`, raised when `lock_timeout` expires. */ +const LOCK_NOT_AVAILABLE = '55P03' + +export interface LockTimeoutRetryAttempt { + /** The attempt that just failed, starting at 1. */ + attempt: number + delayMs: number + elapsedMs: number + budgetMs: number +} + +export interface LockTimeoutRetryOptions { + /** + * Time, measured on a monotonic clock from the first attempt, within which + * attempts may start. No attempt starts once the budget has elapsed, whether + * the backoff delay would end past it or a timer resolved late; the last lock + * timeout is thrown instead. An attempt that starts in time can still run for + * up to one `lock_timeout` past the budget. + */ + budgetMs: number + backoff: { baseMs: number; maxMs: number } + onRetry?: (attempt: LockTimeoutRetryAttempt) => void + /** Monotonic milliseconds; defaults to `performance.now`, immune to wall-clock corrections. */ + now?: () => number + sleep?: (ms: number) => Promise +} + +/** + * Run `attempt` until it succeeds, retrying only lock timeouts (55P03, found + * anywhere in the wrapped `cause` chain) for up to `budgetMs`. + * + * DDL on a hot table needs an ACCESS EXCLUSIVE lock, which it can only take in + * a moment when no transaction holds any lock on the table. Each attempt must + * keep a short `lock_timeout`, because a queued ACCESS EXCLUSIVE request blocks + * every later query on the table for as long as it waits. Many short attempts + * spread over a long budget find such a moment without ever stalling traffic + * for more than one `lock_timeout`; a fixed attempt count gives up after a few + * minutes whenever the table is continuously held by transactions that each + * outlive the timeout. Any other error is thrown immediately. + */ +export async function retryOnLockTimeout( + attempt: (attemptNumber: number) => Promise, + options: LockTimeoutRetryOptions +): Promise { + const now = options.now ?? (() => performance.now()) + const sleep = options.sleep ?? defaultSleep + const startedAt = now() + for (let attemptNumber = 1; ; attemptNumber++) { + try { + return await attempt(attemptNumber) + } catch (error) { + if (getPostgresErrorCode(error) !== LOCK_NOT_AVAILABLE) throw error + const delayMs = backoffWithJitter(attemptNumber, null, options.backoff) + const elapsedMs = now() - startedAt + if (elapsedMs + delayMs >= options.budgetMs) throw error + options.onRetry?.({ + attempt: attemptNumber, + delayMs, + elapsedMs, + budgetMs: options.budgetMs, + }) + await sleep(delayMs) + if (now() - startedAt >= options.budgetMs) throw error + } + } +} diff --git a/packages/db/scripts/migrate.ts b/packages/db/scripts/migrate.ts index b66f0484760..c760107eb5c 100644 --- a/packages/db/scripts/migrate.ts +++ b/packages/db/scripts/migrate.ts @@ -5,6 +5,7 @@ import { drizzle } from 'drizzle-orm/postgres-js' import { migrate } from 'drizzle-orm/postgres-js/migrator' import postgres from 'postgres' import { runScriptMigrations } from '../script-migrations/index' +import { retryOnLockTimeout } from './lock-timeout-retry' /** * Concurrent-index convention: plain `CREATE INDEX` write-blocks large/hot @@ -77,7 +78,14 @@ const LOCK_RETRY_INTERVAL_MS = 5_000 * query on the table behind it — a table-wide stall for the whole wait. */ const DDL_LOCK_TIMEOUT = '5s' -const MAX_MIGRATE_ATTEMPTS = 8 +/** + * Total time to keep retrying lock timeouts. A table held continuously by + * transactions that each outlive `DDL_LOCK_TIMEOUT` frees up only in short + * windows, so the budget is time-based rather than a small attempt count. It + * stays under `LOCK_ACQUIRE_DEADLINE_MS` so a runner waiting on the advisory + * lock sees this one finish, one way or the other, before its own deadline. + */ +const MIGRATE_LOCK_RETRY_BUDGET_MS = 20 * 60_000 const MIGRATE_RETRY_BACKOFF = { baseMs: 2_000, maxMs: 30_000 } as const const CONNECT_MAX_ATTEMPTS = 10 @@ -182,14 +190,6 @@ async function acquireMigrationLock(): Promise { } } -/** - * Run pending migrations, retrying on lock timeout (55P03, found anywhere in - * the wrapped `cause` chain). Each attempt re-verifies the lock session (pid) - * and re-asserts the session timeouts — a migration file may have changed them, - * and `SET` cannot be parameterized, hence `client.unsafe` with constants. - * Replays are safe: drizzle rolls the batch back on failure, and post-COMMIT - * CONCURRENTLY statements are idempotent by convention. - */ /** * Verify the session still holds the migration advisory lock: a changed * backend pid means the connection was recycled and the lock silently dropped. @@ -206,25 +206,34 @@ async function assertLockSessionHeld(): Promise { } } +/** + * Run pending migrations, retrying lock timeouts within + * `MIGRATE_LOCK_RETRY_BUDGET_MS` (see `retryOnLockTimeout`). Each attempt re-verifies the lock session (pid) + * and re-asserts the session timeouts — a migration file may have changed them, + * and `SET` cannot be parameterized, hence `client.unsafe` with constants. + * Replays are safe: drizzle rolls the batch back on failure, and post-COMMIT + * CONCURRENTLY statements are idempotent by convention. + */ async function runMigrationsWithRetry(): Promise { - for (let attempt = 1; ; attempt++) { - await assertLockSessionHeld() - await client.unsafe('SET statement_timeout = 0') - await client.unsafe(`SET lock_timeout = '${DDL_LOCK_TIMEOUT}'`) - try { + await retryOnLockTimeout( + async () => { + await assertLockSessionHeld() + await client.unsafe('SET statement_timeout = 0') + await client.unsafe(`SET lock_timeout = '${DDL_LOCK_TIMEOUT}'`) await migrate(drizzle(client), { migrationsFolder: './migrations' }) - return - } catch (error) { - const isLockTimeout = getPostgresErrorCode(error) === '55P03' - if (!isLockTimeout || attempt >= MAX_MIGRATE_ATTEMPTS) throw error - const delayMs = backoffWithJitter(attempt, null, MIGRATE_RETRY_BACKOFF) - console.warn( - `WARN: migration DDL hit lock_timeout (attempt ${attempt}/${MAX_MIGRATE_ATTEMPTS}); ` + - `retrying in ${Math.round(delayMs)}ms.` - ) - await sleep(delayMs) + }, + { + budgetMs: MIGRATE_LOCK_RETRY_BUDGET_MS, + backoff: MIGRATE_RETRY_BACKOFF, + onRetry: ({ attempt, delayMs, elapsedMs, budgetMs }) => { + console.warn( + `WARN: migration DDL hit lock_timeout (attempt ${attempt}, ` + + `${Math.round(elapsedMs / 1000)}s of ${Math.round(budgetMs / 1000)}s budget); ` + + `retrying in ${Math.round(delayMs)}ms.` + ) + }, } - } + ) } /** From 230e42ab0ede4ef321e771b08a2d0dda041cf027 Mon Sep 17 00:00:00 2001 From: Waleed Date: Tue, 22 Sep 2026 23:11:13 -0700 Subject: [PATCH 03/13] fix(desktop): set the browser user agent process-wide so Cloudflare Turnstile passes (#8192) * fix(desktop): set the browser user agent process-wide so Cloudflare Turnstile passes * test(desktop): cover the process-wide user agent from the first launched request --- apps/desktop/e2e/smoke.spec.ts | 19 +++++++++++++++++++ .../src/main/browser-agent/session.test.ts | 8 ++++---- .../desktop/src/main/browser-agent/session.ts | 10 ---------- apps/desktop/src/main/index.ts | 2 ++ .../{browser-agent => }/user-agent.test.ts | 19 ++++++++++--------- .../main/{browser-agent => }/user-agent.ts | 13 ++++++++----- 6 files changed, 43 insertions(+), 28 deletions(-) rename apps/desktop/src/main/{browser-agent => }/user-agent.test.ts (69%) rename apps/desktop/src/main/{browser-agent => }/user-agent.ts (77%) diff --git a/apps/desktop/e2e/smoke.spec.ts b/apps/desktop/e2e/smoke.spec.ts index f6799b7a835..08ace958c77 100644 --- a/apps/desktop/e2e/smoke.spec.ts +++ b/apps/desktop/e2e/smoke.spec.ts @@ -21,9 +21,13 @@ const PAGES: Record = { '/login': '

fixture-login

', } +/** `User-Agent` of every request the fixture origin has served, in arrival order. */ +const requestUserAgents: string[] = [] + function startFixtureServer(): Promise<{ server: Server; origin: string }> { return new Promise((resolvePromise) => { const server = createServer((request, response) => { + requestUserAgents.push(request.headers['user-agent'] ?? '') const path = new URL(request.url ?? '/', 'http://127.0.0.1').pathname const sessionCookie = request.headers.cookie ?.split(';') @@ -85,6 +89,21 @@ test.describe('desktop shell smoke', () => { expect(window.url()).toBe(`${origin}/home`) }) + test('presents one stock Chrome user agent on every request from the first load', async () => { + requestUserAgents.length = 0 + app = await launchApp(origin) + const window = await app.firstWindow() + await expect(window.locator('#app')).toHaveText('fixture-app') + await window.evaluate(() => fetch('/home').then((response) => response.text())) + + const pageUserAgent = await window.evaluate(() => navigator.userAgent) + expect(pageUserAgent).toMatch( + /^Mozilla\/5\.0 \(.+\) AppleWebKit\/537\.36 \(KHTML, like Gecko\) Chrome\/\d+\.0\.0\.0 Safari\/537\.36$/ + ) + expect(requestUserAgents.length).toBeGreaterThanOrEqual(2) + expect(new Set(requestUserAgents)).toEqual(new Set([pageUserAgent])) + }) + test('internal window.open creates an independent full Sim window', async () => { app = await launchApp(origin) const window = await app.firstWindow() diff --git a/apps/desktop/src/main/browser-agent/session.test.ts b/apps/desktop/src/main/browser-agent/session.test.ts index 39e1bc2683b..a8f5603f594 100644 --- a/apps/desktop/src/main/browser-agent/session.test.ts +++ b/apps/desktop/src/main/browser-agent/session.test.ts @@ -37,6 +37,7 @@ interface MockView { session: { setPermissionRequestHandler: ReturnType setPermissionCheckHandler: ReturnType + setUserAgent: ReturnType webRequest: { onBeforeRequest: ReturnType } } on: ReturnType @@ -358,15 +359,14 @@ describe('browser-agent session', () => { expect(onTabNavigated).toHaveBeenCalledWith(contents, true) }) - it('gives every tab a user agent with no Electron token in it', () => { + it('leaves every tab on the process-wide user agent instead of overriding it', () => { const first = session.ensureTab() const second = session.addTab() for (const tab of [first, second]) { const contents = (tab.view as unknown as MockView).webContents - const agent = contents.setUserAgent.mock.calls.at(-1)?.[0] as string | undefined - expect(agent).toMatch(/^Mozilla\/5\.0 \(.+\) .*Chrome\/\d+\.0\.0\.0 Safari\/537\.36$/) - expect(agent).not.toMatch(/Electron|Sim\//) + expect(contents.setUserAgent).not.toHaveBeenCalled() + expect(contents.session.setUserAgent).not.toHaveBeenCalled() } }) diff --git a/apps/desktop/src/main/browser-agent/session.ts b/apps/desktop/src/main/browser-agent/session.ts index e183137c9a0..1e106142eae 100644 --- a/apps/desktop/src/main/browser-agent/session.ts +++ b/apps/desktop/src/main/browser-agent/session.ts @@ -67,7 +67,6 @@ import { isBlockedSubresourceUrl, subresourceNeedsResolution, } from '@/main/browser-agent/url-guard' -import { browserUserAgent } from '@/main/browser-agent/user-agent' import type { BrowserSessionSnapshot } from '@/main/desktop-chat-session-store' import { suggestedFilename, uniqueDownloadPath } from '@/main/downloads' import { @@ -1413,11 +1412,6 @@ function configureAgentPartition(ses: Session): void { } return ALLOWED_SITE_PERMISSIONS.has(permission) }) - // Service workers do not inherit a tab's user agent. With only the tab's set, - // the document request carries the browser string while the worker's own - // script request still announces Electron — and on a site that routes its - // fetches through a worker, that is the one the server sees. - ses.setUserAgent(browserUserAgent()) // SSRF choke point for the agent partition. Document navigations (top-level + // iframes) get the full DNS-resolving check — the one seam every navigation // passes through, including page-initiated ones the driver never sees (server @@ -1965,10 +1959,6 @@ function initializeTabView(view: WebContentsView, scopeId: string): WebContentsV const contents = view.webContents registerAgentWebContents(contents) configureAgentPartition(contents.session) - // The session default does not reach a WebContents that already exists, and - // the first tab is what brings the session into being, so each tab sets its - // own as well — otherwise tab one browses as Electron and the rest as Chrome. - contents.setUserAgent(browserUserAgent()) attachAgentContextMenu(contents, { addToChat: (text) => withBrowserScope(scopeId, () => addPageSelectionToChat(contents, text)), openTab: (url) => withBrowserScope(scopeId, () => openTabWithUrl(url, { agentOwned: false })), diff --git a/apps/desktop/src/main/index.ts b/apps/desktop/src/main/index.ts index 30972db3c9e..18996aedcce 100644 --- a/apps/desktop/src/main/index.ts +++ b/apps/desktop/src/main/index.ts @@ -83,6 +83,7 @@ import { attachTelemetryPolicy } from '@/main/telemetry-policy' import { TerminalRegistry } from '@/main/terminal/registry' import { installTray, type TrayHandle } from '@/main/tray' import { checkForUpdatesInteractive, initUpdater, type UpdaterHandle } from '@/main/updater' +import { installBrowserUserAgent } from '@/main/user-agent' import { createMainWindow, setupPermissionHandlers } from '@/main/window' import { attachWindowOpenPolicy, isPopupContents } from '@/main/windows' @@ -899,6 +900,7 @@ app.setName(APP_NAME_FOR_CHANNEL[channelForOrigin(DEFAULT_ORIGIN)]) if (process.env.SIM_DESKTOP_USER_DATA) { app.setPath('userData', process.env.SIM_DESKTOP_USER_DATA) } +installBrowserUserAgent() // The scheme the offline page and server picker load from must be declared // before the app is ready; the per-session handlers attach later. diff --git a/apps/desktop/src/main/browser-agent/user-agent.test.ts b/apps/desktop/src/main/user-agent.test.ts similarity index 69% rename from apps/desktop/src/main/browser-agent/user-agent.test.ts rename to apps/desktop/src/main/user-agent.test.ts index db63efa7141..423ef6600d1 100644 --- a/apps/desktop/src/main/browser-agent/user-agent.test.ts +++ b/apps/desktop/src/main/user-agent.test.ts @@ -1,11 +1,13 @@ import { app } from 'electron' import { describe, expect, it, vi } from 'vitest' -import { browserUserAgent, stockChromeUserAgent } from '@/main/browser-agent/user-agent' +import { installBrowserUserAgent, stockChromeUserAgent } from '@/main/user-agent' vi.mock('electron', () => import('@/test/electron-mock')) const ELECTRON_DEFAULT = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Sim/1.0.0 Chrome/140.0.7339.207 Electron/43.1.1 Safari/537.36' +const STOCK_CHROME = + 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36' describe('stockChromeUserAgent', () => { it('drops the application and Electron tokens a browser allowlist rejects', () => { @@ -15,9 +17,7 @@ describe('stockChromeUserAgent', () => { }) it('reproduces the desktop string Chrome sends under user-agent reduction', () => { - expect(stockChromeUserAgent(ELECTRON_DEFAULT)).toBe( - 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36' - ) + expect(stockChromeUserAgent(ELECTRON_DEFAULT)).toBe(STOCK_CHROME) }) it('keeps the platform token of the machine it is running on', () => { @@ -32,12 +32,13 @@ describe('stockChromeUserAgent', () => { }) }) -describe('browserUserAgent', () => { - it('derives from the string Electron would otherwise have sent', () => { +describe('installBrowserUserAgent', () => { + it('idempotently makes stock Chrome the process-wide fallback every request path uses', () => { app.userAgentFallback = ELECTRON_DEFAULT - expect(browserUserAgent()).toBe( - 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36' - ) + installBrowserUserAgent() + installBrowserUserAgent() + + expect(app.userAgentFallback).toBe(STOCK_CHROME) }) }) diff --git a/apps/desktop/src/main/browser-agent/user-agent.ts b/apps/desktop/src/main/user-agent.ts similarity index 77% rename from apps/desktop/src/main/browser-agent/user-agent.ts rename to apps/desktop/src/main/user-agent.ts index 9e988c67987..d14e5c80f1d 100644 --- a/apps/desktop/src/main/browser-agent/user-agent.ts +++ b/apps/desktop/src/main/user-agent.ts @@ -1,5 +1,6 @@ /** - * The user agent the browser resource presents to sites. + * The user agent the whole desktop process presents: the embedded browser and + * the app's own windows alike (desktop identity travels in `X-Sim-Client-Info`). * * Electron's default string carries two tokens no browser sends — * `Sim/` and `Electron/`. Chromium's own token sits right @@ -38,9 +39,11 @@ export function stockChromeUserAgent(defaultUserAgent: string): string { } /** - * Derived from the string Electron would otherwise have sent, so the reported - * Chromium version tracks whatever Chromium the app actually ships. + * Sets the stock Chrome identity as `app.userAgentFallback` before any session + * exists. Session and per-tab overrides miss some request paths (a cross-origin + * challenge frame still sends the process default), and a site that sees two + * user agents in one challenge rejects it as a spoof. Idempotent. */ -export function browserUserAgent(): string { - return stockChromeUserAgent(app.userAgentFallback) +export function installBrowserUserAgent(): void { + app.userAgentFallback = stockChromeUserAgent(app.userAgentFallback) } From 22ff39fefd1e414bc2e149b17d4a37e1a26d2d83 Mon Sep 17 00:00:00 2001 From: Waleed Date: Tue, 22 Sep 2026 23:20:03 -0700 Subject: [PATCH 04/13] fix(knowledge): stop holding the connector table lock across the processing commit (#8191) * fix(knowledge): check connector and knowledge base at the end of the processing commit * fix(knowledge): skip index writes for an inactive source and probe connector locks after the insert * test(knowledge): arm the pre-commit source check in the processing utils test --- .github/workflows/test-build.yml | 1 + apps/sim/app/api/knowledge/utils.test.ts | 2 + .../processing-lock-scope.integration.ts | 225 ++++++++++++++++ .../documents/document-indexing-usage.test.ts | 5 +- .../document-processing-source.test.ts | 22 ++ apps/sim/lib/knowledge/documents/service.ts | 251 +++++++++++------- 6 files changed, 405 insertions(+), 101 deletions(-) create mode 100644 apps/sim/lib/knowledge/__integration__/processing-lock-scope.integration.ts diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index d72ced902c8..4736123fd93 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -267,6 +267,7 @@ jobs: lib/knowledge/__integration__/search-source-pagination.integration.ts lib/knowledge/__integration__/search-reference-batching.integration.ts lib/knowledge/__integration__/embedding-insert-batches.integration.ts + lib/knowledge/__integration__/processing-lock-scope.integration.ts lib/knowledge/__integration__/connector-lifecycle-locks.integration.ts lib/knowledge/__integration__/connector-deferral.integration.ts lib/knowledge/__integration__/stored-document-recovery.integration.ts diff --git a/apps/sim/app/api/knowledge/utils.test.ts b/apps/sim/app/api/knowledge/utils.test.ts index d957ff866c3..b0177b13b1e 100644 --- a/apps/sim/app/api/knowledge/utils.test.ts +++ b/apps/sim/app/api/knowledge/utils.test.ts @@ -211,6 +211,8 @@ describe('Knowledge Utils', () => { secretProvenanceVersion: null, }, ]) + /** Pre-commit connector and knowledge-base activity check. */ + queueTableRows(schemaMock.document, [{ id: 'doc1' }]) /** In-transaction active-document recheck. */ queueTableRows(schemaMock.document, [{ id: 'doc1' }]) diff --git a/apps/sim/lib/knowledge/__integration__/processing-lock-scope.integration.ts b/apps/sim/lib/knowledge/__integration__/processing-lock-scope.integration.ts new file mode 100644 index 00000000000..ec37840614d --- /dev/null +++ b/apps/sim/lib/knowledge/__integration__/processing-lock-scope.integration.ts @@ -0,0 +1,225 @@ +/** Relation locks held by the document processing commit while it writes embeddings. */ +import { mkdtempSync } from 'node:fs' +import { rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import path from 'node:path' +import { db } from '@sim/db' +import { document, embedding, knowledgeBase, organization, user, workspace } from '@sim/db/schema' +import { generateId } from '@sim/utils/id' +import { eq, inArray } from 'drizzle-orm' +import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from 'vitest' + +const fixtures = vi.hoisted(() => ({ root: '', process: vi.fn(), embeddings: vi.fn() })) +vi.mock('@/lib/uploads/core/setup.server', () => ({ + get UPLOAD_DIR_SERVER() { + return fixtures.root + }, +})) +vi.mock('@/lib/knowledge/documents/document-processor', () => ({ + processDocument: fixtures.process, +})) +vi.mock('@/lib/knowledge/embeddings', () => ({ generateEmbeddings: fixtures.embeddings })) + +import { resolveBillingAttribution } from '@/lib/billing/core/billing-attribution' +import * as embeddingClient from '@/lib/embeddings/client' +import { + createKnowledgeAclFixtureIds, + seedKnowledgeAclFixture, +} from '@/lib/knowledge/__integration__/seed-source-access-fixture' +import { createContentSyncLease } from '@/lib/knowledge/connectors/sync-lock' +import { addDocument } from '@/lib/knowledge/connectors/sync-persistence' +import { processDocumentAsync } from '@/lib/knowledge/documents/service' + +describe('document processing commit lock scope', () => { + const ids = createKnowledgeAclFixtureIds() + const probe = `fixture_lock_probe_${generateId().replaceAll('-', '')}` + const chunks = Array.from({ length: 3 }, (_, index) => ({ + text: `Synthetic chunk ${index}`, + metadata: { startIndex: index * 20, endIndex: index * 20 + 19 }, + })) + const embeddingResult = { + embeddings: chunks.map(() => Array(1536).fill(0.2)), + billableTokens: 0, + modelName: 'text-embedding-3-small', + pricingId: 'text-embedding-3-small', + } + + beforeAll(async () => { + fixtures.root = mkdtempSync(path.join(tmpdir(), 'sim-processing-lock-scope-')) + await seedKnowledgeAclFixture(ids, { connectorType: 'google_drive' }) + vi.spyOn(embeddingClient, 'assertKnowledgeEmbeddingCapacity').mockResolvedValue(undefined) + fixtures.process.mockResolvedValue({ + chunks, + metadata: { chunkCount: chunks.length, tokenCount: 9, characterCount: 60 }, + }) + fixtures.embeddings.mockResolvedValue(embeddingResult) + }) + + afterEach(async () => { + await db.$client.unsafe(`DROP TRIGGER IF EXISTS ${probe} ON embedding`) + await db.$client.unsafe(`DROP FUNCTION IF EXISTS ${probe}()`) + }) + + afterAll(async () => { + vi.restoreAllMocks() + await db.delete(knowledgeBase).where(eq(knowledgeBase.id, ids.knowledgeBaseId)) + await db.delete(workspace).where(eq(workspace.id, ids.workspaceId)) + await db.delete(organization).where(eq(organization.id, ids.organizationId)) + await db.delete(user).where(inArray(user.id, [ids.aliceId, ids.bobId])) + await rm(fixtures.root, { recursive: true, force: true }) + await db.$client.end() + }) + + async function addConnectorDocument(externalId: string) { + return addDocument( + ids.knowledgeBaseId, + ids.connectorId, + 'google_drive', + { + externalId, + title: `${externalId}.txt`, + content: 'Synthetic source text', + mimeType: 'text/plain', + contentHash: externalId, + }, + { userId: ids.aliceId, workspaceId: ids.workspaceId }, + undefined, + 'workspace', + createContentSyncLease(ids.connectorId, ids.lockId) + ) + } + + /** Installs `body` as a BEFORE INSERT row trigger on embedding for one document. */ + async function installEmbeddingProbe(documentId: string, body: string) { + await db.$client.unsafe(`CREATE FUNCTION ${probe}() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN + IF NEW.document_id = '${documentId}' THEN + ${body} + END IF; + RETURN NEW; + END; + $$`) + await db.$client.unsafe(`CREATE TRIGGER ${probe} BEFORE INSERT ON embedding + FOR EACH ROW EXECUTE FUNCTION ${probe}()`) + } + + /** + * Fails any embedding insert statement whose backend holds a lock on + * `knowledge_connector`. A statement-level AFTER trigger fires once the row + * triggers and foreign-key checks of that statement have run, so it sees + * every lock the insert itself took. + */ + async function installConnectorLockProbe() { + await db.$client.unsafe(`CREATE FUNCTION ${probe}() RETURNS trigger LANGUAGE plpgsql AS $$ + BEGIN + IF EXISTS ( + SELECT 1 FROM pg_locks + WHERE pid = pg_backend_pid() AND relation = 'knowledge_connector'::regclass + ) THEN + RAISE EXCEPTION 'embedding write holds a knowledge_connector lock'; + END IF; + RETURN NULL; + END; + $$`) + await db.$client.unsafe(`CREATE TRIGGER ${probe} AFTER INSERT ON embedding + FOR EACH STATEMENT EXECUTE FUNCTION ${probe}()`) + } + + function billing() { + return resolveBillingAttribution({ actorUserId: ids.aliceId, workspaceId: ids.workspaceId }) + } + + it('holds no knowledge_connector lock while writing embeddings', async () => { + const file = await addConnectorDocument('lock-scope-fixture') + await installConnectorLockProbe() + + const result = await processDocumentAsync( + ids.knowledgeBaseId, + file.documentId, + file, + {}, + await billing() + ) + + expect(result).toEqual({ outcome: 'indexed' }) + expect(await db.select().from(document).where(eq(document.id, file.documentId))).toMatchObject([ + { processingStatus: 'completed', chunkCount: 3 }, + ]) + }) + + it.each([ + ['connector', 'lock-scope-deleted-connector', 'knowledge_connector', () => ids.connectorId], + ['knowledge base', 'lock-scope-deleted-kb', 'knowledge_base', () => ids.knowledgeBaseId], + ])( + 'rolls back the embeddings when the %s is deleted during the embedding writes', + async (_, externalId, table, id) => { + const file = await addConnectorDocument(externalId) + await installEmbeddingProbe( + file.documentId, + `UPDATE ${table} SET deleted_at = now() WHERE id = '${id()}';` + ) + + const result = await processDocumentAsync( + ids.knowledgeBaseId, + file.documentId, + file, + {}, + await billing() + ) + + expect(result).toEqual({ outcome: 'skipped', reason: 'superseded' }) + expect( + await db + .select({ id: embedding.id }) + .from(embedding) + .where(eq(embedding.documentId, file.documentId)) + ).toEqual([]) + for (const table of ['embedding_search', 'embedding_keyword_search']) { + expect( + await db.$client.unsafe(`SELECT id FROM ${table} WHERE document_id = $1`, [ + file.documentId, + ]) + ).toEqual([]) + } + expect( + await db.select().from(document).where(eq(document.id, file.documentId)) + ).toMatchObject([{ processingStatus: 'processing', chunkCount: 0 }]) + } + ) + it.each([ + ['connector', 'lock-scope-precheck-connector', 'knowledge_connector', () => ids.connectorId], + ['knowledge base', 'lock-scope-precheck-kb', 'knowledge_base', () => ids.knowledgeBaseId], + ])( + 'skips the index writes when the %s went inactive after the claim', + async (_, externalId, table, id) => { + const file = await addConnectorDocument(externalId) + await installEmbeddingProbe( + file.documentId, + `RAISE EXCEPTION 'index writes ran for an inactive source';` + ) + fixtures.embeddings.mockImplementationOnce(async () => { + await db.$client.unsafe(`UPDATE ${table} SET deleted_at = now() WHERE id = $1`, [id()]) + return embeddingResult + }) + + try { + const result = await processDocumentAsync( + ids.knowledgeBaseId, + file.documentId, + file, + {}, + await billing() + ) + expect(result).toEqual({ outcome: 'skipped', reason: 'superseded' }) + } finally { + await db.$client.unsafe(`UPDATE ${table} SET deleted_at = NULL WHERE id = $1`, [id()]) + } + expect( + await db + .select({ id: embedding.id }) + .from(embedding) + .where(eq(embedding.documentId, file.documentId)) + ).toEqual([]) + } + ) +}) diff --git a/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts b/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts index a6a351c5cbf..988220bd59c 100644 --- a/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts +++ b/apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts @@ -186,14 +186,15 @@ const DOC_DATA = { /** * Re-arms the row sets one `processDocumentAsync` call consumes: the KB/document - * context JOIN, the document secret-provenance row, and the in-transaction claim - * re-check that lets the attempt commit. + * context JOIN, the document secret-provenance row, the pre-commit source check, + * and the in-transaction claim re-check that lets the attempt commit. */ function armDocumentReads(): void { dbChainMockFns.limit .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: DOCUMENT_ID }]) + .mockResolvedValueOnce([{ id: DOCUMENT_ID }]) } /** The `sourceReference` of the single embedding charge recorded by call `index`. */ diff --git a/apps/sim/lib/knowledge/documents/document-processing-source.test.ts b/apps/sim/lib/knowledge/documents/document-processing-source.test.ts index 32e7eed699e..4c6ecfa113f 100644 --- a/apps/sim/lib/knowledge/documents/document-processing-source.test.ts +++ b/apps/sim/lib/knowledge/documents/document-processing-source.test.ts @@ -219,6 +219,7 @@ describe('knowledge document processing source', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) mockGetFileMetadataByKeys.mockImplementation(async (_keys: string[], context: string) => context === 'workspace' ? [SOURCE_BINDING] : [] @@ -333,6 +334,7 @@ describe('knowledge document processing source', () => { .mockResolvedValueOnce([{ ...PERSISTED_CONTEXT, connectorId: 'connector-1' }]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) await processDocumentAsync( 'knowledge-base-1', @@ -445,6 +447,7 @@ describe('knowledge document processing source', () => { .mockResolvedValueOnce([{ ...PERSISTED_CONTEXT, fileUrl: executionUrl }]) .mockResolvedValueOnce([{ ...PERSISTED_PROVENANCE_ROW, fileUrl: executionUrl }]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockImplementation(async (_keys: string[], context: string) => context === 'execution' ? [executionBinding] : [] ) @@ -571,6 +574,7 @@ describe('knowledge document processing source', () => { .mockResolvedValueOnce([{ ...PERSISTED_CONTEXT, processingStatus: 'processing' }]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) dbChainMockFns.returning.mockReset().mockResolvedValue([{ id: 'document-1' }]) await processDocumentAsync( @@ -603,6 +607,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -643,6 +648,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -673,6 +679,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -707,6 +714,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -916,6 +924,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -961,6 +970,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -1003,6 +1013,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -1063,6 +1074,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) // The guarded claim matched no rows: another pass owns this document. dbChainMockFns.returning.mockReset().mockResolvedValue([]) @@ -1115,6 +1127,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -1191,6 +1204,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -1235,6 +1249,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -1278,6 +1293,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -1311,6 +1327,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -1527,6 +1544,7 @@ describe('processDocumentAsync write guards', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -1600,6 +1618,7 @@ describe('in-process quota continuation dispatch', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockCheckAttributedUsageLimits.mockResolvedValue({ isExceeded: false }) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( @@ -1767,6 +1786,7 @@ describe('in-process quota continuation dispatch', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) dbChainMockFns.set.mockClear() mockGenerateEmbeddings.mockResolvedValue({ embeddings: [Array(1536).fill(0)], @@ -1811,6 +1831,7 @@ describe('in-process quota continuation dispatch', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) mockGetFileMetadataByKeys.mockResolvedValue([SOURCE_BINDING]) mockGetBoundWorkspaceFileSecretProvenanceByMetadata.mockResolvedValue( new Map([[SOURCE_BINDING.id, { status: 'exact', entries: [] }]]) @@ -1881,6 +1902,7 @@ describe('in-process quota continuation dispatch', () => { .mockResolvedValueOnce([PERSISTED_CONTEXT]) .mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW]) .mockResolvedValueOnce([{ id: 'document-1' }]) + .mockResolvedValueOnce([{ id: 'document-1' }]) await expect( processDocumentAsync( diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts index 36732b8e036..3be65ef600a 100644 --- a/apps/sim/lib/knowledge/documents/service.ts +++ b/apps/sim/lib/knowledge/documents/service.ts @@ -214,6 +214,27 @@ export class KnowledgeBaseFileOwnershipError extends OrchestrationError { } } +/** + * Rolls back a processing pass whose completion write matched no row: the + * document, its connector, or its knowledge base stopped being active before + * the completion write, so none of the pass's output may commit. + */ +class SupersededProcessingOutput extends Error { + constructor() { + super('Document processing output was superseded before commit') + this.name = 'SupersededProcessingOutput' + } +} + +/** The document's knowledge base has not been deleted. */ +function knowledgeBaseIsActive() { + return sql`EXISTS ( + SELECT 1 FROM ${knowledgeBase} + WHERE ${knowledgeBase.id} = ${document.knowledgeBaseId} + AND ${knowledgeBase.deletedAt} IS NULL + )` +} + /** Internal KB uploads require the workspace's trusted file binding; external ingestion URLs do not. */ function getKnowledgeBaseStorageKeys(fileUrls: readonly string[]): string[] { return [ @@ -1952,114 +1973,146 @@ export async function processDocumentAsync( })) signal.throwIfAborted() - processingCommitted = await db.transaction(async (tx) => { - signal.throwIfAborted() - const activeDocument = await tx - .select({ id: document.id }) - .from(document) - .innerJoin(knowledgeBase, eq(document.knowledgeBaseId, knowledgeBase.id)) - .where( - and( - eq(document.id, documentId), - eq(document.processingStatus, 'processing'), - eq(document.processingStartedAt, processingStartedAt), - ...queueGenerationConditions(attemptContext), - eq(document.userExcluded, false), - isNull(document.archivedAt), - isNull(document.deletedAt), - documentConnectorIsActive(), - isNull(knowledgeBase.deletedAt) - ) + /** + * Skips the index writes when the connector or knowledge base went + * inactive after the claim. As its own autocommit statement it + * releases its locks immediately; the completion write inside the + * transaction stays the authoritative check. + */ + const [sourceActive] = await db + .select({ id: document.id }) + .from(document) + .where( + and( + eq(document.id, documentId), + documentConnectorIsActive(), + knowledgeBaseIsActive() ) - .for('update', { of: document }) - .limit(1) + ) + .limit(1) + if (!sourceActive) return + processingCommitted = await db + .transaction(async (tx) => { + signal.throwIfAborted() + /** + * Reads only the document row. Connector activity is checked by + * the completion write at the end instead: reading + * `knowledge_connector` here would hold a lock on it until commit, + * across the embedding writes, so a slow index write would block + * DDL on the connector table for its whole duration. The + * knowledge base table stays locked for the pass regardless, + * through the embedding foreign key and projection triggers. + */ + const activeDocument = await tx + .select({ id: document.id }) + .from(document) + .where( + and( + eq(document.id, documentId), + eq(document.processingStatus, 'processing'), + eq(document.processingStartedAt, processingStartedAt), + ...queueGenerationConditions(attemptContext), + eq(document.userExcluded, false), + isNull(document.archivedAt), + isNull(document.deletedAt) + ) + ) + .for('update') + .limit(1) - if (activeDocument.length === 0) { - return false - } + if (activeDocument.length === 0) { + return false + } - if (embeddingRecords.length > 0) { - await tx.delete(embedding).where(eq(embedding.documentId, documentId)) + if (embeddingRecords.length > 0) { + await tx.delete(embedding).where(eq(embedding.documentId, documentId)) - const insertBatchSize = LARGE_DOC_CONFIG.MAX_CHUNKS_PER_BATCH - const batches: (typeof embeddingRecords)[] = [] - for (let i = 0; i < embeddingRecords.length; i += insertBatchSize) { - batches.push(embeddingRecords.slice(i, i + insertBatchSize)) - } + const insertBatchSize = LARGE_DOC_CONFIG.MAX_CHUNKS_PER_BATCH + const batches: (typeof embeddingRecords)[] = [] + for (let i = 0; i < embeddingRecords.length; i += insertBatchSize) { + batches.push(embeddingRecords.slice(i, i + insertBatchSize)) + } - logger.info(`[${documentId}] Inserting ${embeddingRecords.length} embeddings`) - for (const [batchIndex, batch] of batches.entries()) { - signal.throwIfAborted() - const insertStartedAt = Date.now() - try { - await tx.insert(embedding).values(batch) - } catch (error) { - logger.error(`[${documentId}] Failed to insert embedding batch`, { - knowledgeBaseId, - operation: 'embedding.insert', - batchNumber: batchIndex + 1, - batchSize: batch.length, - totalChunks: embeddingRecords.length, - embeddingModel: kbEmbeddingModel, - embeddingDimensions: kbEmbedding.dimensions, - elapsedMs: Date.now() - insertStartedAt, - diagnostic: getConnectorFailureDiagnostic(error), - }) - throw error + logger.info(`[${documentId}] Inserting ${embeddingRecords.length} embeddings`) + for (const [batchIndex, batch] of batches.entries()) { + signal.throwIfAborted() + const insertStartedAt = Date.now() + try { + await tx.insert(embedding).values(batch) + } catch (error) { + logger.error(`[${documentId}] Failed to insert embedding batch`, { + knowledgeBaseId, + operation: 'embedding.insert', + batchNumber: batchIndex + 1, + batchSize: batch.length, + totalChunks: embeddingRecords.length, + embeddingModel: kbEmbeddingModel, + embeddingDimensions: kbEmbedding.dimensions, + elapsedMs: Date.now() - insertStartedAt, + diagnostic: getConnectorFailureDiagnostic(error), + }) + throw error + } + } + const provenanceRecords = embeddingRecords.flatMap((record, index) => { + const provenance = chunkProvenances[index] + if (!provenance) return [] + return [ + { + embeddingId: record.id, + contentHash: record.chunkHash, + status: provenance.status, + entries: provenance.status === 'exact' ? [...provenance.entries] : [], + updatedAt: now, + }, + ] + }) + for (let i = 0; i < provenanceRecords.length; i += insertBatchSize) { + signal.throwIfAborted() + await tx + .insert(embeddingSecretProvenance) + .values(provenanceRecords.slice(i, i + insertBatchSize)) } } - const provenanceRecords = embeddingRecords.flatMap((record, index) => { - const provenance = chunkProvenances[index] - if (!provenance) return [] - return [ - { - embeddingId: record.id, - contentHash: record.chunkHash, - status: provenance.status, - entries: provenance.status === 'exact' ? [...provenance.entries] : [], - updatedAt: now, - }, - ] - }) - for (let i = 0; i < provenanceRecords.length; i += insertBatchSize) { - signal.throwIfAborted() - await tx - .insert(embeddingSecretProvenance) - .values(provenanceRecords.slice(i, i + insertBatchSize)) - } - } - signal.throwIfAborted() - await tx - .update(document) - .set({ - chunkCount: processed.metadata.chunkCount, - tokenCount: processed.metadata.tokenCount, - characterCount: processed.metadata.characterCount, - processingStatus: 'completed', - processingCompletedAt: now, - processingError: null, - /** A completed pass restores the retry allowance for a future failure. */ - processingAttempts: 0, - processingQueueToken: null, - processingQueuedAt: null, - processingDeferredUntil: null, - }) - .where( - and( - eq(document.id, documentId), - eq(document.processingStatus, 'processing'), - eq(document.processingStartedAt, processingStartedAt), - ...queueGenerationConditions(attemptContext), - eq(document.userExcluded, false), - isNull(document.archivedAt), - isNull(document.deletedAt), - documentConnectorIsActive() + signal.throwIfAborted() + const completed = await tx + .update(document) + .set({ + chunkCount: processed.metadata.chunkCount, + tokenCount: processed.metadata.tokenCount, + characterCount: processed.metadata.characterCount, + processingStatus: 'completed', + processingCompletedAt: now, + processingError: null, + /** A completed pass restores the retry allowance for a future failure. */ + processingAttempts: 0, + processingQueueToken: null, + processingQueuedAt: null, + processingDeferredUntil: null, + }) + .where( + and( + eq(document.id, documentId), + eq(document.processingStatus, 'processing'), + eq(document.processingStartedAt, processingStartedAt), + ...queueGenerationConditions(attemptContext), + eq(document.userExcluded, false), + isNull(document.archivedAt), + isNull(document.deletedAt), + documentConnectorIsActive(), + knowledgeBaseIsActive() + ) ) - ) - signal.throwIfAborted() - return true - }) + .returning({ id: document.id }) + if (completed.length === 0) throw new SupersededProcessingOutput() + signal.throwIfAborted() + return true + }) + .catch((error: unknown) => { + if (error instanceof SupersededProcessingOutput) return false + throw error + }) }, { opaqueInputSafe: From 0c321f75e447524fa62de1e58109d29e4061213d Mon Sep 17 00:00:00 2001 From: Waleed Date: Tue, 22 Sep 2026 23:36:59 -0700 Subject: [PATCH 05/13] improvement(insights): serve org usage from settled segments and redesign the overview (#8193) * improvement(insights): serve org usage from settled segments and redesign the overview * fix(insights): bound usage settling by the stream cap, discard unreadable segments, neutral zero delta * fix(insights): align rolling windows to the viewer's hour and clear stale legend highlights * improvement(emcn): draw each stacked column as one bar split by color * improvement(insights): brand-aligned, colorblind-validated chart palette --- .../[id]/usage/{summary => overview}/route.ts | 15 +- .../[organizationId]/usage/route.test.ts | 6 +- .../components/activity-summary.tsx | 132 ++++---- .../components/usage-consumers.tsx | 14 +- .../components/usage-credits.test.tsx | 60 ++++ .../components/usage-credits.tsx | 184 +++++++++++ .../components/usage-member-avatar.tsx | 24 ++ .../components/usage-monitoring.tsx | 99 +++--- .../components/usage-source-mix.tsx | 50 --- .../components/usage-summary.test.tsx | 28 -- .../components/usage-summary.tsx | 76 ----- .../components/usage-top-cards.tsx | 104 ++++++ apps/sim/ee/organization-usage/constants.ts | 52 +++ .../hooks/use-legend-highlight.ts | 27 ++ apps/sim/hooks/queries/organization-usage.ts | 26 +- apps/sim/hooks/queries/organization.ts | 5 + .../queries/utils/organization-usage-keys.ts | 28 +- .../api/contracts/organization-activity.ts | 3 +- .../lib/api/contracts/organization-usage.ts | 42 ++- .../authorized-organization-usage-use-case.ts | 14 +- .../get-organization-activity.test.ts | 54 ++- .../get-organization-activity.ts | 12 +- .../get-organization-usage-breakdown.ts | 307 ++++++++++-------- .../get-organization-usage-overview.test.ts | 150 +++++++++ .../get-organization-usage-overview.ts | 185 +++++++++++ .../organization-usage/operations.ts | 8 + .../core/organization-activity-queries.ts | 139 ++++++-- .../core/organization-activity-summary.ts | 59 ++++ .../organization-activity.postgres.test.ts | 27 +- .../lib/billing/core/organization-activity.ts | 42 +++ .../usage-analytics-queries.postgres.test.ts | 64 +++- .../billing/core/usage-analytics-queries.ts | 248 ++++++++++++-- .../lib/billing/core/usage-analytics.test.ts | 152 +++++++++ apps/sim/lib/billing/core/usage-analytics.ts | 186 ++++++++++- .../billing/core/usage-segment-cache.test.ts | 136 ++++++++ .../lib/billing/core/usage-segment-cache.ts | 224 +++++++++++++ .../emcn/src/components/charts/bar-chart.tsx | 299 ++++++++++++----- .../components/charts/chart-data-table.tsx | 22 +- .../src/components/charts/chart-format.ts | 20 +- .../components/charts/chart-layout.test.tsx | 178 +++++++--- .../src/components/charts/donut-chart.tsx | 84 ----- packages/emcn/src/components/charts/index.ts | 3 +- 42 files changed, 2827 insertions(+), 761 deletions(-) rename apps/sim/app/api/organizations/[id]/usage/{summary => overview}/route.ts (62%) create mode 100644 apps/sim/ee/organization-usage/components/usage-credits.test.tsx create mode 100644 apps/sim/ee/organization-usage/components/usage-credits.tsx create mode 100644 apps/sim/ee/organization-usage/components/usage-member-avatar.tsx delete mode 100644 apps/sim/ee/organization-usage/components/usage-source-mix.tsx delete mode 100644 apps/sim/ee/organization-usage/components/usage-summary.test.tsx delete mode 100644 apps/sim/ee/organization-usage/components/usage-summary.tsx create mode 100644 apps/sim/ee/organization-usage/components/usage-top-cards.tsx create mode 100644 apps/sim/ee/organization-usage/hooks/use-legend-highlight.ts create mode 100644 apps/sim/lib/billing/application/organization-usage/get-organization-usage-overview.test.ts create mode 100644 apps/sim/lib/billing/application/organization-usage/get-organization-usage-overview.ts create mode 100644 apps/sim/lib/billing/core/organization-activity-summary.ts create mode 100644 apps/sim/lib/billing/core/usage-segment-cache.test.ts create mode 100644 apps/sim/lib/billing/core/usage-segment-cache.ts delete mode 100644 packages/emcn/src/components/charts/donut-chart.tsx diff --git a/apps/sim/app/api/organizations/[id]/usage/summary/route.ts b/apps/sim/app/api/organizations/[id]/usage/overview/route.ts similarity index 62% rename from apps/sim/app/api/organizations/[id]/usage/summary/route.ts rename to apps/sim/app/api/organizations/[id]/usage/overview/route.ts index 4d1111e1962..8294ead619c 100644 --- a/apps/sim/app/api/organizations/[id]/usage/summary/route.ts +++ b/apps/sim/app/api/organizations/[id]/usage/overview/route.ts @@ -1,24 +1,19 @@ -import { getOrganizationUsageSummaryContract } from '@/lib/api/contracts/organization-usage' +import { getOrganizationUsageOverviewContract } from '@/lib/api/contracts/organization-usage' import { defineInternalJsonRoute, internalRateLimits, internalSessionAuth, } from '@/lib/api/server/routes' -import { getOrganizationUsageSummary } from '@/lib/billing/application/organization-usage/get-organization-usage-summary' +import { getOrganizationUsageOverview } from '@/lib/billing/application/organization-usage/get-organization-usage-overview' import { organizationUsageOperations } from '@/lib/billing/application/organization-usage/operations' import { organizationUsageErrorPolicy } from '@/app/api/organizations/[id]/usage/error-policy' export const dynamic = 'force-dynamic' -/** - * Everything above the fold in one round trip. Kept separate from the breakdown - * route because every read here is index-covered, and folding in a dimension that - * heap-scans would put that cost on first paint. - */ export const GET = defineInternalJsonRoute({ - contract: getOrganizationUsageSummaryContract, + contract: getOrganizationUsageOverviewContract, auth: internalSessionAuth, - operation: organizationUsageOperations.readSummary, + operation: organizationUsageOperations.readOverview, rateLimit: internalRateLimits.none({ reason: 'Authenticated org-admin settings read, gated on enterprise entitlement and billing authority', @@ -32,6 +27,6 @@ export const GET = defineInternalJsonRoute({ endDate: query.endDate ? new Date(query.endDate) : undefined, timezone: query.timezone, }), - useCase: getOrganizationUsageSummary, + useCase: getOrganizationUsageOverview, present: (result) => result, }) diff --git a/apps/sim/app/api/v2/organizations/[organizationId]/usage/route.test.ts b/apps/sim/app/api/v2/organizations/[organizationId]/usage/route.test.ts index aa72574f436..0336f0878f4 100644 --- a/apps/sim/app/api/v2/organizations/[organizationId]/usage/route.test.ts +++ b/apps/sim/app/api/v2/organizations/[organizationId]/usage/route.test.ts @@ -57,7 +57,7 @@ vi.mock('@/lib/billing/organizations/member-limits', () => ({ vi.mock('@/lib/billing/core/usage-analytics-queries', () => ({ readUsageTotals: mocks.totals, readUsageTimeSeries: mocks.series, - readUsageBreakdown: mocks.breakdown, + readUsageGroups: mocks.breakdown, readUsageEntityNames: vi.fn().mockResolvedValue(new Map()), })) vi.mock('@/lib/billing/core/usage-log', () => ({ getBillingEntityUsageLogs: mocks.logs })) @@ -431,7 +431,9 @@ describe('organization usage API authorization and bounds', () => { ) const response = await breakdown(request('usage/breakdown?dimension=member'), usageContext) expect(response.status).toBe(413) - expect(mocks.breakdown).toHaveBeenCalledWith(expect.any(Array), 'member', undefined, 10_000) + expect(mocks.breakdown).toHaveBeenCalledWith( + expect.objectContaining({ dimension: 'member', maxRows: 10_000 }) + ) }) }) diff --git a/apps/sim/ee/organization-usage/components/activity-summary.tsx b/apps/sim/ee/organization-usage/components/activity-summary.tsx index 47d42aceed1..75376f8e62d 100644 --- a/apps/sim/ee/organization-usage/components/activity-summary.tsx +++ b/apps/sim/ee/organization-usage/components/activity-summary.tsx @@ -1,12 +1,51 @@ 'use client' import { useMemo } from 'react' -import { BarChart, ChartFrame, DashboardMetric, DonutChart, formatChartLatency } from '@sim/emcn' +import { + BarChart, + type BarChartSeries, + ChartFrame, + ChartLegend, + type ChartLegendItem, + cn, + DashboardMetric, + formatChartLatency, +} from '@sim/emcn' import type { OrganizationActivitySummary } from '@/lib/api/contracts/organization-activity' import { SettingsSection } from '@/app/workspace/[workspaceId]/settings/components/settings-section/settings-section' +import { + USAGE_CHAT_COLOR, + USAGE_OTHER_COLOR, + USAGE_PALETTE_CLASS, +} from '@/ee/organization-usage/constants' +import { useLegendHighlight } from '@/ee/organization-usage/hooks/use-legend-highlight' import { useOrganizationActivitySummary } from '@/hooks/queries/organization-activity' import type { OrganizationUsageWindowKey } from '@/hooks/queries/utils/organization-usage-keys' +const CHART_HEIGHT = 180 + +/** + * Outcome layers, bottom-up. Failed is the status red and sits on the stack where a + * spike reads at a glance; Other (cancelled, paused, unfinished) stays neutral, in a + * gray whose lightness keeps it apart from the red for color-vision deficiency. + */ +const OUTCOMES = [ + { id: 'completed', label: 'Completed', color: 'var(--brand-blue)' }, + { id: 'failed', label: 'Failed', color: 'var(--text-error)' }, + { id: 'other', label: 'Other', color: USAGE_OTHER_COLOR }, +] as const + +const OUTCOME_LEGEND: ChartLegendItem[] = [...OUTCOMES] +const OUTCOME_IDS = OUTCOMES.map((outcome) => outcome.id) + +type ActivityPoint = OrganizationActivitySummary['series'][number] + +const OUTCOME_VALUE: Record<(typeof OUTCOMES)[number]['id'], (point: ActivityPoint) => number> = { + completed: (point) => point.completed, + failed: (point) => point.failed, + other: (point) => Math.max(0, point.workflowRuns - point.completed - point.failed), +} + interface ActivitySummaryProps { summary?: OrganizationActivitySummary loading?: boolean @@ -24,30 +63,26 @@ export function formatFailureRate(rate: number | null): string { } export function ActivitySummary({ summary, loading, error, onRetry }: ActivitySummaryProps) { - const workflowSeries = useMemo( + const highlight = useLegendHighlight(OUTCOME_IDS) + + const outcomeSeries = useMemo( () => - summary?.series.map((point) => ({ - timestamp: point.timestamp, - value: point.workflowRuns, - })) ?? [], + OUTCOMES.map((outcome) => ({ + ...outcome, + data: (summary?.series ?? []).map((point) => ({ + timestamp: point.timestamp, + value: OUTCOME_VALUE[outcome.id](point), + })), + })), [summary?.series] ) + const chatSeries = useMemo( () => - summary?.series.map((point) => ({ - timestamp: point.timestamp, - value: point.chatRuns, - })) ?? [], - [summary?.series] - ) - const failureSeries = useMemo( - () => - summary?.series.map((point) => ({ - timestamp: point.timestamp, - value: point.failed, - })) ?? [], + summary?.series.map((point) => ({ timestamp: point.timestamp, value: point.chatRuns })) ?? [], [summary?.series] ) + const totals = summary?.totals const metrics = [ { @@ -81,18 +116,10 @@ export function ActivitySummary({ summary, loading, error, onRetry }: ActivitySu description: 'Completed and failed workflows with a recorded duration.', }, ] - const outcomes = [ - { label: 'Completed', value: totals?.completed ?? 0, color: 'var(--indicator-seat-filled)' }, - { label: 'Failed', value: totals?.failed ?? 0, color: 'var(--text-error)' }, - { - label: 'Other', - value: totals ? totals.workflowRuns - totals.completed - totals.failed : 0, - color: 'var(--text-muted)', - }, - ] const chartState = { loading, error: error ? "Couldn't load activity." : undefined, onRetry } + return ( -
+
{metrics.map((metric) => (
- +
+ + + + +
+ - - - - - - - - -
) diff --git a/apps/sim/ee/organization-usage/components/usage-consumers.tsx b/apps/sim/ee/organization-usage/components/usage-consumers.tsx index 8a508eaf2c9..aef6f01de4e 100644 --- a/apps/sim/ee/organization-usage/components/usage-consumers.tsx +++ b/apps/sim/ee/organization-usage/components/usage-consumers.tsx @@ -1,6 +1,6 @@ 'use client' -import type { ComponentType } from 'react' +import type { ComponentType, ReactNode } from 'react' import { cn, disclosureChevronClass, formatChartCompactNumber } from '@sim/emcn' import { ArrowRight, ChevronDown } from '@sim/emcn/icons' import { @@ -89,6 +89,8 @@ export const USAGE_PROVIDER_ICON_IDS = Object.keys(PROVIDER_ICONS) interface UsageConsumerRowProps { row: OrganizationUsageBreakdownRow + /** Replaces the provider mark, e.g. with a member's avatar. */ + leading?: ReactNode /** BYOK rows carry no cost, so tokens are the only usage they can show. */ showTokensOnly: boolean onSelect?: (row: OrganizationUsageBreakdownRow) => void @@ -126,6 +128,7 @@ export const USAGE_ROW_CLASSES = 'flex w-full items-center gap-2.5 rounded-lg p- */ function UsageConsumerRow({ row, + leading, showTokensOnly, onSelect, actions, @@ -148,14 +151,15 @@ function UsageConsumerRow({ onSelect && 'transition-colors hover-hover:bg-[var(--surface-active)]' )} > - {ProviderIcon && } + {leading ?? + (ProviderIcon && )} {row.label}