diff --git a/app/w/[id]/workflow.tsx b/app/w/[id]/workflow.tsx index 4c6c9535ca..a80ca6e6e6 100644 --- a/app/w/[id]/workflow.tsx +++ b/app/w/[id]/workflow.tsx @@ -97,8 +97,18 @@ const edgeTypes: EdgeTypes = { custom: CustomEdge } function WorkflowCanvas() { const [selectedBlockId, setSelectedBlockId] = useState(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(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 (
diff --git a/stores/workflow/history-middleware.ts b/stores/workflow/history-middleware.ts new file mode 100644 index 0000000000..414dfb479e --- /dev/null +++ b/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 +): StateCreator => { + 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 + | (( + state: WorkflowStoreWithHistory + ) => Partial), + 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: [], + }, + }) +} \ No newline at end of file diff --git a/stores/workflow/history-types.ts b/stores/workflow/history-types.ts new file mode 100644 index 0000000000..5b951bbdc5 --- /dev/null +++ b/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 +} \ No newline at end of file diff --git a/stores/workflow/state-logger.ts b/stores/workflow/state-logger.ts index 7652c36759..5fe8a5a753 100644 --- a/stores/workflow/state-logger.ts +++ b/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, + } }) }) } \ No newline at end of file diff --git a/stores/workflow/workflow-store.ts b/stores/workflow/workflow-store.ts index 20f730915b..f89fb1682d 100644 --- a/stores/workflow/workflow-store.ts +++ b/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()( +export const useWorkflowStore = create()( 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()( } }) - set((state) => ({ + const newState = { blocks: { - ...state.blocks, + ...get().blocks, [id]: { id, type, @@ -43,7 +72,11 @@ export const useWorkflowStore = create()( : 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()( 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' } ) ) \ No newline at end of file