feat(build): Add undo/redo functionality and integrate custom nodes

Introduce a `useUndoRedo` hook to manage state history with undo/redo, persistence using localStorage, and state validation. Updated the build page to display a flow with initial custom nodes and edges, replacing the previous flow editor implementation.

autogpt_platform/frontend/package.json

@@ -63,6 +63,7 @@
     "framer-motion": "^12.0.11",
     "geist": "^1.3.1",
     "launchdarkly-react-client-sdk": "^3.6.1",
+    "lodash": "^4.17.21",
     "lucide-react": "^0.474.0",
     "moment": "^2.30.1",
     "next": "^14.2.21",
@@ -94,6 +95,7 @@
     "@storybook/react": "^8.3.5",
     "@storybook/test": "^8.3.5",
     "@storybook/test-runner": "^0.21.0",
+    "@types/lodash": "^4.17.15",
     "@types/negotiator": "^0.6.3",
     "@types/node": "^22.13.0",
     "@types/react": "^18",

autogpt_platform/frontend/src/app/build/page.tsx

@@ -3,15 +3,43 @@
 import { useSearchParams } from "next/navigation";
 import { GraphID } from "@/lib/autogpt-server-api/types";
 import FlowEditor from "@/components/Flow";
+import React from "react";
+import BuildFlow from "@/components/build/BuildFlow";
+import { CustomNode } from "@/components/CustomNode";
 
 export default function Home() {
   const query = useSearchParams();
 
+  const nodeTypes = { custom: CustomNode };
+
+  // Define initial nodes and edges
+  const initialNodes = [
+    {
+      id: "1",
+      type: "custom", // Type must match the key in the `nodeTypes` object
+      position: { x: 100, y: 100 },
+      data: { label: "I am a custom node" }, // Must match `CustomNodeData`
+    },
+    {
+      id: "2",
+      type: "custom",
+      position: { x: 400, y: 200 },
+      data: { label: "Another custom node" },
+    },
+  ];
+
+  const initialEdges = [
+    { id: "e1-2", source: "1", target: "2", type: "default" },
+  ];
+
   return (
-    <FlowEditor
-      className="flow-container"
-      flowID={query.get("flowID") as GraphID | null ?? undefined}
-      flowVersion={query.get("flowVersion") ?? undefined}
-    />
+    <div className="flow-container">
+      <BuildFlow
+        id="example-canvas"
+        initialNodes={initialNodes}
+        initialEdges={initialEdges}
+        readOnly={false}
+      />
+    </div>
   );
 }

autogpt_platform/frontend/src/components/build/BuildCustomEdge.tsx

@@ -0,0 +1,37 @@
+import React from "react";
+import { EdgeProps, getBezierPath } from "@xyflow/react";
+
+const BuildCustomEdge: React.FC<EdgeProps> = ({
+  id,
+  sourceX,
+  sourceY,
+  targetX,
+  targetY,
+  sourcePosition,
+  targetPosition,
+  style = {},
+  markerEnd,
+}) => {
+  const [edgePath] = getBezierPath({
+    sourceX,
+    sourceY,
+    sourcePosition,
+    targetX,
+    targetY,
+    targetPosition,
+  });
+
+  return (
+    <>
+      <path
+        id={id}
+        style={{ ...style, stroke: "#888", strokeWidth: 2 }}
+        className="react-flow__edge-path"
+        d={edgePath}
+        markerEnd={markerEnd}
+      />
+    </>
+  );
+};
+
+export default BuildCustomEdge;

autogpt_platform/frontend/src/components/build/BuildCustomNode.tsx

@@ -0,0 +1,35 @@
+import React, { useCallback } from "react";
+import { Handle, Position, useReactFlow, Node } from "@xyflow/react";
+
+export interface BuildCustomNodeData extends Record<string, unknown> {
+  label: string; // Label to be displayed in the node
+}
+
+const BuildCustomNode: React.FC<Node<BuildCustomNodeData>> = ({ id, data }) => {
+  const { fitView } = useReactFlow();
+
+  const handleClick = useCallback(() => {
+    console.log(`Fitting view on Node: ${id}`);
+    fitView({ duration: 800 });
+  }, [id, fitView]);
+
+  return (
+    <div
+      style={{
+        border: "1px solid #ddd",
+        padding: "10px",
+        borderRadius: "4px",
+        background: "white",
+      }}
+    >
+      <Handle type="target" position={Position.Top} />
+      <div>{data.label}</div>
+      <button onClick={handleClick} style={{ marginTop: "5px" }}>
+        Focus Node
+      </button>
+      <Handle type="source" position={Position.Bottom} />
+    </div>
+  );
+};
+
+export default BuildCustomNode;

autogpt_platform/frontend/src/components/build/BuildFlow.tsx

