From 1f9acf313e732bdc38084dd61d8ce0bf24052dbf Mon Sep 17 00:00:00 2001
From: Emir Karabeg <emirkarabeg@berkeley.edu>
Date: Wed, 12 Feb 2025 00:43:28 -0800
Subject: [PATCH] Feature: added max iterations to loop UI; will add adjustable

---
 .../components/loop-input/loop-input.tsx      | 21 ++++++++++++++++
 .../workflow-loop/workflow-loop.tsx           | 25 ++++++++++++++++---
 app/w/[id]/workflow.tsx                       |  2 ++
 3 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/app/w/[id]/components/workflow-loop/components/loop-input/loop-input.tsx b/app/w/[id]/components/workflow-loop/components/loop-input/loop-input.tsx
index e69de29bb2..42d2bac844 100644
--- a/app/w/[id]/components/workflow-loop/components/loop-input/loop-input.tsx
+++ b/app/w/[id]/components/workflow-loop/components/loop-input/loop-input.tsx
@@ -0,0 +1,21 @@
+import { NodeProps } from 'reactflow'
+import { Badge } from '@/components/ui/badge'
+import { useWorkflowStore } from '@/stores/workflow/store'
+
+export function LoopInput({ id }: NodeProps) {
+  // Extract the loop ID from the node ID (removes 'loop-input-' prefix)
+  const loopId = id.replace('loop-input-', '')
+
+  // Get the max iterations from the store for this loop
+  const maxIterations = useWorkflowStore((state) => state.loops[loopId]?.maxIterations ?? 5)
+
+  return (
+    <Badge
+      variant="outline"
+      className="bg-white border-[rgb(203,213,225)] text-gray-700 font-medium px-2 py-0.5 text-sm 
+      hover:bg-gray-50 transition-colors duration-150"
+    >
+      Max Iterations: {maxIterations}
+    </Badge>
+  )
+}
diff --git a/app/w/[id]/components/workflow-loop/workflow-loop.tsx b/app/w/[id]/components/workflow-loop/workflow-loop.tsx
index 2ffad7082f..2ac4bcf526 100644
--- a/app/w/[id]/components/workflow-loop/workflow-loop.tsx
+++ b/app/w/[id]/components/workflow-loop/workflow-loop.tsx
@@ -20,6 +20,22 @@ function createLoopLabelNode(loopId: string, bounds: { x: number; y: number }) {
   }
 }
 
+// Helper function to create loop input node
+function createLoopInputNode(loopId: string, bounds: { x: number; width: number }) {
+  const BADGE_WIDTH = 128
+
+  return {
+    id: `loop-input-${loopId}`,
+    type: 'loopInput',
+    position: { x: bounds.width - BADGE_WIDTH, y: -32 }, // Position from right edge
+    parentNode: `loop-${loopId}`,
+    draggable: false,
+    data: {
+      loopId,
+    },
+  }
+}
+
 function calculateLoopBounds(loop: Loop, blocks: Record<string, any>) {
   // Get all blocks in this loop and filter out any undefined blocks
   const loopBlocks = loop.nodes
@@ -64,7 +80,7 @@ function calculateLoopBounds(loop: Loop, blocks: Record<string, any>) {
   }
 }
 
-// Helper function to create loop node
+// Update the createLoopNode function
 export function createLoopNode({ loopId, loop, blocks }: WorkflowLoopProps) {
   const loopBounds = calculateLoopBounds(loop, blocks)
   if (!loopBounds) return null
@@ -88,11 +104,12 @@ export function createLoopNode({ loopId, loop, blocks }: WorkflowLoopProps) {
     },
   }
 
-  // Create the label node
+  // Create both label and input nodes
   const labelNode = createLoopLabelNode(loopId, loopBounds)
+  const inputNode = createLoopInputNode(loopId, loopBounds)
 
-  // Return both nodes
-  return [loopNode, labelNode]
+  // Return all three nodes
+  return [loopNode, labelNode, inputNode]
 }
 
 // Helper function to calculate relative position for child blocks
diff --git a/app/w/[id]/workflow.tsx b/app/w/[id]/workflow.tsx
index 5bf07acb38..45addefdf9 100644
--- a/app/w/[id]/workflow.tsx
+++ b/app/w/[id]/workflow.tsx
@@ -20,6 +20,7 @@ import { NotificationList } from '@/app/w/components/notifications/notifications
 import { getBlock } from '../../../blocks'
 import { CustomEdge } from './components/custom-edge/custom-edge'
 import { WorkflowBlock } from './components/workflow-block/workflow-block'
+import { LoopInput } from './components/workflow-loop/components/loop-input/loop-input'
 import { LoopLabel } from './components/workflow-loop/components/loop-label/loop-label'
 import { createLoopNode, getRelativeLoopPosition } from './components/workflow-loop/workflow-loop'
 
@@ -27,6 +28,7 @@ import { createLoopNode, getRelativeLoopPosition } from './components/workflow-l
 const nodeTypes: NodeTypes = {
   workflowBlock: WorkflowBlock,
   loopLabel: LoopLabel,
+  loopInput: LoopInput,
 }
 const edgeTypes: EdgeTypes = { custom: CustomEdge }