fix(audit): scope rollups to the run that recorded them and bound flush concurrency

Author: waleedlatif1

apps/sim/lib/table/application/row-audit.integration.ts

@@ -1,4 +1,4 @@
-import { flushAuditRollups } from '@sim/audit'
+import { withAuditRollupScope } from '@sim/audit'
 import { db } from '@sim/db'
 import {
   auditLog,
@@ -206,7 +206,6 @@ describe('table mutation audit history in PostgreSQL', () => {
   }
 
   async function auditRows() {
-    await flushAuditRollups()
     return db
       .select()
       .from(auditLog)
@@ -216,7 +215,9 @@ describe('table mutation audit history in PostgreSQL', () => {
 
   it("merges one execution's repeated row inserts into one row that keeps every row id and count", async () => {
     const rowIds: string[] = []
-    for (const value of ['a', 'b', 'c']) rowIds.push(await insertAsExecution('exec-1', value))
+    await withAuditRollupScope(async () => {
+      for (const value of ['a', 'b', 'c']) rowIds.push(await insertAsExecution('exec-1', value))
+    })
 
     const entries = await auditRows()
     expect(entries).toHaveLength(1)
@@ -242,9 +243,11 @@ describe('table mutation audit history in PostgreSQL', () => {
   })
 
   it('never merges row changes from two executions', async () => {
-    await insertAsExecution('exec-1', 'a')
-    await insertAsExecution('exec-2', 'b')
-    await insertAsExecution('exec-1', 'c')
+    await withAuditRollupScope(async () => {
+      await insertAsExecution('exec-1', 'a')
+      await insertAsExecution('exec-2', 'b')
+      await insertAsExecution('exec-1', 'c')
+    })
 
     const entries = await auditRows()
     expect(entries).toHaveLength(2)
@@ -258,7 +261,7 @@ describe('table mutation audit history in PostgreSQL', () => {
   })
 
   it('writes a lone execution change exactly as an unmerged entry', async () => {
-    const insertedRowId = await insertAsExecution('exec-1', 'a')
+    const insertedRowId = await withAuditRollupScope(() => insertAsExecution('exec-1', 'a'))
 
     const [entry] = await auditRows()
     expect(entry.metadata).toMatchObject({

apps/sim/lib/workflows/deployment-outbox.ts

@@ -3,6 +3,7 @@ import type { PrincipalActor } from '@sim/auth/principal'
 import { db, workflowDeploymentVersion, workflow as workflowTable } from '@sim/db'
 import { outboxEvent, workspaceOperationReceipt } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
+import { toStringOrNull } from '@sim/utils/coerce'
 import { toError } from '@sim/utils/errors'
 import { and, eq } from 'drizzle-orm'
 import { NextRequest } from 'next/server'
@@ -1528,11 +1529,25 @@ function parseOptionalPrincipalActor(value: unknown): PrincipalActor | undefined
       serviceId,
       subjectUserId: parseRequiredString(record.subjectUserId, 'actor.subjectUserId'),
       delegationId: parseRequiredString(record.delegationId, 'actor.delegationId'),
+      ...parseOptionalActorRun(record),
     }
   }
   throw new Error(`Invalid deployment outbox actor kind: ${kind}`)
 }
 
+/** The workflow run an executor delegation acted for, kept so its audit stays traceable. */
+function parseOptionalActorRun(record: Record<string, unknown>): {
+  workflowId?: string
+  executionId?: string
+} {
+  const workflowId = toStringOrNull(record.workflowId)
+  const executionId = toStringOrNull(record.executionId)
+  return {
+    ...(workflowId ? { workflowId } : {}),
+    ...(executionId ? { executionId } : {}),
+  }
+}
+
 function parseDeploymentPreparationCheckpoints(value: unknown): DeploymentPreparationCheckpoints {
   if (!value || typeof value !== 'object' || Array.isArray(value)) return {}
   const record = value as Record<string, unknown>

apps/sim/lib/workflows/executor/execution-core.ts

@@ -3,7 +3,7 @@
  * This is the SINGLE source of truth for workflow execution
  */
 
-import { flushAuditRollups } from '@sim/audit'
+import { withAuditRollupScope } from '@sim/audit'
 import { resolvePrincipalSubject } from '@sim/auth/principal'
 import { db } from '@sim/db'
 import { organization, workspace } from '@sim/db/schema'
@@ -609,17 +609,14 @@ function prefetchExecutionReads(options: ExecuteWorkflowCoreOptions): ExecutionR
  *
  * Runs under `runWorkflowOutsideContentDelivery`, so a run Chat started reads what any run reads.
  *
- * Writes pending audit rollups before returning, since a background run's process
- * may be frozen or ended once it does.
+ * Owns the audit rollups the run and its inline child runs record, and writes
+ * them before returning, since a background run's process may be frozen or
+ * ended once it does.
  */
-export async function executeWorkflowCore(
-  options: ExecuteWorkflowCoreOptions
-): Promise<ExecutionResult> {
-  try {
-    return await runWorkflowOutsideContentDelivery(() => executeWorkflowCoreInRun(options))
-  } finally {
-    await flushAuditRollups()
-  }
+export function executeWorkflowCore(options: ExecuteWorkflowCoreOptions): Promise<ExecutionResult> {
+  return withAuditRollupScope(() =>
+    runWorkflowOutsideContentDelivery(() => executeWorkflowCoreInRun(options))
+  )
 }
 
 async function executeWorkflowCoreInRun(

packages/audit/src/index.ts

@@ -1,7 +1,7 @@
 export type { AuditLogParams } from './log'
 export { recordAudit, recordAuditBatch, recordAuditOnce } from './log'
-export type { AuditRollupEntry, AuditRollupSummary } from './rollup'
-export { flushAuditRollups, recordAuditRollup } from './rollup'
+export type { AuditRollupSummary } from './rollup'
+export { recordAuditRollup, withAuditRollupScope } from './rollup'
 export type { AuditActionType, AuditResourceTypeValue } from './types'
 export { AuditAction, AuditResourceType } from './types'
 export { auditUpdatedFields } from './updated-fields'

packages/audit/src/rollup.ts

@@ -1,3 +1,4 @@
+import { AsyncLocalStorage } from 'node:async_hooks'
 import { createLogger } from '@sim/logger'
 import { type AuditLogParams, ambientSurface, insertDeferredAuditEntry } from './log'
 
@@ -12,6 +13,9 @@ const ROLLUP_MAX_ENTRIES = 100
 /** Pending rollups held by the process; reaching it writes every one of them. */
 const ROLLUP_MAX_PENDING = 1_000
 
+/** Rows inserted at once by one write, so a large flush never floods the pool. */
+const ROLLUP_WRITE_CONCURRENCY = 10
+
 export interface AuditRollupSummary {
   entries: number
   counters: Record<string, number>
@@ -41,7 +45,17 @@ export interface AuditRollupEntry {
   describe(summary: AuditRollupSummary): string
 }
 
+/**
+ * A unit of work that owns the rollups recorded inside it (a workflow run and
+ * the child runs it executes inline), so its end writes only its own rows.
+ */
+interface RollupScope {
+  id: number
+  writes: Set<Promise<void>>
+}
+
 interface PendingRollup {
+  scope: RollupScope | undefined
   first: AuditRollupEntry
   surface: string | undefined
   firstAt: Date
@@ -53,20 +67,22 @@ interface PendingRollup {
 }
 
 const pending = new Map<string, PendingRollup>()
-
-/** Writes started but not finished, so a flush also waits for rows already leaving. */
-const inFlight = new Set<Promise<void>>()
+const scopeStorage = new AsyncLocalStorage<RollupScope>()
+let nextScopeId = 0
 
 /**
  * Records a repetitive audit event, merging it with its repeats. Fire-and-forget
  * like `recordAudit`. Each pending row is written when its window closes, when
- * it reaches {@link ROLLUP_MAX_ENTRIES}, or on {@link flushAuditRollups}. The
- * window timer is not unref'd, so a process that ends naturally still writes it.
+ * it reaches {@link ROLLUP_MAX_ENTRIES}, or when the {@link withAuditRollupScope}
+ * it was recorded in ends. The window timer is not unref'd, so a process that
+ * ends naturally still writes it.
  */
 export function recordAuditRollup(entry: AuditRollupEntry): void {
   const surface = ambientSurface()
+  const scope = scopeStorage.getStore()
   const { params } = entry
   const key = JSON.stringify([
+    scope?.id ?? null,
     surface ?? null,
     params.workspaceId ?? null,
     params.actorId,
@@ -90,6 +106,7 @@ export function recordAuditRollup(entry: AuditRollupEntry): void {
 
   const now = new Date()
   const rollup: PendingRollup = {
+    scope,
     first: entry,
     surface,
     firstAt: now,
@@ -104,19 +121,33 @@ export function recordAuditRollup(entry: AuditRollupEntry): void {
 }
 
 /**
- * Writes every pending row now and waits for writes already started. Work that
- * may be frozen or ended once it returns (a workflow run) awaits this at its
- * end. Never throws.
+ * Runs `work` as the owner of the rollups recorded inside it, then writes them
+ * and waits for those writes before returning, so work whose process may be
+ * frozen or ended once it returns (a background workflow run) leaves nothing
+ * pending. Other scopes' rollups keep their own window. Each write is one
+ * insert, bounded by the database's statement and connect timeouts; a failed
+ * one is logged, never thrown.
  */
-export async function flushAuditRollups(): Promise<void> {
-  if (pending.size > 0) startWrite([...pending.keys()].map(take))
-  await Promise.all(inFlight)
+export async function withAuditRollupScope<T>(work: () => Promise<T>): Promise<T> {
+  const scope: RollupScope = { id: nextScopeId++, writes: new Set() }
+  try {
+    return await scopeStorage.run(scope, work)
+  } finally {
+    const keys = [...pending.entries()]
+      .filter(([, rollup]) => rollup.scope === scope)
+      .map(([key]) => key)
+    if (keys.length > 0) startWrite(keys.map(take))
+    await Promise.all(scope.writes)
+  }
 }
 
 function startWrite(rollups: PendingRollup[]): void {
   const done = writePending(rollups)
-  inFlight.add(done)
-  void done.finally(() => inFlight.delete(done))
+  const scopes = new Set(rollups.map((rollup) => rollup.scope))
+  for (const scope of scopes) scope?.writes.add(done)
+  void done.finally(() => {
+    for (const scope of scopes) scope?.writes.delete(done)
+  })
 }
 
 function mergeInto(rollup: PendingRollup, entry: AuditRollupEntry): void {
@@ -163,19 +194,24 @@ function toAuditParams(rollup: PendingRollup): AuditLogParams {
   }
 }
 
-/** One insert per row, so a row that cannot be written never takes others with it. */
+/**
+ * One insert per row, so a row that cannot be written never takes others with
+ * it, at most {@link ROLLUP_WRITE_CONCURRENCY} at a time.
+ */
 async function writePending(rollups: PendingRollup[]): Promise<void> {
-  await Promise.all(
-    rollups.map(async (rollup) => {
-      try {
-        await insertDeferredAuditEntry(toAuditParams(rollup), rollup.surface)
-      } catch (error) {
-        logger.error('Failed to record audit log rollup', {
-          error,
-          action: rollup.first.params.action,
-          entries: rollup.entries,
-        })
-      }
+  for (let index = 0; index < rollups.length; index += ROLLUP_WRITE_CONCURRENCY) {
+    await Promise.all(rollups.slice(index, index + ROLLUP_WRITE_CONCURRENCY).map(writeRollup))
+  }
+}
+
+async function writeRollup(rollup: PendingRollup): Promise<void> {
+  try {
+    await insertDeferredAuditEntry(toAuditParams(rollup), rollup.surface)
+  } catch (error) {
+    logger.error('Failed to record audit log rollup', {
+      error,
+      action: rollup.first.params.action,
+      entries: rollup.entries,
     })
-  )
+  }
 }

packages/testing/src/mocks/audit.mock.ts

@@ -272,8 +272,8 @@ const AuditResourceType = {
 /**
  * Controllable mock functions for `@sim/audit`.
  *
- * `mockRecordAuditOnce` and `mockFlushAuditRollups` resolve `undefined` by default,
- * like the real awaited writes.
+ * `mockRecordAuditOnce` resolves `undefined` by default, like the real awaited insert.
+ * `mockWithAuditRollupScope` runs the work it is given, like the real scope.
  *
  * @example
  * ```ts
@@ -292,7 +292,7 @@ export const auditMockFns = {
   mockRecordAuditBatch: vi.fn(),
   mockRecordAuditOnce: vi.fn(async (_id: string, _params: unknown): Promise<void> => undefined),
   mockRecordAuditRollup: vi.fn(),
-  mockFlushAuditRollups: vi.fn(async (): Promise<void> => undefined),
+  mockWithAuditRollupScope: vi.fn(<T>(work: () => Promise<T>): Promise<T> => work()),
 }
 
 /**
@@ -309,7 +309,7 @@ export const auditMock = {
   recordAuditBatch: auditMockFns.mockRecordAuditBatch,
   recordAuditOnce: auditMockFns.mockRecordAuditOnce,
   recordAuditRollup: auditMockFns.mockRecordAuditRollup,
-  flushAuditRollups: auditMockFns.mockFlushAuditRollups,
+  withAuditRollupScope: auditMockFns.mockWithAuditRollupScope,
   /**
    * Real implementation, not a stub: callers under test derive their audit
    * metadata through it, so stubbing it would erase what the test asserts.