Revert "Add support for returning control in Handoffs"

This reverts commit 6fea778.

Author: amitcoralogix

src/agents/_run_impl.py

@@ -176,11 +176,6 @@ class NextStepHandoff:
     new_agent: Agent[Any]
 
 
-@dataclass
-class NextStepHandoffReturnControl:
-    previous_agent: Agent[Any]
-
-
 @dataclass
 class NextStepFinalOutput:
     output: Any
@@ -206,9 +201,7 @@ class SingleStepResult:
     new_step_items: list[RunItem]
     """Items generated during this current step."""
 
-    next_step: (
-        NextStepHandoff | NextStepFinalOutput | NextStepRunAgain | NextStepHandoffReturnControl
-    )
+    next_step: NextStepHandoff | NextStepFinalOutput | NextStepRunAgain
     """The next step to take."""
 
     @property
@@ -245,7 +238,6 @@ async def execute_tools_and_side_effects(
         hooks: RunHooks[TContext],
         context_wrapper: RunContextWrapper[TContext],
         run_config: RunConfig,
-        previous_agents: list[Agent],
     ) -> SingleStepResult:
         # Make a copy of the generated items
         pre_step_items = list(pre_step_items)
@@ -294,7 +286,6 @@ async def execute_tools_and_side_effects(
                 hooks=hooks,
                 context_wrapper=context_wrapper,
                 run_config=run_config,
-                previous_agents=previous_agents,
             )
 
         # Next, we'll check if the tool use should result in a final output
@@ -325,7 +316,6 @@ async def execute_tools_and_side_effects(
                 final_output=check_tool_use.final_output,
                 hooks=hooks,
                 context_wrapper=context_wrapper,
-                previous_agents=previous_agents,
             )
 
         # Now we can check if the model also produced a final output
@@ -350,7 +340,6 @@ async def execute_tools_and_side_effects(
                 final_output=final_output,
                 hooks=hooks,
                 context_wrapper=context_wrapper,
-                previous_agents=previous_agents,
             )
         elif (
             not output_schema or output_schema.is_plain_text()
@@ -364,7 +353,6 @@ async def execute_tools_and_side_effects(
                 final_output=potential_final_output_text or "",
                 hooks=hooks,
                 context_wrapper=context_wrapper,
-                previous_agents=previous_agents,
             )
         else:
             # If there's no final output, we can just run again
@@ -675,7 +663,6 @@ async def execute_handoffs(
         hooks: RunHooks[TContext],
         context_wrapper: RunContextWrapper[TContext],
         run_config: RunConfig,
-        previous_agents: list[Agent[TContext]],
     ) -> SingleStepResult:
         # If there is more than one handoff, add tool responses that reject those handoffs
         multiple_handoffs = len(run_handoffs) > 1
@@ -697,8 +684,6 @@ async def execute_handoffs(
         actual_handoff = run_handoffs[0]
         with handoff_span(from_agent=agent.name) as span_handoff:
             handoff = actual_handoff.handoff
-            if handoff.should_return_control:
-                previous_agents.append(agent)
             new_agent: Agent[Any] = await handoff.on_invoke_handoff(
                 context_wrapper, actual_handoff.tool_call.arguments
             )
@@ -840,21 +825,16 @@ async def execute_final_output(
         final_output: Any,
         hooks: RunHooks[TContext],
         context_wrapper: RunContextWrapper[TContext],
-        previous_agents: list[Agent[TContext]],
     ) -> SingleStepResult:
-        is_returning_control = len(previous_agents) > 0
         # Run the on_end hooks
-        await cls.run_final_output_hooks(
-            agent, hooks, context_wrapper, final_output, is_returning_control
-        )
+        await cls.run_final_output_hooks(agent, hooks, context_wrapper, final_output)
+
         return SingleStepResult(
             original_input=original_input,
             model_response=new_response,
             pre_step_items=pre_step_items,
             new_step_items=new_step_items,
-            next_step=NextStepHandoffReturnControl(previous_agents.pop())
-            if is_returning_control
-            else NextStepFinalOutput(final_output),
+            next_step=NextStepFinalOutput(final_output),
         )
 
     @classmethod
@@ -864,19 +844,13 @@ async def run_final_output_hooks(
         hooks: RunHooks[TContext],
         context_wrapper: RunContextWrapper[TContext],
         final_output: Any,
-        is_returning_control: bool,
     ):
-        # If the agent is not returning control, run the hooks
-        if not is_returning_control:
-            await asyncio.gather(
-                hooks.on_agent_end(context_wrapper, agent, final_output),
-                agent.hooks.on_end(context_wrapper, agent, final_output)
-                if agent.hooks
-                else _coro.noop_coroutine(),
-            )
-        # If the agent is returning control, only run the current agent's hooks
-        elif agent.hooks:
-            await agent.hooks.on_end(context_wrapper, agent, final_output)
+        await asyncio.gather(
+            hooks.on_agent_end(context_wrapper, agent, final_output),
+            agent.hooks.on_end(context_wrapper, agent, final_output)
+            if agent.hooks
+            else _coro.noop_coroutine(),
+        )
 
     @classmethod
     async def run_single_input_guardrail(

src/agents/handoffs.py

@@ -105,12 +105,6 @@ 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})
 
@@ -134,7 +128,6 @@ def handoff(
     tool_description_override: str | None = None,
     input_filter: Callable[[HandoffInputData], HandoffInputData] | None = None,
     is_enabled: bool | Callable[[RunContextWrapper[Any], Agent[Any]], MaybeAwaitable[bool]] = True,
-    should_return_control: bool = False,
 ) -> Handoff[TContext]: ...
 
 
@@ -148,7 +141,6 @@ def handoff(
     tool_name_override: str | None = None,
     input_filter: Callable[[HandoffInputData], HandoffInputData] | None = None,
     is_enabled: bool | Callable[[RunContextWrapper[Any], Agent[Any]], MaybeAwaitable[bool]] = True,
-    should_return_control: bool = False,
 ) -> Handoff[TContext]: ...
 
 
@@ -161,7 +153,6 @@ def handoff(
     tool_name_override: str | None = None,
     input_filter: Callable[[HandoffInputData], HandoffInputData] | None = None,
     is_enabled: bool | Callable[[RunContextWrapper[Any], Agent[Any]], MaybeAwaitable[bool]] = True,
-    should_return_control: bool = False,
 ) -> Handoff[TContext]: ...
 
 
@@ -173,7 +164,6 @@ def handoff(
     input_type: type[THandoffInput] | None = None,
     input_filter: Callable[[HandoffInputData], HandoffInputData] | None = None,
     is_enabled: bool | Callable[[RunContextWrapper[Any], Agent[Any]], MaybeAwaitable[bool]] = True,
-    should_return_control: bool = False,
 ) -> Handoff[TContext]:
     """Create a handoff from an agent.
 
@@ -191,7 +181,7 @@ def handoff(
             hidden from the LLM at runtime.
     """
     assert (on_handoff and input_type) or not (on_handoff and input_type), (
-        "You must provide either both on_input and input_type, or neither"
+        "You must provide either both on_handoff and input_type, or neither"
     )
     type_adapter: TypeAdapter[Any] | None
     if input_type is not None:
@@ -257,5 +247,4 @@ async def _invoke_handoff(
         input_filter=input_filter,
         agent_name=agent.name,
         is_enabled=is_enabled,
-        should_return_control=should_return_control,
     )