Added support for "return" handoffs (#1)

Author: niv-hertz

src/agents/handoffs.py

@@ -111,6 +111,12 @@ class Handoff(Generic[TContext]):
     agent and returns whether the handoff is enabled. You can use this to dynamically enable/disable
     a handoff based on your context/state."""
 
+    should_return_control: bool = False
+    """Whether the Agent that receives control during a handoff should return control to the
+    original (previous) Agent upon completion of its work. If False, after the Agent that received
+    the handoff completes its work, the interaction will end.
+    """
+
     def get_transfer_message(self, agent: Agent[Any]) -> str:
         return json.dumps({"assistant": agent.name})
 
@@ -135,6 +141,7 @@ def handoff(
     input_filter: Callable[[HandoffInputData], HandoffInputData] | None = None,
     should_return_control: bool = False,
     is_enabled: bool | Callable[[RunContextWrapper[Any], Agent[Any]], MaybeAwaitable[bool]] = True,
+    should_return_control: bool = False,
 ) -> Handoff[TContext]: ...
 
 
@@ -149,6 +156,7 @@ def handoff(
     input_filter: Callable[[HandoffInputData], HandoffInputData] | None = None,
     should_return_control: bool = False,
     is_enabled: bool | Callable[[RunContextWrapper[Any], Agent[Any]], MaybeAwaitable[bool]] = True,
+    should_return_control: bool = False,
 ) -> Handoff[TContext]: ...
 
 
@@ -162,6 +170,7 @@ def handoff(
     input_filter: Callable[[HandoffInputData], HandoffInputData] | None = None,
     should_return_control: bool = False,
     is_enabled: bool | Callable[[RunContextWrapper[Any], Agent[Any]], MaybeAwaitable[bool]] = True,
+    should_return_control: bool = False,
 ) -> Handoff[TContext]: ...
 
 
@@ -174,6 +183,7 @@ def handoff(
     input_filter: Callable[[HandoffInputData], HandoffInputData] | None = None,
     should_return_control: bool = False,
     is_enabled: bool | Callable[[RunContextWrapper[Any], Agent[Any]], MaybeAwaitable[bool]] = True,
+    should_return_control: bool = False,
 ) -> Handoff[TContext]:
     """Create a handoff from an agent.
 
@@ -258,4 +268,5 @@ async def _invoke_handoff(
         agent_name=agent.name,
         should_return_control=should_return_control,
         is_enabled=is_enabled,
+        should_return_control=should_return_control,
     )

tests/test_run_step_execution.py

@@ -371,6 +371,7 @@ async def get_execute_result(
     output_schema = AgentRunner._get_output_schema(agent)
     previous_agents = previous_agents if previous_agents is not None else []
     handoffs = await AgentRunner._get_handoffs(agent, context_wrapper or RunContextWrapper(None))
+    previous_agents = previous_agents if previous_agents is not None else []
 
     processed_response = RunImpl.process_model_response(
         agent=agent,