From bddc6cb7f79b18296b8bae34c26c5634ce30c025 Mon Sep 17 00:00:00 2001
From: Emir Karabeg <emirkarabeg@berkeley.edu>
Date: Tue, 11 Feb 2025 14:49:33 -0800
Subject: [PATCH] Fixed delete block in loop error

---
 .../components/workflow-loop/workflow-loop.tsx | 10 ++++++++--
 app/w/[id]/workflow.tsx                        |  1 -
 components/ui/command.tsx                      | 15 ---------------
 stores/workflow/store.ts                       | 18 ++++++++++++++++++
 4 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/app/w/[id]/components/workflow-loop/workflow-loop.tsx b/app/w/[id]/components/workflow-loop/workflow-loop.tsx
index 18bda5c36f..592005e79e 100644
--- a/app/w/[id]/components/workflow-loop/workflow-loop.tsx
+++ b/app/w/[id]/components/workflow-loop/workflow-loop.tsx
@@ -8,8 +8,14 @@ interface WorkflowLoopProps {
 
 // Pure calculation function - no hooks
 function calculateLoopBounds(loop: Loop, blocks: Record<string, any>) {
-  // Get all blocks in this loop
-  const loopBlocks = loop.nodes.map((id) => blocks[id])
+  // Get all blocks in this loop and filter out any undefined blocks
+  const loopBlocks = loop.nodes
+    .map((id) => blocks[id])
+    .filter(
+      (block): block is NonNullable<typeof block> =>
+        block !== undefined && block.position !== undefined
+    )
+
   if (!loopBlocks.length) return null
 
   // Calculate bounds of all blocks in loop
diff --git a/app/w/[id]/workflow.tsx b/app/w/[id]/workflow.tsx
index ac621008d7..f99b1493d6 100644
--- a/app/w/[id]/workflow.tsx
+++ b/app/w/[id]/workflow.tsx
@@ -18,7 +18,6 @@ import { useWorkflowRegistry } from '@/stores/workflow/registry'
 import { useWorkflowStore } from '@/stores/workflow/store'
 import { NotificationList } from '@/app/w/components/notifications/notifications'
 import { getBlock } from '../../../blocks'
-import { useWorkflowExecution } from '../hooks/use-workflow-execution'
 import { CustomEdge } from './components/custom-edge/custom-edge'
 import { WorkflowBlock } from './components/workflow-block/workflow-block'
 import { createLoopNode, getRelativeLoopPosition } from './components/workflow-loop/workflow-loop'
diff --git a/components/ui/command.tsx b/components/ui/command.tsx
index d6b78224e2..b8adf23b03 100644
--- a/components/ui/command.tsx
+++ b/components/ui/command.tsx
@@ -12,21 +12,6 @@ import { cn } from '@/lib/utils'
 // This file is not typed correctly from shadcn, so we're disabling the type checker
 // @ts-nocheck
 
-// This file is not typed correctly from shadcn, so we're disabling the type checker
-// @ts-nocheck
-
-// This file is not typed correctly from shadcn, so we're disabling the type checker
-// @ts-nocheck
-
-// This file is not typed correctly from shadcn, so we're disabling the type checker
-// @ts-nocheck
-
-// This file is not typed correctly from shadcn, so we're disabling the type checker
-// @ts-nocheck
-
-// This file is not typed correctly from shadcn, so we're disabling the type checker
-// @ts-nocheck
-
 const Command = React.forwardRef<
   React.ElementRef<typeof CommandPrimitive>,
   React.ComponentPropsWithoutRef<typeof CommandPrimitive> & {
diff --git a/stores/workflow/store.ts b/stores/workflow/store.ts
index b05c5c019d..a018965d62 100644
--- a/stores/workflow/store.ts
+++ b/stores/workflow/store.ts
@@ -175,6 +175,24 @@ export const useWorkflowStore = create<WorkflowStoreWithHistory>()(
           edges: [...get().edges].filter((edge) => edge.source !== id && edge.target !== id),
           loops: { ...get().loops },
         }
+
+        // Remove the block from any loops that contain it
+        Object.entries(newState.loops).forEach(([loopId, loop]) => {
+          if (loop.nodes.includes(id)) {
+            // If the loop would only have 1 or 0 nodes after removal, delete the loop
+            if (loop.nodes.length <= 2) {
+              delete newState.loops[loopId]
+            } else {
+              // Otherwise, just remove the node from the loop
+              newState.loops[loopId] = {
+                ...loop,
+                nodes: loop.nodes.filter((nodeId) => nodeId !== id)
+              }
+            }
+          }
+        })
+
+        // Delete the block itself
         delete newState.blocks[id]
 
         set(newState)