Skip to content

Commit ec381ea

Browse files
camielvscamielvs
committed
Refactor addAndConnectNode
1 parent f7f0199 commit ec381ea

File tree

1 file changed

+112
-135
lines changed

1 file changed

+112
-135
lines changed

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

Lines changed: 112 additions & 135 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import type { Connection, Handle } from "@xyflow/react";
22

3-
import type { NodeManager, NodeType } from "@/nodeManager";
3+
import type { HandleInfo, NodeManager, NodeType } from "@/nodeManager";
44
import {
55
type ComponentReference,
66
type ComponentSpec,
@@ -9,10 +9,6 @@ import {
99
type TypeSpecType,
1010
} from "@/utils/componentSpec";
1111
import { DEFAULT_NODE_DIMENSIONS } from "@/utils/constants";
12-
import {
13-
inputIdToInputName,
14-
outputIdToOutputName,
15-
} from "@/utils/nodes/conversions";
1612

1713
import addTask from "./addTask";
1814
import { handleConnection } from "./handleConnection";
@@ -25,180 +21,161 @@ type AddAndConnectNodeParams = {
2521
nodeManager: NodeManager;
2622
};
2723

24+
/*
25+
Add a new task node to the graph and connect it to an existing handle
26+
- ComponentRef: The component reference for the new task
27+
- fromHandle: The handle from which the connection originates (input or output handle)
28+
- position: The position to place the new task node
29+
- componentSpec: The current component specification containing the graph
30+
- nodeManager: The NodeManager instance for managing node IDs
31+
*/
32+
2833
export function addAndConnectNode({
2934
componentRef,
3035
fromHandle,
3136
position,
3237
componentSpec,
3338
nodeManager,
3439
}: AddAndConnectNodeParams): ComponentSpec {
35-
// 1. Add the new node
36-
const taskSpec: TaskSpec = {
37-
annotations: {},
38-
componentRef: { ...componentRef },
39-
};
40-
41-
if (!isGraphImplementation(componentSpec.implementation)) {
42-
return componentSpec;
43-
}
44-
45-
const oldGraphSpec = componentSpec.implementation.graph;
46-
47-
if (!fromHandle?.id) {
48-
return componentSpec;
49-
}
50-
51-
const fromNodeId = fromHandle.nodeId;
52-
const fromNodeType = nodeManager.getNodeType(fromNodeId);
53-
const fromTaskId = nodeManager.getRefId(fromNodeId);
54-
55-
if (!fromTaskId) {
56-
return componentSpec;
57-
}
58-
59-
let fromHandleType: NodeType | undefined;
60-
let fromHandleName: string | undefined;
61-
62-
if (fromNodeType === "task") {
63-
const fromHandleInfo = nodeManager.getHandleInfo(fromHandle.id);
64-
fromHandleName = fromHandleInfo?.handleName;
65-
fromHandleType = nodeManager.getNodeType(fromHandle.id);
66-
} else if (fromNodeType === "input") {
67-
fromHandleType = "outputHandle";
68-
fromHandleName = inputIdToInputName(fromTaskId);
69-
} else if (fromNodeType === "output") {
70-
fromHandleType = "inputHandle";
71-
fromHandleName = outputIdToOutputName(fromTaskId);
72-
} else {
40+
if (!isGraphImplementation(componentSpec.implementation) || !fromHandle?.id) {
7341
return componentSpec;
7442
}
7543

76-
if (!fromHandleName) {
77-
return componentSpec;
78-
}
44+
const handleInfo = nodeManager.getHandleInfo(fromHandle.id);
45+
const fromHandleType = handleInfo?.handleType;
7946

8047
if (
48+
!handleInfo ||
8149
!fromHandleType ||
82-
(fromHandleType !== "inputHandle" && fromHandleType !== "outputHandle")
50+
!["inputHandle", "outputHandle"].includes(fromHandleType)
8351
) {
8452
return componentSpec;
8553
}
8654

55+
// 1. Create new task
56+
const newTaskSpec: TaskSpec = {
57+
annotations: {},
58+
componentRef: { ...componentRef },
59+
};
60+
8761
const adjustedPosition =
8862
fromHandleType === "inputHandle"
8963
? { ...position, x: position.x - DEFAULT_NODE_DIMENSIONS.w }
9064
: position;
9165

9266
const newComponentSpec = addTask(
9367
"task",
94-
taskSpec,
68+
newTaskSpec,
9569
adjustedPosition,
9670
componentSpec,
9771
);
9872

99-
// 2. Find the new node
10073
if (!isGraphImplementation(newComponentSpec.implementation)) {
101-
return newComponentSpec;
74+
return componentSpec;
10275
}
10376

104-
const graphSpec = newComponentSpec.implementation.graph;
105-
106-
const newTaskId = Object.keys(graphSpec.tasks).find(
107-
(key) => !(key in oldGraphSpec.tasks),
108-
);
77+
const oldTasks = componentSpec.implementation.graph.tasks;
78+
const newTasks = newComponentSpec.implementation.graph.tasks;
79+
const newTaskId = Object.keys(newTasks).find((key) => !(key in oldTasks));
10980

110-
if (!newTaskId) {
111-
return newComponentSpec;
112-
}
113-
114-
const newNodeId = nodeManager.getNodeId(newTaskId, "task");
81+
if (!newTaskId) return newComponentSpec;
11582

116-
// 3. Determine the connection data type and find the first matching handle on the new node
117-
let fromComponentSpec: ComponentSpec | undefined;
118-
119-
if (fromNodeType === "task") {
120-
// Get spec from task
121-
const fromTaskSpec = graphSpec.tasks[fromTaskId];
122-
fromComponentSpec = fromTaskSpec?.componentRef.spec;
123-
} else {
124-
// For IO nodes, get spec from component spec
125-
fromComponentSpec = componentSpec;
126-
}
127-
128-
let connectionType: TypeSpecType | undefined;
129-
if (fromHandleType === "inputHandle") {
130-
connectionType = fromComponentSpec?.inputs?.find(
131-
(io) => io.name === fromHandleName,
132-
)?.type;
133-
} else if (fromHandleType === "outputHandle") {
134-
connectionType = fromComponentSpec?.outputs?.find(
135-
(io) => io.name === fromHandleName,
136-
)?.type;
137-
}
83+
// 2. Find the first matching handle on the new task
84+
const connectionType = getConnectionType(
85+
handleInfo,
86+
fromHandleType,
87+
componentSpec,
88+
newTasks,
89+
);
90+
if (!connectionType) return newComponentSpec;
13891

139-
// Find the first matching handle on the new node
14092
const toHandleType =
14193
fromHandleType === "inputHandle" ? "outputHandle" : "inputHandle";
94+
const toHandleName = findCompatibleHandle(
95+
componentRef,
96+
toHandleType,
97+
connectionType,
98+
);
14299

143-
const inputHandleName = componentRef.spec?.inputs?.find(
144-
(io) => io.type === connectionType,
145-
)?.name;
146-
147-
const outputHandleName = componentRef.spec?.outputs?.find(
148-
(io) => io.type === connectionType,
149-
)?.name;
150-
151-
const toHandleName =
152-
toHandleType === "inputHandle" ? inputHandleName : outputHandleName;
153-
154-
if (!toHandleName) {
155-
return newComponentSpec;
156-
}
100+
if (!toHandleName) return newComponentSpec;
157101

102+
// 3. Create connection
103+
const newNodeId = nodeManager.getNodeId(newTaskId, "task");
158104
const targetHandleId = nodeManager.getHandleNodeId(
159105
newTaskId,
160106
toHandleName,
161107
toHandleType,
162108
);
109+
const isReversed =
110+
fromHandleType === "inputHandle" && toHandleType === "outputHandle";
111+
112+
const connection: Connection = isReversed
113+
? {
114+
source: newNodeId,
115+
sourceHandle: targetHandleId,
116+
target: fromHandle.nodeId,
117+
targetHandle: fromHandle.id,
118+
}
119+
: {
120+
source: fromHandle.nodeId,
121+
sourceHandle: fromHandle.id,
122+
target: newNodeId,
123+
targetHandle: targetHandleId,
124+
};
125+
126+
const updatedGraphSpec = handleConnection(
127+
newComponentSpec.implementation.graph,
128+
connection,
129+
nodeManager,
130+
);
163131

164-
// 4. Build a Connection object and use handleConnection to add the edge
165-
if (targetHandleId) {
166-
const fromNodeId = fromHandle.nodeId;
167-
const fromHandleId = fromHandle.id;
168-
169-
const isReversedConnection =
170-
fromHandleType === "inputHandle" && toHandleType === "outputHandle";
171-
172-
const connection: Connection = isReversedConnection
173-
? // Drawing from an input handle to a new output handle
174-
{
175-
source: newNodeId,
176-
sourceHandle: targetHandleId,
177-
target: fromNodeId,
178-
targetHandle: fromHandleId,
179-
}
180-
: // Drawing from an output handle to a new input handle
181-
{
182-
source: fromNodeId,
183-
sourceHandle: fromHandleId,
184-
target: newNodeId,
185-
targetHandle: targetHandleId,
186-
};
187-
188-
const updatedGraphSpec = handleConnection(
189-
graphSpec,
190-
connection,
191-
nodeManager,
192-
);
193-
194-
return {
195-
...newComponentSpec,
196-
implementation: {
197-
...newComponentSpec.implementation,
198-
graph: updatedGraphSpec,
199-
},
200-
};
132+
return {
133+
...newComponentSpec,
134+
implementation: {
135+
...newComponentSpec.implementation,
136+
graph: updatedGraphSpec,
137+
},
138+
};
139+
}
140+
141+
function getConnectionType(
142+
handleInfo: HandleInfo,
143+
handleType: NodeType,
144+
componentSpec: ComponentSpec,
145+
tasks: Record<string, TaskSpec>,
146+
): TypeSpecType | undefined {
147+
let targetSpec: ComponentSpec | undefined;
148+
149+
const taskSpec = tasks[handleInfo.parentRefId];
150+
151+
if (taskSpec) {
152+
targetSpec = taskSpec.componentRef.spec;
153+
} else {
154+
targetSpec = componentSpec;
155+
}
156+
157+
if (!targetSpec) return undefined;
158+
159+
if (handleType === "inputHandle") {
160+
return targetSpec.inputs?.find((io) => io.name === handleInfo.handleName)
161+
?.type;
162+
} else {
163+
return targetSpec.outputs?.find((io) => io.name === handleInfo.handleName)
164+
?.type;
201165
}
166+
}
202167

203-
return newComponentSpec;
168+
function findCompatibleHandle(
169+
componentRef: ComponentReference,
170+
handleType: "inputHandle" | "outputHandle",
171+
connectionType: TypeSpecType,
172+
): string | undefined {
173+
const spec = componentRef.spec;
174+
if (!spec) return undefined;
175+
176+
if (handleType === "inputHandle") {
177+
return spec.inputs?.find((io) => io.type === connectionType)?.name;
178+
} else {
179+
return spec.outputs?.find((io) => io.type === connectionType)?.name;
180+
}
204181
}

0 commit comments

Comments
 (0)