Python: Allow union types in FanIn edge group (#868)

* Improve type utils

* Add sample

* Add Union

* Add more test cases

* Add more test cases

* Fix RequestResponse typing to only coerce mapping original_request

---------

Co-authored-by: Evan Mattson <evan.mattson@microsoft.com>

python/packages/main/agent_framework/_workflow/_typing_utils.py

@@ -1,7 +1,9 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import logging
+from collections.abc import Mapping
 from dataclasses import fields, is_dataclass
+from types import UnionType
 from typing import Any, Union, get_args, get_origin
 
 logger = logging.getLogger(__name__)
@@ -54,7 +56,7 @@ def _coerce_to_type(value: Any, target_type: type) -> Any | None:
     return None
 
 
-def is_instance_of(data: Any, target_type: type) -> bool:
+def is_instance_of(data: Any, target_type: type | UnionType | Any) -> bool:
     """Check if the data is an instance of the target type.
 
     Args:
@@ -77,17 +79,28 @@ def is_instance_of(data: Any, target_type: type) -> bool:
 
     # Case 2: target_type is Optional[T] or Union[T1, T2, ...]
     # Optional[T] is really just as Union[T, None]
+    if origin is UnionType:
+        return any(is_instance_of(data, arg) for arg in args)
+
+    # Case 2b: Handle typing.Union (legacy Union syntax)
     if origin is Union:
         return any(is_instance_of(data, arg) for arg in args)
 
     # Case 3: target_type is a generic type
     if origin in [list, set]:
-        return isinstance(data, origin) and all(is_instance_of(item, args[0]) for item in data)  # type: ignore
+        return isinstance(data, origin) and (
+            not args or all(any(is_instance_of(item, arg) for arg in args) for item in data)
+        )  # type: ignore
 
     # Case 4: target_type is a tuple
     if origin is tuple:
+        if len(args) == 2 and args[1] is Ellipsis:  # Tuple[T, ...] case
+            element_type = args[0]
+            return isinstance(data, tuple) and all(is_instance_of(item, element_type) for item in data)
         if len(args) == 1 and args[0] is Ellipsis:  # Tuple[...] case
             return isinstance(data, tuple)
+        if len(args) == 0:
+            return isinstance(data, tuple)
         return (
             isinstance(data, tuple)
             and len(data) == len(args)  # type: ignore
@@ -96,9 +109,12 @@ def is_instance_of(data: Any, target_type: type) -> bool:
 
     # Case 5: target_type is a dict
     if origin is dict:
-        return isinstance(data, dict) and all(
-            is_instance_of(key, args[0]) and is_instance_of(value, args[1])
-            for key, value in data.items()  # type: ignore
+        return isinstance(data, dict) and (
+            not args
+            or all(
+                is_instance_of(key, args[0]) and is_instance_of(value, args[1])
+                for key, value in data.items()  # type: ignore
+            )
         )
 
     # Case 6: target_type is RequestResponse[T, U] - validate generic parameters
@@ -114,10 +130,17 @@ def is_instance_of(data: Any, target_type: type) -> bool:
                 and data.original_request is not None
                 and not is_instance_of(data.original_request, request_type)
             ):
-                coerced = _coerce_to_type(data.original_request, request_type)
-                if coerced is None:
+                # Checkpoint decoding can leave original_request as a plain mapping. In that
+                # case we coerce it back into the expected request type so downstream handlers
+                # and validators still receive a fully typed RequestResponse instance.
+                original_request = data.original_request
+                if isinstance(original_request, Mapping):
+                    coerced = _coerce_to_type(dict(original_request), request_type)
+                    if coerced is None or not isinstance(coerced, request_type):
+                        return False
+                    data.original_request = coerced
+                else:
                     return False
-                data.original_request = coerced
             if hasattr(data, "data") and data.data is not None and not is_instance_of(data.data, response_type):
                 return False
         return True

python/packages/main/tests/workflow/test_edge.py

@@ -79,6 +79,33 @@ class MockAggregator(Executor):
         self.call_count += 1
         self.last_message = message
 
+    @handler
+    async def mock_aggregator_handler_secondary(
+        self,
+        message: list[MockMessageSecondary],
+        ctx: WorkflowContext,
+    ) -> None:
+        """A mock aggregator handler that does nothing."""
+        self.call_count += 1
+        self.last_message = message
+
+
+class MockAggregatorSecondary(Executor):
+    """A mock aggregator that has a handler for a union type for testing purposes."""
+
+    call_count: int = 0
+    last_message: Any = None
+
+    @handler
+    async def mock_aggregator_handler_combine(
+        self,
+        message: list[MockMessage | MockMessageSecondary],
+        ctx: WorkflowContext,
+    ) -> None:
+        """A mock aggregator handler that does nothing."""
+        self.call_count += 1
+        self.last_message = message
+
 
 # region Edge
 
@@ -1062,6 +1089,73 @@ async def test_fan_in_edge_group_tracing_type_mismatch(span_exporter) -> None:
     assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value
 
 
+async def test_fan_in_edge_group_with_multiple_message_types() -> None:
+    source1 = MockExecutor(id="source_executor_1")
+    source2 = MockExecutor(id="source_executor_2")
+    target = MockAggregatorSecondary(id="target_executor")
+
+    edge_group = FanInEdgeGroup(source_ids=[source1.id, source2.id], target_id=target.id)
+
+    executors: dict[str, Executor] = {source1.id: source1, source2.id: source2, target.id: target}
+    edge_runner = create_edge_runner(edge_group, executors)
+
+    shared_state = SharedState()
+    ctx = InProcRunnerContext()
+
+    data = MockMessage(data="test")
+
+    success = await edge_runner.send_message(
+        Message(data=data, source_id=source1.id),
+        shared_state,
+        ctx,
+    )
+    assert success
+
+    data2 = MockMessageSecondary(data="test")
+    success = await edge_runner.send_message(
+        Message(data=data2, source_id=source2.id),
+        shared_state,
+        ctx,
+    )
+    assert success
+
+
+async def test_fan_in_edge_group_with_multiple_message_types_failed() -> None:
+    source1 = MockExecutor(id="source_executor_1")
+    source2 = MockExecutor(id="source_executor_2")
+    target = MockAggregator(id="target_executor")
+
+    edge_group = FanInEdgeGroup(source_ids=[source1.id, source2.id], target_id=target.id)
+
+    executors: dict[str, Executor] = {source1.id: source1, source2.id: source2, target.id: target}
+    edge_runner = create_edge_runner(edge_group, executors)
+
+    shared_state = SharedState()
+    ctx = InProcRunnerContext()
+
+    data = MockMessage(data="test")
+
+    success = await edge_runner.send_message(
+        Message(data=data, source_id=source1.id),
+        shared_state,
+        ctx,
+    )
+    assert success
+
+    with pytest.raises(RuntimeError):
+        # Although `MockAggregator` can handle `list[MockMessage]` and `list[MockMessageSecondary]`
+        # separately (i.e., it has handlers for each type individually), it cannot handle
+        # `list[MockMessage | MockMessageSecondary]` (a list containing a mix of both types).
+        # With the fan-in edge group, the target executor must handle all message types from the
+        # source executors as a union.
+        data2 = MockMessageSecondary(data="test")
+        _ = await edge_runner.send_message(
+            Message(data=data2, source_id=source2.id),
+            shared_state,
+            ctx,
+        )
+
+
 # endregion FanInEdgeGroup
 
 # region SwitchCaseEdgeGroup

python/packages/main/tests/workflow/test_typing_utils.py

@@ -0,0 +1,134 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar, Union
+
+from agent_framework._workflow import RequestInfoMessage, RequestResponse
+from agent_framework._workflow._typing_utils import is_instance_of
+
+
+def test_basic_types() -> None:
+    """Test basic built-in types."""
+    assert is_instance_of(5, int)
+    assert is_instance_of("hello", str)
+    assert is_instance_of(None, type(None))
+
+
+def test_union_types() -> None:
+    """Test union types (|) and optional types."""
+    assert is_instance_of(5, int | str)
+    assert is_instance_of("hello", int | str)
+    assert is_instance_of(5, Union[int, str])
+    assert not is_instance_of(5.0, int | str)
+
+
+def test_list_types() -> None:
+    """Test list types with various element types."""
+    assert is_instance_of([], list)
+    assert is_instance_of([1, 2, 3], list)
+    assert is_instance_of([1, 2, 3], list[int])
+    assert is_instance_of([1, 2, 3], list[int | str])
+    assert is_instance_of([1, "a", 3], list[int | str])
+    assert is_instance_of([1, "a", 3], list[Union[int, str]])
+    assert not is_instance_of([1, 2.0, 3], dict)
+    assert not is_instance_of([1, 2.0, 3], list[int | str])
+
+
+def test_tuple_types() -> None:
+    """Test tuple types with fixed and variable lengths."""
+    assert is_instance_of((1, "a"), tuple)
+    assert is_instance_of((1, "a"), tuple[int, str])
+    assert is_instance_of((1, "a", 3), tuple[int | str, ...])
+    assert is_instance_of((1, 2.0, "a"), tuple[...])  # type: ignore
+    assert not is_instance_of((1, 2.0, 3), tuple[int | str, ...])
+    assert not is_instance_of((1, 2.0, 3), dict)
+
+
+def test_dict_types() -> None:
+    """Test dictionary types with typed keys and values."""
+    assert is_instance_of({"key": "value"}, dict)
+    assert is_instance_of({"key": "value"}, dict[str, str])
+    assert is_instance_of({"key": 5, "another_key": "value"}, dict[str, int | str])
+    assert not is_instance_of({"key": 5, "another_key": 3.0}, dict[str, int | str])
+    assert not is_instance_of({"key": 5, "another_key": 3.0}, list)
+
+
+def test_set_types() -> None:
+    """Test set types with various element types."""
+    assert is_instance_of({1, 2, 3}, set)
+    assert is_instance_of({1, 2, 3}, set[int])
+    assert is_instance_of({1, 2, 3}, set[int | str])
+    assert is_instance_of({1, "a", 3}, set[int | str])
+    assert is_instance_of({1, "a", 3}, set[Union[int, str]])
+    assert is_instance_of(set(), set[int])
+    assert not is_instance_of({1, 2.0, 3}, set[int | str])
+    assert not is_instance_of({1, 2, 3}, list)
+    assert not is_instance_of({1, 2, 3}, dict)
+
+
+def test_any_type() -> None:
+    """Test Any type - should accept all values."""
+    assert is_instance_of(5, Any)
+    assert is_instance_of("hello", Any)
+    assert is_instance_of([1, 2, 3], Any)
+
+
+def test_nested_types() -> None:
+    """Test complex nested type structures."""
+    assert is_instance_of([{"key": [1, 2]}, {"another_key": [3]}], list[dict[str, list[int]]])
+    assert not is_instance_of([{"key": [1, 2]}, {"another_key": [3.0]}], list[dict[str, list[int]]])
+
+
+def test_custom_type() -> None:
+    """Test custom object type checking."""
+
+    @dataclass
+    class CustomClass:
+        value: int
+
+    instance = CustomClass(10)
+    assert is_instance_of(instance, CustomClass)
+    assert not is_instance_of(instance, dict)
+
+
+def test_request_response_type() -> None:
+    """Test RequestResponse generic type checking."""
+
+    request_instance = RequestResponse[RequestInfoMessage, str](
+        is_handled=False,
+        original_request=RequestInfoMessage(),
+    )
+
+    class CustomRequestInfoMessage(RequestInfoMessage):
+        info: str
+
+    assert is_instance_of(request_instance, RequestResponse[RequestInfoMessage, str])
+    assert not is_instance_of(request_instance, RequestResponse[CustomRequestInfoMessage, str])
+
+
+def test_custom_generic_type() -> None:
+    """Test custom generic type checking."""
+
+    T = TypeVar("T")
+    U = TypeVar("U")
+
+    class CustomClass(Generic[T, U]):
+        def __init__(self, request: T, response: U, extra: Any | None = None) -> None:
+            self.request = request
+            self.response = response
+            self.extra = extra
+
+    instance = CustomClass[int, str](request=5, response="response")
+
+    assert is_instance_of(instance, CustomClass[int, str])
+    # Generic parameters are not strictly enforced at runtime
+    assert is_instance_of(instance, CustomClass[str, str])
+
+
+def test_edge_cases() -> None:
+    """Test edge cases and unusual scenarios."""
+    assert is_instance_of([], list[int])  # Empty list should be valid
+    assert is_instance_of((), tuple[int, ...])  # Empty tuple should be valid
+    assert is_instance_of({}, dict[str, int])  # Empty dict should be valid
+    assert is_instance_of(None, int | None)  # Optional type with None
+    assert not is_instance_of(5, str | None)  # Optional type without matching type

python/samples/getting_started/workflow/README.md

@@ -87,6 +87,7 @@ Once comfortable with these, explore the rest of the samples below.
 | Sample | File | Concepts |
 |---|---|---|
 | Concurrent (Fan-out/Fan-in) | [parallelism/fan_out_fan_in_edges.py](./parallelism/fan_out_fan_in_edges.py) | Dispatch to multiple executors and aggregate results |
+| Aggregate Results of Different Types | [parallelism/aggregate_results_of_different_types.py](./parallelism/aggregate_results_of_different_types.py) | Handle results of different types from multiple concurrent executors |
 | Map-Reduce with Visualization | [parallelism/map_reduce_and_visualization.py](./parallelism/map_reduce_and_visualization.py) | Fan-out/fan-in pattern with diagram export |
 
 ### state-management

python/samples/getting_started/workflow/parallelism/aggregate_results_of_different_types.py

@@ -0,0 +1,100 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import random
+
+from agent_framework import Executor, WorkflowBuilder, WorkflowCompletedEvent, WorkflowContext, handler
+
+"""
+Sample: Concurrent fan out and fan in with two different tasks that output results of different types.
+
+Purpose:
+Show how to construct a parallel branch pattern in workflows. Demonstrate:
+- Fan out by targeting multiple executors from one dispatcher.
+- Fan in by collecting a list of results from the executors.
+- Simple tracing using AgentRunEvent to observe execution order and progress.
+
+Prerequisites:
+- Familiarity with WorkflowBuilder, executors, edges, events, and streaming runs.
+"""
+
+
+class Dispatcher(Executor):
+    """
+    The sole purpose of this decorator is to dispatch the input of the workflow to
+    other executors.
+    """
+
+    @handler
+    async def handle(self, numbers: list[int], ctx: WorkflowContext[list[int]]):
+        if not numbers:
+            raise RuntimeError("Input must be a valid list of integers.")
+
+        await ctx.send_message(numbers)
+
+
+class Average(Executor):
+    """Calculate the average of a list of integers."""
+
+    @handler
+    async def handle(self, numbers: list[int], ctx: WorkflowContext[float]):
+        average: float = sum(numbers) / len(numbers)
+        await ctx.send_message(average)
+
+
+class Sum(Executor):
+    """Calculate the sum of a list of integers."""
+
+    @handler
+    async def handle(self, numbers: list[int], ctx: WorkflowContext[int]):
+        total: int = sum(numbers)
+        await ctx.send_message(total)
+
+
+class Aggregator(Executor):
+    """Aggregate the results from the different tasks and emit the `WorkflowCompletedEvent`."""
+
+    @handler
+    async def handle(self, results: list[int | float], ctx: WorkflowContext[None]):
+        """Receive the results from the source executors.
+
+        The framework will automatically collect messages from the source executors
+        and deliver them as a list.
+
+        Args:
+            results (list[int | float]): execution results from upstream executors.
+                The type annotation must be a list of union types that the upstream
+                executors will produce.
+            cts (WorkflowContext[None]): A workflow context.
+        """
+        await ctx.add_event(WorkflowCompletedEvent(data=results))
+
+
+async def main() -> None:
+    # 1) Create the executors
+    dispatcher = Dispatcher(id="dispatcher")
+    average = Average(id="average")
+    summation = Sum(id="summation")
+    aggregator = Aggregator(id="aggregator")
+
+    # 2) Build a simple fan out and fan in workflow
+    workflow = (
+        WorkflowBuilder()
+        .set_start_executor(dispatcher)
+        .add_fan_out_edges(dispatcher, [average, summation])
+        .add_fan_in_edges([average, summation], aggregator)
+        .build()
+    )
+
+    # 3) Run the workflow
+    completion: WorkflowCompletedEvent | None = None
+    async for event in workflow.run_stream([random.randint(1, 100) for _ in range(10)]):
+        if isinstance(event, WorkflowCompletedEvent):
+            completion = event
+
+    if completion is not None:
+        print(completion.data)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

python/uv.lock

@@ -3578,7 +3578,7 @@ wheels = [
 
 [[package]]
 name = "posthog"
-version = "6.7.5"
+version = "6.7.6"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "backoff", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
@@ -3588,9 +3588,9 @@ dependencies = [
     { name = "six", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
     { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/1f/f3/d9055bcd190980730bdc600318b34290e85fcb0afb1e31f83cdc33f92615/posthog-6.7.5.tar.gz", hash = "sha256:f4f32b4a4b0df531ae8f80f255a33a49e8880c8c1b62712e6b640535e33a905f", size = 118558, upload-time = "2025-09-16T12:40:34.431Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/e2/ce/11d6fa30ab517018796e1d675498992da585479e7079770ec8fa99a61561/posthog-6.7.6.tar.gz", hash = "sha256:ee5c5ad04b857d96d9b7a4f715e23916a2f206bfcf25e5a9d328a3d27664b0d3", size = 119129, upload-time = "2025-09-22T18:11:12.365Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/42/b4/b40f8467252b4ff481e54a9767b211b4ff83114e6d0b6f481852d0ef3e46/posthog-6.7.5-py3-none-any.whl", hash = "sha256:95b00f915365939e63fa183635bad1caaf89cf4a24b63c8bb6983f2a22a56cb3", size = 136766, upload-time = "2025-09-16T12:40:32.741Z" },
+    { url = "https://files.pythonhosted.org/packages/de/84/586422d8861b5391c8414360b10f603c0b7859bb09ad688e64430ed0df7b/posthog-6.7.6-py3-none-any.whl", hash = "sha256:b09a7e65a042ec416c28874b397d3accae412a80a8b0ef3fa686fbffc99e4d4b", size = 137348, upload-time = "2025-09-22T18:11:10.807Z" },
 ]
 
 [[package]]