From 20198e12c61604c71004b2aa4025e25cf5d98635 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Mon, 6 Jul 2026 15:11:56 +0900 Subject: [PATCH] feat: add invalid final output recovery handler --- src/agents/run.py | 1 + src/agents/run_error_handlers.py | 5 +- src/agents/run_internal/error_handlers.py | 12 +- src/agents/run_internal/run_loop.py | 1 + src/agents/run_internal/turn_resolution.py | 86 ++++- tests/test_invalid_final_output_handler.py | 372 +++++++++++++++++++++ tests/test_run_internal_error_handlers.py | 3 + 7 files changed, 470 insertions(+), 10 deletions(-) create mode 100644 tests/test_invalid_final_output_handler.py diff --git a/src/agents/run.py b/src/agents/run.py index 014271a5ea..02ba5c5dd7 100644 --- a/src/agents/run.py +++ b/src/agents/run.py @@ -1073,6 +1073,7 @@ def _finalize_result(result: RunResult) -> RunResult: ) handler_result = await resolve_run_error_handler_result( error_handlers=error_handlers, + error_kind="max_turns", error=max_turns_error, context_wrapper=context_wrapper, run_data=run_error_data, diff --git a/src/agents/run_error_handlers.py b/src/agents/run_error_handlers.py index 6f345852eb..cdf83f3149 100644 --- a/src/agents/run_error_handlers.py +++ b/src/agents/run_error_handlers.py @@ -7,7 +7,7 @@ from typing_extensions import TypedDict from .agent import Agent -from .exceptions import MaxTurnsExceeded, ModelRefusalError +from .exceptions import MaxTurnsExceeded, ModelBehaviorError, ModelRefusalError from .items import ModelResponse, RunItem, TResponseInputItem from .run_context import RunContextWrapper, TContext from .util._types import MaybeAwaitable @@ -27,7 +27,7 @@ class RunErrorData: @dataclass class RunErrorHandlerInput(Generic[TContext]): - error: MaxTurnsExceeded | ModelRefusalError + error: MaxTurnsExceeded | ModelRefusalError | ModelBehaviorError context: RunContextWrapper[TContext] run_data: RunErrorData @@ -52,6 +52,7 @@ class RunErrorHandlers(TypedDict, Generic[TContext], total=False): max_turns: RunErrorHandler[TContext] model_refusal: RunErrorHandler[TContext] + invalid_final_output: RunErrorHandler[TContext] __all__ = [ diff --git a/src/agents/run_internal/error_handlers.py b/src/agents/run_internal/error_handlers.py index 81a94a2002..39a291d4f9 100644 --- a/src/agents/run_internal/error_handlers.py +++ b/src/agents/run_internal/error_handlers.py @@ -2,7 +2,7 @@ import inspect import json -from typing import Any +from typing import Any, Literal from openai.types.responses import ResponseOutputMessage, ResponseOutputText @@ -27,6 +27,8 @@ from .items import ReasoningItemIdPolicy, run_item_to_input_item from .turn_preparation import get_output_schema +RunErrorHandlerKind = Literal["max_turns", "model_refusal", "invalid_final_output"] + def build_run_error_data( *, @@ -128,16 +130,14 @@ def create_message_output_item(agent: Agent[Any], output_text: str) -> MessageOu async def resolve_run_error_handler_result( *, error_handlers: RunErrorHandlers[TContext] | None, - error: MaxTurnsExceeded | ModelRefusalError, + error_kind: RunErrorHandlerKind, + error: MaxTurnsExceeded | ModelRefusalError | ModelBehaviorError, context_wrapper: RunContextWrapper[TContext], run_data: RunErrorData, ) -> RunErrorHandlerResult | None: if not error_handlers: return None - if isinstance(error, ModelRefusalError): - handler = error_handlers.get("model_refusal") - else: - handler = error_handlers.get("max_turns") + handler = error_handlers.get(error_kind) if handler is None: return None handler_input = RunErrorHandlerInput( diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index 45f09c0fa0..d0a84a7809 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -901,6 +901,7 @@ async def _save_stream_items_without_count( ) handler_result = await resolve_run_error_handler_result( error_handlers=error_handlers, + error_kind="max_turns", error=max_turns_error, context_wrapper=context_wrapper, run_data=run_error_data, diff --git a/src/agents/run_internal/turn_resolution.py b/src/agents/run_internal/turn_resolution.py index 2c95cf2e13..33e1bee8e2 100644 --- a/src/agents/run_internal/turn_resolution.py +++ b/src/agents/run_internal/turn_resolution.py @@ -368,6 +368,44 @@ async def execute_final_output( ) +async def _resolve_invalid_final_output( + *, + error_handlers: RunErrorHandlers[TContext] | None, + error: ModelBehaviorError, + public_agent: Agent[TContext], + original_input: str | list[TResponseInputItem], + new_response: ModelResponse, + new_items: list[RunItem], + context_wrapper: RunContextWrapper[TContext], +) -> tuple[Any, MessageOutputItem | None] | None: + run_error_data = build_run_error_data( + input=original_input, + new_items=new_items, + raw_responses=[new_response], + last_agent=public_agent, + ) + handler_result = await resolve_run_error_handler_result( + error_handlers=error_handlers, + error_kind="invalid_final_output", + error=error, + context_wrapper=context_wrapper, + run_data=run_error_data, + ) + if handler_result is None: + return None + + final_output = validate_handler_final_output(public_agent, handler_result.final_output) + message_item = ( + create_message_output_item( + public_agent, + format_final_output_text(public_agent, final_output), + ) + if handler_result.include_in_history + else None + ) + return final_output, message_item + + def _resolve_server_managed_handoff_behavior( *, handoff: Handoff[Any, Agent[Any]], @@ -781,6 +819,7 @@ async def execute_tools_and_side_effects( ) handler_result = await resolve_run_error_handler_result( error_handlers=error_handlers, + error_kind="model_refusal", error=refusal_error, context_wrapper=context_wrapper, run_data=run_error_data, @@ -806,8 +845,51 @@ async def execute_tools_and_side_effects( tool_input_guardrail_results=tool_input_guardrail_results, tool_output_guardrail_results=tool_output_guardrail_results, ) - if output_schema and not output_schema.is_plain_text() and potential_final_output_text: - final_output = output_schema.validate_json(potential_final_output_text) + if output_schema and not output_schema.is_plain_text(): + if potential_final_output_text: + try: + final_output = output_schema.validate_json(potential_final_output_text) + except ModelBehaviorError as error: + resolved_handler_output = await _resolve_invalid_final_output( + error_handlers=error_handlers, + error=error, + public_agent=public_agent, + original_input=original_input, + new_response=new_response, + new_items=pre_step_items + new_step_items, + context_wrapper=context_wrapper, + ) + if resolved_handler_output is None: + raise + final_output, message_item = resolved_handler_output + if message_item is not None: + new_step_items.append(message_item) + else: + resolved_handler_output = await _resolve_invalid_final_output( + error_handlers=error_handlers, + error=ModelBehaviorError( + "Model returned no final output for the structured output type." + ), + public_agent=public_agent, + original_input=original_input, + new_response=new_response, + new_items=pre_step_items + new_step_items, + context_wrapper=context_wrapper, + ) + if resolved_handler_output is None: + return SingleStepResult( + original_input=original_input, + model_response=new_response, + pre_step_items=pre_step_items, + new_step_items=new_step_items, + next_step=NextStepRunAgain(), + tool_input_guardrail_results=tool_input_guardrail_results, + tool_output_guardrail_results=tool_output_guardrail_results, + ) + final_output, message_item = resolved_handler_output + if message_item is not None: + new_step_items.append(message_item) + return await execute_final_output_call( public_agent=public_agent, original_input=original_input, diff --git a/tests/test_invalid_final_output_handler.py b/tests/test_invalid_final_output_handler.py new file mode 100644 index 0000000000..b4519debdc --- /dev/null +++ b/tests/test_invalid_final_output_handler.py @@ -0,0 +1,372 @@ +from __future__ import annotations + +import json +from typing import Any + +import pytest +from openai.types.responses import ResponseOutputMessage +from pydantic import BaseModel + +from agents import ( + Agent, + AgentHookContext, + GuardrailFunctionOutput, + ItemHelpers, + MessageOutputItem, + ModelBehaviorError, + OutputGuardrail, + RunContextWrapper, + RunErrorHandlerInput, + RunErrorHandlerResult, + RunErrorHandlers, + RunHooks, + Runner, + UserError, + function_tool, +) +from agents.items import TResponseInputItem, TResponseOutputItem +from agents.stream_events import RunItemStreamEvent + +from .fake_model import FakeModel +from .test_responses import get_function_tool_call, get_text_message +from .utils.simple_session import SimpleListSession + + +class FinalOutput(BaseModel): + summary: str + + +class RecordingRunHooks(RunHooks[None]): + def __init__(self) -> None: + self.final_outputs: list[Any] = [] + + async def on_agent_end( + self, + context: AgentHookContext[None], + agent: Agent[None], + output: Any, + ) -> None: + self.final_outputs.append(output) + + +def _message_texts(items: list[TResponseInputItem]) -> list[str]: + texts: list[str] = [] + for item in items: + if not isinstance(item, dict) or item.get("type") != "message": + continue + message = ResponseOutputMessage.model_validate(item) + texts.append(ItemHelpers.extract_text(message) or "") + return texts + + +@pytest.mark.asyncio +async def test_invalid_final_output_raises_without_handler() -> None: + model = FakeModel(initial_output=[get_text_message("not valid json")]) + agent = Agent(name="test", model=model, output_type=FinalOutput) + + with pytest.raises(ModelBehaviorError, match="Invalid JSON"): + await Runner.run(agent, input="user_message") + + +@pytest.mark.asyncio +async def test_invalid_final_output_handler_returns_validated_fallback() -> None: + model = FakeModel(initial_output=[get_text_message("not valid json")]) + agent = Agent(name="test", model=model, output_type=FinalOutput) + + def handler(data: RunErrorHandlerInput[None]) -> FinalOutput: + assert isinstance(data.error, ModelBehaviorError) + assert data.run_data.raw_responses + assert ItemHelpers.text_message_outputs(data.run_data.new_items) == "not valid json" + return FinalOutput(summary="safe fallback") + + result = await Runner.run( + agent, + input="user_message", + error_handlers={"invalid_final_output": handler}, + ) + + assert result.final_output == FinalOutput(summary="safe fallback") + assert [ + ItemHelpers.text_message_output(item) + for item in result.new_items + if isinstance(item, MessageOutputItem) + ] == ["not valid json", '{"summary":"safe fallback"}'] + + +@pytest.mark.asyncio +async def test_invalid_final_output_handler_can_skip_fallback_history() -> None: + model = FakeModel(initial_output=[get_text_message("not valid json")]) + agent = Agent(name="test", model=model, output_type=FinalOutput) + + result = await Runner.run( + agent, + input="user_message", + error_handlers={ + "invalid_final_output": lambda _data: RunErrorHandlerResult( + final_output=FinalOutput(summary="safe fallback"), + include_in_history=False, + ) + }, + ) + + assert result.final_output == FinalOutput(summary="safe fallback") + assert ItemHelpers.text_message_outputs(result.new_items) == "not valid json" + + +@pytest.mark.asyncio +async def test_invalid_final_output_handler_rejects_invalid_fallback() -> None: + model = FakeModel(initial_output=[get_text_message("not valid json")]) + agent = Agent(name="test", model=model, output_type=FinalOutput) + + with pytest.warns(UserWarning, match="Pydantic serializer warnings"): + with pytest.raises(UserError, match="Invalid run error handler final_output"): + await Runner.run( + agent, + input="user_message", + error_handlers={"invalid_final_output": lambda _data: {"unexpected": "value"}}, + ) + + +@pytest.mark.asyncio +async def test_invalid_final_output_handler_can_decline_recovery() -> None: + model = FakeModel(initial_output=[get_text_message("not valid json")]) + agent = Agent(name="test", model=model, output_type=FinalOutput) + + with pytest.raises(ModelBehaviorError, match="Invalid JSON"): + await Runner.run( + agent, + input="user_message", + error_handlers={"invalid_final_output": lambda _data: None}, + ) + + +@pytest.mark.asyncio +async def test_invalid_final_output_handler_does_not_catch_other_model_behavior_errors() -> None: + model = FakeModel(initial_output=[get_function_tool_call("missing_tool")]) + agent = Agent(name="test", model=model, output_type=FinalOutput) + handler_called = False + + def handler(_data: RunErrorHandlerInput[None]) -> FinalOutput: + nonlocal handler_called + handler_called = True + return FinalOutput(summary="safe fallback") + + with pytest.raises(ModelBehaviorError, match="not found"): + await Runner.run( + agent, + input="user_message", + error_handlers={"invalid_final_output": handler}, + ) + + assert not handler_called + + +@pytest.mark.asyncio +@pytest.mark.parametrize("invalid_output", [[], [get_text_message("")]]) +async def test_empty_structured_output_handler_avoids_another_model_turn( + invalid_output: list[TResponseOutputItem], +) -> None: + model = FakeModel() + model.add_multiple_turn_outputs([invalid_output, [get_text_message('{"summary":"unused"}')]]) + agent = Agent(name="test", model=model, output_type=FinalOutput) + + def handler(data: RunErrorHandlerInput[None]) -> FinalOutput: + assert isinstance(data.error, ModelBehaviorError) + assert data.error.message == ( + "Model returned no final output for the structured output type." + ) + return FinalOutput(summary="safe fallback") + + result = await Runner.run( + agent, + input="user_message", + error_handlers={"invalid_final_output": handler}, + ) + + assert result.final_output == FinalOutput(summary="safe fallback") + assert len(model.turn_outputs) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error_handlers", + [None, {"invalid_final_output": lambda _data: None}], +) +async def test_empty_structured_output_without_fallback_keeps_existing_next_turn_behavior( + error_handlers: RunErrorHandlers[None] | None, +) -> None: + model = FakeModel() + model.add_multiple_turn_outputs([[], [get_text_message('{"summary":"second turn"}')]]) + agent = Agent(name="test", model=model, output_type=FinalOutput) + + result = await Runner.run(agent, input="user_message", error_handlers=error_handlers) + + assert result.final_output == FinalOutput(summary="second turn") + assert not model.turn_outputs + + +@pytest.mark.asyncio +async def test_streamed_invalid_final_output_emits_exact_fallback_item() -> None: + model = FakeModel(initial_output=[get_text_message("not valid json")]) + agent = Agent(name="test", model=model, output_type=FinalOutput) + session = SimpleListSession() + + result = Runner.run_streamed( + agent, + input="user_message", + session=session, + error_handlers={"invalid_final_output": lambda _data: FinalOutput(summary="safe fallback")}, + ) + events = [event async for event in result.stream_events()] + + assert result.final_output == FinalOutput(summary="safe fallback") + fallback_events = [ + event + for event in events + if isinstance(event, RunItemStreamEvent) + and event.name == "message_output_created" + and isinstance(event.item, MessageOutputItem) + and ItemHelpers.text_message_output(event.item) == '{"summary":"safe fallback"}' + ] + assert len(fallback_events) == 1 + assert [ + ItemHelpers.text_message_output(item) + for item in result.new_items + if isinstance(item, MessageOutputItem) + ] == ["not valid json", '{"summary":"safe fallback"}'] + assert _message_texts(await session.get_items()) == [ + "not valid json", + '{"summary":"safe fallback"}', + ] + + +@pytest.mark.asyncio +async def test_streamed_empty_structured_output_handler_avoids_another_model_turn() -> None: + model = FakeModel() + model.add_multiple_turn_outputs([[], [get_text_message('{"summary":"unused"}')]]) + agent = Agent(name="test", model=model, output_type=FinalOutput) + + result = Runner.run_streamed( + agent, + input="user_message", + error_handlers={"invalid_final_output": lambda _data: FinalOutput(summary="safe fallback")}, + ) + events = [event async for event in result.stream_events()] + + assert result.final_output == FinalOutput(summary="safe fallback") + assert len(model.turn_outputs) == 1 + assert any( + isinstance(event, RunItemStreamEvent) + and event.name == "message_output_created" + and isinstance(event.item, MessageOutputItem) + and ItemHelpers.text_message_output(event.item) == '{"summary":"safe fallback"}' + for event in events + ) + + +@pytest.mark.asyncio +async def test_invalid_final_output_fallback_runs_hooks_and_output_guardrails() -> None: + model = FakeModel(initial_output=[get_text_message("not valid json")]) + hooks = RecordingRunHooks() + guarded_outputs: list[Any] = [] + + def record_output( + context: RunContextWrapper[None], + agent: Agent[Any], + output: Any, + ) -> GuardrailFunctionOutput: + guarded_outputs.append(output) + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=False) + + agent = Agent( + name="test", + model=model, + output_type=FinalOutput, + output_guardrails=[OutputGuardrail(guardrail_function=record_output)], + ) + + result = await Runner.run( + agent, + input="user_message", + hooks=hooks, + error_handlers={"invalid_final_output": lambda _data: FinalOutput(summary="safe fallback")}, + ) + + expected = FinalOutput(summary="safe fallback") + assert result.final_output == expected + assert hooks.final_outputs == [expected] + assert guarded_outputs == [expected] + assert len(result.output_guardrail_results) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streamed", [False, True]) +async def test_invalid_final_output_fallback_does_not_retry_or_replay_tools( + streamed: bool, +) -> None: + side_effects: list[str] = [] + + @function_tool + async def record_side_effect(value: str) -> str: + side_effects.append(value) + return f"recorded:{value}" + + model = FakeModel() + model.add_multiple_turn_outputs( + [ + [ + get_function_tool_call( + "record_side_effect", + json.dumps({"value": "once"}), + call_id="first_call", + ) + ], + [get_text_message("not valid json")], + [ + get_function_tool_call( + "record_side_effect", + json.dumps({"value": "replayed"}), + call_id="replayed_call", + ) + ], + [get_text_message('{"summary":"unexpected retry"}')], + ] + ) + agent = Agent( + name="test", + model=model, + tools=[record_side_effect], + output_type=FinalOutput, + ) + error_handlers: RunErrorHandlers[None] = { + "invalid_final_output": lambda _data: FinalOutput(summary="safe fallback") + } + + if streamed: + streamed_result = Runner.run_streamed( + agent, + input="user_message", + error_handlers=error_handlers, + ) + events = [event async for event in streamed_result.stream_events()] + final_output = streamed_result.final_output + fallback_events = [ + event + for event in events + if isinstance(event, RunItemStreamEvent) + and event.name == "message_output_created" + and isinstance(event.item, MessageOutputItem) + and ItemHelpers.text_message_output(event.item) == '{"summary":"safe fallback"}' + ] + assert len(fallback_events) == 1 + else: + result = await Runner.run( + agent, + input="user_message", + error_handlers=error_handlers, + ) + final_output = result.final_output + + assert final_output == FinalOutput(summary="safe fallback") + assert side_effects == ["once"] + assert len(model.turn_outputs) == 2 diff --git a/tests/test_run_internal_error_handlers.py b/tests/test_run_internal_error_handlers.py index 48574ded65..ba1efd8e24 100644 --- a/tests/test_run_internal_error_handlers.py +++ b/tests/test_run_internal_error_handlers.py @@ -92,6 +92,7 @@ async def test_resolve_run_error_handler_result_covers_async_and_validation_path no_handler = await run_error_handlers.resolve_run_error_handler_result( error_handlers={}, + error_kind="max_turns", error=error, context_wrapper=context_wrapper, run_data=run_data, @@ -103,6 +104,7 @@ async def async_handler(_handler_input: Any) -> None: async_none = await run_error_handlers.resolve_run_error_handler_result( error_handlers={"max_turns": async_handler}, + error_kind="max_turns", error=error, context_wrapper=context_wrapper, run_data=run_data, @@ -114,6 +116,7 @@ async def async_handler(_handler_input: Any) -> None: error_handlers={ "max_turns": lambda _handler_input: {"final_output": "x", "extra": "y"} }, + error_kind="max_turns", error=error, context_wrapper=context_wrapper, run_data=run_data,