Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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')

Expand Down Expand Up @@ -212,16 +214,19 @@ 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

const accessiblePrefixes = useAccessibleReferencePrefixes(blockId)
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 =
Expand Down Expand Up @@ -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: () => {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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)
Expand All @@ -539,7 +554,7 @@ export const Code = memo(function Code({
if (!isPreview && !readOnly) {
setCode(newValue)
emitTagSelection(newValue)
hasEditedSinceFocusRef.current = true
recordChange(newValue)
}
setShowEnvVars(false)

Expand Down Expand Up @@ -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) {
Expand All @@ -646,7 +661,7 @@ export const Code = memo(function Code({
}
}
},
[isAiStreaming, isPreview, disabled, readOnly, setStoreValue]
[isAiStreaming, isPreview, disabled, readOnly, recordChange, setStoreValue]
)

const handleKeyDown = useCallback(
Expand All @@ -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.
Expand Down Expand Up @@ -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 })}
/>
Expand Down
239 changes: 239 additions & 0 deletions apps/sim/hooks/use-code-undo-redo.ts
Original file line number Diff line number Diff line change
@@ -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,
}
}
Loading