pydantic--pydantic-ai/tests/models/test_fallback.py

from __future__ import annotations

import dataclasses
import json
import sys
from collections.abc import AsyncIterator
from datetime import timezone
from typing import Any, Literal, cast

import pytest
from dirty_equals import IsJson
from pydantic import BaseModel
from pydantic_core import to_json
from typing_extensions import TypedDict

from pydantic_ai import (
    Agent,
    ModelAPIError,
    ModelHTTPError,
    ModelMessage,
    ModelProfile,
    ModelRequest,
    ModelResponse,
    TextPart,
    ToolCallPart,
    ToolDefinition,
    UserPromptPart,
)
from pydantic_ai.capabilities.instrumentation import Instrumentation
from pydantic_ai.messages import InstructionPart, NativeToolCallPart, NativeToolReturnPart
from pydantic_ai.models import ModelRequestParameters
from pydantic_ai.models.fallback import FallbackModel, ResponseRejected
from pydantic_ai.models.function import AgentInfo, FunctionModel
from pydantic_ai.models.instrumented import InstrumentationSettings, InstrumentedModel
from pydantic_ai.output import OutputObjectDefinition
from pydantic_ai.settings import ModelSettings
from pydantic_ai.usage import RequestUsage

from .._inline_snapshot import snapshot
from ..conftest import IsDatetime, IsFloat, IsNow, IsStr, strip_logfire_metrics, try_import

with try_import() as openai_imports_successful:
    from pydantic_ai.models.openai import OpenAIChatModel
    from pydantic_ai.providers.openai import OpenAIProvider

requires_openai = pytest.mark.skipif(not openai_imports_successful(), reason='openai not installed')

if sys.version_info < (3, 11):
    from exceptiongroup import ExceptionGroup as ExceptionGroup  # pragma: lax no cover
else:
    ExceptionGroup = ExceptionGroup  # pragma: lax no cover

with try_import() as logfire_imports_successful:
    from logfire.testing import CaptureLogfire


pytestmark = pytest.mark.anyio


def success_response(_model_messages: list[ModelMessage], _agent_info: AgentInfo) -> ModelResponse:
    return ModelResponse(parts=[TextPart('success')])


def failure_response(_model_messages: list[ModelMessage], _agent_info: AgentInfo) -> ModelResponse:
    raise ModelHTTPError(status_code=500, model_name='test-function-model', body={'error': 'test error'})


success_model = FunctionModel(success_response)
failure_model = FunctionModel(failure_response)


def test_init() -> None:
    fallback_model = FallbackModel(failure_model, success_model)
    assert fallback_model.model_name == snapshot('fallback:function:failure_response:,function:success_response:')
    assert fallback_model.model_id == snapshot(
        'fallback:function:function:failure_response:,function:function:success_response:'
    )
    assert fallback_model.system == 'fallback:function,function'
    assert fallback_model.base_url is None


def test_all_fields_are_accessible() -> None:
    """Every declared dataclass field must be a real attribute on the instance.

    Regression: `_model_name` was declared as a field but never assigned (`model_name` is a
    computed property), so generic dataclass introspection — e.g. Prefect's `visit_collection`
    during durable execution, which does `getattr(model, f.name)` for each field — crashed with
    `AttributeError`.
    """
    fallback_model = FallbackModel(failure_model, success_model)
    for f in dataclasses.fields(fallback_model):
        getattr(fallback_model, f.name)  # must not raise


def test_first_successful() -> None:
    fallback_model = FallbackModel(success_model, failure_model)
    agent = Agent(model=fallback_model)
    result = agent.run_sync('hello')
    assert result.output == snapshot('success')
    assert result.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[
                    UserPromptPart(content='hello', timestamp=IsNow(tz=timezone.utc)),
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='success')],
                usage=RequestUsage(input_tokens=51, output_tokens=1),
                model_name='function:success_response:',
                timestamp=IsNow(tz=timezone.utc),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )


def test_first_failed() -> None:
    fallback_model = FallbackModel(failure_model, success_model)
    agent = Agent(model=fallback_model)
    result = agent.run_sync('hello')
    assert result.output == snapshot('success')
    assert result.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[
                    UserPromptPart(
                        content='hello',
                        timestamp=IsNow(tz=timezone.utc),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='success')],
                usage=RequestUsage(input_tokens=51, output_tokens=1),
                model_name='function:success_response:',
                timestamp=IsNow(tz=timezone.utc),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )


@pytest.mark.skipif(not logfire_imports_successful(), reason='logfire not installed')
def test_first_failed_instrumented(capfire: CaptureLogfire) -> None:
    fallback_model = FallbackModel(failure_model, success_model)
    agent = Agent(model=fallback_model, capabilities=[Instrumentation(settings=InstrumentationSettings())])
    result = agent.run_sync('hello')
    assert result.output == snapshot('success')
    assert result.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[
                    UserPromptPart(
                        content='hello',
                        timestamp=IsNow(tz=timezone.utc),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='success')],
                usage=RequestUsage(input_tokens=51, output_tokens=1),
                model_name='function:success_response:',
                timestamp=IsNow(tz=timezone.utc),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )
    assert strip_logfire_metrics(capfire.exporter.exported_spans_as_dict(parse_json_attributes=True)) == snapshot(
        [
            {
                'name': 'chat function:success_response:',
                'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
                'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                'start_time': 2000000000,
                'end_time': 3000000000,
                'attributes': {
                    'gen_ai.operation.name': 'chat',
                    'model_request_parameters': {
                        'function_tools': [],
                        'native_tools': [],
                        'output_mode': 'text',
                        'output_object': None,
                        'output_tools': [],
                        'prompted_output_template': None,
                        'allow_text_output': True,
                        'allow_image_output': False,
                        'instruction_parts': None,
                        'thinking': None,
                    },
                    'logfire.span_type': 'span',
                    'gen_ai.conversation.id': IsStr(),
                    'gen_ai.agent.name': 'agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.provider.name': 'function',
                    'logfire.msg': 'chat fallback:function:failure_response:,function:success_response:',
                    'gen_ai.system': 'function',
                    'gen_ai.request.model': 'function:success_response:',
                    'gen_ai.input.messages': [{'role': 'user', 'parts': [{'type': 'text', 'content': 'hello'}]}],
                    'gen_ai.output.messages': [
                        {'role': 'assistant', 'parts': [{'type': 'text', 'content': 'success'}]}
                    ],
                    'gen_ai.usage.input_tokens': 51,
                    'gen_ai.usage.output_tokens': 1,
                    'gen_ai.response.model': 'function:success_response:',
                    'logfire.json_schema': {
                        'type': 'object',
                        'properties': {
                            'gen_ai.input.messages': {'type': 'array'},
                            'gen_ai.output.messages': {'type': 'array'},
                            'model_request_parameters': {'type': 'object'},
                        },
                    },
                },
            },
            {
                'name': 'invoke_agent agent',
                'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                'parent': None,
                'start_time': 1000000000,
                'end_time': 4000000000,
                'attributes': {
                    'model_name': 'fallback:function:failure_response:,function:success_response:',
                    'agent_name': 'agent',
                    'gen_ai.agent.name': 'agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.conversation.id': IsStr(),
                    'gen_ai.operation.name': 'invoke_agent',
                    'logfire.msg': 'agent run',
                    'logfire.span_type': 'span',
                    'gen_ai.aggregated_usage.input_tokens': 51,
                    'gen_ai.aggregated_usage.output_tokens': 1,
                    'pydantic_ai.all_messages': [
                        {'role': 'user', 'parts': [{'type': 'text', 'content': 'hello'}]},
                        {'role': 'assistant', 'parts': [{'type': 'text', 'content': 'success'}]},
                    ],
                    'final_result': 'success',
                    'logfire.json_schema': {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'final_result': {'type': 'object'},
                        },
                    },
                },
            },
        ]
    )


