Add undo & redo functionality

packages/toolpad-app/src/toolpad/AppEditor/PageEditor/ComponentEditor.tsx

@@ -149,7 +149,13 @@ export default function ComponentEditor({ className }: ComponentEditorProps) {
 
   const { selection } = editor;
 
-  const selectedNode = selection ? appDom.getNode(dom, selection) : null;
+  let selectedNode;
+
+  try {
+    selectedNode = selection ? appDom.getNode(dom, selection) : null;
+  } catch (error) {
+    console.warn('Selected node does not exists');
+  }
 
   return (
     <ComponentEditorRoot className={className} data-testid="component-editor">

packages/toolpad-app/src/toolpad/AppEditor/PageEditor/EditorCanvasHost.tsx

@@ -12,6 +12,8 @@ import { ToolpadBridge } from '../../../canvas';
 import useEvent from '../../../utils/useEvent';
 import { LogEntry } from '../../../components/Console';
 import { Maybe } from '../../../utils/types';
+import { useDomApi } from '../../DomLoader';
+import { hasFieldFocus } from '../../../utils/fields';
 
 type IframeContentWindow = Window & typeof globalThis;
 
@@ -79,6 +81,7 @@ export default React.forwardRef<EditorCanvasHostHandle, EditorCanvasHostProps>(
     forwardedRef,
   ) {
     const frameRef = React.useRef<HTMLIFrameElement>(null);
+    const domApi = useDomApi();
 
     const [bridge, setBridge] = React.useState<ToolpadBridge | null>(null);
 
@@ -145,10 +148,44 @@ export default React.forwardRef<EditorCanvasHostHandle, EditorCanvasHostProps>(
 
     const handleRuntimeEvent = useEvent(onRuntimeEvent);
 
+    const iframeKeyDownHandler = React.useCallback(
+      (iframeDocument: Document) => {
+        return (event: KeyboardEvent) => {
+          if (hasFieldFocus(iframeDocument)) {
+            return;
+          }
+
+          const { code, metaKey, shiftKey } = event;
+          const undoShortcut = code === 'KeyZ' && metaKey;
+          const redoShortcut = undoShortcut && shiftKey;
+
+          if (redoShortcut) {
+            domApi.redo();
+          } else if (undoShortcut) {
+            domApi.undo();
+          }
+        };
+      },
+      [domApi],
+    );
+
     const handleFrameLoad = React.useCallback(() => {
       invariant(frameRef.current, 'Iframe ref not attached');
-      setContentWindow(frameRef.current.contentWindow);
-    }, []);
+
+      const iframeWindow = frameRef.current.contentWindow;
+      setContentWindow(iframeWindow);
+
+      if (!iframeWindow) {
+        return;
+      }
+
+      const keyDownHandler = iframeKeyDownHandler(iframeWindow.document);
+
+      iframeWindow?.addEventListener('keydown', keyDownHandler);
+      iframeWindow?.addEventListener('unload', () => {
+        iframeWindow?.removeEventListener('keydown', keyDownHandler);
+      });
+    }, [iframeKeyDownHandler]);
 
     React.useEffect(() => {
       if (!contentWindow) {

packages/toolpad-app/src/toolpad/AppEditor/PageEditor/RenderPanel/RenderOverlay.tsx

@@ -163,7 +163,12 @@ export default function RenderOverlay({ canvasHostRef }: RenderOverlayProps) {
     return [pageNode, ...appDom.getDescendants(dom, pageNode)];
   }, [dom, pageNode]);
 
-  const selectedNode = selection && appDom.getNode(dom, selection);
+  let selectedNode: appDom.AppDom | null;
+  try {
+    selectedNode = selection && appDom.getNode(dom, selection);
+  } catch (error) {
+    console.warn('Selected node does not exists');
+  }
 
   const overlayRef = React.useRef<HTMLDivElement | null>(null);
 

packages/toolpad-app/src/toolpad/AppEditor/PageEditor/index.tsx

@@ -6,13 +6,15 @@ import SplitPane from '../../../components/SplitPane';
 import RenderPanel from './RenderPanel';
 import ComponentPanel from './ComponentPanel';
 import { PageEditorProvider } from './PageEditorProvider';
-import { useDom } from '../../DomLoader';
+import { useDom, useDomApi } from '../../DomLoader';
 import * as appDom from '../../../appDom';
 import ComponentCatalog from './ComponentCatalog';
 import NotFoundEditor from '../NotFoundEditor';
 import usePageTitle from '../../../utils/usePageTitle';
 import useLocalStorageState from '../../../utils/useLocalStorageState';
 import useDebouncedHandler from '../../../utils/useDebouncedHandler';
+import useShortcut from '../../../utils/useShortcut';
+import { hasFieldFocus } from '../../../utils/fields';
 
 const classes = {
   renderPanel: 'Toolpad_RenderPanel',
@@ -70,9 +72,23 @@ interface PageEditorProps {
 
 export default function PageEditor({ appId }: PageEditorProps) {
   const dom = useDom();
+  const domApi = useDomApi();
   const { nodeId } = useParams();
   const pageNode = appDom.getMaybeNode(dom, nodeId as NodeId, 'page');
 
+  useShortcut({ code: 'KeyZ', metaKey: true, preventDefault: false }, () => {
+    if (hasFieldFocus()) {
+      return;
+    }
+    domApi.undo();
+  });
+  useShortcut({ code: 'KeyZ', metaKey: true, shiftKey: true, preventDefault: false }, () => {
+    if (hasFieldFocus()) {
+      return;
+    }
+    domApi.redo();
+  });
+
   return pageNode ? (
     <PageEditorContent appId={appId} node={pageNode} />
   ) : (

packages/toolpad-app/src/toolpad/DomLoader.tsx

@@ -1,6 +1,7 @@
 import * as React from 'react';
 import { NodeId, BindableAttrValue, BindableAttrValues } from '@mui/toolpad-core';
 import invariant from 'invariant';
+import { throttle } from 'lodash-es';
 import * as appDom from '../appDom';
 import { update } from '../utils/immutability';
 import client from '../api';
@@ -12,6 +13,12 @@ import insecureHash from '../utils/insecureHash';
 import { NodeHashes } from '../types';
 
 export type DomAction =
+  | {
+      type: 'DOM_UNDO';
+    }
+  | {
+      type: 'DOM_REDO';
+    }
   | {
       type: 'DOM_SAVING';
     }
@@ -68,8 +75,43 @@ export type DomAction =
       node: appDom.AppDomNode;
     };
 
+const undoStack: appDom.AppDom[] = [];
+const redoStack: appDom.AppDom[] = [];
+
+const updateUndoStack = throttle((dom: appDom.AppDom) => {
+  undoStack.push(dom);
+}, 1000);
+
+const SKIP_UNDO_ACTIONS = ['DOM_UNDO', 'DOM_REDO', 'DOM_SAVED', 'DOM_SAVING', 'DOM_SAVING_ERROR'];
+
 export function domReducer(dom: appDom.AppDom, action: DomAction): appDom.AppDom {
+  if (!SKIP_UNDO_ACTIONS.includes(action.type)) {
+    updateUndoStack(dom);
+  }
+
   switch (action.type) {
+    case 'DOM_UNDO': {
+      const undoDom = undoStack.pop();
+
+      if (!undoDom) {
+        return dom;
+      }
+
+      redoStack.push(dom);
+
+      return undoDom;
+    }
+    case 'DOM_REDO': {
+      const redoDom = redoStack.pop();
+
+      if (!redoDom) {
+        return dom;
+      }
+
+      undoStack.push(dom);
+
+      return redoDom;
+    }
     case 'DOM_SET_NODE_NAME': {
       // TODO: Also update all bindings on the page that use this name
       const node = appDom.getNode(dom, action.nodeId);
@@ -158,6 +200,12 @@ export function domLoaderReducer(state: DomLoader, action: DomAction): DomLoader
 
 function createDomApi(dispatch: React.Dispatch<DomAction>) {
   return {
+    undo() {
+      dispatch({ type: 'DOM_UNDO' });
+    },
+    redo() {
+      dispatch({ type: 'DOM_REDO' });
+    },
     setNodeName(nodeId: NodeId, name: string) {
       dispatch({ type: 'DOM_SET_NODE_NAME', nodeId, name });
     },

packages/toolpad-app/src/utils/fields.ts

@@ -0,0 +1,14 @@
+export const hasFieldFocus = (documentTarget = document) => {
+  const activeElement = documentTarget.activeElement as HTMLElement;
+
+  if (!activeElement) {
+    return false;
+  }
+  const { nodeName, contentEditable } = activeElement;
+
+  const focusedInput = nodeName === 'INPUT';
+  const focusedTextarea = nodeName === 'TEXTAREA';
+  const focusedContentEditable = contentEditable === 'true';
+
+  return focusedInput || focusedTextarea || focusedContentEditable;
+};

packages/toolpad-app/src/utils/useShortcut.ts

@@ -3,11 +3,13 @@ import * as React from 'react';
 export interface ShortCut {
   code: string;
   metaKey?: boolean;
+  shiftKey?: boolean;
   disabled?: boolean;
+  preventDefault?: boolean;
 }
 
 export default function useShortcut(
-  { code, metaKey = false, disabled = false }: ShortCut,
+  { code, metaKey = false, disabled = false, shiftKey = false, preventDefault = true }: ShortCut,
   handler: () => void,
 ) {
   React.useEffect(() => {
@@ -16,13 +18,15 @@ export default function useShortcut(
     }
 
     const handleKeydown = (event: KeyboardEvent) => {
-      if (event.code === code && event.metaKey === metaKey) {
+      if (event.code === code && event.metaKey === metaKey && event.shiftKey === shiftKey) {
         handler();
-        event.preventDefault();
+        if (preventDefault) {
+          event.preventDefault();
+        }
       }
     };
 
     document.addEventListener('keydown', handleKeydown);
     return () => document.removeEventListener('keydown', handleKeydown);
-  }, [code, metaKey, handler, disabled]);
+  }, [code, metaKey, shiftKey, handler, disabled, preventDefault]);
 }