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 })}
           />