From dbfc9df44f646e882f98953b18d13a4a0ec89848 Mon Sep 17 00:00:00 2001
From: Emir Karabeg <emirkarabeg@berkeley.edu>
Date: Tue, 11 Feb 2025 20:03:44 -0800
Subject: [PATCH] Added loop label

---
 .../components/loop-input/loop-input.tsx      |  0
 .../components/loop-label/loop-label.tsx      | 14 +++++++++++
 .../workflow-loop/workflow-loop.tsx           | 23 +++++++++++++++++--
 app/w/[id]/workflow.tsx                       | 15 ++++++++----
 4 files changed, 46 insertions(+), 6 deletions(-)
 create mode 100644 app/w/[id]/components/workflow-loop/components/loop-input/loop-input.tsx
 create mode 100644 app/w/[id]/components/workflow-loop/components/loop-label/loop-label.tsx

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
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/app/w/[id]/components/workflow-loop/components/loop-label/loop-label.tsx b/app/w/[id]/components/workflow-loop/components/loop-label/loop-label.tsx
new file mode 100644
index 0000000000..9a3c23c2f5
--- /dev/null
+++ b/app/w/[id]/components/workflow-loop/components/loop-label/loop-label.tsx
@@ -0,0 +1,14 @@
+import { NodeProps } from 'reactflow'
+import { Badge } from '@/components/ui/badge'
+
+export function LoopLabel({ data }: NodeProps) {
+  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 shadow-sm transition-colors duration-150"
+    >
+      {data.label}
+    </Badge>
+  )
+}
diff --git a/app/w/[id]/components/workflow-loop/workflow-loop.tsx b/app/w/[id]/components/workflow-loop/workflow-loop.tsx
index 592005e79e..0424069781 100644
--- a/app/w/[id]/components/workflow-loop/workflow-loop.tsx
+++ b/app/w/[id]/components/workflow-loop/workflow-loop.tsx
@@ -6,7 +6,20 @@ interface WorkflowLoopProps {
   blocks: Record<string, any>
 }
 
-// Pure calculation function - no hooks
+// Helper function to create loop label node
+function createLoopLabelNode(loopId: string, bounds: { x: number; y: number }) {
+  return {
+    id: `loop-label-${loopId}`,
+    type: 'loopLabel',
+    position: { x: 0, y: -32 },
+    parentNode: `loop-${loopId}`,
+    draggable: false,
+    data: {
+      label: 'Loop',
+    },
+  }
+}
+
 function calculateLoopBounds(loop: Loop, blocks: Record<string, any>) {
   // Get all blocks in this loop and filter out any undefined blocks
   const loopBlocks = loop.nodes
@@ -45,7 +58,7 @@ export function createLoopNode({ loopId, loop, blocks }: WorkflowLoopProps) {
   const loopBounds = calculateLoopBounds(loop, blocks)
   if (!loopBounds) return null
 
-  return {
+  const loopNode = {
     id: `loop-${loopId}`,
     type: 'group',
     position: { x: loopBounds.x, y: loopBounds.y },
@@ -63,6 +76,12 @@ export function createLoopNode({ loopId, loop, blocks }: WorkflowLoopProps) {
       label: 'Loop',
     },
   }
+
+  // Create the label node
+  const labelNode = createLoopLabelNode(loopId, loopBounds)
+
+  // Return both nodes
+  return [loopNode, labelNode]
 }
 
 // Helper function to calculate relative position for child blocks
diff --git a/app/w/[id]/workflow.tsx b/app/w/[id]/workflow.tsx
index f99b1493d6..5bf07acb38 100644
--- a/app/w/[id]/workflow.tsx
+++ b/app/w/[id]/workflow.tsx
@@ -20,10 +20,14 @@ 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 { LoopLabel } from './components/workflow-loop/components/loop-label/loop-label'
 import { createLoopNode, getRelativeLoopPosition } from './components/workflow-loop/workflow-loop'
 
 // Define custom node and edge types
-const nodeTypes: NodeTypes = { workflowBlock: WorkflowBlock }
+const nodeTypes: NodeTypes = {
+  workflowBlock: WorkflowBlock,
+  loopLabel: LoopLabel,
+}
 const edgeTypes: EdgeTypes = { custom: CustomEdge }
 
 function WorkflowContent() {
@@ -96,10 +100,13 @@ function WorkflowContent() {
   const nodes = useMemo(() => {
     const nodeArray: any[] = []
 
-    // Add loop group nodes
+    // Add loop group nodes and their labels
     Object.entries(loops).forEach(([loopId, loop]) => {
-      const loopNode = createLoopNode({ loopId, loop, blocks })
-      if (loopNode) nodeArray.push(loopNode)
+      const loopNodes = createLoopNode({ loopId, loop, blocks })
+      if (loopNodes) {
+        // Add both the loop node and its label node
+        nodeArray.push(...loopNodes)
+      }
     })
 
     // Add block nodes