fix(core): validate tool_use_behavior callback results (#5173)

fix: validate tool_use_behavior callback results

Author: jbeckwith-oai

src/agents/run_internal/turn_resolution.py

@@ -781,7 +781,11 @@ async def check_for_final_output_from_tools(
     elif callable(agent.tool_use_behavior):
         result = agent.tool_use_behavior(context_wrapper, tool_results)
         if inspect.isawaitable(result):
-            return await result
+            result = await result
+        if not isinstance(result, ToolsToFinalOutputResult):
+            raise UserError(
+                "Agent tool_use_behavior callable must return ToolsToFinalOutputResult."
+            )
         return result
 
     logger.error("Invalid tool_use_behavior: %s", agent.tool_use_behavior)

tests/test_tool_use_behavior.py

@@ -11,13 +11,15 @@
     Agent,
     FunctionToolResult,
     RunContextWrapper,
+    Runner,
     ToolCallOutputItem,
     ToolsToFinalOutputResult,
     UserError,
     function_tool,
     tool_namespace,
 )
 from agents.run_internal import run_loop
+from agents.testing import ScriptedModel, function_call
 
 from .test_responses import get_function_tool
 
@@ -175,6 +177,54 @@ async def __call__(
     assert behavior.calls == 1
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("streamed", [False, True])
+@pytest.mark.parametrize("async_behavior", [False, True])
+@pytest.mark.parametrize("valid_result", [False, True])
+async def test_custom_tool_use_behavior_return_contract(
+    streamed: bool, async_behavior: bool, valid_result: bool
+) -> None:
+    calls = 0
+
+    def behavior(context: RunContextWrapper, results: list[FunctionToolResult]) -> Any:
+        nonlocal calls
+        calls += 1
+        assert len(results) == 1
+        assert results[0].output == "pong"
+        if valid_result:
+            return ToolsToFinalOutputResult(is_final_output=True, final_output="done")
+        return "done"
+
+    async def async_callback(context: RunContextWrapper, results: list[FunctionToolResult]) -> Any:
+        return behavior(context, results)
+
+    model = ScriptedModel([[function_call("ping", {}, call_id="call_1")]])
+    agent = Agent(
+        name="test",
+        tools=[get_function_tool("ping", return_value="pong")],
+        model=model,
+        tool_use_behavior=async_callback if async_behavior else behavior,
+    )
+
+    async def run() -> Any:
+        if streamed:
+            result = Runner.run_streamed(agent, "test")
+            async for _ in result.stream_events():
+                pass
+            return result.final_output
+        return (await Runner.run(agent, "test")).final_output
+
+    if valid_result:
+        assert await run() == "done"
+    else:
+        with pytest.raises(
+            UserError, match="tool_use_behavior callable must return ToolsToFinalOutputResult"
+        ):
+            await run()
+    assert calls == 1
+    assert len(model.calls) == 1
+
+
 @pytest.mark.asyncio
 async def test_invalid_tool_use_behavior_raises() -> None:
     """If tool_use_behavior is invalid, we should raise a UserError."""