Implement Node Manager

Author: camielvs

src/components/Editor/IOEditor/InputValueEditor/InputValueEditor.tsx

@@ -9,13 +9,13 @@ import { Icon } from "@/components/ui/icon";
 import { BlockStack } from "@/components/ui/layout";
 import { Heading, Paragraph } from "@/components/ui/typography";
 import useConfirmationDialog from "@/hooks/useConfirmationDialog";
+import { useNodeManager } from "@/hooks/useNodeManager";
 import { useNodeSelectionTransfer } from "@/hooks/useNodeSelectionTransfer";
 import useToastNotification from "@/hooks/useToastNotification";
 import { useComponentSpec } from "@/providers/ComponentSpecProvider";
 import { useContextPanel } from "@/providers/ContextPanelProvider";
 import { type InputSpec } from "@/utils/componentSpec";
 import { checkInputConnectionToRequiredFields } from "@/utils/inputConnectionUtils";
-import { inputNameToNodeId } from "@/utils/nodes/nodeIdUtils";
 
 import { NameField, TextField, TypeField } from "./FormFields/FormFields";
 import { checkNameCollision } from "./FormFields/utils";
@@ -30,8 +30,11 @@ export const InputValueEditor = ({
   input,
   disabled = false,
 }: InputValueEditorProps) => {
+  const { getInputNodeId } = useNodeManager();
+
   const notify = useToastNotification();
-  const { transferSelection } = useNodeSelectionTransfer(inputNameToNodeId);
+
+  const { transferSelection } = useNodeSelectionTransfer(getInputNodeId);
   const { componentSpec, setComponentSpec } = useComponentSpec();
   const { clearContent } = useContextPanel();
   const {

src/components/Editor/IOEditor/OutputNameEditor/OutputNameEditor.tsx

@@ -7,11 +7,11 @@ import { Icon } from "@/components/ui/icon";
 import { BlockStack, InlineStack } from "@/components/ui/layout";
 import { Heading, Paragraph } from "@/components/ui/typography";
 import useConfirmationDialog from "@/hooks/useConfirmationDialog";
+import { useNodeManager } from "@/hooks/useNodeManager";
 import { useNodeSelectionTransfer } from "@/hooks/useNodeSelectionTransfer";
 import { useComponentSpec } from "@/providers/ComponentSpecProvider";
 import { useContextPanel } from "@/providers/ContextPanelProvider";
 import { type OutputSpec } from "@/utils/componentSpec";
-import { outputNameToNodeId } from "@/utils/nodes/nodeIdUtils";
 
 import { type OutputConnectedDetails } from "../../utils/getOutputConnectedDetails";
 import { updateOutputNameOnComponentSpec } from "../../utils/updateOutputNameOnComponentSpec";
@@ -29,7 +29,9 @@ export const OutputNameEditor = ({
   disabled,
   connectedDetails,
 }: OutputNameEditorProps) => {
-  const { transferSelection } = useNodeSelectionTransfer(outputNameToNodeId);
+  const { getOutputNodeId } = useNodeManager();
+
+  const { transferSelection } = useNodeSelectionTransfer(getOutputNodeId);
   const { setComponentSpec, componentSpec } = useComponentSpec();
   const { clearContent } = useContextPanel();
   const {

src/components/shared/ReactFlow/FlowCanvas/FlowCanvas.tsx

@@ -115,6 +115,7 @@ const FlowCanvas = ({
     graphSpec,
     updateGraphSpec,
     currentSubgraphPath,
+    nodeManager,
   } = useComponentSpec();
   const { preserveIOSelectionOnSpecChange, resetPrevSpec } =
     useIOSelectionPersistence();
@@ -286,15 +287,23 @@ const FlowCanvas = ({
       let updatedComponentSpec = { ...componentSpec };
 
       for (const edge of params.edges) {
-        updatedComponentSpec = removeEdge(edge, updatedComponentSpec);
+        updatedComponentSpec = removeEdge(
+          edge,
+          updatedComponentSpec,
+          nodeManager,
+        );
       }
       for (const node of params.nodes) {
-        updatedComponentSpec = removeNode(node, updatedComponentSpec);
+        updatedComponentSpec = removeNode(
+          node,
+          updatedComponentSpec,
+          nodeManager,
+        );
       }
 
       setComponentSpec(updatedComponentSpec);
     },
-    [componentSpec, setComponentSpec],
+    [componentSpec, nodeManager, setComponentSpec],
   );
 
   const nodeCallbacks = useNodeCallbacks({
@@ -308,15 +317,20 @@ const FlowCanvas = ({
       connectable: !readOnly && !!nodesConnectable,
       readOnly,
       callbacks: nodeCallbacks,
+      nodeManager,
     }),
-    [readOnly, nodesConnectable, nodeCallbacks],
+    [readOnly, nodesConnectable, nodeCallbacks, nodeManager],
   );
 
   const onConnect = useCallback(
     (connection: Connection) => {
       if (connection.source === connection.target) return;
 
-      const updatedGraphSpec = handleConnection(graphSpec, connection);
+      const updatedGraphSpec = handleConnection(
+        graphSpec,
+        connection,
+        nodeManager,
+      );
       updateGraphSpec(updatedGraphSpec);
     },
     [graphSpec, handleConnection, updateGraphSpec],
@@ -356,19 +370,30 @@ const FlowCanvas = ({
         );
 
       if (existingInputEdge) {
-        newComponentSpec = removeEdge(existingInputEdge, newComponentSpec);
+        newComponentSpec = removeEdge(
+          existingInputEdge,
+          newComponentSpec,
+          nodeManager,
+        );
       }
 
       const updatedComponentSpec = addAndConnectNode({
         componentRef,
         fromHandle,
         position,
         componentSpec: newComponentSpec,
+        nodeManager,
       });
 
       setComponentSpec(updatedComponentSpec);
     },
-    [reactFlowInstance, componentSpec, setComponentSpec, updateOrAddNodes],
+    [
+      reactFlowInstance,
+      componentSpec,
+      nodeManager,
+      setComponentSpec,
+      updateOrAddNodes,
+    ],
   );
 
   useEffect(() => {
@@ -615,14 +640,15 @@ const FlowCanvas = ({
           const updatedComponentSpec = updateNodePositions(
             updatedNodes,
             componentSpec,
+            nodeManager,
           );
           setComponentSpec(updatedComponentSpec);
         }
       }
 
       onNodesChange(changes);
     },
-    [nodes, componentSpec, setComponentSpec, onNodesChange],
+    [nodes, componentSpec, nodeManager, setComponentSpec, onNodesChange],
   );
 
   const handleBeforeDelete = async (params: NodesAndEdges) => {
@@ -650,6 +676,7 @@ const FlowCanvas = ({
     const { updatedComponentSpec, newNodes, updatedNodes } = duplicateNodes(
       componentSpec,
       selectedNodes,
+      nodeManager,
       { selected: true },
     );
 
@@ -659,7 +686,7 @@ const FlowCanvas = ({
       updatedNodes,
       newNodes,
     });
-  }, [componentSpec, selectedNodes, setComponentSpec, setNodes]);
+  }, [componentSpec, selectedNodes, nodeManager, setComponentSpec, setNodes]);
 
   const onUpgradeNodes = useCallback(async () => {
     let newGraphSpec = graphSpec;
@@ -828,6 +855,7 @@ const FlowCanvas = ({
           const { newNodes, updatedComponentSpec } = duplicateNodes(
             componentSpec,
             nodesToPaste,
+            nodeManager,
             { position: reactFlowCenter, connection: "internal" },
           );
 
@@ -853,6 +881,7 @@ const FlowCanvas = ({
     nodes,
     reactFlowInstance,
     store,
+    nodeManager,
     updateOrAddNodes,
     setComponentSpec,
     readOnly,

src/components/shared/ReactFlow/FlowCanvas/IONode/IONode.tsx

@@ -12,16 +12,13 @@ import {
   TooltipTrigger,
 } from "@/components/ui/tooltip";
 import { Paragraph } from "@/components/ui/typography";
+import { useNodeManager } from "@/hooks/useNodeManager";
 import { cn } from "@/lib/utils";
 import { useComponentSpec } from "@/providers/ComponentSpecProvider";
 import { useContextPanel } from "@/providers/ContextPanelProvider";
 import type { IONodeData } from "@/types/nodes";
 import { isInputSpec, typeSpecToString } from "@/utils/componentSpec";
 import { ENABLE_DEBUG_MODE } from "@/utils/constants";
-import {
-  inputNameToNodeId,
-  outputNameToNodeId,
-} from "@/utils/nodes/nodeIdUtils";
 
 interface IONodeProps {
   type: "input" | "output";
@@ -31,6 +28,7 @@ interface IONodeProps {
 }
 
 const IONode = ({ type, data, selected = false }: IONodeProps) => {
+  const { getNodeId, getHandleNodeId } = useNodeManager();
   const { graphSpec, componentSpec } = useComponentSpec();
   const { setContent, clearContent } = useContextPanel();
 
@@ -58,11 +56,9 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
     [componentSpec.outputs, spec.name],
   );
 
-  const nodeId = isInput
-    ? inputNameToNodeId(spec.name)
-    : outputNameToNodeId(spec.name);
-
-  const nodeHandleId = `${nodeId}_handle`;
+  const nodeId = getNodeId(spec.name, type);
+  const handleNodeType = isInput ? "outputHandle" : "inputHandle";
+  const nodeHandleId = getHandleNodeId(spec.name, spec.name, handleNodeType);
 
   useEffect(() => {
     if (selected) {
@@ -163,6 +159,7 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
         <Tooltip>
           <TooltipTrigger asChild>
             <Handle
+              id={nodeHandleId}
               type={handleType}
               position={handlePosition}
               className={cn(handleDefaultClassName, handleClassName)}

src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskNodeCard/Handles.tsx

@@ -11,6 +11,7 @@ import {
   TooltipContent,
   TooltipTrigger,
 } from "@/components/ui/tooltip";
+import { useNodeManager } from "@/hooks/useNodeManager";
 import { cn } from "@/lib/utils";
 import { useTaskNode } from "@/providers/TaskNodeProvider";
 import type { InputSpec, OutputSpec } from "@/utils/componentSpec";
@@ -33,7 +34,8 @@ export const InputHandle = ({
   onLabelClick,
   onHandleSelectionChange,
 }: InputHandleProps) => {
-  const { nodeId, state, name } = useTaskNode();
+  const { getInputHandleNodeId } = useNodeManager();
+  const { taskId, nodeId, state, name } = useTaskNode();
 
   const fromHandle = useConnection((connection) => connection.fromHandle?.id);
   const toHandle = useConnection((connection) => connection.toHandle?.id);
@@ -45,7 +47,7 @@ export const InputHandle = ({
   const [selected, setSelected] = useState(false);
   const [active, setActive] = useState(false);
 
-  const handleId = getInputHandleId(input.name);
+  const handleId = getInputHandleNodeId(taskId, input.name);
 
   const missing = invalid ? "bg-red-700!" : "bg-gray-500!";
   const hasValue = value !== undefined && value !== null;
@@ -229,7 +231,8 @@ export const OutputHandle = ({
   onLabelClick,
   onHandleSelectionChange,
 }: OutputHandleProps) => {
-  const { nodeId, state, name } = useTaskNode();
+  const { getOutputHandleNodeId } = useNodeManager();
+  const { taskId, nodeId, state, name } = useTaskNode();
 
   const fromHandle = useConnection((connection) => connection.fromHandle?.id);
   const toHandle = useConnection((connection) => connection.toHandle?.id);
@@ -241,7 +244,7 @@ export const OutputHandle = ({
   const [selected, setSelected] = useState(false);
   const [active, setActive] = useState(false);
 
-  const handleId = getOutputHandleId(output.name);
+  const handleId = getOutputHandleNodeId(taskId, output.name);
   const hasValue = value !== undefined && value !== "" && value !== null;
 
   const handleHandleClick = useCallback(
@@ -376,14 +379,6 @@ export const OutputHandle = ({
   );
 };
 
-const getOutputHandleId = (outputName: string) => {
-  return `output_${outputName}`;
-};
-
-const getInputHandleId = (inputName: string) => {
-  return `input_${inputName}`;
-};
-
 const skipHandleDeselect = (e: MouseEvent) => {
   let el = e.target as HTMLElement | null;
   while (el) {

src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskNodeCard/TaskNodeInputs.test.tsx

@@ -17,6 +17,15 @@ vi.mock("@/providers/ComponentLibraryProvider", () => ({
   }),
 }));
 
+vi.mock("@/providers/ComponentLibraryProvider/ForcedSearchProvider", () => ({
+  useForcedSearchContext: () => ({
+    highlightSearchFilter: vi.fn(),
+    resetSearchFilter: vi.fn(),
+    currentSearchFilter: { searchTerm: "", filters: [] },
+    highlightSearchResults: false,
+  }),
+}));
+
 vi.mock("@/providers/ComponentSpecProvider", () => ({
   useComponentSpec: () => ({
     graphSpec: {
@@ -25,6 +34,18 @@ vi.mock("@/providers/ComponentSpecProvider", () => ({
   }),
 }));
 
+vi.mock("@/hooks/useNodeManager", () => ({
+  useNodeManager: () => ({
+    getInputHandleNodeId: vi.fn(
+      (_refId: string, inputName: string) => `input-handle-${inputName}`,
+    ),
+    getOutputHandleNodeId: vi.fn(),
+    getNodeId: vi.fn(),
+    getHandleNodeId: vi.fn(),
+    nodeManager: {},
+  }),
+}));
+
 vi.mock("@/providers/TaskNodeProvider");
 
 const TestWrapper = ReactFlowProvider;
@@ -39,6 +60,7 @@ describe("<TaskNodeInputs />", () => {
       state: { readOnly: false },
       select: vi.fn(),
       nodeId: "test-node",
+      taskId: "test-task",
     } as any);
   };
 

src/components/shared/ReactFlow/FlowCanvas/TaskNode/TaskNodeCard/TaskNodeInputs.tsx

@@ -2,6 +2,7 @@ import { useConnection } from "@xyflow/react";
 import { AlertCircle } from "lucide-react";
 import { type MouseEvent, useCallback, useEffect, useState } from "react";
 
+import { useNodeManager } from "@/hooks/useNodeManager";
 import { cn } from "@/lib/utils";
 import { useForcedSearchContext } from "@/providers/ComponentLibraryProvider/ForcedSearchProvider";
 import { isValidFilterRequest } from "@/providers/ComponentLibraryProvider/types";
@@ -10,7 +11,6 @@ import { useTaskNode } from "@/providers/TaskNodeProvider";
 import { inputsWithInvalidArguments } from "@/services/componentService";
 import type { InputSpec } from "@/utils/componentSpec";
 import { ComponentSearchFilter } from "@/utils/constants";
-import { inputNameToNodeId } from "@/utils/nodes/nodeIdUtils";
 import { checkArtifactMatchesSearchFilters } from "@/utils/searchUtils";
 
 import { InputHandle } from "./Handles";
@@ -27,7 +27,8 @@ export function TaskNodeInputs({
   expanded,
   onBackgroundClick,
 }: TaskNodeInputsProps) {
-  const { inputs, taskSpec, state, select } = useTaskNode();
+  const { getInputHandleNodeId } = useNodeManager();
+  const { taskId, inputs, taskSpec, state, select } = useTaskNode();
   const { graphSpec } = useComponentSpec();
   const {
     highlightSearchFilter,
@@ -145,7 +146,7 @@ export function TaskNodeInputs({
     }
 
     const input = inputs.find(
-      (i) => inputNameToNodeId(i.name) === fromHandle?.id,
+      (i) => getInputHandleNodeId(taskId, i.name) === fromHandle?.id,
     );
 
     if (!input) return;