From 84bae0f42a041d88280f72567919fb57e11df760 Mon Sep 17 00:00:00 2001
From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com>
Date: Fri, 13 Mar 2026 08:17:24 +0000
Subject: [PATCH] Python: Fix type hint for `Case` and `Default` (#3985)

* Fix type hint for `Case` and `Default`

* Add test

---------

Co-authored-by: Evan Mattson <35585003+moonbox3@users.noreply.github.com>
---
 .../core/agent_framework/_workflows/_edge.py  |  5 ++--
 .../tests/workflow/test_workflow_builder.py   | 25 +++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/python/packages/core/agent_framework/_workflows/_edge.py b/python/packages/core/agent_framework/_workflows/_edge.py
index 02544ad3df..b9dbd266ec 100644
--- a/python/packages/core/agent_framework/_workflows/_edge.py
+++ b/python/packages/core/agent_framework/_workflows/_edge.py
@@ -9,6 +9,7 @@ from collections.abc import Awaitable, Callable, Sequence
 from dataclasses import dataclass, field
 from typing import Any, ClassVar, TypeAlias, TypeVar
 
+from .._agents import SupportsAgentRun
 from ._const import INTERNAL_SOURCE_ID
 from ._executor import Executor
 from ._model_utils import DictConvertible, encode_value
@@ -264,7 +265,7 @@ class Case:
     """
 
     condition: Callable[[Any], bool]
-    target: Executor | str
+    target: Executor | SupportsAgentRun
 
 
 @dataclass
@@ -287,7 +288,7 @@ class Default:
             assert fallback.target.id == "dead_letter"
     """
 
-    target: Executor | str
+    target: Executor | SupportsAgentRun
 
 
 @dataclass(init=False)
diff --git a/python/packages/core/tests/workflow/test_workflow_builder.py b/python/packages/core/tests/workflow/test_workflow_builder.py
index 3a7b719530..29a4bf0292 100644
--- a/python/packages/core/tests/workflow/test_workflow_builder.py
+++ b/python/packages/core/tests/workflow/test_workflow_builder.py
@@ -13,6 +13,8 @@ from agent_framework import (
     AgentRunInputs,
     AgentSession,
     BaseAgent,
+    Case,
+    Default,
     Executor,
     Message,
     ResponseStream,
@@ -223,6 +225,29 @@ def test_add_edge_with_condition():
     assert "Target" in workflow.executors
 
 
+def test_switch_case_with_agents():
+    """Test add_switch_case_edge_group with Case and Default edges using agents."""
+    router = DummyAgent(id="router_agent", name="router")
+    handler = DummyAgent(id="handler", name="handler")
+    fallback = DummyAgent(id="fallback_agent", name="fallback")
+
+    workflow = (
+        WorkflowBuilder(start_executor=router)
+        .add_switch_case_edge_group(
+            router,
+            [
+                Case(condition=lambda _: True, target=handler),
+                Default(target=fallback),
+            ],
+        )
+        .build()
+    )
+
+    # All three agents should be AgentExecutor wrappers
+    agent_executors = [e for e in workflow.executors.values() if isinstance(e, AgentExecutor)]
+    assert len(agent_executors) == 3
+
+
 # region with_output_from tests