@pytest.mark.skipif(not logfire_imports_successful(), reason='logfire not installed')
async def test_first_failed_instrumented_stream(capfire: CaptureLogfire) -> None:
    fallback_model = FallbackModel(failure_model_stream, success_model_stream)
    agent = Agent(model=fallback_model, capabilities=[Instrumentation(settings=InstrumentationSettings())])
    async with agent.run_stream('input') as result:
        assert [c async for c in result.stream_response(debounce_by=None)] == snapshot(
            [
                ModelResponse(
                    parts=[TextPart(content='hello ')],
                    usage=RequestUsage(input_tokens=50, output_tokens=1),
                    model_name='function::success_response_stream',
                    timestamp=IsNow(tz=timezone.utc),
                    state='incomplete',
                ),
                ModelResponse(
                    parts=[TextPart(content='hello world')],
                    usage=RequestUsage(input_tokens=50, output_tokens=2),
                    model_name='function::success_response_stream',
                    timestamp=IsNow(tz=timezone.utc),
                    state='incomplete',
                ),
                ModelResponse(
                    parts=[TextPart(content='hello world')],
                    usage=RequestUsage(input_tokens=50, output_tokens=2),
                    model_name='function::success_response_stream',
                    timestamp=IsNow(tz=timezone.utc),
                    state='incomplete',
                ),
                ModelResponse(
                    parts=[TextPart(content='hello world')],
                    usage=RequestUsage(input_tokens=50, output_tokens=2),
                    model_name='function::success_response_stream',
                    timestamp=IsDatetime(),
                    run_id=IsStr(),
                    conversation_id=IsStr(),
                    state='complete',
                ),
            ]
        )
        assert result.is_complete

    assert strip_logfire_metrics(capfire.exporter.exported_spans_as_dict(parse_json_attributes=True)) == snapshot(
        [
            {
                'name': 'chat function::success_response_stream',
                'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
                'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                'start_time': 2000000000,
                'end_time': 3000000000,
                'attributes': {
                    'gen_ai.operation.name': 'chat',
                    'model_request_parameters': {
                        'function_tools': [],
                        'native_tools': [],
                        'output_mode': 'text',
                        'output_object': None,
                        'output_tools': [],
                        'prompted_output_template': None,
                        'allow_text_output': True,
                        'allow_image_output': False,
                        'instruction_parts': None,
                        'thinking': None,
                    },
                    'logfire.span_type': 'span',
                    'gen_ai.conversation.id': IsStr(),
                    'gen_ai.agent.name': 'agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.provider.name': 'function',
                    'logfire.msg': 'chat fallback:function::failure_response_stream,function::success_response_stream',
                    'gen_ai.system': 'function',
                    'gen_ai.request.model': 'function::success_response_stream',
                    'gen_ai.input.messages': [{'role': 'user', 'parts': [{'type': 'text', 'content': 'input'}]}],
                    'gen_ai.output.messages': [
                        {'role': 'assistant', 'parts': [{'type': 'text', 'content': 'hello world'}]}
                    ],
                    'gen_ai.usage.input_tokens': 50,
                    'gen_ai.usage.output_tokens': 2,
                    'gen_ai.response.model': 'function::success_response_stream',
                    'gen_ai.client.operation.time_to_first_chunk': IsFloat(),
                    'logfire.json_schema': {
                        'type': 'object',
                        'properties': {
                            'gen_ai.input.messages': {'type': 'array'},
                            'gen_ai.output.messages': {'type': 'array'},
                            'model_request_parameters': {'type': 'object'},
                        },
                    },
                },
            },
            {
                'name': 'invoke_agent agent',
                'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                'parent': None,
                'start_time': 1000000000,
                'end_time': 4000000000,
                'attributes': {
                    'model_name': 'fallback:function::failure_response_stream,function::success_response_stream',
                    'agent_name': 'agent',
                    'gen_ai.agent.name': 'agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.conversation.id': IsStr(),
                    'gen_ai.operation.name': 'invoke_agent',
                    'logfire.msg': 'agent run',
                    'logfire.span_type': 'span',
                    'final_result': 'hello world',
                    'gen_ai.aggregated_usage.input_tokens': 50,
                    'gen_ai.aggregated_usage.output_tokens': 2,
                    'pydantic_ai.all_messages': [
                        {'role': 'user', 'parts': [{'type': 'text', 'content': 'input'}]},
                        {'role': 'assistant', 'parts': [{'type': 'text', 'content': 'hello world'}]},
                    ],
                    'logfire.json_schema': {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'final_result': {'type': 'object'},
                        },
                    },
                },
            },
        ]
    )


def test_all_failed() -> None:
    fallback_model = FallbackModel(failure_model, failure_model)
    agent = Agent(model=fallback_model)
    with pytest.raises(ExceptionGroup) as exc_info:
        agent.run_sync('hello')
    assert 'All models from FallbackModel failed' in exc_info.value.args[0]
    exceptions = exc_info.value.exceptions
    assert len(exceptions) == 2
    assert isinstance(exceptions[0], ModelHTTPError)
    assert exceptions[0].status_code == 500
    assert exceptions[0].model_name == 'test-function-model'
    assert exceptions[0].body == {'error': 'test error'}


def add_missing_response_model(spans: list[dict[str, Any]]) -> list[dict[str, Any]]:
    for span in spans:
        attrs = span.setdefault('attributes', {})
        if 'gen_ai.request.model' in attrs:
            attrs.setdefault('gen_ai.response.model', attrs['gen_ai.request.model'])
    return spans


@pytest.mark.skipif(not logfire_imports_successful(), reason='logfire not installed')
def test_all_failed_instrumented(capfire: CaptureLogfire) -> None:
    fallback_model = FallbackModel(failure_model, failure_model)
    agent = Agent(model=fallback_model, capabilities=[Instrumentation(settings=InstrumentationSettings())])
    with pytest.raises(ExceptionGroup) as exc_info:
        agent.run_sync('hello')
    assert 'All models from FallbackModel failed' in exc_info.value.args[0]
    exceptions = exc_info.value.exceptions
    assert len(exceptions) == 2
    assert isinstance(exceptions[0], ModelHTTPError)
    assert exceptions[0].status_code == 500
    assert exceptions[0].model_name == 'test-function-model'
    assert exceptions[0].body == {'error': 'test error'}
    assert add_missing_response_model(capfire.exporter.exported_spans_as_dict(parse_json_attributes=True)) == snapshot(
        [
            {
                'name': 'chat fallback:function:failure_response:,function:failure_response:',
                'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
                'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                'start_time': 2000000000,
                'end_time': 4000000000,
                'attributes': {
                    'gen_ai.operation.name': 'chat',
                    'gen_ai.provider.name': 'fallback:function,function',
                    'gen_ai.system': 'fallback:function,function',
                    'gen_ai.request.model': 'fallback:function:failure_response:,function:failure_response:',
                    'model_request_parameters': {
                        'function_tools': [],
                        'native_tools': [],
                        'output_mode': 'text',
                        'output_object': None,
                        'output_tools': [],
                        'prompted_output_template': None,
                        'allow_text_output': True,
                        'allow_image_output': False,
                        'instruction_parts': None,
                        'thinking': None,
                    },
                    'logfire.json_schema': {
                        'type': 'object',
                        'properties': {'model_request_parameters': {'type': 'object'}},
                    },
                    'logfire.span_type': 'span',
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.msg': 'chat fallback:function:failure_response:,function:failure_response:',
                    'gen_ai.agent.name': 'agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'logfire.level_num': 17,
                    'gen_ai.response.model': 'fallback:function:failure_response:,function:failure_response:',
                },
                'events': [
                    {
                        'name': 'exception',
                        'timestamp': 3000000000,
                        'attributes': {
                            'exception.type': 'pydantic_ai.exceptions.FallbackExceptionGroup',
                            'exception.message': 'All models from FallbackModel failed (2 sub-exceptions)',
                            'exception.stacktrace': '+------------------------------------',
                            'exception.escaped': 'False',
                        },
                    }
                ],
            },
            {
                'name': 'invoke_agent agent',
                'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                'parent': None,
                'start_time': 1000000000,
                'end_time': 6000000000,
                'attributes': {
                    'model_name': 'fallback:function:failure_response:,function:failure_response:',
                    'agent_name': 'agent',
                    'gen_ai.agent.name': 'agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.conversation.id': IsStr(),
                    'gen_ai.operation.name': 'invoke_agent',
                    'logfire.msg': 'agent run',
                    'logfire.span_type': 'span',
                    'logfire.exception.fingerprint': '0000000000000000000000000000000000000000000000000000000000000000',
                    'pydantic_ai.all_messages': [{'role': 'user', 'parts': [{'type': 'text', 'content': 'hello'}]}],
                    'logfire.json_schema': {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'final_result': {'type': 'object'},
                        },
                    },
                    'logfire.level_num': 17,
                },
                'events': [
                    {
                        'name': 'exception',
                        'timestamp': 5000000000,
                        'attributes': {
                            'exception.type': 'pydantic_ai.exceptions.FallbackExceptionGroup',
                            'exception.message': 'All models from FallbackModel failed (2 sub-exceptions)',
                            'exception.stacktrace': '+------------------------------------',
                            'exception.escaped': 'False',
                        },
                    }
                ],
            },
        ]
    )


