Fix graph stall for collectors downstream of empty iterators

[JPPhoto]

Summary

Fixes graph execution stalling when a Collect node is downstream of an iterator over an empty collection.

Zero-iteration expansions are now explicitly recorded as complete. This allows downstream collectors to execute with an empty collection, typed consumers to receive [], and the session to report completion.

Collector output validation now handles materialized collectors with no runtime input edges. In this case, the already-validated source graph remains authoritative for type compatibility.

Related Issues / Discussions

Closes #9347

QA Instructions

Added coverage for:

• IntegerCollection([]) -> Iterate -> Add -> Collect -> typed consumer
• Nested empty Iterate -> Collect -> Iterate -> Collect chains
• Collector output and typed consumer both receiving []
• GraphExecutionState.is_complete() returning True

Merge Plan

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[lstein]

I ran a review of this branch at 495656c (multi-agent find + adversarial verify; every finding below was checked against the code, and the main one was reproduced by running it). The flat empty-iterator fix itself holds up well — completion logic, JSON round-trip/resume, event consumers, and the frontend all tolerate the empty-marked nodes (the If-branch skip path already set that precedent). But the fallback is wrong for nested iteration.

1. Confirmed bug: [((), [])] fallback produces wrong results when the empty collector sits inside a non-empty outer iterator

graph.py:704 — when _get_collect_iteration_mapping_groups returns [], the fallback hardcodes a single collect exec node at iteration path (). If the collector is nested under a non-empty outer iterator (outer iterates [a, b], inner iterator gets an empty collection for each outer item), the collector should materialize once per outer iteration at paths (0,) and (1,). Instead one node at () is created, downstream per-outer-iteration consumers run once instead of twice, and the outer collector under-counts.

Observed on this branch (repro below):

=== (a) ALL-EMPTY: outer=[a,b], inner=[] for both ===
is_complete: True
  inner_collect: 1 exec node(s), iteration paths=[()], results=[[]]
  per_outer_consumer: 1 exec node(s), iteration paths=[()], results=[[]]
  outer_collect: 1 exec node(s), iteration paths=[()], results=[[[]]]   <-- should be [[], []]

On main this graph stalls (the bug this PR fixes), so the PR converts a visible stall into a silently wrong result for the nested case. The fallback path set needs to be derived from the collector's enclosing iterators' prepared exec nodes rather than hardcoded to ().

Related but pre-existing (identical output on main, so not caused by this PR): in the mixed case — inner collection empty for outer item a, non-empty for b — _get_collect_iteration_mapping_groups returns only b's group, the fallback never fires, and branch a's collector instance is silently missing (outer_collect = [['b']] instead of [[], ['b']]). A path-aware fallback would likely fix both.

Repro script (run from repo root with PYTHONPATH=.)

"""Repro: nested iterator with empty inner collection under a non-empty outer iterator."""

from invokeai.app.invocations.baseinvocation import BaseInvocation, InvocationContext
from invokeai.app.services.shared.graph import (
    CollectInvocation,
    Graph,
    GraphExecutionState,
    IterateInvocation,
)
from tests.test_nodes import (
    PromptCollectionTestInvocation,
    PromptCollectionTestInvocationOutput,
    PromptTestInvocation,
    create_edge,
    run_session_with_mock_context,
)


class EmptyIfATestInvocation(BaseInvocation):
    """Outputs [] if prompt == 'a' (or always), else [prompt]."""

    prompt: str = ""
    always_empty: bool = True

    def invoke(self, context: InvocationContext) -> PromptCollectionTestInvocationOutput:
        if self.always_empty or self.prompt == "a":
            return PromptCollectionTestInvocationOutput(collection=[])
        return PromptCollectionTestInvocationOutput(collection=[self.prompt])


def build_and_run(always_empty: bool, label: str):
    g = Graph()
    n1 = PromptCollectionTestInvocation(id="outer_src", collection=["a", "b"])
    n2 = IterateInvocation(id="outer_iter")
    n3 = EmptyIfATestInvocation(id="inner_src", always_empty=always_empty)
    n4 = IterateInvocation(id="inner_iter")
    n5 = PromptTestInvocation(id="body")
    n6 = CollectInvocation(id="inner_collect")
    n7 = PromptCollectionTestInvocation(id="per_outer_consumer")  # runs per outer iteration
    n8 = CollectInvocation(id="outer_collect")

    for node in [n1, n2, n3, n4, n5, n6, n7, n8]:
        g.add_node(node)

    g.add_edge(create_edge(n1.id, "collection", n2.id, "collection"))
    g.add_edge(create_edge(n2.id, "item", n3.id, "prompt"))
    g.add_edge(create_edge(n3.id, "collection", n4.id, "collection"))
    g.add_edge(create_edge(n4.id, "item", n5.id, "prompt"))
    g.add_edge(create_edge(n5.id, "prompt", n6.id, "item"))
    g.add_edge(create_edge(n6.id, "collection", n7.id, "collection"))
    g.add_edge(create_edge(n7.id, "collection", n8.id, "item"))

    session = GraphExecutionState(graph=g)
    run_session_with_mock_context(session)

    print(f"\n=== {label} ===")
    print(f"is_complete: {session.is_complete()}")
    for src in ["inner_collect", "per_outer_consumer", "outer_collect"]:
        prepared = sorted(session.source_prepared_mapping.get(src, set()))
        paths = [session._get_iteration_path(p) for p in prepared]
        results = [getattr(session.results.get(p), "collection", None) for p in prepared]
        print(f"  {src}: {len(prepared)} exec node(s), iteration paths={paths}, results={results}")


