fix(logs): keep snapshot dedup a single atomic upsert (no select race)

Addresses review: the DO NOTHING + follow-up select could fail if cleanup
deletes the conflicting (orphaned, aged) snapshot between the no-op insert and
the select. Revert to one atomic upsert but SET only state_hash, so RETURNING
always yields the row (no race) while the unchanged TOASTed state_data jsonb is
still not rewritten under MVCC — keeping the per-execution write tiny.

Author: waleedlatif1

apps/sim/lib/logs/execution/snapshot/service.test.ts

@@ -381,36 +381,24 @@ describe('SnapshotService', () => {
       createdAt: Date
     }
 
-    /** Mock the insert → values → onConflictDoNothing → returning chain. */
-    function mockInsertReturning(rows: SnapshotRow[]) {
+    /** Mock the insert → values → onConflictDoUpdate → returning chain. */
+    function mockUpsertReturning(rows: SnapshotRow[]) {
       let capturedConflictConfig: Record<string, unknown> | undefined
-      const onConflictDoNothing = vi.fn().mockImplementation((config: Record<string, unknown>) => {
+      const onConflictDoUpdate = vi.fn().mockImplementation((config: Record<string, unknown>) => {
         capturedConflictConfig = config
         return { returning: vi.fn().mockResolvedValue(rows) }
       })
-      const values = vi.fn().mockReturnValue({ onConflictDoNothing })
+      const values = vi.fn().mockReturnValue({ onConflictDoUpdate })
       databaseMock.db.insert = vi.fn().mockReturnValue({ values })
-      return {
-        values,
-        onConflictDoNothing,
-        getConflictConfig: () => capturedConflictConfig,
-      }
-    }
-
-    /** Mock the select → from → where → limit chain used on the reuse path. */
-    function mockSelectReturning(rows: SnapshotRow[]) {
-      const limit = vi.fn().mockResolvedValue(rows)
-      const where = vi.fn().mockReturnValue({ limit })
-      const from = vi.fn().mockReturnValue({ where })
-      databaseMock.db.select = vi.fn().mockReturnValue({ from })
-      return databaseMock.db.select
+      databaseMock.db.select = vi.fn()
+      return { values, onConflictDoUpdate, getConflictConfig: () => capturedConflictConfig }
     }
 
-    it('inserts a new snapshot via onConflictDoNothing without a follow-up select', async () => {
+    it('inserts a new snapshot in a single atomic upsert', async () => {
       const service = new SnapshotService()
       const workflowId = 'wf-123'
 
-      const { values } = mockInsertReturning([
+      const { values } = mockUpsertReturning([
         {
           id: 'generated-uuid-1',
           workflowId,
@@ -419,7 +407,6 @@ describe('SnapshotService', () => {
           createdAt: new Date('2026-02-19T00:00:00Z'),
         },
       ])
-      const select = mockSelectReturning([])
 
       const result = await service.createSnapshotWithDeduplication(workflowId, mockState)
 
@@ -428,57 +415,57 @@ describe('SnapshotService', () => {
       )
       expect(result.snapshot.id).toBe('generated-uuid-1')
       expect(result.isNew).toBe(true)
-      // New row returned by the insert → no extra read needed.
-      expect(select).not.toHaveBeenCalled()
+      // Single atomic statement — never a follow-up select (which would race with cleanup).
+      expect(databaseMock.db.select).not.toHaveBeenCalled()
     })
 
-    it('does NOT rewrite state_data on conflict (onConflictDoNothing, no set clause)', async () => {
+    it('reuses the existing snapshot atomically when the returned id differs', async () => {
       const service = new SnapshotService()
       const workflowId = 'wf-123'
 
-      const { onConflictDoNothing, getConflictConfig } = mockInsertReturning([
+      mockUpsertReturning([
         {
-          id: 'generated-uuid-1',
+          id: 'existing-snapshot-id',
           workflowId,
           stateHash: 'abc123',
           stateData: mockState,
           createdAt: new Date('2026-02-19T00:00:00Z'),
         },
       ])
-      mockSelectReturning([])
 
-      await service.createSnapshotWithDeduplication(workflowId, mockState)
+      const result = await service.createSnapshotWithDeduplication(workflowId, mockState)
 
-      expect(onConflictDoNothing).toHaveBeenCalledTimes(1)
-      const config = getConflictConfig()
-      expect(config?.target).toBeDefined()
-      // The whole point of this change: no SET clause, so the large jsonb is never rewritten.
-      expect(config).not.toHaveProperty('set')
+      expect(result.snapshot.id).toBe('existing-snapshot-id')
+      expect(result.isNew).toBe(false)
+      expect(databaseMock.db.select).not.toHaveBeenCalled()
     })
 
-    it('reuses the existing snapshot via a follow-up select when the insert no-ops', async () => {
+    it('SET targets only state_hash on conflict, never the large state_data', async () => {
       const service = new SnapshotService()
       const workflowId = 'wf-123'
 
-      mockInsertReturning([]) // conflict → insert returns nothing
-      const select = mockSelectReturning([
+      const { onConflictDoUpdate, getConflictConfig } = mockUpsertReturning([
         {
-          id: 'existing-snapshot-id',
+          id: 'generated-uuid-1',
           workflowId,
           stateHash: 'abc123',
           stateData: mockState,
           createdAt: new Date('2026-02-19T00:00:00Z'),
         },
       ])
 
-      const result = await service.createSnapshotWithDeduplication(workflowId, mockState)
+      await service.createSnapshotWithDeduplication(workflowId, mockState)
 
-      expect(result.snapshot.id).toBe('existing-snapshot-id')
-      expect(result.isNew).toBe(false)
-      expect(select).toHaveBeenCalledTimes(1)
+      expect(onConflictDoUpdate).toHaveBeenCalledTimes(1)
+      const config = getConflictConfig()
+      expect(config?.target).toBeDefined()
+      // The crux of this change: the SET touches state_hash only, so the unchanged
+      // TOASTed state_data jsonb is never rewritten.
+      expect(config?.set).toHaveProperty('stateHash')
+      expect(config?.set).not.toHaveProperty('stateData')
     })
 
-    it('does not throw on concurrent inserts with the same hash (loser falls back to select)', async () => {
+    it('does not throw on concurrent inserts with the same hash', async () => {
       const service = new SnapshotService()
       const workflowId = 'wf-123'
 
@@ -491,39 +478,24 @@ describe('SnapshotService', () => {
       }
       const existingRow: SnapshotRow = { ...newRow, id: 'existing-snapshot-id' }
 
-      // First caller wins the insert; second caller's insert no-ops and selects.
-      let insertCall = 0
+      let upsertCall = 0
       databaseMock.db.insert = vi.fn().mockImplementation(() => ({
         values: vi.fn().mockReturnValue({
-          onConflictDoNothing: vi.fn().mockReturnValue({
-            returning: vi.fn().mockResolvedValue(insertCall++ === 0 ? [newRow] : []),
+          onConflictDoUpdate: vi.fn().mockReturnValue({
+            returning: vi.fn().mockResolvedValue(upsertCall++ === 0 ? [newRow] : [existingRow]),
           }),
         }),
       }))
-      mockSelectReturning([existingRow])
 
       const [result1, result2] = await Promise.all([
         service.createSnapshotWithDeduplication(workflowId, mockState),
         service.createSnapshotWithDeduplication(workflowId, mockState),
       ])
 
-      const byId = [result1, result2].sort((a, b) => a.snapshot.id.localeCompare(b.snapshot.id))
-      expect(byId[0].snapshot.id).toBe('existing-snapshot-id')
-      expect(byId[0].isNew).toBe(false)
-      expect(byId[1].snapshot.id).toBe('generated-uuid-1')
-      expect(byId[1].isNew).toBe(true)
-    })
-
-    it('throws a descriptive error when neither the insert nor the select yields a row', async () => {
-      const service = new SnapshotService()
-      const workflowId = 'wf-123'
-
-      mockInsertReturning([])
-      mockSelectReturning([])
-
-      await expect(service.createSnapshotWithDeduplication(workflowId, mockState)).rejects.toThrow(
-        /Failed to create or load execution snapshot/
-      )
+      expect(result1.snapshot.id).toBe('generated-uuid-1')
+      expect(result1.isNew).toBe(true)
+      expect(result2.snapshot.id).toBe('existing-snapshot-id')
+      expect(result2.isNew).toBe(false)
     })
   })
 

apps/sim/lib/logs/execution/snapshot/service.ts

@@ -38,51 +38,31 @@ export class SnapshotService implements ISnapshotService {
     }
 
     /**
-     * Insert the snapshot, or do nothing if a row already exists for this
-     * (workflowId, stateHash). The hash is a sha256 of the normalized state, so
-     * an existing row's stateData is byte-identical — there is nothing to update.
+     * Insert the snapshot, or — when an identical (workflowId, stateHash) row
+     * already exists — return it without rewriting the large stateData jsonb.
      *
-     * The previous implementation used onConflictDoUpdate(set state_data), which
-     * rewrote the full (tens-of-KB) state jsonb on every execution. Under Postgres
-     * MVCC that churned a dead tuple + TOAST/WAL per run for no change.
-     * onConflictDoNothing avoids the write entirely on the reuse path.
+     * The hash is a sha256 of the normalized state, so an existing row's stateData
+     * is byte-identical; there is nothing to update. The previous implementation
+     * SET state_data on conflict, which rewrote the full (tens-of-KB) jsonb every
+     * run. We keep a single atomic upsert — so RETURNING always yields the row and
+     * there is no race with snapshot cleanup (unlike DO NOTHING + a follow-up
+     * select) — but SET only the small state_hash column to itself. Under Postgres
+     * MVCC the unchanged, TOASTed stateData is not rewritten: its existing
+     * out-of-line storage is reused, so the per-execution write drops from the
+     * full blob to a tiny heap tuple.
      */
-    const [insertedSnapshot] = await db
+    const [upsertedSnapshot] = await db
       .insert(workflowExecutionSnapshots)
       .values(snapshotData)
-      .onConflictDoNothing({
+      .onConflictDoUpdate({
         target: [workflowExecutionSnapshots.workflowId, workflowExecutionSnapshots.stateHash],
+        set: {
+          stateHash: sql`excluded.state_hash`,
+        },
       })
       .returning()
 
-    const isNew = Boolean(insertedSnapshot)
-
-    /**
-     * On conflict the insert returns no row, so load the existing snapshot by its
-     * unique (workflowId, stateHash). A freshly created snapshot cannot be removed
-     * in this window — cleanupOrphanedSnapshots only targets rows older than its
-     * cutoff — so this lookup is guaranteed to resolve.
-     */
-    const snapshotRow =
-      insertedSnapshot ??
-      (
-        await db
-          .select()
-          .from(workflowExecutionSnapshots)
-          .where(
-            and(
-              eq(workflowExecutionSnapshots.workflowId, workflowId),
-              eq(workflowExecutionSnapshots.stateHash, stateHash)
-            )
-          )
-          .limit(1)
-      )[0]
-
-    if (!snapshotRow) {
-      throw new Error(
-        `Failed to create or load execution snapshot for workflow ${workflowId} (hash: ${stateHash.slice(0, 12)}...)`
-      )
-    }
+    const isNew = upsertedSnapshot.id === snapshotData.id
 
     logger.info(
       isNew
@@ -92,9 +72,9 @@ export class SnapshotService implements ISnapshotService {
 
     return {
       snapshot: {
-        ...snapshotRow,
-        stateData: snapshotRow.stateData as WorkflowState,
-        createdAt: snapshotRow.createdAt.toISOString(),
+        ...upsertedSnapshot,
+        stateData: upsertedSnapshot.stateData as WorkflowState,
+        createdAt: upsertedSnapshot.createdAt.toISOString(),
       },
       isNew,
     }