From 06a360b2c88f94a752226515a2b75cac0a49f7b8 Mon Sep 17 00:00:00 2001
From: Emir Karabeg <emirkarabeg@berkeley.edu>
Date: Tue, 11 Feb 2025 21:49:41 -0800
Subject: [PATCH] Added max iterations update to state for loop

---
 stores/workflow/store.ts | 23 +++++++++++++++++++++++
 stores/workflow/types.ts |  2 ++
 2 files changed, 25 insertions(+)

diff --git a/stores/workflow/store.ts b/stores/workflow/store.ts
index 3c0f03d71d..c5e9dffb19 100644
--- a/stores/workflow/store.ts
+++ b/stores/workflow/store.ts
@@ -114,6 +114,7 @@ export const useWorkflowStore = create<WorkflowStoreWithHistory>()(
             },
           }
         })
+        get().updateLastSaved()
       },
 
       addBlock: (id: string, type: string, name: string, position: Position) => {
@@ -229,6 +230,7 @@ export const useWorkflowStore = create<WorkflowStoreWithHistory>()(
               newLoops[loopId] = {
                 id: loopId,
                 nodes: path,
+                maxIterations: 5,
               }
               processedPaths.add(canonicalPath)
             }
@@ -265,6 +267,7 @@ export const useWorkflowStore = create<WorkflowStoreWithHistory>()(
               newLoops[loopId] = {
                 id: loopId,
                 nodes: path,
+                maxIterations: 5,
               }
               processedPaths.add(canonicalPath)
             }
@@ -414,6 +417,7 @@ export const useWorkflowStore = create<WorkflowStoreWithHistory>()(
           edges: [...state.edges],
           loops: { ...get().loops },
         }))
+        get().updateLastSaved()
       },
 
       updateBlockHeight: (id: string, height: number) => {
@@ -427,6 +431,25 @@ export const useWorkflowStore = create<WorkflowStoreWithHistory>()(
           },
           edges: [...state.edges],
         }))
+        get().updateLastSaved()
+      },
+
+      updateLoopMaxIterations: (loopId: string, maxIterations: number) => {
+        const newState = {
+          blocks: { ...get().blocks },
+          edges: [...get().edges],
+          loops: {
+            ...get().loops,
+            [loopId]: {
+              ...get().loops[loopId],
+              maxIterations: Math.max(1, Math.min(50, maxIterations)), // Clamp between 1-50
+            },
+          },
+        }
+
+        set(newState)
+        pushHistory(set, get, newState, 'Update loop max iterations')
+        get().updateLastSaved()
       },
     })),
     { name: 'workflow-store' }
diff --git a/stores/workflow/types.ts b/stores/workflow/types.ts
index a6d0c4de61..5847f7cf20 100644
--- a/stores/workflow/types.ts
+++ b/stores/workflow/types.ts
@@ -28,6 +28,7 @@ export interface SubBlockState {
 export interface Loop {
   id: string
   nodes: string[]
+  maxIterations: number
 }
 
 export interface WorkflowState {
@@ -52,6 +53,7 @@ export interface WorkflowActions {
   updateBlockName: (id: string, name: string) => void
   toggleBlockWide: (id: string) => void
   updateBlockHeight: (id: string, height: number) => void
+  updateLoopMaxIterations: (loopId: string, maxIterations: number) => void
 }
 
 export type WorkflowStore = WorkflowState & WorkflowActions