pydantic-ai/tests/conftest.py

from __future__ import annotations as _annotations

import asyncio
import dataclasses
import gc
import importlib.util
import logging
import os
import re
import secrets
import sys
import warnings
from collections.abc import AsyncIterator, Callable, Generator, Iterator, Sequence
from contextlib import contextmanager, suppress
from dataclasses import dataclass, field
from datetime import datetime
from decimal import Decimal
from functools import cache
from pathlib import Path
from types import ModuleType
from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, cast, overload

import httpx
import httpx2
import pytest
from _pytest.assertion.rewrite import AssertionRewritingHook
from cassetter import Cassette, RecordMode
from pydantic import JsonValue, TypeAdapter
from pytest_mock import MockerFixture

import pydantic_ai._http
import pydantic_ai.models
from pydantic_ai import Agent, BinaryContent, BinaryImage, Embedder, ImageGenerator
from pydantic_ai._cache_health import ConversationCacheMarkStore
from pydantic_ai.capabilities import instrumentation as instrumentation_capability
from pydantic_ai.messages import (
    DocumentUrl,
    FilePart,
    ImageUrl,
    ModelMessage,
    ModelRequest,
    ModelResponse,
    NativeToolCallPart,
    NativeToolReturnPart,
    RetryPromptPart,
    SystemPromptPart,
    TextPart,
    ThinkingPart,
    ToolCallPart,
    ToolReturnPart,
    UserPromptPart,
    VideoUrl,
)
from pydantic_ai.models import Model
from pydantic_ai.usage import RequestUsage, RunUsage

from . import cassette_hooks, cost_guards
from ._inline_snapshot import Builder, Custom, customize
from .cassette_utils import check_cache_prefix_stability

# `logfire` builds its JSON schema lookup table on first use with a lazy `import pandas`
# (`logfire/_internal/json_schema.py`), and importing pandas reads timezone data from disk. When
# that first import happens inside an async test with the detector configured below armed, the read
# is reported as blocking and aborts the import, leaving a partially initialized `pandas` in
# `sys.modules` that every later import in the session trips over. Importing it here means the
# detector never sees a first import. The read genuinely blocks, so this stays regardless of what
# BlockBuster reports: https://github.com/pydantic/pydantic-ai/issues/7446.
with suppress(ImportError):
    import pandas  # pyright: ignore[reportUnusedImport] # noqa: F401

T = TypeVar('T')

# Like `try_import` for a whole directory: these suites need their workspace member installed, which
# the `--package` matrix cells for other members do not do.
collect_ignore = [
    directory
    for directory, package in (('harness', 'pydantic_ai_harness'), ('clai2', 'pydantic_clai2'))
    if importlib.util.find_spec(package) is None
]

__all__ = (
    'IsDatetime',
    'IsDecimal',
    'IsFloat',
    'IsNow',
    'IsStr',
    'IsBytes',
    'IsInt',
    'IsInstance',
    'IsList',
    'TestEnv',
    'try_import',
    'SNAPSHOT_BYTES_COLLAPSE_THRESHOLD',
    'strip_logfire_metrics',
    'remove_schema_descriptions',
    'message',
    'message_part',
)

pydantic_ai.models.ALLOW_MODEL_REQUESTS = False

os.environ.setdefault('HF_HUB_DISABLE_PROGRESS_BARS', '1')


def pytest_configure(config: pytest.Config) -> None:
    config.pluginmanager.register(cost_guards, cost_guards.PLUGIN_NAME)
    config.addinivalue_line(
        'markers',
        'moves_cache_prefix(reason): recorded conversation deliberately moves the cache prefix; reason required',
    )
    config.addinivalue_line(
        'markers',
        'realtime_ws_hold_open: keep a replay WebSocket open after its last recorded frame',
    )


if TYPE_CHECKING:
    from blockbuster import BlockBuster
    from pluggy import Result

    from pydantic_ai.providers.bedrock import BedrockProvider
    from pydantic_ai.providers.xai import XaiProvider

    def IsInstance(arg: type[T]) -> T: ...
    def IsDatetime(*args: Any, **kwargs: Any) -> datetime: ...
    def IsDecimal(*args: Any, **kwargs: Any) -> Decimal: ...
    def IsFloat(*args: Any, **kwargs: Any) -> float: ...
    def IsInt(*args: Any, **kwargs: Any) -> int: ...
    def IsNow(*args: Any, **kwargs: Any) -> datetime: ...
    def IsStr(*args: Any, **kwargs: Any) -> str: ...
    def IsSameStr(*args: Any, **kwargs: Any) -> str: ...
    def IsBytes(*args: Any, **kwargs: Any) -> bytes: ...
    def IsList(*args: T, **kwargs: Any) -> list[T]: ...
else:
    from dirty_equals import IsBytes, IsDatetime, IsFloat, IsInstance, IsInt, IsList, IsNow as _IsNow, IsNumeric, IsStr

    class IsDecimal(IsNumeric[Decimal]):
        allowed_types = Decimal

    def IsNow(*args: Any, **kwargs: Any):
        # Increase the default value of `delta` to 10 to reduce test flakiness on overburdened machines
        if 'delta' not in kwargs:  # pragma: no branch
            kwargs['delta'] = 10
        return _IsNow(*args, **kwargs)

    class IsSameStr(IsStr):
        """
        Checks if the value is a string, and that subsequent uses have the same value as the first one.

        Example:
        ```python {test="skip"}
        assert events == [
            {
                'type': 'RUN_STARTED',
                'threadId': (thread_id := IsSameStr()),
                'runId': (run_id := IsSameStr()),
            },
            {'type': 'TEXT_MESSAGE_START', 'messageId': (message_id := IsSameStr()), 'role': 'assistant'},
            {'type': 'TEXT_MESSAGE_CONTENT', 'messageId': message_id, 'delta': 'success '},
            {
                'type': 'TEXT_MESSAGE_CONTENT',
                'messageId': message_id,
                'delta': '(no tool calls)',
            },
            {'type': 'TEXT_MESSAGE_END', 'messageId': message_id},
            {
                'type': 'RUN_FINISHED',
                'threadId': thread_id,
                'runId': run_id,
            },
        ]
        ```
        """

        _first_other: str | None = None

        def equals(self, other: Any) -> bool:
            if self._first_other is None:
                self._first_other = other
                return super().equals(other)
            else:
                return other == self._first_other


JsonSchemaValue: TypeAlias = 'dict[str, JsonSchemaValue] | list[JsonSchemaValue] | str | int | float | bool | None'


def remove_schema_descriptions(schema: JsonSchemaValue) -> JsonSchemaValue:
    """Recursively drop docstring-derived `description` keys from a JSON schema for version-tolerant comparison.

    Pydantic 2.13 emits class and stdlib-dataclass docstrings as schema `description`s, where 2.12
    suppressed stdlib-dataclass docstrings (see pydantic#12812). Dropping descriptions keeps these
    schema snapshots stable across the supported pydantic range while still asserting structure.

    Only string-valued `description` keys (the docstring/annotation form) are dropped; a `description`
    key whose value is a sub-schema is a property literally named `description` (e.g. a capability's
    `description` field) and must be preserved.
    """
    if isinstance(schema, dict):
        return {
            k: remove_schema_descriptions(v)
            for k, v in schema.items()
            if not (k == 'description' and isinstance(v, str))
        }
    if isinstance(schema, list):
        return [remove_schema_descriptions(v) for v in schema]
    return schema