async def success_response_stream(_model_messages: list[ModelMessage], _agent_info: AgentInfo) -> AsyncIterator[str]:
    yield 'hello '
    yield 'world'


async def failure_response_stream(_model_messages: list[ModelMessage], _agent_info: AgentInfo) -> AsyncIterator[str]:
    # Note: exception-based fallback for streaming only catches errors during stream initialization
    raise ModelHTTPError(status_code=500, model_name='test-function-model', body={'error': 'test error'})
    yield 'uh oh... '


success_model_stream = FunctionModel(stream_function=success_response_stream)
failure_model_stream = FunctionModel(stream_function=failure_response_stream)


async def test_first_success_streaming() -> None:
    fallback_model = FallbackModel(success_model_stream, failure_model_stream)
    agent = Agent(model=fallback_model)
    async with agent.run_stream('input') as result:
        assert [c async for c in result.stream_response(debounce_by=None)] == snapshot(
            [
                ModelResponse(
                    parts=[TextPart(content='hello ')],
                    usage=RequestUsage(input_tokens=50, output_tokens=1),
                    model_name='function::success_response_stream',
                    timestamp=IsNow(tz=timezone.utc),
                    state='incomplete',
                ),
                ModelResponse(
                    parts=[TextPart(content='hello world')],
                    usage=RequestUsage(input_tokens=50, output_tokens=2),
                    model_name='function::success_response_stream',
                    timestamp=IsNow(tz=timezone.utc),
                    state='incomplete',
                ),
                ModelResponse(
                    parts=[TextPart(content='hello world')],
                    usage=RequestUsage(input_tokens=50, output_tokens=2),
                    model_name='function::success_response_stream',
                    timestamp=IsNow(tz=timezone.utc),
                    state='incomplete',
                ),
                ModelResponse(
                    parts=[TextPart(content='hello world')],
                    usage=RequestUsage(input_tokens=50, output_tokens=2),
                    model_name='function::success_response_stream',
                    timestamp=IsDatetime(),
                    run_id=IsStr(),
                    conversation_id=IsStr(),
                    state='complete',
                ),
            ]
        )
        assert result.is_complete


async def test_first_failed_streaming() -> None:
    fallback_model = FallbackModel(failure_model_stream, success_model_stream)
    agent = Agent(model=fallback_model)
    async with agent.run_stream('input') as result:
        assert [c async for c in result.stream_response(debounce_by=None)] == snapshot(
            [
                ModelResponse(
                    parts=[TextPart(content='hello ')],
                    usage=RequestUsage(input_tokens=50, output_tokens=1),
                    model_name='function::success_response_stream',
                    timestamp=IsNow(tz=timezone.utc),
                    state='incomplete',
                ),
                ModelResponse(
                    parts=[TextPart(content='hello world')],
                    usage=RequestUsage(input_tokens=50, output_tokens=2),
                    model_name='function::success_response_stream',
                    timestamp=IsNow(tz=timezone.utc),
                    state='incomplete',
                ),
                ModelResponse(
                    parts=[TextPart(content='hello world')],
                    usage=RequestUsage(input_tokens=50, output_tokens=2),
                    model_name='function::success_response_stream',
                    timestamp=IsNow(tz=timezone.utc),
                    state='incomplete',
                ),
                ModelResponse(
                    parts=[TextPart(content='hello world')],
                    usage=RequestUsage(input_tokens=50, output_tokens=2),
                    model_name='function::success_response_stream',
                    timestamp=IsDatetime(),
                    run_id=IsStr(),
                    conversation_id=IsStr(),
                    state='complete',
                ),
            ]
        )
        assert result.is_complete


async def test_all_failed_streaming() -> None:
    fallback_model = FallbackModel(failure_model_stream, failure_model_stream)
    agent = Agent(model=fallback_model)
    with pytest.raises(ExceptionGroup) as exc_info:
        async with agent.run_stream('hello') as result:
            [c async for c in result.stream_response(debounce_by=None)]  # pragma: lax no cover
    assert 'All models from FallbackModel failed' in exc_info.value.args[0]
    exceptions = exc_info.value.exceptions
    assert len(exceptions) == 2
    assert isinstance(exceptions[0], ModelHTTPError)
    assert exceptions[0].status_code == 500
    assert exceptions[0].model_name == 'test-function-model'
    assert exceptions[0].body == {'error': 'test error'}


async def test_fallback_condition_override() -> None:
    def should_fallback(exc: Exception) -> bool:
        return False

    fallback_model = FallbackModel(failure_model, success_model, fallback_on=should_fallback)
    agent = Agent(model=fallback_model)
    with pytest.raises(ModelHTTPError):
        await agent.run('hello')


class PotatoException(Exception): ...


def potato_exception_response(_model_messages: list[ModelMessage], _agent_info: AgentInfo) -> ModelResponse:
    raise PotatoException()


async def test_fallback_condition_tuple() -> None:
    potato_model = FunctionModel(potato_exception_response)
    fallback_model = FallbackModel(potato_model, success_model, fallback_on=(PotatoException, ModelHTTPError))
    agent = Agent(model=fallback_model)

    response = await agent.run('hello')
    assert response.output == 'success'


async def test_fallback_connection_error() -> None:
    def connection_error_response(_model_messages: list[ModelMessage], _agent_info: AgentInfo) -> ModelResponse:
        raise ModelAPIError(model_name='test-connection-model', message='Connection timed out')

    connection_error_model = FunctionModel(connection_error_response)
    fallback_model = FallbackModel(connection_error_model, success_model)
    agent = Agent(model=fallback_model)

    response = await agent.run('hello')
    assert response.output == 'success'


