refactor: custom nodes and edges (#929)

* refactor: custom nodes and edges

* fix: error

* fix: reset workspace after app loaded

Author: shczhen

core/src/ten_manager/designer_frontend/src/App.tsx

@@ -18,9 +18,7 @@ import { ThemeProvider } from "@/components/ThemeProvider";
 import AppBar from "@/components/AppBar";
 import StatusBar from "@/components/StatusBar";
 import FlowCanvas, { type FlowCanvasRef } from "@/flow/FlowCanvas";
-import { CustomNodeType } from "@/flow/CustomNode";
-import { CustomEdgeType } from "@/flow/CustomEdge";
-import { getLayoutedElements } from "@/flow/graph";
+import { generateNodesAndEdges } from "@/flow/graph";
 import {
   ResizableHandle,
   ResizablePanel,
@@ -37,6 +35,8 @@ import { usePreferences } from "@/api/services/common";
 import { SpinnerLoading } from "@/components/Status/Loading";
 import { PREFERENCES_SCHEMA_LOG } from "@/types/apps";
 
+import type { TCustomEdge, TCustomNode } from "@/types/flow";
+
 const App: React.FC = () => {
   const { nodes, setNodes, edges, setEdges, setNodesAndEdges } = useFlowStore();
   const [resizablePanelMode, setResizablePanelMode] = useState<
@@ -63,16 +63,14 @@ const App: React.FC = () => {
   );
 
   const performAutoLayout = useCallback(() => {
-    const { nodes: layoutedNodes, edges: layoutedEdges } = getLayoutedElements(
-      nodes,
-      edges
-    );
+    const { nodes: layoutedNodes, edges: layoutedEdges } =
+      generateNodesAndEdges(nodes, edges);
     setNodesAndEdges(layoutedNodes, layoutedEdges);
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [nodes, edges]);
 
   const handleNodesChange = useCallback(
-    (changes: NodeChange<CustomNodeType>[]) => {
+    (changes: NodeChange<TCustomNode>[]) => {
       const newNodes = applyNodeChanges(changes, nodes);
       setNodes(newNodes);
     },
@@ -81,7 +79,7 @@ const App: React.FC = () => {
   );
 
   const handleEdgesChange = useCallback(
-    (changes: EdgeChange<CustomEdgeType>[]) => {
+    (changes: EdgeChange<TCustomEdge>[]) => {
       const newEdges = applyEdgeChanges(changes, edges);
       setEdges(newEdges);
     },

core/src/ten_manager/designer_frontend/src/components/DataTable/ConnectionTable.tsx

@@ -108,7 +108,7 @@ export const commonConnectionColumns: ColumnDef<TConnection>[] = [
     cell: ({ row }) => {
       const name = row.getValue("name") as string;
       if (!name) return null;
-      return <div className="text-xs">{name}</div>;
+      return name;
     },
   },
 ];

core/src/ten_manager/designer_frontend/src/components/Popup/AppPopup.tsx

@@ -48,14 +48,18 @@ export const AppFolderPopup = () => {
 
   const { removeWidget } = useWidgetStore();
   const { setNodesAndEdges } = useFlowStore();
-  const { folderPath } = useAppStore();
+  const { folderPath, updateCurrentWorkspace } = useAppStore();
 
   const { mutate: mutateApps } = useApps();
 
   const handleSetBaseDir = async (folderPath: string) => {
     try {
       await postLoadDir(folderPath.trim());
       setNodesAndEdges([], []); // Clear the contents of the FlowCanvas.
+      updateCurrentWorkspace({
+        baseDir: folderPath.trim(),
+        graphName: null,
+      });
       mutateApps();
       toast.success(t("header.menuApp.loadAppSuccess"));
     } catch (error: unknown) {

core/src/ten_manager/designer_frontend/src/components/Popup/GraphSelectPopup.tsx

@@ -10,11 +10,11 @@ import { toast } from "sonner";
 import { CheckIcon } from "lucide-react";
 
 import {
-  enhanceNodesWithCommands,
-  fetchAddonInfoForNodes,
-  getLayoutedElements,
-  processConnections,
-  processNodes,
+  generateRawNodes,
+  generateRawEdges,
+  updateNodesWithConnections,
+  updateNodesWithAddonInfo,
+  generateNodesAndEdges,
 } from "@/flow/graph";
 import { Popup } from "@/components/Popup/Popup";
 import {
@@ -38,8 +38,6 @@ import { useApps } from "@/api/services/apps";
 import { useWidgetStore, useFlowStore, useAppStore } from "@/store";
 import { GRAPH_SELECT_POPUP_ID } from "@/constants/widgets";
 
-import type { CustomNodeType } from "@/flow/CustomNode";
-
 export function GraphSelectPopup() {
   const { t } = useTranslation();
 
@@ -79,43 +77,20 @@ export function GraphSelectPopup() {
           graphName,
           baseDir
         );
-
-        let initialNodes: CustomNodeType[] = processNodes(backendNodes);
-
-        const {
-          initialEdges,
-          nodeSourceCmdMap,
-          nodeSourceDataMap,
-          nodeSourceAudioFrameMap,
-          nodeSourceVideoFrameMap,
-          nodeTargetCmdMap,
-          nodeTargetDataMap,
-          nodeTargetAudioFrameMap,
-          nodeTargetVideoFrameMap,
-        } = processConnections(backendConnections);
-
-        // Write back the cmd information to nodes, so that CustomNode could
-        // generate corresponding handles.
-        initialNodes = enhanceNodesWithCommands(initialNodes, {
-          nodeSourceCmdMap,
-          nodeTargetCmdMap,
-          nodeSourceDataMap,
-          nodeTargetDataMap,
-          nodeSourceAudioFrameMap,
-          nodeSourceVideoFrameMap,
-          nodeTargetAudioFrameMap,
-          nodeTargetVideoFrameMap,
-        });
-
-        // Fetch additional addon information for each node.
-        const nodesWithAddonInfo = await fetchAddonInfoForNodes(
+        const rawNodes = generateRawNodes(backendNodes);
+        const [rawEdges, rawEdgeAddressMap] =
+          generateRawEdges(backendConnections);
+        const nodesWithConnections = updateNodesWithConnections(
+          rawNodes,
+          rawEdgeAddressMap
+        );
+        const nodesWithAddonInfo = await updateNodesWithAddonInfo(
           baseDir ?? "",
-          initialNodes
+          nodesWithConnections
         );
 
-        // Auto-layout the nodes and edges.
         const { nodes: layoutedNodes, edges: layoutedEdges } =
-          getLayoutedElements(nodesWithAddonInfo, initialEdges);
+          generateNodesAndEdges(nodesWithAddonInfo, rawEdges);
 
         setNodesAndEdges(layoutedNodes, layoutedEdges);
       } catch (err: unknown) {

core/src/ten_manager/designer_frontend/src/flow/ContextMenu/EdgeContextMenu.tsx

@@ -9,14 +9,15 @@ import { useTranslation } from "react-i18next";
 import { PencilIcon, ListCollapseIcon, TrashIcon } from "lucide-react";
 
 import ContextMenu, { ContextMenuItem } from "@/flow/ContextMenu/ContextMenu";
-import { CustomEdgeType } from "@/flow/CustomEdge";
 import { dispatchCustomNodeActionPopup } from "@/utils/popup";
 
+import type { TCustomEdge } from "@/types/flow";
+
 interface EdgeContextMenuProps {
   visible: boolean;
   x: number;
   y: number;
-  edge: CustomEdgeType;
+  edge: TCustomEdge;
   onClose: () => void;
 }
 

core/src/ten_manager/designer_frontend/src/flow/ContextMenu/NodeContextMenu.tsx

@@ -14,15 +14,15 @@ import {
 } from "lucide-react";
 
 import ContextMenu, { ContextMenuItem } from "@/flow/ContextMenu/ContextMenu";
-import { CustomNodeType } from "@/flow/CustomNode";
 
 import type { TerminalData, EditorData } from "@/types/widgets";
+import type { TCustomNode } from "@/types/flow";
 
 interface NodeContextMenuProps {
   visible: boolean;
   x: number;
   y: number;
-  node: CustomNodeType;
+  node: TCustomNode;
   onClose: () => void;
   onLaunchTerminal: (data: TerminalData) => void;
   onLaunchEditor: (data: EditorData) => void;
@@ -49,7 +49,7 @@ const NodeContextMenu: React.FC<NodeContextMenuProps> = ({
         onClose();
         if (node?.data.url)
           onLaunchEditor({
-            title: `${node.data.addon} manifest.json`,
+            title: `${node.data.name} manifest.json`,
             content: "",
             url: `${node.data.url}/manifest.json`,
           });
@@ -62,7 +62,7 @@ const NodeContextMenu: React.FC<NodeContextMenuProps> = ({
         onClose();
         if (node?.data.url)
           onLaunchEditor({
-            title: `${node.data.addon} property.json`,
+            title: `${node.data.name} property.json`,
             content: "",
             url: `${node.data.url}/property.json`,
           });

core/src/ten_manager/designer_frontend/src/flow/CustomEdge.tsx

@@ -6,25 +6,12 @@
 //
 import {
   type EdgeProps,
-  type Edge,
   BaseEdge,
   getSmoothStepPath,
   Position,
 } from "@xyflow/react";
 
-import { type EConnectionType } from "@/types/graphs";
-
-export type CustomEdgeType = Edge<
-  {
-    labelOffsetX: number;
-    labelOffsetY: number;
-    connectionType: EConnectionType;
-    name: string;
-    srcApp: string;
-    destApp: string;
-  },
-  "customEdge"
->;
+import type { TCustomEdge } from "@/types/flow";
 
 export function CustomEdge({
   sourceX,
@@ -35,7 +22,7 @@ export function CustomEdge({
   style,
   selected,
   markerEnd,
-}: EdgeProps<CustomEdgeType>) {
+}: EdgeProps<TCustomEdge>) {
   const [path] = getSmoothStepPath({
     sourceX: sourceX,
     sourceY: sourceY,

core/src/ten_manager/designer_frontend/src/flow/CustomNode.tsx

@@ -5,7 +5,7 @@
 // Refer to the "LICENSE" file in the root directory for more information.
 //
 import * as React from "react";
-import { Position, NodeProps, Connection, Edge, Node } from "@xyflow/react";
+import { Position, NodeProps, Connection, Edge } from "@xyflow/react";
 import {
   BlocksIcon as ExtensionIcon,
   TerminalIcon,
@@ -18,28 +18,14 @@ import { Separator } from "@/components/ui/Separator";
 import { cn } from "@/lib/utils";
 import { dispatchCustomNodeActionPopup } from "@/utils/popup";
 import CustomHandle from "@/flow/CustomHandle";
-import { EConnectionType, type TConnectionItem } from "@/types/graphs";
+import { EConnectionType } from "@/types/graphs";
+
+import type { TCustomNode, TCustomNodeData } from "@/types/flow";
 
 const onConnect = (params: Connection | Edge) =>
   console.log("Handle onConnect", params);
 
-export type TCustomNodeData = {
-  name: string;
-  addon: string;
-  sourceCmds: TConnectionItem[];
-  targetCmds: TConnectionItem[];
-  sourceData: TConnectionItem[];
-  targetData: TConnectionItem[];
-  sourceAudioFrame: TConnectionItem[];
-  targetAudioFrame: TConnectionItem[];
-  sourceVideoFrame: TConnectionItem[];
-  targetVideoFrame: TConnectionItem[];
-  url?: string;
-};
-
-export type CustomNodeType = Node<TCustomNodeData, "customNode">;
-
-export function CustomNode({ data, isConnectable }: NodeProps<CustomNodeType>) {
+export function CustomNode({ data, isConnectable }: NodeProps<TCustomNode>) {
   return (
     <>
       <div
@@ -119,7 +105,7 @@ const HandleGroupItem = (props: {
       target?: boolean;
     }) =>
     () => {
-      dispatchCustomNodeActionPopup("connections", data.addon, undefined, {
+      dispatchCustomNodeActionPopup("connections", data.name, undefined, {
         filters: {
           type,
           source,
@@ -132,11 +118,11 @@ const HandleGroupItem = (props: {
     <div className="flex items-center gap-x-4 justify-between">
       <div className="flex items-center gap-x-2">
         <CustomHandle
-          key={`target-${data.addon}-${connectionType}`}
+          key={`target-${data.name}-${connectionType}`}
           type="target"
           position={Position.Left}
-          id={`target-${data.addon}-${connectionType}`}
-          label={data.addon}
+          id={`target-${data.name}-${connectionType}`}
+          label={data.name}
           labelOffsetX={0}
           isConnectable={isConnectable}
           onConnect={onConnect}
@@ -148,16 +134,16 @@ const HandleGroupItem = (props: {
           })}
         >
           {connectionType === EConnectionType.CMD && (
-            <span>{data.targetCmds?.length || 0}</span>
+            <span>{data.src[connectionType]?.length || 0}</span>
           )}
           {connectionType === EConnectionType.DATA && (
-            <span>{data.targetData?.length || 0}</span>
+            <span>{data.src[connectionType]?.length || 0}</span>
           )}
           {connectionType === EConnectionType.AUDIO_FRAME && (
-            <span>{data.targetAudioFrame?.length || 0}</span>
+            <span>{data.src[connectionType]?.length || 0}</span>
           )}
           {connectionType === EConnectionType.VIDEO_FRAME && (
-            <span>{data.targetVideoFrame?.length || 0}</span>
+            <span>{data.src[connectionType]?.length || 0}</span>
           )}
         </ConnectionCount>
       </div>
@@ -191,24 +177,24 @@ const HandleGroupItem = (props: {
           })}
         >
           {connectionType === EConnectionType.CMD && (
-            <span>{data.sourceCmds?.length || 0}</span>
+            <span>{data.target[connectionType]?.length || 0}</span>
           )}
           {connectionType === EConnectionType.DATA && (
-            <span>{data.sourceData?.length || 0}</span>
+            <span>{data.target[connectionType]?.length || 0}</span>
           )}
           {connectionType === EConnectionType.AUDIO_FRAME && (
-            <span>{data.sourceAudioFrame?.length || 0}</span>
+            <span>{data.target[connectionType]?.length || 0}</span>
           )}
           {connectionType === EConnectionType.VIDEO_FRAME && (
-            <span>{data.sourceVideoFrame?.length || 0}</span>
+            <span>{data.target[connectionType]?.length || 0}</span>
           )}
         </ConnectionCount>
         <CustomHandle
-          key={`source-${data.addon}-${connectionType}`}
+          key={`source-${data.name}-${connectionType}`}
           type="source"
           position={Position.Right}
-          id={`source-${data.addon}-${connectionType}`}
-          label={data.addon}
+          id={`source-${data.name}-${connectionType}`}
+          label={data.name}
           labelOffsetX={0}
           isConnectable={isConnectable}
         />