Tighten lazy If runtime edge-case handling

Author: JPPhoto

invokeai/app/services/shared/README.md

@@ -96,7 +96,10 @@ mutation helpers. Those helpers reject changes once the affected nodes have alre
 ### 4.1 Data
 
 - `graph: Graph` - source graph for the run; treated as stable during normal execution.
-- `execution_graph: Graph` - materialized runtime nodes/edges.
+- `execution_graph: Graph` - materialized runtime nodes/edges. This is mutable runtime state, not an immutable audit
+  log. Lazy `If` pruning may remove unselected input edges during execution, so persisted failed/completed session
+  snapshots can contain a structurally pruned execution graph. Retry paths rebuild from `graph`, not from a previously
+  persisted `execution_graph`.
 - `executed: set[str]`, `executed_history: list[str]`.
 - `results: dict[str, AnyInvocationOutput]`, `errors: dict[str, str]`.
 - `prepared_source_mapping: dict[str, str]` - exec id -> source id.
@@ -179,9 +182,9 @@ Run `C` -> `D:0` -> enqueue `D`. Run `D` -> done.
 - For **CollectInvocation**: gather all incoming `item` values into `collection`, sorting inputs by iteration path so
   collected results are stable across expanded iterations. Incoming `collection` values are merged first, then incoming
   `item` values are appended.
-- For **IfInvocation**: hydrate only `condition` and the selected branch input. If the selected branch's upstream exec
-  node was skipped and therefore produced no runtime output, the branch input is left at its default value (typically
-  `None`) instead of raising during hydration.
+- For **IfInvocation**: hydrate only `condition` and the selected branch input. As a defensive guard against
+  inconsistent runtime or deserialized session state, the runtime raises if the selected input edge points at an exec
+  node with no stored runtime output. In normal scheduling this path should be unreachable.
 - For all others: deep-copy each incoming edge's value into the destination field. This prevents cross-node mutation
   through shared references.
 

invokeai/app/services/shared/graph.py

@@ -243,6 +243,9 @@ def is_deferred_by_unresolved_if(self, exec_node_id: str) -> bool:
         return False
 
     def mark_exec_node_skipped(self, exec_node_id: str) -> None:
+        if self._state._get_prepared_exec_metadata(exec_node_id).state == "executed":
+            return
+
         self._state._remove_from_ready_queues(exec_node_id)
         self._state._set_prepared_exec_state(exec_node_id, "skipped")
         self._state.executed.add(exec_node_id)
@@ -356,7 +359,7 @@ def _initialize_execution_node(self, exec_node_id: str) -> None:
     def _get_collect_iteration_mappings(self, parent_node_ids: list[str]) -> list[tuple[str, str]]:
         all_iteration_mappings: list[tuple[str, str]] = []
         for source_node_id in parent_node_ids:
-            prepared_nodes = self._state.source_prepared_mapping[source_node_id]
+            prepared_nodes = self._get_prepared_nodes_for_source(source_node_id)
             all_iteration_mappings.extend((source_node_id, prepared_id) for prepared_id in prepared_nodes)
         return all_iteration_mappings
 
@@ -475,10 +478,12 @@ def get_iteration_node(
         prepared_iterator_nodes: list[str],
     ) -> Optional[str]:
         prepared_nodes = self._get_prepared_nodes_for_source(source_node_id)
-        if len(prepared_nodes) == 1:
-            return next(iter(prepared_nodes))
-
         parent_iterators = self._get_parent_iterator_exec_nodes(source_node_id, graph, prepared_iterator_nodes)
