From 295548928287ee37137dc0c689909bd5c4c0dcce Mon Sep 17 00:00:00 2001 From: unmarco Date: Wed, 14 Oct 2026 12:42:01 +0200 Subject: [PATCH 1/2] fix(graph): reject mismatched collector items for union collection inputs Collector output validation took the first type argument of the destination as its item type, so Union[list[X], X] and Optional[list[X]] inputs raised a TypeError on mismatched items and enqueue returned 500. Read item types from the list branches instead. --- .../app/services/shared/graph_validation.py | 18 ++++++-- tests/test_node_graph.py | 45 +++++++++++++++++++ tests/test_nodes.py | 8 ++++ 3 files changed, 68 insertions(+), 3 deletions(-) diff --git a/invokeai/app/services/shared/graph_validation.py b/invokeai/app/services/shared/graph_validation.py index 801b97b83e1..0ccb15e11fe 100644 --- a/invokeai/app/services/shared/graph_validation.py +++ b/invokeai/app/services/shared/graph_validation.py @@ -2315,19 +2315,31 @@ def _validate_collector_output_types( if not all(is_list_or_contains_list(t) or is_any(t) for t in output_field_types): return "Collector output must connect to a collection input" + # Read item types from the list branch: the first type argument of `Union[list[X], X]` or + # `Optional[list[X]]` is `list[X]` itself, which issubclass() rejects with a TypeError. if input_root_type is not None: if not all( is_any(t) - or is_union_subtype(input_root_type, get_args(t)[0]) - or issubclass(input_root_type, get_args(t)[0]) + or any( + self._collector_item_matches(input_root_type, item_type) + for item_type in extract_collection_item_types(t) + ) for t in output_field_types ): return "Collector outputs must connect to a collection input with a matching type" - elif any(not is_any(t) and get_args(t)[0] != Any for t in output_field_types): + elif any(not is_any(t) and Any not in extract_collection_item_types(t) for t in output_field_types): return "Collector outputs must connect to a collection input with a matching type" return None + @staticmethod + def _collector_item_matches(input_root_type: Any, item_type: Any) -> bool: + if is_any(item_type) or is_union_subtype(input_root_type, item_type): + return True + return ( + isinstance(input_root_type, type) and isinstance(item_type, type) and issubclass(input_root_type, item_type) + ) + def _validate_downstream_collector_outputs( self, outputs: list[EdgeConnection], diff --git a/tests/test_node_graph.py b/tests/test_node_graph.py index d06a1905cb5..d0e343a1fa4 100644 --- a/tests/test_node_graph.py +++ b/tests/test_node_graph.py @@ -21,6 +21,7 @@ from invokeai.app.invocations.logic import IfInvocation from invokeai.app.invocations.loops import ForInvocation, ForReturnInvocation from invokeai.app.invocations.math import AddInvocation +from invokeai.app.invocations.metadata import MetadataInvocation, MetadataItemInvocation from invokeai.app.invocations.primitives import ( ColorInvocation, FloatCollectionInvocation, @@ -50,6 +51,7 @@ AnyTypeTestInvocation, ImageToImageTestInvocation, ListPassThroughInvocation, + OptionalCollectionTestInvocation, PolymorphicStringTestInvocation, PromptCollectionTestInvocation, PromptTestInvocation, @@ -2199,6 +2201,49 @@ def test_graph_collector_rejects_mismatched_item_with_union_collection_input(): g.add_edge(create_edge("3", "color", "2", "item")) +# A collection input whose annotation is not a bare list: the item type lives in a union branch. +COLLECTION_TARGETS_WITH_LIST_BRANCH = [ + pytest.param(lambda: MetadataInvocation(id="target"), "items", id="list-first-union"), + pytest.param(lambda: OptionalCollectionTestInvocation(id="target"), "value", id="optional-list"), +] + + +@pytest.mark.parametrize(("make_target", "field"), COLLECTION_TARGETS_WITH_LIST_BRANCH) +def test_graph_collector_rejects_mismatched_items_for_collection_union_input(make_target, field: str): + g = Graph() + g.add_node(IntegerInvocation(id="1", value=1)) + g.add_node(IntegerInvocation(id="2", value=2)) + g.add_node(CollectInvocation(id="collect")) + g.add_node(make_target()) + g.add_edge(create_edge("1", "value", "collect", "item")) + g.add_edge(create_edge("2", "value", "collect", "item")) + + with pytest.raises(InvalidEdgeError, match="matching type"): + g.add_edge(create_edge("collect", "collection", "target", field)) + + +def test_graph_collector_accepts_matching_items_for_list_first_union_input(): + g = Graph() + g.add_node(MetadataItemInvocation(id="1", label="a", value=1)) + g.add_node(MetadataItemInvocation(id="2", label="b", value=2)) + g.add_node(CollectInvocation(id="collect")) + g.add_node(MetadataInvocation(id="target")) + g.add_edge(create_edge("1", "item", "collect", "item")) + g.add_edge(create_edge("2", "item", "collect", "item")) + + g.add_edge(create_edge("collect", "collection", "target", "items")) + + +def test_graph_collector_accepts_matching_items_for_optional_list_input(): + g = Graph() + g.add_node(StringInvocation(id="1", value="a")) + g.add_node(CollectInvocation(id="collect")) + g.add_node(OptionalCollectionTestInvocation(id="target")) + g.add_edge(create_edge("1", "value", "collect", "item")) + + g.add_edge(create_edge("collect", "collection", "target", "value")) + + def test_graph_connects_iterator(): g = Graph() n1 = ListPassThroughInvocation(id="1") diff --git a/tests/test_nodes.py b/tests/test_nodes.py index 6368e416c12..aaf50f35fc4 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -142,6 +142,14 @@ def invoke(self, context: InvocationContext) -> UnionCollectionTestInvocationOut return UnionCollectionTestInvocationOutput(value=self.value) +@invocation("test_optional_collection", version="1.0.0") +class OptionalCollectionTestInvocation(BaseInvocation): + value: Union[list[str], None] = InputField(default=None) + + def invoke(self, context: InvocationContext) -> PromptCollectionTestInvocationOutput: + return PromptCollectionTestInvocationOutput(collection=self.value or []) + + # Importing these must happen after test invocations are defined or they won't register from invokeai.app.services.events.events_base import EventServiceBase # noqa: E402 from invokeai.app.services.shared.graph import Edge, EdgeConnection, GraphExecutionState # noqa: E402 From 796b0f83b54c6261ec1e1f4dd832ef4bf2974e8c Mon Sep 17 00:00:00 2001 From: unmarco Date: Sat, 17 Oct 2026 09:10:02 +0200 Subject: [PATCH 2/2] fix(graph): accept subclass collector items for union item types A collection input like list[Union[Base, str]] now matches a collector of Derived(Base) again: the class guard checks issubclass against each union member. Adds Graph.add_edge regressions for typing.Union and PEP 604 unions. --- .../app/services/shared/graph_validation.py | 9 +++-- tests/test_node_graph.py | 35 ++++++++++++++++++ tests/test_nodes.py | 37 +++++++++++++++++++ 3 files changed, 78 insertions(+), 3 deletions(-) diff --git a/invokeai/app/services/shared/graph_validation.py b/invokeai/app/services/shared/graph_validation.py index 0ccb15e11fe..1430c1b0722 100644 --- a/invokeai/app/services/shared/graph_validation.py +++ b/invokeai/app/services/shared/graph_validation.py @@ -5,6 +5,7 @@ from contextvars import ContextVar from dataclasses import dataclass from functools import wraps +from types import UnionType from typing import ( TYPE_CHECKING, Any, @@ -2336,9 +2337,11 @@ def _validate_collector_output_types( def _collector_item_matches(input_root_type: Any, item_type: Any) -> bool: if is_any(item_type) or is_union_subtype(input_root_type, item_type): return True - return ( - isinstance(input_root_type, type) and isinstance(item_type, type) and issubclass(input_root_type, item_type) - ) + if not isinstance(input_root_type, type): + return False + # A union item type (e.g. `list[Union[Base, str]]`) matches a subclass of any of its members. + candidates = get_args(item_type) if get_origin(item_type) in (Union, UnionType) else (item_type,) + return any(isinstance(c, type) and issubclass(input_root_type, c) for c in candidates) def _validate_downstream_collector_outputs( self, diff --git a/tests/test_node_graph.py b/tests/test_node_graph.py index d0e343a1fa4..5e13b4544ee 100644 --- a/tests/test_node_graph.py +++ b/tests/test_node_graph.py @@ -49,15 +49,18 @@ ) from tests.test_nodes import ( AnyTypeTestInvocation, + DerivedFieldTestInvocation, ImageToImageTestInvocation, ListPassThroughInvocation, OptionalCollectionTestInvocation, + Pep604UnionItemCollectionTestInvocation, PolymorphicStringTestInvocation, PromptCollectionTestInvocation, PromptTestInvocation, PromptTestInvocationOutput, TextToImageTestInvocation, UnionCollectionTestInvocation, + UnionItemCollectionTestInvocation, get_single_output_from_session, run_session_with_mock_context, ) @@ -2244,6 +2247,38 @@ def test_graph_collector_accepts_matching_items_for_optional_list_input(): g.add_edge(create_edge("collect", "collection", "target", "value")) +# Collection inputs whose item type is a union, spelled with typing.Union and with PEP 604. +UNION_ITEM_COLLECTION_TARGETS = [ + pytest.param(UnionItemCollectionTestInvocation, id="typing-union"), + pytest.param(Pep604UnionItemCollectionTestInvocation, id="pep604-union"), +] + + +@pytest.mark.parametrize("target_type", UNION_ITEM_COLLECTION_TARGETS) +def test_graph_collector_accepts_subclass_items_for_union_item_collection_input(target_type): + g = Graph() + g.add_node(DerivedFieldTestInvocation(id="1")) + g.add_node(DerivedFieldTestInvocation(id="2")) + g.add_node(CollectInvocation(id="collect")) + g.add_node(target_type(id="target")) + g.add_edge(create_edge("1", "value", "collect", "item")) + g.add_edge(create_edge("2", "value", "collect", "item")) + + g.add_edge(create_edge("collect", "collection", "target", "value")) + + +@pytest.mark.parametrize("target_type", UNION_ITEM_COLLECTION_TARGETS) +def test_graph_collector_rejects_mismatched_items_for_union_item_collection_input(target_type): + g = Graph() + g.add_node(IntegerInvocation(id="1", value=1)) + g.add_node(CollectInvocation(id="collect")) + g.add_node(target_type(id="target")) + g.add_edge(create_edge("1", "value", "collect", "item")) + + with pytest.raises(InvalidEdgeError, match="matching type"): + g.add_edge(create_edge("collect", "collection", "target", "value")) + + def test_graph_connects_iterator(): g = Graph() n1 = ListPassThroughInvocation(id="1") diff --git a/tests/test_nodes.py b/tests/test_nodes.py index aaf50f35fc4..71d3ffc1216 100644 --- a/tests/test_nodes.py +++ b/tests/test_nodes.py @@ -1,6 +1,8 @@ from typing import Any, Callable, Union from unittest.mock import MagicMock +from pydantic import BaseModel + from invokeai.app.invocations.baseinvocation import ( BaseInvocation, BaseInvocationOutput, @@ -150,6 +152,41 @@ def invoke(self, context: InvocationContext) -> PromptCollectionTestInvocationOu return PromptCollectionTestInvocationOutput(collection=self.value or []) +class BaseTestField(BaseModel): + name: str = "" + + +class DerivedTestField(BaseTestField): + pass + + +@invocation_output("test_derived_field_output") +class DerivedFieldTestInvocationOutput(BaseInvocationOutput): + value: DerivedTestField = OutputField() + + +@invocation("test_derived_field", version="1.0.0") +class DerivedFieldTestInvocation(BaseInvocation): + def invoke(self, context: InvocationContext) -> DerivedFieldTestInvocationOutput: + return DerivedFieldTestInvocationOutput(value=DerivedTestField()) + + +@invocation("test_union_item_collection", version="1.0.0") +class UnionItemCollectionTestInvocation(BaseInvocation): + value: list[Union[BaseTestField, str]] = InputField(default=[]) + + def invoke(self, context: InvocationContext) -> PromptCollectionTestInvocationOutput: + return PromptCollectionTestInvocationOutput(collection=[]) + + +@invocation("test_pep604_union_item_collection", version="1.0.0") +class Pep604UnionItemCollectionTestInvocation(BaseInvocation): + value: list[BaseTestField | str] = InputField(default=[]) + + def invoke(self, context: InvocationContext) -> PromptCollectionTestInvocationOutput: + return PromptCollectionTestInvocationOutput(collection=[]) + + # Importing these must happen after test invocations are defined or they won't register from invokeai.app.services.events.events_base import EventServiceBase # noqa: E402 from invokeai.app.services.shared.graph import Edge, EdgeConnection, GraphExecutionState # noqa: E402