async def test_fallback_model_settings_merge():
    """Test that FallbackModel properly merges model settings from wrapped model and runtime settings."""

    def return_settings(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        return ModelResponse(parts=[TextPart(to_json(info.model_settings).decode())])

    base_model = FunctionModel(return_settings, settings=ModelSettings(temperature=0.1, max_tokens=1024))
    fallback_model = FallbackModel(base_model)

    # Test that base model settings are preserved when no additional settings are provided
    agent = Agent(fallback_model)
    result = await agent.run('Hello')
    assert result.output == IsJson({'max_tokens': 1024, 'temperature': 0.1})

    # Test that runtime model_settings are merged with base settings
    agent_with_settings = Agent(fallback_model, model_settings=ModelSettings(temperature=0.5, parallel_tool_calls=True))
    result = await agent_with_settings.run('Hello')
    expected = {'max_tokens': 1024, 'temperature': 0.5, 'parallel_tool_calls': True}
    assert result.output == IsJson(expected)

    # Test that run-time model_settings override both base and agent settings
    result = await agent_with_settings.run(
        'Hello', model_settings=ModelSettings(temperature=0.9, extra_headers={'runtime_setting': 'runtime_value'})
    )
    expected = {
        'max_tokens': 1024,
        'temperature': 0.9,
        'parallel_tool_calls': True,
        'extra_headers': {
            'runtime_setting': 'runtime_value',
        },
    }
    assert result.output == IsJson(expected)


async def test_fallback_model_settings_merge_streaming():
    """Test that FallbackModel properly merges model settings in streaming mode."""

    async def return_settings_stream(_: list[ModelMessage], info: AgentInfo):
        # Yield the merged settings as JSON to verify they were properly combined
        yield to_json(info.model_settings).decode()

    base_model = FunctionModel(
        stream_function=return_settings_stream,
        settings=ModelSettings(temperature=0.1, extra_headers={'anthropic-beta': 'context-1m-2025-08-07'}),
    )
    fallback_model = FallbackModel(base_model)

    # Test that base model settings are preserved in streaming mode
    agent = Agent(fallback_model)
    async with agent.run_stream('Hello') as result:
        output = await result.get_output()

    assert json.loads(output) == {'extra_headers': {'anthropic-beta': 'context-1m-2025-08-07'}, 'temperature': 0.1}

    # Test that runtime model_settings are merged with base settings in streaming mode
    agent_with_settings = Agent(fallback_model, model_settings=ModelSettings(temperature=0.5))
    async with agent_with_settings.run_stream('Hello') as result:
        output = await result.get_output()

    expected = {'extra_headers': {'anthropic-beta': 'context-1m-2025-08-07'}, 'temperature': 0.5}
    assert json.loads(output) == expected


async def test_fallback_thinking_idempotent_across_heterogeneous_models() -> None:
    """`thinking='high'` flows correctly through a FallbackModel whose inner models disagree on thinking support.

    `FallbackModel.request` calls each inner model's `prepare_request` once (for span attributes), then
    `model.request` re-runs it — so `prepare_request` runs twice per inner model. This locks that the double-run is
    idempotent and leaks nothing across runs or across inner models: the reasoning model still sees `thinking='high'`
    lifted into its request parameters, the non-reasoning fallback has it gated out, and the caller's `model_settings`
    is left untouched.
    """
    seen_params: dict[str, ModelRequestParameters] = {}
    seen_settings: dict[str, ModelSettings | None] = {}

    def reasoning(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        seen_params['reasoning'] = info.model_request_parameters
        seen_settings['reasoning'] = info.model_settings
        raise ModelHTTPError(status_code=500, model_name='reasoning', body=None)

    def non_reasoning(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        seen_params['non_reasoning'] = info.model_request_parameters
        seen_settings['non_reasoning'] = info.model_settings
        return ModelResponse(parts=[TextPart('success')])

    reasoning_model = FunctionModel(reasoning, profile=ModelProfile(supports_thinking=True))
    non_reasoning_model = FunctionModel(non_reasoning, profile=ModelProfile(supports_thinking=False))
    fallback_model = FallbackModel(reasoning_model, non_reasoning_model)

    settings = ModelSettings(thinking='high')
    agent = Agent(fallback_model, model_settings=settings)
    result = await agent.run('Hello')

    assert result.output == 'success'
    # Reasoning model: unified `thinking` lifted into request parameters and stripped from `model_settings`.
    assert seen_params['reasoning'].thinking == 'high'
    assert seen_settings['reasoning'] is None
    # Non-reasoning fallback: `thinking` gated out at the profile, never reaching request parameters.
    assert seen_params['non_reasoning'].thinking is None
    assert seen_settings['non_reasoning'] is None
    # The caller's settings object is not mutated by the double `prepare_request` run.
    assert settings == {'thinking': 'high'}


async def test_fallback_model_structured_output():
    class Foo(BaseModel):
        bar: str

    def tool_output_func(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        nonlocal enabled_model
        if enabled_model != 'tool':
            raise ModelHTTPError(status_code=500, model_name='tool-model', body=None)

        assert info.model_request_parameters == snapshot(
            ModelRequestParameters(
                output_mode='tool',
                output_tools=[
                    ToolDefinition(
                        name='final_result',
                        parameters_json_schema={
                            'properties': {'bar': {'type': 'string'}},
                            'required': ['bar'],
                            'title': 'Foo',
                            'type': 'object',
                        },
                        description='The final response which ends this conversation',
                        kind='output',
                        defer_loading=False,
                    )
                ],
                allow_text_output=False,
            )
        )

        args = Foo(bar='baz').model_dump()
        assert info.output_tools
        return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args)])

    def native_output_func(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        nonlocal enabled_model
        if enabled_model != 'native':
            raise ModelHTTPError(status_code=500, model_name='native-model', body=None)

        assert info.model_request_parameters == snapshot(
            ModelRequestParameters(
                output_mode='native',
                output_object=OutputObjectDefinition(
                    json_schema={
                        'properties': {'bar': {'type': 'string'}},
                        'required': ['bar'],
                        'title': 'Foo',
                        'type': 'object',
                    },
                    name='Foo',
                ),
            )
        )

        text = Foo(bar='baz').model_dump_json()
        return ModelResponse(parts=[TextPart(content=text)])

    def prompted_output_func(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        nonlocal enabled_model
        if enabled_model != 'prompted':
            raise ModelHTTPError(status_code=500, model_name='prompted-model', body=None)  # pragma: lax no cover

        assert info.model_request_parameters == snapshot(
            ModelRequestParameters(
                output_mode='prompted',
                output_object=OutputObjectDefinition(
                    json_schema={
                        'properties': {'bar': {'type': 'string'}},
                        'required': ['bar'],
                        'title': 'Foo',
                        'type': 'object',
                    },
                    name='Foo',
                ),
                prompted_output_template="""\

Always respond with a JSON object that's compatible with this schema:

{schema}

Don't include any text or Markdown fencing before or after.
""",
                instruction_parts=[
                    InstructionPart(
                        content="""\

Always respond with a JSON object that's compatible with this schema:

{"properties": {"bar": {"type": "string"}}, "required": ["bar"], "title": "Foo", "type": "object"}

Don't include any text or Markdown fencing before or after.
"""
                    )
                ],
            )
        )

        text = Foo(bar='baz').model_dump_json()
        return ModelResponse(parts=[TextPart(content=text)])

    tool_model = FunctionModel(
        tool_output_func, profile=ModelProfile(default_structured_output_mode='tool', supports_tools=True)
    )
    native_model = FunctionModel(
        native_output_func,
        profile=ModelProfile(default_structured_output_mode='native', supports_json_schema_output=True),
    )
    prompted_model = FunctionModel(
        prompted_output_func, profile=ModelProfile(default_structured_output_mode='prompted')
    )

    fallback_model = FallbackModel(tool_model, native_model, prompted_model)
    agent = Agent(fallback_model, output_type=Foo)

    enabled_model: Literal['tool', 'native', 'prompted'] = 'tool'
    tool_result = await agent.run('hello')
    assert tool_result.output == snapshot(Foo(bar='baz'))

    enabled_model = 'native'
    tool_result = await agent.run('hello')
    assert tool_result.output == snapshot(Foo(bar='baz'))

    enabled_model = 'prompted'
    tool_result = await agent.run('hello')
    assert tool_result.output == snapshot(Foo(bar='baz'))


@pytest.mark.skipif(not logfire_imports_successful(), reason='logfire not installed')
async def test_fallback_model_structured_output_instrumented(capfire: CaptureLogfire) -> None:
    class Foo(BaseModel):
        bar: str

    def tool_output_func(_: list[ModelMessage], _info: AgentInfo) -> ModelResponse:
        raise ModelHTTPError(status_code=500, model_name='tool-model', body=None)

    def prompted_output_func(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        assert info.model_request_parameters == snapshot(
            ModelRequestParameters(
                output_mode='prompted',
                output_object=OutputObjectDefinition(
                    json_schema={
                        'properties': {'bar': {'type': 'string'}},
                        'required': ['bar'],
                        'title': 'Foo',
                        'type': 'object',
                    },
                    name='Foo',
                ),
                prompted_output_template="""\

Always respond with a JSON object that's compatible with this schema:

{schema}

Don't include any text or Markdown fencing before or after.
""",
                instruction_parts=[
                    InstructionPart(content='Be kind'),
                    InstructionPart(
                        content="""\

Always respond with a JSON object that's compatible with this schema:

{"properties": {"bar": {"type": "string"}}, "required": ["bar"], "title": "Foo", "type": "object"}

Don't include any text or Markdown fencing before or after.
"""
                    ),
                ],
            )
        )

        text = Foo(bar='baz').model_dump_json()
        return ModelResponse(parts=[TextPart(content=text)])

    tool_model = FunctionModel(
        tool_output_func, profile=ModelProfile(default_structured_output_mode='tool', supports_tools=True)
    )
    prompted_model = FunctionModel(
        prompted_output_func, profile=ModelProfile(default_structured_output_mode='prompted')
    )
    fallback_model = FallbackModel(tool_model, prompted_model)
    agent = Agent(
        model=fallback_model,
        capabilities=[Instrumentation(settings=InstrumentationSettings())],
        output_type=Foo,
        instructions='Be kind',
    )
    result = await agent.run('hello')
    assert result.output == snapshot(Foo(bar='baz'))
    assert result.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[
                    UserPromptPart(
                        content='hello',
                        timestamp=IsNow(tz=timezone.utc),
                    )
                ],
                timestamp=IsDatetime(),
                instructions='Be kind',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='{"bar":"baz"}')],
                usage=RequestUsage(input_tokens=51, output_tokens=4),
                model_name='function:prompted_output_func:',
                timestamp=IsNow(tz=timezone.utc),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )
    assert strip_logfire_metrics(capfire.exporter.exported_spans_as_dict(parse_json_attributes=True)) == snapshot(
        [
            {
                'name': 'chat function:prompted_output_func:',
                'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
                'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                'start_time': 2000000000,
                'end_time': 3000000000,
                'attributes': {
                    'gen_ai.operation.name': 'chat',
                    'gen_ai.tool.definitions': [
                        {
                            'type': 'function',
                            'name': 'final_result',
                            'description': 'The final response which ends this conversation',
                            'parameters': {
                                'properties': {'bar': {'type': 'string'}},
                                'required': ['bar'],
                                'title': 'Foo',
                                'type': 'object',
                            },
                        }
                    ],
                    'model_request_parameters': {
                        'function_tools': [],
                        'native_tools': [],
                        'output_mode': 'prompted',
                        'output_object': {
                            'json_schema': {
                                'properties': {'bar': {'type': 'string'}},
                                'required': ['bar'],
                                'title': 'Foo',
                                'type': 'object',
                            },
                            'name': 'Foo',
                            'description': None,
                            'strict': None,
                        },
                        'output_tools': [],
                        'prompted_output_template': """\

Always respond with a JSON object that's compatible with this schema:

{schema}

Don't include any text or Markdown fencing before or after.
""",
                        'allow_text_output': True,
                        'allow_image_output': False,
                        'instruction_parts': [
                            {'content': 'Be kind', 'dynamic': False, 'part_kind': 'instruction'},
                            {
                                'content': """\

Always respond with a JSON object that's compatible with this schema:

{"properties": {"bar": {"type": "string"}}, "required": ["bar"], "title": "Foo", "type": "object"}

Don't include any text or Markdown fencing before or after.
""",
                                'dynamic': False,
                                'part_kind': 'instruction',
                            },
                        ],
                        'thinking': None,
                    },
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.span_type': 'span',
                    'gen_ai.agent.name': 'agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.provider.name': 'function',
                    'logfire.msg': 'chat fallback:function:tool_output_func:,function:prompted_output_func:',
                    'gen_ai.system': 'function',
                    'gen_ai.request.model': 'function:prompted_output_func:',
                    'gen_ai.input.messages': [{'role': 'user', 'parts': [{'type': 'text', 'content': 'hello'}]}],
                    'gen_ai.output.messages': [
                        {'role': 'assistant', 'parts': [{'type': 'text', 'content': '{"bar":"baz"}'}]}
                    ],
                    'gen_ai.system_instructions': [{'type': 'text', 'content': 'Be kind'}],
                    'gen_ai.usage.input_tokens': 51,
                    'gen_ai.usage.output_tokens': 4,
                    'gen_ai.response.model': 'function:prompted_output_func:',
                    'logfire.json_schema': {
                        'type': 'object',
                        'properties': {
                            'gen_ai.input.messages': {'type': 'array'},
                            'gen_ai.output.messages': {'type': 'array'},
                            'gen_ai.system_instructions': {'type': 'array'},
                            'model_request_parameters': {'type': 'object'},
                        },
                    },
                },
            },
            {
                'name': 'invoke_agent agent',
                'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                'parent': None,
                'start_time': 1000000000,
                'end_time': 4000000000,
                'attributes': {
                    'model_name': 'fallback:function:tool_output_func:,function:prompted_output_func:',
                    'agent_name': 'agent',
                    'gen_ai.agent.name': 'agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.conversation.id': IsStr(),
                    'gen_ai.operation.name': 'invoke_agent',
                    'logfire.msg': 'agent run',
                    'logfire.span_type': 'span',
                    'gen_ai.aggregated_usage.input_tokens': 51,
                    'gen_ai.aggregated_usage.output_tokens': 4,
                    'pydantic_ai.all_messages': [
                        {'role': 'user', 'parts': [{'type': 'text', 'content': 'hello'}]},
                        {'role': 'assistant', 'parts': [{'type': 'text', 'content': '{"bar":"baz"}'}]},
                    ],
                    'final_result': {'bar': 'baz'},
                    'gen_ai.system_instructions': [{'type': 'text', 'content': 'Be kind'}],
                    'logfire.json_schema': {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'gen_ai.system_instructions': {'type': 'array'},
                            'final_result': {'type': 'object'},
                        },
                    },
                },
            },
        ]
    )


