fix: include downstream nodes in loop iteration so observers collect each pass

Nodes downstream of the loop body (e.g. outputGallery connected to a looped
generateVideo) were classified as suffix and only executed once after all
iterations. Now the loop iteration set expands via forward-edge traversal
to include downstream observers, so they execute each pass and can accumulate
results (e.g. collecting each generated video into the gallery).

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

src/store/__tests__/loopEdge.integration.test.ts

@@ -119,39 +119,48 @@ describe("executeWorkflow with loop edges", () => {
     copyLoopOutputSpy.mockRestore();
   });
 
-  it("executes non-loop prefix and suffix nodes once", async () => {
+  it("executes prefix nodes once and downstream nodes each iteration", async () => {
-    // Create PREFIX(imageInput) → A(prompt) → B(prompt) → SUFFIX(output)
+    // Create PREFIX(imageInput) → A(prompt) → B(prompt) → DOWNSTREAM(output)
     // with loop edge B→A (loopCount: 3)
+    // PREFIX is not downstream of the loop body — it only feeds INTO A
+    // DOWNSTREAM is connected from B (loop body) — it should iterate with the loop
     const prefix = makeNode("prefix", "imageInput", { image: "data:image/png;base64,test" });
     const nodeA = makeNode("a", "prompt", { prompt: "loop start" });
     const nodeB = makeNode("b", "prompt", { prompt: "loop end" });
-    const suffix = makeNode("suffix", "output");
+    const downstream = makeNode("downstream", "output");
 
     const edges = [
       makeEdge("prefix", "a"),
       makeEdge("a", "b"),
-      makeEdge("b", "suffix"),
+      makeEdge("b", "downstream"),
       makeEdge("b", "a", { isLoop: true, loopCount: 3 }),
     ];
 
-    setupStore([prefix, nodeA, nodeB, suffix], edges);
+    setupStore([prefix, nodeA, nodeB, downstream], edges);
-
+
-    // Spy on updateNodeData to track execution
+    // Track which nodes are executed by subscribing to currentNodeIds changes.
-    const updateNodeDataSpy = vi.spyOn(useWorkflowStore.getState(), "updateNodeData");
+    // Only push when the reference changes (Zustand preserves refs for unchanged fields).
+    const executedNodeIds: string[] = [];
+    let prevRef: string[] = [];
+    const unsubscribe = useWorkflowStore.subscribe((state) => {
+      if (state.currentNodeIds !== prevRef && state.currentNodeIds.length > 0) {
+        executedNodeIds.push(...state.currentNodeIds);
+        prevRef = state.currentNodeIds;
+      }
+    });
 
     await useWorkflowStore.getState().executeWorkflow();
+    unsubscribe();
+
+    const countExecutions = (id: string) => executedNodeIds.filter((nid) => nid === id).length;
 
-    // Check that prefix and suffix nodes were updated less frequently than loop nodes
+    // Prefix feeds into the loop but is not downstream — executes once
-    // (This is a simplified check - in real execution, loop nodes would be updated more)
+    expect(countExecutions("prefix")).toBe(1);
-    const prefixUpdates = updateNodeDataSpy.mock.calls.filter((call) => call[0] === "prefix");
+    // Loop body nodes execute 3 times (loopCount: 3)
-    const suffixUpdates = updateNodeDataSpy.mock.calls.filter((call) => call[0] === "suffix");
+    expect(countExecutions("a")).toBe(3);
-    const loopNodeUpdates = updateNodeDataSpy.mock.calls.filter(
+    expect(countExecutions("b")).toBe(3);
-      (call) => call[0] === "a" || call[0] === "b"
+    // Downstream is connected from loop body node B — iterates with the loop
-    );
+    expect(countExecutions("downstream")).toBe(3);
-
-    // Prefix and suffix should have fewer updates than loop nodes
-    expect(prefixUpdates.length).toBeLessThan(loopNodeUpdates.length);
-    expect(suffixUpdates.length).toBeLessThan(loopNodeUpdates.length);
   });
 
   it("does not affect workflows with no loop edges", async () => {

src/store/workflowStore.ts

@@ -1497,33 +1497,34 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
         const loopEdge = loopEdges[0];
         const loopBodyIds = new Set(findLoopSubgraph(loopEdge.source, loopEdge.target, forwardEdges));
 
-        // Partition levels into prefix (before loop), loop body, and suffix (after loop)
+        // Expand loop body to include downstream nodes that depend on loop output.
+        // These nodes (e.g. outputGallery connected to a looped generateVideo) should
+        // execute each iteration so they can collect results from every pass.
+        const loopIterationIds = new Set(loopBodyIds);
+        let expanded = true;
+        while (expanded) {
+          expanded = false;
+          for (const edge of forwardEdges) {
+            if (loopIterationIds.has(edge.source) && !loopIterationIds.has(edge.target)) {
+              loopIterationIds.add(edge.target);
+              expanded = true;
+            }
+          }
+        }
+
+        // Partition levels: iterating nodes run each loop pass, non-iterating run once as prefix
         const prefixLevels: typeof levels = [];
         const loopLevels: typeof levels = [];
-        const suffixLevels: typeof levels = [];
-
-        let foundLoop = false;
 
         for (const level of levels) {
-          const hasLoopNode = level.nodeIds.some(id => loopBodyIds.has(id));
+          const iteratingIds = level.nodeIds.filter(id => loopIterationIds.has(id));
-
+          const nonIteratingIds = level.nodeIds.filter(id => !loopIterationIds.has(id));
-          if (!foundLoop && !hasLoopNode) {
+
-            prefixLevels.push(level);
+          if (iteratingIds.length > 0) {
-          } else if (hasLoopNode) {
+            loopLevels.push({ level: level.level, nodeIds: iteratingIds });
-            foundLoop = true;
+          }
-            // Filter level to only include loop body nodes
+          if (nonIteratingIds.length > 0) {
-            const loopNodeIds = level.nodeIds.filter(id => loopBodyIds.has(id));
+            prefixLevels.push({ level: level.level, nodeIds: nonIteratingIds });
-            const nonLoopNodeIds = level.nodeIds.filter(id => !loopBodyIds.has(id));
-
-            if (loopNodeIds.length > 0) {
-              loopLevels.push({ level: level.level, nodeIds: loopNodeIds });
-            }
-            // Non-loop nodes at the same level as loop nodes go to suffix
-            if (nonLoopNodeIds.length > 0) {
-              suffixLevels.push({ level: level.level, nodeIds: nonLoopNodeIds });
-            }
-          } else if (foundLoop) {
-            suffixLevels.push(level);
           }
         }
 
@@ -1554,15 +1555,9 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
 
           logger.info('node.execution', `Loop iteration ${i + 1}/${loopCount}`);
 
-          // Execute loop body levels with fresh node state
+          // Execute loop body + downstream levels with fresh node state
           await executeLevels(loopLevels);
         }
-
-        if (!abortController.signal.aborted) {
-          // Execute suffix once
-          logger.info('node.execution', 'Executing suffix levels', { count: suffixLevels.length });
-          await executeLevels(suffixLevels);
-        }
       }
 
       // Check if we completed or were aborted