diff --git a/frontend/src/lib/components/flows/map/FlowModuleSchemaMap.svelte b/frontend/src/lib/components/flows/map/FlowModuleSchemaMap.svelte
index 2111eddce1..e0a1d2222b 100644
--- a/frontend/src/lib/components/flows/map/FlowModuleSchemaMap.svelte
+++ b/frontend/src/lib/components/flows/map/FlowModuleSchemaMap.svelte
@@ -16,7 +16,6 @@
 	import FlowConstantsItem from './FlowConstantsItem.svelte'
 
 	import { dfs } from '../dfs'
-	import { FlowGraph } from '$lib/components/graph'
 	import FlowErrorHandlerItem from './FlowErrorHandlerItem.svelte'
 	import { push } from '$lib/history'
 	import ConfirmationModal from '$lib/components/common/confirmationModal/ConfirmationModal.svelte'
@@ -29,6 +28,7 @@
 	import FlowTutorials from '$lib/components/FlowTutorials.svelte'
 	import { ignoredTutorials } from '$lib/components/tutorials/ignoredTutorials'
 	import { tutorialInProgress } from '$lib/tutorialUtils'
+	import FlowGraphV2 from '$lib/components/graph/FlowGraphV2.svelte'
 
 	export let modules: FlowModule[] | undefined
 	export let sidebarSize: number | undefined = undefined
@@ -251,7 +251,7 @@
 	</div>
 
 	<div class="z-10 flex-auto grow" bind:clientHeight={minHeight}>
-		<FlowGraph
+		<FlowGraphV2
 			{disableAi}
 			insertable
 			scroll
