From 3f49b7c293f41f8418e7c856ecdd4109cdf42733 Mon Sep 17 00:00:00 2001 From: shrimbly Date: Fri, 3 Apr 2026 00:18:40 +1300 Subject: [PATCH] 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 --- .../__tests__/loopEdge.integration.test.ts | 47 +++++++++------- src/store/workflowStore.ts | 55 +++++++++---------- 2 files changed, 53 insertions(+), 49 deletions(-) diff --git a/src/store/__tests__/loopEdge.integration.test.ts b/src/store/__tests__/loopEdge.integration.test.ts index 40e886f5..3cccca45 100644 --- a/src/store/__tests__/loopEdge.integration.test.ts +++ b/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 () => { - // Create PREFIX(imageInput) → A(prompt) → B(prompt) → SUFFIX(output) + it("executes prefix nodes once and downstream nodes each iteration", async () => { + // 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); - - // Spy on updateNodeData to track execution - const updateNodeDataSpy = vi.spyOn(useWorkflowStore.getState(), "updateNodeData"); + setupStore([prefix, nodeA, nodeB, downstream], edges); + + // Track which nodes are executed by subscribing to currentNodeIds changes. + // 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 - // (This is a simplified check - in real execution, loop nodes would be updated more) - const prefixUpdates = updateNodeDataSpy.mock.calls.filter((call) => call[0] === "prefix"); - const suffixUpdates = updateNodeDataSpy.mock.calls.filter((call) => call[0] === "suffix"); - const loopNodeUpdates = updateNodeDataSpy.mock.calls.filter( - (call) => call[0] === "a" || call[0] === "b" - ); - - // Prefix and suffix should have fewer updates than loop nodes - expect(prefixUpdates.length).toBeLessThan(loopNodeUpdates.length); - expect(suffixUpdates.length).toBeLessThan(loopNodeUpdates.length); + // Prefix feeds into the loop but is not downstream — executes once + expect(countExecutions("prefix")).toBe(1); + // Loop body nodes execute 3 times (loopCount: 3) + expect(countExecutions("a")).toBe(3); + expect(countExecutions("b")).toBe(3); + // Downstream is connected from loop body node B — iterates with the loop + expect(countExecutions("downstream")).toBe(3); }); it("does not affect workflows with no loop edges", async () => { diff --git a/src/store/workflowStore.ts b/src/store/workflowStore.ts index b839666e..a4a0775f 100644 --- a/src/store/workflowStore.ts +++ b/src/store/workflowStore.ts @@ -1497,33 +1497,34 @@ const workflowStoreImpl: StateCreator = (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)); - - if (!foundLoop && !hasLoopNode) { - prefixLevels.push(level); - } else if (hasLoopNode) { - foundLoop = true; - // Filter level to only include loop body nodes - const loopNodeIds = level.nodeIds.filter(id => loopBodyIds.has(id)); - 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); + const iteratingIds = level.nodeIds.filter(id => loopIterationIds.has(id)); + const nonIteratingIds = level.nodeIds.filter(id => !loopIterationIds.has(id)); + + if (iteratingIds.length > 0) { + loopLevels.push({ level: level.level, nodeIds: iteratingIds }); + } + if (nonIteratingIds.length > 0) { + prefixLevels.push({ level: level.level, nodeIds: nonIteratingIds }); } } @@ -1554,15 +1555,9 @@ const workflowStoreImpl: StateCreator = (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