diff --git a/frontend/src/lib/components/graph/FlowGraphV2.svelte b/frontend/src/lib/components/graph/FlowGraphV2.svelte
index 2f749424ae..d31e728b97 100644
--- a/frontend/src/lib/components/graph/FlowGraphV2.svelte
+++ b/frontend/src/lib/components/graph/FlowGraphV2.svelte
@@ -21,7 +21,7 @@
 	import graphBuilder from './graphBuilder'
 	import ModuleNode from './renderers/nodes/ModuleNode.svelte'
 	import InputNode from './renderers/nodes/InputNode.svelte'
-	import BranchAllStart from './renderers/nodes/BranchOneStart.svelte'
+	import BranchAllStart from './renderers/nodes/BranchAllStart.svelte'
 	import BranchAllEndNode from './renderers/nodes/branchAllEndNode.svelte'
 	import ForLoopEndNode from './renderers/nodes/ForLoopEndNode.svelte'
 	import ForLoopStartNode from './renderers/nodes/ForLoopStartNode.svelte'
diff --git a/frontend/src/lib/components/graph/graphBuilder.ts b/frontend/src/lib/components/graph/graphBuilder.ts
index 1fdbf17ecd..75d11a86d3 100644
--- a/frontend/src/lib/components/graph/graphBuilder.ts
+++ b/frontend/src/lib/components/graph/graphBuilder.ts
@@ -67,6 +67,12 @@ export default function graphBuilder(
 	) {
 		parents[targetId] = [...(parents[targetId] ?? []), sourceId]
 
+		// Find the index of the target module in the modules array
+		const mods = options?.subModules ?? modules
+
+		// Index of the target module in the modules array
+		let index = mods?.findIndex((m) => m.id === targetId) ?? -1
+
 		edges.push({
 			id: options?.customId || `edge:${sourceId}->${targetId}`,
 			source: sourceId,
@@ -80,7 +86,8 @@ export default function graphBuilder(
 				offset: options?.offset,
 				moving,
 				eventHandlers,
-				index: options?.subModules?.length ?? modules?.length
+				// If the index is -1, it means that the target module is not in the modules array, so we set it to the length of the array
+				index: index >= 0 ? index : mods?.length ?? 0
 			}
 		})
 	}
@@ -120,12 +127,12 @@ export default function graphBuilder(
 		let previousId: string | undefined = undefined
 
 		if (modules.length === 0) {
-			addEdge(beforeNode.id, nextNode.id, { subModules: modules })
+			addEdge(beforeNode.id, nextNode.id, { subModules: modules, offset: currentOffset })
 		} else {
 			modules.forEach((module, index) => {
 				// Add the edge between the previous node and the current one
 				if (index > 0 && previousId) {
-					addEdge(previousId, module.id)
+					addEdge(previousId, module.id, { subModules: modules, offset: currentOffset })
 				}
 
 				if (module.value.type === 'branchall') {
@@ -167,9 +174,9 @@ export default function graphBuilder(
 
 						nodes.push(startNode)
 
-						addEdge(module.id, startNode.id, { type: 'empty' })
+						addEdge(module.id, startNode.id, { type: 'empty', offset: currentOffset })
 
-						processModules(branch.modules, startNode, endNode)
+						processModules(branch.modules, startNode, endNode, currentOffset)
 					})
 
 					previousId = endNode.id
@@ -190,7 +197,7 @@ export default function graphBuilder(
 						type: 'forLoopStart'
 					}
 
-					addEdge(module.id, startNode.id, { type: 'empty' })
+					addEdge(module.id, startNode.id, { type: 'empty', offset: currentOffset })
 
 					const endNode = {
 						id: `${module.id}-end`,
@@ -310,11 +317,11 @@ export default function graphBuilder(
 				}
 
 				if (index === 0) {
-					addEdge(beforeNode.id, module.id)
+					addEdge(beforeNode.id, module.id, { subModules: modules })
 				}
 
 				if (index === modules.length - 1 && previousId) {
-					addEdge(previousId, nextNode.id)
+					addEdge(previousId, nextNode.id, { subModules: modules })
 				}
 			})
 		}
diff --git a/frontend/src/lib/components/graph/renderers/edges/BaseEdge.svelte b/frontend/src/lib/components/graph/renderers/edges/BaseEdge.svelte
index feb4b52a25..9fb6997c8e 100644
--- a/frontend/src/lib/components/graph/renderers/edges/BaseEdge.svelte
+++ b/frontend/src/lib/components/graph/renderers/edges/BaseEdge.svelte
@@ -6,6 +6,7 @@
 	import { getContext } from 'svelte'
 	import type { Writable } from 'svelte/store'
 	import type { GraphEventHandlers } from '../../graphBuilder'
+	import { getStraightLinePath } from '../utils'
 
 	export let sourceX: number
 	export let sourceY: number
@@ -29,16 +30,13 @@
 		sourcePosition,
 		targetX,
 		targetY,
-		targetPosition
+		targetPosition,
+		curvature: 0.25
 	})
 
