Skip to content

Commit cf466ba

Browse files
committed
Fix IO Node Copy + Paste
1 parent c596d45 commit cf466ba

File tree

3 files changed

+124
-91
lines changed

3 files changed

+124
-91
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,
@@ -570,7 +571,7 @@ describe("duplicateNodes", () => {
570571
connection: "all",
571572
});
572573

573-
if ("graph" in result.updatedComponentSpec.implementation!) {
574+
if (isGraphImplementation(result.updatedComponentSpec.implementation)) {
574575
const duplicatedTask =
575576
result.updatedComponentSpec.implementation.graph.tasks["task1 2"];
576577
expect(duplicatedTask.arguments?.input1).toEqual({
@@ -582,16 +583,35 @@ describe("duplicateNodes", () => {
582583
});
583584

584585
it("should handle graph output connections", () => {
585-
const taskSpec: TaskSpec = {
586-
...mockTaskSpec,
587-
arguments: {},
588-
};
589-
590586
const outputSpec: OutputSpec = {
591587
...mockOutputSpec,
592588
name: "graph-output",
593589
};
594590

591+
const taskComponentSpec: ComponentSpec = {
592+
name: "task-component",
593+
inputs: [],
594+
outputs: [
595+
{
596+
name: "graph-output",
597+
type: "String",
598+
annotations: {},
599+
},
600+
],
601+
implementation: {
602+
container: { image: "task-image" },
603+
},
604+
};
605+
606+
const taskSpec: TaskSpec = {
607+
...mockTaskSpec,
608+
arguments: {},
609+
componentRef: {
610+
name: "task-component",
611+
spec: taskComponentSpec,
612+
},
613+
};
614+
595615
const componentSpec = createMockComponentSpecWithOutputs(
596616
{ task1: taskSpec },
597617
[],
@@ -617,7 +637,7 @@ describe("duplicateNodes", () => {
617637
expect(outputValues?.["graph-output 2"]).toEqual({
618638
taskOutput: {
619639
taskId: "task1 2",
620-
outputName: "result",
640+
outputName: "graph-output",
621641
},
622642
});
623643
}

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

Lines changed: 85 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -70,12 +70,7 @@ export const duplicateNodes = (
7070
const oldNodeId = node.id;
7171

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

@@ -95,16 +90,14 @@ export const duplicateNodes = (
9590
};
9691
newTasks[newTaskId] = newTaskSpec;
9792
} else if (isInputNode(node)) {
98-
const inputSpec = componentSpec.inputs?.find(
99-
(input) => input.name === node.data.label,
100-
);
93+
const inputSpec = node.data.spec;
10194

102-
const newInputName = getUniqueInputName(componentSpec, inputSpec?.name);
95+
const newInputName = getUniqueInputName(componentSpec, inputSpec.name);
10396
const newNodeId = nodeManager.getNodeId(newInputName, "input");
10497

10598
nodeIdMap[oldNodeId] = newNodeId;
10699

107-
const annotations = inputSpec?.annotations || {};
100+
const annotations = inputSpec.annotations || {};
108101

109102
const updatedAnnotations = setPositionInAnnotations(annotations, {
110103
x: node.position.x + OFFSET,
@@ -119,19 +112,14 @@ export const duplicateNodes = (
119112

120113
newInputs[newInputName] = newInputSpec;
121114
} else if (isOutputNode(node)) {
122-
const outputSpec = componentSpec.outputs?.find(
123-
(output) => output.name === node.data.label,
124-
);
115+
const outputSpec = node.data.spec;
125116

126-
const newOutputName = getUniqueOutputName(
127-
componentSpec,
128-
outputSpec?.name,
129-
);
117+
const newOutputName = getUniqueOutputName(componentSpec, outputSpec.name);
130118
const newNodeId = nodeManager.getNodeId(newOutputName, "output");
131119

132120
nodeIdMap[oldNodeId] = newNodeId;
133121

134-
const annotations = outputSpec?.annotations || {};
122+
const annotations = outputSpec.annotations || {};
135123

136124
const updatedAnnotations = setPositionInAnnotations(annotations, {
137125
x: node.position.x + OFFSET,
@@ -186,10 +174,11 @@ export const duplicateNodes = (
186174
}
187175
});
188176

189-
// Outputs are defined in the graph spec
190177
const updatedGraphOutputs = { ...graphSpec.outputValues };
191-
if (connection !== "none") {
192-
/* Reconfigure Outputs */
178+
179+
// Output Nodes can only have one input, so the "external" connection mode does not apply to them.
180+
if (connection !== "none" && connection !== "external") {
181+
/* Reconfigure Output Nodes */
193182
Object.entries(newOutputs).forEach((output) => {
194183
const [outputName] = output;
195184
const newNodeId = nodeManager.getNodeId(outputName, "output");
@@ -201,57 +190,62 @@ export const duplicateNodes = (
201190
return;
202191
}
203192

204-
const oldOutputName = nodeManager.getRefId(oldNodeId);
193+
const originalOutputNode = nodesToDuplicate.find(
194+
(node) => node.id === oldNodeId && isOutputNode(node),
195+
);
205196

206-
if (!graphSpec.outputValues || !oldOutputName) {
197+
if (!originalOutputNode || !isOutputNode(originalOutputNode)) {
207198
return;
208199
}
209200

210-
const outputValue = graphSpec.outputValues[oldOutputName];
201+
const oldOutputName = originalOutputNode.data.spec.name;
202+
203+
const originalOutputValue = graphSpec.outputValues?.[oldOutputName];
211204

212-
if (!outputValue) {
205+
if (!originalOutputValue) {
206+
// 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)
207+
console.warn(
208+
`No output value found for output ${oldOutputName} in graph spec.`,
209+
);
213210
return;
214211
}
215212

216-
const updatedOutputValue = { ...outputValue };
217-
218-
// If the outputvalue references a task that was also duplicated (internal connection), we need to update it to refer to the duplicated task id
219-
let isInternal = false;
220-
if (
221-
typeof updatedOutputValue === "object" &&
222-
updatedOutputValue !== null &&
223-
connection !== "external"
224-
) {
225-
if ("taskOutput" in updatedOutputValue) {
226-
const oldTaskId = updatedOutputValue.taskOutput.taskId;
227-
const oldTaskNodeId = nodeManager.getNodeId(oldTaskId, "task");
228-
if (oldTaskNodeId in nodeIdMap) {
229-
const newTaskId = nodeManager.getRefId(nodeIdMap[oldTaskNodeId]);
230-
if (!newTaskId) {
231-
return;
232-
}
233-
234-
updatedOutputValue.taskOutput = {
235-
...updatedOutputValue.taskOutput,
236-
taskId: newTaskId,
237-
};
238-
239-
isInternal = true;
240-
}
241-
}
213+
if (!("taskOutput" in originalOutputValue)) {
214+
console.warn(
215+
`Output ${oldOutputName} is not connected to a task output`,
216+
);
217+
return;
242218
}
243219

244-
if (
245-
(isInternal && connection === "internal") ||
246-
(!isInternal && connection === "external") ||
247-
connection === "all"
248-
) {
249-
updatedGraphOutputs[outputName] = updatedOutputValue;
220+
const connectedTaskId = originalOutputValue.taskOutput.taskId;
221+
const connectedOutputName = originalOutputValue.taskOutput.outputName;
222+
223+
const originalNodeId = nodeManager.getNodeId(connectedTaskId, "task");
224+
225+
const newTaskNodeId = nodeIdMap[originalNodeId];
226+
if (!newTaskNodeId) {
227+
console.warn(`No mapping found for task node ${originalNodeId}`);
228+
return;
250229
}
230+
231+
const newTaskId = nodeManager.getRefId(newTaskNodeId);
232+
if (!newTaskId) {
233+
console.warn(`Could not get new task ID for node ${newTaskNodeId}`);
234+
return;
235+
}
236+
237+
const outputValue: TaskOutputArgument = {
238+
taskOutput: {
239+
taskId: newTaskId,
240+
outputName: connectedOutputName,
241+
},
242+
};
243+
244+
updatedGraphOutputs[outputName] = outputValue;
251245
});
252246
}
253247

254-
/* Update the Graph Spec & Inputs */
248+
/* Update the Graph Spec */
255249
const updatedTasks = { ...graphSpec.tasks, ...newTasks };
256250
const updatedGraphSpec = {
257251
...graphSpec,
@@ -507,43 +501,61 @@ function reconfigureConnections(
507501

508502
if ("taskOutput" in argument) {
509503
const oldTaskId = argument.taskOutput.taskId;
510-
oldNodeId = nodeManager.getNodeId(oldTaskId, "task");
511504

512-
if (!isGraphImplementation(componentSpec.implementation)) {
513-
throw new Error("ComponentSpec does not contain a graph implementation.");
505+
const taskNode = nodes.find(
506+
(node) => isTaskNode(node) && node.data.taskId === oldTaskId,
507+
);
508+
509+
if (taskNode) {
510+
oldNodeId = taskNode.id;
511+
isExternal = false;
512+
} else {
513+
if (!isGraphImplementation(componentSpec.implementation)) {
514+
throw new Error(
515+
"ComponentSpec does not contain a graph implementation.",
516+
);
517+
}
518+
519+
const graphSpec = componentSpec.implementation.graph;
520+
isExternal = oldTaskId in graphSpec.tasks;
514521
}
515522

516-
const graphSpec = componentSpec.implementation.graph;
517-
isExternal = oldTaskId in graphSpec.tasks;
523+
if (!oldNodeId) {
524+
return reconfigureExternalConnection(taskSpec, argKey, mode);
525+
}
518526

519527
const newNodeId = nodeIdMap[oldNodeId];
520-
521528
if (!newNodeId) {
522529
return reconfigureExternalConnection(taskSpec, argKey, mode);
523530
}
524531

525532
const newTaskId = nodeManager.getRefId(newNodeId);
526-
527533
newArgId = newTaskId;
528534
} else if ("graphInput" in argument) {
529535
const oldInputName = argument.graphInput.inputName;
530-
oldNodeId = nodeManager.getNodeId(oldInputName, "input");
531536

532-
if (!("inputs" in componentSpec)) {
533-
throw new Error("ComponentSpec does not contain inputs.");
537+
const inputNode = nodes.find(
538+
(node) => isInputNode(node) && node.data.spec.name === oldInputName,
539+
);
540+
541+
if (inputNode) {
542+
oldNodeId = inputNode.id;
543+
isExternal = false;
544+
} else {
545+
const inputs = componentSpec.inputs || [];
546+
isExternal = inputs.some((input) => input.name === oldInputName);
534547
}
535548

536-
const inputs = componentSpec.inputs || [];
537-
isExternal = inputs.some((input) => input.name === oldInputName);
549+
if (!oldNodeId) {
550+
return reconfigureExternalConnection(taskSpec, argKey, mode);
551+
}
538552

539553
const newNodeId = nodeIdMap[oldNodeId];
540-
541554
if (!newNodeId) {
542555
return reconfigureExternalConnection(taskSpec, argKey, mode);
543556
}
544557

545558
const newInputName = nodeManager.getRefId(newNodeId);
546-
547559
newArgId = newInputName;
548560
}
549561

src/utils/unique.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,12 @@ const makeNameUniqueByAddingIndex = (
44
name: string,
55
existingNames: Set<string>,
66
): string => {
7-
let finalName = name;
7+
const baseName = name.replace(/ \(\d+\)$/, "");
8+
let finalName = baseName;
89
let index = 1;
910
while (existingNames.has(finalName)) {
1011
index++;
11-
finalName = name + " " + index.toString();
12+
finalName = baseName + " (" + index.toString() + ")";
1213
}
1314
return finalName;
1415
};

0 commit comments

Comments
 (0)