if __name__ == "__main__":
    build_and_run(always_empty=True, label="(a) ALL-EMPTY: outer=[a,b], inner=[] for both")
    build_and_run(always_empty=False, label="(b) MIXED: inner=[] for 'a', ['b'] for 'b'")

2. Minor / non-blocking

• graph.py:1539 — the and self._get_input_edges(source_node.id) guard applies to user-authored source graphs too, not just materialized execution graphs. Adding collect.collection → consumer before any collector input edge is now silently accepted at add_edge time; the error only surfaces later, at enqueue via _validate_special_nodes (unrelaxed), or attributed to a subsequently-added input edge. Verified not a correctness hole — every path to execution re-validates, and nothing ever runs validate_self on the execution graph (including JSON round-trip resume) — but scoping the exemption to execution-graph wiring would preserve immediate source-graph validation. Also, _get_input_edges builds a full O(E) list just for truthiness on a path that runs per attached edge during materialization; any(e.destination.node_id == source_node.id for e in self.edges) short-circuits. (Folding the skip into _is_collector_connection_valid's zero-inputs check is not equivalent — _validate_special_nodes calls it with no new edge and would then pass inputless collectors at enqueue.)
• workflow_call_runtime.py:145 — behavior change worth knowing: a saved workflow whose workflow_return sits directly under a zero-iteration iterator (no intervening collector) now completes empty-marked, so the parent fails with ValueError("...unsupported number of workflow_return executions") instead of hanging forever. This is a strict improvement (the ≥2-iteration case already produced the same error), but a hint like "place a collector before workflow_return" would help.
• graph.py:704 — the empty-group fallback lives at the call site; the sibling _get_parent_iteration_mappings_without_iterators handles its zero-iteration default internally (if not iteration_paths: iteration_paths = [()]). Returning the empty group from _get_collect_iteration_mapping_groups itself would be behavior-identical and protect any future caller.
• graph.py:720 — with the new early-return above it, the None default in return next(iter(new_node_ids), None) is unreachable; return new_node_ids[0] says what it means. (The adjacent if create_results is not None guards are also dead — create_execution_node returns list[str], never None — but those predate this PR.)

Also verified as not problems, in case anyone else wonders: prepare() now returning a source-node id is inert (its only caller uses it as a loop-continuation sentinel); the if not new_node_ids catch-all can't fire for non-empty-iteration causes given the topological selection guards; and empty-marked nodes appearing in executed/executed_history without results entries is tolerated by every production reader (same as If-branch skips).

🤖 Generated with Claude Code

[JPPhoto]

@lstein Thanks for the review. I believe I've addressed those issues with the latest commits.

[Pfannkuchensack]

Findings

Low: invokeai/app/services/shared/graph.py:1539 weakens collector edge validation on the user-facing source graph, not just the execution graph

The new guard and self._get_input_edges(source_node.id) was added to Graph._validate_collector_edge_rules so a materialized empty collector (legitimately input-less in the execution graph) can have its output edge attached. But _validate_collector_edge_rules is a method on the shared Graph class and runs via add_edge -> _validate_edge for BOTH the execution graph and the user-submitted source graph.

Evidence chain:

1. Pre-patch, adding a collector output edge while the collector had zero input edges raised InvalidEdgeError("Collector must have at least one item or collection input edge") (reproduced by running the pre-patch condition against invokeai/app/services/shared/graph.py).
2. Post-patch, the same edge is silently accepted because self._get_input_edges(source_node.id) is empty/falsy, short-circuiting the entire collector-output validation.
3. Scenario: any caller/graph builder that attaches a collector's collection output before wiring the collector's inputs now bypasses the eager structural check at add_edge time.

This is bounded, not critical: validate_self() -> _validate_special_nodes() at invokeai/app/services/shared/graph.py:1424-1427 still calls _is_collector_connection_valid(node.id) for every collector, so a permanently input-less collector is still rejected at GraphExecutionState construction (graph_is_valid validator). The regression is that the per-edge structural guard is deferred to whole-graph validation, and the fix is applied at the wrong layer (shared Graph, affecting source graphs) rather than scoped to execution materialization.

To expose this issue, add a test that constructs a source Graph, calls add_edge for a collector's collection output before any collector input edge, and asserts the previously-raised InvalidEdgeError behavior is intentional; and a companion test asserting validate_self() still rejects a fully input-less collector so the boundary that bounds this finding is locked in.

Open Questions

• invokeai/app/services/shared/graph.py:716-718: _prepare_until_node_ready at invokeai/app/services/shared/graph.py:2063-2073 relies on prepare() returning a truthy value (now the source id via _mark_source_node_empty) to keep iterating and a falsy value to stop. This terminates correctly because _mark_source_node_empty inserts the node into source_prepared_mapping, excluding it from the next topological scan (invokeai/app/services/shared/graph.py:680). Confirmed no infinite loop in the tested cases, but there is no test asserting termination/idempotency when multiple empty branches exist in one graph.

[JPPhoto]

@Pfannkuchensack Thanks, I'd appreciate another review.