def primary_response(_model_messages: list[ModelMessage], _agent_info: AgentInfo) -> ModelResponse:
    return ModelResponse(parts=[TextPart('primary response')])


def fallback_response(_model_messages: list[ModelMessage], _agent_info: AgentInfo) -> ModelResponse:
    return ModelResponse(parts=[TextPart('fallback response')])


primary_model = FunctionModel(primary_response)
fallback_model_impl = FunctionModel(fallback_response)


async def test_response_handler_triggered() -> None:
    """Test that a response handler can trigger fallback based on response content."""

    def should_fallback_on_primary(response: ModelResponse) -> bool:
        part = response.parts[0] if response.parts else None
        return isinstance(part, TextPart) and 'primary' in part.content

    # Auto-detected as response handler via type hint
    fallback = FallbackModel(
        primary_model,
        fallback_model_impl,
        fallback_on=should_fallback_on_primary,
    )
    agent = Agent(model=fallback)

    result = await agent.run('hello')
    assert result.output == snapshot('fallback response')
    assert result.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[
                    UserPromptPart(content='hello', timestamp=IsNow(tz=timezone.utc)),
                ],
                timestamp=IsNow(tz=timezone.utc),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='fallback response')],
                usage=RequestUsage(input_tokens=51, output_tokens=2),
                model_name='function:fallback_response:',
                timestamp=IsNow(tz=timezone.utc),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )


