Merge pull request #501 from refly-ai/feat/canvas-undo

feat(canvas): undo and redo management for canvas

packages/ai-workspace-common/src/components/canvas/index.tsx

@@ -54,6 +54,7 @@ import { useUserStore } from '@refly-packages/ai-workspace-common/stores/user';
 import { useUpdateSettings } from '@refly-packages/ai-workspace-common/queries';
 import { IconCreateDocument } from '@refly-packages/ai-workspace-common/components/common/icon';
 import { useUploadImage } from '@refly-packages/ai-workspace-common/hooks/use-upload-image';
+import { useCanvasSync } from '@refly-packages/ai-workspace-common/hooks/canvas/use-canvas-sync';
 
 const selectionStyles = `
   .react-flow__selection {
@@ -501,6 +502,7 @@ const Flow = memo(({ canvasId }: { canvasId: string }) => {
   );
 
   const { handleUploadImage } = useUploadImage();
+  const { undoManager } = useCanvasSync();
 
   // Add drag and drop handlers
   const handleDragOver = useCallback((event: React.DragEvent) => {
@@ -521,6 +523,41 @@ const Flow = memo(({ canvasId }: { canvasId: string }) => {
     [addNode, reactFlowInstance],
   );
 
+  const handleKeyDown = (e: KeyboardEvent) => {
+    const target = e.target as HTMLElement;
+
+    // Ignore input, textarea and contentEditable elements
+    if (
+      target.tagName === 'INPUT' ||
+      target.tagName === 'TEXTAREA' ||
+      target.contentEditable === 'true'
+    ) {
+      return;
+    }
+
+    // Check for mod key (Command on Mac, Ctrl on Windows/Linux)
+    const isModKey = e.metaKey || e.ctrlKey;
+
+    if (isModKey && e.key.toLowerCase() === 'z') {
+      e.preventDefault();
+      if (e.shiftKey) {
+        // Mod+Shift+Z for Redo
+        undoManager.redo();
+      } else {
+        // Mod+Z for Undo
+        undoManager.undo();
+      }
+    }
+  };
+
+  // Set up keyboard shortcuts for undo/redo
+  useEffect(() => {
+    if (!undoManager) return;
+
+    window.addEventListener('keydown', handleKeyDown);
+    return () => window.removeEventListener('keydown', handleKeyDown);
+  }, [undoManager]);
+
   return (
     <Spin
       className="w-full h-full"

packages/ai-workspace-common/src/components/canvas/layout-control/index.tsx

@@ -1,10 +1,11 @@
 import React, { useState, useEffect, useCallback, memo, useMemo, useRef } from 'react';
 import { Button, Dropdown, Space, Divider, Tooltip } from 'antd';
+import { UndoManager } from 'yjs';
 import { LuCompass, LuLayoutDashboard, LuLightbulb, LuShipWheel } from 'react-icons/lu';
 import { RiFullscreenFill } from 'react-icons/ri';
 import { FiHelpCircle } from 'react-icons/fi';
 import { useTranslation } from 'react-i18next';
-import { LuZoomIn, LuZoomOut } from 'react-icons/lu';
+import { LuUndo, LuRedo, LuZoomIn, LuZoomOut } from 'react-icons/lu';
 import {
   IconDocumentation,
   IconDown,
@@ -22,6 +23,7 @@ import { IconExpand, IconShrink } from '@refly-packages/ai-workspace-common/comp
 
 import './index.scss';
 import { useUserStoreShallow } from '@refly-packages/ai-workspace-common/stores/user';
+import { useCanvasSync } from '@refly-packages/ai-workspace-common/hooks/canvas/use-canvas-sync';
 
 interface LayoutControlProps {
   mode: 'mouse' | 'touchpad';
@@ -46,6 +48,7 @@ interface ActionButtonsProps {
   onLayout: (direction: 'TB' | 'LR') => void;
   onToggleSizeMode: () => void;
   nodeSizeMode: 'compact' | 'adaptive';
+  undoManager: UndoManager;
   t: TFunction;
 }
 
@@ -78,8 +81,26 @@ const TooltipButton = memo(({ tooltip, children, ...buttonProps }: TooltipButton
 
 // Update component definitions
 const ActionButtons = memo(
-  ({ onFitView, onLayout, onToggleSizeMode, nodeSizeMode, t }: ActionButtonsProps) => (
+  ({ onFitView, onLayout, onToggleSizeMode, nodeSizeMode, undoManager, t }: ActionButtonsProps) => (
     <>
+      <TooltipButton
+        tooltip={t('canvas.toolbar.tooltip.undo')}
+        onClick={() => undoManager?.undo()}
+        className={buttonClass}
+      >
+        <LuUndo className={iconClass} size={16} />
+      </TooltipButton>
+
+      <TooltipButton
+        tooltip={t('canvas.toolbar.tooltip.redo')}
+        onClick={() => undoManager?.redo()}
+        className={buttonClass}
+      >
+        <LuRedo className={iconClass} size={16} />
+      </TooltipButton>
+
+      <Divider type="vertical" className="h-full" />
+
       <TooltipButton
         tooltip={t('canvas.toolbar.tooltip.fitView')}
         onClick={onFitView}
@@ -272,6 +293,7 @@ export const LayoutControl: React.FC<LayoutControlProps> = memo(({ mode, changeM
     setNodeSizeMode: state.setNodeSizeMode,
   }));
   const { updateAllNodesSizeMode } = useNodeOperations();
+  const { undoManager } = useCanvasSync();
 
   // Add handler for size mode toggle
   const handleToggleSizeMode = useCallback(() => {
@@ -329,6 +351,7 @@ export const LayoutControl: React.FC<LayoutControlProps> = memo(({ mode, changeM
           onLayout={onLayout}
           onToggleSizeMode={handleToggleSizeMode}
           nodeSizeMode={nodeSizeMode}
+          undoManager={undoManager}
           t={t}
         />
 

packages/ai-workspace-common/src/hooks/canvas/use-canvas-sync.ts

@@ -3,11 +3,25 @@ import { useCanvasContext } from '../../context/canvas';
 import { useThrottledCallback } from 'use-debounce';
 import { Edge } from '@xyflow/react';
 import { CanvasNode } from '../../components/canvas/nodes';
+import { UndoManager } from 'yjs';
+import { omit } from '@refly/utils';
 
 export const useCanvasSync = () => {
   const { provider } = useCanvasContext();
   const ydoc = provider.document;
 
+  const undoManager = useMemo(() => {
+    if (!ydoc) return null;
+
+    // Create UndoManager tracking title, nodes and edges
+    return new UndoManager(
+      [ydoc.getText('title'), ydoc.getArray('nodes'), ydoc.getArray('edges')],
+      {
+        captureTimeout: 1000,
+      },
+    );
+  }, [ydoc]);
+
   const syncFunctions = useMemo(() => {
     const syncTitleToYDoc = (title: string) => {
       ydoc?.transact(() => {
@@ -29,7 +43,7 @@ export const useCanvasSync = () => {
       ydoc?.transact(() => {
         const yEdges = ydoc?.getArray('edges');
         yEdges?.delete(0, yEdges?.length ?? 0);
-        yEdges?.push(edges);
+        yEdges?.push(edges.map((edge) => omit(edge, ['style'])));
       });
     };
 
@@ -54,5 +68,6 @@ export const useCanvasSync = () => {
     ...syncFunctions,
     throttledSyncNodesToYDoc,
     throttledSyncEdgesToYDoc,
+    undoManager,
   };
 };

packages/i18n/src/en-US/ui.ts

@@ -672,6 +672,8 @@ const translations = {
       mouse: 'Mouse Mode',
       touchpad: 'Touchpad Mode',
       tooltip: {
+        undo: 'Undo',
+        redo: 'Redo',
         zoom: 'Zoom Percentage',
         zoomIn: 'Zoom In',
         zoomOut: 'Zoom Out',

packages/i18n/src/zh-Hans/ui.ts

@@ -665,6 +665,8 @@ const translations = {
       mouse: '鼠标模式',
       touchpad: '触控板模式',
       tooltip: {
+        undo: '撤销',
+        redo: '重做',
         zoom: '缩放百分比',
         zoomIn: '放大',
         zoomOut: '缩小',