-	function getStraightLinePath({ sourceX, sourceY, targetX, targetY }) {
-		return `M${sourceX},${sourceY} L${sourceX},${targetY - 100}`
-	}
-
 	$: completeEdge =
 		targetY - sourceY > 100
-			? `${edgePath} ${getStraightLinePath({ sourceX, sourceY, targetX, targetY })}`
+			? `${edgePath} ${getStraightLinePath({ sourceX, sourceY, targetY })}`
 			: edgePath
 
 	const { useDataflow } = getContext<{
diff --git a/frontend/src/lib/components/graph/renderers/edges/DataflowEdge.svelte b/frontend/src/lib/components/graph/renderers/edges/DataflowEdge.svelte
index 84073a0bab..d13bb2236d 100644
--- a/frontend/src/lib/components/graph/renderers/edges/DataflowEdge.svelte
+++ b/frontend/src/lib/components/graph/renderers/edges/DataflowEdge.svelte
@@ -1,6 +1,7 @@
 <script lang="ts">
 	import type { FlowModule } from '$lib/gen/models/FlowModule'
 	import { getBezierPath, BaseEdge, type Position, EdgeLabelRenderer } from '@xyflow/svelte'
+	import { getStraightLinePath } from '../utils'
 
 	export let sourceX: number
 	export let sourceY: number
@@ -23,22 +24,13 @@
 		targetX,
 		targetY,
 		targetPosition,
-		curvature: 0.1
+		curvature: 0.25
 	})
-
-	function getStraightLinePath({ sourceX, sourceY, targetX, targetY }) {
-		return `M${sourceX},${sourceY} L${sourceX},${targetY - 100}`
-	}
-
-	$: completeEdge =
-		targetY - sourceY > 100
-			? `${edgePath} ${getStraightLinePath({ sourceX, sourceY, targetX, targetY })}`
-			: edgePath
 </script>
 
 <EdgeLabelRenderer>
 	<div
-		class="edgeButtonContainer nodrag nopan bg-surface-selected p-1 border text-xs"
+		class="absolute cursor-pointer nodrag nopan bg-surface-selected p-1 border text-xs"
 		style:transform="translate(-50%, -50%) translate({labelX}px,{labelY}px)"
 	>
 		{data.sourceId} -> {data.targetId}
@@ -46,17 +38,9 @@
 </EdgeLabelRenderer>
 
 <BaseEdge
-	path={completeEdge}
+	path={targetY - sourceY > 100
+		? `${edgePath} ${getStraightLinePath({ sourceX, sourceY, targetY })}`
+		: edgePath}
 	{markerEnd}
 	style={`animation:dashdraw 0.5s linear infinite; stroke-dasharray: 5px;`}
 />
-
-<style>
-	.edgeButtonContainer {
-		position: absolute;
-		font-size: 12pt;
-		/* everything inside EdgeLabelRenderer has no pointer events by default */
-		/* if you have an interactive element, set pointer-events: all */
-		pointer-events: all;
-	}
-</style>
diff --git a/frontend/src/lib/components/graph/renderers/edges/EmptyEdge.svelte b/frontend/src/lib/components/graph/renderers/edges/EmptyEdge.svelte
index 768a5e664f..bc10f9d1f9 100644
--- a/frontend/src/lib/components/graph/renderers/edges/EmptyEdge.svelte
+++ b/frontend/src/lib/components/graph/renderers/edges/EmptyEdge.svelte
@@ -21,7 +21,8 @@
 		sourcePosition,
 		targetX,
 		targetY,