async def test_response_handler_not_triggered() -> None:
    """Test that response handler returning False allows the response through."""

    def never_fallback(response: ModelResponse) -> bool:
        return False

    # Auto-detected as response handler via type hint
    fallback = FallbackModel(
        primary_model,
        fallback_model_impl,
        fallback_on=never_fallback,
    )
    agent = Agent(model=fallback)

    result = await agent.run('hello')
    assert result.output == snapshot('primary response')


async def test_response_handler_all_fail() -> None:
    """Test that when all models are rejected by response handler, an error is raised."""

    def always_fallback(response: ModelResponse) -> bool:
        return True

    # Auto-detected as response handler via type hint
    fallback = FallbackModel(
        primary_model,
        fallback_model_impl,
        fallback_on=always_fallback,
    )
    agent = Agent(model=fallback)

    with pytest.raises(ExceptionGroup) as exc_info:
        await agent.run('hello')
    assert 'All models from FallbackModel failed' in exc_info.value.args[0]
    assert len(exc_info.value.exceptions) == 1
    assert isinstance(exc_info.value.exceptions[0], ResponseRejected)
    assert 'rejected by fallback_on' in str(exc_info.value.exceptions[0])


async def test_mixed_exception_and_response_handlers() -> None:
    """Test combining exception types and response handlers in a list."""
    call_order: list[str] = []

    def first_fails_with_exception(_: list[ModelMessage], __: AgentInfo) -> ModelResponse:
        call_order.append('first')
        raise ModelHTTPError(status_code=500, model_name='first', body=None)

    def second_fails_response_check(_: list[ModelMessage], __: AgentInfo) -> ModelResponse:
        call_order.append('second')
        return ModelResponse(parts=[TextPart('bad response')])

    def third_succeeds(_: list[ModelMessage], __: AgentInfo) -> ModelResponse:
        call_order.append('third')
        return ModelResponse(parts=[TextPart('good response')])

    def reject_bad_response(response: ModelResponse) -> bool:
        part = response.parts[0] if response.parts else None
        return isinstance(part, TextPart) and 'bad' in part.content

    first_model = FunctionModel(first_fails_with_exception)
    second_model = FunctionModel(second_fails_response_check)
    third_model = FunctionModel(third_succeeds)

    # Use a list to combine exception type and response handler (auto-detected via type hint)
    fallback = FallbackModel(
        first_model,
        second_model,
        third_model,
        fallback_on=[ModelHTTPError, reject_bad_response],
    )
    agent = Agent(model=fallback)

    result = await agent.run('hello')

    assert result.output == snapshot('good response')
    assert call_order == snapshot(['first', 'second', 'third'])


async def test_mixed_failures_all_fail() -> None:
    """Test error reporting when both exceptions and response rejections occur."""
    call_order: list[str] = []

    def first_fails_with_exception(_: list[ModelMessage], __: AgentInfo) -> ModelResponse:
        call_order.append('first')
        raise ModelHTTPError(status_code=500, model_name='first', body=None)

    def second_fails_response_check(_: list[ModelMessage], __: AgentInfo) -> ModelResponse:
        call_order.append('second')
        return ModelResponse(parts=[TextPart('bad response')])

    def reject_bad_response(response: ModelResponse) -> bool:
        part = response.parts[0] if response.parts else None
        return isinstance(part, TextPart) and 'bad' in part.content

    first_model = FunctionModel(first_fails_with_exception)
    second_model = FunctionModel(second_fails_response_check)

    # Auto-detected via type hint
    fallback = FallbackModel(
        first_model,
        second_model,
        fallback_on=[ModelHTTPError, reject_bad_response],
    )
    agent = Agent(model=fallback)

    with pytest.raises(ExceptionGroup) as exc_info:
        await agent.run('hello')

    assert 'All models from FallbackModel failed' in exc_info.value.args[0]
    assert len(exc_info.value.exceptions) == 2
    assert isinstance(exc_info.value.exceptions[0], ModelHTTPError)
    assert isinstance(exc_info.value.exceptions[1], ResponseRejected)
    assert 'rejected by fallback_on' in str(exc_info.value.exceptions[1])
    assert call_order == ['first', 'second']


async def test_web_fetch_scenario() -> None:
    """Test real-world scenario: fallback when web_fetch builtin tool fails.

    This matches the actual Google SDK structure where content is a list of
    UrlMetadata dicts with 'retrieved_url' and 'url_retrieval_status' fields.
    """

    def google_web_fetch_fails(_: list[ModelMessage], __: AgentInfo) -> ModelResponse:
        # Content is a list of UrlMetadata dicts, matching google.genai.types.UrlMetadata.model_dump()
        # Include multiple items to cover loop iteration branch
        return ModelResponse(
            parts=[
                NativeToolCallPart(tool_name='web_fetch', args={'urls': ['https://example.com']}, tool_call_id='1'),
                NativeToolReturnPart(
                    tool_name='web_fetch',
                    tool_call_id='1',
                    content=[
                        {'retrieved_url': 'https://ok.com', 'url_retrieval_status': 'URL_RETRIEVAL_STATUS_SUCCESS'},
                        {'retrieved_url': 'https://example.com', 'url_retrieval_status': 'URL_RETRIEVAL_STATUS_FAILED'},
                    ],
                ),
                TextPart('Could not fetch URL'),
            ]
        )

    def anthropic_succeeds(_: list[ModelMessage], __: AgentInfo) -> ModelResponse:
        return ModelResponse(parts=[TextPart('Successfully fetched and summarized the content')])

    class UrlMetadataDict(TypedDict):
        retrieved_url: str
        url_retrieval_status: str

    def web_fetch_failed(response: ModelResponse) -> bool:
        for call, result in response.native_tool_calls:  # pragma: no branch
            if call.tool_name != 'web_fetch':
                continue  # pragma: lax no cover
            if not isinstance(result.content, list):
                continue  # pragma: lax no cover
            # Cast needed because result.content is typed as Any
            items = cast(list[UrlMetadataDict], result.content)  # pyright: ignore[reportUnknownMemberType]
            for item in items:  # pragma: no branch
                if item['url_retrieval_status'] != 'URL_RETRIEVAL_STATUS_SUCCESS':
                    return True
        return False

    google_model = FunctionModel(google_web_fetch_fails)
    anthropic_model = FunctionModel(anthropic_succeeds)

    # Auto-detected via type hint
    fallback = FallbackModel(
        google_model,
        anthropic_model,
        fallback_on=web_fetch_failed,
    )
    agent = Agent(model=fallback)

    result = await agent.run('Summarize https://example.com')
    assert result.output == 'Successfully fetched and summarized the content'