@@ -0,0 +1,163 @@
+"use client";
+
+import React, { useCallback, useMemo, useState } from "react";
+import {
+  ReactFlow,
+  ReactFlowProvider,
+  Background,
+  Controls,
+  addEdge,
+  applyNodeChanges,
+  applyEdgeChanges,
+  Connection,
+  Edge,
+  Node,
+  NodeChange,
+  EdgeChange,
+} from "@xyflow/react";
+import { Undo2, Redo2, RefreshCw } from "lucide-react";
+import ControlPanel from "@/components/edit/control/ControlPanel";
+import useUndoRedo from "@/hooks/build/useUndoRedo";
+
+export interface CanvasProps {
+  id?: string;
+  initialNodes?: Node<any>[];
+  initialEdges?: Edge<any>[];
+  readOnly?: boolean;
+}
+
+const CanvasInner = ({
+  id = "canvas",
+  initialNodes = [],
+  initialEdges = [],
+  readOnly = false,
+}: CanvasProps) => {
+  const { current, addState, undo, redo, canUndo, canRedo, reset } =
+    useUndoRedo({
+      nodes: initialNodes,
+      edges: initialEdges,
+    });
+
+  const [tempNodes, setTempNodes] = useState<Node<any>[]>(current.nodes);
+  const [isDragging, setIsDragging] = useState(false);
+
+  const onNodesChange = useCallback(
+    (changes: NodeChange[]) => {
+      const updatedNodes = applyNodeChanges(changes, tempNodes);
+      setTempNodes(updatedNodes);
+
+      // Only save state for significant changes
+      const hasSignificantChanges = changes.some(
+        (change) =>
+          change.type === "position" ||
+          change.type === "remove" ||
+          change.type === "add",
+      );
+
+      if (!isDragging && hasSignificantChanges) {
+        const cleanNodes = updatedNodes.map((node) => {
+          const { selected, measured, dragging, ...cleanNode } = node;
+          return cleanNode;
+        });
+        addState({ nodes: cleanNodes, edges: current.edges });
+      }
+    },
+    [tempNodes, current.edges, addState, isDragging],
+  );
+
+  const onNodeDragStart = useCallback(() => {
+    setIsDragging(true);
+  }, []);
+
+  const onNodeDragStop = useCallback(() => {
+    setIsDragging(false);
+    // Clean nodes before saving
+    const cleanNodes = tempNodes.map((node) => {
+      const { selected, measured, dragging, ...cleanNode } = node;
+      return cleanNode;
+    });
+
+    if (JSON.stringify(cleanNodes) !== JSON.stringify(current.nodes)) {
+      addState({ nodes: cleanNodes, edges: current.edges });
+    }
+  }, [tempNodes, current.nodes, current.edges, addState]);
+
+  const onEdgesChange = useCallback(
+    (changes: EdgeChange[]) => {
+      const updatedEdges = applyEdgeChanges(changes, current.edges);
+      if (changes.length > 0) {
+        addState({ nodes: current.nodes, edges: updatedEdges });
+      }
+    },
+    [current.nodes, current.edges, addState],
+  );
+
+  const onConnect = useCallback(
+    (connection: Connection) => {
+      const updatedEdges = addEdge(connection, current.edges);
+      addState({ nodes: current.nodes, edges: updatedEdges });
+    },
+    [current.nodes, current.edges, addState],
+  );
+
+  // Sync tempNodes when current.nodes changes (e.g., after undo/redo)
+  React.useEffect(() => {
+    setTempNodes(current.nodes);
+  }, [current.nodes]);
+
+  const controls = useMemo(
+    () => [
+      {
+        icon: <Undo2 className="h-4 w-4" />,
+        label: "Undo",
+        disabled: !canUndo,
+        onClick: undo,
+      },
+      {
+        icon: <Redo2 className="h-4 w-4" />,
+        label: "Redo",
+        disabled: !canRedo,
+        onClick: redo,
+      },
+      {
+        icon: <RefreshCw className="h-4 w-4" />,
+        label: "Reset",
+        onClick: reset,
+      },
+    ],
+    [canUndo, canRedo, undo, redo, reset],
+  );
+
+  return (
+    <div className="relative flex h-full w-full">
+      <ControlPanel controls={controls} className="z-10" />
+      <div className="flex-1">
+        <ReactFlow
+          id={id}
+          nodes={tempNodes}
+          edges={current.edges}
+          onNodesChange={onNodesChange}
+          onNodeDragStart={onNodeDragStart}
+          onNodeDragStop={onNodeDragStop}
+          onEdgesChange={onEdgesChange}
+          onConnect={onConnect}
+          nodesDraggable={!readOnly}
+          nodesConnectable={!readOnly}
+          fitView
+          className={readOnly ? "read-only-mode" : ""}
+        >
+          <Controls />
+          <Background gap={16} color="#ddd" />
+        </ReactFlow>
+      </div>
+    </div>
+  );
+};
+
+const BuildFlow = (props: CanvasProps) => (
+  <ReactFlowProvider>
+    <CanvasInner {...props} />
+  </ReactFlowProvider>
+);
+
+export default BuildFlow;

