feat: add invalid final output recovery handler (#3736)

2026-07-06 15:25:54 +09:00
parent 1643dbe617
commit 3d07ff1687
7 changed files with 470 additions and 10 deletions
@@ -1073,6 +1073,7 @@ class AgentRunner:
                        )
                        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,
@@ -7,7 +7,7 @@ from typing import Any, Generic
 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__ = [
@@ -2,7 +2,7 @@ from __future__ import annotations

 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 ..run_error_handlers import (
 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(
@@ -901,6 +901,7 @@ async def start_streaming(
                )
                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,
@@ -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,
@@ -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
@@ -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 test_resolve_run_error_handler_result_covers_async_and_validation_path

    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 test_resolve_run_error_handler_result_covers_async_and_validation_path
            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,