Skip to content

Commit 51891da

Browse files
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
1 parent 9ee5dfe commit 51891da

File tree

6 files changed

+477
-12
lines changed

6 files changed

+477
-12
lines changed

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

Lines changed: 45 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ import { normalizeName } from '@/executor/constants'
3939
import { createEnvVarPattern, createReferencePattern } from '@/executor/utils/reference-validation'
4040
import { useTagSelection } from '@/hooks/kb/use-tag-selection'
4141
import { createShouldHighlightEnvVar, useAvailableEnvVarKeys } from '@/hooks/use-available-env-vars'
42+
import { useCodeUndoRedo } from '@/hooks/use-code-undo-redo'
43+
import { useWorkflowStore } from '@/stores/workflows/workflow/store'
4244

4345
const logger = createLogger('Code')
4446

@@ -212,16 +214,19 @@ export const Code = memo(function Code({
212214
const handleStreamStartRef = useRef<() => void>(() => {})
213215
const handleGeneratedContentRef = useRef<(generatedCode: string) => void>(() => {})
214216
const handleStreamChunkRef = useRef<(chunk: string) => void>(() => {})
215-
const hasEditedSinceFocusRef = useRef(false)
216217
const codeRef = useRef(code)
217218
codeRef.current = code
218219

219220
const accessiblePrefixes = useAccessibleReferencePrefixes(blockId)
220221
const emitTagSelection = useTagSelection(blockId, subBlockId)
221222
const [languageValue] = useSubBlockValue<string>(blockId, 'language')
222223
const availableEnvVars = useAvailableEnvVarKeys(workspaceId)
224+
const blockType = useWorkflowStore(
225+
useCallback((state) => state.blocks?.[blockId]?.type, [blockId])
226+
)
223227

224228
const effectiveLanguage = (languageValue as 'javascript' | 'python' | 'json') || language
229+
const isFunctionCode = blockType === 'function' && subBlockId === 'code'
225230

226231
const trimmedCode = code.trim()
227232
const containsReferencePlaceholders =
@@ -296,6 +301,15 @@ export const Code = memo(function Code({
296301
const updatePromptValue = wandHook?.updatePromptValue || (() => {})
297302
const cancelGeneration = wandHook?.cancelGeneration || (() => {})
298303

304+
const { recordChange, recordReplace, flushPending, startSession, undo, redo } = useCodeUndoRedo({
305+
blockId,
306+
subBlockId,
307+
value: code,
308+
enabled: isFunctionCode,
309+
isReadOnly: readOnly || disabled || isPreview,
310+
isStreaming: isAiStreaming,
311+
})
312+
299313
const [storeValue, setStoreValue] = useSubBlockValue(blockId, subBlockId, false, {
300314
isStreaming: isAiStreaming,
301315
onStreamingEnd: () => {
@@ -347,9 +361,10 @@ export const Code = memo(function Code({
347361
setCode(generatedCode)
348362
if (!isPreview && !disabled) {
349363
setStoreValue(generatedCode)
364+
recordReplace(generatedCode)
350365
}
351366
}
352-
}, [isPreview, disabled, setStoreValue])
367+
}, [disabled, isPreview, recordReplace, setStoreValue])
353368

354369
useEffect(() => {
355370
if (!editorRef.current) return
@@ -492,7 +507,7 @@ export const Code = memo(function Code({
492507

493508
setCode(newValue)
494509
setStoreValue(newValue)
495-
hasEditedSinceFocusRef.current = true
510+
recordChange(newValue)
496511
const newCursorPosition = dropPosition + 1
497512
setCursorPosition(newCursorPosition)
498513

@@ -521,7 +536,7 @@ export const Code = memo(function Code({
521536
if (!isPreview && !readOnly) {
522537
setCode(newValue)
523538
emitTagSelection(newValue)
524-
hasEditedSinceFocusRef.current = true
539+
recordChange(newValue)
525540
}
526541
setShowTags(false)
527542
setActiveSourceBlockId(null)
@@ -539,7 +554,7 @@ export const Code = memo(function Code({
539554
if (!isPreview && !readOnly) {
540555
setCode(newValue)
541556
emitTagSelection(newValue)
542-
hasEditedSinceFocusRef.current = true
557+
recordChange(newValue)
543558
}
544559
setShowEnvVars(false)
545560

@@ -625,9 +640,9 @@ export const Code = memo(function Code({
625640
const handleValueChange = useCallback(
626641
(newCode: string) => {
627642
if (!isAiStreaming && !isPreview && !disabled && !readOnly) {
628-
hasEditedSinceFocusRef.current = true
629643
setCode(newCode)
630644
setStoreValue(newCode)
645+
recordChange(newCode)
631646

632647
const textarea = editorRef.current?.querySelector('textarea')
633648
if (textarea) {
@@ -646,7 +661,7 @@ export const Code = memo(function Code({
646661
}
647662
}
648663
},
649-
[isAiStreaming, isPreview, disabled, readOnly, setStoreValue]
664+
[isAiStreaming, isPreview, disabled, readOnly, recordChange, setStoreValue]
650665
)
651666

652667
const handleKeyDown = useCallback(
@@ -657,21 +672,39 @@ export const Code = memo(function Code({
657672
}
658673
if (isAiStreaming) {
659674
e.preventDefault()
675+
return
660676
}
661-
if (e.key === 'z' && (e.metaKey || e.ctrlKey) && !hasEditedSinceFocusRef.current) {
677+
if (!isFunctionCode) return
678+
const isUndo = (e.key === 'z' || e.key === 'Z') && (e.metaKey || e.ctrlKey) && !e.shiftKey
679+
const isRedo =
680+
((e.key === 'z' || e.key === 'Z') && (e.metaKey || e.ctrlKey) && e.shiftKey) ||
681+
(e.key === 'y' && (e.metaKey || e.ctrlKey))
682+
if (isUndo) {
662683
e.preventDefault()
684+
e.stopPropagation()
685+
undo()
686+
return
687+
}
688+
if (isRedo) {
689+
e.preventDefault()
690+
e.stopPropagation()
691+
redo()
663692
}
664693
},
665-
[isAiStreaming]
694+
[isAiStreaming, isFunctionCode, redo, undo]
666695
)
667696

668697
const handleEditorFocus = useCallback(() => {
669-
hasEditedSinceFocusRef.current = false
698+
startSession(codeRef.current)
670699
if (!isPreview && !disabled && !readOnly && codeRef.current.trim() === '') {
671700
setShowTags(true)
672701
setCursorPosition(0)
673702
}
674-
}, [isPreview, disabled, readOnly])
703+
}, [disabled, isPreview, readOnly, startSession])
704+
705+
const handleEditorBlur = useCallback(() => {
706+
flushPending()
707+
}, [flushPending])
675708

676709
/**
677710
* Renders the line numbers, aligned with wrapped visual lines and highlighting the active line.
@@ -791,6 +824,7 @@ export const Code = memo(function Code({
791824
onValueChange={handleValueChange}
792825
onKeyDown={handleKeyDown}
793826
onFocus={handleEditorFocus}
827+
onBlur={handleEditorBlur}
794828
highlight={highlightCode}
795829
{...getCodeEditorProps({ isStreaming: isAiStreaming, isPreview, disabled })}
796830
/>
Lines changed: 239 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,239 @@
1+
import { useCallback, useEffect, useMemo, useRef } from 'react'
2+
import { createLogger } from '@sim/logger'
3+
import { useShallow } from 'zustand/react/shallow'
4+
import { useCollaborativeWorkflow } from '@/hooks/use-collaborative-workflow'
5+
import { useCodeUndoRedoStore } from '@/stores/undo-redo'
6+
import { useWorkflowDiffStore } from '@/stores/workflow-diff/store'
7+
import { useWorkflowRegistry } from '@/stores/workflows/registry/store'
8+
9+
const logger = createLogger('CodeUndoRedo')
10+
11+
interface UseCodeUndoRedoOptions {
12+
blockId: string
13+
subBlockId: string
14+
value: string
15+
enabled?: boolean
16+
isReadOnly?: boolean
17+
isStreaming?: boolean
18+
debounceMs?: number
19+
}
20+
21+
export function useCodeUndoRedo({
22+
blockId,
23+
subBlockId,
24+
value,
25+
enabled = true,
26+
isReadOnly = false,
27+
isStreaming = false,
28+
debounceMs = 500,
29+
}: UseCodeUndoRedoOptions) {
30+
const { collaborativeSetSubblockValue } = useCollaborativeWorkflow()
31+
const activeWorkflowId = useWorkflowRegistry((state) => state.activeWorkflowId)
32+
const { isShowingDiff, hasActiveDiff } = useWorkflowDiffStore(
33+
useShallow((state) => ({
34+
isShowingDiff: state.isShowingDiff,
35+
hasActiveDiff: state.hasActiveDiff,
36+
}))
37+
)
38+
39+
const isBaselineView = hasActiveDiff && !isShowingDiff
40+
const isEnabled = useMemo(
41+
() => Boolean(enabled && activeWorkflowId && !isReadOnly && !isStreaming && !isBaselineView),
42+
[enabled, activeWorkflowId, isReadOnly, isStreaming, isBaselineView]
43+
)
44+
const isReplaceEnabled = useMemo(
45+
() => Boolean(enabled && activeWorkflowId && !isReadOnly && !isBaselineView),
46+
[enabled, activeWorkflowId, isReadOnly, isBaselineView]
47+
)
48+
49+
const lastCommittedValueRef = useRef<string>(value ?? '')
50+
const pendingBeforeRef = useRef<string | null>(null)
51+
const pendingAfterRef = useRef<string | null>(null)
52+
const timeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null)
53+
const isApplyingRef = useRef(false)
54+
55+
const clearTimer = useCallback(() => {
56+
if (timeoutRef.current) {
57+
clearTimeout(timeoutRef.current)
58+
timeoutRef.current = null
59+
}
60+
}, [])
61+
62+
const resetPending = useCallback(() => {
63+
pendingBeforeRef.current = null
64+
pendingAfterRef.current = null
65+
}, [])
66+
67+
const commitPending = useCallback(() => {
68+
if (!isEnabled || !activeWorkflowId) {
69+
clearTimer()
70+
resetPending()
71+
return
72+
}
73+
74+
const before = pendingBeforeRef.current
75+
const after = pendingAfterRef.current
76+
if (before === null || after === null) return
77+
78+
if (before === after) {
79+
lastCommittedValueRef.current = after
80+
clearTimer()
81+
resetPending()
82+
return
83+
}
84+
85+
useCodeUndoRedoStore.getState().push({
86+
id: crypto.randomUUID(),
87+
createdAt: Date.now(),
88+
workflowId: activeWorkflowId,
89+
blockId,
90+
subBlockId,
91+
before,
92+
after,
93+
})
94+
95+
lastCommittedValueRef.current = after
96+
clearTimer()
97+
resetPending()
98+
}, [activeWorkflowId, blockId, clearTimer, isEnabled, resetPending, subBlockId])
99+
100+
const recordChange = useCallback(
101+
(nextValue: string) => {
102+
if (!isEnabled || isApplyingRef.current) return
103+
104+
if (pendingBeforeRef.current === null) {
105+
pendingBeforeRef.current = lastCommittedValueRef.current ?? ''
106+
}
107+
108+
pendingAfterRef.current = nextValue
109+
clearTimer()
110+
timeoutRef.current = setTimeout(commitPending, debounceMs)
111+
},
112+
[clearTimer, commitPending, debounceMs, isEnabled]
113+
)
114+
115+
const recordReplace = useCallback(
116+
(nextValue: string) => {
117+
if (!isReplaceEnabled || isApplyingRef.current || !activeWorkflowId) return
118+
119+
if (pendingBeforeRef.current !== null) {
120+
commitPending()
121+
}
122+
123+
const before = lastCommittedValueRef.current ?? ''
124+
if (before === nextValue) {
125+
lastCommittedValueRef.current = nextValue
126+
resetPending()
127+
return
128+
}
129+
130+
useCodeUndoRedoStore.getState().push({
131+
id: crypto.randomUUID(),
132+
createdAt: Date.now(),
133+
workflowId: activeWorkflowId,
134+
blockId,
135+
subBlockId,
136+
before,
137+
after: nextValue,
138+
})
139+
140+
lastCommittedValueRef.current = nextValue
141+
clearTimer()
142+
resetPending()
143+
},
144+
[
145+
activeWorkflowId,
146+
blockId,
147+
clearTimer,
148+
commitPending,
149+
isReplaceEnabled,
150+
resetPending,
151+
subBlockId,
152+
]
153+
)
154+
155+
const flushPending = useCallback(() => {
156+
if (pendingBeforeRef.current === null) return
157+
clearTimer()
158+
commitPending()
159+
}, [clearTimer, commitPending])
160+
161+
const startSession = useCallback(
162+
(currentValue: string) => {
163+
clearTimer()
164+
resetPending()
165+
lastCommittedValueRef.current = currentValue ?? ''
166+
},
167+
[clearTimer, resetPending]
168+
)
169+
170+
const applyValue = useCallback(
171+
(nextValue: string) => {
172+
if (!isEnabled) return
173+
isApplyingRef.current = true
174+
try {
175+
collaborativeSetSubblockValue(blockId, subBlockId, nextValue)
176+
} finally {
177+
isApplyingRef.current = false
178+
}
179+
lastCommittedValueRef.current = nextValue
180+
clearTimer()
181+
resetPending()
182+
},
183+
[blockId, clearTimer, collaborativeSetSubblockValue, isEnabled, resetPending, subBlockId]
184+
)
185+
186+
const undo = useCallback(() => {
187+
if (!activeWorkflowId || !isEnabled) return
188+
if (pendingBeforeRef.current !== null) {
189+
flushPending()
190+
}
191+
const entry = useCodeUndoRedoStore.getState().undo(activeWorkflowId, blockId, subBlockId)
192+
if (!entry) return
193+
logger.debug('Undo code edit', { blockId, subBlockId })
194+
applyValue(entry.before)
195+
}, [activeWorkflowId, applyValue, blockId, flushPending, isEnabled, subBlockId])
196+
197+
const redo = useCallback(() => {
198+
if (!activeWorkflowId || !isEnabled) return
199+
if (pendingBeforeRef.current !== null) {
200+
flushPending()
201+
}
202+
const entry = useCodeUndoRedoStore.getState().redo(activeWorkflowId, blockId, subBlockId)
203+
if (!entry) return
204+
logger.debug('Redo code edit', { blockId, subBlockId })
205+
applyValue(entry.after)
206+
}, [activeWorkflowId, applyValue, blockId, flushPending, isEnabled, subBlockId])
207+
208+
useEffect(() => {
209+
if (isApplyingRef.current || isStreaming) return
210+
211+
const nextValue = value ?? ''
212+
213+
if (pendingBeforeRef.current !== null) {
214+
if (pendingAfterRef.current !== nextValue) {
215+
clearTimer()
216+
resetPending()
217+
lastCommittedValueRef.current = nextValue
218+
}
219+
return
220+
}
221+
222+
lastCommittedValueRef.current = nextValue
223+
}, [clearTimer, isStreaming, resetPending, value])
224+
225+
useEffect(() => {
226+
return () => {
227+
flushPending()
228+
}
229+
}, [flushPending])
230+
231+
return {
232+
recordChange,
233+
recordReplace,
234+
flushPending,
235+
startSession,
236+
undo,
237+
redo,
238+
}
239+
}

0 commit comments

Comments
 (0)