Skip to content

Commit fe2444d

Browse files
camielvscamielvs
committed
Fix IO Node Copy + Paste
1 parent 802180e commit fe2444d

File tree

3 files changed

+138
-68
lines changed

3 files changed

+138
-68
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ const IONode = ({ type, data, selected = false }: IONodeProps) => {
8181
);
8282
}
8383
}
84-
}, [selected, readOnly, input, output, isInput, isOutput, graphSpec]);
84+
}, [selected, readOnly, input, output, isInput, graphSpec]);
8585

8686
useEffect(() => {
8787
return () => {

‎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: 101 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -72,12 +72,7 @@ export const duplicateNodes = (
7272
const oldNodeId = node.id;
7373

7474
if (isTaskNode(node)) {
75-
const oldTaskId = nodeManager.getRefId(oldNodeId);
76-
if (!oldTaskId) {
77-
console.warn("Could not find taskId for node:", node);
78-
return;
79-
}
80-
75+
const oldTaskId = node.data.taskId as string;
8176
const oldTaskName = taskIdToTaskName(oldTaskId);
8277
const newTaskName = getUniqueTaskName(graphSpec, oldTaskName);
8378
const newTaskId = taskNameToTaskId(newTaskName);
@@ -104,16 +99,14 @@ export const duplicateNodes = (
10499
};
105100
newTasks[newTaskId] = newTaskSpec;
106101
} else if (isInputNode(node)) {
107-
const inputSpec = componentSpec.inputs?.find(
108-
(input) => input.name === node.data.label,
109-
);
102+
const inputSpec = node.data.spec;
110103

111-
const newInputName = getUniqueInputName(componentSpec, inputSpec?.name);
104+
const newInputName = getUniqueInputName(componentSpec, inputSpec.name);
112105
const newNodeId = nodeManager.getNodeId(newInputName, "input");
113106

114107
nodeIdMap[oldNodeId] = newNodeId;
115108

116-
const annotations = inputSpec?.annotations || {};
109+
const annotations = inputSpec.annotations || {};
117110

118111
const updatedAnnotations = setPositionInAnnotations(annotations, {
119112
x: node.position.x + OFFSET,
@@ -128,19 +121,14 @@ export const duplicateNodes = (
128121

129122
newInputs[newInputName] = newInputSpec;
130123
} else if (isOutputNode(node)) {
131-
const outputSpec = componentSpec.outputs?.find(
132-
(output) => output.name === node.data.label,
133-
);
124+
const outputSpec = node.data.spec;
134125

135-
const newOutputName = getUniqueOutputName(
136-
componentSpec,
137-
outputSpec?.name,
138-
);
126+
const newOutputName = getUniqueOutputName(componentSpec, outputSpec.name);
139127
const newNodeId = nodeManager.getNodeId(newOutputName, "output");
140128

141129
nodeIdMap[oldNodeId] = newNodeId;
142130

143-
const annotations = outputSpec?.annotations || {};
131+
const annotations = outputSpec.annotations || {};
144132

145133
const updatedAnnotations = setPositionInAnnotations(annotations, {
146134
x: node.position.x + OFFSET,
@@ -195,13 +183,13 @@ export const duplicateNodes = (
195183
}
196184
});
197185

198-
// Outputs are defined in the graph spec
199186
const updatedGraphOutputs = { ...graphSpec.outputValues };
200187
if (connection !== "none") {
201-
/* Reconfigure Outputs */
188+
/* Reconfigure Output Nodes */
202189
Object.entries(newOutputs).forEach((output) => {
203190
const [outputName] = output;
204191
const newNodeId = nodeManager.getNodeId(outputName, "output");
192+
205193
const oldNodeId = Object.keys(nodeIdMap).find(
206194
(key) => nodeIdMap[key] === newNodeId,
207195
);
@@ -210,15 +198,52 @@ export const duplicateNodes = (
210198
return;
211199
}
212200

213-
const oldOutputName = nodeManager.getRefId(oldNodeId);
201+
const originalOutputNode = nodesToDuplicate.find(
202+
(node) => node.id === oldNodeId && isOutputNode(node),
203+
);
214204

215-
if (!graphSpec.outputValues || !oldOutputName) {
205+
if (!originalOutputNode) {
216206
return;
217207
}
218208

219-
const outputValue = graphSpec.outputValues[oldOutputName];
209+
const oldOutputName = (
210+
originalOutputNode.data.spec as OutputSpec
211+
).name.toLowerCase();
212+
213+
let outputValue: TaskOutputArgument | null = null;
214+
let connectedTaskId: string | null = null;
215+
let outputArgName: string | null = null;
216+
217+
for (const node of nodesToDuplicate) {
218+
if (isTaskNode(node)) {
219+
const taskData = node.data;
220+
const taskId = taskData.taskId;
221+
const taskOutputs =
222+
taskData.taskSpec.componentRef.spec?.outputs || [];
223+
224+
for (const taskOutput of taskOutputs) {
225+
if (taskOutput.name === oldOutputName) {
226+
connectedTaskId = taskId;
227+
outputArgName = taskOutput.name;
228+
229+
outputValue = {
230+
taskOutput: {
231+
taskId: connectedTaskId,
232+
outputName: outputArgName,
233+
},
234+
};
235+
break;
236+
}
237+
}
238+
239+
if (outputValue) break;
240+
}
241+
}
220242

221-
if (!outputValue) {
243+
if (!outputValue || !connectedTaskId) {
244+
console.warn(
245+
`No connecting task found for output ${oldOutputName} in duplicated nodes`,
246+
);
222247
return;
223248
}
224249

@@ -233,19 +258,26 @@ export const duplicateNodes = (
233258
) {
234259
if ("taskOutput" in updatedOutputValue) {
235260
const oldTaskId = updatedOutputValue.taskOutput.taskId;
236-
const oldTaskNodeId = nodeManager.getNodeId(oldTaskId, "task");
237-
if (oldTaskNodeId in nodeIdMap) {
238-
const newTaskId = nodeManager.getRefId(nodeIdMap[oldTaskNodeId]);
239-
if (!newTaskId) {
240-
return;
241-
}
242261

243-
updatedOutputValue.taskOutput = {
244-
...updatedOutputValue.taskOutput,
245-
taskId: newTaskId,
246-
};
262+
const taskNode = nodesToDuplicate.find(
263+
(node) => isTaskNode(node) && node.data.taskId === oldTaskId,
264+
);
247265

248-
isInternal = true;
266+
if (taskNode) {
267+
const newTaskNodeId = nodeIdMap[taskNode.id];
268+
if (newTaskNodeId) {
269+
const newTaskId = nodeManager.getRefId(newTaskNodeId);
270+
if (!newTaskId) {
271+
return;
272+
}
273+
274+
updatedOutputValue.taskOutput = {
275+
...updatedOutputValue.taskOutput,
276+
taskId: newTaskId,
277+
};
278+
279+
isInternal = true;
280+
}
249281
}
250282
}
251283
}
@@ -260,7 +292,7 @@ export const duplicateNodes = (
260292
});
261293
}
262294

263-
/* Update the Graph Spec & Inputs */
295+
/* Update the Graph Spec */
264296
const updatedTasks = { ...graphSpec.tasks, ...newTasks };
265297
const updatedGraphSpec = {
266298
...graphSpec,
@@ -516,43 +548,61 @@ function reconfigureConnections(
516548

517549
if ("taskOutput" in argument) {
518550
const oldTaskId = argument.taskOutput.taskId;
519-
oldNodeId = nodeManager.getNodeId(oldTaskId, "task");
520551

521-
if (!isGraphImplementation(componentSpec.implementation)) {
522-
throw new Error("ComponentSpec does not contain a graph implementation.");
552+
const taskNode = nodes.find(
553+
(node) => isTaskNode(node) && node.data.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;
523568
}
524569

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

528574
const newNodeId = nodeIdMap[oldNodeId];
529-
530575
if (!newNodeId) {
531576
return reconfigureExternalConnection(taskSpec, argKey, mode);
532577
}
533578

534579
const newTaskId = nodeManager.getRefId(newNodeId);
535-
536580
newArgId = newTaskId;
537581
} else if ("graphInput" in argument) {
538582
const oldInputName = argument.graphInput.inputName;
539-
oldNodeId = nodeManager.getNodeId(oldInputName, "input");
540583

541-
if (!("inputs" in componentSpec)) {
542-
throw new Error("ComponentSpec does not contain inputs.");
584+
const inputNode = nodes.find(
585+
(node) => isInputNode(node) && node.data.spec.name === oldInputName,
586+
);
587+
588+
if (inputNode) {
589+
oldNodeId = inputNode.id;
590+
isExternal = false;
591+
} else {
592+
const inputs = componentSpec.inputs || [];
593+
isExternal = inputs.some((input) => input.name === oldInputName);
543594
}
544595

545-
const inputs = componentSpec.inputs || [];
546-
isExternal = inputs.some((input) => input.name === oldInputName);
596+
if (!oldNodeId) {
597+
return reconfigureExternalConnection(taskSpec, argKey, mode);
598+
}
547599

548600
const newNodeId = nodeIdMap[oldNodeId];
549-
550601
if (!newNodeId) {
551602
return reconfigureExternalConnection(taskSpec, argKey, mode);
552603
}
553604

554605
const newInputName = nodeManager.getRefId(newNodeId);
555-
556606
newArgId = newInputName;
557607
}
558608

0 commit comments

Comments
 (0)