+        if len(prepared_nodes) == 1:
+            prepared_node_id = next(iter(prepared_nodes))
+            if self._matches_parent_iterators(prepared_node_id, parent_iterators, execution_graph):
+                return prepared_node_id
+            return None
 
         direct_iterator_match = self._get_direct_prepared_iterator_match(
             prepared_nodes, prepared_iterator_nodes, parent_iterators, execution_graph
@@ -778,9 +783,9 @@ def _prepare_if_inputs(self, node: IfInvocation, input_edges: list[Edge]) -> Non
 
             found_value, copied_value = self._try_get_copied_result_value(edge)
             if not found_value:
-                # A skipped branch-local exec node is considered executed for scheduling purposes, but it does not
-                # produce an output payload. Leave the optional branch input at its default None instead of crashing.
-                continue
+                raise RuntimeError(
+                    "IfInvocation selected input edge points at an exec node with no stored result output"
+                )
 
             setattr(node, edge.destination.field, copied_value)
 

tests/test_graph_execution_state.py

@@ -8,7 +8,12 @@
 from invokeai.app.invocations.collections import RangeInvocation
 from invokeai.app.invocations.logic import IfInvocation, IfInvocationOutput
 from invokeai.app.invocations.math import AddInvocation, MultiplyInvocation
-from invokeai.app.invocations.primitives import BooleanCollectionInvocation, BooleanInvocation, BooleanOutput
+from invokeai.app.invocations.primitives import (
+    BooleanCollectionInvocation,
+    BooleanCollectionOutput,
+    BooleanInvocation,
+    BooleanOutput,
+)
 from invokeai.app.services.shared.graph import (
     CollectInvocation,
     Graph,
@@ -886,7 +891,7 @@ def test_if_graph_optimized_behavior_handles_selected_false_branch_with_shared_t
     assert "true_collect" not in executed_source_ids
 
 
-def test_prepare_if_inputs_ignores_selected_branch_sources_without_results():
+def test_prepare_if_inputs_raises_when_selected_branch_source_has_no_result():
     graph = Graph()
     graph.add_node(BooleanInvocation(id="condition", value=True))
     graph.add_node(PromptTestInvocation(id="true_value", prompt="true branch"))
@@ -910,10 +915,23 @@ def test_prepare_if_inputs_ignores_selected_branch_sources_without_results():
     g._resolved_if_exec_branches[if_exec_id] = "true_input"
 
     if_node = g.execution_graph.get_node(if_exec_id)
-    g._prepare_inputs(if_node)
+    with pytest.raises(RuntimeError, match="selected input edge"):
+        g._prepare_inputs(if_node)
+
+
+def test_get_collect_iteration_mappings_ignores_skipped_prepared_exec_nodes():
+    graph = Graph()
+    graph.add_node(AnyTypeTestInvocation(id="parent", value="value"))
+
+    g = GraphExecutionState(graph=graph)
+
+    skipped_exec_id = g._create_execution_node("parent", [])[0]
+    active_exec_id = g._create_execution_node("parent", [])[0]
+    g._set_prepared_exec_state(skipped_exec_id, "skipped")
+
+    mappings = g._materializer()._get_collect_iteration_mappings(["parent"])
 
-    assert if_node.condition is True
-    assert if_node.true_input is None
+    assert mappings == [("parent", active_exec_id)]
 
 
 def test_get_iteration_node_ignores_skipped_prepared_exec_nodes():
@@ -958,6 +976,65 @@ def test_get_iteration_node_returns_none_when_only_skipped_prepared_exec_nodes_e
     assert selected_exec_id is None
 
 
+def test_get_iteration_node_does_not_reuse_wrong_iterator_when_only_other_iteration_is_live():
+    graph = Graph()
+    graph.add_node(BooleanCollectionInvocation(id="conditions", collection=[True, False]))
+    graph.add_node(IterateInvocation(id="condition_iter"))
+    graph.add_node(AnyTypeTestInvocation(id="value"))
+
+    graph.add_edge(create_edge("conditions", "collection", "condition_iter", "collection"))
+    graph.add_edge(create_edge("condition_iter", "item", "value", "value"))
+
+    g = GraphExecutionState(graph=graph)
+
+    conditions_exec_id = g._create_execution_node("conditions", [])[0]
+    g.executed.add(conditions_exec_id)
+    g.results[conditions_exec_id] = BooleanCollectionOutput(collection=[True, False])
+
+    iterator_exec_ids = g._create_execution_node("condition_iter", [("conditions", conditions_exec_id)])
+    assert len(iterator_exec_ids) == 2
+    iterator_exec_ids_by_index = {g.execution_graph.get_node(exec_id).index: exec_id for exec_id in iterator_exec_ids}
+    first_iter_exec_id = iterator_exec_ids_by_index[0]
+    second_iter_exec_id = iterator_exec_ids_by_index[1]
+
+    value_exec_ids = []
+    value_exec_ids.extend(g._create_execution_node("value", [("condition_iter", first_iter_exec_id)]))
+    value_exec_ids.extend(g._create_execution_node("value", [("condition_iter", second_iter_exec_id)]))
+    assert len(value_exec_ids) == 2
+
+    for exec_id in value_exec_ids:
+        if g._get_iteration_path(exec_id) == (1,):
+            active_value_exec_id = exec_id
+        else:
+            skipped_value_exec_id = exec_id
+
+    g._set_prepared_exec_state(skipped_value_exec_id, "skipped")
+
+    selected_exec_id = g._get_iteration_node(
+        "value", graph.nx_graph_flat(), g.execution_graph.nx_graph_flat(), [first_iter_exec_id]
+    )
+
+    assert selected_exec_id is None
+    assert active_value_exec_id != skipped_value_exec_id
+
+
+def test_mark_exec_node_skipped_does_not_hide_already_executed_results():
+    graph = Graph()
+    graph.add_node(AnyTypeTestInvocation(id="value", value="value"))
+
+    g = GraphExecutionState(graph=graph)
+
+    exec_id = g._create_execution_node("value", [])[0]
+    g.results[exec_id] = AnyTypeTestInvocation(id="result", value="value").invoke(Mock(InvocationContext))
+    g.executed.add(exec_id)
+    g._set_prepared_exec_state(exec_id, "executed")
+
+    g._if_scheduler().mark_exec_node_skipped(exec_id)
+
+    assert g._get_prepared_exec_metadata(exec_id).state == "executed"
+    assert g.results[exec_id].value == "value"
+
+
 def test_are_connection_types_compatible_accepts_subclass_to_base():
     """A subclass output should be connectable to a base-class input.