Fix #1238 by enhancing HandoffInputData and enable passing async functions (#1302)

This pull request resolves #1238

Author: seratch

examples/handoffs/message_filter.py

@@ -24,6 +24,7 @@ def spanish_handoff_message_filter(handoff_message_data: HandoffInputData) -> Ha
         else handoff_message_data.input_history
     )
 
+    # or, you can use the HandoffInputData.clone(kwargs) method
     return HandoffInputData(
         input_history=history,
         pre_handoff_items=tuple(handoff_message_data.pre_handoff_items),

examples/handoffs/message_filter_streaming.py

@@ -24,6 +24,7 @@ def spanish_handoff_message_filter(handoff_message_data: HandoffInputData) -> Ha
         else handoff_message_data.input_history
     )
 
+    # or, you can use the HandoffInputData.clone(kwargs) method
     return HandoffInputData(
         input_history=history,
         pre_handoff_items=tuple(handoff_message_data.pre_handoff_items),

src/agents/_run_impl.py

@@ -774,6 +774,7 @@ async def execute_handoffs(
                     else original_input,
                     pre_handoff_items=tuple(pre_step_items),
                     new_items=tuple(new_step_items),
+                    run_context=context_wrapper,
                 )
                 if not callable(input_filter):
                     _error_tracing.attach_error_to_span(
@@ -785,6 +786,8 @@ async def execute_handoffs(
                     )
                     raise UserError(f"Invalid input filter: {input_filter}")
                 filtered = input_filter(handoff_input_data)
+                if inspect.isawaitable(filtered):
+                    filtered = await filtered
                 if not isinstance(filtered, HandoffInputData):
                     _error_tracing.attach_error_to_span(
                         span_handoff,

src/agents/extensions/handoff_filters.py

@@ -29,6 +29,7 @@ def remove_all_tools(handoff_input_data: HandoffInputData) -> HandoffInputData:
         input_history=filtered_history,
         pre_handoff_items=filtered_pre_handoff_items,
         new_items=filtered_new_items,
+        run_context=handoff_input_data.run_context,
     )
 
 

src/agents/handoffs.py

@@ -3,7 +3,7 @@
 import inspect
 import json
 from collections.abc import Awaitable
-from dataclasses import dataclass
+from dataclasses import dataclass, replace as dataclasses_replace
 from typing import TYPE_CHECKING, Any, Callable, Generic, cast, overload
 
 from pydantic import TypeAdapter
@@ -49,8 +49,24 @@ class HandoffInputData:
     handoff and the tool output message representing the response from the handoff output.
     """
 
+    run_context: RunContextWrapper[Any] | None = None
+    """
+    The run context at the time the handoff was invoked.
+    Note that, since this property was added later on, it's optional for backwards compatibility.
+    """
+
+    def clone(self, **kwargs: Any) -> HandoffInputData:
+        """
+        Make a copy of the handoff input data, with the given arguments changed. For example, you
+        could do:
+        ```
+        new_handoff_input_data = handoff_input_data.clone(new_items=())
+        ```
+        """
+        return dataclasses_replace(self, **kwargs)
+
 
-HandoffInputFilter: TypeAlias = Callable[[HandoffInputData], HandoffInputData]
+HandoffInputFilter: TypeAlias = Callable[[HandoffInputData], MaybeAwaitable[HandoffInputData]]
 """A function that filters the input data passed to the next agent."""
 
 

tests/test_agent_runner.py

@@ -224,6 +224,7 @@ def remove_new_items(handoff_input_data: HandoffInputData) -> HandoffInputData:
         input_history=handoff_input_data.input_history,
         pre_handoff_items=(),
         new_items=(),
+        run_context=handoff_input_data.run_context,
     )
 
 
@@ -262,7 +263,7 @@ async def test_handoff_filters():
 
 
 @pytest.mark.asyncio
-async def test_async_input_filter_fails():
+async def test_async_input_filter_supported():
     # DO NOT rename this without updating pyproject.toml
 
     model = FakeModel()
@@ -274,7 +275,7 @@ async def test_async_input_filter_fails():
     async def on_invoke_handoff(_ctx: RunContextWrapper[Any], _input: str) -> Agent[Any]:
         return agent_1
 
-    async def invalid_input_filter(data: HandoffInputData) -> HandoffInputData:
+    async def async_input_filter(data: HandoffInputData) -> HandoffInputData:
         return data  # pragma: no cover
 
     agent_2 = Agent[None](
@@ -287,8 +288,7 @@ async def invalid_input_filter(data: HandoffInputData) -> HandoffInputData:
                 input_json_schema={},
                 on_invoke_handoff=on_invoke_handoff,
                 agent_name=agent_1.name,
-                # Purposely ignoring the type error here to simulate invalid input
-                input_filter=invalid_input_filter,  # type: ignore
+                input_filter=async_input_filter,
             )
         ],
     )
@@ -300,8 +300,8 @@ async def invalid_input_filter(data: HandoffInputData) -> HandoffInputData:
         ]
     )
 
-    with pytest.raises(UserError):
-        await Runner.run(agent_2, input="user_message")
+    result = await Runner.run(agent_2, input="user_message")
+    assert result.final_output == "last"
 
 
 @pytest.mark.asyncio

tests/test_agent_runner_streamed.py

@@ -241,6 +241,7 @@ def remove_new_items(handoff_input_data: HandoffInputData) -> HandoffInputData:
         input_history=handoff_input_data.input_history,
         pre_handoff_items=(),
         new_items=(),
+        run_context=handoff_input_data.run_context,
     )
 
 
@@ -281,7 +282,7 @@ async def test_handoff_filters():
 
 
 @pytest.mark.asyncio
-async def test_async_input_filter_fails():
+async def test_async_input_filter_supported():
     # DO NOT rename this without updating pyproject.toml
 
     model = FakeModel()
@@ -293,7 +294,7 @@ async def test_async_input_filter_fails():
     async def on_invoke_handoff(_ctx: RunContextWrapper[Any], _input: str) -> Agent[Any]:
         return agent_1
 
-    async def invalid_input_filter(data: HandoffInputData) -> HandoffInputData:
+    async def async_input_filter(data: HandoffInputData) -> HandoffInputData:
         return data  # pragma: no cover
 
     agent_2 = Agent[None](
@@ -306,8 +307,7 @@ async def invalid_input_filter(data: HandoffInputData) -> HandoffInputData:
                 input_json_schema={},
                 on_invoke_handoff=on_invoke_handoff,
                 agent_name=agent_1.name,
-                # Purposely ignoring the type error here to simulate invalid input
-                input_filter=invalid_input_filter,  # type: ignore
+                input_filter=async_input_filter,
             )
         ],
     )
@@ -319,10 +319,9 @@ async def invalid_input_filter(data: HandoffInputData) -> HandoffInputData:
         ]
     )
 
-    with pytest.raises(UserError):
-        result = Runner.run_streamed(agent_2, input="user_message")
-        async for _ in result.stream_events():
-            pass
+    result = Runner.run_streamed(agent_2, input="user_message")
+    async for _ in result.stream_events():
+        pass
 
 
 @pytest.mark.asyncio

tests/test_extension_filters.py

@@ -1,6 +1,6 @@
 from openai.types.responses import ResponseOutputMessage, ResponseOutputText
 
-from agents import Agent, HandoffInputData
+from agents import Agent, HandoffInputData, RunContextWrapper
 from agents.extensions.handoff_filters import remove_all_tools
 from agents.items import (
     HandoffOutputItem,
@@ -78,13 +78,23 @@ def _get_handoff_output_run_item(content: str) -> HandoffOutputItem:
 
 
 def test_empty_data():
-    handoff_input_data = HandoffInputData(input_history=(), pre_handoff_items=(), new_items=())
+    handoff_input_data = HandoffInputData(
+        input_history=(),
+        pre_handoff_items=(),
+        new_items=(),
+        run_context=RunContextWrapper(context=()),
+    )
     filtered_data = remove_all_tools(handoff_input_data)
     assert filtered_data == handoff_input_data
 
 
 def test_str_historyonly():
-    handoff_input_data = HandoffInputData(input_history="Hello", pre_handoff_items=(), new_items=())
+    handoff_input_data = HandoffInputData(
+        input_history="Hello",
+        pre_handoff_items=(),
+        new_items=(),
+        run_context=RunContextWrapper(context=()),
+    )
     filtered_data = remove_all_tools(handoff_input_data)
     assert filtered_data == handoff_input_data
 
@@ -94,6 +104,7 @@ def test_str_history_and_list():
         input_history="Hello",
         pre_handoff_items=(),
         new_items=(_get_message_output_run_item("Hello"),),
+        run_context=RunContextWrapper(context=()),
     )
     filtered_data = remove_all_tools(handoff_input_data)
     assert filtered_data == handoff_input_data
@@ -104,6 +115,7 @@ def test_list_history_and_list():
         input_history=(_get_message_input_item("Hello"),),
         pre_handoff_items=(_get_message_output_run_item("123"),),
         new_items=(_get_message_output_run_item("World"),),
+        run_context=RunContextWrapper(context=()),
     )
     filtered_data = remove_all_tools(handoff_input_data)
     assert filtered_data == handoff_input_data
@@ -121,6 +133,7 @@ def test_removes_tools_from_history():
             _get_message_output_run_item("123"),
         ),
         new_items=(_get_message_output_run_item("World"),),
+        run_context=RunContextWrapper(context=()),
     )
     filtered_data = remove_all_tools(handoff_input_data)
     assert len(filtered_data.input_history) == 2
@@ -136,6 +149,7 @@ def test_removes_tools_from_new_items():
             _get_message_output_run_item("Hello"),
             _get_tool_output_run_item("World"),
         ),
+        run_context=RunContextWrapper(context=()),
     )
     filtered_data = remove_all_tools(handoff_input_data)
     assert len(filtered_data.input_history) == 0
@@ -158,6 +172,7 @@ def test_removes_tools_from_new_items_and_history():
             _get_message_output_run_item("Hello"),
             _get_tool_output_run_item("World"),
         ),
+        run_context=RunContextWrapper(context=()),
     )
     filtered_data = remove_all_tools(handoff_input_data)
     assert len(filtered_data.input_history) == 2
@@ -181,6 +196,7 @@ def test_removes_handoffs_from_history():
             _get_tool_output_run_item("World"),
             _get_handoff_output_run_item("World"),
         ),
+        run_context=RunContextWrapper(context=()),
     )
     filtered_data = remove_all_tools(handoff_input_data)
     assert len(filtered_data.input_history) == 1

tests/test_handoff_tool.py

@@ -221,13 +221,15 @@ def test_handoff_input_data():
         input_history="",
         pre_handoff_items=(),
         new_items=(),
+        run_context=RunContextWrapper(context=()),
     )
     assert get_len(data) == 1
 
     data = HandoffInputData(
         input_history=({"role": "user", "content": "foo"},),
         pre_handoff_items=(),
         new_items=(),
+        run_context=RunContextWrapper(context=()),
     )
     assert get_len(data) == 1
 
@@ -238,6 +240,7 @@ def test_handoff_input_data():
         ),
         pre_handoff_items=(),
         new_items=(),
+        run_context=RunContextWrapper(context=()),
     )
     assert get_len(data) == 2
 
@@ -251,6 +254,7 @@ def test_handoff_input_data():
             message_item("bar", agent),
             message_item("baz", agent),
         ),
+        run_context=RunContextWrapper(context=()),
     )
     assert get_len(data) == 5
 
@@ -264,6 +268,7 @@ def test_handoff_input_data():
             message_item("baz", agent),
             message_item("qux", agent),
         ),
+        run_context=RunContextWrapper(context=()),
     )
 
     assert get_len(data) == 5