feat(workflow): enhance Loop node with dynamic iteration control

PriNova · PriNova · commit 6b06a3fd5dfc · 2025-02-14T11:57:22.000Z
- Add iterations-override input handle for external loop control
- Implement executeLoopStartNode for improved execution logic
- Add getLoopIterations helper for dynamic iteration counting
- Update node sorting to handle iteration override inputs
- Refactor loop initialization and state management
- Clean up IfElse node display formatting
- Add overrideIterations flag to Loop initialization

The changes enable dynamic adjustment of loop iterations through external inputs, improving workflow flexibility and control flow management.
diff --git a/vscode/package.json b/vscode/package.json
@@ -3,7 +3,7 @@
   "name": "cody-ai",
   "private": true,
   "displayName": "Cody: AI Code Assistant",
-  "version": "1.67.0+3",
+  "version": "1.67.0+4",
   "publisher": "sourcegraph",
   "license": "Apache-2.0",
   "icon": "resources/sourcegraph.png",
diff --git a/vscode/src/workflow/node-sorting.ts b/vscode/src/workflow/node-sorting.ts
@@ -572,8 +572,7 @@ function processLoopWithCycles(
         const preLoopNodeIds = new Set(preLoopNodes.map(node => node.id))
         const nodesInsideLoop = findLoopNodes(loopStart, nodes, edges, preLoopNodeIds)
 
-        // Process loop iterations
-        const iterations = shouldIterateLoops ? (loopStart as LoopStartNode).data.iterations : 1
+        const iterations = getLoopIterations(loopStart, nodes, edges, shouldIterateLoops)
 
         for (let i = 0; i < iterations; i++) {
             processedNodes.push({ ...loopStart })
@@ -597,3 +596,27 @@ function processLoopWithCycles(
 
     return processedNodes
 }
+
+function getLoopIterations(
+    loopStart: WorkflowNodes,
+    nodes: WorkflowNodes[],
+    edges: Edge[],
+    shouldIterateLoops: boolean
+): number {
+    if (!shouldIterateLoops) {
+        return 1
+    }
+
+    const iterationOverrideEdges = edges.filter(
+        edge => edge.target === loopStart.id && edge.targetHandle === 'iterations-override'
+    )
+
+    if (iterationOverrideEdges.length === 0) {
+        return (loopStart as LoopStartNode).data.iterations
+    }
+
+    const sourceNode = nodes.find(n => n.id === iterationOverrideEdges[0].source)
+    const overrideValue = Number.parseInt(sourceNode?.data.content || '', 10)
+
+    return !Number.isNaN(overrideValue) ? overrideValue : (loopStart as LoopStartNode).data.iterations
+}
diff --git a/vscode/src/workflow/workflow-executor.ts b/vscode/src/workflow/workflow-executor.ts
@@ -100,15 +100,6 @@ export async function executeWorkflow(
         type: 'execution_started',
     } as ExtensionToWorkflow)
 
-    for (const node of sortedNodes) {
-        if (node.type === NodeType.LOOP_START) {
-            context.loopStates.set(node.id, {
-                currentIteration: -1,
-                maxIterations: (node as LoopStartNode).data.iterations,
-                variable: (node as LoopStartNode).data.loopVariable,
-            })
-        }
-    }
     for (const node of sortedNodes) {
         if (skipNodes.has(node.id)) {
             continue
@@ -261,22 +252,7 @@ export async function executeWorkflow(
                     break
                 }
                 case NodeType.LOOP_START: {
-                    const inputs = combineParentOutputsByConnectionOrder(node.id, context)
-
-                    const loopState = context.loopStates.get(node.id) || {
-                        currentIteration: -1,
-                        maxIterations: (node as LoopStartNode).data.iterations,
-                        variable: (node as LoopStartNode).data.loopVariable,
-                    }
-
-                    // Only increment for next iteration if not at max
-                    if (loopState.currentIteration < loopState.maxIterations) {
-                        context.loopStates.set(node.id, {
-                            ...loopState,
-                            currentIteration: loopState.currentIteration + 1,
-                        })
-                    }
-                    result = await executePreviewNode(inputs.join('\n'), node.id, webview, context)
+                    result = await executeLoopStartNode(node as LoopStartNode, context)
                     break
                 }
                 case NodeType.LOOP_END: {
@@ -913,3 +889,51 @@ const commandsNotAllowed = [
     'lsusb',
     'lspci',
 ]
+
+async function executeLoopStartNode(
+    node: LoopStartNode,
+    context: IndexedExecutionContext
+): Promise<string> {
+    // Helper function to filter edges by handle type
+    const getInputsByHandle = (handleType: string) =>
+        combineParentOutputsByConnectionOrder(node.id, {
+            ...context,
+            edgeIndex: {
+                ...context.edgeIndex,
+                byTarget: new Map([
+                    [
+                        node.id,
+                        (context.edgeIndex.byTarget.get(node.id) || []).filter(
+                            edge => edge.targetHandle === handleType
+                        ),
+                    ],
+                ]),
+            },
+        })
+
+    const mainInputs = getInputsByHandle('main')
+    const iterationOverrides = getInputsByHandle('iterations-override')
+
+    let loopState = context.loopStates.get(node.id)
+
+    if (!loopState) {
+        const maxIterations =
+            iterationOverrides.length > 0
+                ? Number.parseInt(iterationOverrides[0], 10) || node.data.iterations
+                : node.data.iterations
+
+        loopState = {
+            currentIteration: 0, // Start at 0 for clearer iteration counting
+            maxIterations,
+            variable: node.data.loopVariable,
+        }
+        context.loopStates.set(node.id, loopState)
+    } else if (loopState.currentIteration < loopState.maxIterations - 1) {
+        context.loopStates.set(node.id, {
+            ...loopState,
+            currentIteration: loopState.currentIteration + 1,
+        })
+    }
+
+    return mainInputs.join('\n')
+}
diff --git a/vscode/webviews/workflow/components/nodes/IfElse_Node.tsx b/vscode/webviews/workflow/components/nodes/IfElse_Node.tsx
@@ -49,9 +49,12 @@ export const IfElseNode: React.FC<BaseNodeProps> = ({ data, selected }) => (
             >
                 IF...ELSE
             </div>
-            <div className="tw-flex tw-items-center tw-justify-center tw-p-2 tw-border tw-border-[var(--vscode-input-border)] tw-rounded">
-                <span>{data.condition || 'Condition'}</span>
+            <div className="tw-flex tw-items-center tw-justify-center">
+                <span>{data.title}</span>
             </div>
+            {/*<div className="tw-flex tw-items-center tw-justify-center tw-p-2 tw-border tw-border-[var(--vscode-input-border)] tw-rounded">
+                <span>{data.condition || 'Condition'}</span>
+            </div>*/}
             <div className="tw-flex tw-justify-between tw-mt-2">
                 <div className="tw-text-sm tw-opacity-75">True</div>
                 <div className="tw-text-sm tw-opacity-75">False</div>
diff --git a/vscode/webviews/workflow/components/nodes/LoopStart_Node.tsx b/vscode/webviews/workflow/components/nodes/LoopStart_Node.tsx
@@ -14,6 +14,7 @@ export type LoopStartNode = Omit<WorkflowNode, 'data'> & {
     data: BaseNodeData & {
         iterations: number
         loopVariable: string
+        overrideIterations?: boolean
     }
 }
 
@@ -33,6 +34,7 @@ export const LoopStartNode: React.FC<BaseNodeProps> = ({ data, selected }) => (
         }}
     >
         <Handle type="target" position={Position.Top} />
+        <Handle type="target" position={Position.Left} id="iterations-override" style={{ top: '83%' }} />
         <div className="tw-flex tw-flex-col tw-gap-2">
             <div
                 className="tw-text-center tw-py-1 tw-mb-2 tw-rounded-t-sm tw-font-bold"
@@ -44,7 +46,7 @@ export const LoopStartNode: React.FC<BaseNodeProps> = ({ data, selected }) => (
                         selected,
                         interrupted: data.interrupted,
                         active: data.active,
-                    })}`,
+                    })})`,
                     color: ' #1e1e1e',
                     marginLeft: '-0.5rem',
                     marginRight: '-0.5rem',
@@ -56,6 +58,7 @@ export const LoopStartNode: React.FC<BaseNodeProps> = ({ data, selected }) => (
             <div className="tw-flex tw-flex-col tw-justify-center">
                 <span>{data.title}</span>
                 <span className="tw-text-sm tw-opacity-70">Iterations: {data.iterations || 1}</span>
+                <span className="tw-text-xs tw-opacity-50">← Iterations Override</span>
             </div>
         </div>
         <Handle type="source" position={Position.Bottom} />
diff --git a/vscode/webviews/workflow/components/nodes/Nodes.tsx b/vscode/webviews/workflow/components/nodes/Nodes.tsx
@@ -141,6 +141,15 @@ export const createNode = (node: Omit<WorkflowNodes, 'id'>): WorkflowNodes => {
                     falsePathActive: false,
                 },
             } as IfElseNode
+        case NodeType.LOOP_START:
+            return {
+                ...node,
+                id,
+                data: {
+                    ...node.data,
+                    overrideIterations: false,
+                },
+            } as LoopStartNode
         default:
             return {
                 ...node,
