diff --git a/frontend/src/lib/components/graph/FlowGraphV2.svelte b/frontend/src/lib/components/graph/FlowGraphV2.svelte
index bb88abfd98..7ea4710f7e 100644
--- a/frontend/src/lib/components/graph/FlowGraphV2.svelte
+++ b/frontend/src/lib/components/graph/FlowGraphV2.svelte
@@ -145,8 +145,8 @@
 
 				dispatch('delete', detail)
 			},
-			newBranch: (detail, label) => {
-				dispatch('newBranch', detail)
+			newBranch: (module) => {
+				dispatch('newBranch', module)
 			},
 			move: (module, modules) => {
 				dispatch('move', { module, modules })
@@ -154,7 +154,8 @@
 			selectIteration: (detail, moduleId) => {
 				dispatch('selectIteration', { ...detail, moduleId: moduleId })
 			}
-		}
+		},
+		success
 	)
 
 	const nodes = writable<Node[]>([])
@@ -163,14 +164,10 @@
 	function updateStores() {
 		$nodes = layoutNodes(graph?.nodes)
 		$edges = graph.edges
-
-		renderCount++
 	}
 
 	$: graph && updateStores()
 
-	let renderCount: number = 0
-
 	const nodeTypes = {
 		input2: InputNode,
 		module: ModuleNode,
@@ -210,34 +207,32 @@
 	bind:clientWidth={width}
 	class="h-full"
 >
-	{#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}
+	<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>
 </div>
 
 <style>
diff --git a/frontend/src/lib/components/graph/graphBuilder.ts b/frontend/src/lib/components/graph/graphBuilder.ts
index 8ad81232bf..dd408cb010 100644
--- a/frontend/src/lib/components/graph/graphBuilder.ts
+++ b/frontend/src/lib/components/graph/graphBuilder.ts
@@ -6,8 +6,8 @@ export type GraphEventHandlers = {
 	deleteBranch: (detail, label: string) => void
 	select: (modId: string) => void
 	delete: (detail, label: string) => void
-	newBranch: (detail, label: string) => void
-	move: (module: FlowModule, modules: FlowModule[], index: number) => void
+	newBranch: (module: FlowModule) => void
+	move: (module: FlowModule, modules: FlowModule[]) => void
 	selectIteration: (detail, moduleId: string) => void
 }
 
@@ -15,7 +15,8 @@ export default function graphBuilder(
 	modules: FlowModule[] | undefined,
 	extra: Record<string, any>,
 	failureModule: FlowModule | undefined,
-	eventHandlers: GraphEventHandlers
+	eventHandlers: GraphEventHandlers,
+	success: boolean | undefined
 ): {
 	nodes: Node[]
 	edges: Edge[]
@@ -68,14 +69,19 @@ export default function graphBuilder(
 		position: { x: -1, y: -1 },
 		type: 'input2',
 		data: {
-			parentIds: []
+			eventHandlers: eventHandlers,
+			modules: modules,
+			...extra
 		}
 	}
 
 	const resultNode: Node = {
 		id: 'result',
 		data: {
-			parentIds: []
+			eventHandlers: eventHandlers,
+			modules: modules,
+			success: success,
+			...extra
 		},
 		position: { x: -1, y: -1 },
 		type: 'result'
diff --git a/frontend/src/lib/components/graph/renderers/nodes/ForLoopEndNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/ForLoopEndNode.svelte
index 5fcda8ec97..7d8e4ea582 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/ForLoopEndNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/ForLoopEndNode.svelte
@@ -4,6 +4,9 @@
 	import { Handle, NodeToolbar, Position, type NodeProps } from '@xyflow/svelte'
 	import NodeWrapper from './NodeWrapper.svelte'
 	import type { FlowModule } from '$lib/gen/models/FlowModule'
+	import type { GraphModuleState } from '../../model'
+	import { getStateColor } from '../../util'
+	import type { GraphEventHandlers } from '../../graphBuilder'
 
 	export let data: {
 		offset: number
@@ -11,6 +14,8 @@
 		time: number | undefined
 		insertable: boolean
 		modules: FlowModule[]
+		flowModuleStates: Record<string, GraphModuleState> | undefined
+		eventHandlers: GraphEventHandlers
 	}
 </script>
 
@@ -20,13 +25,19 @@
 	</NodeToolbar>
 {/if}
 
-<NodeWrapper offset={data.offset} enableSourceHandle enableTargetHandle>
+<NodeWrapper offset={data.offset} enableSourceHandle enableTargetHandle let:darkMode>
 	<VirtualItem
 		label={'Collect result of each iteration'}
 		modules={data.modules}
-		index={1}
-		selectable
+		index={data.modules.findIndex((m) => m.id == data.id) + 1}
+		selectable={true}
 		selected={false}
 		insertable={data.insertable}
+		id={data.id}
+		bgColor={getStateColor(undefined, darkMode)}
+		borderColor={getStateColor(data.flowModuleStates?.[data.id]?.type, darkMode)}
+		on:select={(e) => {
+			data?.eventHandlers?.select(e.detail)
+		}}
 	/>
 </NodeWrapper>
diff --git a/frontend/src/lib/components/graph/renderers/nodes/ForLoopStartNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/ForLoopStartNode.svelte
index 2e5b0e4f98..bfa9a099b3 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/ForLoopStartNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/ForLoopStartNode.svelte
@@ -3,29 +3,33 @@
 	// @ts-ignore
 	import { Handle, NodeToolbar, Position, type NodeProps } from '@xyflow/svelte'
 	import NodeWrapper from './NodeWrapper.svelte'
+	import type { FlowModule } from '$lib/gen/models/FlowModule'
+	import { getStateColor } from '../../util'
+	import type { GraphModuleState } from '../../model'
+	import type { GraphEventHandlers } from '../../graphBuilder'
 
 	export let data: {
 		offset: number
 		id: string
-		time: number | undefined
+		modules: FlowModule[]
+		flowModuleStates: Record<string, GraphModuleState> | undefined
+		eventHandlers: GraphEventHandlers
 	}
 </script>
 
-{#if data.time}
-	<NodeToolbar isVisible position={Position.Top} align="end">
-		<span class="text-xs">0.01s</span>
-	</NodeToolbar>
-{/if}
-
 <NodeWrapper let:darkMode offset={data.offset} enableSourceHandle enableTargetHandle>
 	<VirtualItem
 		label={'Do one iteration'}
-		modules={[]}
-		index={1}
-		selectable
-		selected={false}
+		modules={data.modules}
+		index={0}
+		selectable={false}
 		insertable={false}
+		selected={false}
 		id={data.id}
-		bgColor={darkMode ? '#2e3440' : '#dfe6ee'}
+		bgColor={getStateColor(undefined, darkMode)}
+		borderColor={getStateColor(data.flowModuleStates?.[data.id]?.type, darkMode)}
+		on:select={(e) => {
+			data?.eventHandlers?.select(e.detail)
+		}}
 	/>
 </NodeWrapper>
diff --git a/frontend/src/lib/components/graph/renderers/nodes/InputNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/InputNode.svelte
index db3519d3bb..d716886cb2 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/InputNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/InputNode.svelte
@@ -2,21 +2,75 @@
 	import VirtualItem from '$lib/components/flows/map/VirtualItem.svelte'
 	// @ts-ignore
 	import { Handle, NodeToolbar, Position, type NodeProps } from '@xyflow/svelte'
-	import { Wand2 } from 'lucide-svelte'
+	import { ExternalLink, Wand2 } from 'lucide-svelte'
 	import NodeWrapper from './NodeWrapper.svelte'
 	import Popover from '$lib/components/Popover.svelte'
+	import type { GraphEventHandlers } from '../../graphBuilder'
+	import type { FlowModule } from '$lib/gen/models/FlowModule'
+	import { getStateColor } from '../../util'
+	import Menu from '$lib/components/common/menu/Menu.svelte'
+	import type { FlowCopilotContext } from '$lib/components/copilot/flow'
+	import { getContext } from 'svelte'
+	import { copilotInfo } from '$lib/stores'
+	import { base } from '$lib/base'
+	import type { Writable } from 'svelte/store'
 
 	export let data: {
 		insertable: boolean
+		eventHandlers: GraphEventHandlers
+		modules: FlowModule[]
+		disableAi: boolean
 	}
+	let openNoCopilot = false
+
+	const { drawerStore: copilotDrawerStore } =
+		getContext<FlowCopilotContext | undefined>('FlowCopilotContext') || {}
+
+	const { selectedId } = getContext<{
+		selectedId: Writable<string | undefined>
+	}>('FlowGraphContext')
 </script>
 
-{#if data.insertable}
+{#if data.insertable && !data.disableAi}
 	<NodeToolbar isVisible position={Position.Top}>
 		<Popover>
-			<button class="rounded-full border p-2 hover:bg-surface-hover bg-surface text-indigo-400">
-				<Wand2 size={16} />
-			</button>
+			<div class={openNoCopilot ? 'z-10' : ''}>
+				<Menu pointerDown noMinW placement="bottom-center" let:close bind:show={openNoCopilot}>
+					<button
+						title="AI Flow Builder"
+						on:pointerdown={$copilotInfo.exists_openai_resource_path
+							? (ev) => {
+									ev.preventDefault()
+									ev.stopPropagation()
+									$copilotDrawerStore?.openDrawer()
+							  }
+							: undefined}
+						slot="trigger"
+						type="button"
+						class=" bg-surface text-violet-800 dark:text-violet-400 border mx-0.5 focus:outline-none hover:bg-surface-hover focus:ring-4 focus:ring-gray-200 font-medium rounded-full text-sm w-8 h-8 flex items-center justify-center"
+					>
+						<Wand2 size={16} />
+					</button>
+					{#if !$copilotInfo.exists_openai_resource_path}
+						<div class="text-primary p-4">
+							<p class="text-sm w-80">
+								Enable Windmill AI in the
+								<a
+									href="{base}/workspace_settings?tab=openai"
+									target="_blank"
+									class="inline-flex flex-row items-center gap-1"
+									on:click={() => {
+										close()
+									}}
+								>
+									workspace settings
+									<ExternalLink size={16} />
+								</a>
+							</p>
+						</div>
+					{/if}
+				</Menu>
+			</div>
 
 			<svelte:fragment slot="text">AI Flow builder</svelte:fragment>
 		</Popover>
@@ -26,11 +80,17 @@
 <NodeWrapper enableTargetHandle={false} let:darkMode>
 	<VirtualItem
 		label="Input"
-		modules={[]}
-		index={1}
+		modules={data.modules}
+		index={0}
 		selectable
-		selected={false}
+		selected={$selectedId === 'Input'}
 		insertable={false}
-		bgColor={darkMode ? '#2e3440' : '#dfe6ee'}
+		bgColor={getStateColor(undefined, darkMode)}
+		on:insert={(e) => {
+			data.eventHandlers?.insert(e.detail)
+		}}
+		on:select={(e) => {
+			data.eventHandlers?.select(e.detail)
+		}}
 	/>
 </NodeWrapper>
diff --git a/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte
index 795751543c..3851070256 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/ModuleNode.svelte
@@ -17,7 +17,6 @@
 		trigger: boolean
 		insertable: boolean
 		insertableEnd: boolean
-		annotation: string | undefined
 		branchable: boolean
 		bgColor: string
 		modules: FlowModule[]
@@ -39,10 +38,19 @@
 	const dispatch = createEventDispatcher()
 	let openMenu: boolean | undefined = undefined
 
-	let type = data.flowModuleStates?.[data.module.id]?.type
+	$: type = data.flowModuleStates?.[data.module.id]?.type
 	if (!type && data.flowJobs) {
 		type = 'InProgress'
 	}
+
+	const state = data.flowModuleStates?.[data.module.id]
+	const flowJobs = state?.flow_jobs
+		? {
+				flowJobs: state?.flow_jobs,
+				selected: state?.selectedForloopIndex ?? -1,
+				flowJobsSuccess: state?.flow_jobs_success
+		  }
+		: (undefined as any)
 </script>
 
 <NodeWrapper offset={data.offset} let:darkMode>
@@ -51,7 +59,9 @@
 		trigger={data.trigger}
 		insertable={data.insertable}
 		insertableEnd={data.insertableEnd}
-		annotation={data.annotation}
+		annotation={state?.flow_jobs
+			? 'ierations ' + state?.flow_jobs?.length + '/' + (state?.iteration_total ?? '?')
+			: ''}
 		branchable={data.branchable}
 		bgColor={getStateColor(type, darkMode)}
 		modules={data.modules ?? []}
@@ -60,17 +70,25 @@
 		disableAi={data.disableAi}
 		wrapperId={data.wrapperId}
 		retries={data.retries}
-		flowJobs={data.flowJobs}
+		{flowJobs}
 		on:delete={(e) => {
 			data.eventHandlers.delete(e.detail, '')
 		}}
 		on:insert={(e) => {
 			data.eventHandlers.insert(e.detail)
 		}}
-		on:move
-		on:newBranch
-		on:select
-		on:selectedIteration
+		on:move={(e) => {
+			data.eventHandlers.move(data.module, data.modules)
+		}}
+		on:newBranch={(e) => {
+			data.eventHandlers.newBranch(data.module)
+		}}
+		on:select={(e) => {
+			data.eventHandlers.select(e.detail)
+		}}
+		on:selectedIteration={(e) => {
+			data.eventHandlers.selectIteration(e.detail, data.module.id)
+		}}
 	/>
 </NodeWrapper>
 
diff --git a/frontend/src/lib/components/graph/renderers/nodes/ResultNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/ResultNode.svelte
index c8e337d6e6..57c946ba56 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/ResultNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/ResultNode.svelte
@@ -3,16 +3,39 @@
 	// @ts-ignore
 	import { Handle, Position, type NodeProps } from '@xyflow/svelte'
 	import NodeWrapper from './NodeWrapper.svelte'
+	import type { FlowModule } from '$lib/gen/models/FlowModule'
+	import type { GraphEventHandlers } from '../../graphBuilder'
+	import { getStateColor } from '../../util'
+	import type { Writable } from 'svelte/store'
+	import { getContext } from 'svelte'
+
+	export let data: {
+		eventHandlers: GraphEventHandlers
+		modules: FlowModule[]
+		success: boolean | undefined
+	}
+
+	const { selectedId } = getContext<{
+		selectedId: Writable<string | undefined>
+	}>('FlowGraphContext')
 </script>
 
 <NodeWrapper let:darkMode enableSourceHandle={false}>
 	<VirtualItem
+		id={'Result'}
 		label={'Result'}
-		modules={[]}
-		index={1}
-		selectable
-		selected={false}
+		modules={undefined}
+		index={data.modules.length}
+		selectable={true}
+		selected={$selectedId === 'Result'}
 		insertable={false}
-		bgColor={darkMode ? '#2e3440' : '#dfe6ee'}
+		hideId={true}
+		bgColor={getStateColor(
+			data.success == undefined ? undefined : data.success ? 'Success' : 'Failure',
+			darkMode
+		)}
+		on:select={(e) => {
+			data.eventHandlers?.select(e.detail)
+		}}
 	/>
 </NodeWrapper>
diff --git a/frontend/src/lib/components/graph/renderers/nodes/branchAllEndNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/branchAllEndNode.svelte
index 22871529d3..4114d6c1eb 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/branchAllEndNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/branchAllEndNode.svelte
@@ -4,30 +4,29 @@
 	import { Handle, NodeToolbar, Position, type NodeProps } from '@xyflow/svelte'
 	import NodeWrapper from './NodeWrapper.svelte'
 	import type { FlowModule } from '$lib/gen/models/FlowModule'
+	import type { GraphEventHandlers } from '../../graphBuilder'
+	import { getStateColor } from '../../util'
 
 	export let data: {
 		id: string
-		time: number | undefined
 		insertable: boolean
 		modules: FlowModule[]
+		eventHandlers: GraphEventHandlers
 	}
 </script>
 
-{#if data.time}
-	<NodeToolbar isVisible position={Position.Top} align="end">
-		<span class="text-xs">0.01s</span>
-	</NodeToolbar>
-{/if}
-
 <NodeWrapper let:darkMode enableSourceHandle enableTargetHandle>
 	<VirtualItem
 		label={'Collect result from all branches'}
-		modules={[]}
+		modules={data.modules}
 		index={data.modules.findIndex((m) => m.id == data.id) + 1}
 		id={data.id}
-		selectable
+		selectable={true}
 		selected={false}
 		insertable={false}
-		bgColor={darkMode ? '#2e3440' : '#dfe6ee'}
+		bgColor={getStateColor(undefined, darkMode)}
+		on:select={(e) => {
+			data?.eventHandlers?.select(e.detail)
+		}}
 	/>
 </NodeWrapper>