Fix IO Node Copy + Paste

Author: camielvs

src/components/shared/ReactFlow/FlowCanvas/utils/duplicateNodes.test.ts

@@ -3,12 +3,13 @@ import { describe, expect, it, vi } from "vitest";
 
 import { NodeManager } from "@/nodeManager";
 import type { TaskNodeData } from "@/types/nodes";
-import type {
-  ComponentSpec,
-  InputSpec,
-  OutputSpec,
-  TaskOutputArgument,
-  TaskSpec,
+import {
+  type ComponentSpec,
+  type InputSpec,
+  isGraphImplementation,
+  type OutputSpec,
+  type TaskOutputArgument,
+  type TaskSpec,
 } from "@/utils/componentSpec";
 
 import { duplicateNodes } from "./duplicateNodes";
@@ -63,7 +64,7 @@ const createMockComponentSpecWithOutputs = (
           acc[output.name] = {
             taskOutput: {
               taskId: "task1",
-              outputName: "result",
+              outputName: output.name,
             },
           };
           return acc;
@@ -123,7 +124,7 @@ const createMockInputNode = (
     position,
     data: {
       label: inputName,
-      inputSpec: { ...mockInputSpec, name: inputName },
+      spec: { ...mockInputSpec, name: inputName },
     },
     selected: false,
     dragging: false,
@@ -144,7 +145,7 @@ const createMockOutputNode = (
     position,
     data: {
       label: outputName,
-      outputSpec: { ...mockOutputSpec, name: outputName },
+      spec: { ...mockOutputSpec, name: outputName },
     },
     selected: false,
     dragging: false,
@@ -606,7 +607,7 @@ describe("duplicateNodes", () => {
         connection: "all",
       });
 
-      if ("graph" in result.updatedComponentSpec.implementation!) {
+      if (isGraphImplementation(result.updatedComponentSpec.implementation)) {
         const duplicatedTask =
           result.updatedComponentSpec.implementation.graph.tasks["task1 2"];
         expect(duplicatedTask.arguments?.input1).toEqual({
@@ -618,16 +619,35 @@ describe("duplicateNodes", () => {
     });
 
     it("should handle graph output connections", () => {
-      const taskSpec: TaskSpec = {
-        ...mockTaskSpec,
-        arguments: {},
-      };
-
       const outputSpec: OutputSpec = {
         ...mockOutputSpec,
         name: "graph-output",
       };
 
+      const taskComponentSpec: ComponentSpec = {
+        name: "task-component",
+        inputs: [],
+        outputs: [
+          {
+            name: "graph-output",
+            type: "String",
+            annotations: {},
+          },
+        ],
+        implementation: {
+          container: { image: "task-image" },
+        },
+      };
+
+      const taskSpec: TaskSpec = {
+        ...mockTaskSpec,
+        arguments: {},
+        componentRef: {
+          name: "task-component",
+          spec: taskComponentSpec,
+        },
+      };
+
       const componentSpec = createMockComponentSpecWithOutputs(
         { task1: taskSpec },
         [],
@@ -653,7 +673,7 @@ describe("duplicateNodes", () => {
         expect(outputValues?.["graph-output 2"]).toEqual({
           taskOutput: {
             taskId: "task1 2",
-            outputName: "result",
+            outputName: "graph-output",
           },
         });
       }

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

@@ -71,12 +71,7 @@ export const duplicateNodes = (
     const oldNodeId = node.id;
 
     if (isTaskNode(node)) {
-      const oldTaskId = nodeManager.getRefId(oldNodeId);
-      if (!oldTaskId) {
-        console.warn("Could not find taskId for node:", node);
-        return;
-      }
-
+      const oldTaskId = node.data.taskId;
       const newTaskId = getUniqueTaskId(graphSpec, oldTaskId);
       const newNodeId = nodeManager.getNodeId(newTaskId, "task");
 
@@ -101,16 +96,14 @@ export const duplicateNodes = (
       };
       newTasks[newTaskId] = newTaskSpec;
     } else if (isInputNode(node)) {
-      const inputSpec = componentSpec.inputs?.find(
-        (input) => input.name === node.data.label,
-      );
+      const inputSpec = node.data.spec;
 
-      const newInputName = getUniqueInputName(componentSpec, inputSpec?.name);
+      const newInputName = getUniqueInputName(componentSpec, inputSpec.name);
       const newNodeId = nodeManager.getNodeId(newInputName, "input");
 
       nodeIdMap[oldNodeId] = newNodeId;
 
-      const annotations = inputSpec?.annotations || {};
+      const annotations = inputSpec.annotations || {};
 
       const updatedAnnotations = setPositionInAnnotations(annotations, {
         x: node.position.x + OFFSET,
@@ -125,19 +118,14 @@ export const duplicateNodes = (
 
       newInputs[newInputName] = newInputSpec;
     } else if (isOutputNode(node)) {
-      const outputSpec = componentSpec.outputs?.find(
-        (output) => output.name === node.data.label,
-      );
+      const outputSpec = node.data.spec;
 
-      const newOutputName = getUniqueOutputName(
-        componentSpec,
-        outputSpec?.name,
-      );
+      const newOutputName = getUniqueOutputName(componentSpec, outputSpec.name);
       const newNodeId = nodeManager.getNodeId(newOutputName, "output");
 
       nodeIdMap[oldNodeId] = newNodeId;
 
-      const annotations = outputSpec?.annotations || {};
+      const annotations = outputSpec.annotations || {};
 
       const updatedAnnotations = setPositionInAnnotations(annotations, {
         x: node.position.x + OFFSET,
@@ -192,10 +180,11 @@ export const duplicateNodes = (
     }
   });
 
-  // Outputs are defined in the graph spec
   const updatedGraphOutputs = { ...graphSpec.outputValues };
-  if (connection !== "none") {
-    /* Reconfigure Outputs */
+
+  // Output Nodes can only have one input, so the "external" connection mode does not apply to them.
+  if (connection !== "none" && connection !== "external") {
+    /* Reconfigure Output Nodes */
     Object.entries(newOutputs).forEach((output) => {
       const [outputName] = output;
       const newNodeId = nodeManager.getNodeId(outputName, "output");
@@ -207,57 +196,62 @@ export const duplicateNodes = (
         return;
       }
 
-      const oldOutputName = nodeManager.getRefId(oldNodeId);
+      const originalOutputNode = nodesToDuplicate.find(
+        (node) => node.id === oldNodeId && isOutputNode(node),
+      );
 
-      if (!graphSpec.outputValues || !oldOutputName) {
+      if (!originalOutputNode || !isOutputNode(originalOutputNode)) {
         return;
       }
 
-      const outputValue = graphSpec.outputValues[oldOutputName];
+      const oldOutputName = originalOutputNode.data.spec.name;
+
+      const originalOutputValue = graphSpec.outputValues?.[oldOutputName];
 
-      if (!outputValue) {
+      if (!originalOutputValue) {
+        // Todo: Handle cross-instance copy + paste for output nodes (we don't have the original graphSpec available so we can't look up the output value)
+        console.warn(
+          `No output value found for output ${oldOutputName} in graph spec.`,
+        );
         return;
       }
 
-      const updatedOutputValue = { ...outputValue };
-
-      // If the outputvalue references a task that was also duplicated (internal connection), we need to update it to refer to the duplicated task id
-      let isInternal = false;
-      if (
-        typeof updatedOutputValue === "object" &&
-        updatedOutputValue !== null &&
-        connection !== "external"
-      ) {
-        if ("taskOutput" in updatedOutputValue) {
-          const oldTaskId = updatedOutputValue.taskOutput.taskId;
-          const oldTaskNodeId = nodeManager.getNodeId(oldTaskId, "task");
-          if (oldTaskNodeId in nodeIdMap) {
-            const newTaskId = nodeManager.getRefId(nodeIdMap[oldTaskNodeId]);
-            if (!newTaskId) {
-              return;
-            }
-
-            updatedOutputValue.taskOutput = {
-              ...updatedOutputValue.taskOutput,
-              taskId: newTaskId,
-            };
-
-            isInternal = true;
-          }
-        }
+      if (!("taskOutput" in originalOutputValue)) {
+        console.warn(
+          `Output ${oldOutputName} is not connected to a task output`,
+        );
+        return;
       }
 
-      if (
-        (isInternal && connection === "internal") ||
-        (!isInternal && connection === "external") ||
-        connection === "all"
-      ) {
-        updatedGraphOutputs[outputName] = updatedOutputValue;
+      const connectedTaskId = originalOutputValue.taskOutput.taskId;
+      const connectedOutputName = originalOutputValue.taskOutput.outputName;
+
+      const originalNodeId = nodeManager.getNodeId(connectedTaskId, "task");
+
+      const newTaskNodeId = nodeIdMap[originalNodeId];
+      if (!newTaskNodeId) {
+        console.warn(`No mapping found for task node ${originalNodeId}`);
+        return;
       }
+
+      const newTaskId = nodeManager.getRefId(newTaskNodeId);
+      if (!newTaskId) {
+        console.warn(`Could not get new task ID for node ${newTaskNodeId}`);
+        return;
+      }
+
+      const outputValue: TaskOutputArgument = {
+        taskOutput: {
+          taskId: newTaskId,
+          outputName: connectedOutputName,
+        },
+      };
+
+      updatedGraphOutputs[outputName] = outputValue;
     });
   }
 
-  /* Update the Graph Spec & Inputs */
+  /* Update the Graph Spec */
   const updatedTasks = { ...graphSpec.tasks, ...newTasks };
   const updatedGraphSpec = {
     ...graphSpec,
@@ -513,43 +507,61 @@ function reconfigureConnections(
 
   if ("taskOutput" in argument) {
     const oldTaskId = argument.taskOutput.taskId;
-    oldNodeId = nodeManager.getNodeId(oldTaskId, "task");
 
-    if (!isGraphImplementation(componentSpec.implementation)) {
-      throw new Error("ComponentSpec does not contain a graph implementation.");
+    const taskNode = nodes.find(
+      (node) => isTaskNode(node) && node.data.taskId === oldTaskId,
+    );
+
+    if (taskNode) {
+      oldNodeId = taskNode.id;
+      isExternal = false;
+    } else {
+      if (!isGraphImplementation(componentSpec.implementation)) {
+        throw new Error(
+          "ComponentSpec does not contain a graph implementation.",
+        );
+      }
+
+      const graphSpec = componentSpec.implementation.graph;
+      isExternal = oldTaskId in graphSpec.tasks;
     }
 
-    const graphSpec = componentSpec.implementation.graph;
-    isExternal = oldTaskId in graphSpec.tasks;
+    if (!oldNodeId) {
+      return reconfigureExternalConnection(taskSpec, argKey, mode);
+    }
 
     const newNodeId = nodeIdMap[oldNodeId];
-
     if (!newNodeId) {
       return reconfigureExternalConnection(taskSpec, argKey, mode);
     }
 
     const newTaskId = nodeManager.getRefId(newNodeId);
-
     newArgId = newTaskId;
   } else if ("graphInput" in argument) {
     const oldInputName = argument.graphInput.inputName;
-    oldNodeId = nodeManager.getNodeId(oldInputName, "input");
 
-    if (!("inputs" in componentSpec)) {
-      throw new Error("ComponentSpec does not contain inputs.");
+    const inputNode = nodes.find(
+      (node) => isInputNode(node) && node.data.spec.name === oldInputName,
+    );
+
+    if (inputNode) {
+      oldNodeId = inputNode.id;
+      isExternal = false;
+    } else {
+      const inputs = componentSpec.inputs || [];
+      isExternal = inputs.some((input) => input.name === oldInputName);
     }
 
-    const inputs = componentSpec.inputs || [];
-    isExternal = inputs.some((input) => input.name === oldInputName);
+    if (!oldNodeId) {
+      return reconfigureExternalConnection(taskSpec, argKey, mode);
+    }
 
     const newNodeId = nodeIdMap[oldNodeId];
-
     if (!newNodeId) {
       return reconfigureExternalConnection(taskSpec, argKey, mode);
     }
 
     const newInputName = nodeManager.getRefId(newNodeId);
-
     newArgId = newInputName;
   }
 

src/utils/unique.ts

@@ -4,11 +4,12 @@ const makeNameUniqueByAddingIndex = (
   name: string,
   existingNames: Set<string>,
 ): string => {
-  let finalName = name;
+  const baseName = name.replace(/ \(\d+\)$/, "");
+  let finalName = baseName;
   let index = 1;
   while (existingNames.has(finalName)) {
     index++;
-    finalName = name + " " + index.toString();
+    finalName = baseName + " (" + index.toString() + ")";
   }
   return finalName;
 };