-		targetPosition
+		targetPosition,
+		curvature: 0.25
 	})
 </script>
 
diff --git a/frontend/src/lib/components/graph/renderers/nodes/BranchAllStart.svelte b/frontend/src/lib/components/graph/renderers/nodes/BranchAllStart.svelte
index 6952696dae..ad2dee4be8 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/BranchAllStart.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/BranchAllStart.svelte
@@ -1,14 +1,12 @@
 <script lang="ts">
 	import VirtualItem from '$lib/components/flows/map/VirtualItem.svelte'
-	import { NodeToolbar, Position } from '@xyflow/svelte'
 	import NodeWrapper from './NodeWrapper.svelte'
 	import { X } from 'lucide-svelte'
-	import Popover from '$lib/components/Popover.svelte'
-	import type { FlowStatusModule } from '$lib/gen'
 	import type { GraphModuleState } from '../../model'
 	import { getStateColor } from '../../util'
 	import type { FlowModule } from '$lib/gen/models/FlowModule'
 	import type { GraphEventHandlers } from '../../graphBuilder'
+	import { computeBorderStatus } from '../utils'
 
 	export let data: {
 		label: string
@@ -19,48 +17,13 @@
 		modules: FlowModule[]
 		selected: boolean
 		eventHandlers: GraphEventHandlers
+		offset: number
 	}
 
-	let borderStatus: FlowStatusModule['type'] | undefined = undefined
-
-	let flow_jobs_success = data?.flowModuleStates?.[data?.id]?.flow_jobs_success
-	if (!flow_jobs_success) {
-		borderStatus = 'WaitingForPriorSteps'
-	} else {
-		let status = flow_jobs_success?.[data.branchIndex]
-		if (status == undefined) {
-			borderStatus = 'WaitingForExecutor'
-		} else {
-			borderStatus = status ? 'Success' : 'Failure'
-		}
-	}
+	$: borderStatus = computeBorderStatus(data.branchIndex, data.flowModuleStates?.[data.id])
 </script>
 
-<NodeToolbar isVisible position={Position.Top} align="center">
-	{#if data.insertable}
-		<Popover>
-			<button
-				class="rounded-full border p-1 hover:bg-surface-hover bg-surface"
-				on:click={() => {
-					data.eventHandlers.deleteBranch(
-						{
-							module: data.modules.find((m) => m.id === data.id),
-							index: data.branchIndex
-						},
-						data.label
-					)
-				}}
-			>
-				<X size={16} />
-			</button>
-
-			<svelte:fragment slot="text">Delete branch</svelte:fragment>
-		</Popover>
-	{/if}
-</NodeToolbar>
-
-<NodeWrapper let:darkMode
-	>s
+<NodeWrapper let:darkMode offset={data.offset}>
 	<VirtualItem
 		label={data.label}
 		modules={data.modules}
@@ -81,4 +44,20 @@
 			data.eventHandlers.insert(e.detail)
 		}}
 	/>
+	<button
+		class="z-50 absolute -top-[10px] -right-[10px] rounded-full h-[20px] w-[20px] center-center text-primary
+border-[1.5px] border-gray-700 bg-surface duration-150 hover:bg-red-400 hover:text-white
+hover:border-red-700"
+		on:click|preventDefault|stopPropagation={() => {
+			data.eventHandlers.deleteBranch(
+				{
+					module: data.modules.find((m) => m.id === data.id),
+					index: data.branchIndex
+				},
+				data.label
+			)
+		}}
+	>
+		<X size={12} />
+	</button>
 </NodeWrapper>
diff --git a/frontend/src/lib/components/graph/renderers/nodes/BranchOneStart.svelte b/frontend/src/lib/components/graph/renderers/nodes/BranchOneStart.svelte
index bcbe335f1a..148b06a140 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/BranchOneStart.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/BranchOneStart.svelte
@@ -1,14 +1,12 @@
 <script lang="ts">
 	import VirtualItem from '$lib/components/flows/map/VirtualItem.svelte'
-	import { NodeToolbar, Position } from '@xyflow/svelte'
 	import NodeWrapper from './NodeWrapper.svelte'
 	import { X } from 'lucide-svelte'
