From dbfaa07e030c1c5cb9124b0cd6ecfc0ec25a384d Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Thu, 16 May 2024 21:22:46 +1000
Subject: [PATCH] feat(ui): add checks for undo/redo actions

---
 .../features/nodes/components/flow/Flow.tsx   | 24 +++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx b/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx
index 177abd57f5..b47c5c38f7 100644
--- a/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx
+++ b/invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx
@@ -66,6 +66,8 @@ export const Flow = memo(() => {
   const nodes = useAppSelector((s) => s.nodes.present.nodes);
   const edges = useAppSelector((s) => s.nodes.present.edges);
   const viewport = useStore($viewport);
+  const mayUndo = useAppSelector((s) => s.nodes.past.length > 0);
+  const mayRedo = useAppSelector((s) => s.nodes.future.length > 0);
   const shouldSnapToGrid = useAppSelector((s) => s.workflowSettings.shouldSnapToGrid);
   const selectionMode = useAppSelector((s) => s.workflowSettings.selectionMode);
   const { onConnectStart, onConnect, onConnectEnd } = useConnection();
@@ -192,13 +194,13 @@ export const Flow = memo(() => {
   const { copySelection, pasteSelection } = useCopyPaste();
 
   useHotkeys(['Ctrl+c', 'Meta+c'], (e) => {
-    e.preventDefault();
-    copySelection();
+      e.preventDefault();
+      copySelection();
   });
 
   useHotkeys(['Ctrl+a', 'Meta+a'], (e) => {
-    e.preventDefault();
-    dispatch(selectedAll());
+      e.preventDefault();
+      dispatch(selectedAll());
   });
 
   useHotkeys(['Ctrl+v', 'Meta+v'], (e) => {
@@ -222,6 +224,20 @@ export const Flow = memo(() => {
     [dispatch]
   );
 
+  const onUndoHotkey = useCallback(() => {
+    if (mayUndo) {
+      dispatch(undo());
+    }
+  }, [dispatch, mayUndo]);
+  useHotkeys(['meta+z', 'ctrl+z'], onUndoHotkey);
+
+  const onRedoHotkey = useCallback(() => {
+    if (mayRedo) {
+      dispatch(redo());
+    }
+  }, [dispatch, mayRedo]);
+  useHotkeys(['meta+shift+z', 'ctrl+shift+z'], onRedoHotkey);
+
   return (
     <ReactFlow
       id="workflow-editor"