diff --git a/frontend/src/lib/components/graph/FlowGraphV2.svelte b/frontend/src/lib/components/graph/FlowGraphV2.svelte
index 18a64886ab..c5ac094c36 100644
--- a/frontend/src/lib/components/graph/FlowGraphV2.svelte
+++ b/frontend/src/lib/components/graph/FlowGraphV2.svelte
@@ -12,7 +12,8 @@
 		Position,
 		type Node,
 		type Edge,
-		ConnectionLineType
+		ConnectionLineType,
+		MiniMap
 		// @ts-ignore
 	} from '@xyflow/svelte'
 	// @ts-ignore
@@ -26,6 +27,7 @@
 	import ForLoopStartNode from './renderers/nodes/ForLoopStartNode.svelte'
 	import ResultNode from './renderers/nodes/ResultNode.svelte'
 	import BaseEdge from './renderers/edges/BaseEdge.svelte'
+	import EmptyEdge from './renderers/edges/EmptyEdge.svelte'
 
 	export let success: boolean | undefined = undefined
 	export let modules: FlowModule[] | undefined = []
@@ -53,11 +55,10 @@
 		flowInputsStore: Writable<FlowInput | undefined>
 	}>('FlowGraphContext', { selectedId, flowInputsStore })
 
-	// XYFLOW
-
 	const nodeWidth = 300
 	const nodeHeight = 36
 	const dagreGraph = new dagre.graphlib.Graph()
+
 	dagreGraph.setDefaultEdgeLabel(() => ({}))
 
 	function getLayoutedElements(nodes: Node[], edges: Edge[], direction = 'TB') {
@@ -89,10 +90,7 @@
 		return { nodes, edges }
 	}
 
-	const { nodes: initialNodes, edges: initialEdges } = graphBuilder(modules, {
-		insertable,
-		disableAi
-	})
+	const { nodes: initialNodes, edges: initialEdges } = graphBuilder(modules)
 
 	const { nodes: layoutedNodes, edges: layoutedEdges } = getLayoutedElements(
 		initialNodes,
@@ -114,8 +112,11 @@
 		whileLoopEnd: ForLoopEndNode
 	}
 	const edgeTypes = {
-		edge: BaseEdge
+		edge: BaseEdge,
+		empty: EmptyEdge
 	}
+
+	const proOptions = { hideAttribution: true }
 </script>
 
 <div
@@ -132,7 +133,9 @@
 		connectionLineType={ConnectionLineType.SmoothStep}
 		defaultEdgeOptions={{ type: 'smoothstep' }}
 		fitView
+		{proOptions}
 	>
 		<Background class="!bg-surface-secondary" />
+		<MiniMap nodeStrokeWidth={3} />
 	</SvelteFlow>
 </div>
diff --git a/frontend/src/lib/components/graph/graphBuilder.ts b/frontend/src/lib/components/graph/graphBuilder.ts
index 1cd0153a7f..5ca9086180 100644
--- a/frontend/src/lib/components/graph/graphBuilder.ts
+++ b/frontend/src/lib/components/graph/graphBuilder.ts
@@ -28,12 +28,12 @@ export default function graphBuilder(modules: FlowModule[] | undefined): {
 		return module.id
 	}
 
-	function addEdge(sourceId: string, targetId: string, customId?: string) {
+	function addEdge(sourceId: string, targetId: string, customId?: string, type?: string) {
 		edges.push({
 			id: customId || `edge:${sourceId}->${targetId}`,
 			source: sourceId,
 			target: targetId,
-			type: 'edge'
+			type: type ?? 'edge'
 		})
 	}
 
@@ -96,7 +96,7 @@ export default function graphBuilder(modules: FlowModule[] | undefined): {
 
 					nodes.push(startNode)
 
-					addEdge(module.id, startNode.id)
+					addEdge(module.id, startNode.id, undefined, 'empty')
 
 					processModules(branch.modules, startNode, endNode)
 				})
@@ -112,7 +112,7 @@ export default function graphBuilder(modules: FlowModule[] | undefined): {
 					type: 'forLoopStart'
 				}
 
