diff --git a/frontend/src/lib/components/graph/FlowGraphV2.svelte b/frontend/src/lib/components/graph/FlowGraphV2.svelte
index 0c8feb80e2..bbc102dece 100644
--- a/frontend/src/lib/components/graph/FlowGraphV2.svelte
+++ b/frontend/src/lib/components/graph/FlowGraphV2.svelte
@@ -105,14 +105,20 @@
 		branchAllEnd: BranchAllEndNode,
 		forLoopEnd: ForLoopEndNode,
 		forLoopStart: ForLoopStartNode,
-		result: ResultNode
+		result: ResultNode,
+		whileLoopStart: ForLoopStartNode,
+		whileLoopEnd: ForLoopEndNode
 	}
 	const edgeTypes = {
 		edge: BaseEdge
 	}
 </script>
 
-<div class="h-screen">
+<div
+	style={`min-height: ${minHeight}px; max-height: ${
+		maxHeight ? maxHeight + 'px' : 'none'
+	}; height:100%;`}
+>
 	<SvelteFlow
 		{nodes}
 		{edges}
diff --git a/frontend/src/lib/components/graph/graphBuilder.ts b/frontend/src/lib/components/graph/graphBuilder.ts
index 58352d5d9f..4a170d4f25 100644
--- a/frontend/src/lib/components/graph/graphBuilder.ts
+++ b/frontend/src/lib/components/graph/graphBuilder.ts
@@ -135,6 +135,7 @@ export default function graphBuilder(modules: FlowModule[] | undefined): {
 					position: { x: 0, y: 0 },
 					type: 'whileLoopStart'
 				}
+				addEdge(module.id, startNode.id)
 
 				const endNode = {
 					id: `${module.id}-end`,
@@ -147,6 +148,7 @@ export default function graphBuilder(modules: FlowModule[] | undefined): {
 				nodes.push(endNode)
 
 				processModules(module.value.modules, startNode, endNode, currentOffset + 50)
+				previousId = endNode.id
 			} else if (module.value.type === 'branchone') {
 				// Start
 				addNode(module, currentOffset, 'module')
@@ -179,12 +181,10 @@ export default function graphBuilder(modules: FlowModule[] | undefined): {
 			}
 
 			if (index === 0) {
-				console.log('Adding edge between beforeNode and module', beforeNode.id, module.id)
 				addEdge(beforeNode.id, module.id)
 			}
 
 			if (index === modules.length - 1 && previousId) {
-				console.log(nextNode, previousId)
 				addEdge(previousId, nextNode.id)
 			}
 		})