fix(graph): reject mismatched collector items for union collection inputs

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,
@@ -2315,19 +2316,33 @@ 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
+        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,
         outputs: list[EdgeConnection],

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,
@@ -48,14 +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,
 )
@@ -2222,6 +2227,81 @@ 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"))
+
+
+# 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")

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,
@@ -142,6 +144,49 @@ 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 [])
+
+
+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