Skip to content

Commit d874723

Browse files
camielvscamielvs
committed
Fix IO Node Copy + Paste
1 parent 3d8a4ff commit d874723

File tree

2 files changed

+143
-67
lines changed

2 files changed

+143
-67
lines changed

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

Lines changed: 36 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,13 @@ import { describe, expect, it, vi } from "vitest";
33

44
import { NodeManager } from "@/nodeManager";
55
import type { TaskNodeData } from "@/types/nodes";
6-
import type {
7-
ComponentSpec,
8-
InputSpec,
9-
OutputSpec,
10-
TaskOutputArgument,
11-
TaskSpec,
6+
import {
7+
type ComponentSpec,
8+
type InputSpec,
9+
isGraphImplementation,
10+
type OutputSpec,
11+
type TaskOutputArgument,
12+
type TaskSpec,
1213
} from "@/utils/componentSpec";
1314

1415
import { duplicateNodes } from "./duplicateNodes";
@@ -63,7 +64,7 @@ const createMockComponentSpecWithOutputs = (
6364
acc[output.name] = {
6465
taskOutput: {
6566
taskId: "task1",
66-
outputName: "result",
67+
outputName: output.name,
6768
},
6869
};
6970
return acc;
@@ -123,7 +124,7 @@ const createMockInputNode = (
123124
position,
124125
data: {
125126
label: inputName,
126-
inputSpec: { ...mockInputSpec, name: inputName },
127+
spec: { ...mockInputSpec, name: inputName },
127128
},
128129
selected: false,
129130
dragging: false,
@@ -144,7 +145,7 @@ const createMockOutputNode = (
144145
position,
145146
data: {
146147
label: outputName,
147-
outputSpec: { ...mockOutputSpec, name: outputName },
148+
spec: { ...mockOutputSpec, name: outputName },
148149
},
149150
selected: false,
150151
dragging: false,
@@ -606,7 +607,7 @@ describe("duplicateNodes", () => {
606607
connection: "all",
607608
});
608609

609-
if ("graph" in result.updatedComponentSpec.implementation!) {
610+
if (isGraphImplementation(result.updatedComponentSpec.implementation)) {
610611
const duplicatedTask =
611612
result.updatedComponentSpec.implementation.graph.tasks["task1 2"];
612613
expect(duplicatedTask.arguments?.input1).toEqual({
@@ -618,16 +619,35 @@ describe("duplicateNodes", () => {
618619
});
619620

620621
it("should handle graph output connections", () => {
621-
const taskSpec: TaskSpec = {
622-
...mockTaskSpec,
623-
arguments: {},
624-
};
625-
626622
const outputSpec: OutputSpec = {
627623
...mockOutputSpec,
628624
name: "graph-output",
629625
};
630626

627+
const taskComponentSpec: ComponentSpec = {
628+
name: "task-component",
629+
inputs: [],
630+
outputs: [
631+
{
632+
name: "graph-output",
633+
type: "String",
634+
annotations: {},
635+
},
636+
],
637+
implementation: {
638+
container: { image: "task-image" },
639+
},
640+
};
641+
642+
const taskSpec: TaskSpec = {
643+
...mockTaskSpec,
644+
arguments: {},
645+
componentRef: {
646+
name: "task-component",
647+
spec: taskComponentSpec,
648+
},
649+
};
650+
631651
const componentSpec = createMockComponentSpecWithOutputs(
632652
{ task1: taskSpec },
633653
[],
@@ -653,7 +673,7 @@ describe("duplicateNodes", () => {
653673
expect(outputValues?.["graph-output 2"]).toEqual({
654674
taskOutput: {
655675
taskId: "task1 2",
656-
outputName: "result",
676+
outputName: "graph-output",
657677
},
658678
});
659679
}

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

Lines changed: 107 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -67,12 +67,7 @@ export const duplicateNodes = (
6767
const oldNodeId = node.id;
6868

6969
if (node.type === "task") {
70-
const oldTaskId = nodeManager.getRefId(oldNodeId);
71-
if (!oldTaskId) {
72-
console.warn("Could not find taskId for node:", node);
73-
return;
74-
}
75-
70+
const oldTaskId = node.data.taskId as string;
7671
const oldTaskName = taskIdToTaskName(oldTaskId);
7772
const newTaskName = getUniqueTaskName(graphSpec, oldTaskName);
7873
const newTaskId = taskNameToTaskId(newTaskName);
@@ -99,16 +94,14 @@ export const duplicateNodes = (
9994
};
10095
newTasks[newTaskId] = newTaskSpec;
10196
} else if (node.type === "input") {
102-
const inputSpec = componentSpec.inputs?.find(
103-
(input) => input.name === node.data.label,
104-
);
97+
const inputSpec = node.data.spec as InputSpec;
10598

106-
const newInputName = getUniqueInputName(componentSpec, inputSpec?.name);
99+
const newInputName = getUniqueInputName(componentSpec, inputSpec.name);
107100
const newNodeId = nodeManager.getNodeId(newInputName, "input");
108101

109102
nodeIdMap[oldNodeId] = newNodeId;
110103

111-
const annotations = inputSpec?.annotations || {};
104+
const annotations = inputSpec.annotations || {};
112105

113106
const updatedAnnotations = setPositionInAnnotations(annotations, {
114107
x: node.position.x + OFFSET,
@@ -123,19 +116,14 @@ export const duplicateNodes = (
123116

124117
newInputs[newInputName] = newInputSpec;
125118
} else if (node.type === "output") {
126-
const outputSpec = componentSpec.outputs?.find(
127-
(output) => output.name === node.data.label,
128-
);
119+
const outputSpec = node.data.spec as OutputSpec;
129120

130-
const newOutputName = getUniqueOutputName(
131-
componentSpec,
132-
outputSpec?.name,
133-
);
121+
const newOutputName = getUniqueOutputName(componentSpec, outputSpec.name);
134122
const newNodeId = nodeManager.getNodeId(newOutputName, "output");
135123

136124
nodeIdMap[oldNodeId] = newNodeId;
137125

138-
const annotations = outputSpec?.annotations || {};
126+
const annotations = outputSpec.annotations || {};
139127

140128
const updatedAnnotations = setPositionInAnnotations(annotations, {
141129
x: node.position.x + OFFSET,
@@ -190,13 +178,13 @@ export const duplicateNodes = (
190178
}
191179
});
192180

193-
// Outputs are defined in the graph spec
194181
const updatedGraphOutputs = { ...graphSpec.outputValues };
195182
if (connection !== "none") {
196-
/* Reconfigure Outputs */
183+
/* Reconfigure Output Nodes */
197184
Object.entries(newOutputs).forEach((output) => {
198185
const [outputName] = output;
199186
const newNodeId = nodeManager.getNodeId(outputName, "output");
187+
200188
const oldNodeId = Object.keys(nodeIdMap).find(
201189
(key) => nodeIdMap[key] === newNodeId,
202190
);
@@ -205,15 +193,52 @@ export const duplicateNodes = (
205193
return;
206194
}
207195

208-
const oldOutputName = nodeManager.getRefId(oldNodeId);
196+
const originalOutputNode = nodesToDuplicate.find(
197+
(node) => node.id === oldNodeId && node.type === "output",
198+
);
209199

210-
if (!graphSpec.outputValues || !oldOutputName) {
200+
if (!originalOutputNode) {
211201
return;
212202
}
213203

214-
const outputValue = graphSpec.outputValues[oldOutputName];
204+
const oldOutputName = (
205+
originalOutputNode.data.spec as OutputSpec
206+
).name.toLowerCase();
207+
208+
let outputValue: TaskOutputArgument | null = null;
209+
let connectedTaskId: string | null = null;
210+
let outputArgName: string | null = null;
211+
212+
for (const node of nodesToDuplicate) {
213+
if (node.type === "task") {
214+
const taskData = node.data as TaskNodeData;
215+
const taskId = taskData.taskId;
216+
const taskOutputs =
217+
taskData.taskSpec.componentRef.spec?.outputs || [];
218+
219+
for (const taskOutput of taskOutputs) {
220+
if (taskOutput.name === oldOutputName) {
221+
connectedTaskId = taskId;
222+
outputArgName = taskOutput.name;
223+
224+
outputValue = {
225+
taskOutput: {
226+
taskId: connectedTaskId,
227+
outputName: outputArgName,
228+
},
229+
};
230+
break;
231+
}
232+
}
215233

216-
if (!outputValue) {
234+
if (outputValue) break;
235+
}
236+
}
237+
238+
if (!outputValue || !connectedTaskId) {
239+
console.warn(
240+
`No connecting task found for output ${oldOutputName} in duplicated nodes`,
241+
);
217242
return;
218243
}
219244

@@ -228,19 +253,28 @@ export const duplicateNodes = (
228253
) {
229254
if ("taskOutput" in updatedOutputValue) {
230255
const oldTaskId = updatedOutputValue.taskOutput.taskId;
231-
const oldTaskNodeId = nodeManager.getNodeId(oldTaskId, "task");
232-
if (oldTaskNodeId in nodeIdMap) {
233-
const newTaskId = nodeManager.getRefId(nodeIdMap[oldTaskNodeId]);
234-
if (!newTaskId) {
235-
return;
236-
}
237256

238-
updatedOutputValue.taskOutput = {
239-
...updatedOutputValue.taskOutput,
240-
taskId: newTaskId,
241-
};
257+
const taskNode = nodesToDuplicate.find(
258+
(node) =>
259+
node.type === "task" &&
260+
(node.data as TaskNodeData).taskId === oldTaskId,
261+
);
262+
263+
if (taskNode) {
264+
const newTaskNodeId = nodeIdMap[taskNode.id];
265+
if (newTaskNodeId) {
266+
const newTaskId = nodeManager.getRefId(newTaskNodeId);
267+
if (!newTaskId) {
268+
return;
269+
}
242270

243-
isInternal = true;
271+
updatedOutputValue.taskOutput = {
272+
...updatedOutputValue.taskOutput,
273+
taskId: newTaskId,
274+
};
275+
276+
isInternal = true;
277+
}
244278
}
245279
}
246280
}
@@ -255,7 +289,7 @@ export const duplicateNodes = (
255289
});
256290
}
257291

258-
/* Update the Graph Spec & Inputs */
292+
/* Update the Graph Spec */
259293
const updatedTasks = { ...graphSpec.tasks, ...newTasks };
260294
const updatedGraphSpec = {
261295
...graphSpec,
@@ -512,43 +546,65 @@ function reconfigureConnections(
512546

513547
if ("taskOutput" in argument) {
514548
const oldTaskId = argument.taskOutput.taskId;
515-
oldNodeId = nodeManager.getNodeId(oldTaskId, "task");
516549

517-
if (!isGraphImplementation(componentSpec.implementation)) {
518-
throw new Error("ComponentSpec does not contain a graph implementation.");
550+
const taskNode = nodes.find(
551+
(node) =>
552+
node.type === "task" &&
553+
(node.data as TaskNodeData).taskId === oldTaskId,
554+
);
555+
556+
if (taskNode) {
557+
oldNodeId = taskNode.id;
558+
isExternal = false;
559+
} else {
560+
if (!isGraphImplementation(componentSpec.implementation)) {
561+
throw new Error(
562+
"ComponentSpec does not contain a graph implementation.",
563+
);
564+
}
565+
566+
const graphSpec = componentSpec.implementation.graph;
567+
isExternal = oldTaskId in graphSpec.tasks;
519568
}
520569

521-
const graphSpec = componentSpec.implementation.graph;
522-
isExternal = oldTaskId in graphSpec.tasks;
570+
if (!oldNodeId) {
571+
return reconfigureExternalConnection(taskSpec, argKey, mode);
572+
}
523573

524574
const newNodeId = nodeIdMap[oldNodeId];
525-
526575
if (!newNodeId) {
527576
return reconfigureExternalConnection(taskSpec, argKey, mode);
528577
}
529578

530579
const newTaskId = nodeManager.getRefId(newNodeId);
531-
532580
newArgId = newTaskId;
533581
} else if ("graphInput" in argument) {
534582
const oldInputName = argument.graphInput.inputName;
535-
oldNodeId = nodeManager.getNodeId(oldInputName, "input");
536583

537-
if (!("inputs" in componentSpec)) {
538-
throw new Error("ComponentSpec does not contain inputs.");
584+
const inputNode = nodes.find(
585+
(node) =>
586+
node.type === "input" &&
587+
(node.data.spec as InputSpec).name === oldInputName,
588+
);
589+
590+
if (inputNode) {
591+
oldNodeId = inputNode.id;
592+
isExternal = false;
593+
} else {
594+
const inputs = componentSpec.inputs || [];
595+
isExternal = inputs.some((input) => input.name === oldInputName);
539596
}
540597

541-
const inputs = componentSpec.inputs || [];
542-
isExternal = inputs.some((input) => input.name === oldInputName);
598+
if (!oldNodeId) {
599+
return reconfigureExternalConnection(taskSpec, argKey, mode);
600+
}
543601

544602
const newNodeId = nodeIdMap[oldNodeId];
545-
546603
if (!newNodeId) {
547604
return reconfigureExternalConnection(taskSpec, argKey, mode);
548605
}
549606

550607
const newInputName = nodeManager.getRefId(newNodeId);
551-
552608
newArgId = newInputName;
553609
}
554610

0 commit comments

Comments
 (0)