-	import Popover from '$lib/components/Popover.svelte'
-	import type { FlowStatusModule } from '$lib/gen'
 	import type { GraphModuleState } from '../../model'
 	import { getStateColor } from '../../util'
 	import type { FlowModule } from '$lib/gen/models/FlowModule'
 	import type { GraphEventHandlers } from '../../graphBuilder'
+	import { computeBorderStatus } from '../utils'
 
 	export let data: {
 		label: string
@@ -19,47 +17,13 @@
 		selected: boolean
 		eventHandlers: GraphEventHandlers
 		branchIndex: number
+		offset: number
 	}
 
-	let borderStatus: FlowStatusModule['type'] | undefined = undefined
-
-	let flow_jobs_success = data?.flowModuleStates?.[data?.id]?.flow_jobs_success
-	if (!flow_jobs_success) {
-		borderStatus = 'WaitingForPriorSteps'
-	} else {
-		let status = flow_jobs_success?.[data.branchIndex]
-		if (status == undefined) {
-			borderStatus = 'WaitingForExecutor'
-		} else {
-			borderStatus = status ? 'Success' : 'Failure'
-		}
-	}
+	$: borderStatus = computeBorderStatus(data.branchIndex, data.flowModuleStates?.[data.id])
 </script>
 
-<NodeToolbar isVisible position={Position.Top}>
-	{#if data.insertable}
-		<Popover>
-			<button
-				class="rounded-full border p-1 hover:bg-surface-hover bg-surface"
-				on:click={(e) => {
-					data.eventHandlers.deleteBranch(
-						{
-							module: data.modules.find((m) => m.id === data.id),
-							index: data.branchIndex
-						},
-						data.label
-					)
-				}}
-			>
-				<X size={16} />
-			</button>
-
-			<svelte:fragment slot="text">Delete branch</svelte:fragment>
-		</Popover>
-	{/if}
-</NodeToolbar>
-
-<NodeWrapper let:darkMode>
+<NodeWrapper let:darkMode offset={data.offset}>
 	<VirtualItem
 		label={data.label}
 		modules={data.modules}
@@ -74,4 +38,20 @@
 			data.eventHandlers.select(data.id)
 		}}
 	/>
+	<button
+		class="z-50 absolute -top-[10px] -right-[10px] rounded-full h-[20px] w-[20px] center-center text-primary
+border-[1.5px] border-gray-700 bg-surface duration-150 hover:bg-red-400 hover:text-white
+hover:border-red-700"
+		on:click|preventDefault|stopPropagation={() => {
+			data.eventHandlers.deleteBranch(
+				{
+					module: data.modules.find((m) => m.id === data.id),
+					index: data.branchIndex
+				},
+				data.label
+			)
+		}}
+	>
+		<X size={12} />
+	</button>
 </NodeWrapper>
diff --git a/frontend/src/lib/components/graph/renderers/nodes/ForLoopEndNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/ForLoopEndNode.svelte
index c07015a4b3..f1d842281c 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/ForLoopEndNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/ForLoopEndNode.svelte
@@ -1,6 +1,5 @@
 <script lang="ts">
 	import VirtualItem from '$lib/components/flows/map/VirtualItem.svelte'
-	import { NodeToolbar, Position } from '@xyflow/svelte'
 	import NodeWrapper from './NodeWrapper.svelte'
 	import type { FlowModule } from '$lib/gen/models/FlowModule'
 	import type { GraphModuleState } from '../../model'
@@ -10,7 +9,6 @@
 	export let data: {
 		offset: number
 		id: string
-		time: number | undefined
 		insertable: boolean
 		modules: FlowModule[]
 		flowModuleStates: Record<string, GraphModuleState> | undefined
@@ -18,13 +16,7 @@
 	}
 </script>
 
-{#if data.time}
-	<NodeToolbar isVisible position={Position.Top} align="end">
-		<span class="text-xs">0.01s</span>
-	</NodeToolbar>
-{/if}
-
-<NodeWrapper offset={data.offset} enableSourceHandle enableTargetHandle let:darkMode>
+<NodeWrapper offset={data.offset} let:darkMode>
 	<VirtualItem
 		label={'Collect result of each iteration'}
 		modules={data.modules}
diff --git a/frontend/src/lib/components/graph/renderers/nodes/ForLoopStartNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/ForLoopStartNode.svelte
index bce3d35306..9c68f4d8ea 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/ForLoopStartNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/ForLoopStartNode.svelte
@@ -15,7 +15,7 @@
 	}
 </script>
 
