Skip to content

Commit f31471f

Browse files
camielvscamielvs
committed
Fix IO Node Copy + Paste
1 parent ab05c72 commit f31471f

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
@@ -71,12 +71,7 @@ export const duplicateNodes = (
7171
const oldNodeId = node.id;
7272

7373
if (isTaskNode(node)) {
74-
const oldTaskId = nodeManager.getRefId(oldNodeId);
75-
if (!oldTaskId) {
76-
console.warn("Could not find taskId for node:", node);
77-
return;
78-
}
79-
74+
const oldTaskId = node.data.taskId;
8075
const newTaskId = getUniqueTaskId(graphSpec, oldTaskId);
8176
const newNodeId = nodeManager.getNodeId(newTaskId, "task");
8277

@@ -101,16 +96,14 @@ export const duplicateNodes = (
10196
};
10297
newTasks[newTaskId] = newTaskSpec;
10398
} else if (isInputNode(node)) {
104-
const inputSpec = componentSpec.inputs?.find(
105-
(input) => input.name === node.data.label,
106-
);
99+
const inputSpec = node.data.spec;
107100

108-
const newInputName = getUniqueInputName(componentSpec, inputSpec?.name);
101+
const newInputName = getUniqueInputName(componentSpec, inputSpec.name);
109102
const newNodeId = nodeManager.getNodeId(newInputName, "input");
110103

111104
nodeIdMap[oldNodeId] = newNodeId;
112105

113-
const annotations = inputSpec?.annotations || {};
106+
const annotations = inputSpec.annotations || {};
114107

115108
const updatedAnnotations = setPositionInAnnotations(annotations, {
116109
x: node.position.x + OFFSET,
@@ -125,19 +118,14 @@ export const duplicateNodes = (
125118

126119
newInputs[newInputName] = newInputSpec;
127120
} else if (isOutputNode(node)) {
128-
const outputSpec = componentSpec.outputs?.find(
129-
(output) => output.name === node.data.label,
130-
);
121+
const outputSpec = node.data.spec;
131122

132-
const newOutputName = getUniqueOutputName(
133-
componentSpec,
134-
outputSpec?.name,
135-
);
123+
const newOutputName = getUniqueOutputName(componentSpec, outputSpec.name);
136124
const newNodeId = nodeManager.getNodeId(newOutputName, "output");
137125

138126
nodeIdMap[oldNodeId] = newNodeId;
139127

140-
const annotations = outputSpec?.annotations || {};
128+
const annotations = outputSpec.annotations || {};
141129

142130
const updatedAnnotations = setPositionInAnnotations(annotations, {
143131
x: node.position.x + OFFSET,
@@ -192,13 +180,13 @@ export const duplicateNodes = (
192180
}
193181
});
194182

195-
// Outputs are defined in the graph spec
196183
const updatedGraphOutputs = { ...graphSpec.outputValues };
197184
if (connection !== "none") {
198-
/* Reconfigure Outputs */
185+
/* Reconfigure Output Nodes */
199186
Object.entries(newOutputs).forEach((output) => {
200187
const [outputName] = output;
201188
const newNodeId = nodeManager.getNodeId(outputName, "output");
189+
202190
const oldNodeId = Object.keys(nodeIdMap).find(
203191
(key) => nodeIdMap[key] === newNodeId,
204192
);
@@ -207,15 +195,52 @@ export const duplicateNodes = (
207195
return;
208196
}
209197

210-
const oldOutputName = nodeManager.getRefId(oldNodeId);
198+
const originalOutputNode = nodesToDuplicate.find(
199+
(node) => node.id === oldNodeId && isOutputNode(node),
200+
);
211201

212-
if (!graphSpec.outputValues || !oldOutputName) {
202+
if (!originalOutputNode) {
213203
return;
214204
}
215205

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

218-
if (!outputValue) {
240+
if (!outputValue || !connectedTaskId) {
241+
console.warn(
242+
`No connecting task found for output ${oldOutputName} in duplicated nodes`,
243+
);
219244
return;
220245
}
221246

@@ -230,19 +255,26 @@ export const duplicateNodes = (
230255
) {
231256
if ("taskOutput" in updatedOutputValue) {
232257
const oldTaskId = updatedOutputValue.taskOutput.taskId;
233-
const oldTaskNodeId = nodeManager.getNodeId(oldTaskId, "task");
234-
if (oldTaskNodeId in nodeIdMap) {
235-
const newTaskId = nodeManager.getRefId(nodeIdMap[oldTaskNodeId]);
236-
if (!newTaskId) {
237-
return;
238-
}
239258

240-
updatedOutputValue.taskOutput = {
241-
...updatedOutputValue.taskOutput,
242-
taskId: newTaskId,
243-
};
259+
const taskNode = nodesToDuplicate.find(
260+
(node) => isTaskNode(node) && node.data.taskId === oldTaskId,
261+
);
244262

245-
isInternal = true;
263+
if (taskNode) {
264+
const newTaskNodeId = nodeIdMap[taskNode.id];
265+
if (newTaskNodeId) {
266+
const newTaskId = nodeManager.getRefId(newTaskNodeId);
267+
if (!newTaskId) {
268+
return;
269+
}
270+
271+
updatedOutputValue.taskOutput = {
272+
...updatedOutputValue.taskOutput,
273+
taskId: newTaskId,
274+
};
275+
276+
isInternal = true;
277+
}
246278
}
247279
}
248280
}
@@ -257,7 +289,7 @@ export const duplicateNodes = (
257289
});
258290
}
259291

260-
/* Update the Graph Spec & Inputs */
292+
/* Update the Graph Spec */
261293
const updatedTasks = { ...graphSpec.tasks, ...newTasks };
262294
const updatedGraphSpec = {
263295
...graphSpec,
@@ -513,43 +545,61 @@ function reconfigureConnections(
513545

514546
if ("taskOutput" in argument) {
515547
const oldTaskId = argument.taskOutput.taskId;
516-
oldNodeId = nodeManager.getNodeId(oldTaskId, "task");
517548

518-
if (!isGraphImplementation(componentSpec.implementation)) {
519-
throw new Error("ComponentSpec does not contain a graph implementation.");
549+
const taskNode = nodes.find(
550+
(node) => isTaskNode(node) && node.data.taskId === oldTaskId,
551+
);
552+
553+
if (taskNode) {
554+
oldNodeId = taskNode.id;
555+
isExternal = false;
556+
} else {
557+
if (!isGraphImplementation(componentSpec.implementation)) {
558+
throw new Error(
559+
"ComponentSpec does not contain a graph implementation.",
560+
);
561+
}
562+
563+
const graphSpec = componentSpec.implementation.graph;
564+
isExternal = oldTaskId in graphSpec.tasks;
520565
}
521566

522-
const graphSpec = componentSpec.implementation.graph;
523-
isExternal = oldTaskId in graphSpec.tasks;
567+
if (!oldNodeId) {
568+
return reconfigureExternalConnection(taskSpec, argKey, mode);
569+
}
524570

525571
const newNodeId = nodeIdMap[oldNodeId];
526-
527572
if (!newNodeId) {
528573
return reconfigureExternalConnection(taskSpec, argKey, mode);
529574
}
530575

531576
const newTaskId = nodeManager.getRefId(newNodeId);
532-
533577
newArgId = newTaskId;
534578
} else if ("graphInput" in argument) {
535579
const oldInputName = argument.graphInput.inputName;
536-
oldNodeId = nodeManager.getNodeId(oldInputName, "input");
537580

538-
if (!("inputs" in componentSpec)) {
539-
throw new Error("ComponentSpec does not contain inputs.");
581+
const inputNode = nodes.find(
582+
(node) => isInputNode(node) && node.data.spec.name === oldInputName,
583+
);
584+
585+
if (inputNode) {
586+
oldNodeId = inputNode.id;
587+
isExternal = false;
588+
} else {
589+
const inputs = componentSpec.inputs || [];
590+
isExternal = inputs.some((input) => input.name === oldInputName);
540591
}
541592

542-
const inputs = componentSpec.inputs || [];
543-
isExternal = inputs.some((input) => input.name === oldInputName);
593+
if (!oldNodeId) {
594+
return reconfigureExternalConnection(taskSpec, argKey, mode);
595+
}
544596

545597
const newNodeId = nodeIdMap[oldNodeId];
546-
547598
if (!newNodeId) {
548599
return reconfigureExternalConnection(taskSpec, argKey, mode);
549600
}
550601

551602
const newInputName = nodeManager.getRefId(newNodeId);
552-
553603
newArgId = newInputName;
554604
}
555605

0 commit comments

Comments
 (0)