fix: preserve handoff target resolution compatibility in run-state agent maps (#2423)

Author: seratch

src/agents/handoffs/__init__.py

@@ -2,8 +2,9 @@
 
 import inspect
 import json
+import weakref
 from collections.abc import Awaitable
-from dataclasses import dataclass, replace as dataclasses_replace
+from dataclasses import dataclass, field, replace as dataclasses_replace
 from typing import TYPE_CHECKING, Any, Callable, Generic, cast, overload
 
 from pydantic import TypeAdapter
@@ -148,6 +149,11 @@ class Handoff(Generic[TContext, TAgent]):
     context or state.
     """
 
+    _agent_ref: weakref.ReferenceType[AgentBase[Any]] | None = field(
+        default=None, init=False, repr=False
+    )
+    """Weak reference to the target agent when constructed via `handoff()`."""
+
     def get_transfer_message(self, agent: AgentBase[Any]) -> str:
         return json.dumps({"assistant": agent.name})
 
@@ -300,7 +306,7 @@ async def _is_enabled(ctx: RunContextWrapper[Any], agent_base: AgentBase[Any]) -
             return await result
         return bool(result)
 
-    return Handoff(
+    handoff_obj = Handoff(
         tool_name=tool_name,
         tool_description=tool_description,
         input_json_schema=input_json_schema,
@@ -310,6 +316,8 @@ async def _is_enabled(ctx: RunContextWrapper[Any], agent_base: AgentBase[Any]) -
         agent_name=agent.name,
         is_enabled=_is_enabled if callable(is_enabled) else is_enabled,
     )
+    handoff_obj._agent_ref = weakref.ref(agent)
+    return handoff_obj
 
 
 __all__ = [

src/agents/run_state.py

@@ -1898,11 +1898,63 @@ def _build_agent_map(initial_agent: Agent[Any]) -> dict[str, Agent[Any]]:
         agent_map[current.name] = current
 
         # Add handoff agents to the queue
-        for handoff in current.handoffs:
-            # Handoff can be either an Agent or a Handoff object with an .agent attribute
-            handoff_agent = handoff if not hasattr(handoff, "agent") else handoff.agent
-            if handoff_agent and handoff_agent.name not in agent_map:  # type: ignore[union-attr]
-                queue.append(handoff_agent)  # type: ignore[arg-type]
+        for handoff_item in current.handoffs:
+            handoff_agent: Any | None = None
+            handoff_agent_name: str | None = None
+
+            if isinstance(handoff_item, Handoff):
+                # Some custom/mocked Handoff subclasses bypass dataclass initialization.
+                # Prefer agent_name, then legacy name fallback used in tests.
+                candidate_name = getattr(handoff_item, "agent_name", None) or getattr(
+                    handoff_item, "name", None
+                )
+                if isinstance(candidate_name, str):
+                    handoff_agent_name = candidate_name
+                    if handoff_agent_name in agent_map:
+                        continue
+
+                handoff_ref = getattr(handoff_item, "_agent_ref", None)
+                handoff_agent = handoff_ref() if callable(handoff_ref) else None
+                if handoff_agent is None:
+                    # Backward-compatibility fallback for custom legacy handoff objects that store
+                    # the target directly on `.agent`. New code should prefer `handoff()` objects.
+                    legacy_agent = getattr(handoff_item, "agent", None)
+                    if legacy_agent is not None:
+                        handoff_agent = legacy_agent
+                        logger.debug(
+                            "Using legacy handoff `.agent` fallback while building agent map. "
+                            "This compatibility path is not recommended for new code."
+                        )
+                if handoff_agent_name is None:
+                    candidate_name = getattr(handoff_agent, "name", None)
+                    handoff_agent_name = candidate_name if isinstance(candidate_name, str) else None
+                if handoff_agent is None or not hasattr(handoff_agent, "handoffs"):
+                    if handoff_agent_name:
+                        logger.debug(
+                            "Skipping unresolved handoff target while building agent map: %s",
+                            handoff_agent_name,
+                        )
+                    continue
+            else:
+                # Backward-compatibility fallback for custom legacy handoff wrappers that expose
+                # the target directly on `.agent` without inheriting from `Handoff`.
+                legacy_agent = getattr(handoff_item, "agent", None)
+                if legacy_agent is not None:
+                    handoff_agent = legacy_agent
+                    logger.debug(
+                        "Using legacy non-`Handoff` `.agent` fallback while building agent map."
+                    )
+                else:
+                    handoff_agent = handoff_item
+                candidate_name = getattr(handoff_agent, "name", None)
+                handoff_agent_name = candidate_name if isinstance(candidate_name, str) else None
+
+            if (
+                handoff_agent is not None
+                and handoff_agent_name
+                and handoff_agent_name not in agent_map
+            ):
+                queue.append(cast(Any, handoff_agent))
 
         # Include agent-as-tool instances so nested approvals can be restored.
         tools = getattr(current, "tools", None)

tests/test_run_state.py

@@ -876,6 +876,70 @@ def test_build_agent_map_handles_complex_handoff_graphs(self):
         assert len(agent_map) == 4
         assert all(agent_map.get(name) is not None for name in ["A", "B", "C", "D"])
 
+    def test_build_agent_map_handles_handoff_objects(self):
+        """Test that buildAgentMap resolves handoff() objects via weak references."""
+        agent_a = Agent(name="AgentA")
+        agent_b = Agent(name="AgentB")
+        agent_a.handoffs = [handoff(agent_b)]
+
+        agent_map = _build_agent_map(agent_a)
+
+        assert sorted(agent_map.keys()) == ["AgentA", "AgentB"]
+
+    def test_build_agent_map_supports_legacy_handoff_agent_attribute(self):
+        """Test that buildAgentMap keeps legacy custom handoffs with `.agent` targets working."""
+        agent_a = Agent(name="AgentA")
+        agent_b = Agent(name="AgentB")
+
+        class LegacyHandoff(Handoff):
+            def __init__(self, target: Agent[Any]):
+                # Legacy custom handoff shape supported only for backward compatibility.
+                self.agent = target
+                self.agent_name = target.name
+                self.name = "legacy_handoff"
+
+        agent_a.handoffs = [LegacyHandoff(agent_b)]
+
+        agent_map = _build_agent_map(agent_a)
+
+        assert sorted(agent_map.keys()) == ["AgentA", "AgentB"]
+
+    def test_build_agent_map_supports_legacy_non_handoff_agent_wrapper(self):
+        """Test that buildAgentMap supports legacy non-Handoff wrappers with `.agent` targets."""
+        agent_a = Agent(name="AgentA")
+        agent_b = Agent(name="AgentB")
+
+        class LegacyWrapper:
+            def __init__(self, target: Agent[Any]):
+                self.agent = target
+
+        agent_a.handoffs = [LegacyWrapper(agent_b)]  # type: ignore[list-item]
+
+        agent_map = _build_agent_map(agent_a)
+
+        assert sorted(agent_map.keys()) == ["AgentA", "AgentB"]
+
+    def test_build_agent_map_skips_unresolved_handoff_objects(self):
+        """Test that buildAgentMap skips custom handoffs without target agent references."""
+        agent_a = Agent(name="AgentA")
+        agent_b = Agent(name="AgentB")
+
+        async def _invoke_handoff(_ctx: RunContextWrapper[Any], _input: str) -> Agent[Any]:
+            return agent_b
+
+        detached_handoff = Handoff(
+            tool_name="transfer_to_agent_b",
+            tool_description="Transfer to AgentB.",
+            input_json_schema={},
+            on_invoke_handoff=_invoke_handoff,
+            agent_name=agent_b.name,
+        )
+        agent_a.handoffs = [detached_handoff]
+
+        agent_map = _build_agent_map(agent_a)
+
+        assert sorted(agent_map.keys()) == ["AgentA"]
+
 
 class TestSerializationRoundTrip:
     """Test that serialization and deserialization preserve state correctly."""