-				addEdge(module.id, startNode.id)
+				addEdge(module.id, startNode.id, undefined, 'empty')
 
 				const endNode = {
 					id: `${module.id}-end`,
@@ -136,7 +136,7 @@ export default function graphBuilder(modules: FlowModule[] | undefined): {
 					position: { x: 0, y: 0 },
 					type: 'whileLoopStart'
 				}
-				addEdge(module.id, startNode.id)
+				addEdge(module.id, startNode.id, undefined, 'empty')
 
 				const endNode = {
 					id: `${module.id}-end`,
diff --git a/frontend/src/lib/components/graph/renderers/edges/BaseEdge.svelte b/frontend/src/lib/components/graph/renderers/edges/BaseEdge.svelte
index 061fd57008..1bae3e31ce 100644
--- a/frontend/src/lib/components/graph/renderers/edges/BaseEdge.svelte
+++ b/frontend/src/lib/components/graph/renderers/edges/BaseEdge.svelte
@@ -1,9 +1,16 @@
 <script lang="ts">
-	// @ts-ignore
-	import { getBezierPath, BaseEdge, type EdgeProps } from '@xyflow/svelte'
+	import InsertModuleButton from '$lib/components/flows/map/InsertModuleButton.svelte'
+	import {
+		getBezierPath,
+		BaseEdge,
+		type EdgeProps,
+		EdgeLabelRenderer
+		// @ts-ignore
+	} from '@xyflow/svelte'
 
 	type $$Props = EdgeProps
 
+	export let id: $$Props['id']
 	export let sourceX: $$Props['sourceX']
 	export let sourceY: $$Props['sourceY']
 	export let sourcePosition: $$Props['sourcePosition']
@@ -13,7 +20,7 @@
 	export let markerEnd: $$Props['markerEnd'] = undefined
 	export let style: $$Props['style'] = undefined
 
-	$: [edgePath] = getBezierPath({
+	$: [edgePath, labelX, labelY] = getBezierPath({
 		sourceX,
 		sourceY,
 		sourcePosition,
@@ -24,3 +31,36 @@
 </script>
 
 <BaseEdge path={edgePath} {markerEnd} {style} />
+<EdgeLabelRenderer>
+	<div
+		class="edgeButtonContainer nodrag nopan"
+		style:transform="translate(-50%, -50%) translate({labelX}px,{labelY}px)"
+	>
+		<InsertModuleButton index={0} modules={[]} />
+	</div>
+</EdgeLabelRenderer>
+
+<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;
+	}
+
+	.edgeButton {
+		width: 20px;
+		height: 20px;
+		background: #eee;
+		border: 1px solid #fff;
+		cursor: pointer;
+		border-radius: 50%;
+		font-size: 12px;
+		line-height: 1;
+	}
+
+	.edgeButton:hover {
+		box-shadow: 0 0 6px 2px rgba(0, 0, 0, 0.08);
+	}
+</style>
diff --git a/frontend/src/lib/components/graph/renderers/edges/EmptyEdge.svelte b/frontend/src/lib/components/graph/renderers/edges/EmptyEdge.svelte
new file mode 100644
index 0000000000..ee504504f2
--- /dev/null
+++ b/frontend/src/lib/components/graph/renderers/edges/EmptyEdge.svelte
@@ -0,0 +1,30 @@
+<script lang="ts">
+	import {
+		getBezierPath,
+		BaseEdge,
+		type EdgeProps
+		// @ts-ignore
+	} from '@xyflow/svelte'
+
+	type $$Props = EdgeProps
+
+	export let sourceX: $$Props['sourceX']
+	export let sourceY: $$Props['sourceY']
+	export let sourcePosition: $$Props['sourcePosition']
+	export let targetX: $$Props['targetX']
+	export let targetY: $$Props['targetY']
+	export let targetPosition: $$Props['targetPosition']
+	export let markerEnd: $$Props['markerEnd'] = undefined
+	export let style: $$Props['style'] = undefined
+
+	$: [edgePath] = getBezierPath({
+		sourceX,
+		sourceY,
+		sourcePosition,
+		targetX,
+		targetY,
+		targetPosition
+	})
+</script>
+
+<BaseEdge path={edgePath} {markerEnd} {style} />
diff --git a/frontend/src/lib/components/graph/renderers/nodes/BranchAllStart.svelte b/frontend/src/lib/components/graph/renderers/nodes/BranchAllStart.svelte
index c45294d82a..d4214bc78d 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/BranchAllStart.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/BranchAllStart.svelte
@@ -5,6 +5,7 @@
 	import NodeWrapper from './NodeWrapper.svelte'
 	import { X } from 'lucide-svelte'
 	import { createEventDispatcher } from 'svelte'
+	import Popover from '$lib/components/Popover.svelte'
 
 	export let data: {
 		label: string
@@ -15,20 +16,21 @@
 </script>
 
 <NodeToolbar isVisible position={Position.Top}>
-	<button
-		class="rounded-full border bg-white p-1 hover:bg-gray-200"
-		on:click={() => {
-			dispatch('deleteBranch', { branchIndex: data.branchIndex })
-		}}
-	>
-		<X size={16} />
-	</button>
+	<Popover>
+		<button
+			class="rounded-full border p-1 hover:bg-surface-hover bg-surface"
+			on:click={() => {
+				dispatch('deleteBranch', { branchIndex: data.branchIndex })
+			}}
+		>
+			<X size={16} />
+		</button>
+
+		<svelte:fragment slot="text">Delete branch</svelte:fragment>
+	</Popover>
 </NodeToolbar>
 
-<NodeWrapper>
-	<svelte:fragment slot="header">
-		<span class="text-xs">0.01s</span>
-	</svelte:fragment>
+<NodeWrapper let:darkMode>
 	<VirtualItem
 		label={data.label}
 		modules={[]}
@@ -36,5 +38,6 @@
 		selectable
 		selected={false}
 		insertable={false}
+		bgColor={darkMode ? '#2e3440' : '#dfe6ee'}
 	/>
 </NodeWrapper>
diff --git a/frontend/src/lib/components/graph/renderers/nodes/ForLoopEndNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/ForLoopEndNode.svelte
index b14efbfeb4..fce031c5ad 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/ForLoopEndNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/ForLoopEndNode.svelte
@@ -6,7 +6,6 @@
 	type $$Props = NodeProps
 
 	export let data: $$Props['data']
-	$$restProps
 </script>
 
 {#if data.time}
diff --git a/frontend/src/lib/components/graph/renderers/nodes/ForLoopStartNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/ForLoopStartNode.svelte
index a742bc0969..2e5b0e4f98 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/ForLoopStartNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/ForLoopStartNode.svelte
@@ -1,24 +1,23 @@
 <script lang="ts">
-	import DarkModeObserver from '$lib/components/DarkModeObserver.svelte'
 	import VirtualItem from '$lib/components/flows/map/VirtualItem.svelte'
 	// @ts-ignore
 	import { Handle, NodeToolbar, Position, type NodeProps } from '@xyflow/svelte'
+	import NodeWrapper from './NodeWrapper.svelte'
 
-	type $$Props = NodeProps
-
-	export let data: $$Props['data']
-	let darkMode: boolean = false
+	export let data: {
+		offset: number
+		id: string
+		time: number | undefined
+	}
 </script>
 
-<DarkModeObserver bind:darkMode />
-
 {#if data.time}
 	<NodeToolbar isVisible position={Position.Top} align="end">
 		<span class="text-xs">0.01s</span>
 	</NodeToolbar>
 {/if}
 
-<div style={`margin-left: ${data.offset}px;`}>
+<NodeWrapper let:darkMode offset={data.offset} enableSourceHandle enableTargetHandle>
 	<VirtualItem
 		label={'Do one iteration'}
 		modules={[]}
@@ -29,7 +28,4 @@
 		id={data.id}
 		bgColor={darkMode ? '#2e3440' : '#dfe6ee'}
 	/>
-</div>
-
-<Handle type="source" position={Position.Bottom} style={`margin-left: ${data.offset / 2}px;`} />
-<Handle type="target" position={Position.Top} style={`margin-left: ${data.offset / 2}px;`} />
+</NodeWrapper>
diff --git a/frontend/src/lib/components/graph/renderers/nodes/InputNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/InputNode.svelte
index 66636c33ec..c70dbf6cc3 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/InputNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/InputNode.svelte
@@ -3,16 +3,28 @@
 	// @ts-ignore
 	import { Handle, NodeToolbar, Position, type NodeProps } from '@xyflow/svelte'
 	import { Wand2 } from 'lucide-svelte'
-
-	type $$Props = NodeProps
+	import NodeWrapper from './NodeWrapper.svelte'
+	import Popover from '$lib/components/Popover.svelte'
 </script>
 
 <NodeToolbar isVisible position={Position.Top}>
-	<button class="rounded-full border bg-indigo-100 p-2 hover:bg-indigo-200">
-		<Wand2 size={16} />
-	</button>
+	<Popover>
+		<button class="rounded-full border p-2 hover:bg-surface-hover bg-surface text-indigo-400">
+			<Wand2 size={16} />
+		</button>
+
+		<svelte:fragment slot="text">AI Flow builder</svelte:fragment>
+	</Popover>
 </NodeToolbar>
 
-<VirtualItem label="Input" modules={[]} index={1} selectable selected={false} insertable={false} />
-
-<Handle type="source" position={Position.Bottom} />
+<NodeWrapper enableTargetHandle={false} let:darkMode>
+	<VirtualItem
+		label="Input"
+		modules={[]}
+		index={1}
+		selectable
+		selected={false}
+		insertable={false}
+		bgColor={darkMode ? '#2e3440' : '#dfe6ee'}
+	/>
+</NodeWrapper>
diff --git a/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte
index ace79b5969..783c4d9bdb 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte
@@ -1,4 +1,5 @@
 <script lang="ts">
+	// @ts-ignore
 	import { NodeToolbar, Position } from '@xyflow/svelte'
 	import MapItem from '$lib/components/flows/map/MapItem.svelte'
 	import type { FlowModule, FlowModuleValue } from '$lib/gen/types.gen'
@@ -105,7 +106,7 @@
 <NodeToolbar isVisible position={Position.Bottom} align="center">
 	{#if data.value.type === 'branchall'}
 		<button
-			class="rounded-full border bg-white p-1 hover:bg-gray-200"
+			class="rounded-full border hover:bg-surface-hover bg-surface p-1"
 			on:click={() => {
 				dispatch('addBranch')
 			}}
diff --git a/frontend/src/lib/components/graph/renderers/nodes/NodeWrapper.svelte b/frontend/src/lib/components/graph/renderers/nodes/NodeWrapper.svelte
index db18e2e063..c1f0d0ffb7 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/NodeWrapper.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/NodeWrapper.svelte
@@ -1,18 +1,23 @@
 <script lang="ts">
+	import DarkModeObserver from '$lib/components/DarkModeObserver.svelte'
+
 	// @ts-ignore
 	import { Handle, NodeToolbar, Position, type NodeProps } from '@xyflow/svelte'
 
 	export let enableSourceHandle: boolean = true
 	export let enableTargetHandle: boolean = true
 	export let offset: number = 0
+	let darkMode: boolean = false
 </script>
 
-<div class="relative" style={`margin-left: ${offset}px;`}>
+<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 />
+	<slot {darkMode} />
 </div>
 
 {#if enableSourceHandle}
diff --git a/frontend/src/lib/components/graph/renderers/nodes/ResultNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/ResultNode.svelte
index 56226dba00..c8e337d6e6 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/ResultNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/ResultNode.svelte
@@ -1,24 +1,18 @@
 <script lang="ts">
-	import DarkModeObserver from '$lib/components/DarkModeObserver.svelte'
 	import VirtualItem from '$lib/components/flows/map/VirtualItem.svelte'
 	// @ts-ignore
 	import { Handle, Position, type NodeProps } from '@xyflow/svelte'
-
-	type $$Props = NodeProps
-
-	let darkMode: boolean = false
+	import NodeWrapper from './NodeWrapper.svelte'
 </script>
 
-<DarkModeObserver bind:darkMode />
-
-<VirtualItem
-	label={'Result'}
-	modules={[]}
-	index={1}
-	selectable
-	selected={false}
-	insertable={false}
-	bgColor={darkMode ? '#2e3440' : '#dfe6ee'}
-/>
-
-<Handle type="target" position={Position.Top} />
+<NodeWrapper let:darkMode enableSourceHandle={false}>
+	<VirtualItem
+		label={'Result'}
+		modules={[]}
+		index={1}
+		selectable
+		selected={false}
+		insertable={false}
+		bgColor={darkMode ? '#2e3440' : '#dfe6ee'}
+	/>
+</NodeWrapper>
diff --git a/frontend/src/lib/components/graph/renderers/nodes/branchAllEndNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/branchAllEndNode.svelte
index 2562b91486..567813afb7 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/branchAllEndNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/branchAllEndNode.svelte
@@ -2,6 +2,7 @@
 	import VirtualItem from '$lib/components/flows/map/VirtualItem.svelte'
 	// @ts-ignore
 	import { Handle, NodeToolbar, Position, type NodeProps } from '@xyflow/svelte'
+	import NodeWrapper from './NodeWrapper.svelte'
 
 	type $$Props = NodeProps
 
@@ -14,15 +15,15 @@
 	</NodeToolbar>
 {/if}
 
-<VirtualItem
-	label={'Collect result from all branches'}
-	modules={[]}
-	index={1}
-	id={data.id}
-	selectable
-	selected={false}
-	insertable={false}
-/>
-
-<Handle type="source" position={Position.Bottom} />
-<Handle type="target" position={Position.Top} />
+<NodeWrapper let:darkMode enableSourceHandle enableTargetHandle>
+	<VirtualItem
+		label={'Collect result from all branches'}
+		modules={[]}
+		index={1}
+		id={data.id}
+		selectable
+		selected={false}
+		insertable={false}
+		bgColor={darkMode ? '#2e3440' : '#dfe6ee'}
+	/>
+</NodeWrapper>