Skip to content

Commit eaa16e9

Browse files
committed
Refactor Switch node branches with unique identifiers, reusable helper functions, and standardized handle aliasing logic
1 parent 783ead5 commit eaa16e9

14 files changed

Lines changed: 260 additions & 52 deletions

File tree

src/components/workflow/properties/use-node-properties-form.ts

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,18 @@ import { useForm, useWatch } from "react-hook-form";
55
import { zodResolver } from "@hookform/resolvers/zod";
66
import type { NodeType, WorkflowNodeData } from "@/types/workflow";
77
import { nodeSchemaMap, NODE_REGISTRY } from "@/lib/node-registry";
8+
import { WorkflowNodeType } from "@/types/workflow";
9+
import { normalizeSwitchBranches } from "@/nodes/switch/branches";
10+
11+
function normalizeNodeFormData(data: WorkflowNodeData | undefined) {
12+
if (!data) return undefined;
13+
if (data.type !== WorkflowNodeType.Switch) return data as Record<string, unknown>;
14+
15+
return {
16+
...data,
17+
branches: normalizeSwitchBranches(data.branches),
18+
} as Record<string, unknown>;
19+
}
820

921
interface UseNodePropertiesFormOptions {
1022
selectedNodeId: string | null;
@@ -35,16 +47,26 @@ export function useNodePropertiesForm({
3547

3648
const watchedValues = useWatch({ control: form.control });
3749
const readyNodeIdRef = useRef<string | null>(null);
50+
const previousSelectedNodeIdRef = useRef<string | null>(null);
3851

3952
useEffect(() => {
53+
const previousSelectedNodeId = previousSelectedNodeIdRef.current;
54+
const selectedNodeChanged = previousSelectedNodeId !== selectedNodeId;
4055
readyNodeIdRef.current = null;
4156

4257
if (nodeData && registryEntry) {
4358
const defaults = registryEntry.defaultData() as Record<string, unknown>;
44-
const merged = { ...defaults, ...nodeData } as Record<string, unknown>;
45-
form.reset(merged);
59+
const normalizedNodeData = normalizeNodeFormData(nodeData);
60+
const merged = { ...defaults, ...normalizedNodeData } as Record<string, unknown>;
61+
const nextSignature = JSON.stringify(merged);
62+
const currentSignature = JSON.stringify(form.getValues());
63+
if (selectedNodeChanged || currentSignature !== nextSignature) {
64+
form.reset(merged);
65+
}
4666
}
4767

68+
previousSelectedNodeIdRef.current = selectedNodeId ?? null;
69+
4870
const nextReadyNodeId = selectedNodeId ?? null;
4971
const handle = requestAnimationFrame(() => {
5072
readyNodeIdRef.current = nextReadyNodeId;

src/hooks/use-auto-layout.ts

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
} from "@/types/workflow";
1212
import { NODE_REGISTRY } from "@/lib/node-registry";
1313
import { NodeSize, NODE_SIZE_DIMENSIONS } from "@/nodes/shared/node-size";
14+
import { findSwitchBranchIndexByHandle } from "@/nodes/switch/branches";
1415

1516
const LAYOUT_DURATION = 400;
1617

@@ -78,19 +79,23 @@ function fixBranchOrdering(
7879
orderedTargetIds = sorted.map((e) => e.target);
7980
}
8081
} else if (nodeType === WorkflowNodeType.Switch) {
81-
const branches = d.branches as Array<{ label: string }> | undefined;
82+
const branches = d.branches as import("@/types/workflow").SwitchBranch[] | undefined;
8283
if (branches && branches.length >= 2) {
83-
// Build label → edge target lookup
84-
const labelToTarget = new Map<string, string>();
85-
for (const e of branchEdges) {
86-
if (e.sourceHandle) labelToTarget.set(e.sourceHandle, e.target);
87-
}
88-
// Order targets by branch definition order
8984
const ordered: string[] = [];
90-
for (const branch of branches) {
91-
const target = labelToTarget.get(branch.label);
85+
const indexedTargets = new Map<number, string>();
86+
87+
for (const edge of branchEdges) {
88+
const branchIndex = findSwitchBranchIndexByHandle(branches, edge.sourceHandle);
89+
if (branchIndex !== -1 && !indexedTargets.has(branchIndex)) {
90+
indexedTargets.set(branchIndex, edge.target);
91+
}
92+
}
93+
94+
for (let index = 0; index < branches.length; index++) {
95+
const target = indexedTargets.get(index);
9296
if (target) ordered.push(target);
9397
}
98+
9499
if (ordered.length >= 2) orderedTargetIds = ordered;
95100
}
96101
} else if (nodeType === WorkflowNodeType.ParallelAgent) {

src/lib/__tests__/workflow-connections.test.ts

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import {
1010
SKILL_TARGET_HANDLE,
1111
normalizeWorkflowConnection,
1212
} from "../workflow-connections";
13+
import { createSwitchBranch } from "@/nodes/switch/branches";
1314

1415
function connect(
1516
connection: Partial<Connection>,
@@ -198,5 +199,45 @@ describe("workflow connections", () => {
198199
connect({ source: "parallel-1", target: "agent-1", sourceHandle: "branch-0" }, [parallelAgent, agent]),
199200
).toHaveLength(1);
200201
});
202+
203+
it("canonicalizes switch branch handles to the stable branch handle id", () => {
204+
const branches = [
205+
createSwitchBranch({ id: "switch-branch-case-1", label: "Pending", condition: "" }),
206+
createSwitchBranch({ id: "switch-branch-default", label: "default", condition: "Other cases" }),
207+
];
208+
const switchNode = makeWorkflowNode({
209+
id: "switch-1",
210+
type: WorkflowNodeType.Switch,
211+
data: {
212+
type: WorkflowNodeType.Switch,
213+
label: "Switch",
214+
name: "switch-1",
215+
evaluationTarget: "status",
216+
branches,
217+
},
218+
});
219+
const endA = makeWorkflowNode({
220+
id: "end-a",
221+
type: WorkflowNodeType.End,
222+
data: { type: WorkflowNodeType.End, label: "End A", name: "end-a" },
223+
});
224+
const endB = makeWorkflowNode({
225+
id: "end-b",
226+
type: WorkflowNodeType.End,
227+
data: { type: WorkflowNodeType.End, label: "End B", name: "end-b" },
228+
});
229+
230+
const next = connect(
231+
{ source: "switch-1", target: "end-b", sourceHandle: "Pending" },
232+
[switchNode, endA, endB],
233+
[makeWorkflowEdge({ id: "edge-1", source: "switch-1", target: "end-a", sourceHandle: "switch-branch-case-1" })],
234+
);
235+
236+
expect(next).toHaveLength(1);
237+
expect(next?.[0]).toMatchObject({
238+
sourceHandle: "switch-branch-case-1",
239+
target: "end-b",
240+
});
241+
});
201242
});
202243

src/lib/workflow-connections.ts

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ import {
66
type WorkflowEdge,
77
type WorkflowNode,
88
} from "@/types/workflow";
9+
import { findSwitchBranchIndexByHandle, getSwitchBranchHandleId } from "@/nodes/switch/branches";
10+
import type { SwitchNodeData } from "@/nodes/switch/types";
911

1012
export const WORKFLOW_EDGE_TYPE = "deletable" as const;
1113
export const SCRIPT_SOURCE_HANDLE = "script-out";
@@ -35,12 +37,26 @@ export function normalizeWorkflowConnection({
3537

3638
const sourceType = findNodeType(nodes, connection.source);
3739
const targetType = findNodeType(nodes, connection.target);
40+
let normalizedConnection = connection;
41+
42+
if (sourceType === WorkflowNodeType.Switch) {
43+
const sourceNode = nodes.find((node) => node.id === connection.source);
44+
const branches = (sourceNode?.data as SwitchNodeData | undefined)?.branches;
45+
const branchIndex = branches ? findSwitchBranchIndexByHandle(branches, connection.sourceHandle) : -1;
46+
47+
if (branches && branchIndex !== -1) {
48+
normalizedConnection = {
49+
...connection,
50+
sourceHandle: getSwitchBranchHandleId(branches[branchIndex], branchIndex, branches.length),
51+
};
52+
}
53+
}
3854

3955
if (sourceType === WorkflowNodeType.Script) {
4056
if (targetType !== WorkflowNodeType.Skill) return null;
4157
return addEdge(
4258
{
43-
...connection,
59+
...normalizedConnection,
4460
sourceHandle: SCRIPT_SOURCE_HANDLE,
4561
targetHandle: SCRIPT_TARGET_HANDLE,
4662
type: WORKFLOW_EDGE_TYPE,
@@ -53,7 +69,7 @@ export function normalizeWorkflowConnection({
5369
if (!targetType || !AGENT_LIKE_NODE_TYPES.has(targetType)) return null;
5470
return addEdge(
5571
{
56-
...connection,
72+
...normalizedConnection,
5773
targetHandle: SKILL_TARGET_HANDLE,
5874
type: WORKFLOW_EDGE_TYPE,
5975
},
@@ -65,7 +81,7 @@ export function normalizeWorkflowConnection({
6581
if (!targetType || !AGENT_LIKE_NODE_TYPES.has(targetType)) return null;
6682
return addEdge(
6783
{
68-
...connection,
84+
...normalizedConnection,
6985
targetHandle: DOCUMENT_TARGET_HANDLE,
7086
type: WORKFLOW_EDGE_TYPE,
7187
},
@@ -87,16 +103,16 @@ export function normalizeWorkflowConnection({
87103

88104
if (
89105
sourceType === WorkflowNodeType.ParallelAgent
90-
&& connection.sourceHandle?.startsWith("branch-")
106+
&& normalizedConnection.sourceHandle?.startsWith("branch-")
91107
&& targetType !== WorkflowNodeType.Agent
92108
) {
93109
return null;
94110
}
95111

96112
const filteredEdges = edges.filter(
97-
(edge) => !(edge.source === connection.source && edge.sourceHandle === connection.sourceHandle),
113+
(edge) => !(edge.source === normalizedConnection.source && edge.sourceHandle === normalizedConnection.sourceHandle),
98114
);
99115

100-
return addEdge({ ...connection, type: WORKFLOW_EDGE_TYPE }, filteredEdges);
116+
return addEdge({ ...normalizedConnection, type: WORKFLOW_EDGE_TYPE }, filteredEdges);
101117
}
102118

src/lib/workflow-generation/shared.ts

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import { generator as switchGen } from "@/nodes/switch/generator";
1414
import { generator as askUserGen } from "@/nodes/ask-user/generator";
1515
import type { NodeGeneratorModule } from "@/nodes/shared/registry-types";
1616
import { mermaidId, mermaidLabel } from "@/nodes/shared/mermaid-utils";
17+
import { getSwitchBranchLabelFromHandle } from "@/nodes/switch/branches";
1718

1819
export interface GeneratedFile {
1920
path: string;
@@ -101,6 +102,17 @@ export function mermaidEdge(
101102
}
102103
}
103104

105+
if (nodeById) {
106+
const srcNode = nodeById.get(edge.source);
107+
if (srcNode?.data?.type === WorkflowNodeType.Switch) {
108+
const d = srcNode.data as import("@/types/workflow").SwitchNodeData;
109+
const switchLabel = getSwitchBranchLabelFromHandle(d.branches ?? [], raw);
110+
if (switchLabel) {
111+
raw = switchLabel;
112+
}
113+
}
114+
}
115+
104116
const displayLabel = boolHandles.has(raw)
105117
? raw.charAt(0).toUpperCase() + raw.slice(1)
106118
: raw;

src/nodes/switch/branches.ts

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
import { customAlphabet } from "nanoid";
2+
import type { SwitchBranch } from "./types";
3+
4+
const nanoid = customAlphabet("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789", 8);
5+
6+
export function createSwitchBranchId() {
7+
return `switch-branch-${nanoid(8)}`;
8+
}
9+
10+
export function createSwitchBranch(branch: Partial<SwitchBranch> = {}): SwitchBranch {
11+
return {
12+
id: branch.id ?? createSwitchBranchId(),
13+
label: branch.label ?? "",
14+
condition: branch.condition ?? "",
15+
};
16+
}
17+
18+
export function createDefaultSwitchBranches(): SwitchBranch[] {
19+
return [
20+
createSwitchBranch({ label: "Case 1", condition: "" }),
21+
createSwitchBranch({ label: "Case 2", condition: "" }),
22+
createSwitchBranch({ label: "default", condition: "Other cases" }),
23+
];
24+
}
25+
26+
export function normalizeSwitchBranches(branches?: SwitchBranch[]): SwitchBranch[] {
27+
if (!branches?.length) return createDefaultSwitchBranches();
28+
return branches.map((branch) => createSwitchBranch(branch));
29+
}
30+
31+
export function isDefaultSwitchBranch(branch: Pick<SwitchBranch, "label">, index: number, total: number) {
32+
return index === total - 1 || branch.label === "default";
33+
}
34+
35+
export function getSwitchBranchHandleId(
36+
branch: Pick<SwitchBranch, "id" | "label">,
37+
index: number,
38+
total: number,
39+
) {
40+
if (branch.id) return branch.id;
41+
if (isDefaultSwitchBranch(branch, index, total)) return "default";
42+
return `branch-${index}`;
43+
}
44+
45+
export function getSwitchBranchHandleAliases(
46+
branch: Pick<SwitchBranch, "id" | "label">,
47+
index: number,
48+
total: number,
49+
) {
50+
const aliases = new Set<string>();
51+
const primary = getSwitchBranchHandleId(branch, index, total);
52+
const indexHandle = `branch-${index}`;
53+
54+
if (indexHandle !== primary) aliases.add(indexHandle);
55+
if (isDefaultSwitchBranch(branch, index, total) && primary !== "default") aliases.add("default");
56+
57+
const trimmedLabel = branch.label.trim();
58+
if (trimmedLabel && trimmedLabel !== primary) aliases.add(trimmedLabel);
59+
60+
return [...aliases];
61+
}
62+
63+
export function findSwitchBranchIndexByHandle(
64+
branches: Array<Pick<SwitchBranch, "id" | "label">>,
65+
handle: string | null | undefined,
66+
) {
67+
if (!handle) return -1;
68+
69+
return branches.findIndex((branch, index) => {
70+
const primary = getSwitchBranchHandleId(branch, index, branches.length);
71+
if (handle === primary) return true;
72+
return getSwitchBranchHandleAliases(branch, index, branches.length).includes(handle);
73+
});
74+
}
75+
76+
export function getSwitchBranchLabelFromHandle(
77+
branches: Array<Pick<SwitchBranch, "id" | "label">>,
78+
handle: string | null | undefined,
79+
) {
80+
const index = findSwitchBranchIndexByHandle(branches, handle);
81+
if (index === -1) return null;
82+
83+
const branch = branches[index];
84+
if (branch.label.trim()) return branch.label;
85+
return isDefaultSwitchBranch(branch, index, branches.length) ? "default" : `Case ${index + 1}`;
86+
}
87+

src/nodes/switch/constants.ts

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import type { NodeRegistryEntry } from "@/nodes/shared/registry-types";
55
import { NODE_ACCENT } from "@/lib/node-colors";
66
import { WorkflowNodeType } from "@/types/workflow";
77
import type { SwitchNodeData } from "./types";
8+
import { createDefaultSwitchBranches } from "./branches";
89

910
export const switchRegistryEntry: NodeRegistryEntry = {
1011
type: WorkflowNodeType.Switch,
@@ -19,11 +20,7 @@ export const switchRegistryEntry: NodeRegistryEntry = {
1920
label: "Switch",
2021
name: "",
2122
evaluationTarget: "",
22-
branches: [
23-
{ label: "Case 1", condition: "" },
24-
{ label: "Case 2", condition: "" },
25-
{ label: "default", condition: "Other cases" },
26-
],
23+
branches: createDefaultSwitchBranches(),
2724
}),
2825
};
2926

src/nodes/switch/fields.tsx

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { Plus, X } from "lucide-react";
1010
import { NODE_ACCENT, BRANCH_DEFAULT } from "@/lib/node-colors";
1111
import type { FormRegister, FormControl } from "@/nodes/shared/form-types";
1212
import { RequiredIndicator } from "@/nodes/shared/required-indicator";
13+
import { createSwitchBranch } from "./branches";
1314

1415
interface SwitchFieldsProps {
1516
register: FormRegister;
@@ -27,7 +28,7 @@ export function Fields({ register, control }: SwitchFieldsProps) {
2728

2829
const handleAddBranch = () => {
2930
// Insert before the default branch (last item)
30-
insert(editableCount, { label: `Case ${editableCount + 1}`, condition: "" } as never);
31+
insert(editableCount, createSwitchBranch({ label: `Case ${editableCount + 1}`, condition: "" }) as never);
3132
};
3233

3334
return (
@@ -169,4 +170,4 @@ export function Fields({ register, control }: SwitchFieldsProps) {
169170
</div>
170171
</>
171172
);
172-
}
173+
}

0 commit comments

Comments
 (0)