def strip_logfire_metrics(spans: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Drop the version-dependent `logfire.metrics` span attribute before snapshotting.

    logfire 4.3x+ attaches the aggregated `gen_ai.client.token.usage` metric to spans; older
    versions don't. Stripping it keeps span snapshots stable across the supported logfire range
    (the token usage itself is still asserted via the `gen_ai.usage.*` attributes / collected metrics).
    """
    for span in spans:
        span.get('attributes', {}).pop('logfire.metrics', None)
    return spans


SNAPSHOT_BYTES_COLLAPSE_THRESHOLD = 50


def sanitize_filename(name: str, max_len: int) -> str:
    """Sanitize a string for safe use as a filename across platforms."""
    # Windows does not allow these characters in paths. Linux bans slashes only.
    return re.sub('[' + re.escape('<>:"/\\|?*') + ']', '-', name)[:max_len]


@customize
def binary_handler(value: Any) -> Any | None:  # pragma: no cover
    # Use IsBytes() for large byte sequences in snapshots.
    if isinstance(value, bytes) and len(value) > SNAPSHOT_BYTES_COLLAPSE_THRESHOLD:
        return IsBytes()


@customize
def isdatetime_handler(value: Any, builder: Builder) -> Any | None:  # pragma: no cover
    # Use IsDatetime() for datetime values in snapshots.
    if isinstance(value, datetime):
        return IsDatetime()


@customize
def usage_handler(value: Any, builder: Builder) -> Custom | None:  # pragma: no cover
    if isinstance(value, (RequestUsage, RunUsage)):
        # Usage objects accept arbitrary fields that inline-snapshot's default dataclass handler cannot see.
        kwargs = value.__dict__.copy()
        for field in dataclasses.fields(value):
            if field.name not in kwargs:
                continue
            if field.default is not dataclasses.MISSING:
                kwargs[field.name] = builder.with_default(kwargs[field.name], field.default)
            elif field.default_factory is not dataclasses.MISSING:
                kwargs[field.name] = builder.with_default(kwargs[field.name], field.default_factory())
        return builder.create_call(type(value), [], kwargs)


@customize
def content_handler(value: Any, builder: Builder) -> Custom | None:  # pragma: no cover
    # special handler for types which need an identifier argument for __init__ but declare an _identifier in the class
    if isinstance(value, BinaryImage):
        return builder.create_call(
            BinaryImage,
            [],
            {
                # prevent generation of IsBytes() because it does not work together with Pydantic models
                'data': builder.create_code(f'{value.data!r}'),
                'media_type': value.media_type,
                'identifier': builder.with_default(value.identifier, None),
                'vendor_metadata': builder.with_default(value.vendor_metadata, None),
                # kind is always "binary"
                # 'kind': builder.with_default(value.kind, 'binary'),
            },
        )

    if isinstance(value, BinaryContent):
        return builder.create_call(
            BinaryContent,
            [],
            {
                # prevent generation of IsBytes() because it does not work together with Pydantic models
                'data': builder.create_code(f'{value.data!r}'),
                'media_type': value.media_type,
                'identifier': builder.with_default(value.identifier, None),
                'vendor_metadata': builder.with_default(value.vendor_metadata, None),
                'kind': builder.with_default(value.kind, 'binary'),
            },
        )

    for cls, kind in [(VideoUrl, 'video-url'), (DocumentUrl, 'document-url'), (ImageUrl, 'image-url')]:
        if type(value) is cls:
            return builder.create_call(
                cls,
                [],
                {
                    'url': value.url,
                    'media_type': builder.with_default(value.media_type, None),
                    # TODO: identifier is not used for == comparison should we ignore it?
                    'identifier': builder.with_default(value.identifier, None),
                    'force_download': builder.with_default(value.force_download, False),
                    'vendor_metadata': builder.with_default(value.vendor_metadata, None),
                    'kind': builder.with_default(value.kind, kind),
                },
            )


@customize
def variable_handler(value: Any, builder: Builder, local_vars: dict[str, Any]) -> Custom | None:  # pragma: no cover
    for name, local_variable in local_vars.items():
        # use local_function.__qualname__ when there exist a local_function with the wanted name
        if hasattr(local_variable, '__qualname__') and value == local_variable.__qualname__:
            return builder.create_code(f'{name}.__qualname__')

        # use `part.tool_call_id` when there is a local variable part with the wanted id
        if name == 'part' and hasattr(local_variable, 'tool_call_id') and local_variable.tool_call_id == value:
            return builder.create_code(f'{name}.tool_call_id')

        # skip IsSameStr variables that haven't been compared yet (no value captured)
        if hasattr(local_variable, '_first_other') and local_variable._first_other is None:
            continue

        # uses local variables like part* *_content or thread_id when their value is equal to the wanted value in the snapshot
        if (
            (name.startswith('part') or name.endswith('_content') or name in ('thread_id',))
            and name != 'parts'
            and local_variable == value
        ):
            return builder.create_code(name)

        # match the local var like `var := IsSameStr()` a second time
        if type(local_variable) is IsSameStr and local_variable == value:
            return builder.create_code(name)


class TestEnv:
    __test__ = False

    def __init__(self):
        self.envars: dict[str, str | None] = {}

    def set(self, name: str, value: str) -> None:
        self.envars[name] = os.getenv(name)
        os.environ[name] = value

    def remove(self, name: str) -> None:
        self.envars[name] = os.environ.pop(name, None)

    def reset(self) -> None:
        for name, value in self.envars.items():
            if value is None:
                os.environ.pop(name, None)
            else:
                os.environ[name] = value  # pragma: lax no cover


@pytest.fixture
def env() -> Iterator[TestEnv]:
    test_env = TestEnv()

    yield test_env

    test_env.reset()


@pytest.fixture(scope='session')
def anyio_backend(pytestconfig: pytest.Config) -> str:
    backend = pytestconfig.getoption('--anyio-backend')
    assert isinstance(backend, str)
    return backend


# Calls that are allowed to block in the event loop, as (blockbuster function, file, functions).
# Each entry should say why the blocking call is acceptable; anything not listed here should be
# fixed (e.g. offloaded to a thread with `anyio.to_thread.run_sync`) rather than exempted.
BLOCKBUSTER_EXEMPTIONS: list[tuple[str, str, str | tuple[str, ...]]] = [
    # pytest-examples locates the source line of a captured `print()` with `Path.samefile`, so an
    # example printing from inside a running event loop trips the detector on the harness's own
    # `os.stat`. Exempting the capture entry point keeps `os.stat` calls from example and library
    # code detectable.
    ('os.stat', 'pytest_examples/run_code.py', '__call__'),
    # The first-run banner asks whether the harness is installed and what version it is, to name it
    # in the banner. That happens once per process, before the first request.
    ('os.stat', 'pydantic_ai/_display.py', '_version_line'),
    ('os.listdir', 'pydantic_ai/_display.py', '_version_line'),
    ('io.TextIOWrapper.read', 'pydantic_ai/_display.py', '_version_line'),
    ('io.BufferedReader.read', 'pydantic_ai/_display.py', '_version_line'),
    # `load_mcp_toolsets` is a sync config-file loader; reading the file is its documented job.
    ('os.stat', 'pydantic_ai/mcp.py', 'load_mcp_toolsets'),
    ('io.BufferedReader.read', 'pydantic_ai/mcp.py', 'load_mcp_toolsets'),
    # fastmcp's transport inference stats candidate paths when the client is constructed.
    ('os.stat', 'pydantic_ai/mcp.py', '__init__'),
    # boto3/botocore read AWS config files and service models during sync, setup-time client construction.
    ('os.stat', 'pydantic_ai/providers/bedrock.py', '__init__'),
    ('os.listdir', 'pydantic_ai/providers/bedrock.py', '__init__'),
    ('io.TextIOWrapper.read', 'pydantic_ai/providers/bedrock.py', '__init__'),
    ('io.BufferedReader.read', 'pydantic_ai/providers/bedrock.py', '__init__'),
    # google-genai's sync client constructor has google-auth read mTLS/credential config files.
    ('os.stat', 'pydantic_ai/providers/google_cloud.py', '__init__'),
    ('io.TextIOWrapper.read', 'pydantic_ai/providers/google_cloud.py', '__init__'),
    ('io.BufferedReader.read', 'pydantic_ai/providers/google_cloud.py', '__init__'),
    # Anthropic's async Bedrock clients synchronously resolve/sign AWS credentials
    # (https://github.com/anthropics/anthropic-sdk-python/issues/1770); remove after upstream
    # covers both auth paths.
    ('os.stat', 'anthropic/lib/bedrock/_auth.py', 'get_auth_headers'),
    ('io.TextIOWrapper.read', 'anthropic/lib/bedrock/_auth.py', 'get_auth_headers'),
    ('io.BufferedReader.read', 'anthropic/lib/bedrock/_auth.py', 'get_auth_headers'),
    ('os.stat', 'anthropic/lib/aws/_auth.py', 'get_auth_headers'),
    ('io.TextIOWrapper.read', 'anthropic/lib/aws/_auth.py', 'get_auth_headers'),
    ('io.BufferedReader.read', 'anthropic/lib/aws/_auth.py', 'get_auth_headers'),
    # Decoding the first stream from an Anthropic Bedrock client loads botocore's `bedrock-runtime` service model
    # from disk, once per process (`lru_cache`).
    ('os.stat', 'anthropic/lib/bedrock/_stream_decoder.py', 'get_response_stream_shape'),
    ('os.listdir', 'anthropic/lib/bedrock/_stream_decoder.py', 'get_response_stream_shape'),
    ('io.TextIOWrapper.read', 'anthropic/lib/bedrock/_stream_decoder.py', 'get_response_stream_shape'),
    ('io.BufferedReader.read', 'anthropic/lib/bedrock/_stream_decoder.py', 'get_response_stream_shape'),
    # pydantic extracts field docstrings from source (`inspect`/`linecache`) the first time a
    # tool schema is built, which can happen during an agent run.
    ('os.stat', 'pydantic_ai/_function_schema.py', 'function_schema'),
    ('io.TextIOWrapper.read', 'pydantic_ai/_function_schema.py', 'function_schema'),
    # Enum member docstrings are read from source the same way, once per enum, when a schema is built.
    ('os.stat', 'pydantic_ai/_utils.py', 'enum_member_docstrings'),
    ('os.getcwd', 'pydantic_ai/_utils.py', 'enum_member_docstrings'),
    ('io.TextIOWrapper.read', 'pydantic_ai/_utils.py', 'enum_member_docstrings'),
    ('io.BufferedReader.read', 'pydantic_ai/_utils.py', 'enum_member_docstrings'),
    # A local workspace built with a relative `working_dir` resolves it against the current directory
    # once, at construction, which may happen in async code.
    ('os.getcwd', 'pydantic_ai/workspaces/local.py', '__init__'),
    # Prefect's `Task` reads its function's source for display when core builds a task, which the
    # durable workspace tests do while constructing an agent inside a running test.
    ('os.stat', 'prefect/tasks.py', '__init__'),
    # logfire resolves the current working directory while classifying user stack frames.
    ('os.getcwd', 'logfire/_internal/stack_info.py', 'is_user_code'),
    # `Dataset.to_file`/`from_file` and schema saving are sync serialization APIs; file I/O is
    # their documented job, even when called from async user code.
    ('os.stat', 'pydantic_evals/dataset.py', ('to_file', 'from_file', '_save_schema')),
    ('io.TextIOWrapper.read', 'pydantic_evals/dataset.py', ('to_file', 'from_file', '_save_schema')),
    ('io.TextIOWrapper.write', 'pydantic_evals/dataset.py', ('to_file', 'from_file', '_save_schema')),
    ('io.BufferedReader.read', 'pydantic_evals/dataset.py', ('to_file', 'from_file', '_save_schema')),
    ('io.BufferedWriter.write', 'pydantic_evals/dataset.py', ('to_file', 'from_file', '_save_schema')),
]


def _configure_blockbuster(
    exemptions: Sequence[tuple[str, str, str | tuple[str, ...]]] = BLOCKBUSTER_EXEMPTIONS,
    excluded_modules: tuple[str, ...] = (),
) -> BlockBuster:
    # BlockBuster imports ForbiddenFruit, which mutates builtins during import. Disabled CI lanes
    # must remain unaffected by that instrumentation.
    from blockbuster import BlockBuster

    # The harness isn't installed in every CI lane (the `pydantic-evals` one, say).
    harness = ['pydantic_ai_harness'] if importlib.util.find_spec('pydantic_ai_harness') is not None else []
    bb = BlockBuster(
        ['pydantic_ai', *harness, 'pydantic_graph', 'pydantic_evals', 'clai'],
        excluded_modules=excluded_modules or None,
    )
    for func, filename, functions in exemptions:
        bb.functions[func].can_block_in(filename, functions)
    return bb


@cache
def _configured_blockbuster(excluded_modules: tuple[str, ...]) -> BlockBuster:
    """Create one configured detector for each module-exclusion set in a pytest worker."""
    return _configure_blockbuster(excluded_modules=excluded_modules)


@pytest.fixture
def blockbuster_excluded_modules() -> tuple[str, ...]:
    return ()


@pytest.fixture
def blockbuster_enabled() -> bool:
    return True


@contextmanager
def _activated_blockbuster(blockbuster: BlockBuster) -> Generator[BlockBuster]:
    try:
        blockbuster.activate()
        yield blockbuster
    finally:
        blockbuster.deactivate()


@pytest.fixture(autouse=True)
def blockbuster(
    blockbuster_enabled: bool,
    blockbuster_excluded_modules: tuple[str, ...],
) -> Iterator[BlockBuster | None]:
    """Raise `BlockingError` when library code makes a blocking call inside the event loop.

    Scanned modules are the shipped packages only, so test-only and third-party stacks (e.g. VCR
    reading cassettes) are ignored. Calls from user callbacks and tools remain covered when a
    scanned library frame is below them. CI enables the detector on one complete Python-version
    lane to avoid multiplying its stack-inspection overhead across the version matrix.
    """
    if os.getenv('BLOCKBUSTER_ENABLED') == 'false' or not blockbuster_enabled:
        yield None
        return

    bb = _configured_blockbuster(blockbuster_excluded_modules)
    with _activated_blockbuster(bb):
        yield bb


def detach_dbos_logging() -> None:
    """Detach the OTel `LoggingHandler` and `DBOSLogTransformer` filter `DBOS.destroy()` leaves behind.

    DBOS attaches them to the root logger and every registered logger (the filter always, the handler
    with `enable_otlp=True`) and does not detach them on destroy. Their emit path imports
    `dbos._context`, which imports `http.server`: fatal inside the Temporal workflow sandbox when a
    later test in the same process runs a Temporal workflow.
    TODO(dsfaccini): Drop once DBOS cleans up its handlers on destroy.
    https://github.com/dbos-inc/dbos-transact-py/issues/871
    """
    from dbos import _logger as dbos_logger_module
    from opentelemetry.sdk._logs import LoggingHandler

    handler_types: tuple[type[logging.Handler], ...] = (LoggingHandler,)
    # DBOS 2.31+ builds this one instead of the deprecated SDK handler.
    with suppress(ImportError):
        from opentelemetry.instrumentation.logging.handler import LoggingHandler as InstrumentationLoggingHandler

        handler_types += (InstrumentationLoggingHandler,)

    for logger in [logging.root, *(logging.getLogger(name) for name in logging.root.manager.loggerDict)]:
        for handler in [h for h in logger.handlers if isinstance(h, handler_types)]:
            logger.removeHandler(handler)
        for log_filter in [f for f in logger.filters if isinstance(f, dbos_logger_module.DBOSLogTransformer)]:
            logger.removeFilter(log_filter)


@pytest.fixture
def allow_model_requests():
    with pydantic_ai.models.override_allow_model_requests(True):
        yield


# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false
@pytest.fixture
def create_module(tmp_path: Path, request: pytest.FixtureRequest) -> Callable[[str], Any]:
    """Taken from `pydantic/tests/conftest.py`, create module object, execute and return it."""

    def run(
        source_code: str,
        rewrite_assertions: bool = True,
        module_name_prefix: str | None = None,
    ) -> ModuleType:
        """Create module object, execute and return it.

        Can be used as a decorator of the function from the source code of which the module will be constructed.

        Args:
            source_code: Python source code of the module
            rewrite_assertions: whether to rewrite assertions in module or not
            module_name_prefix: string prefix to use in the name of the module, does not affect the name of the file.

        """

        # Max path length in Windows is 260. Leaving some buffer here
        max_name_len = 240 - len(str(tmp_path))
        sanitized_name = sanitize_filename(request.node.name, max_name_len)
        module_name = f'{sanitized_name}_{secrets.token_hex(5)}'
        path = tmp_path / f'{module_name}.py'
        path.write_text(source_code, encoding='utf-8')
        filename = str(path)

        if module_name_prefix:  # pragma: no cover
            module_name = module_name_prefix + module_name

        if rewrite_assertions:
            loader = AssertionRewritingHook(config=request.config)
            loader.mark_rewrite(module_name)
        else:  # pragma: no cover
            loader = None

        spec = importlib.util.spec_from_file_location(module_name, filename, loader=loader)
        sys.modules[module_name] = module = importlib.util.module_from_spec(spec)  # pyright: ignore[reportArgumentType]
        spec.loader.exec_module(module)  # pyright: ignore[reportOptionalMemberAccess]
        return module

    return run


@contextmanager
def try_import() -> Generator[Callable[[], bool]]:
    import_success = False

    def check_import() -> bool:
        return import_success

    try:
        yield check_import
    except ImportError:
        pass
    else:
        import_success = True


@pytest.fixture(scope='session', autouse=True)
def event_loop() -> Iterator[None]:
    new_loop = asyncio.new_event_loop()
    asyncio.set_event_loop(new_loop)
    yield
    new_loop.close()


class UndrivableEventLoop(asyncio.AbstractEventLoop):
    """An event loop that inherits `AbstractEventLoop`'s unimplemented `run_until_complete()`.

    This is how Temporal's workflow event loop behaves: it subclasses `asyncio.AbstractEventLoop` and never
    implements `run_until_complete()`, so calling that raises a bare `NotImplementedError` from CPython.
    """


@contextmanager
def undrivable_event_loop() -> Generator[None]:
    """Make the current event loop one that can't be driven by the caller."""
    previous = asyncio.get_event_loop()
    asyncio.set_event_loop(UndrivableEventLoop())
    try:
        yield
    finally:
        asyncio.set_event_loop(previous)


@pytest.fixture
def closed_event_loop() -> Iterator[asyncio.AbstractEventLoop]:
    original_loop = asyncio.get_event_loop()
    closed_loop = asyncio.new_event_loop()
    asyncio.set_event_loop(closed_loop)
    closed_loop.close()

    try:
        yield closed_loop
    finally:
        asyncio.get_event_loop().close()
        asyncio.set_event_loop(original_loop)


@pytest.fixture
def missing_event_loop() -> Iterator[asyncio.AbstractEventLoop]:
    """Empty the thread's event loop slot, yielding the loop that was installed before."""
    original_loop = asyncio.get_event_loop()
    asyncio.set_event_loop(None)

    try:
        yield original_loop
    finally:
        with suppress(RuntimeError):
            asyncio.get_event_loop().close()
        asyncio.set_event_loop(original_loop)


@pytest.fixture
def young_gc() -> Iterator[None]:
    """Limit the test's `gc.collect()` calls to the objects it creates.

    A full collection walks every object the worker has accumulated, which takes seconds per call
    late in a CI run. Freezing everything that exists before the test keeps it out of the collector's
    reach until teardown, so collections stay cheap, and garbage leaked by an earlier test can't be
    finalized (and raise an unraisable exception or `ResourceWarning`) in the middle of this one.
    """
    gc.freeze()
    try:
        yield
    finally:
        gc.unfreeze()


@pytest.fixture(autouse=True)
def no_instrumentation_by_default():
    Agent.instrument_all(False)
    Embedder.instrument_all(False)
    ImageGenerator.instrument_all(False)


@pytest.fixture(autouse=True)
def fresh_cache_mark_store(monkeypatch: pytest.MonkeyPatch) -> None:
    """Prompt-cache marks are kept process-wide per conversation; tests that reuse a fixed conversation
    id (or pin the clock) must not see each other's."""
    monkeypatch.setattr(instrumentation_capability, '_conversation_cache_marks', ConversationCacheMarkStore())


try:
    import logfire
    from opentelemetry import context as otel_context
    from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor

    logfire.DEFAULT_LOGFIRE_INSTANCE.config.ignore_no_config = True

    _httpx_instrumentor = HTTPXClientInstrumentor()

    @pytest.fixture(autouse=True)
    def fresh_logfire():
        logfire.shutdown(flush=False)
        # `test_examples.py` runs doc snippets that call the process-global `logfire.instrument_httpx()`,
        # which patches httpx via OTel and is never torn down. Reset it so it can't leak request spans
        # into other tests sharing the xdist worker (e.g. stray `POST` spans in `tests/durable_exec/temporal` snapshots).
        if _httpx_instrumentor._is_instrumented_by_opentelemetry:  # pyright: ignore[reportPrivateUsage]
            _httpx_instrumentor.uninstrument()
        # The worker's main-thread OTel context also persists across tests: an `attach` without a
        # matching `detach` leaves its values (an active span, suppression flags) visible to every
        # later test in the worker. A leaked *non-sampled* span is the nasty case: parent-based
        # sampling then marks all descendant spans unsampled and exporters silently drop them
        # (seen as `context_subtree()` returning an empty tree in `tests/evals/test_otel.py`).
        # The detach below resets the context to its pre-test snapshot, so no leaked `attach`
        # outlives its test. (The clean context itself only shields sync test bodies: async bodies
        # run in anyio's runner task, whose context is copied before this fixture attaches.)
        token = otel_context.attach(otel_context.Context())
        try:
            yield
        finally:
            otel_context.detach(token)

except ImportError:
    pass


@pytest.fixture(scope='session')
def prefect_test_server() -> Iterator[None]:
    """A Prefect test server with an isolated database, shared by every Prefect test in the worker.

    Starting one takes 15-25s on a CI runner, so tests that run Prefect flows request this fixture
    and share the `prefect` xdist group, starting a single server per job instead of one per module
    or test. The implicit ephemeral server would instead use the shared default `PREFECT_HOME` and a
    short connect timeout that flakes on slow runners.
    """
    pytest.importorskip('prefect')
    from prefect.settings import PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED, temporary_settings
    from prefect.testing.utilities import prefect_test_harness

    # The task-run recorder is a background writer against the same sqlite file the flows write to.
    # Prefect PRAGMAs a 60s `busy_timeout` onto every connection, and under CI contention the
    # recorder's bulk inserts exhaust it, failing the flow whose state it was recording. Nothing
    # here reads what it records: task run states reach the API through the task engine.
    with (
        temporary_settings({PREFECT_SERVER_SERVICES_TASK_RUN_RECORDER_ENABLED: False}),
        prefect_test_harness(server_startup_timeout=120),
    ):
        yield
    # Prefect's test server leaves client sockets for GC on Python 3.14; collect them here, where the
    # warning is expected, rather than in whichever test the collector happens to run in.
    with warnings.catch_warnings():
        warnings.filterwarnings('ignore', message='unclosed.*socket', category=ResourceWarning)
        gc.collect()


def raise_if_exception(e: Any) -> None:
    if isinstance(e, Exception):
        raise e


def pytest_addoption(parser: Any) -> None:
    parser.addoption(
        '--anyio-backend',
        choices=('asyncio', 'trio'),
        default='asyncio',
        help='Select the async test backend without duplicating the suite (default: asyncio).',
    )
    parser.addoption(
        '--xai-proto-include-json',
        action='store_true',
        default=True,
        dest='xai_proto_include_json',
        help='Include JSON representations in xAI proto cassette YAML files.',
    )
    parser.addoption(
        '--run-gateway-live',
        action='store_true',
        default=False,
        help='Run live gateway smoke tests that make real paid model requests.',
    )
    parser.addoption(
        '--strict-vcr-cassette-usage',
        action='store_true',
        help='Fail when a loaded VCR cassette has no interactions played, not only when playback leaves a stale tail.',
    )


@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_makereport(
    item: pytest.Item, call: pytest.CallInfo[None]
) -> Generator[None, Result[pytest.TestReport], None]:
    outcome = yield
    report = outcome.get_result()
    setattr(item, f'rep_{report.when}', report)


@pytest.fixture(scope='module')
def vcr_config() -> dict[str, Any]:
    """Module `conftest.py` overrides should extend this (`{**vcr_config, ...}`) so the hooks stay on."""
    return {
        'ignore_localhost': True,
        'before_record_request': cassette_hooks.before_record_request,
        'before_record_response': cassette_hooks.before_record_response,
        'uri_normalizer': cassette_hooks.normalize_uri,
    }


def check_vcr_cassette_usage(vcr: Cassette, strict_usage: bool) -> None:
    if vcr.play_count == 0 and not strict_usage:
        return

    # Each protocol numbers its interactions from 0, and `play_counts` only covers HTTP.
    unused = {
        'HTTP': [index for index in range(len(vcr.interactions)) if vcr.play_counts.get(index, 0) == 0],
        'gRPC': [index for index, played in enumerate(vcr.grpc_played_indices) if not played],
        'WebSocket': [index for index, played in enumerate(vcr.ws_played_indices) if not played],
    }
    if any(unused.values()):
        details = '; '.join(f'unused {protocol} indexes: {indexes}' for protocol, indexes in unused.items() if indexes)
        pytest.fail(f'Cassette {vcr.path} did not play all interactions: {details}')


@pytest.fixture(autouse=True)
def fail_partially_used_vcr_cassettes(request: pytest.FixtureRequest, vcr: Cassette | None) -> Iterator[None]:
    yield
    setup_report = getattr(request.node, 'rep_setup', None)
    call_report = getattr(request.node, 'rep_call', None)
    if any(
        getattr(report, 'skipped', False) or getattr(report, 'failed', False) for report in (setup_report, call_report)
    ):
        return
    if vcr is None or vcr.record_mode != RecordMode.NONE or vcr.all_played:
        return

    strict_usage = bool(request.config.getoption('--strict-vcr-cassette-usage'))
    check_vcr_cassette_usage(vcr, strict_usage)


@pytest.fixture(autouse=True)
def fail_cache_prefix_violations(request: pytest.FixtureRequest, vcr: Cassette | None) -> Iterator[None]:
    """Check final recorded conversations during playback; recording leaves the on-disk cassette unfinished."""
    yield
    setup_report = getattr(request.node, 'rep_setup', None)
    call_report = getattr(request.node, 'rep_call', None)
    if any(
        getattr(report, 'skipped', False) or getattr(report, 'failed', False) for report in (setup_report, call_report)
    ):
        return
    if vcr is None or vcr.record_mode != RecordMode.NONE:
        return
    # Playback never records, so the interactions cassetter loaded are exactly the file's contents.
    check_cache_prefix_stability(request.node, vcr)


# `validate_json` parses through pydantic-core rather than the stdlib, and types the result without a cast.
_REQUEST_BODY_ADAPTER = TypeAdapter(dict[str, JsonValue])

# What `httpx2.AsyncClient()` uses when no timeout is passed.
_HTTPX_DEFAULT_TIMEOUT = 5.0


@dataclass
class RequestCapture:
    """Outbound request bodies, as the live code built them.

    A cassette records what was sent when it was recorded, and the default matchers ignore the body,
    so a request whose payload has since drifted still replays against its recording. httpx event
    hooks run inside `AsyncClient.send`, above the transport VCR patches, so they fire on replay too
    and see what is actually going out. Pass `capture.client` as a provider's `http_client` and
    snapshot a projection of `capture.body(...)` to pin the fields a test's claim rests on.
    A provider that needs a longer deadline than httpx's default takes `capture.http_client(...)`.
    """

    paths: list[str] = field(default_factory=list[str])
    raw_bodies: list[bytes] = field(default_factory=list[bytes])
    headers: list[httpx2.Headers] = field(default_factory=list[httpx2.Headers])
    client: httpx2.AsyncClient = field(init=False)

    def __post_init__(self) -> None:
        self.client = httpx2.AsyncClient(event_hooks={'request': [self._record]})

    def http_client(self, *, timeout: float = _HTTPX_DEFAULT_TIMEOUT) -> httpx2.AsyncClient:
        """`client`, with `timeout` as the read timeout, for a provider that reads it back out.

        `GoogleProvider` reads the injected client's read timeout into google-genai's
        `HttpOptions.timeout`, which becomes both the per-request httpx timeout and the
        `X-Server-Timeout` deadline the API enforces. Gemini rejects a deadline under 10 seconds, and
        httpx defaults to 5. The default here is httpx's own, so it leaves the client as `client` has
        it; only a caller that needs a longer deadline passes a `timeout`.
        """
        self.client.timeout = timeout
        return self.client

    async def _record(self, request: httpx2.Request) -> None:
        # Only the raw bytes are kept here: the hook runs on every request of every test that asks
        # for a capture, while a test typically inspects one of them. Parsing happens in `body`.
        self.paths.append(request.url.path)
        self.raw_bodies.append(request.read())
        # The cassette serializer strips `anthropic-*` headers, so the wire is the only place a test
        # can see beta gating.
        self.headers.append(request.headers)

    def bodies(self, path_suffix: str = '') -> list[dict[str, JsonValue]]:
        """Every captured body whose URL path ends with `path_suffix`, parsed on demand."""
        return [
            _REQUEST_BODY_ADAPTER.validate_json(raw)
            for path, raw in zip(self.paths, self.raw_bodies)
            if path.endswith(path_suffix)
        ]

    def body(self, path_suffix: str = '', index: int = 0) -> dict[str, JsonValue]:
        """The `index`th captured body whose URL path ends with `path_suffix`, parsed on demand."""
        matches = self.bodies(path_suffix)
        assert matches, f'no captured request matching {path_suffix!r}; saw {self.paths}'
        return matches[index]


@pytest.fixture
async def request_capture(anyio_backend: str) -> AsyncIterator[RequestCapture]:
    capture = RequestCapture()
    yield capture
    # Built directly rather than through `create_async_httpx2_client`, so the autouse
    # `close_httpx_clients` tracker never sees it and its pool would otherwise leak per test.
    await capture.client.aclose()


_HttpClient: TypeAlias = 'httpx.AsyncClient | httpx2.AsyncClient'
_HttpClientCache: TypeAlias = 'dict[tuple[str, str], _HttpClient]'


@pytest.fixture(autouse=True)
def track_httpx_clients(monkeypatch: pytest.MonkeyPatch) -> Iterator[_HttpClientCache]:
    """Monkeypatch the HTTP client factories in all loaded modules and track created clients.

    Within a single test, calls with the same arguments reuse the same client. On teardown, all clients are closed — no process-global state leaks.

    This is a sync fixture so it applies to both sync and async tests. For async tests, the
    companion `close_httpx_clients` fixture handles async cleanup first.
    """
    cache: _HttpClientCache = {}
    original_httpx = pydantic_ai.models.create_async_http_client
    original_httpx2 = pydantic_ai._http.create_async_httpx2_client

    def make_cached(
        family: str, factory: Callable[..., _HttpClient], expected: type[_HttpClient]
    ) -> Callable[..., _HttpClient]:
        def cached_per_test(**kwargs: Any) -> _HttpClient:
            # `repr`, because `Timeout` and `Limits` arguments compare by value but aren't hashable.
            key = (family, repr(sorted(kwargs.items())))
            if key not in cache or cache[key].is_closed:
                cache[key] = factory(**kwargs)
            client = cache[key]
            assert isinstance(client, expected)
            return client

        return cached_per_test

    cached_per_test = make_cached('httpx', original_httpx, httpx.AsyncClient)
    cached_httpx2_per_test = make_cached('httpx2', original_httpx2, httpx2.AsyncClient)

    for mod in list(sys.modules.values()):
        # Read the module's own namespace via `__dict__` rather than `getattr`: some
        # modules (e.g. `transformers` submodules) define a lazy PEP 562 `__getattr__`
        # that imports submodules on attribute access, and probing every loaded module
        # with `getattr` would trigger those unrelated (and possibly failing) imports.
        mod_dict: dict[str, object] = getattr(mod, '__dict__', None) or {}
        if mod_dict.get('create_async_http_client', None) is original_httpx:
            monkeypatch.setattr(mod, 'create_async_http_client', cached_per_test)
        if mod_dict.get('create_async_httpx2_client', None) is original_httpx2:
            monkeypatch.setattr(mod, 'create_async_httpx2_client', cached_httpx2_per_test)

    yield cache

    unclosed = [c for c in cache.values() if not c.is_closed]
    if unclosed:  # pragma: no cover

        async def _close_all() -> None:
            for client in unclosed:
                await client.aclose()

        asyncio.run(_close_all())


@pytest.fixture(autouse=True)
async def close_httpx_clients(anyio_backend: str, track_httpx_clients: _HttpClientCache) -> AsyncIterator[None]:
    """Close tracked HTTP clients after async tests."""
    yield
    for client in track_httpx_clients.values():
        if not client.is_closed:
            await client.aclose()


try:
    from huggingface_hub.inference._providers._common import (
        _fetch_inference_provider_mapping as _hf_provider_mapping_func,  # pyright: ignore[reportPrivateUsage]
    )
except (ImportError, AttributeError):
    _hf_provider_mapping_func = None


@pytest.fixture(autouse=True)
def clear_huggingface_provider_cache():
    """Clear HuggingFace SDK's LRU cache after each test.

    The huggingface_hub library caches _fetch_inference_provider_mapping() with
    @lru_cache(maxsize=None), causing issues with VCR cassettes. The first test
    records the GET request, but subsequent tests skip it because the result is
    cached. This fixture ensures a fresh cache state for subsequent tests.
    """
    yield

    if _hf_provider_mapping_func is not None:
        _hf_provider_mapping_func.cache_clear()


@pytest.fixture(autouse=True, scope='session')
def patch_google_genai_gc_crash():
    """Work around google-genai BaseApiClient GC crash.

    BaseApiClient.__del__ schedules aclose() during GC, which crashes when the
    object was only partially initialized (_async_httpx_client never set).
    Remove when https://github.com/googleapis/python-genai/issues/2023 closes.
    """
    try:
        from google.genai._api_client import BaseApiClient
    except ImportError:
        yield
        return

    original_aclose = BaseApiClient.aclose

    async def safe_aclose(self: BaseApiClient) -> None:
        if hasattr(self, '_async_httpx_client'):
            await original_aclose(self)
        else:  # pragma: lax no cover
            # In some test runs, the `if` above will always run, so we get an `if -> exit` branch coverage miss.
            # This is a workaround to specify that the `else` branch may not be hit (as we don't have `lax no branch`)
            pass

    BaseApiClient.aclose = safe_aclose
    yield
    BaseApiClient.aclose = original_aclose


@pytest.fixture(scope='session')
def assets_path() -> Path:
    return Path(__file__).parent / 'assets'


@pytest.fixture(scope='session')
def audio_content(assets_path: Path) -> BinaryContent:
    audio_bytes = assets_path.joinpath('marcelo.mp3').read_bytes()
    return BinaryContent(data=audio_bytes, media_type='audio/mpeg')


@pytest.fixture(scope='session')
def image_content(assets_path: Path) -> BinaryImage:
    image_bytes = assets_path.joinpath('kiwi.jpg').read_bytes()
    return BinaryImage(data=image_bytes, media_type='image/jpeg')


@pytest.fixture(scope='session')
def video_content(assets_path: Path) -> BinaryContent:
    video_bytes = assets_path.joinpath('small_video.mp4').read_bytes()
    return BinaryContent(data=video_bytes, media_type='video/mp4')


@pytest.fixture(scope='session')
def document_content(assets_path: Path) -> BinaryContent:
    pdf_bytes = assets_path.joinpath('dummy.pdf').read_bytes()
    return BinaryContent(data=pdf_bytes, media_type='application/pdf')


@pytest.fixture(scope='session')
def text_document_content(assets_path: Path) -> BinaryContent:
    content = assets_path.joinpath('dummy.txt').read_text(encoding='utf-8')
    bin_content = BinaryContent(data=content.encode(), media_type='text/plain')
    return bin_content


# `tiny_*` fixtures: opaque 3-byte payloads for roundtrip / serialization tests where the bytes
# are never decoded by a model and assertions only check structure (e.g. `IsStr()` on base64).
# Prefer these whenever a real, decodable image is not required: keeps inline-snapshot diffs
# small and avoids CI errors that explode when failure traces include large base64 payloads.
# When the model has to actually see the content (e.g. VCR tests against provider APIs), use
# the KB-scale session fixtures above (`image_content`, `audio_content`, `video_content`).
@pytest.fixture
def tiny_image() -> BinaryImage:
    return BinaryImage(data=b'\x00\x01\x02', media_type='image/jpeg')


@pytest.fixture
def tiny_audio() -> BinaryContent:
    return BinaryContent(data=b'\x10\x11\x12', media_type='audio/mpeg')


@pytest.fixture
def tiny_video() -> BinaryContent:
    return BinaryContent(data=b'\x20\x21\x22', media_type='video/mp4')


os.environ.pop('OPENAI_BASE_URL', None)
os.environ.pop('ANTHROPIC_BASE_URL', None)
os.environ.pop('LOGFIRE_EMIT_CONFIGURATION_SPAN', None)


@pytest.fixture(scope='session')
def deepseek_api_key() -> str:
    return os.getenv('DEEPSEEK_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def openai_api_key() -> str:
    return os.getenv('OPENAI_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def azure_api_key() -> str:
    return os.getenv('AZURE_OPENAI_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def gemini_api_key() -> str:
    return os.getenv('GEMINI_API_KEY', os.getenv('GOOGLE_API_KEY', 'mock-api-key'))


@pytest.fixture(scope='session')
def groq_api_key() -> str:
    return os.getenv('GROQ_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def typesafe_api_key() -> str:
    return os.getenv('TYPESAFE_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def anthropic_api_key() -> str:
    return os.getenv('ANTHROPIC_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def gateway_api_key() -> str | None:
    return os.getenv('PYDANTIC_AI_GATEWAY_API_KEY', os.getenv('PAIG_API_KEY'))


@pytest.fixture(scope='session')
def co_api_key() -> str:
    return os.getenv('CO_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def voyage_api_key() -> str:
    return os.getenv('VOYAGE_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def mistral_api_key() -> str:
    return os.getenv('MISTRAL_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def openrouter_api_key() -> str:
    return os.getenv('OPENROUTER_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def ollama_api_key() -> str:
    return os.getenv('OLLAMA_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def huggingface_api_key() -> str:
    return os.getenv('HF_TOKEN', 'hf_token')


@pytest.fixture(autouse=True, scope='session')
def _patch_hf_provider_mappings():
    """Populate the SDK's hardcoded model mappings to avoid sync HTTP calls during VCR tests.

    The HuggingFace SDK makes a synchronous HTTP call to resolve provider mappings at request time,
    which is incompatible with VCR's async test infrastructure.
    """
    try:
        from huggingface_hub.hf_api import InferenceProviderMapping
        from huggingface_hub.inference._providers._common import HARDCODED_MODEL_INFERENCE_MAPPING
    except ImportError:
        return

    models: list[tuple[str, str, str]] = [
        ('together', 'deepseek-ai/DeepSeek-R1', 'conversational'),
        ('together', 'meta-llama/Llama-4-Scout-17B-16E-Instruct', 'conversational'),
        ('nebius', 'Qwen/Qwen2.5-VL-72B-Instruct', 'conversational'),
        ('nebius', 'Qwen/Qwen2.5-72B-Instruct', 'conversational'),
    ]

    for provider, model_id, task in models:
        HARDCODED_MODEL_INFERENCE_MAPPING[provider][model_id] = InferenceProviderMapping(
            provider=provider,
            hf_model_id=model_id,
            providerId=model_id,
            status='live',
            task=task,
        )


@pytest.fixture(scope='session')
def heroku_inference_key() -> str:
    return os.getenv('HEROKU_INFERENCE_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def cerebras_api_key() -> str:
    return os.getenv('CEREBRAS_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def xai_api_key() -> str:
    return os.getenv('XAI_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def tavily_api_key() -> str:
    return os.getenv('TAVILY_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def zai_api_key() -> str:
    return os.getenv('ZAI_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def crusoe_api_key() -> str:
    return os.getenv('CRUSOE_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def github_copilot_api_key() -> str:
    return os.getenv('GITHUB_COPILOT_API_KEY', 'mock-api-key')


@pytest.fixture(scope='session')
def snowflake_account() -> str:
    return os.getenv('SNOWFLAKE_ACCOUNT', 'myorg-myaccount')


@pytest.fixture(scope='session')
def snowflake_token() -> str:
    return os.getenv('SNOWFLAKE_TOKEN', 'mock-api-key')


@pytest.fixture(scope='function')  # Needs to be function scoped to get the request node name
async def xai_provider(request: pytest.FixtureRequest) -> AsyncIterator[XaiProvider | None]:
    """xAI provider fixture backed by protobuf cassettes.

    Mirrors the `bedrock_provider` pattern: yields a provider, and callers can use `provider.client`.
    Returns None for non-xAI tests to avoid loading cassettes unnecessarily.
    """
    if 'xai' not in request.node.name:
        yield None
        return

    try:
        from pydantic_ai.providers.xai import XaiProvider
        from tests.models.xai_proto_cassettes import xai_proto_cassette_session
    except ImportError:
        pytest.skip('xai_sdk not installed')

    cassette_name = sanitize_filename(request.node.name, 240)
    test_module = cast(str, request.node.fspath.basename.replace('.py', ''))
    cassette_path = Path(request.node.fspath).parent / 'cassettes' / test_module / f'{cassette_name}.xai.yaml'
    record_mode: str | None
    try:
        # Provided by cassetter's pytest plugin as `--record-mode=...`.
        record_mode = cast(Any, request.config).getoption('record_mode')
    except Exception:  # pragma: no cover
        record_mode = None
    include_debug_json = bool(cast(Any, request.config).getoption('xai_proto_include_json'))
    session = xai_proto_cassette_session(
        cassette_path,
        record_mode=record_mode,
        include_debug_json=include_debug_json,
    )
    provider = XaiProvider(xai_client=cast(Any, session.client))
    try:
        yield provider
    finally:
        session.dump_if_recording()
        await session.aclose()


@pytest.fixture(scope='session')
def bedrock_provider():
    try:
        import boto3

        from pydantic_ai.providers.bedrock import BedrockProvider

        bearer_token = os.getenv('AWS_BEARER_TOKEN_BEDROCK')
        if bearer_token:  # pragma: no cover
            provider = BedrockProvider(
                api_key=bearer_token,
                region_name=os.getenv('AWS_REGION', 'us-east-1'),
            )
            yield provider
            provider.client.close()
        else:  # pragma: lax no cover
            if os.getenv('AWS_PROFILE'):
                bedrock_client = boto3.client(
                    'bedrock-runtime',
                    region_name=os.getenv('AWS_REGION', 'us-east-1'),
                )
            else:
                bedrock_client = boto3.client(
                    'bedrock-runtime',
                    region_name=os.getenv('AWS_REGION', 'us-east-1'),
                    aws_access_key_id=os.getenv('AWS_ACCESS_KEY_ID', 'AKIA6666666666666666'),
                    aws_secret_access_key=os.getenv(
                        'AWS_SECRET_ACCESS_KEY', '6666666666666666666666666666666666666666'
                    ),
                    aws_session_token=os.getenv('AWS_SESSION_TOKEN', None),
                )
            yield BedrockProvider(bedrock_client=bedrock_client)
            bedrock_client.close()
    except ImportError:  # pragma: lax no cover
        pytest.skip('boto3 is not installed')


@pytest.fixture()
def vertex_provider_auth(mocker: MockerFixture) -> None:  # pragma: lax no cover
    # Locally, we authenticate via `gcloud` CLI, so we don't need to patch anything.
    if not os.getenv('CI', False):
        return  # pragma: lax no cover

    try:
        from google.genai import _api_client
    except ImportError:
        return  # do nothing if this isn't installed

    @dataclass
    class NoOpCredentials:
        token = 'my-token'
        quota_project_id = 'pydantic-ai'

        def refresh(self, request: httpx.Request): ...

        def expired(self) -> bool:
            return False

    return_value = (NoOpCredentials(), 'pydantic-ai')
    mocker.patch.object(_api_client, 'load_auth', return_value=return_value)


@pytest.fixture()
async def vertex_provider(vertex_provider_auth: None):  # pragma: lax no cover
    # NOTE: You need to comment out this line to rewrite the cassettes locally.
    if not os.getenv('CI', False):
        pytest.skip('Requires properly configured local google vertex config to pass')

    try:
        from pydantic_ai.providers.google import GoogleCloudLocation
        from pydantic_ai.providers.google_cloud import GoogleCloudProvider
    except ImportError:  # pragma: lax no cover
        pytest.skip('google is not installed')

    project = os.getenv('GOOGLE_PROJECT', 'pydantic-ai')
    location = os.getenv('GOOGLE_LOCATION', 'global')
    yield GoogleCloudProvider(project=project, location=cast(GoogleCloudLocation, location))


@pytest.fixture()
def model(
    request: pytest.FixtureRequest,
    openai_api_key: str,
    anthropic_api_key: str,
    mistral_api_key: str,
    groq_api_key: str,
    co_api_key: str,
    gemini_api_key: str,
    huggingface_api_key: str,
    bedrock_provider: BedrockProvider,
) -> Model:  # pragma: lax no cover
    try:
        if request.param == 'test':
            from pydantic_ai.models.test import TestModel

            return TestModel()
        elif request.param == 'openai':
            from pydantic_ai.models.openai import OpenAIChatModel
            from pydantic_ai.providers.openai import OpenAIProvider

            return OpenAIChatModel('o3-mini', provider=OpenAIProvider(api_key=openai_api_key))
        elif request.param == 'anthropic':
            from pydantic_ai.models.anthropic import AnthropicModel
            from pydantic_ai.providers.anthropic import AnthropicProvider

            return AnthropicModel('claude-sonnet-4-5', provider=AnthropicProvider(api_key=anthropic_api_key))
        elif request.param == 'mistral':
            from pydantic_ai.models.mistral import MistralModel
            from pydantic_ai.providers.mistral import MistralProvider

            return MistralModel('ministral-8b-latest', provider=MistralProvider(api_key=mistral_api_key))
        elif request.param == 'groq':
            from pydantic_ai.models.groq import GroqModel
            from pydantic_ai.providers.groq import GroqProvider

            return GroqModel('llama3-8b-8192', provider=GroqProvider(api_key=groq_api_key))
        elif request.param == 'cohere':
            from pydantic_ai.models.cohere import CohereModel
            from pydantic_ai.providers.cohere import CohereProvider

            return CohereModel('command-r-plus', provider=CohereProvider(api_key=co_api_key))
        elif request.param == 'google':
            from pydantic_ai.models.google import GoogleModel
            from pydantic_ai.providers.google import GoogleProvider

            return GoogleModel('gemini-1.5-flash', provider=GoogleProvider(api_key=gemini_api_key))
        elif request.param == 'bedrock':
            from pydantic_ai.models.bedrock import BedrockConverseModel

            return BedrockConverseModel('us.amazon.nova-micro-v1:0', provider=bedrock_provider)
        elif request.param == 'huggingface':
            from pydantic_ai.models.huggingface import HuggingFaceModel
            from pydantic_ai.providers.huggingface import HuggingFaceProvider

            return HuggingFaceModel(
                'Qwen/Qwen2.5-72B-Instruct',
                provider=HuggingFaceProvider(provider_name='nebius', api_key=huggingface_api_key),
            )
        else:
            raise ValueError(f'Unknown model: {request.param}')
    except ImportError:
        pytest.skip(f'{request.param} is not installed')


@pytest.fixture
def disable_ssrf_protection_for_vcr():
    """Disable SSRF protection for VCR compatibility.

    VCR cassettes record requests with the original hostname. Since SSRF protection
    resolves hostnames to IPs before making requests, we need to disable the validation
    for VCR tests to match the pre-recorded cassettes.

    This fixture patches validate_and_resolve_url to return the hostname in place
    of the resolved IP, allowing the request URL to use the original hostname.
    """
    from unittest.mock import patch

    from pydantic_ai._ssrf import ResolvedUrl, extract_host_and_port

    async def mock_validate_and_resolve(url: str, allow_local: bool) -> ResolvedUrl:
        hostname, path, port, is_https = extract_host_and_port(url)
        # Return hostname in place of resolved IP - this allows VCR matching
        return ResolvedUrl(resolved_ip=hostname, hostname=hostname, port=port, is_https=is_https, path=path)

    with patch('pydantic_ai._ssrf.validate_and_resolve_url', mock_validate_and_resolve):
        yield


_RequestPartT = TypeVar('_RequestPartT', bound=SystemPromptPart | UserPromptPart | ToolReturnPart | RetryPromptPart)
_ResponsePartT = TypeVar(
    '_ResponsePartT',
    bound=TextPart | ToolCallPart | NativeToolCallPart | NativeToolReturnPart | ThinkingPart | FilePart,
)


@overload
def iter_message_parts(
    messages: Sequence[ModelMessage],
    message_type: type[ModelRequest],
    part_type: type[_RequestPartT],
) -> Iterator[_RequestPartT]: ...


@overload
def iter_message_parts(
    messages: Sequence[ModelMessage],
    message_type: type[ModelResponse],
    part_type: type[_ResponsePartT],
) -> Iterator[_ResponsePartT]: ...


def iter_message_parts(
    messages: Sequence[ModelMessage],
    message_type: type[ModelRequest] | type[ModelResponse],
    part_type: type[_RequestPartT] | type[_ResponsePartT],
) -> Iterator[_RequestPartT | _ResponsePartT]:
    """Iterate over all parts of a given type in messages of a given type."""
    for msg in messages:  # pragma: no branch
        if isinstance(msg, message_type):
            for part in msg.parts:
                if isinstance(part, part_type):
                    yield part


def message(messages: Sequence[ModelMessage], message_type: type[T], *, index: int = 0) -> T:
    """Return `messages[index]`, asserting it is an instance of `message_type`."""
    msg = messages[index]
    assert isinstance(msg, message_type)
    return msg


def message_part(
    messages: Sequence[ModelMessage], part_type: type[T], *, message_index: int = 0, part_index: int = 0
) -> T:
    """Return `messages[message_index].parts[part_index]`, asserting it is an instance of `part_type`."""
    msg = messages[message_index]
    part = msg.parts[part_index]
    assert isinstance(part, part_type)
    return part


# endregion