Migrate Legacy NodeId Utilities to NodeManager

Author: camielvs

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

@@ -9,13 +9,14 @@ 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 { inputNameToInputId } from "@/utils/nodes/conversions";
 
 import { NameField, TextField, TypeField } from "./FormFields/FormFields";
 import { checkNameCollision } from "./FormFields/utils";
@@ -30,6 +31,16 @@ export const InputValueEditor = ({
   input,
   disabled = false,
 }: InputValueEditorProps) => {
+  const { getInputNodeId } = useNodeManager();
+
+  const inputNameToNodeId = useCallback(
+    (inputName: string): string => {
+      const inputId = inputNameToInputId(inputName);
+      return getInputNodeId(inputId);
+    },
+    [getInputNodeId],
+  );
+
   const notify = useToastNotification();
   const { transferSelection } = useNodeSelectionTransfer(inputNameToNodeId);
   const { componentSpec, setComponentSpec } = useComponentSpec();

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,6 +29,16 @@ export const OutputNameEditor = ({
   disabled,
   connectedDetails,
 }: OutputNameEditorProps) => {
+  const { getOutputNodeId } = useNodeManager();
+
+  const outputNameToNodeId = useCallback(
+    (outputName: string): string => {
+      const outputId = outputNameToNodeId(outputName);
+      return getOutputNodeId(outputId);
+    },
+    [getOutputNodeId],
+  );
+
   const { transferSelection } = useNodeSelectionTransfer(outputNameToNodeId);
   const { setComponentSpec, componentSpec } = useComponentSpec();
   const { clearContent } = useContextPanel();

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

@@ -303,15 +303,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({
@@ -334,7 +342,11 @@ const FlowCanvas = ({
     (connection: Connection) => {
       if (connection.source === connection.target) return;
 
-      const updatedGraphSpec = handleConnection(graphSpec, connection);
+      const updatedGraphSpec = handleConnection(
+        graphSpec,
+        connection,
+        nodeManager,
+      );
       updateGraphSpec(updatedGraphSpec);
     },
     [graphSpec, handleConnection, updateGraphSpec],
@@ -374,19 +386,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(() => {
@@ -633,14 +656,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) => {

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,7 @@ 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 { inputNameToInputId } from "@/utils/nodes/conversions";
 import { checkArtifactMatchesSearchFilters } from "@/utils/searchUtils";
 
 import { InputHandle } from "./Handles";
@@ -27,6 +28,7 @@ export function TaskNodeInputs({
   expanded,
   onBackgroundClick,
 }: TaskNodeInputsProps) {
+  const { getInputNodeId } = useNodeManager();
   const { inputs, taskSpec, state, select } = useTaskNode();
   const { graphSpec } = useComponentSpec();
   const {
@@ -145,7 +147,7 @@ export function TaskNodeInputs({
     }
 
     const input = inputs.find(
-      (i) => inputNameToNodeId(i.name) === fromHandle?.id,
+      (i) => getInputNodeId(inputNameToInputId(i.name)) === fromHandle?.id,
     );
 
     if (!input) return;

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

@@ -1,13 +1,14 @@
 import { useConnection, useEdges } from "@xyflow/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";
 import { useTaskNode } from "@/providers/TaskNodeProvider";
 import type { OutputSpec } from "@/utils/componentSpec";
 import { ComponentSearchFilter } from "@/utils/constants";
-import { outputNameToNodeId } from "@/utils/nodes/nodeIdUtils";
+import { outputNameToOutputId } from "@/utils/nodes/conversions";
 import { checkArtifactMatchesSearchFilters } from "@/utils/searchUtils";
 
 import { OutputHandle } from "./Handles";
@@ -23,6 +24,7 @@ export function TaskNodeOutputs({
   expanded,
   onBackgroundClick,
 }: TaskNodeOutputsProps) {
+  const { getOutputNodeId } = useNodeManager();
   const { nodeId, outputs, state, select } = useTaskNode();
   const {
     highlightSearchFilter,
@@ -40,7 +42,8 @@ export function TaskNodeOutputs({
     edges.some(
       (edge) =>
         edge.source === nodeId &&
-        edge.sourceHandle === outputNameToNodeId(output.name),
+        edge.sourceHandle ===
+          getOutputNodeId(outputNameToOutputId(output.name)),
     ),
   );
 
@@ -138,7 +141,7 @@ export function TaskNodeOutputs({
     }
 
     const output = outputs.find(
-      (o) => outputNameToNodeId(o.name) === fromHandle?.id,
+      (o) => getOutputNodeId(outputNameToOutputId(o.name)) === fromHandle?.id,
     );
 
     if (!output) return;

src/components/shared/ReactFlow/FlowCanvas/utils/addAndConnectNode.ts

@@ -1,18 +1,18 @@
 import type { Connection, Handle } from "@xyflow/react";
 
-import type {
-  ComponentReference,
-  ComponentSpec,
-  TaskSpec,
-  TypeSpecType,
+import type { NodeManager } from "@/nodeManager";
+import {
+  type ComponentReference,
+  type ComponentSpec,
+  isGraphImplementation,
+  type TaskSpec,
+  type TypeSpecType,
 } from "@/utils/componentSpec";
 import { DEFAULT_NODE_DIMENSIONS } from "@/utils/constants";
 import {
-  inputNameToNodeId,
-  nodeIdToTaskId,
-  outputNameToNodeId,
-  taskIdToNodeId,
-} from "@/utils/nodes/nodeIdUtils";
+  inputNameToInputId,
+  outputNameToOutputId,
+} from "@/utils/nodes/conversions";
 
 import addTask from "./addTask";
 import { handleConnection } from "./handleConnection";
@@ -22,21 +22,23 @@ type AddAndConnectNodeParams = {
   fromHandle: Handle | null;
   position: { x: number; y: number };
   componentSpec: ComponentSpec;
+  nodeManager: NodeManager;
 };
 
 export function addAndConnectNode({
   componentRef,
   fromHandle,
   position,
   componentSpec,
+  nodeManager,
 }: AddAndConnectNodeParams): ComponentSpec {
   // 1. Add the new node
   const taskSpec: TaskSpec = {
     annotations: {},
     componentRef: { ...componentRef },
   };
 
-  if (!("graph" in componentSpec.implementation)) {
+  if (!isGraphImplementation(componentSpec.implementation)) {
     return componentSpec;
   }
 
@@ -58,7 +60,7 @@ export function addAndConnectNode({
   );
 
   // 2. Find the new node
-  if (!("graph" in newComponentSpec.implementation)) {
+  if (!isGraphImplementation(newComponentSpec.implementation)) {
     return newComponentSpec;
   }
 
@@ -72,14 +74,17 @@ export function addAndConnectNode({
     return newComponentSpec;
   }
 
-  const newNodeId = taskIdToNodeId(newTaskId);
+  const newNodeId = nodeManager.getNodeId(newTaskId, "task");
 
   // 3. Determine the connection data type and find the first matching handle on the new node
   if (!fromHandle) {
     return newComponentSpec;
   }
 
-  const fromTaskId = nodeIdToTaskId(fromHandle.nodeId);
+  const fromTaskId = nodeManager.getTaskId(fromHandle.nodeId);
+  if (!fromTaskId) {
+    return newComponentSpec;
+  }
 
   const fromTaskSpec = graphSpec.tasks[fromTaskId];
   const fromComponentSpec = fromTaskSpec?.componentRef.spec;
@@ -112,7 +117,8 @@ export function addAndConnectNode({
       return newComponentSpec;
     }
 
-    targetHandleId = inputNameToNodeId(handleName);
+    const inputId = inputNameToInputId(handleName);
+    targetHandleId = nodeManager.getNodeId(inputId, "input");
   } else if (toHandleType === "output") {
     const handleName = componentRef.spec?.outputs?.find(
       (io) => io.type === connectionType,
@@ -121,7 +127,8 @@ export function addAndConnectNode({
       return newComponentSpec;
     }
 
-    targetHandleId = outputNameToNodeId(handleName);
+    const outputId = outputNameToOutputId(handleName);
+    targetHandleId = nodeManager.getNodeId(outputId, "output");
   }
 
   // 4. Build a Connection object and use handleConnection to add the edge
@@ -144,7 +151,11 @@ export function addAndConnectNode({
           targetHandle: targetHandleId,
         };
 
-    const updatedGraphSpec = handleConnection(graphSpec, connection);
+    const updatedGraphSpec = handleConnection(
+      graphSpec,
+      connection,
+      nodeManager,
+    );
 
     return {
       ...newComponentSpec,

src/components/shared/ReactFlow/FlowCanvas/utils/addTask.ts

@@ -9,6 +9,7 @@ import {
   type OutputSpec,
   type TaskSpec,
 } from "@/utils/componentSpec";
+import { taskNameToTaskId } from "@/utils/nodes/conversions";
 import {
   getUniqueInputName,
   getUniqueOutputName,
@@ -61,11 +62,13 @@ const addTask = (
       arguments: taskArguments ?? {},
     };
 
-    const taskId = getUniqueTaskName(
+    const uniqueTaskName = getUniqueTaskName(
       graphSpec,
       taskSpec.componentRef.spec?.name ?? "Task",
     );
 
+    const taskId = taskNameToTaskId(uniqueTaskName);
+
     const newGraphSpec: GraphSpec = {
       ...graphSpec,
       tasks: {
@@ -78,9 +81,9 @@ const addTask = (
   }
 
   if (taskType === "input") {
-    const inputId = getUniqueInputName(componentSpec);
+    const uniqueInputName = getUniqueInputName(componentSpec);
     const inputSpec: InputSpec = {
-      name: inputId,
+      name: uniqueInputName,
       annotations: positionAnnotations,
     };
     const inputs = (componentSpec.inputs ?? []).concat([inputSpec]);
@@ -89,9 +92,9 @@ const addTask = (
   }
 
   if (taskType === "output") {
-    const outputId = getUniqueOutputName(componentSpec);
+    const uniqueOutputName = getUniqueOutputName(componentSpec);
     const outputSpec: OutputSpec = {
-      name: outputId,
+      name: uniqueOutputName,
       annotations: positionAnnotations,
     };