diff --git a/frontend/src/lib/components/FlowStatusViewerInner.svelte b/frontend/src/lib/components/FlowStatusViewerInner.svelte
index cbc7ece547..e881dad539 100644
--- a/frontend/src/lib/components/FlowStatusViewerInner.svelte
+++ b/frontend/src/lib/components/FlowStatusViewerInner.svelte
@@ -197,6 +197,7 @@
 								args: job?.args,
 								tag: job?.tag
 							}
+
 							setModuleState(mod.id ?? '', newState)
 						})
 						.catch((e) => {
@@ -207,7 +208,6 @@
 					(mod.type == 'Success' || mod.type == 'Failure') &&
 					!['Success', 'Failure'].includes($localModuleStates?.[mod.id ?? '']?.type)
 				) {
-					// console.log(mod.id, 'FOO')
 					setModuleState(
 						mod.id ?? '',
 						{
@@ -423,6 +423,11 @@
 					started_at
 				})
 			} else {
+				const parent_module = mod['parent_module']
+
+				// Delete existing failure node attached to the same parent module
+				removeFailureNode(mod.id, parent_module)
+
 				setModuleState(
 					mod.id,
 					{
@@ -432,7 +437,7 @@
 						result: job['result'],
 						job_id: job.id,
 						tag: job.tag,
-						parent_module: mod['parent_module'],
+						parent_module,
 						duration_ms: job['duration_ms'],
 						started_at: started_at,
 						flow_jobs: mod.flow_jobs,
@@ -444,6 +449,7 @@
 					},
 					force
 				)
+
 				setDurationStatusByJob(mod.id, job.id, {
 					created_at: job.created_at ? new Date(job.created_at).getTime() : undefined,
 					started_at,
@@ -608,6 +614,23 @@
 
 	let storedListJobs: Record<number, Job> = {}
 	let wrapperHeight: number = 0
+
+	function removeFailureNode(id: string, parent_module: any) {
+		if (id?.startsWith('failure-') && parent_module) {
+			;[...globalModuleStates, localModuleStates].forEach((stateMapStore) => {
+				stateMapStore.update((stateMap) => {
+					if (id) {
+						Object.keys(stateMap).forEach((key) => {
+							if (stateMap[key]?.parent_module == parent_module) {
+								delete stateMap[key]
+							}
+						})
+					}
+					return stateMap
+				})
+			})
+		}
+	}
 </script>
 
 {#if notAnonynmous}
@@ -857,7 +880,7 @@
 											{childFlow}
 											globalModuleStates={[localModuleStates, ...globalModuleStates]}
 											globalDurationStatuses={[localDurationStatuses, ...globalDurationStatuses]}
-											render={failedRetry == retry_selected}
+											render={failedRetry == retry_selected && render}
 											reducedPolling={false}
 											{workspaceId}
 											jobId={failedRetry}
diff --git a/frontend/src/lib/components/graph/FlowGraphV2.svelte b/frontend/src/lib/components/graph/FlowGraphV2.svelte
index b9c90f499f..e7a48ac84f 100644
--- a/frontend/src/lib/components/graph/FlowGraphV2.svelte
+++ b/frontend/src/lib/components/graph/FlowGraphV2.svelte
@@ -37,6 +37,7 @@
 	import { Alert, Drawer } from '../common'
 	import Button from '../common/button/Button.svelte'
 	import FlowYamlEditor from '../flows/header/FlowYamlEditor.svelte'
+	import BranchOneEndNode from './renderers/nodes/branchOneEndNode.svelte'
 	export let success: boolean | undefined = undefined
 	export let modules: FlowModule[] | undefined = []
 	export let failureModule: FlowModule | undefined = undefined
@@ -215,7 +216,7 @@
 		whileLoopStart: ForLoopStartNode,
 		whileLoopEnd: ForLoopEndNode,
 		branchOneStart: BranchOneStart,
-		branchOneEnd: BranchAllEndNode,
+		branchOneEnd: BranchOneEndNode,
 		noBranch: NoBranchNode,
 		trigger: TriggersNode
 	} as any
@@ -266,7 +267,7 @@
 
 <div style={`height: ${height}px; max-height: ${maxHeight}px;`} bind:clientWidth={width}>
 	{#if graph?.error}
-		<div class="center-center">
+		<div class="center-center p-2">
 			<Alert title="Error parsing the flow" type="error" class="max-w-1/2">
 				{graph.error}
 
diff --git a/frontend/src/lib/components/graph/graphBuilder.ts b/frontend/src/lib/components/graph/graphBuilder.ts
index 573d2da2d5..e63d6c2c51 100644
--- a/frontend/src/lib/components/graph/graphBuilder.ts
+++ b/frontend/src/lib/components/graph/graphBuilder.ts
@@ -36,14 +36,15 @@ export default function graphBuilder(
 } {
 	const nodes: Node[] = []
 	const edges: Edge[] = []
-
 	try {
 		if (!modules) {
 			return { nodes, edges }
 		}
 
 		function addNode(module: FlowModule, offset: number, type: string, subModules?: FlowModule[]) {
-			if (nodes.some((n) => n.id === module.id)) {
+			const duplicated = nodes.find((n) => n.id === module.id)
+			if (duplicated) {
+				console.log('Duplicated node detected: ', module, duplicated)
 				throw new Error(`Duplicated node detected: ${module.id}`)
 			}
 
@@ -199,7 +200,8 @@ export default function graphBuilder(
 			nextNode: Node,
 			currentOffset = 0,
 			disableMoveIds: string[] = [],
-			parentIndex?: string
+			parentIndex?: string,
+			branchChosen?: boolean
 		) {
 			let previousId: string | undefined = undefined
 
@@ -428,13 +430,15 @@ export default function graphBuilder(
 
 						addEdge(module.id, defaultBranch.id, { type: 'empty' })
 
+						const branchChosen = extra.flowModuleStates?.[module.id]?.branchChosen
 						processModules(
 							module.value.default,
 							defaultBranch,
 							endNode,
 							currentOffset,
 							localDisableMoveIds,
-							parentIndex ? `${parentIndex}-${index}` : index.toString()
+							parentIndex ? `${parentIndex}-${index}` : index.toString(),
+							branchChosen == 0
 						)
 
 						module.value.branches.forEach((branch, branchIndex) => {
@@ -466,7 +470,8 @@ export default function graphBuilder(
 								endNode,
 								currentOffset,
 								localDisableMoveIds,
-								parentIndex ? `${parentIndex}-${index}` : index.toString()
+								parentIndex ? `${parentIndex}-${index}` : index.toString(),
+								branchChosen == branchIndex + 1
 							)
 						})
 
@@ -475,6 +480,8 @@ export default function graphBuilder(
 						addNode(module, currentOffset, 'module', modules)
 
 						previousId = module.id
+
+
 					}
 
 					if (index === 0) {
@@ -492,27 +499,29 @@ export default function graphBuilder(
 					}
 				})
 
-				if (failureModule) {
-					const id = parentIndex ? `failure-${parentIndex}` : 'failure'
-					const failureState = extra.flowModuleStates?.[id] as GraphModuleState | undefined
 
-					if (failureState && failureState.parent_module) {
-						addNode(
-							{
-								...failureModule,
-								id: id
-							},
-							0,
-							'module'
-						)
-						addEdge(failureState.parent_module, id, { type: 'empty' })
-					}
-				}
 			}
 		}
 
 		processModules(modules, inputNode, resultNode)
 
+		if (failureModule) {
+			let toAdd: Record<string, string> = {}
+			Object.keys(extra.flowModuleStates ?? {}).forEach((id) => {
+				if (id.startsWith('failure-')) {
+					const failureState = extra.flowModuleStates?.[id] as GraphModuleState | undefined
+					if (failureState?.parent_module) {
+						toAdd[failureState.parent_module] = id
+					}
+				}
+			})
+
+			Object.entries(toAdd).forEach((x) => {
+				addNode({ ...failureModule, id: x[1] }, 0, 'module')
+				addEdge(x[0], x[1], { type: 'empty' })
+			})
+		}
+
 		if (preprocessorModule) {
 			addNode(preprocessorModule, 0, 'module')
 			const id = JSON.parse(JSON.stringify(preprocessorModule.id))
diff --git a/frontend/src/lib/components/graph/renderers/nodes/branchOneEndNode.svelte b/frontend/src/lib/components/graph/renderers/nodes/branchOneEndNode.svelte
index ab446c74d3..3f9622dbf0 100644
--- a/frontend/src/lib/components/graph/renderers/nodes/branchOneEndNode.svelte
+++ b/frontend/src/lib/components/graph/renderers/nodes/branchOneEndNode.svelte
@@ -16,7 +16,7 @@
 
 <NodeWrapper let:darkMode offset={data.offset}>
 	<VirtualItem
-		label={'Collect result from all branches'}
+		label={'Collect result from chosen branch'}
 		id={data.id}
 		selectable={true}
 		selected={false}