@@ -267,7 +267,11 @@
 				dependents = getDependentComponents(e.id, $flowStore)
 				const cb = () => {
 					push(history, $flowStore)
+
+					console.log('deleting node', e.id)
+
 					selectNextId(e.id)
+
 					removeAtId($flowStore.value.modules, e.id)
 
 					if ($flowInputsStore) {
diff --git a/frontend/src/lib/components/flows/map/MapItem.svelte b/frontend/src/lib/components/flows/map/MapItem.svelte
index a617520781..34c545f6d5 100644
--- a/frontend/src/lib/components/flows/map/MapItem.svelte
+++ b/frontend/src/lib/components/flows/map/MapItem.svelte
@@ -74,41 +74,25 @@
 </script>
 
 {#if mod}
-	{#if insertable}
-		<div
-			class="{openMenu
-				? 'z-20'
-				: ''} w-[27px] absolute -top-[35px] left-[50%] right-[50%] -translate-x-1/2"
-		>
-			{#if moving}
-				<button
-					title="Add branch"
-					on:click={() => {
-						dispatch('insert', { modules, index: idx, detail: 'move' })
-					}}
-					type="button"
-					disabled={wrapperId === moving}
-					class=" text-primary bg-surface border mx-[1px] border-gray-300 dark:border-gray-500 focus:outline-none hover:bg-gray-100 focus:ring-4 focus:ring-gray-200 font-medium rounded-full text-sm w-[25px] h-[25px] flex items-center justify-center"
-				>
-					<ClipboardCopy class="m-[5px]" size={15} />
-				</button>
-			{:else}
-				<InsertModuleButton
-					{disableAi}
-					bind:open={openMenu}
-					{trigger}
-					on:insert={(e) => {
-						dispatch('insert', { modules, index: idx + 1, detail: 'script', script: e.detail })
-					}}
-					on:new={(e) => {
-						dispatch('insert', { modules, index: idx, detail: e.detail })
-					}}
-					index={idx}
-					{modules}
-				/>
-			{/if}
-		</div>
-	{/if}
+	<div
+		class="{openMenu
+			? 'z-20'
+			: ''} w-[27px] absolute -top-[35px] left-[50%] right-[50%] -translate-x-1/2"
+	>
+		{#if moving}
+			<button
+				title="Add branch"
+				on:click={() => {
+					dispatch('insert', { modules, index: idx, detail: 'move' })
+				}}
+				type="button"
+				disabled={wrapperId === moving}
+				class=" text-primary bg-surface border mx-[1px] border-gray-300 dark:border-gray-500 focus:outline-none hover:bg-gray-100 focus:ring-4 focus:ring-gray-200 font-medium rounded-full text-sm w-[25px] h-[25px] flex items-center justify-center"
+			>
+				<ClipboardCopy class="m-[5px]" size={15} />
+			</button>
+		{/if}
+	</div>
 	<div class="relative">
 		{#if moving == mod.id}
 			<div class="absolute z-10 right-20 top-0.5 center-center">
diff --git a/frontend/src/lib/components/flows/map/VirtualItem.svelte b/frontend/src/lib/components/flows/map/VirtualItem.svelte
index a7993f636d..a77dfb1ab3 100644
--- a/frontend/src/lib/components/flows/map/VirtualItem.svelte
+++ b/frontend/src/lib/components/flows/map/VirtualItem.svelte
@@ -5,7 +5,6 @@
 	import { classNames } from '$lib/utils'
 	import { ClipboardCopy, ExternalLink, Wand2, X } from 'lucide-svelte'
 	import { createEventDispatcher, getContext } from 'svelte'
-	import InsertModuleButton from './InsertModuleButton.svelte'
 	import type { FlowCopilotContext } from '$lib/components/copilot/flow'
 	import { copilotInfo } from '$lib/stores'
 	import Menu from '$lib/components/common/menu/Menu.svelte'
@@ -103,7 +102,7 @@
 	</div>
 </div>
 
-{#if insertable && modules && (label != 'Input' || modules.length == 0)}
+{#if modules && (label != 'Input' || modules.length == 0)}
 	<div
 		class="{openMenu
 			? 'z-20'
@@ -127,33 +126,6 @@
 			>
 				<ClipboardCopy size={12} />
 			</button>
-		{:else}
-			<InsertModuleButton
-				{disableAi}
-				bind:open={openMenu}
-				trigger={label == 'Input'}
-				on:insert={(e) => {
-					if (modules) {
-						dispatch('insert', {
-							modules,
-							index,
-							detail: 'script',
-							script: e.detail
-						})
-					}
-				}}
-				on:new={(e) => {
-					if (modules) {
-						dispatch('insert', {
-							modules,
-							index: index,
-							detail: e.detail
-						})
-					}
-				}}
-				{index}
-				modules={modules ?? []}
-			/>
 		{/if}
 	</div>
 {/if}
diff --git a/frontend/src/lib/components/graph/FlowGraphV2.svelte b/frontend/src/lib/components/graph/FlowGraphV2.svelte
index baf0f7a1fd..593060db3f 100644
--- a/frontend/src/lib/components/graph/FlowGraphV2.svelte
+++ b/frontend/src/lib/components/graph/FlowGraphV2.svelte
@@ -122,13 +122,98 @@
 			disableAi,
 			insertable
 		},
-		failureModule
+		failureModule,
+		{
+			deleteBranch: (detail, label) => {
+				$selectedId = label
+				dispatch('deleteBranch', detail)
+			},
+			insert: (detail) => {
+				dispatch('insert', detail)
+			},
+			select: (mod) => {
+				if (!notSelectable) {
+					if ($selectedId != mod.id) {
+						$selectedId = mod.id
+					}
+					dispatch('select', mod)
+				}
+			},
+			delete: (detail, label) => {
+				$selectedId = label
+
+				dispatch('delete', detail)
+			},
+			newBranch: (detail, label) => {
+				dispatch('newBranch', detail)
+			},
+			move: (module, modules) => {
+				dispatch('move', { module, modules })
+			},
+			selectIteration: (detail, moduleId) => {
+				dispatch('selectIteration', { ...detail, moduleId: moduleId })
+			}
+		}
 	)
 	const { nodes: layoutedNodes, edges: layoutedEdges } = layoutNodes(initialNodes, initialEdges)
 
 	const nodes = writable<Node[]>(layoutedNodes)
 	const edges = writable<Edge[]>(layoutedEdges)
 
+	let renderCount: number = 0
+
+	function updateStores() {
+		const { nodes: initialNodes, edges: initialEdges } = graphBuilder(
+			modules,
+			{
+				disableAi,
+				insertable
+			},
+			failureModule,
+			{
+				deleteBranch: (detail, label) => {
+					$selectedId = label
+					dispatch('deleteBranch', detail)
+				},
+				insert: (detail) => {
+					dispatch('insert', detail)
+				},
+				select: (mod) => {
+					if (!notSelectable) {
+						if ($selectedId != mod.id) {
+							$selectedId = mod.id
+						}
+						dispatch('select', mod)
+					}
+				},
+				delete: (detail, label) => {
+					$selectedId = label
+
+					dispatch('delete', detail)
+				},
+				newBranch: (detail, label) => {
+					dispatch('newBranch', detail)
+				},
+				move: (module, modules) => {
+					dispatch('move', { module, modules })
+				},
+				selectIteration: (detail, moduleId) => {
+					dispatch('selectIteration', { ...detail, moduleId: moduleId })
+				}
+			}
+		)
+		const { nodes: layoutedNodes, edges: layoutedEdges } = layoutNodes(initialNodes, initialEdges)
+
+		$nodes = layoutedNodes
+		$edges = layoutedEdges
+
+		renderCount++
+	}
+
+	$: console.log('rebuildOnChange', rebuildOnChange)
+
+	$: rebuildOnChange && updateStores()
+
 	const nodeTypes = {
 		input2: InputNode,
 		module: ModuleNode,
@@ -168,29 +253,34 @@
 	bind:clientWidth={width}
 	class="h-full"
 >
-	<SvelteFlow
-		{nodes}
-		{edges}
-		{edgeTypes}
-		{nodeTypes}
-		minZoom={1}
-		connectionLineType={ConnectionLineType.SmoothStep}
-		defaultEdgeOptions={{ type: 'smoothstep' }}
-		fitView
-		{proOptions}
-		nodesDraggable={false}
-		preventScrolling={!scroll}
-		on:nodeclick={(e) => handleNodeClick(e)}
-	>
-		<Background class="!bg-surface-secondary" />
-		<Controls position="top-right" orientation="horizontal" showLock={false}>
-			{#if download}
-				<ControlButton on:click={() => dispatch('expand')} class="!bg-surface">
-					<Expand size="14" />
-				</ControlButton>
-			{/if}
-		</Controls>
-	</SvelteFlow>
+	{#key renderCount}
+		<SvelteFlow
+			{nodes}
+			{edges}
+			{edgeTypes}
+			{nodeTypes}
+			minZoom={1}
+			connectionLineType={ConnectionLineType.SmoothStep}
+			defaultEdgeOptions={{ type: 'smoothstep' }}
+			fitView
+			{proOptions}
+			nodesDraggable={false}
+			preventScrolling={!scroll}
+			on:nodeclick={(e) => handleNodeClick(e)}
+			on:delete={(e) => {
+				console.log('delete', e.detail)
+			}}
+		>
+			<Background class="!bg-surface-secondary" />
+			<Controls position="top-right" orientation="horizontal" showLock={false}>
+				{#if download}
+					<ControlButton on:click={() => dispatch('expand')} class="!bg-surface">
+						<Expand size="14" />
+					</ControlButton>
+				{/if}
+			</Controls>
+		</SvelteFlow>
+	{/key}
 </div>
 
 <style>
diff --git a/frontend/src/lib/components/graph/graphBuilder.ts b/frontend/src/lib/components/graph/graphBuilder.ts
index e3fb7f999e..52cbab0a82 100644
--- a/frontend/src/lib/components/graph/graphBuilder.ts
+++ b/frontend/src/lib/components/graph/graphBuilder.ts
@@ -1,11 +1,21 @@
 import type { FlowModule } from '$lib/gen'
 import { type Node, type Edge } from '@xyflow/svelte'
-import { add } from 'date-fns'
+
+export type GraphEventHandlers = {
+	insert: (detail) => void
+	deleteBranch: (detail, label: string) => void
+	select: (mod: FlowModule) => void
+	delete: (detail, label: string) => void
+	newBranch: (detail, label: string) => void
+	move: (module: FlowModule, modules: FlowModule[], index: number) => void
+	selectIteration: (detail, moduleId: string) => void
+}
 
 export default function graphBuilder(
 	modules: FlowModule[] | undefined,
 	extra: Record<string, any>,
-	failureModule: FlowModule | undefined
+	failureModule: FlowModule | undefined,
+	eventHandlers: GraphEventHandlers
 ): {
 	nodes: Node[]
 	edges: Edge[]
@@ -26,6 +36,7 @@ export default function graphBuilder(
 				module: module,
 				modules: modules,
 				parentIds: [],
+				eventHandlers: eventHandlers,
 				...extra
 			},
 			position: { x: -1, y: -1 },
@@ -129,7 +140,11 @@ export default function graphBuilder(
 
 					addEdge(module.id, startNode.id, undefined, 'empty')
 
-					processModules(branch.modules, startNode, endNode)
+					if (branch.modules.length === 0) {
+						addEdge(startNode.id, endNode.id, undefined, 'empty')
+					} else {
+						processModules(branch.modules, startNode, endNode)
+					}
 				})
 
 				previousId = endNode.id
diff --git a/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte
index c5e9688d7b..f0afefd406 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte
@@ -6,7 +6,7 @@
 	import { ClipboardCopy, GitBranchPlus } from 'lucide-svelte'
 	import { createEventDispatcher } from 'svelte'
 	import NodeWrapper from './NodeWrapper.svelte'
-	import InsertModuleButton from '$lib/components/flows/map/InsertModuleButton.svelte'
+	import type { GraphEventHandlers } from '../../graphBuilder'
 
 	export let data: {
 		offset: number
@@ -27,6 +27,7 @@
 		flowJobs:
 			| { flowJobs: string[]; selected: number; flowJobsSuccess: (boolean | undefined)[] }
 			| undefined
+		eventHandlers: GraphEventHandlers
 	}
 
 	$: idx = data.modules?.findIndex((m) => m.id === data.module.id)
@@ -51,7 +52,9 @@
 		wrapperId={data.wrapperId}
 		retries={data.retries}
 		flowJobs={data.flowJobs}
-		on:delete
+		on:delete={(e) => {
+			data.eventHandlers.delete(e.detail, '')
+		}}
 		on:insert
 		on:move
 		on:newBranch
@@ -79,25 +82,6 @@
 				>
 					<ClipboardCopy class="m-[5px]" size={15} />
 				</button>
-			{:else}
-				<InsertModuleButton
-					bind:disableAi={data.disableAi}
-					bind:open={openMenu}
-					bind:trigger={data.trigger}
-					on:insert={(e) => {
-						dispatch('insert', {
-							modules: data.modules,
-							index: idx + 1,
-							detail: 'script',
-							script: e.detail
-						})
-					}}
-					on:new={(e) => {
-						dispatch('insert', { modules: data.modules, index: idx, detail: e.detail })
-					}}
-					index={idx}
-					modules={data.modules}
-				/>
 			{/if}
 		</div>
 	{/if}