diff --git a/src/components/WorkflowCanvas.tsx b/src/components/WorkflowCanvas.tsx index 2ae6d4fd..ca3693fa 100644 --- a/src/components/WorkflowCanvas.tsx +++ b/src/components/WorkflowCanvas.tsx @@ -281,6 +281,10 @@ export function WorkflowCanvas() { // Check if a node was dropped into a group and add it to that group + const handleNodeDragStart = useCallback(() => { + useWorkflowStore.temporal.getState().pause(); + }, []); + const handleNodeDragStop = useCallback( (_event: React.MouseEvent, node: Node) => { // Skip if it's a group node @@ -313,6 +317,11 @@ export function WorkflowCanvas() { if (targetGroupId !== currentGroupId) { setNodeGroupId(node.id, targetGroupId); } + + // Resume undo tracking after drag and capture final position + const temporal = useWorkflowStore.temporal.getState(); + temporal.resume(); + useWorkflowStore.getState()._nudgeForSnapshot(); }, [groups, nodes, setNodeGroupId] ); @@ -1081,6 +1090,39 @@ export function WorkflowCanvas() { return; } + // Handle undo (Ctrl/Cmd + Z, but NOT Ctrl/Cmd + Shift + Z) + if ((event.ctrlKey || event.metaKey) && (event.key === "z" || event.key === "Z") && !event.shiftKey) { + event.preventDefault(); + const { undo, pastStates } = useWorkflowStore.temporal.getState(); + if (pastStates.length > 0) { + undo(); + useWorkflowStore.setState({ hasUnsavedChanges: true }); + } + return; + } + + // Handle redo (Ctrl/Cmd + Shift + Z) + if ((event.ctrlKey || event.metaKey) && (event.key === "z" || event.key === "Z") && event.shiftKey) { + event.preventDefault(); + const { redo, futureStates } = useWorkflowStore.temporal.getState(); + if (futureStates.length > 0) { + redo(); + useWorkflowStore.setState({ hasUnsavedChanges: true }); + } + return; + } + + // Handle redo alt (Ctrl + Y — Windows convention) + if ((event.ctrlKey || event.metaKey) && event.key === "y") { + event.preventDefault(); + const { redo, futureStates } = useWorkflowStore.temporal.getState(); + if (futureStates.length > 0) { + redo(); + useWorkflowStore.setState({ hasUnsavedChanges: true }); + } + return; + } + // Handle copy (Ctrl/Cmd + C) if ((event.ctrlKey || event.metaKey) && event.key === "c") { event.preventDefault(); @@ -1643,6 +1685,7 @@ export function WorkflowCanvas() { onEdgesChange={onEdgesChange} onConnect={handleConnect} onConnectEnd={handleConnectEnd} + onNodeDragStart={handleNodeDragStart} onNodeDragStop={handleNodeDragStop} onSelectionChange={handleSelectionChange} nodeTypes={nodeTypes} diff --git a/src/store/workflowStore.ts b/src/store/workflowStore.ts index 8df3f97e..705479c9 100644 --- a/src/store/workflowStore.ts +++ b/src/store/workflowStore.ts @@ -1,5 +1,6 @@ -import { create } from "zustand"; +import { create, StateCreator } from "zustand"; import { useShallow } from "zustand/shallow"; +import { temporal, TemporalState } from "zundo"; import { Connection, EdgeChange, @@ -43,6 +44,7 @@ import { getCanvasNavigationSettings, saveCanvasNavigationSettings, } from "./utils/localStorage"; +import { partializeForUndo, undoStateEquality } from "./undoUtils"; import { createDefaultNodeData, defaultNodeDimensions, @@ -333,8 +335,14 @@ interface WorkflowStore { // Canvas navigation settings actions updateCanvasNavigationSettings: (settings: CanvasNavigationSettings) => void; + + // Undo/Redo helper actions + _nudgeForSnapshot: () => void; } +// Extend with TemporalState for undo/redo functionality +export type WorkflowStoreWithTemporal = WorkflowStore & TemporalState; + let nodeIdCounter = 0; let groupIdCounter = 0; let autoSaveIntervalId: ReturnType | null = null; @@ -369,7 +377,8 @@ async function waitForPendingImageSyncs(timeout: number = 60000): Promise export { generateWorkflowId, saveGenerateImageDefaults, saveNanoBananaDefaults } from "./utils/localStorage"; export { GROUP_COLORS } from "./utils/nodeDefaults"; -export const useWorkflowStore = create((set, get) => ({ +// Store implementation with temporal middleware for undo/redo +const workflowStoreImpl: StateCreator = (set, get) => ({ nodes: [], edges: [], edgeStyle: "curved" as EdgeStyle, @@ -2111,7 +2120,23 @@ export const useWorkflowStore = create((set, get) => ({ set({ canvasNavigationSettings: settings }); saveCanvasNavigationSettings(settings); }, -})); + + // Undo/Redo helper actions + _nudgeForSnapshot: () => { + set((state) => ({ nodes: state.nodes })); + }, +}); + +export const useWorkflowStore = create()( + temporal( + workflowStoreImpl, + { + partialize: partializeForUndo as any, + limit: 20, + equality: undoStateEquality as any, + } + ) +); /** * Stable hook for provider API keys.