feat(code): undo-redo state (#3018)

* feat(code): undo-redo state

* address greptile

* address bugbot comments

* fix debounce flush

* inc debounce time

* fix wand case

* address comments

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/code/code.tsx

@@ -39,6 +39,8 @@ import { normalizeName } from '@/executor/constants'
 import { createEnvVarPattern, createReferencePattern } from '@/executor/utils/reference-validation'
 import { useTagSelection } from '@/hooks/kb/use-tag-selection'
 import { createShouldHighlightEnvVar, useAvailableEnvVarKeys } from '@/hooks/use-available-env-vars'
+import { useCodeUndoRedo } from '@/hooks/use-code-undo-redo'
+import { useWorkflowStore } from '@/stores/workflows/workflow/store'
 
 const logger = createLogger('Code')
 
@@ -212,7 +214,6 @@ export const Code = memo(function Code({
   const handleStreamStartRef = useRef<() => void>(() => {})
   const handleGeneratedContentRef = useRef<(generatedCode: string) => void>(() => {})
   const handleStreamChunkRef = useRef<(chunk: string) => void>(() => {})
-  const hasEditedSinceFocusRef = useRef(false)
   const codeRef = useRef(code)
   codeRef.current = code
 
@@ -220,8 +221,12 @@ export const Code = memo(function Code({
   const emitTagSelection = useTagSelection(blockId, subBlockId)
   const [languageValue] = useSubBlockValue<string>(blockId, 'language')
   const availableEnvVars = useAvailableEnvVarKeys(workspaceId)
+  const blockType = useWorkflowStore(
+    useCallback((state) => state.blocks?.[blockId]?.type, [blockId])
+  )
 
   const effectiveLanguage = (languageValue as 'javascript' | 'python' | 'json') || language
+  const isFunctionCode = blockType === 'function' && subBlockId === 'code'
 
   const trimmedCode = code.trim()
   const containsReferencePlaceholders =
@@ -296,6 +301,15 @@ export const Code = memo(function Code({
   const updatePromptValue = wandHook?.updatePromptValue || (() => {})
   const cancelGeneration = wandHook?.cancelGeneration || (() => {})
 
+  const { recordChange, recordReplace, flushPending, startSession, undo, redo } = useCodeUndoRedo({
+    blockId,
+    subBlockId,
+    value: code,
+    enabled: isFunctionCode,
+    isReadOnly: readOnly || disabled || isPreview,
+    isStreaming: isAiStreaming,
+  })
+
   const [storeValue, setStoreValue] = useSubBlockValue(blockId, subBlockId, false, {
     isStreaming: isAiStreaming,
     onStreamingEnd: () => {
@@ -347,9 +361,10 @@ export const Code = memo(function Code({
       setCode(generatedCode)
       if (!isPreview && !disabled) {
         setStoreValue(generatedCode)
+        recordReplace(generatedCode)
       }
     }
-  }, [isPreview, disabled, setStoreValue])
+  }, [disabled, isPreview, recordReplace, setStoreValue])
 
   useEffect(() => {
     if (!editorRef.current) return
@@ -492,7 +507,7 @@ export const Code = memo(function Code({
 
       setCode(newValue)
       setStoreValue(newValue)
-      hasEditedSinceFocusRef.current = true
+      recordChange(newValue)
       const newCursorPosition = dropPosition + 1
       setCursorPosition(newCursorPosition)
 
@@ -521,7 +536,7 @@ export const Code = memo(function Code({
     if (!isPreview && !readOnly) {
       setCode(newValue)
       emitTagSelection(newValue)
-      hasEditedSinceFocusRef.current = true
+      recordChange(newValue)
     }
     setShowTags(false)
     setActiveSourceBlockId(null)
@@ -539,7 +554,7 @@ export const Code = memo(function Code({
     if (!isPreview && !readOnly) {
       setCode(newValue)
       emitTagSelection(newValue)
-      hasEditedSinceFocusRef.current = true
+      recordChange(newValue)
     }
     setShowEnvVars(false)
 
@@ -625,9 +640,9 @@ export const Code = memo(function Code({
   const handleValueChange = useCallback(
     (newCode: string) => {
       if (!isAiStreaming && !isPreview && !disabled && !readOnly) {
-        hasEditedSinceFocusRef.current = true
         setCode(newCode)
         setStoreValue(newCode)
+        recordChange(newCode)
 
         const textarea = editorRef.current?.querySelector('textarea')
         if (textarea) {
@@ -646,7 +661,7 @@ export const Code = memo(function Code({
         }
       }
     },
-    [isAiStreaming, isPreview, disabled, readOnly, setStoreValue]
+    [isAiStreaming, isPreview, disabled, readOnly, recordChange, setStoreValue]
   )
 
   const handleKeyDown = useCallback(
@@ -657,21 +672,39 @@ export const Code = memo(function Code({
       }
       if (isAiStreaming) {
         e.preventDefault()
+        return
       }
-      if (e.key === 'z' && (e.metaKey || e.ctrlKey) && !hasEditedSinceFocusRef.current) {
+      if (!isFunctionCode) return
+      const isUndo = (e.key === 'z' || e.key === 'Z') && (e.metaKey || e.ctrlKey) && !e.shiftKey
+      const isRedo =
+        ((e.key === 'z' || e.key === 'Z') && (e.metaKey || e.ctrlKey) && e.shiftKey) ||
+        (e.key === 'y' && (e.metaKey || e.ctrlKey))
+      if (isUndo) {
         e.preventDefault()
+        e.stopPropagation()
+        undo()
+        return
+      }
+      if (isRedo) {
+        e.preventDefault()
+        e.stopPropagation()
+        redo()
       }
     },
-    [isAiStreaming]
+    [isAiStreaming, isFunctionCode, redo, undo]
   )
 
   const handleEditorFocus = useCallback(() => {
-    hasEditedSinceFocusRef.current = false
+    startSession(codeRef.current)
     if (!isPreview && !disabled && !readOnly && codeRef.current.trim() === '') {
       setShowTags(true)
       setCursorPosition(0)
     }
-  }, [isPreview, disabled, readOnly])
+  }, [disabled, isPreview, readOnly, startSession])
+
+  const handleEditorBlur = useCallback(() => {
+    flushPending()
+  }, [flushPending])
 
   /**
    * Renders the line numbers, aligned with wrapped visual lines and highlighting the active line.
@@ -791,6 +824,7 @@ export const Code = memo(function Code({
             onValueChange={handleValueChange}
             onKeyDown={handleKeyDown}
             onFocus={handleEditorFocus}
+            onBlur={handleEditorBlur}
             highlight={highlightCode}
             {...getCodeEditorProps({ isStreaming: isAiStreaming, isPreview, disabled })}
           />

apps/sim/hooks/use-code-undo-redo.ts

@@ -0,0 +1,239 @@
+import { useCallback, useEffect, useMemo, useRef } from 'react'
+import { createLogger } from '@sim/logger'
+import { useShallow } from 'zustand/react/shallow'
+import { useCollaborativeWorkflow } from '@/hooks/use-collaborative-workflow'
+import { useCodeUndoRedoStore } from '@/stores/undo-redo'
+import { useWorkflowDiffStore } from '@/stores/workflow-diff/store'
+import { useWorkflowRegistry } from '@/stores/workflows/registry/store'
+
+const logger = createLogger('CodeUndoRedo')
+
+interface UseCodeUndoRedoOptions {
+  blockId: string
+  subBlockId: string
+  value: string
+  enabled?: boolean
+  isReadOnly?: boolean
+  isStreaming?: boolean
+  debounceMs?: number
+}
+
+export function useCodeUndoRedo({
+  blockId,
+  subBlockId,
+  value,
+  enabled = true,
+  isReadOnly = false,
+  isStreaming = false,
+  debounceMs = 500,
+}: UseCodeUndoRedoOptions) {
+  const { collaborativeSetSubblockValue } = useCollaborativeWorkflow()
+  const activeWorkflowId = useWorkflowRegistry((state) => state.activeWorkflowId)
+  const { isShowingDiff, hasActiveDiff } = useWorkflowDiffStore(
+    useShallow((state) => ({
+      isShowingDiff: state.isShowingDiff,
+      hasActiveDiff: state.hasActiveDiff,
+    }))
+  )
+
+  const isBaselineView = hasActiveDiff && !isShowingDiff
+  const isEnabled = useMemo(
+    () => Boolean(enabled && activeWorkflowId && !isReadOnly && !isStreaming && !isBaselineView),
+    [enabled, activeWorkflowId, isReadOnly, isStreaming, isBaselineView]
+  )
+  const isReplaceEnabled = useMemo(
+    () => Boolean(enabled && activeWorkflowId && !isReadOnly && !isBaselineView),
+    [enabled, activeWorkflowId, isReadOnly, isBaselineView]
+  )
+
+  const lastCommittedValueRef = useRef<string>(value ?? '')
+  const pendingBeforeRef = useRef<string | null>(null)
+  const pendingAfterRef = useRef<string | null>(null)
+  const timeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null)
+  const isApplyingRef = useRef(false)
+
+  const clearTimer = useCallback(() => {
+    if (timeoutRef.current) {
+      clearTimeout(timeoutRef.current)
+      timeoutRef.current = null
+    }
+  }, [])
+
+  const resetPending = useCallback(() => {
+    pendingBeforeRef.current = null
+    pendingAfterRef.current = null
+  }, [])
+
+  const commitPending = useCallback(() => {
+    if (!isEnabled || !activeWorkflowId) {
+      clearTimer()
+      resetPending()
+      return
+    }
+
+    const before = pendingBeforeRef.current
+    const after = pendingAfterRef.current
+    if (before === null || after === null) return
+
+    if (before === after) {
+      lastCommittedValueRef.current = after
+      clearTimer()
+      resetPending()
+      return
+    }
+
+    useCodeUndoRedoStore.getState().push({
+      id: crypto.randomUUID(),
+      createdAt: Date.now(),
+      workflowId: activeWorkflowId,
+      blockId,
+      subBlockId,
+      before,
+      after,
+    })
+
+    lastCommittedValueRef.current = after
+    clearTimer()
+    resetPending()
+  }, [activeWorkflowId, blockId, clearTimer, isEnabled, resetPending, subBlockId])
+
+  const recordChange = useCallback(
+    (nextValue: string) => {
+      if (!isEnabled || isApplyingRef.current) return
+
+      if (pendingBeforeRef.current === null) {
+        pendingBeforeRef.current = lastCommittedValueRef.current ?? ''
+      }
+
+      pendingAfterRef.current = nextValue
+      clearTimer()
+      timeoutRef.current = setTimeout(commitPending, debounceMs)
+    },
+    [clearTimer, commitPending, debounceMs, isEnabled]
+  )
+
+  const recordReplace = useCallback(
+    (nextValue: string) => {
+      if (!isReplaceEnabled || isApplyingRef.current || !activeWorkflowId) return
+
+      if (pendingBeforeRef.current !== null) {
+        commitPending()
+      }
+
+      const before = lastCommittedValueRef.current ?? ''
+      if (before === nextValue) {
+        lastCommittedValueRef.current = nextValue
+        resetPending()
+        return
+      }
+
+      useCodeUndoRedoStore.getState().push({
+        id: crypto.randomUUID(),
+        createdAt: Date.now(),
+        workflowId: activeWorkflowId,
+        blockId,
+        subBlockId,
+        before,
+        after: nextValue,
+      })
+
+      lastCommittedValueRef.current = nextValue
+      clearTimer()
+      resetPending()
+    },
+    [
+      activeWorkflowId,
+      blockId,
+      clearTimer,
+      commitPending,
+      isReplaceEnabled,
+      resetPending,
+      subBlockId,
+    ]
+  )
+
+  const flushPending = useCallback(() => {
+    if (pendingBeforeRef.current === null) return
+    clearTimer()
+    commitPending()
+  }, [clearTimer, commitPending])
+
+  const startSession = useCallback(
+    (currentValue: string) => {
+      clearTimer()
+      resetPending()
+      lastCommittedValueRef.current = currentValue ?? ''
+    },
+    [clearTimer, resetPending]
+  )
+
+  const applyValue = useCallback(
+    (nextValue: string) => {
+      if (!isEnabled) return
+      isApplyingRef.current = true
+      try {
+        collaborativeSetSubblockValue(blockId, subBlockId, nextValue)
+      } finally {
+        isApplyingRef.current = false
+      }
+      lastCommittedValueRef.current = nextValue
+      clearTimer()
+      resetPending()
+    },
+    [blockId, clearTimer, collaborativeSetSubblockValue, isEnabled, resetPending, subBlockId]
+  )
+
+  const undo = useCallback(() => {
+    if (!activeWorkflowId || !isEnabled) return
+    if (pendingBeforeRef.current !== null) {
+      flushPending()
+    }
+    const entry = useCodeUndoRedoStore.getState().undo(activeWorkflowId, blockId, subBlockId)
+    if (!entry) return
+    logger.debug('Undo code edit', { blockId, subBlockId })
+    applyValue(entry.before)
+  }, [activeWorkflowId, applyValue, blockId, flushPending, isEnabled, subBlockId])
+
+  const redo = useCallback(() => {
+    if (!activeWorkflowId || !isEnabled) return
+    if (pendingBeforeRef.current !== null) {
+      flushPending()
+    }
+    const entry = useCodeUndoRedoStore.getState().redo(activeWorkflowId, blockId, subBlockId)
+    if (!entry) return
+    logger.debug('Redo code edit', { blockId, subBlockId })
+    applyValue(entry.after)
+  }, [activeWorkflowId, applyValue, blockId, flushPending, isEnabled, subBlockId])
+
+  useEffect(() => {
+    if (isApplyingRef.current || isStreaming) return
+
+    const nextValue = value ?? ''
+
+    if (pendingBeforeRef.current !== null) {
+      if (pendingAfterRef.current !== nextValue) {
+        clearTimer()
+        resetPending()
+        lastCommittedValueRef.current = nextValue
+      }
+      return
+    }
+
+    lastCommittedValueRef.current = nextValue
+  }, [clearTimer, isStreaming, resetPending, value])
+
+  useEffect(() => {
+    return () => {
+      flushPending()
+    }
+  }, [flushPending])
+
+  return {
+    recordChange,
+    recordReplace,
+    flushPending,
+    startSession,
+    undo,
+    redo,
+  }
+}

apps/sim/hooks/use-collaborative-workflow.ts

@@ -20,7 +20,7 @@ import {
 import { useNotificationStore } from '@/stores/notifications'
 import { registerEmitFunctions, useOperationQueue } from '@/stores/operation-queue/store'
 import { usePanelEditorStore, useVariablesStore } from '@/stores/panel'
-import { useUndoRedoStore } from '@/stores/undo-redo'
+import { useCodeUndoRedoStore, useUndoRedoStore } from '@/stores/undo-redo'
 import { useWorkflowDiffStore } from '@/stores/workflow-diff/store'
 import { useWorkflowRegistry } from '@/stores/workflows/registry/store'
 import { useSubBlockStore } from '@/stores/workflows/subblock/store'
@@ -449,6 +449,10 @@ export function useCollaborativeWorkflow() {
       try {
         // The setValue function automatically uses the active workflow ID
         useSubBlockStore.getState().setValue(blockId, subblockId, value)
+        const blockType = useWorkflowStore.getState().blocks?.[blockId]?.type
+        if (activeWorkflowId && blockType === 'function' && subblockId === 'code') {
+          useCodeUndoRedoStore.getState().clear(activeWorkflowId, blockId, subblockId)
+        }
       } catch (error) {
         logger.error('Error applying remote subblock update:', error)
       } finally {

apps/sim/stores/undo-redo/code-storage.ts

@@ -0,0 +1,36 @@
+import { createLogger } from '@sim/logger'
+import { del, get, set } from 'idb-keyval'
+import type { StateStorage } from 'zustand/middleware'
+
+const logger = createLogger('CodeUndoRedoStorage')
+
+export const codeUndoRedoStorage: StateStorage = {
+  getItem: async (name: string): Promise<string | null> => {
+    if (typeof window === 'undefined') return null
+    try {
+      const value = await get<string>(name)
+      return value ?? null
+    } catch (error) {
+      logger.warn('IndexedDB read failed', { name, error })
+      return null
+    }
+  },
+
+  setItem: async (name: string, value: string): Promise<void> => {
+    if (typeof window === 'undefined') return
+    try {
+      await set(name, value)
+    } catch (error) {
+      logger.warn('IndexedDB write failed', { name, error })
+    }
+  },
+
+  removeItem: async (name: string): Promise<void> => {
+    if (typeof window === 'undefined') return
+    try {
+      await del(name)
+    } catch (error) {
+      logger.warn('IndexedDB delete failed', { name, error })
+    }
+  },
+}

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

@@ -0,0 +1,151 @@
+import { create } from 'zustand'
+import { createJSONStorage, devtools, persist } from 'zustand/middleware'
+import { codeUndoRedoStorage } from '@/stores/undo-redo/code-storage'
+
+interface CodeUndoRedoEntry {
+  id: string
+  createdAt: number
+  workflowId: string
+  blockId: string
+  subBlockId: string
+  before: string
+  after: string
+}
+
+interface CodeUndoRedoStack {
+  undo: CodeUndoRedoEntry[]
+  redo: CodeUndoRedoEntry[]
+  lastUpdated?: number
+}
+
+interface CodeUndoRedoState {
+  stacks: Record<string, CodeUndoRedoStack>
+  capacity: number
+  push: (entry: CodeUndoRedoEntry) => void
+  undo: (workflowId: string, blockId: string, subBlockId: string) => CodeUndoRedoEntry | null
+  redo: (workflowId: string, blockId: string, subBlockId: string) => CodeUndoRedoEntry | null
+  clear: (workflowId: string, blockId: string, subBlockId: string) => void
+}
+
+const DEFAULT_CAPACITY = 500
+const MAX_STACKS = 50
+
+function getStackKey(workflowId: string, blockId: string, subBlockId: string): string {
+  return `${workflowId}:${blockId}:${subBlockId}`
+}
+
+const initialState = {
+  stacks: {} as Record<string, CodeUndoRedoStack>,
+  capacity: DEFAULT_CAPACITY,
+}
+
+export const useCodeUndoRedoStore = create<CodeUndoRedoState>()(
+  devtools(
+    persist(
+      (set, get) => ({
+        ...initialState,
+        push: (entry) => {
+          if (entry.before === entry.after) return
+
+          const state = get()
+          const key = getStackKey(entry.workflowId, entry.blockId, entry.subBlockId)
+          const currentStacks = { ...state.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 stackKey of stackKeys) {
+              const t = currentStacks[stackKey].lastUpdated ?? 0
+              if (t < oldestTime) {
+                oldestTime = t
+                oldestKey = stackKey
+              }
+            }
+
+            if (oldestKey) {
+              delete currentStacks[oldestKey]
+            }
+          }
+
+          const stack = currentStacks[key] || { undo: [], redo: [] }
+
+          const newUndo = [...stack.undo, entry]
+          if (newUndo.length > state.capacity) {
+            newUndo.shift()
+          }
+
+          currentStacks[key] = {
+            undo: newUndo,
+            redo: [],
+            lastUpdated: Date.now(),
+          }
+
+          set({ stacks: currentStacks })
+        },
+        undo: (workflowId, blockId, subBlockId) => {
+          const key = getStackKey(workflowId, blockId, subBlockId)
+          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]
+
+          set({
+            stacks: {
+              ...state.stacks,
+              [key]: {
+                undo: newUndo,
+                redo: newRedo.slice(-state.capacity),
+                lastUpdated: Date.now(),
+              },
+            },
+          })
+
+          return entry
+        },
+        redo: (workflowId, blockId, subBlockId) => {
+          const key = getStackKey(workflowId, blockId, subBlockId)
+          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]
+
+          set({
+            stacks: {
+              ...state.stacks,
+              [key]: {
+                undo: newUndo.slice(-state.capacity),
+                redo: newRedo,
+                lastUpdated: Date.now(),
+              },
+            },
+          })
+
+          return entry
+        },
+        clear: (workflowId, blockId, subBlockId) => {
+          const key = getStackKey(workflowId, blockId, subBlockId)
+          const state = get()
+          const { [key]: _, ...rest } = state.stacks
+          set({ stacks: rest })
+        },
+      }),
+      {
+        name: 'code-undo-redo-store',
+        storage: createJSONStorage(() => codeUndoRedoStorage),
+        partialize: (state) => ({
+          stacks: state.stacks,
+          capacity: state.capacity,
+        }),
+      }
+    ),
+    { name: 'code-undo-redo-store' }
+  )
+)

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

@@ -1,3 +1,4 @@
+export { useCodeUndoRedoStore } from './code-store'
 export { runWithUndoRedoRecordingSuspended, useUndoRedoStore } from './store'
 export * from './types'
 export * from './utils'