From 51891daf9a56da87c8f536053df9a64417a67010 Mon Sep 17 00:00:00 2001 From: Vikhyath Mondreti Date: Mon, 26 Jan 2026 19:40:40 -0800 Subject: [PATCH] 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 --- .../sub-block/components/code/code.tsx | 56 +++- apps/sim/hooks/use-code-undo-redo.ts | 239 ++++++++++++++++++ apps/sim/hooks/use-collaborative-workflow.ts | 6 +- apps/sim/stores/undo-redo/code-storage.ts | 36 +++ apps/sim/stores/undo-redo/code-store.ts | 151 +++++++++++ apps/sim/stores/undo-redo/index.ts | 1 + 6 files changed, 477 insertions(+), 12 deletions(-) create mode 100644 apps/sim/hooks/use-code-undo-redo.ts create mode 100644 apps/sim/stores/undo-redo/code-storage.ts create mode 100644 apps/sim/stores/undo-redo/code-store.ts diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/code/code.tsx b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/code/code.tsx index c67cb5528..6c43b2e30 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/components/panel/components/editor/components/sub-block/components/code/code.tsx +++ b/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(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 })} /> diff --git a/apps/sim/hooks/use-code-undo-redo.ts b/apps/sim/hooks/use-code-undo-redo.ts new file mode 100644 index 000000000..d0edbb535 --- /dev/null +++ b/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(value ?? '') + const pendingBeforeRef = useRef(null) + const pendingAfterRef = useRef(null) + const timeoutRef = useRef | 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, + } +} diff --git a/apps/sim/hooks/use-collaborative-workflow.ts b/apps/sim/hooks/use-collaborative-workflow.ts index 89aff1d12..5f5721549 100644 --- a/apps/sim/hooks/use-collaborative-workflow.ts +++ b/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 { diff --git a/apps/sim/stores/undo-redo/code-storage.ts b/apps/sim/stores/undo-redo/code-storage.ts new file mode 100644 index 000000000..05f86c82b --- /dev/null +++ b/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 => { + if (typeof window === 'undefined') return null + try { + const value = await get(name) + return value ?? null + } catch (error) { + logger.warn('IndexedDB read failed', { name, error }) + return null + } + }, + + setItem: async (name: string, value: string): Promise => { + if (typeof window === 'undefined') return + try { + await set(name, value) + } catch (error) { + logger.warn('IndexedDB write failed', { name, error }) + } + }, + + removeItem: async (name: string): Promise => { + if (typeof window === 'undefined') return + try { + await del(name) + } catch (error) { + logger.warn('IndexedDB delete failed', { name, error }) + } + }, +} diff --git a/apps/sim/stores/undo-redo/code-store.ts b/apps/sim/stores/undo-redo/code-store.ts new file mode 100644 index 000000000..c421126d5 --- /dev/null +++ b/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 + 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, + capacity: DEFAULT_CAPACITY, +} + +export const useCodeUndoRedoStore = create()( + 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' } + ) +) diff --git a/apps/sim/stores/undo-redo/index.ts b/apps/sim/stores/undo-redo/index.ts index d97adabaf..5c9d815fc 100644 --- a/apps/sim/stores/undo-redo/index.ts +++ b/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'