autogpt_platform/frontend/src/hooks/build/useCanvasMapping.ts

@@ -0,0 +1,34 @@
+import { useMemo } from "react";
+import { Node, Edge } from "@xyflow/react";
+
+interface UseCanvasMappingProps<T extends Record<string, unknown> = any> {
+  nodes: Node<T>[];
+  edges: Edge[];
+}
+
+export const useCanvasMapping = <T extends Record<string, unknown>>({
+  nodes,
+  edges,
+}: UseCanvasMappingProps<T>) => {
+  // Map or enhance nodes
+  const mappedNodes = useMemo(
+    () =>
+      nodes.map((node) => ({
+        ...node,
+        data: { ...node.data, label: `Mapped Node - ${node.id}` },
+      })),
+    [nodes],
+  );
+
+  // Map or enhance edges
+  const mappedEdges = useMemo(
+    () =>
+      edges.map((edge) => ({
+        ...edge,
+        style: { ...edge.style, stroke: "#00f" },
+      })),
+    [edges],
+  );
+
+  return { mappedNodes, mappedEdges };
+};

autogpt_platform/frontend/src/hooks/build/useUndoRedo.ts

@@ -0,0 +1,111 @@
+import { useCallback, useEffect, useState, useRef } from "react";
+import { Edge, Node } from "@xyflow/react";
+import isEqual from "lodash/isEqual";
+
+const LOCAL_STORAGE_KEY = "build-state";
+const MAX_HISTORY_LENGTH = 50;
+
+interface CanvasState {
+  nodes: Node<any>[];
+  edges: Edge<any>[];
+}
+
+function useUndoRedo(initialState: CanvasState) {
+  const [history, setHistory] = useState<CanvasState[]>([initialState]);
+  const [pointer, setPointer] = useState(0);
+  const isFirstLoad = useRef(true);
+  const lastSavedState = useRef<string | null>(null);
+
+  // Load state from localStorage only once on mount
+  useEffect(() => {
+    const cachedState = localStorage.getItem(LOCAL_STORAGE_KEY);
+    if (cachedState && isFirstLoad.current) {
+      try {
+        const parsedHistory: CanvasState[] = JSON.parse(cachedState);
+        if (parsedHistory.every(validateCanvasState)) {
+          setHistory(parsedHistory);
+          setPointer(parsedHistory.length - 1);
+          lastSavedState.current = cachedState;
+        }
+      } catch (error) {
+        console.error("Failed to parse cached state:", error);
+      }
+    }
+    isFirstLoad.current = false;
+  }, []);
+
+  // Save state to localStorage
+  useEffect(() => {
+    if (isFirstLoad.current) return;
+
+    const newState = JSON.stringify(history.slice(0, pointer + 1));
+    if (newState !== lastSavedState.current) {
+      lastSavedState.current = newState;
+      localStorage.setItem(LOCAL_STORAGE_KEY, newState);
+    }
+  }, [history, pointer]);
+
+  const addState = useCallback(
+    (newState: CanvasState) => {
+      if (isFirstLoad.current) return;
+
+      setHistory((prevHistory) => {
+        // If the new state is the same as the current state, don't add it
+        const currentState = prevHistory[pointer];
+        if (isEqual(currentState, newState)) return prevHistory;
+
+        // Truncate the history to the current pointer and add the new state
+        const newHistory = [...prevHistory.slice(0, pointer + 1), newState];
+
+        // Maintain the maximum history length
+        return newHistory.slice(-MAX_HISTORY_LENGTH);
+      });
+
+      // Move the pointer to the new end of history
+      setPointer((prev) => {
+        const newPointer = prev + 1;
+        return Math.min(newPointer, MAX_HISTORY_LENGTH - 1);
+      });
+    },
+    [pointer],
+  );
+
+  const undo = useCallback(() => {
+    if (pointer > 0) {
+      setPointer((prev) => prev - 1);
+    }
+  }, [pointer]);
+
+  const redo = useCallback(() => {
+    if (pointer < history.length - 1) {
+      setPointer((prev) => prev + 1);
+    }
+  }, [pointer, history.length]);
+
+  const reset = useCallback(() => {
+    setHistory([initialState]);
+    setPointer(0);
+  }, [initialState]);
+
+  return {
+    history,
+    current: history[pointer],
+    undo,
+    redo,
+    canUndo: pointer > 0,
+    canRedo: pointer < history.length - 1,
+    addState,
+    reset,
+  };
+}
+
+const validateCanvasState = (state: CanvasState): boolean => {
+  return (
+    Array.isArray(state.nodes) &&
+    Array.isArray(state.edges) &&
+    state.nodes.every((node) => node && node.id && node.position) &&
+    state.edges.every((edge) => edge && edge.id && edge.source && edge.target)
+  );
+};
+
+export default useUndoRedo;