-<NodeWrapper let:darkMode offset={data.offset} enableSourceHandle enableTargetHandle>
+<NodeWrapper let:darkMode offset={data.offset}>
 	<VirtualItem
 		label={'Do one iteration'}
 		modules={data.modules}
diff --git a/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte
index 1bdf031a7a..a89a406c60 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte
@@ -79,7 +79,7 @@
 </NodeWrapper>
 
 <NodeToolbar isVisible position={Position.Bottom} align="center">
-	{#if data.value.type === 'branchall' && data.insertable}
+	{#if (data.value.type === 'branchall' || data.value.type === 'branchone') && data.insertable}
 		<button
 			class="rounded-full border hover:bg-surface-hover bg-surface p-1"
 			on:click={() => {
diff --git a/frontend/src/lib/components/graph/renderers/nodes/NodeWrapper.svelte b/frontend/src/lib/components/graph/renderers/nodes/NodeWrapper.svelte
index 185a752a64..6bc6427a62 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/NodeWrapper.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/NodeWrapper.svelte
@@ -12,10 +12,6 @@
 <DarkModeObserver bind:darkMode />
 
 <div class="relative shadow-md" style={`margin-left: ${offset}px;`}>
-	<div class="absolute -top-6 z-50">
-		<slot name="header" />
-	</div>
-
 	<slot {darkMode} />
 </div>
 
diff --git a/frontend/src/lib/components/graph/renderers/nodes/branchAllEndNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/branchAllEndNode.svelte
index b48b15abfe..d03cd4901f 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/branchAllEndNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/branchAllEndNode.svelte
@@ -12,10 +12,11 @@
 		modules: FlowModule[]
 		eventHandlers: GraphEventHandlers
 		flowModuleStates: Record<string, GraphModuleState> | undefined
+		offset: number
 	}
 </script>
 
-<NodeWrapper let:darkMode enableSourceHandle enableTargetHandle>
+<NodeWrapper offset={data.offset} let:darkMode enableSourceHandle enableTargetHandle>
 	<VirtualItem
 		label={'Collect result from all branches'}
 		modules={data.modules}
diff --git a/frontend/src/lib/components/graph/renderers/nodes/branchOneEndNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/branchOneEndNode.svelte
index b48b15abfe..616dfdeffe 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/branchOneEndNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/branchOneEndNode.svelte
@@ -12,10 +12,11 @@
 		modules: FlowModule[]
 		eventHandlers: GraphEventHandlers
 		flowModuleStates: Record<string, GraphModuleState> | undefined
+		offset: number
 	}
 </script>
 
-<NodeWrapper let:darkMode enableSourceHandle enableTargetHandle>
+<NodeWrapper let:darkMode offset={data.offset}>
 	<VirtualItem
 		label={'Collect result from all branches'}
 		modules={data.modules}
diff --git a/frontend/src/lib/components/graph/renderers/utils.ts b/frontend/src/lib/components/graph/renderers/utils.ts
new file mode 100644
index 0000000000..432787dd2e
--- /dev/null
+++ b/frontend/src/lib/components/graph/renderers/utils.ts
@@ -0,0 +1,25 @@
+import type { FlowStatusModule } from '$lib/gen/types.gen'
+import type { GraphModuleState } from '../model'
+
+export function getStraightLinePath({ sourceX, sourceY, targetY }) {
+	return `M${sourceX},${sourceY} L${sourceX},${targetY - 100}`
+}
+
+export function computeBorderStatus(
+	branchIndex: number,
+	graphModuleState: GraphModuleState | undefined
+): FlowStatusModule['type'] | undefined {
+	const flowJobSuccess = graphModuleState?.flow_jobs_success
+
+	if (!Array.isArray(flowJobSuccess)) {
+		return 'WaitingForPriorSteps'
+	}
+
+	const jobStatus = flowJobSuccess[branchIndex]
+
+	if (jobStatus === undefined) {
+		return 'WaitingForExecutor'
+	}
+
+	return jobStatus ? 'Success' : 'Failure'
+}