def test_response_handler_sync() -> None:
    """Test response handler with synchronous run."""

    def should_fallback(response: ModelResponse) -> bool:
        part = response.parts[0] if response.parts else None
        return isinstance(part, TextPart) and 'primary' in part.content

    # Auto-detected via type hint
    fallback = FallbackModel(
        primary_model,
        fallback_model_impl,
        fallback_on=should_fallback,
    )
    agent = Agent(model=fallback)

    result = agent.run_sync('hello')
    assert result.output == 'fallback response'


def test_fallback_on_list_of_exception_types() -> None:
    """Test fallback_on with a list containing individual exception types."""

    class CustomError(Exception):
        pass

    def raises_custom_error(_: list[ModelMessage], __: AgentInfo) -> ModelResponse:
        raise CustomError('custom error')

    custom_error_model = FunctionModel(raises_custom_error)

    # List with individual exception types (not a tuple)
    fallback = FallbackModel(
        custom_error_model,
        success_model,
        fallback_on=[CustomError, ModelHTTPError],
    )
    agent = Agent(model=fallback)

    result = agent.run_sync('hello')
    assert result.output == 'success'


def test_fallback_on_single_response_handler() -> None:
    """Test fallback_on with a single response handler (auto-detected via type hint)."""

    def reject_primary(response: ModelResponse) -> bool:
        part = response.parts[0] if response.parts else None
        return isinstance(part, TextPart) and 'primary' in part.content

    # Auto-detected as response handler via type hint
    fallback = FallbackModel(
        primary_model,
        fallback_model_impl,
        fallback_on=reject_primary,
    )
    agent = Agent(model=fallback)

    result = agent.run_sync('hello')
    assert result.output == 'fallback response'


def test_fallback_on_single_exception_handler() -> None:
    """Test fallback_on with a single exception handler (auto-detected by type hint)."""

    def custom_exception_handler(exc: Exception) -> bool:
        return isinstance(exc, ModelHTTPError) and exc.status_code == 500

    # Auto-detected as exception handler via type hint (first param is Exception, not ModelResponse)
    fallback = FallbackModel(
        failure_model,
        success_model,
        fallback_on=custom_exception_handler,
    )
    agent = Agent(model=fallback)

    result = agent.run_sync('hello')
    assert result.output == 'success'


def test_fallback_on_mixed_list() -> None:
    """Test fallback_on with a mixed list of exception types, exception handlers, and response handlers."""

    class CustomError(Exception):
        pass

    def custom_exception_handler(exc: Exception) -> bool:  # pragma: no cover
        return isinstance(exc, ModelHTTPError) and exc.status_code == 503

    def reject_bad_response(response: ModelResponse) -> bool:
        part = response.parts[0] if response.parts else None
        return isinstance(part, TextPart) and 'bad' in part.content

    def bad_response_model(_: list[ModelMessage], __: AgentInfo) -> ModelResponse:
        return ModelResponse(parts=[TextPart('bad response')])

    bad_model = FunctionModel(bad_response_model)

    # Mix of exception type, exception handler, and response handler (auto-detected via type hints)
    fallback = FallbackModel(
        bad_model,
        fallback_model_impl,
        fallback_on=[CustomError, custom_exception_handler, reject_bad_response],
    )
    agent = Agent(model=fallback)

    # Should fallback because response contains 'bad'
    result = agent.run_sync('hello')
    assert result.output == 'fallback response'


def test_fallback_on_lambda_exception_handler() -> None:
    """Test that lambdas with 1 param are detected as exception handlers."""
    fallback = FallbackModel(
        failure_model,
        success_model,
        fallback_on=lambda e: isinstance(e, ModelHTTPError),
    )
    agent = Agent(model=fallback)

    result = agent.run_sync('hello')
    assert result.output == 'success'


async def test_async_exception_handler() -> None:
    """Test that async exception handlers work correctly."""

    async def async_exc_handler(exc: Exception) -> bool:
        return isinstance(exc, ModelHTTPError)

    fallback = FallbackModel(
        failure_model,
        success_model,
        fallback_on=async_exc_handler,
    )
    agent = Agent(model=fallback)

    result = await agent.run('hello')
    assert result.output == 'success'


async def test_async_response_handler() -> None:
    """Test that async response handlers work correctly."""

    async def async_response_handler(response: ModelResponse) -> bool:
        # Reject if 'primary' in response
        part = response.parts[0] if response.parts else None
        return isinstance(part, TextPart) and 'primary' in part.content

    fallback = FallbackModel(
        primary_model,
        fallback_model_impl,
        fallback_on=async_response_handler,
    )
    agent = Agent(model=fallback)

    result = await agent.run('hello')
    assert result.output == 'fallback response'


def test_fallback_on_invalid_type() -> None:
    """Test that invalid fallback_on types raise AssertionError via assert_never."""
    with pytest.raises(AssertionError, match='Expected code to be unreachable'):
        FallbackModel(success_model, failure_model, fallback_on='invalid')  # type: ignore


def test_fallback_on_invalid_list_item() -> None:
    """Test that invalid items in fallback_on list raise AssertionError via assert_never."""
    with pytest.raises(AssertionError, match='Expected code to be unreachable'):
        FallbackModel(success_model, failure_model, fallback_on=['invalid'])  # type: ignore


def test_response_handler_only_exception_propagates() -> None:
    """Test that exceptions propagate when only response handlers are configured.

    This documents the expected behavior: if you only configure response handlers
    (no exception types or exception handlers), exceptions are not caught and will
    propagate to the caller.
    """

    def response_check(response: ModelResponse) -> bool:  # pragma: no cover
        return False  # Never reject based on response

    # Auto-detected as response handler via type hint - only a response handler, no exception handling
    fallback = FallbackModel(
        failure_model,  # This will raise ModelHTTPError
        success_model,
        fallback_on=response_check,
    )
    agent = Agent(model=fallback)

    # Exception should propagate since no exception handler is configured
    with pytest.raises(ModelHTTPError):
        agent.run_sync('hello')


def test_callable_class_response_handler() -> None:
    """Test that callable classes with __call__(ModelResponse) trigger response-based fallback."""

    class RejectPrimary:
        def __call__(self, response: ModelResponse) -> bool:
            part = response.parts[0] if response.parts else None
            return isinstance(part, TextPart) and 'primary' in part.content

    fallback = FallbackModel(
        primary_model,
        fallback_model_impl,
        fallback_on=RejectPrimary(),
    )
    agent = Agent(model=fallback)
    result = agent.run_sync('hello')
    assert result.output == 'fallback response'


def test_callable_class_exception_handler() -> None:
    """Test that callable classes with __call__(Exception) trigger exception-based fallback."""

    class HandleHTTPError:
        def __call__(self, exc: Exception) -> bool:
            return isinstance(exc, ModelHTTPError)

    fallback = FallbackModel(
        failure_model,
        success_model,
        fallback_on=HandleHTTPError(),
    )
    agent = Agent(model=fallback)
    result = agent.run_sync('hello')
    assert result.output == 'success'


def test_unresolvable_forward_ref_treated_as_exception_handler() -> None:
    """A handler with unresolvable forward refs is treated as an exception handler."""
    # Create a function whose type hints can't be resolved (triggers except branch in get_first_param_type)
    exec_globals: dict[str, object] = {}
    exec(  # nosec - test-only dynamic function creation for unresolvable forward ref
        """
def handler(x: "NonExistentType") -> bool:
    return isinstance(x, Exception)
""",
        exec_globals,
    )
    handler = exec_globals['handler']

    # Classified as exception handler (forward ref can't resolve), so responses pass through
    fallback = FallbackModel(
        primary_model,
        fallback_model_impl,
        fallback_on=handler,  # type: ignore[arg-type]
    )
    agent = Agent(model=fallback)
    result = agent.run_sync('hello')
    assert result.output == 'primary response'


