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 { updateSubgraphSpec } from "@/utils/subgraphUtils";
 
 import { NameField, TextField, TypeField } from "./FormFields/FormFields";
@@ -31,8 +31,10 @@ export const InputValueEditor = ({
   input,
   disabled = false,
 }: InputValueEditorProps) => {
+  const { getInputNodeId } = useNodeManager();
+
   const notify = useToastNotification();
-  const { transferSelection } = useNodeSelectionTransfer(inputNameToNodeId);
+  const { transferSelection } = useNodeSelectionTransfer(getInputNodeId);
   const {
     componentSpec,
     setComponentSpec,

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 { updateSubgraphSpec } from "@/utils/subgraphUtils";
 
 import { type OutputConnectedDetails } from "../../utils/getOutputConnectedDetails";
@@ -30,7 +30,8 @@ export const OutputNameEditor = ({
   disabled,
   connectedDetails,
 }: OutputNameEditorProps) => {
-  const { transferSelection } = useNodeSelectionTransfer(outputNameToNodeId);
+  const { getOutputNodeId } = useNodeManager();
+  const { transferSelection } = useNodeSelectionTransfer(getOutputNodeId);
   const {
     setComponentSpec,
     componentSpec,

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

@@ -119,6 +119,7 @@ const FlowCanvas = ({
     currentSubgraphSpec,
     updateGraphSpec,
     currentSubgraphPath,
+    nodeManager,
   } = useComponentSpec();
   const { preserveIOSelectionOnSpecChange, resetPrevSpec } =
     useIOSelectionPersistence();
@@ -285,10 +286,18 @@ const FlowCanvas = ({
       let updatedSubgraphSpec = { ...currentSubgraphSpec };
 
       for (const edge of params.edges) {
-        updatedSubgraphSpec = removeEdge(edge, updatedSubgraphSpec);
+        updatedSubgraphSpec = removeEdge(
+          edge,
+          updatedSubgraphSpec,
+          nodeManager,
+        );
       }
       for (const node of params.nodes) {
-        updatedSubgraphSpec = removeNode(node, updatedSubgraphSpec);
+        updatedSubgraphSpec = removeNode(
+          node,
+          updatedSubgraphSpec,
+          nodeManager,
+        );
       }
 
       const updatedRootSpec = updateSubgraphSpec(
@@ -299,7 +308,13 @@ const FlowCanvas = ({
 
       setComponentSpec(updatedRootSpec);
     },
-    [componentSpec, currentSubgraphSpec, currentSubgraphPath, setComponentSpec],
+    [
+      componentSpec,
+      currentSubgraphSpec,
+      currentSubgraphPath,
+      nodeManager,
+      setComponentSpec,
+    ],
   );
 
   const nodeCallbacks = useNodeCallbacks({
@@ -313,18 +328,23 @@ 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(currentGraphSpec, connection);
+      const updatedGraphSpec = handleConnection(
+        currentGraphSpec,
+        connection,
+        nodeManager,
+      );
       updateGraphSpec(updatedGraphSpec);
     },
-    [currentGraphSpec, handleConnection, updateGraphSpec],
+    [currentGraphSpec, nodeManager, handleConnection, updateGraphSpec],
   );
 
   const onConnectEnd = useCallback(
@@ -361,19 +381,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(() => {
@@ -627,6 +658,7 @@ const FlowCanvas = ({
           const updatedSubgraphSpec = updateNodePositions(
             updatedNodes,
             currentSubgraphSpec,
+            nodeManager,
           );
 
           const updatedRootSpec = updateSubgraphSpec(
@@ -646,6 +678,7 @@ const FlowCanvas = ({
       componentSpec,
       currentSubgraphSpec,
       currentSubgraphPath,
+      nodeManager,
       setComponentSpec,
       onNodesChange,
     ],
@@ -677,7 +710,9 @@ const FlowCanvas = ({
       updatedComponentSpec: updatedSubgraphSpec,
       newNodes,
       updatedNodes,
-    } = duplicateNodes(currentSubgraphSpec, selectedNodes, { selected: true });
+    } = duplicateNodes(currentSubgraphSpec, selectedNodes, nodeManager, {
+      selected: true,
+    });
 
     const updatedRootSpec = updateSubgraphSpec(
       componentSpec,
@@ -696,6 +731,7 @@ const FlowCanvas = ({
     currentSubgraphSpec,
     currentSubgraphPath,
     selectedNodes,
+    nodeManager,
     setComponentSpec,
     setNodes,
   ]);
@@ -873,6 +909,7 @@ const FlowCanvas = ({
           const { newNodes, updatedComponentSpec } = duplicateNodes(
             componentSpec,
             nodesToPaste,
+            nodeManager,
             { position: reactFlowCenter, connection: "internal" },
           );
 
@@ -898,6 +935,7 @@ const FlowCanvas = ({
     nodes,
     reactFlowInstance,
     store,
+    nodeManager,
     updateOrAddNodes,
     setComponentSpec,
     readOnly,

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

@@ -7,6 +7,7 @@ import { getOutputConnectedDetails } from "@/components/Editor/utils/getOutputCo
 import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
 import { BlockStack, InlineStack } from "@/components/ui/layout";
 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";
@@ -23,6 +24,7 @@ interface IONodeProps {
 const IONode = ({ type, data, selected = false }: IONodeProps) => {
   const { currentGraphSpec, currentSubgraphSpec } = useComponentSpec();
   const { setContent, clearContent } = useContextPanel();
+  const { getHandleNodeId } = useNodeManager();
 
   const { spec, readOnly } = data;
 
@@ -49,6 +51,9 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
     [currentSubgraphSpec.outputs, spec.name],
   );
 
+  const handleNodeType = isInput ? "handle-out" : "handle-in";
+  const nodeHandleId = getHandleNodeId(spec.name, spec.name, handleNodeType);
+
   useEffect(() => {
     if (selected) {
       if (input && isInput) {
@@ -144,6 +149,7 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
           </InlineStack>
         </BlockStack>
         <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";
@@ -32,7 +33,8 @@ export const InputHandle = ({
   onLabelClick,
   onHandleSelectionChange,
 }: InputHandleProps) => {
-  const { nodeId, state } = useTaskNode();
+  const { getInputHandleNodeId } = useNodeManager();
+  const { taskId, nodeId, state } = useTaskNode();
 
   const fromHandle = useConnection((connection) => connection.fromHandle?.id);
   const toHandle = useConnection((connection) => connection.toHandle?.id);
@@ -44,7 +46,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;
@@ -218,7 +220,8 @@ export const OutputHandle = ({
   onLabelClick,
   onHandleSelectionChange,
 }: OutputHandleProps) => {
-  const { nodeId, state } = useTaskNode();
+  const { getOutputHandleNodeId } = useNodeManager();
+  const { taskId, nodeId, state } = useTaskNode();
 
   const fromHandle = useConnection((connection) => connection.fromHandle?.id);
   const toHandle = useConnection((connection) => connection.toHandle?.id);
@@ -230,7 +233,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(
@@ -355,14 +358,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,
@@ -147,7 +148,7 @@ export function TaskNodeInputs({
     }
 
     const input = inputs.find(
-      (i) => inputNameToNodeId(i.name) === fromHandle?.id,
+      (i) => getInputHandleNodeId(taskId, i.name) === fromHandle?.id,
     );
 
     if (!input) return;