chore: import upstream snapshot with attribution

apps/sim/stores/undo-redo/store.ts

@@ -0,0 +1,511 @@
+import { createLogger } from '@sim/logger'
+import { UNDO_REDO_OPERATIONS } from '@sim/realtime-protocol/constants'
+import type { Edge } from 'reactflow'
+import { create } from 'zustand'
+import { createJSONStorage, persist } from 'zustand/middleware'
+import type {
+  BatchAddBlocksOperation,
+  BatchAddEdgesOperation,
+  BatchMoveBlocksOperation,
+  BatchRemoveBlocksOperation,
+  BatchRemoveEdgesOperation,
+  BatchUpdateParentOperation,
+  Operation,
+  OperationEntry,
+  UndoRedoState,
+} from '@/stores/undo-redo/types'
+import type { BlockState } from '@/stores/workflows/workflow/types'
+
+const logger = createLogger('UndoRedoStore')
+const DEFAULT_CAPACITY = 100
+const MAX_STACKS = 5
+
+let recordingSuspendDepth = 0
+
+function isRecordingSuspended(): boolean {
+  return recordingSuspendDepth > 0
+}
+
+/**
+ * Temporarily suspends undo/redo recording while the provided callback runs.
+ *
+ * @param callback - Function to execute while recording is disabled.
+ * @returns The callback result.
+ */
+export async function runWithUndoRedoRecordingSuspended<T>(
+  callback: () => Promise<T> | T
+): Promise<T> {
+  recordingSuspendDepth += 1
+  try {
+    return await Promise.resolve(callback())
+  } finally {
+    recordingSuspendDepth = Math.max(0, recordingSuspendDepth - 1)
+  }
+}
+
+function getStackKey(workflowId: string, userId: string): string {
+  return `${workflowId}:${userId}`
+}
+
+/**
+ * Custom storage adapter for Zustand's persist middleware.
+ * We need this wrapper to gracefully handle 'QuotaExceededError' when localStorage is full,
+ * Without this, the default storage engine would throw and crash the application.
+ * and to properly handle SSR/Node.js environments.
+ */
+const safeStorageAdapter = {
+  getItem: (name: string): string | null => {
+    if (typeof localStorage === 'undefined') return null
+    try {
+      return localStorage.getItem(name)
+    } catch (e) {
+      logger.warn('Failed to read from localStorage', e)
+      return null
+    }
+  },
+  setItem: (name: string, value: string): void => {
+    if (typeof localStorage === 'undefined') return
+    try {
+      localStorage.setItem(name, value)
+    } catch (e) {
+      // Log warning but don't crash - this handles QuotaExceededError
+      logger.warn('Failed to save to localStorage', e)
+    }
+  },
+  removeItem: (name: string): void => {
+    if (typeof localStorage === 'undefined') return
+    try {
+      localStorage.removeItem(name)
+    } catch (e) {
+      logger.warn('Failed to remove from localStorage', e)
+    }
+  },
+}
+
+function isOperationApplicable(
+  operation: Operation,
+  graph: { blocksById: Record<string, BlockState>; edgesById: Record<string, Edge> }
+): boolean {
+  switch (operation.type) {
+    case UNDO_REDO_OPERATIONS.BATCH_REMOVE_BLOCKS: {
+      const op = operation as BatchRemoveBlocksOperation
+      return op.data.blockSnapshots.every((block) => Boolean(graph.blocksById[block.id]))
+    }
+    case UNDO_REDO_OPERATIONS.BATCH_ADD_BLOCKS: {
+      const op = operation as BatchAddBlocksOperation
+      return op.data.blockSnapshots.every((block) => !graph.blocksById[block.id])
+    }
+    case UNDO_REDO_OPERATIONS.BATCH_MOVE_BLOCKS: {
+      const op = operation as BatchMoveBlocksOperation
+      return op.data.moves.every((move) => Boolean(graph.blocksById[move.blockId]))
+    }
+    case UNDO_REDO_OPERATIONS.UPDATE_PARENT: {
+      const blockId = operation.data.blockId
+      return Boolean(graph.blocksById[blockId])
+    }
+    case UNDO_REDO_OPERATIONS.BATCH_UPDATE_PARENT: {
+      const op = operation as BatchUpdateParentOperation
+      return op.data.updates.every((u) => Boolean(graph.blocksById[u.blockId]))
+    }
+    case UNDO_REDO_OPERATIONS.BATCH_REMOVE_EDGES: {
+      const op = operation as BatchRemoveEdgesOperation
+      return op.data.edgeSnapshots.every((edge) => Boolean(graph.edgesById[edge.id]))
+    }
+    case UNDO_REDO_OPERATIONS.BATCH_ADD_EDGES: {
+      const op = operation as BatchAddEdgesOperation
+      return op.data.edgeSnapshots.every((edge) => !graph.edgesById[edge.id])
+    }
+    default:
+      return true
+  }
+}
+
+export const useUndoRedoStore = create<UndoRedoState>()(
+  persist(
+    (set, get) => ({
+      stacks: {},
+      capacity: DEFAULT_CAPACITY,
+
+      push: (workflowId: string, userId: string, entry: OperationEntry) => {
+        if (isRecordingSuspended()) {
+          logger.debug('Skipped push while undo/redo recording suspended', {
+            workflowId,
+            userId,
+            operationType: entry.operation.type,
+          })
+          return
+        }
+
+        const key = getStackKey(workflowId, userId)
+        const state = get()
+        const currentStacks = { ...state.stacks }
+
+        // Limit number of stacks
+        const stackKeys = Object.keys(currentStacks)
+        if (stackKeys.length >= MAX_STACKS && !currentStacks[key]) {
+          let oldestKey: string | null = null
+          let oldestTime = Number.POSITIVE_INFINITY
+
+          for (const k of stackKeys) {
+            const t = currentStacks[k].lastUpdated ?? 0
+            if (t < oldestTime) {
+              oldestTime = t
+              oldestKey = k
+            }
+          }
+
+          if (oldestKey) {
+            delete currentStacks[oldestKey]
+          }
+        }
+
+        const stack = currentStacks[key] || { undo: [], redo: [] }
+
+        // Prevent duplicate diff operations (apply-diff, accept-diff, reject-diff)
+        if (['apply-diff', 'accept-diff', 'reject-diff'].includes(entry.operation.type)) {
+          const lastEntry = stack.undo[stack.undo.length - 1]
+          if (lastEntry && lastEntry.operation.type === entry.operation.type) {
+            // Check if it's a duplicate by comparing the relevant state data
+            const lastData = lastEntry.operation.data as any
+            const newData = entry.operation.data as any
+
+            // For each diff operation type, check the relevant state
+            let isDuplicate = false
+            if (entry.operation.type === 'apply-diff') {
+              isDuplicate =
+                JSON.stringify(lastData.baselineSnapshot?.blocks) ===
+                  JSON.stringify(newData.baselineSnapshot?.blocks) &&
+                JSON.stringify(lastData.proposedState?.blocks) ===
+                  JSON.stringify(newData.proposedState?.blocks)
+            } else if (entry.operation.type === 'accept-diff') {
+              isDuplicate =
+                JSON.stringify(lastData.afterAccept?.blocks) ===
+                JSON.stringify(newData.afterAccept?.blocks)
+            } else if (entry.operation.type === 'reject-diff') {
+              isDuplicate =
+                JSON.stringify(lastData.afterReject?.blocks) ===
+                JSON.stringify(newData.afterReject?.blocks)
+            }
+
+            if (isDuplicate) {
+              logger.debug('Skipping duplicate diff operation', {
+                type: entry.operation.type,
+                workflowId,
+                userId,
+              })
+              return
+            }
+          }
+        }
+
+        // Coalesce consecutive batch-move-blocks operations for overlapping blocks
+        if (entry.operation.type === 'batch-move-blocks') {
+          const incoming = entry.operation as BatchMoveBlocksOperation
+          const last = stack.undo[stack.undo.length - 1]
+
+          // Skip no-op moves (all moves have same before/after)
+          const allNoOp = incoming.data.moves.every((move) => {
+            const sameParent = (move.before.parentId ?? null) === (move.after.parentId ?? null)
+            return move.before.x === move.after.x && move.before.y === move.after.y && sameParent
+          })
+          if (allNoOp) {
+            logger.debug('Skipped no-op batch move push')
+            return
+          }
+
+          if (
+            last &&
+            last.operation.type === 'batch-move-blocks' &&
+            last.inverse.type === 'batch-move-blocks'
+          ) {
+            const prev = last.operation as BatchMoveBlocksOperation
+            const prevBlockIds = new Set(prev.data.moves.map((m) => m.blockId))
+            const incomingBlockIds = new Set(incoming.data.moves.map((m) => m.blockId))
+
+            // Check if same set of blocks
+            const sameBlocks =
+              prevBlockIds.size === incomingBlockIds.size &&
+              [...prevBlockIds].every((id) => incomingBlockIds.has(id))
+
+            if (sameBlocks) {
+              // Merge: keep earliest before, latest after for each block
+              const mergedMoves = incoming.data.moves.map((incomingMove) => {
+                const prevMove = prev.data.moves.find((m) => m.blockId === incomingMove.blockId)!
+                return {
+                  blockId: incomingMove.blockId,
+                  before: prevMove.before,
+                  after: incomingMove.after,
+                }
+              })
+
+              // Check if all moves result in same position (net no-op)
+              const allSameAfter = mergedMoves.every((move) => {
+                const sameParent = (move.before.parentId ?? null) === (move.after.parentId ?? null)
+                return (
+                  move.before.x === move.after.x && move.before.y === move.after.y && sameParent
+                )
+              })
+
+              const newUndoCoalesced: OperationEntry[] = allSameAfter
+                ? stack.undo.slice(0, -1)
+                : (() => {
+                    const op = entry.operation as BatchMoveBlocksOperation
+                    const inv = entry.inverse as BatchMoveBlocksOperation
+                    const newEntry: OperationEntry = {
+                      id: entry.id,
+                      createdAt: entry.createdAt,
+                      operation: {
+                        id: op.id,
+                        type: 'batch-move-blocks',
+                        timestamp: op.timestamp,
+                        workflowId,
+                        userId,
+                        data: { moves: mergedMoves },
+                      },
+                      inverse: {
+                        id: inv.id,
+                        type: 'batch-move-blocks',
+                        timestamp: inv.timestamp,
+                        workflowId,
+                        userId,
+                        data: {
+                          moves: mergedMoves.map((m) => ({
+                            blockId: m.blockId,
+                            before: m.after,
+                            after: m.before,
+                          })),
+                        },
+                      },
+                    }
+                    return [...stack.undo.slice(0, -1), newEntry]
+                  })()
+
+              currentStacks[key] = {
+                undo: newUndoCoalesced,
+                redo: [],
+                lastUpdated: Date.now(),
+              }
+
+              set({ stacks: currentStacks })
+
+              logger.debug('Coalesced consecutive batch move operations', {
+                workflowId,
+                userId,
+                blockCount: mergedMoves.length,
+                undoSize: newUndoCoalesced.length,
+              })
+              return
+            }
+          }
+        }
+
+        const newUndo = [...stack.undo, entry]
+        if (newUndo.length > state.capacity) {
+          newUndo.shift()
+        }
+
+        currentStacks[key] = {
+          undo: newUndo,
+          redo: [],
+          lastUpdated: Date.now(),
+        }
+
+        set({ stacks: currentStacks })
+
+        logger.debug('Pushed operation to undo stack', {
+          workflowId,
+          userId,
+          operationType: entry.operation.type,
+          undoSize: newUndo.length,
+        })
+      },
+
+      undo: (workflowId: string, userId: string) => {
+        const key = getStackKey(workflowId, userId)
+        const state = get()
+        const stack = state.stacks[key]
+
+        if (!stack || stack.undo.length === 0) {
+          return null
+        }
+
+        const entry = stack.undo[stack.undo.length - 1]
+        const newUndo = stack.undo.slice(0, -1)
+        const newRedo = [...stack.redo, entry]
+
+        if (newRedo.length > state.capacity) {
+          newRedo.shift()
+        }
+
+        set({
+          stacks: {
+            ...state.stacks,
+            [key]: {
+              undo: newUndo,
+              redo: newRedo,
+              lastUpdated: Date.now(),
+            },
+          },
+        })
+
+        logger.debug('Undo operation', {
+          workflowId,
+          userId,
+          operationType: entry.operation.type,
+          undoSize: newUndo.length,
+          redoSize: newRedo.length,
+        })
+
+        return entry
+      },
+
+      redo: (workflowId: string, userId: string) => {
+        const key = getStackKey(workflowId, userId)
+        const state = get()
+        const stack = state.stacks[key]
+
+        if (!stack || stack.redo.length === 0) {
+          return null
+        }
+
+        const entry = stack.redo[stack.redo.length - 1]
+        const newRedo = stack.redo.slice(0, -1)
+        const newUndo = [...stack.undo, entry]
+
+        if (newUndo.length > state.capacity) {
+          newUndo.shift()
+        }
+
+        set({
+          stacks: {
+            ...state.stacks,
+            [key]: {
+              undo: newUndo,
+              redo: newRedo,
+              lastUpdated: Date.now(),
+            },
+          },
+        })
+
+        logger.debug('Redo operation', {
+          workflowId,
+          userId,
+          operationType: entry.operation.type,
+          undoSize: newUndo.length,
+          redoSize: newRedo.length,
+        })
+
+        return entry
+      },
+
+      clear: (workflowId: string, userId: string) => {
+        const key = getStackKey(workflowId, userId)
+        const state = get()
+        const { [key]: _, ...rest } = state.stacks
+
+        set({ stacks: rest })
+
+        logger.debug('Cleared undo/redo stacks', { workflowId, userId })
+      },
+
+      clearRedo: (workflowId: string, userId: string) => {
+        const key = getStackKey(workflowId, userId)
+        const state = get()
+        const stack = state.stacks[key]
+
+        if (!stack) return
+
+        set({
+          stacks: {
+            ...state.stacks,
+            [key]: { ...stack, redo: [] },
+          },
+        })
+
+        logger.debug('Cleared redo stack', { workflowId, userId })
+      },
+
+      getStackSizes: (workflowId: string, userId: string) => {
+        const key = getStackKey(workflowId, userId)
+        const state = get()
+        const stack = state.stacks[key]
+
+        if (!stack) {
+          return { undoSize: 0, redoSize: 0 }
+        }
+
+        return {
+          undoSize: stack.undo.length,
+          redoSize: stack.redo.length,
+        }
+      },
+
+      setCapacity: (capacity: number) => {
+        const state = get()
+        const newStacks: typeof state.stacks = {}
+
+        for (const [key, stack] of Object.entries(state.stacks)) {
+          newStacks[key] = {
+            undo: stack.undo.slice(-capacity),
+            redo: stack.redo.slice(-capacity),
+            lastUpdated: stack.lastUpdated,
+          }
+        }
+
+        set({ capacity, stacks: newStacks })
+
+        logger.debug('Set capacity', { capacity })
+      },
+
+      pruneInvalidEntries: (
+        workflowId: string,
+        userId: string,
+        graph: { blocksById: Record<string, BlockState>; edgesById: Record<string, Edge> }
+      ) => {
+        const key = getStackKey(workflowId, userId)
+        const state = get()
+        const stack = state.stacks[key]
+
+        if (!stack) return
+
+        const originalUndoCount = stack.undo.length
+        const originalRedoCount = stack.redo.length
+
+        const validUndo = stack.undo.filter((entry) => isOperationApplicable(entry.inverse, graph))
+
+        const validRedo = stack.redo.filter((entry) =>
+          isOperationApplicable(entry.operation, graph)
+        )
+
+        const prunedUndoCount = originalUndoCount - validUndo.length
+        const prunedRedoCount = originalRedoCount - validRedo.length
+
+        if (prunedUndoCount > 0 || prunedRedoCount > 0) {
+          set({
+            stacks: {
+              ...state.stacks,
+              [key]: { ...stack, undo: validUndo, redo: validRedo },
+            },
+          })
+
+          logger.debug('Pruned invalid entries', {
+            workflowId,
+            userId,
+            prunedUndo: prunedUndoCount,
+            prunedRedo: prunedRedoCount,
+            remainingUndo: validUndo.length,
+            remainingRedo: validRedo.length,
+          })
+        }
+      },
+    }),
+    {
+      name: 'workflow-undo-redo',
+      storage: createJSONStorage(() => safeStorageAdapter),
+      partialize: (state) => ({
+        stacks: state.stacks,
+        capacity: state.capacity,
+      }),
+    }
+  )
+)