def test_fallback_on_single_exception_type_direct() -> None:
    """Test fallback_on with a single exception type (not in tuple/list)."""

    def raises_api_error(_: list[ModelMessage], __: AgentInfo) -> ModelResponse:
        raise ModelAPIError('test-model', 'test error')

    fallback = FallbackModel(
        FunctionModel(raises_api_error),
        success_model,
        fallback_on=ModelAPIError,  # Single type, not tuple
    )
    agent = Agent(model=fallback)
    result = agent.run_sync('hello')
    assert result.output == 'success'


def test_empty_fallback_on_list_error() -> None:
    """Test that empty fallback_on list raises UserError."""
    from pydantic_ai.exceptions import UserError

    with pytest.raises(UserError, match='empty fallback_on'):
        FallbackModel(
            primary_model,
            fallback_model_impl,
            fallback_on=[],
        )


def test_empty_fallback_on_tuple_error() -> None:
    """Test that empty fallback_on tuple raises UserError."""
    from pydantic_ai.exceptions import UserError

    with pytest.raises(UserError, match='empty fallback_on'):
        FallbackModel(
            primary_model,
            fallback_model_impl,
            fallback_on=(),
        )


async def test_response_rejection_error_message() -> None:
    """Test that error message describes response rejections."""

    def always_reject(response: ModelResponse) -> bool:
        return True

    fallback = FallbackModel(
        primary_model,
        fallback_model_impl,
        fallback_on=always_reject,
    )
    agent = Agent(model=fallback)

    with pytest.raises(ExceptionGroup) as exc_info:
        await agent.run('hello')

    # Find the ResponseRejected in the exception group
    rejection_errors = [e for e in exc_info.value.exceptions if isinstance(e, ResponseRejected)]
    assert len(rejection_errors) == 1

    error_msg = str(rejection_errors[0])
    assert 'rejected by fallback_on handler' in error_msg


@requires_openai
async def test_fallback_model_lifecycle_closes_sub_model_clients():
    """FallbackModel propagates __aenter__/__aexit__ to all sub-models' providers.

    Regression test for PR #4421 (provider lifecycle management).
    https://github.com/pydantic/pydantic-ai/pull/4421
    """
    provider1 = OpenAIProvider(api_key='test-key-1')
    provider2 = OpenAIProvider(api_key='test-key-2')
    model1 = OpenAIChatModel('gpt-4o', provider=provider1)
    model2 = OpenAIChatModel('gpt-4o', provider=provider2)

    fallback = FallbackModel(model1, model2)

    async with fallback:
        assert provider1._own_http_client is not None  # pyright: ignore[reportPrivateUsage]
        assert provider2._own_http_client is not None  # pyright: ignore[reportPrivateUsage]
        assert not provider1._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]
        assert not provider2._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]
    assert provider1._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]
    assert provider2._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]


@requires_openai
async def test_fallback_model_lifecycle_via_agent():
    """Agent context manager propagates lifecycle through FallbackModel to sub-models' providers.

    Regression test for PR #4421 (provider lifecycle management).
    https://github.com/pydantic/pydantic-ai/pull/4421
    """
    provider1 = OpenAIProvider(api_key='test-key-1')
    provider2 = OpenAIProvider(api_key='test-key-2')
    model1 = OpenAIChatModel('gpt-4o', provider=provider1)
    model2 = OpenAIChatModel('gpt-4o', provider=provider2)

    fallback = FallbackModel(model1, model2)
    agent = Agent(model=fallback)

    async with agent:
        assert provider1._own_http_client is not None  # pyright: ignore[reportPrivateUsage]
        assert provider2._own_http_client is not None  # pyright: ignore[reportPrivateUsage]
        assert not provider1._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]
        assert not provider2._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]
    assert provider1._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]
    assert provider2._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]


@requires_openai
async def test_fallback_model_reentrant_lifecycle():
    """Reentrant FallbackModel lifecycle keeps sub-models' clients open until outermost exit.

    Regression test for PR #4421 (provider lifecycle management).
    https://github.com/pydantic/pydantic-ai/pull/4421
    """
    provider1 = OpenAIProvider(api_key='test-key-1')
    provider2 = OpenAIProvider(api_key='test-key-2')
    model1 = OpenAIChatModel('gpt-4o', provider=provider1)
    model2 = OpenAIChatModel('gpt-4o', provider=provider2)

    fallback = FallbackModel(model1, model2)

    async with fallback:
        http1 = provider1._own_http_client  # pyright: ignore[reportPrivateUsage]
        http2 = provider2._own_http_client  # pyright: ignore[reportPrivateUsage]
        assert http1 is not None
        assert http2 is not None
        async with fallback:
            assert not http1.is_closed
            assert not http2.is_closed
        assert not http1.is_closed
        assert not http2.is_closed
    assert http1.is_closed
    assert http2.is_closed


@requires_openai
async def test_fallback_model_instrumented_lifecycle():
    """InstrumentedModel wrapping FallbackModel propagates lifecycle to sub-models.

    Regression test for PR #4421 (provider lifecycle management).
    https://github.com/pydantic/pydantic-ai/pull/4421
    """
    provider1 = OpenAIProvider(api_key='test-key-1')
    provider2 = OpenAIProvider(api_key='test-key-2')
    model1 = OpenAIChatModel('gpt-4o', provider=provider1)
    model2 = OpenAIChatModel('gpt-4o', provider=provider2)

    fallback = FallbackModel(model1, model2)
    instrumented = InstrumentedModel(fallback, InstrumentationSettings())

    async with instrumented:
        assert provider1._own_http_client is not None  # pyright: ignore[reportPrivateUsage]
        assert provider2._own_http_client is not None  # pyright: ignore[reportPrivateUsage]
        assert not provider1._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]
        assert not provider2._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]
    assert provider1._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]
    assert provider2._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]


@requires_openai
async def test_fallback_model_concurrent_entry():
    """Concurrent entry to FallbackModel doesn't race on _entered_count / _exit_stack.

    Without a lock, two coroutines can both see _entered_count == 0 when the first
    yields during sub-model entry, causing one exit stack to be overwritten and leaked.

    Regression test for PR #4421 (provider lifecycle management).
    https://github.com/pydantic/pydantic-ai/pull/4421
    """
    import asyncio

    from pydantic_ai.models.wrapper import WrapperModel

    class SlowEnterModel(WrapperModel):
        """Wrapper that yields during __aenter__ to widen the race window."""

        async def __aenter__(self) -> SlowEnterModel:
            await asyncio.sleep(0)
            await self.wrapped.__aenter__()
            return self

    provider1 = OpenAIProvider(api_key='test-key-1')
    provider2 = OpenAIProvider(api_key='test-key-2')
    model1 = SlowEnterModel(OpenAIChatModel('gpt-4o', provider=provider1))
    model2 = SlowEnterModel(OpenAIChatModel('gpt-4o', provider=provider2))

    fallback = FallbackModel(model1, model2)

    async def enter_and_hold(event: asyncio.Event) -> None:
        async with fallback:
            event.set()
            await asyncio.sleep(0.1)

    event1 = asyncio.Event()
    event2 = asyncio.Event()
    task1 = asyncio.create_task(enter_and_hold(event1))
    task2 = asyncio.create_task(enter_and_hold(event2))

    await event1.wait()
    await event2.wait()
    assert provider1._own_http_client is not None  # pyright: ignore[reportPrivateUsage]
    assert provider2._own_http_client is not None  # pyright: ignore[reportPrivateUsage]
    assert not provider1._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]
    assert not provider2._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]

    await task1
    await task2
    assert provider1._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]
    assert provider2._own_http_client.is_closed  # pyright: ignore[reportPrivateUsage]