Added history tracking with middleware and works with workflow

app/w/[id]/workflow.tsx

@@ -97,8 +97,18 @@ const edgeTypes: EdgeTypes = { custom: CustomEdge }
 function WorkflowCanvas() {
   const [selectedBlockId, setSelectedBlockId] = useState<string | null>(null)
 
-  const { blocks, edges, addBlock, updateBlockPosition, addEdge, removeEdge } =
-    useWorkflowStore()
+  const {
+    blocks,
+    edges,
+    addBlock,
+    updateBlockPosition,
+    addEdge,
+    removeEdge,
+    canUndo,
+    canRedo,
+    undo,
+    redo,
+  } = useWorkflowStore()
   const [isExecuting, setIsExecuting] = useState(false)
   const [executionResult, setExecutionResult] = useState<any>(null)
   const { addNotification } = useNotificationStore()
@@ -372,6 +382,23 @@ function WorkflowCanvas() {
     }
   }
 
+  // Add keyboard shortcut handler
+  useEffect(() => {
+    const handleKeyDown = (event: KeyboardEvent) => {
+      if ((event.metaKey || event.ctrlKey) && event.key === 'z') {
+        if (event.shiftKey) {
+          if (canRedo()) redo()
+        } else {
+          if (canUndo()) undo()
+        }
+        event.preventDefault()
+      }
+    }
+
+    window.addEventListener('keydown', handleKeyDown)
+    return () => window.removeEventListener('keydown', handleKeyDown)
+  }, [canUndo, canRedo, undo, redo])
+
   return (
     <div className="relative w-full h-[calc(100vh-56px)]">
       <NotificationList />

stores/workflow/history-middleware.ts

@@ -0,0 +1,104 @@
+import { StateCreator } from 'zustand'
+import { WorkflowState, WorkflowStore } from './types'
+import { HistoryEntry, WorkflowHistory, HistoryActions } from './history-types'
+
+const MAX_HISTORY_LENGTH = 20
+
+export interface WorkflowStoreWithHistory extends WorkflowStore, HistoryActions {
+  history: WorkflowHistory
+}
+
+export const withHistory = (
+  config: StateCreator<WorkflowStoreWithHistory>
+): StateCreator<WorkflowStoreWithHistory> => {
+  return (set, get, api) => {
+    const initialState = config(set, get, api)
+    const initialHistoryEntry: HistoryEntry = {
+      state: {
+        blocks: initialState.blocks,
+        edges: initialState.edges,
+      },
+      timestamp: Date.now(),
+    }
+
+    return {
+      ...initialState,
+      history: {
+        past: [],
+        present: initialHistoryEntry,
+        future: [],
+      },
+
+      canUndo: () => get().history.past.length > 0,
+      canRedo: () => get().history.future.length > 0,
+
+      undo: () => {
+        const { history, ...state } = get()
+        if (history.past.length === 0) return
+
+        const previous = history.past[history.past.length - 1]
+        const newPast = history.past.slice(0, history.past.length - 1)
+
+        set({
+          ...state,
+          ...previous.state,
+          history: {
+            past: newPast,
+            present: previous,
+            future: [history.present, ...history.future],
+          },
+        })
+      },
+
+      redo: () => {
+        const { history, ...state } = get()
+        if (history.future.length === 0) return
+
+        const next = history.future[0]
+        const newFuture = history.future.slice(1)
+
+        set({
+          ...state,
+          ...next.state,
+          history: {
+            past: [...history.past, history.present],
+            present: next,
+            future: newFuture,
+          },
+        })
+      },
+    }
+  }
+}
+
+export const createHistoryEntry = (state: WorkflowState): HistoryEntry => ({
+  state: {
+    blocks: { ...state.blocks },
+    edges: [...state.edges],
+  },
+  timestamp: Date.now(),
+})
+
+export const pushHistory = (
+  set: (
+    partial:
+      | Partial<WorkflowStoreWithHistory>
+      | ((
+          state: WorkflowStoreWithHistory
+        ) => Partial<WorkflowStoreWithHistory>),
+    replace?: boolean
+  ) => void,
+  get: () => WorkflowStoreWithHistory,
+  newState: WorkflowState
+) => {
+  const { history } = get()
+  const newEntry = createHistoryEntry(newState)
+
+  set({
+    history: {
+      past: [...history.past, history.present].slice(-MAX_HISTORY_LENGTH),
+      present: newEntry,
+      future: [],
+    },
+  })
+} 

stores/workflow/history-types.ts

@@ -0,0 +1,19 @@
+import { WorkflowState } from './types'
+
+export interface HistoryEntry {
+  state: WorkflowState
+  timestamp: number
+}
+
+export interface WorkflowHistory {
+  past: HistoryEntry[]
+  present: HistoryEntry
+  future: HistoryEntry[]
+}
+
+export interface HistoryActions {
+  undo: () => void
+  redo: () => void
+  canUndo: () => boolean
+  canRedo: () => boolean
+} 

stores/workflow/state-logger.ts

@@ -3,9 +3,15 @@ import { useWorkflowStore } from './workflow-store'
 export function initializeStateLogger() {
   useWorkflowStore.subscribe((state) => {
     console.log('Workflow State Updated:', {
-      blocks: state.blocks,
-      edges: state.edges,
-      selectedBlockId: state.selectedBlockId,
+      current: {
+        blocks: state.blocks,
+        edges: state.edges,
+      },
+      history: {
+        past: state.history.past,
+        present: state.history.present,
+        future: state.history.future,
+      }
     })
   })
 } 

stores/workflow/workflow-store.ts

@@ -1,18 +1,47 @@
 import { create } from 'zustand'
 import { devtools } from 'zustand/middleware'
 import { Edge } from 'reactflow'
-import { Position, WorkflowStore } from './types'
+import { Position, SubBlockState, WorkflowStore } from './types'
 import { getBlock } from '@/blocks'
+import { withHistory, WorkflowStoreWithHistory, pushHistory } from './history-middleware'
 
 const initialState = {
   blocks: {},
   edges: [],
+  history: {
+    past: [],
+    present: {
+      state: { blocks: {}, edges: [] },
+      timestamp: Date.now(),
+    },
+    future: [],
+  },
 }
 
-export const useWorkflowStore = create<WorkflowStore>()(
+export const useWorkflowStore = create<WorkflowStoreWithHistory>()(
   devtools(
-    (set, get) => ({
+    withHistory((set, get) => ({
       ...initialState,
+      undo: () => {},
+      redo: () => {},
+      canUndo: () => false,
+      canRedo: () => false,
+
+      updateSubBlock: (blockId: string, subBlockId: string, subBlock: SubBlockState) => {
+        set((state) => ({
+          blocks: {
+            ...state.blocks,
+            [blockId]: {
+              ...state.blocks[blockId],
+              subBlocks: {
+                ...state.blocks[blockId].subBlocks,
+                [subBlockId]: subBlock,
+              },
+            },
+          },
+          edges: [...state.edges],
+        }))
+      },
 
       addBlock: (id: string, type: string, name: string, position: Position) => {
         const blockConfig = getBlock(type)
@@ -28,9 +57,9 @@ export const useWorkflowStore = create<WorkflowStore>()(
           }
         })
 
-        set((state) => ({
+        const newState = {
           blocks: {
-            ...state.blocks,
+            ...get().blocks,
             [id]: {
               id,
               type,
@@ -43,7 +72,11 @@ export const useWorkflowStore = create<WorkflowStore>()(
                   : blockConfig.workflow.outputType.default,
             },
           },
-        }))
+          edges: [...get().edges],
+        }
+
+        set(newState)
+        pushHistory(set, get, newState)
       },
 
       updateBlockPosition: (id: string, position: Position) => {
@@ -55,63 +88,56 @@ export const useWorkflowStore = create<WorkflowStore>()(
               position,
             },
           },
-        }))
-      },
-
-      updateSubBlock: (blockId: string, subBlockId: string, value: any) => {
-        set((state) => ({
-          blocks: {
-            ...state.blocks,
-            [blockId]: {
-              ...state.blocks[blockId],
-              subBlocks: {
-                ...state.blocks[blockId].subBlocks,
-                [subBlockId]: {
-                  ...state.blocks[blockId].subBlocks[subBlockId],
-                  value,
-                },
-              },
-            },
-          },
+          edges: [...state.edges],
         }))
       },
 
       removeBlock: (id: string) => {
-        set((state) => {
-          const { [id]: _, ...remainingBlocks } = state.blocks
-          const remainingEdges = state.edges.filter(
+        const newState = {
+          blocks: { ...get().blocks },
+          edges: [...get().edges].filter(
             (edge) => edge.source !== id && edge.target !== id
-          )
-          return {
-            blocks: remainingBlocks,
-            edges: remainingEdges,
-          }
-        })
+          ),
+        }
+        delete newState.blocks[id]
+        
+        set(newState)
+        pushHistory(set, get, newState)
       },
 
       addEdge: (edge: Edge) => {
-        set((state) => ({
+        const newState = {
+          blocks: { ...get().blocks },
           edges: [
-            ...state.edges,
+            ...get().edges,
             {
               id: edge.id || crypto.randomUUID(),
               source: edge.source,
               target: edge.target,
             },
           ],
-        }))
+        }
+        
+        set(newState)
+        pushHistory(set, get, newState)
       },
 
       removeEdge: (edgeId: string) => {
-        set((state) => ({
-          edges: state.edges.filter((edge) => edge.id !== edgeId),
-        }))
+        const newState = {
+          blocks: { ...get().blocks },
+          edges: get().edges.filter((edge) => edge.id !== edgeId),
+        }
+        
+        set(newState)
+        pushHistory(set, get, newState)
       },
 
       clear: () => {
-        set(initialState)
+        const newState = initialState
+        set(newState)
+        pushHistory(set, get, newState)
       },
-    }),
+    })),
     { name: 'workflow-store' }
   )
 )