pydantic--pydantic-ai/tests/test_tool_search.py

"""Tests for tool search functionality.

Unit tests for ToolSearchToolset plus VCR integration tests using pydantic-evals.

NOTE: If you change the search tool description or keyword schema in _tool_search.py,
re-record all cassettes with: uv run pytest tests/test_tool_search.py --record-mode=rewrite
"""

from __future__ import annotations

import json
import os
import re
from collections.abc import AsyncIterable, AsyncIterator, Sequence
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, TypeVar, cast

import pytest
import yaml
from inline_snapshot import snapshot
from pydantic import BaseModel
from typing_extensions import TypedDict

import pydantic_ai.agent as agent_module
from pydantic_ai import Agent, FunctionToolset, ToolCallPart
from pydantic_ai._agent_graph import _clean_message_history  # pyright: ignore[reportPrivateUsage]
from pydantic_ai._deferred_capabilities import DEFERRED_CAPABILITY_TOOL_METADATA_KEY
from pydantic_ai._run_context import RunContext
from pydantic_ai._tool_search import (
    synthesize_local_from_native_call,
    synthesize_local_tool_search_messages,
)
from pydantic_ai.capabilities import CAPABILITY_TYPES
from pydantic_ai.capabilities._ordering import collect_leaves
from pydantic_ai.capabilities._tool_search import ToolSearch
from pydantic_ai.capabilities.abstract import AbstractCapability
from pydantic_ai.capabilities.capability import Capability
from pydantic_ai.capabilities.combined import CombinedCapability
from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior, UserError
from pydantic_ai.messages import (
    AgentStreamEvent,
    LoadCapabilityCallPart,
    LoadCapabilityReturnPart,
    ModelMessage,
    ModelMessagesTypeAdapter,
    ModelRequest,
    ModelResponse,
    NativeToolCallPart,
    NativeToolReturnPart,
    NativeToolSearchCallPart,
    NativeToolSearchReturnPart,
    PartStartEvent,
    TextPart,
    ToolPartKind,
    ToolReturnPart,
    ToolSearchCallPart,
    ToolSearchReturnContent,
    ToolSearchReturnPart,
    UserPromptPart,
    _model_request_part_discriminator,  # pyright: ignore[reportPrivateUsage]
    _model_response_part_discriminator,  # pyright: ignore[reportPrivateUsage]
)
from pydantic_ai.models import ModelRequestParameters, infer_model
from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel
from pydantic_ai.models.test import TestModel
from pydantic_ai.native_tools import SUPPORTED_NATIVE_TOOLS, AbstractNativeTool, WebSearchTool
from pydantic_ai.native_tools._tool_search import ToolSearchMatch, ToolSearchTool
from pydantic_ai.profiles import ModelProfile, merge_profile
from pydantic_ai.run import AgentRunResult
from pydantic_ai.tool_manager import ToolManager
from pydantic_ai.tools import ToolDefinition
from pydantic_ai.toolsets import AbstractToolset
from pydantic_ai.toolsets._deferred_capability_loader import LOAD_CAPABILITY_TOOL_NAME
from pydantic_ai.toolsets._tool_search import (
    _SEARCH_TOOLS_NAME,  # pyright: ignore[reportPrivateUsage]
    ToolSearchToolset,
    keywords_search_fn,
    parse_discovered_tools,
)
from pydantic_ai.usage import RequestUsage, RunUsage

from .conftest import message, message_part, try_import

with try_import() as evals_available:
    from pydantic_evals import Case, Dataset
    from pydantic_evals.evaluators import Evaluator, EvaluatorContext
    from pydantic_evals.reporting import EvaluationReport

with try_import() as anthropic_available:
    import anthropic  # pyright: ignore[reportUnusedImport]  # noqa: F401
    from anthropic.types.beta import (
        BetaServerToolUseBlock,
        BetaTextBlock,
        BetaToolSearchToolResultBlock,
        BetaUsage,
    )
    from anthropic.types.beta.beta_server_tool_use_block import BetaDirectCaller
    from anthropic.types.beta.beta_tool_search_tool_result_error import BetaToolSearchToolResultError

    from pydantic_ai.models.anthropic import (
        AnthropicModel,
        AnthropicModelSettings,
        _build_custom_tool_search_replay_blocks,  # pyright: ignore[reportPrivateUsage]
        _build_tool_search_replay_block,  # pyright: ignore[reportPrivateUsage]
        _collect_orphan_tool_search_call_ids,  # pyright: ignore[reportPrivateUsage]
        _finalize_streamed_tool_search_call_part,  # pyright: ignore[reportPrivateUsage]
        _map_server_tool_use_block,  # pyright: ignore[reportPrivateUsage]
        _map_tool_search_tool_result_block,  # pyright: ignore[reportPrivateUsage]
    )
    from pydantic_ai.providers.anthropic import AnthropicProvider

    from .models.test_anthropic import MockAnthropic, completion_message, get_mock_chat_completion_kwargs

with try_import() as openai_available:
    from openai.types.responses import (
        FunctionTool,
        ResponseFunctionToolCall,
        ResponseFunctionToolCallParam,
        ResponseOutputMessage,
        ResponseOutputText,
        ResponseToolSearchCall,
        ResponseToolSearchOutputItem,
    )
    from openai.types.responses.file_search_tool import FileSearchTool

    from pydantic_ai.models.openai import (
        OpenAIResponsesModel,
        OpenAIResponsesModelSettings,
        _build_tool_search_return_part,  # pyright: ignore[reportPrivateUsage]
        _map_client_tool_search_call,  # pyright: ignore[reportPrivateUsage]
        _map_tool_search_call,  # pyright: ignore[reportPrivateUsage]
        _normalize_tool_search_args,  # pyright: ignore[reportPrivateUsage]
        _tool_search_namespace_for_synthesis,  # pyright: ignore[reportPrivateUsage]
    )
    from pydantic_ai.providers.openai import OpenAIProvider

    from .models.mock_openai import MockOpenAIResponses, get_mock_responses_kwargs, response_message

with try_import() as google_available:
    import google.genai  # pyright: ignore[reportUnusedImport]  # noqa: F401

pytestmark = pytest.mark.anyio

MOCK_API_KEYS: dict[str, str] = {
    'OPENAI_API_KEY': 'mock-api-key',
    'ANTHROPIC_API_KEY': 'mock-api-key',
    # google-gla checks GEMINI_API_KEY only. Mocking GOOGLE_API_KEY would shadow a real
    # GEMINI_API_KEY in .env because the google-genai SDK prefers GOOGLE_API_KEY when both
    # are present, so re-recording against real credentials would silently use the mock.
    'GEMINI_API_KEY': 'mock-api-key',
}


@pytest.fixture(autouse=True)
def _mock_api_keys(monkeypatch: pytest.MonkeyPatch) -> None:
    for key, default in MOCK_API_KEYS.items():
        if not os.getenv(key):  # pragma: no branch
            monkeypatch.setenv(key, default)


# --- Eval types ---


class EvalOutput(BaseModel):
    tool_calls: list[str]
    search_args: list[dict[str, str]]


class EvalMetadata(BaseModel):
    expected_tools: list[str]


# --- Evaluators ---

if evals_available():

    @dataclass(repr=False)
    class UsedSearchTools(Evaluator[str, EvalOutput, EvalMetadata]):
        """Check that the model used search_tools when expected tools exist."""

        def evaluate(self, ctx: EvaluatorContext[str, EvalOutput, EvalMetadata]) -> bool:
            if not ctx.metadata or not ctx.metadata.expected_tools:
                return True
            return 'search_tools' in ctx.output.tool_calls

        def get_default_evaluation_name(self) -> str:
            return 'used_search_tools'

    @dataclass(repr=False)
    class FoundExpectedTools(Evaluator[str, EvalOutput, EvalMetadata]):
        """Check that the model found and called the expected tools."""

        def evaluate(self, ctx: EvaluatorContext[str, EvalOutput, EvalMetadata]) -> bool:
            if not ctx.metadata or not ctx.metadata.expected_tools:
                return True
            return all(t in ctx.output.tool_calls for t in ctx.metadata.expected_tools)

        def get_default_evaluation_name(self) -> str:
            return 'found_expected_tools'

    @dataclass(repr=False)
    class ReasonableToolUsage(Evaluator[str, EvalOutput, EvalMetadata]):
        """Check that the model didn't use an excessive number of tool calls."""

        max_calls: int = 10

        def evaluate(self, ctx: EvaluatorContext[str, EvalOutput, EvalMetadata]) -> bool:
            return len(ctx.output.tool_calls) <= self.max_calls

        def get_default_evaluation_name(self) -> str:
            return 'reasonable_usage'

    @dataclass(repr=False)
    class KeywordCount(Evaluator[str, EvalOutput, EvalMetadata]):
        """Score the number of keywords used in the search query. Best is <= 3."""

        def evaluate(self, ctx: EvaluatorContext[str, EvalOutput, EvalMetadata]) -> int | dict[str, int]:
            if not ctx.output.search_args:
                return {}
            raw: Any = ctx.output.search_args[0].get('queries')
            queries = cast('list[str]', raw) if isinstance(raw, list) else ([str(raw)] if raw else [])
            return len(' '.join(queries).split())

        def get_default_evaluation_name(self) -> str:
            return 'keyword_count'


# --- Helpers ---


def _extract_tool_calls(result: AgentRunResult[str]) -> list[str]:
    """Extract tool-call names across both local and native tool-search paths.

    Normalizes native tool-search calls (`NativeToolSearchCallPart`, `tool_name='tool_search'`)
    to `'search_tools'` so the evaluator sees the same name regardless of which path the
    active provider took.
    """
    tool_calls: list[str] = []
    for msg in result.all_messages():
        if isinstance(msg, ModelResponse):
            for part in msg.parts:
                if isinstance(part, (ToolCallPart, NativeToolCallPart)):
                    name = 'search_tools' if part.tool_kind == 'tool-search' else part.tool_name
                    tool_calls.append(name)
    return tool_calls


def _extract_search_args(result: AgentRunResult[str]) -> list[dict[str, str]]:
    """Extract parsed args dicts from tool-search calls across local and native paths."""
    args_list: list[dict[str, str]] = []
    for msg in result.all_messages():
        if isinstance(msg, ModelResponse):
            for part in msg.parts:
                if (
                    isinstance(part, (ToolCallPart, NativeToolCallPart))
                    and part.tool_kind == 'tool-search'
                    and part.args is not None
                ):
                    parsed = json.loads(part.args) if isinstance(part.args, str) else part.args
                    args_list.append({k: str(v) for k, v in parsed.items()})
    return args_list


def _build_agent(model_name: str) -> Agent[object, str]:
    """Build an agent with a visible tool and several deferred tools for testing.

    Forces the local `search_tools` function-tool path on every provider by removing
    `ToolSearchTool` from the model profile's `supported_native_tools`. This eval
    exercises OUR search-tool prompts and behavior; providers' native tool-search
    paths use the provider's own prompts and aren't under test here.
    """
    model = infer_model(model_name)
    # Override the cached profile to drop ToolSearchTool — forces the local path
    # uniformly across providers with and without native tool-search support.
    setattr(
        model,
        'profile',
        merge_profile(
            model.profile,
            ModelProfile(
                supported_native_tools=model.profile.get('supported_native_tools', SUPPORTED_NATIVE_TOOLS)
                - {ToolSearchTool}
            ),
        ),
    )
    agent = Agent(model=model)

    @agent.tool_plain
    def get_weather(city: str) -> str:  # pragma: no cover
        """Get the current weather for a city."""
        return f'The weather in {city} is sunny and 72°F.'

    @agent.tool_plain(defer_loading=True)
    def get_exchange_rate(from_currency: str, to_currency: str) -> str:
        """Look up the current exchange rate between two currencies."""
        rates: dict[str, float] = {
            'USD_EUR': 0.92,
            'EUR_USD': 1.09,
            'USD_GBP': 0.79,
            'GBP_USD': 1.27,
        }
        key = f'{from_currency}_{to_currency}'
        rate = rates.get(key, 1.0)
        return f'1 {from_currency} = {rate} {to_currency}'

    @agent.tool_plain(defer_loading=True)
    def stock_lookup(symbol: str) -> str:
        """Look up stock price by ticker symbol."""
        return f'Stock {symbol}: $150.00'

    @agent.tool_plain(defer_loading=True)
    def mortgage_calculator(principal: float, rate: float, years: int) -> str:  # pragma: no cover
        """Calculate monthly mortgage payment for a home loan."""
        monthly_rate = rate / 12 / 100
        num_payments = years * 12
        if monthly_rate == 0:
            payment = principal / num_payments
        else:
            payment = (
                principal
                * (monthly_rate * (1 + monthly_rate) ** num_payments)
                / ((1 + monthly_rate) ** num_payments - 1)
            )
        return f'Monthly payment: ${payment:.2f}'

    return agent


if evals_available():

    def _build_dataset() -> Dataset[str, EvalOutput, EvalMetadata]:
        return Dataset[str, EvalOutput, EvalMetadata](
            name='tool_search',
            cases=[
                Case(
                    name='exchange_rate',
                    inputs='What is the current exchange rate from USD to EUR?',
                    metadata=EvalMetadata(expected_tools=['get_exchange_rate']),
                ),
                Case(
                    name='stock_price',
                    inputs='What is the current stock price for AAPL?',
                    metadata=EvalMetadata(expected_tools=['stock_lookup']),
                ),
                Case(
                    name='translation',
                    inputs="Translate 'hello, how are you?' to French.",
                    metadata=EvalMetadata(expected_tools=[]),
                ),
                Case(
                    name='no_matching_tool',
                    inputs='Book a flight from New York to London for next week.',
                    metadata=EvalMetadata(expected_tools=[]),
                ),
            ],
            evaluators=[
                UsedSearchTools(),
                FoundExpectedTools(),
                ReasonableToolUsage(max_calls=5),
                KeywordCount(),
            ],
        )


def _summarize_report(report: EvaluationReport[str, EvalOutput, EvalMetadata]) -> dict[str, ScenarioSummary]:
    """Extract a compact summary from eval report for snapshotting."""
    summary: dict[str, ScenarioSummary] = {}
    for case in report.cases:
        output: EvalOutput = case.output
        keywords: str | None = None
        if output.search_args:
            raw: Any = output.search_args[0].get('queries')
            queries = cast('list[str]', raw) if isinstance(raw, list) else ([str(raw)] if raw else [])
            keywords = ' '.join(queries) or None
        summary[case.name] = ScenarioSummary(keywords=keywords, tool_calls=output.tool_calls)
    return summary


class ScenarioSummary(TypedDict):
    """The search keywords the model chose and the tools it discovered and called."""

    keywords: str | None
    tool_calls: list[str]


@dataclass
class ModelCase:
    model_name: str
    marks: list[pytest.MarkDecorator] = field(default_factory=list[pytest.MarkDecorator])
    scenario_summary: dict[str, ScenarioSummary] = field(default_factory=dict[str, ScenarioSummary])


_CASES = [
    ModelCase(
        model_name='openai-chat:gpt-5.4-mini',
        scenario_summary=snapshot(
            {
                'exchange_rate': {
                    'keywords': "['exchange rate currency USD EUR current']",
                    'tool_calls': ['search_tools', 'get_exchange_rate'],
                },
                'stock_price': {
                    'keywords': "['stock price market quote AAPL current']",
                    'tool_calls': ['search_tools', 'stock_lookup'],
                },
                'translation': {'keywords': None, 'tool_calls': []},
                'no_matching_tool': {
                    'keywords': None,
                    'tool_calls': [],
                },
            }
        ),
    ),
    ModelCase(
        model_name='anthropic:claude-sonnet-4-5',
        marks=[
            pytest.mark.skipif(not anthropic_available(), reason='anthropic not installed'),
        ],
        scenario_summary=snapshot(
            {
                'exchange_rate': {
                    'keywords': "['currency exchange rate', 'USD EUR conversion', 'foreign exchange', 'currency converter']",
                    'tool_calls': ['search_tools', 'get_exchange_rate'],
                },
                'stock_price': {
                    'keywords': "['stock price', 'AAPL', 'ticker symbol', 'market data', 'financial data']",
                    'tool_calls': ['search_tools', 'stock_lookup', 'stock_lookup'],
                },
                'translation': {
                    'keywords': "['translate', 'translation', 'French', 'language']",
                    'tool_calls': ['search_tools'],
                },
                'no_matching_tool': {
                    'keywords': "['book flight', 'flight booking', 'airline reservation', 'travel booking']",
                    'tool_calls': ['search_tools'],
                },
            }
        ),
    ),
    ModelCase(
        model_name='google:gemini-3-flash-preview',
        marks=[pytest.mark.skipif(not google_available(), reason='google-genai not installed')],
        scenario_summary=snapshot(
            {
                'exchange_rate': {
                    'keywords': "['exchange rate', 'currency conversion']",
                    'tool_calls': ['search_tools', 'get_exchange_rate'],
                },
                'stock_price': {
                    'keywords': "['stock price', 'financial data', 'market data']",
                    'tool_calls': ['search_tools', 'stock_lookup'],
                },
                'translation': {'keywords': None, 'tool_calls': []},
                'no_matching_tool': {
                    'keywords': "['flight booking', 'search flights', 'book flight']",
                    'tool_calls': ['search_tools', 'search_tools'],
                },
            }
        ),
    ),
]


@pytest.mark.skipif(not evals_available(), reason='pydantic-evals not installed')
@pytest.mark.vcr
@pytest.mark.parametrize(
    'case',
    [pytest.param(c, id=c.model_name.split(':')[0], marks=c.marks) for c in _CASES],
)
async def test_tool_search_eval(allow_model_requests: None, case: ModelCase) -> None:
    """Evaluate tool search behavior across scenarios using pydantic-evals.

    Runs 4 scenarios per model: exchange_rate, stock_price, translation, no_matching_tool.
    Evaluators check: used_search_tools, found_expected_tools, reasonable_usage, keyword_count.
    """
    agent = _build_agent(case.model_name)

    async def task(prompt: str) -> EvalOutput:
        try:
            result = await agent.run(prompt)
        except UnexpectedModelBehavior:
            return EvalOutput(tool_calls=[], search_args=[])
        return EvalOutput(
            tool_calls=_extract_tool_calls(result),
            search_args=_extract_search_args(result),
        )

    dataset = _build_dataset()
    report = await dataset.evaluate(task, name='tool_search', progress=False, max_concurrency=1)

    assert not report.failures
    for eval_case in report.cases:
        for name, result in eval_case.assertions.items():
            assert result.value, f'{eval_case.name}/{name} failed'

    assert _summarize_report(report) == case.scenario_summary


# --- Unit tests ---

T = TypeVar('T')


class NoNativeToolSearchModel(FunctionModel):
    @classmethod
    def supported_native_tools(cls) -> frozenset[type[AbstractNativeTool]]:
        return frozenset(super().supported_native_tools()) - {ToolSearchTool}


def _build_run_context(
    deps: T,
    run_step: int = 0,
    messages: list[ModelMessage] | None = None,
    capabilities: dict[str, AbstractCapability[T]] | None = None,
    discovered_tool_names: set[str] | None = None,
) -> RunContext[T]:
    """Build a `RunContext` for unit tests using `TestModel`."""
    return RunContext(
        deps=deps,
        model=TestModel(),
        usage=RunUsage(),
        prompt=None,
        messages=messages or [],
        run_step=run_step,
        capabilities=capabilities or {},
        discovered_tool_names=discovered_tool_names or set(),
    )


def _create_function_toolset() -> FunctionToolset:
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain
    def get_weather(city: str) -> str:  # pragma: no cover
        """Get the current weather for a city."""
        return f'Weather in {city}'

    @toolset.tool_plain
    def get_time(timezone: str) -> str:  # pragma: no cover
        """Get the current time in a timezone."""
        return f'Time in {timezone}'

    @toolset.tool_plain(defer_loading=True)
    def calculate_mortgage(principal: float, rate: float, years: int) -> str:
        """Calculate monthly mortgage payment for a loan."""
        return 'Mortgage calculated'

    @toolset.tool_plain(defer_loading=True)
    def stock_price(symbol: str) -> str:  # pragma: no cover
        """Get the current stock price for a symbol."""
        return f'Stock price for {symbol}'

    @toolset.tool_plain(defer_loading=True)
    def crypto_price(coin: str) -> str:  # pragma: no cover
        """Get the current cryptocurrency price."""
        return f'Crypto price for {coin}'

    return toolset


async def test_tool_search_toolset_filters_deferred_tools():
    """On the local path, deferred tools stay hidden until discovered — only the
    visible tools and the `search_tools` function are exposed up front."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    tool_names = list(tools.keys())

    assert tool_names == snapshot(
        [
            'get_weather',
            'get_time',
            'calculate_mortgage',
            'stock_price',
            'crypto_price',
            'search_tools',
        ]
    )


async def test_search_tool_def_description_and_schema():
    """Test that the search tool definition includes deferred count and TypeAdapter-generated schema."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]

    assert search_tool.tool_def.description == snapshot(
        'There are additional tools not yet visible to you. When you need a capability not provided by your current tools, search here by providing one or more queries to discover and activate relevant tools. Each query is tokenized into words; tool names and descriptions are scored by token overlap. If no tools are found, they do not exist — do not retry.'
    )
    assert search_tool.tool_def.parameters_json_schema == snapshot(
        {
            'additionalProperties': False,
            'properties': {
                'queries': {
                    'description': 'List of search queries to match against tool names and descriptions. Use specific words likely to appear in tool names or descriptions to narrow down relevant tools. Each query is independently tokenized; matches across queries are unioned.',
                    'items': {'type': 'string'},
                    'type': 'array',
                }
            },
            'required': ['queries'],
            'type': 'object',
        }
    )


async def test_tool_search_toolset_search_returns_matching_tools():
    """Test that search_tools returns matching deferred tools."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]

    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['mortgage']}, ctx, search_tool)
    assert result == snapshot({'discovered_tools': [{'name': 'calculate_mortgage'}]})


async def test_tool_search_toolset_search_is_case_insensitive():
    """Test that search is case insensitive."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]

    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['STOCK']}, ctx, search_tool)
    rv = cast(ToolSearchReturnContent, result)
    assert len(rv['discovered_tools']) == 1
    assert rv['discovered_tools'][0]['name'] == 'stock_price'


async def test_tool_search_toolset_search_matches_description():
    """Test that search matches tool descriptions."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]

    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['cryptocurrency']}, ctx, search_tool)
    rv = cast(ToolSearchReturnContent, result)
    assert len(rv['discovered_tools']) == 1
    assert rv['discovered_tools'][0]['name'] == 'crypto_price'


async def test_tool_search_toolset_prefers_specific_term_matches():
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain(defer_loading=True)
    def github_get_me() -> str:  # pragma: no cover
        """Get the authenticated GitHub profile."""
        return 'me'

    @toolset.tool_plain(defer_loading=True)
    def github_create_gist() -> str:  # pragma: no cover
        """Create a new GitHub gist."""
        return 'gist'

    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]

    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['github profile']}, ctx, search_tool)
    assert result == snapshot(
        {
            'discovered_tools': [
                {'name': 'github_get_me'},
                {'name': 'github_create_gist'},
            ]
        }
    )


async def test_tool_search_toolset_keeps_lower_scoring_matches_after_top_hits():
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain(defer_loading=True)
    def stock_price() -> str:  # pragma: no cover
        """Get the current stock price."""
        return 'stock'

    @toolset.tool_plain(defer_loading=True)
    def crypto_price() -> str:  # pragma: no cover
        """Get the current cryptocurrency price."""
        return 'crypto'

    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]

    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['stock price']}, ctx, search_tool)
    assert result == snapshot(
        {
            'discovered_tools': [
                {'name': 'stock_price'},
                {'name': 'crypto_price'},
            ]
        }
    )


async def test_tool_search_toolset_does_not_match_substrings_inside_words():
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain(defer_loading=True)
    def github_get_me() -> str:  # pragma: no cover
        """Get my GitHub profile."""
        return 'me'

    @toolset.tool_plain(defer_loading=True)
    def github_add_comment_to_pending_review() -> str:  # pragma: no cover
        """Add a pending review comment on GitHub."""
        return 'comment'

    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]

    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['get me']}, ctx, search_tool)
    assert result == snapshot({'discovered_tools': [{'name': 'github_get_me'}]})


async def test_tool_search_toolset_search_returns_no_matches():
    """Test that search returns empty list when no matches."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]

    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['nonexistent']}, ctx, search_tool)
    assert result == snapshot(
        {'discovered_tools': [], 'message': 'No matching tools found. The tools you need may not be available.'}
    )


async def test_tool_search_toolset_search_empty_query():
    """Test that search with empty query raises ModelRetry."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]

    with pytest.raises(ModelRetry, match=re.escape('Please provide at least one non-empty search query.')):
        await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['']}, ctx, search_tool)


@pytest.mark.parametrize('query', ['   ', '---', '!!!', '...'])
async def test_tool_search_toolset_search_non_tokenizable_query(query: str):
    """Queries that tokenize to an empty set must retry, not match every tool."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]

    with pytest.raises(ModelRetry, match=re.escape('Please provide at least one non-empty search query.')):
        await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': [query]}, ctx, search_tool)


async def test_tool_search_toolset_max_results():
    """Test that results are capped at `_MAX_SEARCH_RESULTS` (10)."""
    toolset: FunctionToolset = FunctionToolset()

    for i in range(15):

        @toolset.tool_plain(defer_loading=True, name=f'tool_{i}')
        def tool_func() -> str:  # pragma: no cover
            """A tool for testing."""
            return 'result'

    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]

    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['tool']}, ctx, search_tool)
    rv = cast(ToolSearchReturnContent, result)
    assert len(rv['discovered_tools']) == 10


async def test_tool_search_toolset_discovered_tools_flip_defer_loading():
    """Discovered tools have `defer_loading=False`; undiscovered ones still have
    `defer_loading=True`. Both stay in the toolset under their real names — the
    wire-side filter in `Model.prepare_request` decides what reaches the model."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)

    messages: list[ModelMessage] = [
        ModelRequest(
            parts=[
                ToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'calculate_mortgage'}]},
                ),
            ]
        )
    ]
    ctx = _build_run_context(None, messages=messages, discovered_tool_names={'calculate_mortgage'})

    tools = await searchable.get_tools(ctx)
    assert tools['calculate_mortgage'].tool_def.defer_loading is False
    assert tools['stock_price'].tool_def.defer_loading is True
    assert tools['crypto_price'].tool_def.defer_loading is True


async def test_tool_search_toolset_keeps_search_tool_after_all_discovered():
    """`search_tools` stays in the request even when every deferred tool is discovered.

    Dropping it would invalidate the cached request prefix on the next turn — keeping
    it preserves prompt caching across discovery steps. The local tool's body is a no-op
    branch in `_search_tools` since the index is empty, and on native paths it's dropped
    by the adapter via its `unless_native='tool_search'` flag anyway.
    """
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)

    messages: list[ModelMessage] = [
        ModelRequest(
            parts=[
                ToolReturnPart(
                    tool_name=_SEARCH_TOOLS_NAME,
                    content={
                        'tools': [
                            {'name': 'calculate_mortgage'},
                            {'name': 'stock_price'},
                            {'name': 'crypto_price'},
                        ]
                    },
                )
            ]
        )
    ]
    ctx = _build_run_context(None, messages=messages)

    tools = await searchable.get_tools(ctx)
    tool_names = list(tools.keys())

    assert tool_names == snapshot(
        [
            'get_weather',
            'get_time',
            'calculate_mortgage',
            'stock_price',
            'crypto_price',
            'search_tools',
        ]
    )


async def test_tool_search_toolset_reserved_name_collision():
    """Test that `UserError` is raised if a tool is named 'search_tools' and deferred tools exist."""
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain
    def search_tools(query: str) -> str:  # pragma: no cover
        """Search for tools."""
        return 'search result'

    @toolset.tool_plain(defer_loading=True)
    def deferred_tool() -> str:  # pragma: no cover
        """A deferred tool to trigger search injection."""
        return 'deferred'

    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    with pytest.raises(UserError, match="Tool name 'search_tools' is reserved"):
        await searchable.get_tools(ctx)


async def test_tool_search_toolset_no_deferred_tools_returns_all():
    """Test that when there are no deferred tools, all tools are returned without search_tools."""
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain
    def get_weather(city: str) -> str:  # pragma: no cover
        """Get the current weather for a city."""
        return f'Weather in {city}'

    @toolset.tool_plain
    def get_time(timezone: str) -> str:  # pragma: no cover
        """Get the current time in a timezone."""
        return f'Time in {timezone}'

    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    tool_names = list(tools.keys())

    assert tool_names == snapshot(['get_weather', 'get_time'])


async def test_tool_search_handles_search_gated_tools_from_eager_capability():
    """Search-gated tools from eager capabilities stay hidden until searched."""
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain(defer_loading=True)
    def capability_search_tool() -> str:
        """Search-gated tool owned by an eager capability."""
        return 'search-gated-result'

    capability = Capability(
        id='example',
        description='Example capability.',
        toolsets=[toolset],
    )

    seen_tool_names: list[list[str]] = []

    def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        seen_tool_names.append([t.name for t in info.function_tools])
        tool_returns = [
            part
            for message in messages
            if isinstance(message, ModelRequest)
            for part in message.parts
            if isinstance(part, ToolReturnPart)
        ]

        if not any(part.tool_name == _SEARCH_TOOLS_NAME for part in tool_returns):
            return ModelResponse(
                parts=[
                    ToolCallPart(
                        tool_name=_SEARCH_TOOLS_NAME,
                        args={'queries': ['eager']},
                        tool_call_id='search-1',
                    )
                ]
            )

        if not any(part.tool_name == 'capability_search_tool' for part in tool_returns):
            return ModelResponse(
                parts=[
                    ToolCallPart(
                        tool_name='capability_search_tool',
                        args={},
                        tool_call_id='call-search-gated',
                    )
                ]
            )

        gated_result = next(part.content for part in tool_returns if part.tool_name == 'capability_search_tool')
        return ModelResponse(parts=[TextPart(content=f'final: {gated_result}')])

    agent = Agent(NoNativeToolSearchModel(model_fn), capabilities=[capability])

    result = await agent.run('find the gated tool')

    assert result.output == 'final: search-gated-result'
    assert seen_tool_names == snapshot(
        [
            ['search_tools'],
            ['capability_search_tool', 'search_tools'],
            ['capability_search_tool', 'search_tools'],
        ]
    )


async def test_tool_search_handles_capability_deferred_and_loaded_tools():
    """Deferred capability tools become visible as a unit after loading."""
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain
    def inherited_tool() -> str:  # pragma: no cover
        """Tool inheriting cap-level defer."""
        return 'inherited-result'

    @toolset.tool_plain(defer_loading=True)
    def also_deferred_tool() -> str:
        """Tool with its own defer flag — shadowed by cap-level defer."""
        return 'also-deferred-result'

    capability = Capability(
        id='example',
        description='Example capability.',
        defer_loading=True,
        toolsets=[toolset],
    )

    seen_tool_names: list[list[str]] = []

    def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        seen_tool_names.append([t.name for t in info.function_tools])
        tool_returns = [
            part
            for message in messages
            if isinstance(message, ModelRequest)
            for part in message.parts
            if isinstance(part, ToolReturnPart)
        ]

        if not any(part.tool_name == LOAD_CAPABILITY_TOOL_NAME for part in tool_returns):
            return ModelResponse(
                parts=[
                    ToolCallPart(
                        tool_name=LOAD_CAPABILITY_TOOL_NAME,
                        args={'id': 'example'},
                        tool_call_id='load-example',
                    )
                ]
            )

        if not any(part.tool_name == 'also_deferred_tool' for part in tool_returns):
            return ModelResponse(
                parts=[
                    ToolCallPart(
                        tool_name='also_deferred_tool',
                        args={},
                        tool_call_id='call-deferred',
                    )
                ]
            )

        gated_result = next(part.content for part in tool_returns if part.tool_name == 'also_deferred_tool')
        return ModelResponse(parts=[TextPart(content=f'final: {gated_result}')])

    agent = Agent(NoNativeToolSearchModel(model_fn), capabilities=[capability])

    result = await agent.run('use the special tool')

    assert result.output == 'final: also-deferred-result'
    assert seen_tool_names == snapshot(
        [
            ['load_capability', 'search_tools'],
            ['load_capability', 'inherited_tool', 'also_deferred_tool', 'search_tools'],
            ['load_capability', 'inherited_tool', 'also_deferred_tool', 'search_tools'],
        ]
    )


async def test_tool_search_ignores_malformed_loaded_capability_history():
    """Malformed `load_capability` results must not unlock capability-owned tools."""
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain
    def inherited_tool() -> str:  # pragma: no cover
        """Inherited deferred tool."""
        return 'inherited'

    capability = Capability(
        id='reports',
        description='Report tools.',
        defer_loading=True,
        toolsets=[toolset],
    )
    cap_toolset = CombinedCapability([capability]).get_toolset()
    assert isinstance(cap_toolset, AbstractToolset)
    cap_toolset = cast(AbstractToolset, cap_toolset)
    searchable: ToolSearchToolset = ToolSearchToolset(wrapped=cap_toolset)

    cases: dict[str, object] = {
        'not_a_dict': 'loaded reports',
        'non_string_instructions': {'instructions': ['bad']},
    }
    tool_defer_state: dict[str, list[tuple[str, bool]]] = {}

    for case_name, content in cases.items():
        messages: list[ModelMessage] = [
            ModelRequest(parts=[ToolReturnPart(tool_name=LOAD_CAPABILITY_TOOL_NAME, content=content)])
        ]
        ctx = _build_run_context(None, messages=messages, capabilities={'reports': capability})
        tools = await searchable.get_tools(ctx)
        tool_defer_state[case_name] = [(name, bool(t.tool_def.defer_loading)) for name, t in tools.items()]

    assert tool_defer_state == snapshot(
        {
            'not_a_dict': [('inherited_tool', True), ('search_tools', False)],
            'non_string_instructions': [('inherited_tool', True), ('search_tools', False)],
        }
    )


async def test_agent_auto_injects_tool_search_capability():
    """Test that agent auto-injects ToolSearch capability, with and without deferred tools."""
    agent_no_deferred = Agent('test')

    @agent_no_deferred.tool_plain
    def get_weather(city: str) -> str:  # pragma: no cover
        """Get the current weather for a city."""
        return f'Weather in {city}'

    leaves = collect_leaves(agent_no_deferred.root_capability)
    assert any(isinstance(leaf, ToolSearch) for leaf in leaves)

    agent_with_deferred = Agent('test')

    @agent_with_deferred.tool_plain
    def get_weather2(city: str) -> str:  # pragma: no cover
        """Get the current weather for a city."""
        return f'Weather in {city}'

    @agent_with_deferred.tool_plain(defer_loading=True)
    def calculate_mortgage(principal: float) -> str:  # pragma: no cover
        """Calculate mortgage payment."""
        return 'Calculated'

    leaves = collect_leaves(agent_with_deferred.root_capability)
    assert any(isinstance(leaf, ToolSearch) for leaf in leaves)


async def test_explicit_tool_search_not_duplicated():
    """Passing ToolSearch explicitly doesn't result in a second auto-injected one."""
    agent = Agent('test', capabilities=[ToolSearch()])

    @agent.tool_plain
    def get_weather(city: str) -> str:  # pragma: no cover
        """Get the current weather for a city."""
        return f'Weather in {city}'

    leaves = collect_leaves(agent.root_capability)
    tool_search_count = sum(1 for leaf in leaves if isinstance(leaf, ToolSearch))
    assert tool_search_count == 1


def test_tool_search_in_capability_registry():
    """ToolSearch is a public, spec-constructible capability."""

    assert ToolSearch.get_serialization_name() == 'ToolSearch'
    assert CAPABILITY_TYPES['ToolSearch'] is ToolSearch


async def test_tool_manager_with_tool_search_toolset_marks_corpus():
    """Every deferred tool appears once under its real name with
    `with_native='tool_search'`. Visible tools and `search_tools` round
    out the dispatch dict. `Model.prepare_request` filters per-model to decide what
    actually reaches the wire."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tool_manager = ToolManager(searchable)
    run_step_toolset = await tool_manager.for_run_step(ctx)

    managed_names = {t.name for t in run_step_toolset.tool_defs if t.with_native == 'tool_search'}
    assert managed_names == {'calculate_mortgage', 'stock_price', 'crypto_price'}

    local_names = [t.name for t in run_step_toolset.tool_defs if not t.with_native]
    assert 'get_weather' in local_names
    assert 'search_tools' in local_names

    # Undiscovered deferred tools are still dispatchable through the toolset under their
    # real name — the wire-side filtering in `prepare_request` decides whether the
    # model can see them, but `ToolManager` doesn't gatekeep dispatch on that.
    result = await run_step_toolset.handle_call(
        ToolCallPart(tool_name='calculate_mortgage', args={'principal': 100.0, 'rate': 5.0, 'years': 30})
    )
    assert 'Mortgage calculated' in str(result)

    # The local search_tools function is also dispatchable.
    result = await run_step_toolset.handle_call(ToolCallPart(tool_name='search_tools', args={'queries': ['mortgage']}))
    assert 'calculate_mortgage' in str(result)


async def test_tool_search_toolset_tool_with_none_description():
    """Test that tools with None description are handled correctly in search."""
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain(defer_loading=True)
    def no_desc_tool() -> str:  # pragma: no cover
        return 'no description'

    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]

    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['no_desc']}, ctx, search_tool)
    assert result == snapshot({'discovered_tools': [{'name': 'no_desc_tool'}]})


async def test_tool_search_toolset_multiple_searches_accumulate():
    """Discovery accumulates across search turns: tools surfaced in any past
    `search_tools` return have `defer_loading=False` on the next step, and
    not-yet-found ones keep `defer_loading=True`."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)

    messages: list[ModelMessage] = [
        ModelRequest(
            parts=[
                ToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'calculate_mortgage'}]},
                ),
            ]
        ),
        ModelRequest(
            parts=[
                ToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'stock_price'}]},
                ),
            ]
        ),
    ]
    ctx = _build_run_context(None, messages=messages, discovered_tool_names={'calculate_mortgage', 'stock_price'})

    tools = await searchable.get_tools(ctx)
    assert tools['calculate_mortgage'].tool_def.defer_loading is False
    assert tools['stock_price'].tool_def.defer_loading is False
    assert tools['crypto_price'].tool_def.defer_loading is True


async def test_function_toolset_all_deferred():
    """Test FunctionToolset with all tools having defer_loading=True."""
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain(defer_loading=True)
    def deferred_tool1() -> str:  # pragma: no cover
        """First deferred tool."""
        return 'result1'

    @toolset.tool_plain(defer_loading=True)
    def deferred_tool2() -> str:  # pragma: no cover
        """Second deferred tool."""
        return 'result2'

    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    tool_names = list(tools.keys())

    assert tool_names == snapshot(['deferred_tool1', 'deferred_tool2', 'search_tools'])


async def test_tool_search_toolset_reads_legacy_metadata_discovered_tools():
    """Pre-typed-content versions of this toolset wrote discovered tool names to
    `ToolReturnPart.metadata['discovered_tools']` instead of the typed
    `ToolSearchReturn` on `content`. Persisted histories from those versions
    must still surface their discoveries on resume; otherwise an agent reloaded from
    a saved transcript would re-emit `search_tools` and the user would see a
    duplicated discovery turn."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)

    messages: list[ModelMessage] = [
        ModelRequest(
            parts=[
                ToolReturnPart(
                    tool_name=_SEARCH_TOOLS_NAME,
                    content='legacy text return',
                    metadata={'discovered_tools': ['stock_price', 'crypto_price']},
                ),
            ]
        ),
        # Malformed legacy: not a list, ignored.
        ModelRequest(
            parts=[
                ToolReturnPart(
                    tool_name=_SEARCH_TOOLS_NAME,
                    content='another',
                    metadata={'discovered_tools': 'not a list'},
                ),
            ]
        ),
        # Malformed legacy: list with non-string entries; the string ones are still picked up.
        ModelRequest(
            parts=[
                ToolReturnPart(
                    tool_name=_SEARCH_TOOLS_NAME,
                    content='third',
                    metadata={'discovered_tools': [123, 'calculate_mortgage', None]},
                ),
            ]
        ),
    ]
    ctx = _build_run_context(None, messages=messages)

    tools = await searchable.get_tools(ctx)
    assert {'stock_price', 'crypto_price', 'calculate_mortgage'} <= set(tools)


def test_parse_discovered_tools_reads_legacy_metadata():
    """`parse_discovered_tools` reads the pre-typed-content legacy sideband directly.

    `get_tools` now consults `ctx.discovered_tool_names` rather than re-parsing history, so the
    legacy `metadata['discovered_tools']` reader is exercised through this classmethod (which
    run preparation calls to populate `discovered_tool_names`). A valid legacy list surfaces its
    names; a malformed one (wrong shape) is silently skipped via the `ValidationError` guard."""
    valid: list[ModelMessage] = [
        ModelRequest(
            parts=[
                ToolReturnPart(
                    tool_name=_SEARCH_TOOLS_NAME,
                    content='legacy text return',
                    metadata={'discovered_tools': ['stock_price', 'crypto_price']},
                ),
            ]
        ),
    ]
    assert parse_discovered_tools(valid) == {'stock_price', 'crypto_price'}

    # Malformed legacy metadata (`discovered_tools` is not a `list[str]`) trips
    # `_LEGACY_METADATA_TA.validate_python`, so `_collect_legacy` hits its `except: return`.
    malformed: list[ModelMessage] = [
        ModelRequest(
            parts=[
                ToolReturnPart(
                    tool_name=_SEARCH_TOOLS_NAME,
                    content='another',
                    metadata={'discovered_tools': 'not a list'},
                ),
            ]
        ),
    ]
    assert parse_discovered_tools(malformed) == set()


async def test_run_context_seeds_discovered_tool_names_from_history_before_first_step():
    """Pre-first-step hooks see tool-search discoveries replayed from message history."""
    seen_discovered_tool_names: list[set[str]] = []

    @dataclass
    class ObserveDiscoveredTools(AbstractCapability):
        async def before_run(self, ctx: RunContext) -> None:
            seen_discovered_tool_names.append(set(ctx.discovered_tool_names))

    history: list[ModelMessage] = [
        ModelRequest(
            parts=[
                ToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'calculate_mortgage'}]},
                ),
            ]
        )
    ]
    agent = Agent(TestModel(), capabilities=[ObserveDiscoveredTools()])

    await agent.run('hello', message_history=history)

    assert seen_discovered_tool_names == [{'calculate_mortgage'}]


async def test_deferred_loading_toolset_marks_all_tools():
    """`DeferredLoadingToolset` (with `tool_names=None`) flips `defer_loading=True`
    on every tool. After wrapping with `ToolSearchToolset`, all of them appear under
    their real name with `defer_loading=True` (visibility hidden until discovered).
    `search_tools` is the only directly-callable tool up front."""
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain
    def tool_a() -> str:  # pragma: no cover
        """Tool A."""
        return 'a'

    @toolset.tool_plain
    def tool_b() -> str:  # pragma: no cover
        """Tool B."""
        return 'b'

    deferred = toolset.defer_loading()
    searchable = ToolSearchToolset(wrapped=deferred)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    assert 'search_tools' in tools
    assert tools['tool_a'].tool_def.defer_loading is True
    assert tools['tool_b'].tool_def.defer_loading is True


async def test_deferred_loading_toolset_marks_specific_tools():
    """`DeferredLoadingToolset` with explicit names only flips `defer_loading=True`
    on the listed tools; others stay visible."""
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain
    def tool_a() -> str:  # pragma: no cover
        """Tool A."""
        return 'a'

    @toolset.tool_plain
    def tool_b() -> str:  # pragma: no cover
        """Tool B."""
        return 'b'

    deferred = toolset.defer_loading(['tool_b'])
    searchable = ToolSearchToolset(wrapped=deferred)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    assert 'search_tools' in tools
    assert tools['tool_a'].tool_def.defer_loading is False
    assert tools['tool_b'].tool_def.defer_loading is True


async def test_tool_search_toolset_marks_corpus_with_native():
    """Every deferred tool keeps its real name in the toolset output and carries
    `with_native='tool_search'` regardless of the current model — the adapter's
    `prepare_request` decides what reaches the wire so the toolset can't commit early
    (e.g. under `FallbackModel`)."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)

    managed = {name: tool.tool_def for name, tool in tools.items() if tool.tool_def.with_native}
    assert set(managed) == {'calculate_mortgage', 'stock_price', 'crypto_price'}
    for tool_def in managed.values():
        assert tool_def.with_native == 'tool_search'
        assert tool_def.defer_loading
    # The local fallback is still present — dropped by the adapter via `unless_native`.
    assert _SEARCH_TOOLS_NAME in tools


async def test_tool_search_toolset_dispatches_by_plain_name_via_tool_manager():
    """The provider calls a deferred tool by its plain name and `ToolManager`
    dispatches directly via the dict key (also the plain name)."""
    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)
    ctx = _build_run_context(None)

    tool_manager = ToolManager(searchable)
    run_step_toolset = await tool_manager.for_run_step(ctx)
    result = await run_step_toolset.handle_call(
        ToolCallPart(tool_name='calculate_mortgage', args={'principal': 100.0, 'rate': 5.0, 'years': 30})
    )
    assert 'Mortgage calculated' in str(result)


async def test_tool_search_toolset_custom_search_fn_is_used():
    """A custom `search_fn` replaces the default keyword-matching algorithm."""
    calls: list[Sequence[str]] = []

    def custom_search(ctx: RunContext, queries: Sequence[str], tools: Sequence[ToolDefinition]) -> list[str]:
        calls.append(queries)
        # Pick anything with 'price' in the name, regardless of query tokens.
        return [t.name for t in tools if 'price' in t.name]

    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset, search_fn=custom_search)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['anything']}, ctx, tools[_SEARCH_TOOLS_NAME])
    assert result == {
        'discovered_tools': [
            {'name': 'stock_price'},
            {'name': 'crypto_price'},
        ]
    }
    assert calls == [['anything']]


async def test_tool_search_toolset_custom_search_fn_still_marks_corpus():
    """A custom `search_fn` handles local discovery, but the toolset still flags every
    deferred tool with `with_native='tool_search'` — when the model supports
    native tool search (including provider-side custom callable modes like Anthropic's
    tool_reference mechanism or OpenAI's `execution='client'`), the adapter keeps them
    and applies `defer_loading` on the wire. Commitment to native-vs-local happens in
    `Model.prepare_request`, not here."""

    def custom_search(
        ctx: RunContext, queries: Sequence[str], tools: Sequence[ToolDefinition]
    ) -> list[str]:  # pragma: no cover
        return []

    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset, search_fn=custom_search)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)

    managed = [t.tool_def.name for t in tools.values() if t.tool_def.with_native == 'tool_search']
    assert set(managed) == {'calculate_mortgage', 'stock_price', 'crypto_price'}
    assert _SEARCH_TOOLS_NAME in tools


@pytest.mark.vcr
async def test_anthropic_native_tool_search_round_trip(allow_model_requests: None, anthropic_api_key: str) -> None:
    """End-to-end against live Anthropic: native BM25 server-side tool search
    populates `NativeToolCallPart` / `NativeToolReturnPart`, the model invokes
    the discovered deferred tool by its plain name, and the wire request carries
    `defer_loading: true` on the corpus tools and the `tool_search_tool_bm25`
    builtin.
    """
    pytest.importorskip('anthropic')

    model = AnthropicModel('claude-sonnet-4-5', provider=AnthropicProvider(api_key=anthropic_api_key))
    agent = Agent(model=model)

    @agent.tool_plain(defer_loading=True)
    def get_exchange_rate(from_currency: str, to_currency: str) -> str:
        """Look up the current exchange rate between two currencies."""
        return f'1 {from_currency} = 0.92 {to_currency}'

    @agent.tool_plain(defer_loading=True)
    def stock_lookup(symbol: str) -> str:  # pragma: no cover
        """Look up stock price by ticker symbol."""
        return f'Stock {symbol}: $150.00'

    result = await agent.run('What is the current USD to EUR exchange rate?')

    # Native server-side tool search auto-promotes to the typed
    # `NativeToolSearchCallPart` / `NativeToolSearchReturnPart` subclasses
    # (which still `isinstance`-match the base `NativeToolCallPart` /
    # `NativeToolReturnPart`).
    builtin_call_parts = [p for m in result.all_messages() for p in m.parts if isinstance(p, NativeToolCallPart)]
    builtin_return_parts = [p for m in result.all_messages() for p in m.parts if isinstance(p, NativeToolReturnPart)]
    assert builtin_call_parts and builtin_return_parts

    # The model's follow-up tool call for the discovered tool dispatches by its plain
    # name — the toolset exposes deferred tools as their regular variant on the native
    # path so the dispatch doesn't fall through to an "unknown tool" retry.
    rate_returns = [
        p
        for m in result.all_messages()
        for p in m.parts
        if isinstance(p, ToolReturnPart) and p.tool_name == 'get_exchange_rate'
    ]
    assert len(rate_returns) == 1
    assert rate_returns[0].content == '1 USD = 0.92 EUR'

    # Wire-level checks against the live cassette.
    cassette_path = (
        Path(__file__).parent / 'cassettes' / 'test_tool_search' / 'test_anthropic_native_tool_search_round_trip.yaml'
    )
    cassette = cast(dict[str, Any], yaml.safe_load(cassette_path.read_text(encoding='utf-8')))
    interactions = cast(list[dict[str, Any]], cassette['interactions'])

    # Initial request: deferred tools ship with `defer_loading: true`, and the BM25
    # builtin is registered alongside.
    first_request = cast(dict[str, Any], interactions[0]['request']['parsed_body'])
    deferred_names = {
        cast(str, t['name'])
        for t in cast(list[dict[str, Any]], first_request['tools'])
        if t.get('defer_loading') is True
    }
    assert deferred_names == {'get_exchange_rate', 'stock_lookup'}
    builtin_tool_types = {
        cast(str, t.get('type'))
        for t in cast(list[dict[str, Any]], first_request['tools'])
        if cast(str, t.get('type', '')).startswith('tool_search_tool_')
    }
    assert builtin_tool_types == {'tool_search_tool_bm25_20251119'}

    # Provisional beta header is rejected by the API — confirm we don't send it.
    assert 'tool-search-tool-2025-11-19' not in (first_request.get('betas') or [])

    # First response contains the server-side tool search round trip.
    first_response_blocks = cast(list[dict[str, Any]], interactions[0]['response']['parsed_body']['content'])
    assert any(
        b.get('type') == 'server_tool_use' and b.get('name') == 'tool_search_tool_bm25' for b in first_response_blocks
    )
    assert any(b.get('type') == 'tool_search_tool_result' for b in first_response_blocks)


@pytest.mark.vcr
async def test_anthropic_custom_callable_round_trip(allow_model_requests: None, anthropic_api_key: str) -> None:
    """End-to-end: a custom callable `ToolSearch` strategy runs locally but still
    surfaces natively on Anthropic — deferred tools ship with `defer_loading: true`,
    the model invokes the regular `search_tools` function tool, and our
    `tool_result` is formatted as `tool_reference` blocks so the discovered tool
    gets unlocked for the next turn."""
    pytest.importorskip('anthropic')

    def match_exchange_rate(ctx: RunContext, queries: Sequence[str], tools: Sequence[ToolDefinition]) -> list[str]:
        # Deterministic: always point the model at `get_exchange_rate` so the cassette
        # replay doesn't depend on the exact keywords the model picks.
        return ['get_exchange_rate']

    model = AnthropicModel('claude-sonnet-4-5', provider=AnthropicProvider(api_key=anthropic_api_key))
    agent = Agent(
        model=model,
        capabilities=[ToolSearch(strategy=match_exchange_rate)],
    )

    @agent.tool_plain
    def get_weather(city: str) -> str:  # pragma: no cover
        """Get the current weather for a city."""
        return f'Weather in {city} is sunny.'

    @agent.tool_plain(defer_loading=True)
    def get_exchange_rate(from_currency: str, to_currency: str) -> str:
        """Look up the current exchange rate between two currencies."""
        return f'1 {from_currency} = 0.92 {to_currency}'

    @agent.tool_plain(defer_loading=True)
    def stock_lookup(symbol: str) -> str:  # pragma: no cover
        """Look up stock price by ticker symbol."""
        return f'Stock {symbol}: $150.00'

    result = await agent.run('What is the USD to EUR exchange rate?')

    # The full sequence: user prompt -> model asks `search_tools` -> our local callable
    # returns discovered tool names -> model follows up with the discovered tool ->
    # we run it -> model replies with final text.
    part_shape = [
        [(type(part).__name__, getattr(part, 'tool_name', None)) for part in msg.parts] for msg in result.all_messages()
    ]
    assert part_shape == snapshot(part_shape)

    # The deferred tool dispatched successfully end-to-end.
    rate_returns = [
        part
        for msg in result.all_messages()
        for part in msg.parts
        if isinstance(part, ToolReturnPart) and part.tool_name == 'get_exchange_rate'
    ]
    assert len(rate_returns) == 1
    assert rate_returns[0].content == '1 USD = 0.92 EUR'

    # Wire-level checks against the cassette: the deferred corpus ships with
    # `defer_loading: true`, the model's `search_tools` call appears in the response,
    # and our tool result is formatted as `tool_reference` blocks (not plain text).

    cassette_path = (
        Path(__file__).parent / 'cassettes' / 'test_tool_search' / 'test_anthropic_custom_callable_round_trip.yaml'
    )
    cassette = cast(dict[str, Any], yaml.safe_load(cassette_path.read_text(encoding='utf-8')))
    interactions = cast(list[dict[str, Any]], cassette['interactions'])

    first_request_tools = cast(list[dict[str, Any]], interactions[0]['request']['parsed_body']['tools'])
    deferred_names = {t['name'] for t in first_request_tools if t.get('defer_loading') is True}
    assert deferred_names == {'get_exchange_rate', 'stock_lookup'}

    first_response_blocks = cast(list[dict[str, Any]], interactions[0]['response']['parsed_body']['content'])
    assert any(b['type'] == 'tool_use' and b['name'] == 'search_tools' for b in first_response_blocks)

    second_request_messages = cast(list[dict[str, Any]], interactions[1]['request']['parsed_body']['messages'])
    tool_result_blocks: list[dict[str, Any]] = [
        block
        for msg in second_request_messages
        if msg['role'] == 'user' and isinstance(msg.get('content'), list)
        for block in cast(list[dict[str, Any]], msg['content'])
        if isinstance(block, dict) and block.get('type') == 'tool_result'
    ]
    assert tool_result_blocks, 'expected at least one tool_result block in the follow-up turn'
    tool_reference_names: set[str] = {
        cast(str, inner['tool_name'])
        for block in tool_result_blocks
        for inner in cast(list[dict[str, Any]], block.get('content', []))
        if isinstance(inner, dict) and inner.get('type') == 'tool_reference'
    }
    assert tool_reference_names == {'get_exchange_rate'}


@pytest.mark.vcr
async def test_anthropic_promotes_local_search_history_round_trip(
    allow_model_requests: None, anthropic_api_key: str
) -> None:
    """End-to-end against live Anthropic: a turn with local-shape `ToolSearch*Part`
    history (from a prior cross-provider turn — e.g. on Google) runs cleanly on
    Anthropic. The adapter promotes the local-shape return into a `tool_result` with
    `tool_reference` content so Anthropic unlocks the discovered tool's schema, and
    the model dispatches the discovered tool directly without issuing a fresh
    `tool_search_tool_*` call.
    """
    pytest.importorskip('anthropic')

    model = AnthropicModel('claude-sonnet-4-5', provider=AnthropicProvider(api_key=anthropic_api_key))
    agent = Agent(model=model, capabilities=[ToolSearch()])

    @agent.tool_plain(defer_loading=True)
    def get_exchange_rate(from_currency: str, to_currency: str) -> str:
        """Look up the current exchange rate between two currencies."""
        return f'1 {from_currency} = 0.92 {to_currency}'

    @agent.tool_plain(defer_loading=True)
    def stock_lookup(symbol: str) -> str:  # pragma: no cover
        """Look up stock price by ticker symbol."""
        return f'Stock {symbol}: $150.00'

    # Synthetic history: a prior turn on a non-supporting provider (Google etc.)
    # discovered `get_exchange_rate` via the local `search_tools` function tool.
    # Carries the local-shape typed parts on a `ToolSearchReturnPart` (sub of
    # `ToolReturnPart`) — exactly what the toolset would emit on the local path.
    prior_history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='I might want to look up exchange rates later.')]),
        ModelResponse(
            parts=[
                ToolSearchCallPart(args={'queries': ['exchange rate']}, tool_call_id='loc_search_1'),
            ],
        ),
        ModelRequest(
            parts=[
                ToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'get_exchange_rate'}]},
                    tool_call_id='loc_search_1',
                ),
            ],
        ),
    ]

    result = await agent.run('What is the USD to EUR exchange rate?', message_history=prior_history)

    # The model uses the discovered tool directly — no fresh `tool_search_tool_*` call
    # was needed because the prior local-shape return got promoted to native shape on
    # the wire, unlocking `get_exchange_rate` server-side.
    rate_returns = [
        part
        for msg in result.all_messages()
        for part in msg.parts
        if isinstance(part, ToolReturnPart) and part.tool_name == 'get_exchange_rate'
    ]
    assert len(rate_returns) == 1
    assert rate_returns[0].content == '1 USD = 0.92 EUR'

    # No fresh native tool_search exchange after the synthetic history.
    fresh_native_search_calls = [
        part for msg in result.all_messages() for part in msg.parts if isinstance(part, NativeToolSearchCallPart)
    ]
    assert fresh_native_search_calls == []

    # Wire-level: cassette confirms the request to Anthropic carried the prior
    # local-shape return as a `tool_result` with `tool_reference` content (NOT a
    # stringified JSON of the discoveries).
    cassette_path = (
        Path(__file__).parent
        / 'cassettes'
        / 'test_tool_search'
        / 'test_anthropic_promotes_local_search_history_round_trip.yaml'
    )
    cassette = cast(dict[str, Any], yaml.safe_load(cassette_path.read_text(encoding='utf-8')))
    interactions = cast(list[dict[str, Any]], cassette['interactions'])

    first_request_messages = cast(list[dict[str, Any]], interactions[0]['request']['parsed_body']['messages'])
    tool_result_contents: list[Any] = [
        block.get('content')
        for msg in first_request_messages
        if msg.get('role') == 'user' and isinstance(msg.get('content'), list)
        for block in cast(list[dict[str, Any]], msg['content'])
        if isinstance(block, dict) and block.get('type') == 'tool_result'
    ]
    # The `tool_reference` array shape proves the promotion fired.
    promoted_names = {
        cast(str, inner.get('tool_name'))
        for content in tool_result_contents
        if isinstance(content, list)
        for inner in cast(list[dict[str, Any]], content)
        if isinstance(inner, dict) and inner.get('type') == 'tool_reference'
    }
    assert promoted_names == {'get_exchange_rate'}


@pytest.mark.vcr
async def test_openai_promotes_local_search_history_round_trip(
    allow_model_requests: None, openai_api_key: str, vcr: Any
) -> None:
    """End-to-end against live OpenAI: a turn with local-shape `ToolSearch*Part`
    history runs cleanly on OpenAI Responses. The adapter promotes the local-shape
    pair into `tool_search_call` + `tool_search_output` items with
    `execution='client'`, and the model dispatches the discovered tool directly.
    """
    pytest.importorskip('openai')

    model = OpenAIResponsesModel('gpt-5.4-mini', provider=OpenAIProvider(api_key=openai_api_key))
    agent = Agent(model=model, capabilities=[ToolSearch()])

    @agent.tool_plain(defer_loading=True)
    def get_exchange_rate(from_currency: str, to_currency: str) -> str:
        """Look up the current exchange rate between two currencies."""
        return f'1 {from_currency} = 0.92 {to_currency}'

    @agent.tool_plain(defer_loading=True)
    def stock_lookup(symbol: str) -> str:  # pragma: no cover
        """Look up stock price by ticker symbol."""
        return f'Stock {symbol}: $150.00'

    prior_history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='I might want to look up exchange rates later.')]),
        ModelResponse(
            parts=[
                ToolSearchCallPart(args={'queries': ['exchange rate']}, tool_call_id='loc_search_1'),
            ],
        ),
        ModelRequest(
            parts=[
                ToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'get_exchange_rate'}]},
                    tool_call_id='loc_search_1',
                ),
            ],
        ),
    ]

    result = await agent.run('What is the USD to EUR exchange rate?', message_history=prior_history)

    rate_returns = [
        part
        for msg in result.all_messages()
        for part in msg.parts
        if isinstance(part, ToolReturnPart) and part.tool_name == 'get_exchange_rate'
    ]
    assert len(rate_returns) == 1
    assert rate_returns[0].content == '1 USD = 0.92 EUR'

    request_inputs = [json.loads(request.body).get('input', []) for request in vcr.requests]

    # Wire-level: cassette confirms the local-shape pair got promoted to
    # `tool_search_call` + `tool_search_output` items with `execution='client'`.
    promoted_calls = [item for item in request_inputs[0] if item.get('type') == 'tool_search_call']
    promoted_outputs = [item for item in request_inputs[0] if item.get('type') == 'tool_search_output']
    assert promoted_calls, 'expected the local-shape call to be promoted to tool_search_call'
    assert promoted_outputs, 'expected the local-shape return to be promoted to tool_search_output'
    assert all(item.get('execution') == 'client' for item in promoted_calls)
    assert all(item.get('execution') == 'client' for item in promoted_outputs)
    promoted_tool_names = {t.get('name') for output in promoted_outputs for t in output.get('tools', [])}
    assert 'get_exchange_rate' in promoted_tool_names

    # Wire-level: once the discovered tool is dispatched, its replayed `function_call` must carry a
    # `namespace` — OpenAI rejects a tool-search-discovered call without one. Guards the namespace
    # emission in `_map_messages` (a regression would only surface as a live 422 otherwise).
    replayed_calls = [
        item
        for request_input in request_inputs
        for item in request_input
        if item.get('type') == 'function_call' and item.get('name') == 'get_exchange_rate'
    ]
    assert replayed_calls, 'expected the discovered tool call to be replayed to OpenAI'
    assert any(item.get('namespace') == 'get_exchange_rate' for item in replayed_calls)


@pytest.mark.vcr
async def test_anthropic_native_tool_search_regex_strategy(allow_model_requests: None, anthropic_api_key: str) -> None:
    """`ToolSearch(strategy='regex')` registers the regex variant of Anthropic's
    native tool search tool rather than BM25, and the live API accepts the request.
    """
    pytest.importorskip('anthropic')

    model = AnthropicModel('claude-sonnet-4-5', provider=AnthropicProvider(api_key=anthropic_api_key))
    agent = Agent(model=model, capabilities=[ToolSearch(strategy='regex')])

    @agent.tool_plain(defer_loading=True)
    def get_exchange_rate(from_currency: str, to_currency: str) -> str:  # pragma: no cover
        """Look up the current exchange rate between two currencies."""
        return f'1 {from_currency} = 0.92 {to_currency}'

    await agent.run('hi, just say hello')

    # The live request carries the regex variant — the mock-only assertion here would
    # only validate that we generate the correct parameter shape, not that Anthropic
    # accepts it.
    cassette_path = (
        Path(__file__).parent
        / 'cassettes'
        / 'test_tool_search'
        / 'test_anthropic_native_tool_search_regex_strategy.yaml'
    )
    cassette = cast(dict[str, Any], yaml.safe_load(cassette_path.read_text(encoding='utf-8')))
    interactions = cast(list[dict[str, Any]], cassette['interactions'])
    request_body = cast(dict[str, Any], interactions[0]['request']['parsed_body'])
    tool_types = [
        cast(str, t.get('type')) for t in cast(list[dict[str, Any]], request_body['tools']) if isinstance(t, dict)
    ]
    assert 'tool_search_tool_regex_20251119' in tool_types
    assert 'tool_search_tool_bm25_20251119' not in tool_types
    # Live API returned 2xx — the absence of a 4xx is the strongest signal that the
    # request shape (no beta header, regex variant) is accepted.
    assert interactions[0]['response']['status']['code'] == 200


async def test_anthropic_regex_strategy_replay_preserves_variant(allow_model_requests: None):
    """History replay must re-emit the exact server-tool variant the provider used —
    downgrading `tool_search_tool_regex` to `tool_search_tool_bm25` on a resend would
    silently run a different algorithm than the earlier turn."""
    pytest.importorskip('anthropic')

    # Provider-side call used the regex variant; the adapter must round-trip that choice.
    # Anthropic's regex variant emits `pattern` (not `query`) in the wire input.
    regex_block = BetaServerToolUseBlock(
        id='srv_r',
        name='tool_search_tool_regex',
        input={'pattern': 'weather.*'},
        type='server_tool_use',
        caller=BetaDirectCaller(type='direct'),
    )
    call_part = _map_server_tool_use_block(regex_block, 'anthropic')
    assert isinstance(call_part, NativeToolCallPart)
    assert call_part.provider_details == {'strategy': 'regex'}
    # Cross-provider canonical shape collects the regex into the `queries` slot.
    assert call_part.args == snapshot({'queries': ['weather.*']})

    # On replay, the adapter should emit `tool_search_tool_regex` (not bm25).
    response = completion_message(
        [BetaTextBlock(text='done', type='text')],
        BetaUsage(input_tokens=5, output_tokens=5),
    )
    mock_client = MockAnthropic.create_mock(response)
    model = AnthropicModel('claude-sonnet-4-5', provider=AnthropicProvider(anthropic_client=mock_client))
    agent = Agent(model=model, capabilities=[ToolSearch(strategy='regex')])

    @agent.tool_plain(defer_loading=True)
    def get_weather(city: str) -> str:  # pragma: no cover
        return f'Weather in {city}.'

    history: list[ModelMessage] = [
        ModelRequest.user_text_prompt('look it up'),
        ModelResponse(
            parts=[
                call_part,
                NativeToolSearchReturnPart(
                    provider_name='anthropic',
                    tool_call_id='srv_r',
                    content={'discovered_tools': [{'name': 'get_weather'}]},
                ),
            ],
            provider_name='anthropic',
        ),
    ]
    await agent.run('and again', message_history=history)
    kwargs = get_mock_chat_completion_kwargs(mock_client)[0]
    # Inspect the replayed Anthropic request. Content blocks are dicts on the request
    # path (params); flatten via comprehension so each replayed call's `name` shows up
    # in `names`.
    blocks = [
        cast('dict[str, Any]', block) for msg in kwargs['messages'] for block in cast('list[Any]', msg['content'])
    ]
    server_blocks = [block for block in blocks if block.get('type') == 'server_tool_use']
    names = [block['name'] for block in server_blocks]
    assert 'tool_search_tool_regex' in names
    assert 'tool_search_tool_bm25' not in names
    # Regex variant must replay with `pattern` (not `query`) — Anthropic 400s otherwise.
    regex_inputs = [block['input'] for block in server_blocks if block['name'] == 'tool_search_tool_regex']
    assert regex_inputs == snapshot([{'pattern': 'weather.*'}])


def test_collect_orphan_tool_search_call_ids_pairs_across_responses() -> None:
    """An orphan is a `NativeToolSearchCallPart` with no matching `NativeToolSearchReturnPart`
    *anywhere* in history. Anthropic sometimes delivers the return in a *later* `ModelResponse`
    (deferred-result behavior on the direct API), so the pairing check must span turns."""
    pytest.importorskip('anthropic')

    history: list[ModelMessage] = [
        ModelRequest.user_text_prompt('do the thing'),
        # Turn 1: orphan call (paired with a client `ToolCallPart` that ate the turn)
        ModelResponse(
            parts=[
                NativeToolSearchCallPart(args={'queries': ['pay.*']}, tool_call_id='srv_orphan'),
                ToolCallPart(tool_name='send_status', args={'message': 'ok'}, tool_call_id='cl_1'),
            ],
        ),
        ModelRequest(parts=[ToolReturnPart(tool_name='send_status', content='ok', tool_call_id='cl_1')]),
        # Turn 2: deferred-result call+return *and* a fresh paired exchange
        ModelResponse(
            parts=[
                # Anthropic delivers the previous turn's missing search result here.
                NativeToolSearchReturnPart(content={'discovered_tools': []}, tool_call_id='srv_paired'),
                # ...along with a fresh search round.
                NativeToolSearchCallPart(args={'queries': ['weather.*']}, tool_call_id='srv_paired_2'),
                NativeToolSearchReturnPart(content={'discovered_tools': []}, tool_call_id='srv_paired_2'),
            ],
        ),
    ]
    # `srv_orphan` has no matching return anywhere; `srv_paired_2` is paired in the same response.
    # `srv_paired` shows up only as a return — that's not an orphan call, so it isn't reported.
    assert _collect_orphan_tool_search_call_ids(history) == {'srv_orphan'}


async def test_anthropic_drops_orphaned_tool_search_call_on_replay(allow_model_requests: None) -> None:
    """Anthropic occasionally emits a `tool_search_tool_*` server tool use alongside a client
    `tool_use` and ends the turn without delivering the corresponding result block (see
    anthropics/anthropic-sdk-python#1325). Bedrock then 400s on the next request:
    `tool use ... was found without a corresponding tool_search_tool_*_tool_result block`.
    The adapter must drop unpaired tool-search calls from the wire payload. Reported by
    @kclisp on PR #5143.
    """
    pytest.importorskip('anthropic')

    response = completion_message(
        [BetaTextBlock(text='ok', type='text')],
        BetaUsage(input_tokens=5, output_tokens=5),
    )
    mock_client = MockAnthropic.create_mock(response)
    model = AnthropicModel('claude-sonnet-4-5', provider=AnthropicProvider(anthropic_client=mock_client))
    agent = Agent(model=model, capabilities=[ToolSearch()])

    @agent.tool_plain
    def send_status(message: str) -> str:  # pragma: no cover
        return 'ok'

    @agent.tool_plain(defer_loading=True)
    def pay_rent() -> str:  # pragma: no cover
        return 'paid'

    history: list[ModelMessage] = [
        ModelRequest.user_text_prompt('pay rent and send status'),
        ModelResponse(
            parts=[
                # Orphan: server tool search emitted in parallel with a client tool, no result delivered.
                NativeToolSearchCallPart(
                    provider_name='anthropic',
                    args={'queries': ['pay.*']},
                    tool_call_id='srv_orphan',
                    provider_details={'strategy': 'regex'},
                ),
                ToolCallPart(tool_name='send_status', args={'message': 'looking'}, tool_call_id='cl_1'),
            ],
            provider_name='anthropic',
        ),
        ModelRequest(parts=[ToolReturnPart(tool_name='send_status', content='ok', tool_call_id='cl_1')]),
    ]
    await agent.run('continue', message_history=history)
    kwargs = get_mock_chat_completion_kwargs(mock_client)[0]
    blocks = [
        cast('dict[str, Any]', block) for msg in kwargs['messages'] for block in cast('list[Any]', msg['content'])
    ]
    server_tool_block_ids = [block.get('id') for block in blocks if block.get('type') == 'server_tool_use']
    assert 'srv_orphan' not in server_tool_block_ids


async def test_anthropic_cache_tool_definitions_skips_deferred_tools(allow_model_requests: None) -> None:
    """`anthropic_cache_tool_definitions=True` must apply `cache_control` to the last
    *non-deferred* tool. Anthropic rejects requests with `cache_control` and
    `defer_loading=True` on the same tool: `Tools with defer_loading cannot use prompt
    caching`. Reported by @kclisp on PR #5143.
    """
    pytest.importorskip('anthropic')

    response = completion_message(
        [BetaTextBlock(text='ok', type='text')],
        BetaUsage(input_tokens=5, output_tokens=5),
    )
    mock_client = MockAnthropic.create_mock(response)
    model = AnthropicModel('claude-sonnet-4-5', provider=AnthropicProvider(anthropic_client=mock_client))
    agent = Agent(
        model=model,
        capabilities=[ToolSearch()],
        model_settings=AnthropicModelSettings(anthropic_cache_tool_definitions=True),
    )

    @agent.tool_plain
    def visible_tool() -> str:  # pragma: no cover
        return 'visible'

    @agent.tool_plain(defer_loading=True)
    def deferred_tool() -> str:  # pragma: no cover
        return 'deferred'

    await agent.run('hi')

    kwargs = get_mock_chat_completion_kwargs(mock_client)[0]
    tools = cast('list[dict[str, Any]]', kwargs['tools'])
    by_name = {tool['name']: tool for tool in tools}
    # The deferred tool must NOT have `cache_control` — pairing it with `defer_loading`
    # is what Anthropic rejects.
    assert 'cache_control' not in by_name['deferred_tool']
    assert by_name['deferred_tool'].get('defer_loading') is True
    # The last non-deferred tool gets the cache breakpoint.
    assert by_name['visible_tool']['cache_control'] == snapshot({'type': 'ephemeral', 'ttl': '5m'})


async def test_anthropic_cache_tool_definitions_skips_when_all_tools_deferred(allow_model_requests: None) -> None:
    """When *every* tool is deferred, there's nothing in the cacheable prompt prefix to
    attach `cache_control` to. The loop must fall through without breaking — applying
    `cache_control` to any deferred tool would 400 the request.
    """
    pytest.importorskip('anthropic')

    response = completion_message(
        [BetaTextBlock(text='ok', type='text')],
        BetaUsage(input_tokens=5, output_tokens=5),
    )
    mock_client = MockAnthropic.create_mock(response)
    model = AnthropicModel('claude-sonnet-4-5', provider=AnthropicProvider(anthropic_client=mock_client))
    agent = Agent(
        model=model,
        capabilities=[ToolSearch()],
        model_settings=AnthropicModelSettings(anthropic_cache_tool_definitions=True),
    )

    @agent.tool_plain(defer_loading=True)
    def deferred_one() -> str:  # pragma: no cover
        return 'one'

    @agent.tool_plain(defer_loading=True)
    def deferred_two() -> str:  # pragma: no cover
        return 'two'

    await agent.run('hi')

    kwargs = get_mock_chat_completion_kwargs(mock_client)[0]
    tools = cast('list[dict[str, Any]]', kwargs['tools'])
    function_tools = [tool for tool in tools if 'input_schema' in tool]
    # No tool ends up with `cache_control` — pairing any deferred tool with it 400s.
    for tool in function_tools:
        assert 'cache_control' not in tool


async def test_openai_rejects_anthropic_named_strategy(allow_model_requests: None):
    """OpenAI Responses has no bm25/regex concept — using one must error loudly rather
    than silently falling through to OpenAI's default server-side tool search."""
    pytest.importorskip('openai')

    mock_client = MockOpenAIResponses.create_mock(response_message([]))
    model = OpenAIResponsesModel('gpt-5.4', provider=OpenAIProvider(openai_client=mock_client))
    agent = Agent(model=model, capabilities=[ToolSearch(strategy='bm25')])

    @agent.tool_plain(defer_loading=True)
    def get_weather(city: str) -> str:  # pragma: no cover
        return f'Weather in {city}.'

    with pytest.raises(UserError, match='Anthropic-native'):
        await agent.run('what should I wear?')


async def test_openai_client_tool_search_maps_to_local_search_call():
    """Client-executed `tool_search_call` items map to a regular `ToolCallPart` against
    the local `search_tools` function. Replay later detects the OpenAI native variant
    via the current request's builtin configuration plus a `provider_name` match."""
    pytest.importorskip('openai')

    call = ResponseToolSearchCall(
        id='ts_1',
        arguments={'queries': ['exchange rate']},
        call_id='call_1',
        execution='client',
        status='completed',
        type='tool_search_call',
    )
    part = _map_client_tool_search_call(call, 'azure')
    assert part.tool_name == _SEARCH_TOOLS_NAME
    # Provider name flows through from the model — important for OpenAI-compatible
    # providers (Azure, gateways) where `self.system` differs from `'openai'`.
    assert part.provider_name == 'azure'
    # No envelope marker any more: replay derives intent from the current request's
    # builtin configuration + a `provider_name` match against `self.system`.
    assert part.provider_details is None


async def test_openai_deferred_capability_tool_reveal_uses_client_tool_search(allow_model_requests: None):
    """A `load_capability` reveal synthesizes tool-search history for newly visible tools.

    OpenAI uses client-executed `tool_search` while deferred capability-owned tools are
    in the tool-search corpus, so the same registration works for initial discovery and
    later replay of the synthetic history.
    """
    pytest.importorskip('openai')

    refunds_toolset = FunctionToolset()

    @refunds_toolset.tool_plain
    def lookup_refund_policy(order_id: str) -> str:  # pragma: no cover
        """Look up the refund policy for an order."""
        return f'{order_id}: refund allowed'

    capability = Capability(
        id='refunds',
        description='Refund policy tools.',
        instructions='Use the refund policy tool before answering refund questions.',
        defer_loading=True,
        toolsets=[refunds_toolset],
    )
    responses = [
        response_message(
            [
                ResponseFunctionToolCall(
                    id='fc_load',
                    arguments='{"id":"refunds"}',
                    call_id='call_load',
                    name=LOAD_CAPABILITY_TOOL_NAME,
                    status='completed',
                    type='function_call',
                )
            ]
        ),
        response_message(
            [
                ResponseOutputMessage(
                    id='msg_done',
                    content=[ResponseOutputText(text='Loaded.', type='output_text', annotations=[])],
                    role='assistant',
                    status='completed',
                    type='message',
                )
            ]
        ),
    ]
    mock_client = MockOpenAIResponses.create_mock(responses)
    model = OpenAIResponsesModel('gpt-5.4', provider=OpenAIProvider(openai_client=mock_client))
    agent: Agent[None, str] = Agent(model=model, capabilities=[capability])

    result = await agent.run('Can I get a refund on order-123?')

    assert result.output == 'Loaded.'
    assert any(
        isinstance(part, ToolSearchReturnPart)
        and [match['name'] for match in part.discovered_tools] == ['lookup_refund_policy']
        for message in result.all_messages()
        for part in message.parts
    )
    [first_request, second_request] = get_mock_responses_kwargs(mock_client)

    for request in (first_request, second_request):
        [tool_search] = [tool for tool in cast(list[dict[str, Any]], request['tools']) if tool['type'] == 'tool_search']
        assert tool_search['execution'] == 'client'
        assert cast(dict[str, Any], tool_search['parameters'])['required'] == ['queries']

    second_input = cast(list[dict[str, Any]], second_request['input'])
    replay_calls = [item for item in second_input if item.get('type') == 'tool_search_call']
    replay_outputs = [item for item in second_input if item.get('type') == 'tool_search_output']
    assert replay_calls and all(item.get('execution') == 'client' for item in replay_calls)
    assert replay_outputs and all(item.get('execution') == 'client' for item in replay_outputs)


async def test_openai_discovered_tool_without_native_tool_search_omits_defer_loading(
    allow_model_requests: None,
):
    """A tool-search corpus member discovered in a prior turn must not carry the wire-side
    `defer_loading` flag on a model without native `tool_search` (e.g. `gpt-5.2`).

    OpenAI's `defer_loading` only travels alongside a native `tool_search` tool; without one the
    provider rejects a lone `defer_loading` (#5938). Once discovered, the corpus member stays
    callable as a plain function tool but must shed its `with_native='tool_search'` marker, so the
    adapter (which derives `defer_loading` purely from `with_native`) stops stamping it.

    This is a unit test, not VCR: the cassette matcher keys only on method and path, so a request
    that regained a stale `defer_loading` (or an over-eager native-tool swap) would still match the
    existing recording and pass green. Only a direct assertion on the emitted payload pins the wire
    invariant the fix is responsible for. The end-to-end deferred-capability flow is covered by
    `test_openai_deferred_capability_runs_on_model_without_native_tool_search`.
    """
    pytest.importorskip('openai')

    final = response_message(
        [
            ResponseOutputMessage(
                id='msg',
                content=[ResponseOutputText(text='Sunny.', type='output_text', annotations=[])],
                role='assistant',
                status='completed',
                type='message',
            )
        ]
    )
    mock_client = MockOpenAIResponses.create_mock(final)
    model = OpenAIResponsesModel('gpt-5.2', provider=OpenAIProvider(openai_client=mock_client))
    agent: Agent[None, str] = Agent(model=model, capabilities=[ToolSearch()])

    @agent.tool_plain(defer_loading=True)
    def get_weather(city: str) -> str:  # pragma: no cover
        return f'Weather in {city}.'

    # `get_weather` was discovered last turn, so it now rides along as a callable tool
    # (`defer_loading=False`, but `with_native='tool_search'` still set).
    history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='I might want the weather later.')]),
        ModelResponse(parts=[ToolSearchCallPart(args={'queries': ['weather']}, tool_call_id='loc_1')]),
        ModelRequest(
            parts=[
                ToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'get_weather'}]},
                    tool_call_id='loc_1',
                )
            ]
        ),
    ]

    await agent.run('Weather in Paris?', message_history=history)

    [request] = get_mock_responses_kwargs(mock_client)
    request_tools = cast(list[dict[str, Any]], request['tools'])
    assert not any(tool['type'] == 'tool_search' for tool in request_tools)
    [weather_tool] = [tool for tool in request_tools if tool.get('name') == 'get_weather']
    assert 'defer_loading' not in weather_tool


@pytest.mark.vcr
async def test_openai_deferred_capability_runs_on_model_without_native_tool_search(
    allow_model_requests: None, openai_api_key: str
) -> None:
    """Loading a deferred `Capability` and calling its tool must complete on a model
    *without* native `tool_search` (#5938).

    `gpt-5.2` predates OpenAI's native `tool_search`, so search falls back to local and no
    `tool_search` builtin is on the wire. The tool revealed by `load_capability` still carries
    `with_native='tool_search'`, which the base-class filter sheds so no adapter emits a
    wire-side `defer_loading` flag with no native `tool_search` tool to pair it with — which
    OpenAI rejects. The invariant is simply that the run completes and `bar` returns.

    The wire-payload shape (no `defer_loading`, no `tool_search` on the wire) is pinned directly by
    `test_openai_discovered_tool_without_native_tool_search_omits_defer_loading`, since the cassette
    matcher isn't body-sensitive and wouldn't catch a regression here on its own.

    This is the tool-search-*unsupported* half of the matrix. OpenAI is the only provider it
    can be recorded against: every non-deprecated Anthropic (and Google) model supports native
    tool search, so there is no equivalent no-native-`tool_search` model to exercise.
    """
    foo = Capability[None](id='foo', description='Use this capability when the user asks for foo.', defer_loading=True)

    @foo.tool_plain
    def bar(x: int) -> int:
        """Return x."""
        return x

    model = OpenAIResponsesModel('gpt-5.2', provider=OpenAIProvider(api_key=openai_api_key))
    agent: Agent[None, str] = Agent(
        model,
        deps_type=type(None),
        instructions="First load capability id 'foo', then call bar.",
        capabilities=[foo],
    )

    result = await agent.run('Use foo with x=1.')

    # The capability loaded and `bar` returned, so the follow-up request carrying the revealed
    # `bar` (the one that used to 400) was accepted and the run finished.
    assert any(isinstance(p, LoadCapabilityReturnPart) for m in result.all_messages() for p in m.parts)
    bar_returns = [
        p for m in result.all_messages() for p in m.parts if isinstance(p, ToolReturnPart) and p.tool_name == 'bar'
    ]
    assert [p.content for p in bar_returns] == [1]


async def test_cross_provider_history_replay_anthropic_to_openai(allow_model_requests: None):
    """A model switch between turns (Anthropic → OpenAI) should replay cleanly: the
    provider-specific Builtin* tool search parts are skipped by the mismatched provider,
    and the agent can still dispatch already-discovered tools by name. This is the
    canonical FallbackModel-style scenario the design calls for."""
    pytest.importorskip('openai')
    pytest.importorskip('anthropic')

    # Prior turn: Anthropic ran a native BM25 search and discovered `get_weather`.
    prior: list[ModelMessage] = [
        ModelRequest.user_text_prompt('weather please'),
        ModelResponse(
            parts=[
                NativeToolCallPart(
                    provider_name='anthropic',
                    tool_name='tool_search',
                    tool_call_id='srv_a',
                    args={'query': 'weather'},
                    provider_details={'strategy': 'bm25'},
                ),
                NativeToolReturnPart(
                    provider_name='anthropic',
                    tool_name='tool_search',
                    tool_call_id='srv_a',
                    content={'discovered_tools': [{'name': 'get_weather'}]},
                ),
            ],
            provider_name='anthropic',
        ),
    ]

    # Switch to OpenAI for the follow-up. The Anthropic builtin parts should be silently
    # skipped (`provider_name` mismatch). `get_weather` was discovered in the prior turn,
    # so `parse_discovered_tools` picks it up and exposes the regular
    # variant on the new provider — the model can call it directly.
    followup = response_message(
        [
            ResponseOutputMessage(
                id='msg_1',
                content=[ResponseOutputText(text='Sunny.', type='output_text', annotations=[])],
                role='assistant',
                status='completed',
                type='message',
            ),
        ],
    )
    mock_client = MockOpenAIResponses.create_mock(followup)
    model = OpenAIResponsesModel('gpt-5.4', provider=OpenAIProvider(openai_client=mock_client))
    agent = Agent(model=model, capabilities=[ToolSearch()])

    @agent.tool_plain(defer_loading=True)
    def get_weather(city: str) -> str:  # pragma: no cover
        return f'Weather in {city}.'

    await agent.run('and what about tomorrow?', message_history=prior)
    kwargs = get_mock_responses_kwargs(mock_client)[0]
    # The Anthropic-generated tool search parts are not echoed back to OpenAI (wrong
    # provider) — the replayed input contains only the user message from the prior turn
    # and the new user prompt, plus no `tool_search_call` items.
    item_types = [cast('dict[str, Any]', item).get('type') for item in kwargs['input']]
    assert 'tool_search_call' not in item_types
    # `get_weather` is visible on this turn because it was discovered in the prior turn's
    # history — the local `ToolSearchToolset` emits its regular variant in the tool
    # list so the OpenAI request carries `get_weather` as a regular function tool.
    tool_names = [cast('dict[str, Any]', tool).get('name') for tool in kwargs['tools']]
    assert 'get_weather' in tool_names


def _trace_capability_messages(messages: list[ModelMessage]) -> list[tuple[str, list[dict[str, Any]]]]:
    """Compact one-line-per-part trace of a deferred-capability conversation.

    Used by the cross-provider replay tests to assert the *story* of the run
    (load → search → tool call → answer) without coupling to provider-specific
    wire shapes."""
    trace: list[tuple[str, list[dict[str, Any]]]] = []
    for msg in messages:
        part_trace: list[dict[str, Any]] = []
        for part in msg.parts:
            if isinstance(part, UserPromptPart):
                part_info: dict[str, Any] = {'type': 'user', 'content': part.content}
            elif isinstance(part, LoadCapabilityCallPart):
                part_info = {'type': 'load_capability_call', 'id': part.capability_id}
            elif isinstance(part, LoadCapabilityReturnPart):
                part_info = {'type': 'load_capability_return', 'instructions': part.instructions}
            elif isinstance(part, ToolSearchCallPart):
                queries = part.args['queries'] if isinstance(part.args, dict) else part.args
                part_info = {'type': 'tool_search_call', 'queries': queries}
            elif isinstance(part, ToolSearchReturnPart):
                part_info = {
                    'type': 'tool_search_return',
                    'tools': [tool['name'] for tool in part.content['discovered_tools']],
                }
            elif isinstance(part, ToolCallPart):
                # Normalize args from JSON string to dict so per-row snapshots don't
                # pin on provider-specific whitespace or key ordering.
                part_info = {'type': 'tool_call', 'tool_name': part.tool_name, 'args': part.args_as_dict()}
            elif isinstance(part, ToolReturnPart):
                part_info = {'type': 'tool_return', 'tool_name': part.tool_name, 'content': part.content}
            elif isinstance(part, TextPart):
                part_info = {'type': 'text'}
            else:
                raise AssertionError(
                    f'cross-provider replay trace helper saw unexpected part type: {type(part).__name__}'
                )  # pragma: no cover
            part_trace.append(part_info)
        # Use a flat lowercase tag so inline-snapshot writes a plain string instead
        # of "helpfully" resolving the class name to `'request'`.
        tag = 'request' if isinstance(msg, ModelRequest) else 'response'
        trace.append((tag, part_trace))
    return trace


# Both turns vary per provider — first-turn trajectory differs (Anthropic prepends
# prose, Google may run native search), and even the resume turn isn't uniform
# (some providers prepend a text part before the tool_call). A single `snapshot()`
# literal can only hold one row's value, so we key both expected traces by
# `(first_model, resume_model)` and let inline-snapshot fill each entry on
# `--record-mode=once --inline-snapshot=create`.
_TraceShape = list[tuple[str, list[dict[str, Any]]]]

_FIRST_TURN_EXPECTED: dict[tuple[str, str], _TraceShape] = {
    ('anthropic:claude-sonnet-4-5', 'openai-responses:gpt-5.4'): snapshot(
        [
            ('request', [{'type': 'user', 'content': 'Can I get a refund on order-123?'}]),
            ('response', [{'type': 'text'}, {'type': 'load_capability_call', 'id': 'refunds'}]),
            (
                'request',
                [
                    {
                        'type': 'load_capability_return',
                        'instructions': 'Use the refund policy tool before answering refund questions.',
                    }
                ],
            ),
            ('response', [{'type': 'tool_search_call', 'queries': ['refunds']}]),
            ('request', [{'type': 'tool_search_return', 'tools': ['lookup_refund_policy']}]),
            (
                'response',
                [{'type': 'tool_call', 'tool_name': 'lookup_refund_policy', 'args': {'order_id': 'order-123'}}],
            ),
            (
                'request',
                [
                    {
                        'type': 'tool_return',
                        'tool_name': 'lookup_refund_policy',
                        'content': 'order-123: refund allowed for 30 days',
                    }
                ],
            ),
            ('response', [{'type': 'text'}]),
        ]
    ),
    ('openai-responses:gpt-5.4', 'anthropic:claude-sonnet-4-5'): snapshot(
        [
            ('request', [{'type': 'user', 'content': 'Can I get a refund on order-123?'}]),
            ('response', [{'type': 'load_capability_call', 'id': 'refunds'}]),
            (
                'request',
                [
                    {
                        'type': 'load_capability_return',
                        'instructions': 'Use the refund policy tool before answering refund questions.',
                    }
                ],
            ),
            ('response', [{'type': 'tool_search_call', 'queries': ['refunds']}]),
            ('request', [{'type': 'tool_search_return', 'tools': ['lookup_refund_policy']}]),
            (
                'response',
                [{'type': 'tool_call', 'tool_name': 'lookup_refund_policy', 'args': {'order_id': 'order-123'}}],
            ),
            (
                'request',
                [
                    {
                        'type': 'tool_return',
                        'tool_name': 'lookup_refund_policy',
                        'content': 'order-123: refund allowed for 30 days',
                    }
                ],
            ),
            ('response', [{'type': 'text'}]),
        ]
    ),
    ('google:gemini-3-flash-preview', 'openai-responses:gpt-5.4'): snapshot(
        [
            ('request', [{'type': 'user', 'content': 'Can I get a refund on order-123?'}]),
            ('response', [{'type': 'load_capability_call', 'id': 'refunds'}]),
            (
                'request',
                [
                    {
                        'type': 'load_capability_return',
                        'instructions': 'Use the refund policy tool before answering refund questions.',
                    }
                ],
            ),
            ('response', [{'type': 'tool_search_call', 'queries': ['refunds']}]),
            ('request', [{'type': 'tool_search_return', 'tools': ['lookup_refund_policy']}]),
            (
                'response',
                [{'type': 'tool_call', 'tool_name': 'lookup_refund_policy', 'args': {'order_id': 'order-123'}}],
            ),
            (
                'request',
                [
                    {
                        'type': 'tool_return',
                        'tool_name': 'lookup_refund_policy',
                        'content': 'order-123: refund allowed for 30 days',
                    }
                ],
            ),
            ('response', [{'type': 'tool_search_call', 'queries': ['order details', 'order status']}]),
            ('request', [{'type': 'tool_search_return', 'tools': []}]),
            ('response', [{'type': 'tool_search_call', 'queries': ['get order', 'order information']}]),
            ('request', [{'type': 'tool_search_return', 'tools': []}]),
            ('response', [{'type': 'tool_search_call', 'queries': ['refund order', 'process refund']}]),
            ('request', [{'type': 'tool_search_return', 'tools': []}]),
            ('response', [{'type': 'text'}]),
        ]
    ),
    ('openai-responses:gpt-5.4', 'google:gemini-3-flash-preview'): snapshot(
        [
            ('request', [{'type': 'user', 'content': 'Can I get a refund on order-123?'}]),
            ('response', [{'type': 'load_capability_call', 'id': 'refunds'}]),
            (
                'request',
                [
                    {
                        'type': 'load_capability_return',
                        'instructions': 'Use the refund policy tool before answering refund questions.',
                    }
                ],
            ),
            ('response', [{'type': 'tool_search_call', 'queries': ['refunds']}]),
            ('request', [{'type': 'tool_search_return', 'tools': ['lookup_refund_policy']}]),
            (
                'response',
                [{'type': 'tool_call', 'tool_name': 'lookup_refund_policy', 'args': {'order_id': 'order-123'}}],
            ),
            (
                'request',
                [
                    {
                        'type': 'tool_return',
                        'tool_name': 'lookup_refund_policy',
                        'content': 'order-123: refund allowed for 30 days',
                    }
                ],
            ),
            ('response', [{'type': 'text'}]),
        ]
    ),
}


# Resume-turn structure also varies per provider — some prepend a text part before the
# tool_call. Keyed like `_FIRST_TURN_EXPECTED` so each row's record stays distinct.
_RESUME_TURN_EXPECTED: dict[tuple[str, str], _TraceShape] = {
    ('anthropic:claude-sonnet-4-5', 'openai-responses:gpt-5.4'): snapshot(
        [
            ('request', [{'type': 'user', 'content': 'And what about order-456?'}]),
            (
                'response',
                [
                    {'type': 'text'},
                    {'type': 'tool_call', 'tool_name': 'lookup_refund_policy', 'args': {'order_id': 'order-456'}},
                ],
            ),
            (
                'request',
                [
                    {
                        'type': 'tool_return',
                        'tool_name': 'lookup_refund_policy',
                        'content': 'order-456: refund allowed for 30 days',
                    }
                ],
            ),
            ('response', [{'type': 'text'}]),
        ]
    ),
    ('openai-responses:gpt-5.4', 'anthropic:claude-sonnet-4-5'): snapshot(
        [
            ('request', [{'type': 'user', 'content': 'And what about order-456?'}]),
            (
                'response',
                [{'type': 'tool_call', 'tool_name': 'lookup_refund_policy', 'args': {'order_id': 'order-456'}}],
            ),
            (
                'request',
                [
                    {
                        'type': 'tool_return',
                        'tool_name': 'lookup_refund_policy',
                        'content': 'order-456: refund allowed for 30 days',
                    }
                ],
            ),
            ('response', [{'type': 'text'}]),
        ]
    ),
    ('google:gemini-3-flash-preview', 'openai-responses:gpt-5.4'): snapshot(
        [
            ('request', [{'type': 'user', 'content': 'And what about order-456?'}]),
            (
                'response',
                [{'type': 'tool_call', 'tool_name': 'lookup_refund_policy', 'args': {'order_id': 'order-456'}}],
            ),
            (
                'request',
                [
                    {
                        'type': 'tool_return',
                        'tool_name': 'lookup_refund_policy',
                        'content': 'order-456: refund allowed for 30 days',
                    }
                ],
            ),
            ('response', [{'type': 'text'}]),
        ]
    ),
    ('openai-responses:gpt-5.4', 'google:gemini-3-flash-preview'): snapshot(
        [
            ('request', [{'type': 'user', 'content': 'And what about order-456?'}]),
            (
                'response',
                [{'type': 'tool_call', 'tool_name': 'lookup_refund_policy', 'args': {'order_id': 'order-456'}}],
            ),
            (
                'request',
                [
                    {
                        'type': 'tool_return',
                        'tool_name': 'lookup_refund_policy',
                        'content': 'order-456: refund allowed for 30 days',
                    }
                ],
            ),
            ('response', [{'type': 'text'}]),
        ]
    ),
}


# Cassette names depend on the parametrize id, so we keep these stable across the
# matrix. Recording: `pytest --record-mode=once --inline-snapshot=create`.
@pytest.mark.parametrize('first_model_name,resume_model_name', list(_FIRST_TURN_EXPECTED.keys()))
@pytest.mark.vcr
@pytest.mark.filterwarnings('ignore:`BuiltinToolCallEvent` is deprecated:DeprecationWarning')
@pytest.mark.filterwarnings('ignore:`BuiltinToolResultEvent` is deprecated:DeprecationWarning')
async def test_cross_provider_capability_replay(
    first_model_name: str,
    resume_model_name: str,
    allow_model_requests: None,
    anthropic_api_key: str,
    openai_api_key: str,
    gemini_api_key: str,
) -> None:
    """A deferred capability registered with the same definition on both turns must
    replay cleanly across any provider pair: the resuming provider re-marks the
    capability-owned tools, treats the prior turn's tool-search history as
    already-discovered, and dispatches `lookup_refund_policy` directly.

    Asserting the trace-level story over both turns is enough — per-provider wire
    shape is already covered by the dedicated OpenAI/Anthropic/Google adapter tests."""
    pytest.importorskip('anthropic')
    pytest.importorskip('openai')
    pytest.importorskip('google.genai')

    refunds_toolset = FunctionToolset()

    @refunds_toolset.tool_plain
    def lookup_refund_policy(order_id: str) -> str:
        """Look up the refund policy for an order."""
        return f'{order_id}: refund allowed for 30 days'

    def make_refunds_cap() -> Capability:
        return Capability(
            id='refunds',
            description='Refund policy tools.',
            instructions='Use the refund policy tool before answering refund questions.',
            toolsets=[refunds_toolset],
            defer_loading=True,
        )

    first_agent: Agent[None, str] = Agent(model=first_model_name, capabilities=[make_refunds_cap()])
    first_result = await first_agent.run('Can I get a refund on order-123?')

    resume_agent: Agent[None, str] = Agent(model=resume_model_name, capabilities=[make_refunds_cap()])
    resume_result = await resume_agent.run(
        'And what about order-456?',
        message_history=first_result.all_messages(),
    )

    # Per-row records pin each provider pair's full trajectory. The resume turn must
    # dispatch `lookup_refund_policy` off the replayed history without a fresh
    # load_capability or tool_search — drift in that contract surfaces as a diff here.
    key = (first_model_name, resume_model_name)
    assert _trace_capability_messages(first_result.all_messages()) == _FIRST_TURN_EXPECTED[key]
    assert _trace_capability_messages(resume_result.new_messages()) == _RESUME_TURN_EXPECTED[key]


@pytest.mark.vcr
@pytest.mark.filterwarnings('ignore:`BuiltinToolCallEvent` is deprecated:DeprecationWarning')
@pytest.mark.filterwarnings('ignore:`BuiltinToolResultEvent` is deprecated:DeprecationWarning')
async def test_anthropic_to_google_deferred_capability_history_replay(
    allow_model_requests: None,
    anthropic_api_key: str,
    gemini_api_key: str,
) -> None:
    """Deferred capability loads replay across native and local tool-search paths."""
    pytest.importorskip('anthropic')
    pytest.importorskip('google.genai')

    refunds_toolset = FunctionToolset()

    @refunds_toolset.tool_plain
    def lookup_refund_policy(order_id: str) -> str:
        """Look up the refund policy for an order."""
        return f'{order_id}: refund allowed for 30 days'

    def make_refunds_cap() -> Capability:
        return Capability(
            id='refunds',
            description='Refund policy tools.',
            instructions='Use the refund policy tool before answering refund questions.',
            toolsets=[refunds_toolset],
            defer_loading=True,
        )

    def trace_messages(messages: list[ModelMessage]) -> list[tuple[str, list[dict[str, Any]]]]:
        trace: list[tuple[str, list[dict[str, Any]]]] = []
        for msg in messages:
            part_trace: list[dict[str, Any]] = []
            for part in msg.parts:
                if isinstance(part, UserPromptPart):
                    part_info: dict[str, Any] = {'type': 'user', 'content': part.content}
                elif isinstance(part, LoadCapabilityCallPart):
                    part_info = {'type': 'load_capability_call', 'id': part.capability_id}
                elif isinstance(part, LoadCapabilityReturnPart):
                    part_info = {'type': 'load_capability_return', 'instructions': part.instructions}
                elif isinstance(part, ToolSearchCallPart):
                    queries = part.args['queries'] if isinstance(part.args, dict) else part.args
                    part_info = {'type': 'tool_search_call', 'queries': queries}
                elif isinstance(part, ToolSearchReturnPart):
                    part_info = {
                        'type': 'tool_search_return',
                        'tools': [tool['name'] for tool in part.content['discovered_tools']],
                    }
                elif isinstance(part, ToolCallPart):
                    part_info = {'type': 'tool_call', 'tool_name': part.tool_name, 'args': part.args}
                elif isinstance(part, ToolReturnPart):
                    part_info = {'type': 'tool_return', 'tool_name': part.tool_name, 'content': part.content}
                elif isinstance(part, TextPart):
                    part_info = {'type': 'text'}
                else:
                    raise AssertionError(
                        f'anthropic→google replay trace helper saw unexpected part type: {type(part).__name__}'
                    )  # pragma: no cover
                part_trace.append(part_info)
            trace.append((type(msg).__name__, part_trace))
        return trace

    anthropic_agent: Agent[None, str] = Agent(
        model='anthropic:claude-sonnet-4-5',
        capabilities=[make_refunds_cap()],
        model_settings=AnthropicModelSettings(anthropic_cache=True),
    )
    anthropic_result = await anthropic_agent.run('Can I get a refund on order-123?')

    assert trace_messages(anthropic_result.all_messages()) == snapshot(
        [
            (
                'ModelRequest',
                [{'type': 'user', 'content': 'Can I get a refund on order-123?'}],
            ),
            (
                'ModelResponse',
                [
                    {'type': 'text'},
                    {'type': 'load_capability_call', 'id': 'refunds'},
                ],
            ),
            (
                'ModelRequest',
                [
                    {
                        'type': 'load_capability_return',
                        'instructions': 'Use the refund policy tool before answering refund questions.',
                    }
                ],
            ),
            (
                'ModelResponse',
                [{'type': 'tool_search_call', 'queries': ['refunds']}],
            ),
            (
                'ModelRequest',
                [{'type': 'tool_search_return', 'tools': ['lookup_refund_policy']}],
            ),
            (
                'ModelResponse',
                [
                    {
                        'type': 'tool_call',
                        'tool_name': 'lookup_refund_policy',
                        'args': {'order_id': 'order-123'},
                    }
                ],
            ),
            (
                'ModelRequest',
                [
                    {
                        'type': 'tool_return',
                        'tool_name': 'lookup_refund_policy',
                        'content': 'order-123: refund allowed for 30 days',
                    }
                ],
            ),
            ('ModelResponse', [{'type': 'text'}]),
        ]
    )

    google_agent: Agent[None, str] = Agent(
        model='google:gemini-3-flash-preview',
        capabilities=[make_refunds_cap()],
    )
    google_result = await google_agent.run(
        'And what about order-456?',
        message_history=anthropic_result.all_messages(),
    )

    assert trace_messages(google_result.new_messages()) == snapshot(
        [
            (
                'ModelRequest',
                [{'type': 'user', 'content': 'And what about order-456?'}],
            ),
            (
                'ModelResponse',
                [
                    {
                        'type': 'tool_call',
                        'tool_name': 'lookup_refund_policy',
                        'args': {'order_id': 'order-456'},
                    }
                ],
            ),
            (
                'ModelRequest',
                [
                    {
                        'type': 'tool_return',
                        'tool_name': 'lookup_refund_policy',
                        'content': 'order-456: refund allowed for 30 days',
                    }
                ],
            ),
            ('ModelResponse', [{'type': 'text'}]),
        ]
    )
    assert google_result.output == snapshot("""\
For order-456, the policy is the same: **a refund is allowed within 30 days** of your purchase.

Is there anything else I can assist you with?\
""")


def test_anthropic_tool_search_result_error_block_mapping():
    """An error result block (no `tool_references`) produces a
    `NativeToolReturnPart` without discovered tools in its metadata."""
    pytest.importorskip('anthropic')

    error_block = BetaToolSearchToolResultBlock(
        tool_use_id='srv_err',
        type='tool_search_tool_result',
        content=BetaToolSearchToolResultError(
            error_code='unavailable',
            error_message='unavailable',
            type='tool_search_tool_result_error',
        ),
    )
    part = _map_tool_search_tool_result_block(error_block, 'anthropic')
    assert part.tool_name == 'tool_search'
    assert part.metadata is None


def test_anthropic_custom_replay_blocks_malformed_content():
    """Custom-callable replay must fall through to text formatting when the persisted
    return content doesn't parse as a `ToolSearchReturnContent` — e.g. older history
    written before the typed shape, or a hand-crafted return — rather than crashing or
    fabricating an empty discovery."""
    pytest.importorskip('anthropic')

    malformed = ToolReturnPart(tool_name='search_tools', content='not a typed return', tool_call_id='c1')
    refs, message = _build_custom_tool_search_replay_blocks(
        malformed, tool_search_active=True, available_tool_names=set()
    )
    assert refs is None and message is None


def test_anthropic_build_tool_search_replay_block_error_branch():
    """Replay reconstruction must round-trip an error result that the parse-time
    mapper stashed on `provider_details` back into the `tool_search_tool_result_error`
    inner block — otherwise a transient provider error on turn N would silently
    flip into a fake successful empty-search on turn N+1's resend.

    The Anthropic SDK's `BetaToolSearchToolResultErrorParam` carries only `error_code`
    on the wire (no `error_message`), so the message stashed on `provider_details`
    is observability-only — verified separately in
    `test_anthropic_tool_search_result_error_block_mapping`.
    """
    pytest.importorskip('anthropic')

    return_part = NativeToolSearchReturnPart(
        provider_name='anthropic',
        tool_call_id='srv_err',
        content={'discovered_tools': []},
        provider_details={'error_code': 'unavailable', 'error_message': 'temporary outage'},
    )
    block = _build_tool_search_replay_block(return_part, 'srv_err', available_tool_names=set())
    assert block == {
        'tool_use_id': 'srv_err',
        'type': 'tool_search_tool_result',
        'content': {
            'type': 'tool_search_tool_result_error',
            'error_code': 'unavailable',
        },
    }


def test_openai_map_tool_search_call_unit():
    """Unit-level: `_map_tool_search_call` and `_build_tool_search_return_part` produce
    populated metadata for various output shapes — useful as a fast deterministic
    gate without burning a live API call. The end-to-end live cassette in
    `test_openai_native_tool_search_round_trip` exercises the same functions with
    real provider responses."""

    call = ResponseToolSearchCall(
        id='ts_1',
        arguments={'paths': ['get_exchange_rate']},
        call_id='call_1',
        execution='server',
        status='completed',
        type='tool_search_call',
    )
    output = ResponseToolSearchOutputItem(
        id='tso_1',
        call_id='call_1',
        execution='server',
        status='completed',
        tools=[
            FunctionTool(name='get_exchange_rate', description='', parameters={}, strict=False, type='function'),
        ],
        type='tool_search_output',
    )
    call_part, return_part = _map_tool_search_call(call, output, 'openai')
    assert isinstance(call_part, NativeToolSearchCallPart)
    assert call_part.tool_name == 'tool_search'
    # OpenAI server-executed `tool_search.arguments` carries `paths`; the adapter
    # normalizes that into the cross-provider `queries` slot.
    assert call_part.args == {'queries': ['get_exchange_rate']}
    assert isinstance(return_part, NativeToolSearchReturnPart)
    assert return_part.content == {'discovered_tools': [{'name': 'get_exchange_rate'}]}
    assert return_part.provider_details == {'status': 'completed'}

    # No output item → empty discovery (streaming start case).
    empty_return = _build_tool_search_return_part('call_1', 'in_progress', None, 'openai')
    assert empty_return.content == {'discovered_tools': []}
    assert empty_return.provider_details == {'status': 'in_progress'}

    # Non-function tools in the output don't have a `name` attribute and are skipped.

    mixed_output = ResponseToolSearchOutputItem(
        id='tso_mix',
        call_id='call_mix',
        execution='server',
        status='completed',
        tools=[
            FunctionTool(name='real', description='', parameters={}, strict=False, type='function'),
            # FileSearchTool doesn't have a `name` — the loop's `isinstance` guard skips it.
            FileSearchTool(type='file_search', vector_store_ids=['vs_1']),
        ],
        type='tool_search_output',
    )
    mixed = _build_tool_search_return_part('call_mix', 'completed', mixed_output, 'openai')
    assert mixed.content == {'discovered_tools': [{'name': 'real'}]}


@pytest.mark.vcr
async def test_openai_native_tool_search_round_trip(allow_model_requests: None, openai_api_key: str) -> None:
    """End-to-end against live OpenAI Responses: native server-executed `tool_search`
    populates `NativeToolCallPart` / `NativeToolReturnPart`, the model invokes the
    discovered deferred tool by its plain name, and the second-turn replay carries
    `defer_loading: true` on the corpus function tool plus a `tool_search_call` item.
    """

    model = OpenAIResponsesModel('gpt-5.4', provider=OpenAIProvider(api_key=openai_api_key))
    agent = Agent(model=model)

    @agent.tool_plain(defer_loading=True)
    def get_exchange_rate(from_currency: str, to_currency: str) -> str:
        """Look up the current exchange rate between two currencies."""
        return f'1 {from_currency} = 0.92 {to_currency}'

    @agent.tool_plain(defer_loading=True)
    def stock_lookup(symbol: str) -> str:  # pragma: no cover
        """Look up stock price by ticker symbol."""
        return f'Stock {symbol}: $150.00'

    result = await agent.run('What is the current USD to EUR exchange rate?')

    assert any(
        isinstance(p, NativeToolCallPart) and p.tool_name == 'tool_search'
        for m in result.all_messages()
        for p in m.parts
    )
    assert any(
        isinstance(p, NativeToolReturnPart) and p.tool_name == 'tool_search'
        for m in result.all_messages()
        for p in m.parts
    )

    rate_returns = [
        p
        for m in result.all_messages()
        for p in m.parts
        if isinstance(p, ToolReturnPart) and p.tool_name == 'get_exchange_rate'
    ]
    assert len(rate_returns) == 1
    assert rate_returns[0].content == '1 USD = 0.92 EUR'

    # Wire-level checks against the live cassette.
    cassette_path = (
        Path(__file__).parent / 'cassettes' / 'test_tool_search' / 'test_openai_native_tool_search_round_trip.yaml'
    )
    cassette = cast(dict[str, Any], yaml.safe_load(cassette_path.read_text(encoding='utf-8')))
    interactions = cast(list[dict[str, Any]], cassette['interactions'])

    # Initial request: deferred tools ship with `defer_loading: true`, and the native
    # `tool_search` builtin is registered alongside.
    first_request = cast(dict[str, Any], interactions[0]['request']['parsed_body'])
    deferred_names = {
        cast(str, t['name'])
        for t in cast(list[dict[str, Any]], first_request['tools'])
        if t.get('defer_loading') is True
    }
    assert deferred_names == {'get_exchange_rate', 'stock_lookup'}
    assert any(t.get('type') == 'tool_search' for t in cast(list[dict[str, Any]], first_request['tools']))
    # Second-turn replay carries the native tool_search_call back; the deferred corpus
    # is preserved with `defer_loading: true`.
    second_request = cast(dict[str, Any], interactions[1]['request']['parsed_body'])
    second_input_types = {
        cast(str, item.get('type'))
        for item in cast(list[dict[str, Any]], second_request['input'])
        if isinstance(item, dict)
    }
    assert 'tool_search_call' in second_input_types
    second_deferred = {
        cast(str, t['name'])
        for t in cast(list[dict[str, Any]], second_request['tools'])
        if t.get('defer_loading') is True
    }
    assert 'get_exchange_rate' in second_deferred


@pytest.mark.vcr
async def test_openai_native_tool_search_gpt_5_6(allow_model_requests: None, openai_api_key: str) -> None:
    """End-to-end against live OpenAI Responses: GPT-5.6 supports the native `tool_search`
    tool with `defer_loading`, backing `supports_tool_search` in its model profile — the
    server-executed search discovers the deferred tool and the model dispatches it.
    """
    model = OpenAIResponsesModel('gpt-5.6-sol', provider=OpenAIProvider(api_key=openai_api_key))
    agent = Agent(model=model)

    @agent.tool_plain(defer_loading=True)
    def get_exchange_rate(from_currency: str, to_currency: str) -> str:
        """Look up the current exchange rate between two currencies."""
        return f'1 {from_currency} = 0.92 {to_currency}'

    @agent.tool_plain(defer_loading=True)
    def stock_lookup(symbol: str) -> str:  # pragma: no cover
        """Look up stock price by ticker symbol."""
        return f'Stock {symbol}: $150.00'

    result = await agent.run('What is the current USD to EUR exchange rate?')

    assert any(
        isinstance(p, NativeToolCallPart) and p.tool_name == 'tool_search'
        for m in result.all_messages()
        for p in m.parts
    )
    rate_returns = [
        p
        for m in result.all_messages()
        for p in m.parts
        if isinstance(p, ToolReturnPart) and p.tool_name == 'get_exchange_rate'
    ]
    assert len(rate_returns) == 1
    assert rate_returns[0].content == '1 USD = 0.92 EUR'


@pytest.mark.vcr
async def test_openai_execution_client_round_trip(allow_model_requests: None, openai_api_key: str, vcr: Any) -> None:
    """End-to-end: a custom callable `ToolSearch` strategy surfaces natively on OpenAI
    Responses as `ToolSearchToolParam(execution='client')` — the provider emits a
    `tool_search_call` with `execution='client'` whose arguments we dispatch to the
    local `search_tools` function, and the resulting `ToolReturnPart` is replayed
    as a `tool_search_output` (execution='client') carrying the discovered tool defs."""

    def match_exchange_rate(ctx: RunContext, queries: Sequence[str], tools: Sequence[ToolDefinition]) -> list[str]:
        # Deterministic: always point the model at `get_exchange_rate` so the cassette
        # replay doesn't depend on the exact keywords the model picks.
        return ['get_exchange_rate']

    model = OpenAIResponsesModel('gpt-5.4', provider=OpenAIProvider(api_key=openai_api_key))
    agent = Agent(
        model=model,
        instructions=(
            'When you need a capability not provided by your visible tools, call the built-in '
            'tool search first to discover and activate the right one before answering.'
        ),
        capabilities=[ToolSearch(strategy=match_exchange_rate)],
    )

    @agent.tool_plain
    def get_weather(city: str) -> str:  # pragma: no cover
        """Get the current weather for a city."""
        return f'Weather in {city} is sunny.'

    @agent.tool_plain(defer_loading=True)
    def get_exchange_rate(from_currency: str, to_currency: str) -> str:
        """Look up the current exchange rate between two currencies."""
        return f'1 {from_currency} = 0.92 {to_currency}'

    @agent.tool_plain(defer_loading=True)
    def stock_lookup(symbol: str) -> str:  # pragma: no cover
        """Look up stock price by ticker symbol."""
        return f'Stock {symbol}: $150.00'

    result = await agent.run('What is the current exchange rate from USD to EUR?')

    tool_call_names = [
        part.tool_name
        for msg in result.all_messages()
        if isinstance(msg, ModelResponse)
        for part in msg.parts
        if isinstance(part, ToolCallPart)
    ]
    # The model called `search_tools` (our local, client-executed search) and then the
    # discovered `get_exchange_rate` — routed through the regular `ToolCallPart` /
    # `ToolReturnPart` path on both sides of the wire.
    assert 'search_tools' in tool_call_names
    assert 'get_exchange_rate' in tool_call_names

    # The local `search_tools` run recorded the discovered tool on `content` as a typed
    # `ToolSearchReturnContent` — this is the same value read back by `ToolSearchToolset`
    # on later turns to unlock the deferred tool on the local path (and round-tripped as
    # `tool_search_output.tools` in the cassette's replay request body).
    search_returns = [
        part
        for msg in result.all_messages()
        for part in msg.parts
        if isinstance(part, ToolReturnPart) and part.tool_name == 'search_tools'
    ]
    assert len(search_returns) == 1
    assert search_returns[0].content == {'discovered_tools': [{'name': 'get_exchange_rate'}]}

    rate_returns = [
        part
        for msg in result.all_messages()
        for part in msg.parts
        if isinstance(part, ToolReturnPart) and part.tool_name == 'get_exchange_rate'
    ]
    assert len(rate_returns) == 1
    assert rate_returns[0].content == '1 USD = 0.92 EUR'

    # Wire-level: with a custom (client-executed) strategy, OpenAI rejects the `tool_search` tool
    # without a non-null `description`, so the adapter always sends one. Confirm the recorded
    # request carries it (guards the description default against a silent regression → live 422).
    request_tools = json.loads(vcr.requests[0].body)['tools']
    client_tool_search = [
        tool for tool in request_tools if tool.get('type') == 'tool_search' and tool.get('execution') == 'client'
    ]
    assert client_tool_search, 'expected a client-executed tool_search tool in the request'
    assert client_tool_search[0].get('description'), 'tool_search description must be non-null'


@pytest.mark.vcr
async def test_anthropic_native_tool_search_streaming(allow_model_requests: None, anthropic_api_key: str) -> None:
    """End-to-end streaming against live Anthropic: native BM25 server-side tool search
    streams `NativeToolSearchCallPart` / `NativeToolSearchReturnPart` through the part
    manager during `agent.iter` + `node.stream`, the model invokes the discovered
    deferred tool by its plain name, and the agent loop runs to a final text response."""
    pytest.importorskip('anthropic')

    model = AnthropicModel('claude-sonnet-4-6', provider=AnthropicProvider(api_key=anthropic_api_key))
    agent = Agent(model=model)

    @agent.tool_plain(defer_loading=True)
    def get_exchange_rate(from_currency: str, to_currency: str) -> str:
        """Look up the current exchange rate between two currencies."""
        return f'1 {from_currency} = 0.92 {to_currency}'

    @agent.tool_plain(defer_loading=True)
    def stock_lookup(symbol: str) -> str:  # pragma: no cover
        """Look up stock price by ticker symbol."""
        return f'Stock {symbol}: $150.00'

    streamed_events: list[Any] = []
    async with agent.iter(user_prompt='What is the current USD to EUR exchange rate?') as agent_run:
        async for node in agent_run:
            if Agent.is_model_request_node(node) or Agent.is_call_tools_node(node):
                async with node.stream(agent_run.ctx) as request_stream:
                    async for event in request_stream:
                        streamed_events.append(event)

    assert agent_run.result is not None

    # The streamed run materializes the same typed builtin parts as the non-streaming
    # round-trip — the part manager promotes them through the discriminator at
    # `content_block_start` time, not just on final response assembly.
    builtin_call_parts = [
        p for m in agent_run.result.all_messages() for p in m.parts if isinstance(p, NativeToolSearchCallPart)
    ]
    builtin_return_parts = [
        p for m in agent_run.result.all_messages() for p in m.parts if isinstance(p, NativeToolSearchReturnPart)
    ]
    assert builtin_call_parts and builtin_return_parts

    # The discovered deferred tool dispatches by its plain name and produces its
    # ToolReturnPart end-to-end.
    rate_returns = [
        p
        for m in agent_run.result.all_messages()
        for p in m.parts
        if isinstance(p, ToolReturnPart) and p.tool_name == 'get_exchange_rate'
    ]
    assert len(rate_returns) == 1
    assert rate_returns[0].content == '1 USD = 0.92 EUR'

    # We received streaming events from both the model-request node and the call-tools
    # node — i.e. the part manager surfaced the builtin tool-search parts as the stream
    # came in (not just on `streamed.get()`).
    assert streamed_events, 'expected streaming events from the request stream'


@pytest.mark.vcr
async def test_openai_native_tool_search_streaming(allow_model_requests: None, openai_api_key: str) -> None:
    """End-to-end streaming against live OpenAI Responses: native server-executed
    `tool_search` streams `NativeToolSearchCallPart` / `NativeToolSearchReturnPart`
    through the part manager during `agent.iter` + `node.stream`, the model invokes
    the discovered deferred tool by its plain name, and the agent loop runs to a final
    text response."""

    model = OpenAIResponsesModel('gpt-5.4', provider=OpenAIProvider(api_key=openai_api_key))
    agent = Agent(model=model)

    @agent.tool_plain(defer_loading=True)
    def get_exchange_rate(from_currency: str, to_currency: str) -> str:
        """Look up the current exchange rate between two currencies."""
        return f'1 {from_currency} = 0.92 {to_currency}'

    @agent.tool_plain(defer_loading=True)
    def stock_lookup(symbol: str) -> str:  # pragma: no cover
        """Look up stock price by ticker symbol."""
        return f'Stock {symbol}: $150.00'

    streamed_events: list[Any] = []
    async with agent.iter(user_prompt='What is the current USD to EUR exchange rate?') as agent_run:
        async for node in agent_run:
            if Agent.is_model_request_node(node) or Agent.is_call_tools_node(node):
                async with node.stream(agent_run.ctx) as request_stream:
                    async for event in request_stream:
                        streamed_events.append(event)

    assert agent_run.result is not None

    builtin_call_parts = [
        p for m in agent_run.result.all_messages() for p in m.parts if isinstance(p, NativeToolSearchCallPart)
    ]
    builtin_return_parts = [
        p for m in agent_run.result.all_messages() for p in m.parts if isinstance(p, NativeToolSearchReturnPart)
    ]
    assert builtin_call_parts and builtin_return_parts

    rate_returns = [
        p
        for m in agent_run.result.all_messages()
        for p in m.parts
        if isinstance(p, ToolReturnPart) and p.tool_name == 'get_exchange_rate'
    ]
    assert len(rate_returns) == 1
    assert rate_returns[0].content == '1 USD = 0.92 EUR'

    assert streamed_events, 'expected streaming events from the request stream'


@pytest.mark.vcr
async def test_openai_client_tool_search_streaming(allow_model_requests: None, openai_api_key: str) -> None:
    """End-to-end streaming against live OpenAI Responses with a custom callable
    `ToolSearch` strategy. The provider emits a `tool_search_call` with
    `execution='client'` whose arguments we dispatch to the local `search_tools`
    function — both events surface through the streaming part manager (the
    `tool_search_call` as a regular `ToolCallPart`), the agent loop runs the
    callable strategy, the model follows up with the discovered deferred tool, and
    the run completes with a final text response."""

    def match_exchange_rate(ctx: RunContext, queries: Sequence[str], tools: Sequence[ToolDefinition]) -> list[str]:
        # Deterministic: always point the model at `get_exchange_rate` so the cassette
        # replay doesn't depend on the exact keywords the model picks.
        return ['get_exchange_rate']

    model = OpenAIResponsesModel('gpt-5.4', provider=OpenAIProvider(api_key=openai_api_key))
    agent = Agent(
        model=model,
        instructions=(
            'When you need a capability not provided by your visible tools, call the built-in '
            'tool search first to discover and activate the right one before answering.'
        ),
        capabilities=[ToolSearch(strategy=match_exchange_rate)],
    )

    @agent.tool_plain
    def get_weather(city: str) -> str:  # pragma: no cover
        """Get the current weather for a city."""
        return f'Weather in {city} is sunny.'

    @agent.tool_plain(defer_loading=True)
    def get_exchange_rate(from_currency: str, to_currency: str) -> str:
        """Look up the current exchange rate between two currencies."""
        return f'1 {from_currency} = 0.92 {to_currency}'

    @agent.tool_plain(defer_loading=True)
    def stock_lookup(symbol: str) -> str:  # pragma: no cover
        """Look up stock price by ticker symbol."""
        return f'Stock {symbol}: $150.00'

    streamed_events: list[Any] = []
    async with agent.iter(user_prompt='What is the current exchange rate from USD to EUR?') as agent_run:
        async for node in agent_run:
            if Agent.is_model_request_node(node) or Agent.is_call_tools_node(node):
                async with node.stream(agent_run.ctx) as request_stream:
                    async for event in request_stream:
                        streamed_events.append(event)

    assert agent_run.result is not None

    tool_call_names = [
        part.tool_name
        for msg in agent_run.result.all_messages()
        if isinstance(msg, ModelResponse)
        for part in msg.parts
        if isinstance(part, ToolCallPart)
    ]
    # Client-executed tool search: the `tool_search_call` is routed to the local
    # `search_tools` function, then the model follows up with the discovered tool.
    assert 'search_tools' in tool_call_names
    assert 'get_exchange_rate' in tool_call_names

    search_returns = [
        part
        for msg in agent_run.result.all_messages()
        for part in msg.parts
        if isinstance(part, ToolReturnPart) and part.tool_name == 'search_tools'
    ]
    assert len(search_returns) == 1
    assert search_returns[0].content == {'discovered_tools': [{'name': 'get_exchange_rate'}]}

    rate_returns = [
        part
        for msg in agent_run.result.all_messages()
        for part in msg.parts
        if isinstance(part, ToolReturnPart) and part.tool_name == 'get_exchange_rate'
    ]
    assert len(rate_returns) == 1
    assert rate_returns[0].content == '1 USD = 0.92 EUR'

    assert streamed_events, 'expected streaming events from the request stream'


async def test_agent_graph_without_builtin_tools(allow_model_requests: None, monkeypatch: pytest.MonkeyPatch) -> None:
    """Covers `_agent_graph`'s empty `ctx.deps.native_tools` branch.

    Auto-inject always adds `ToolSearchTool`, so the only way to exercise the empty
    branch is to disable auto-inject in the test.
    """

    monkeypatch.setattr(agent_module, '_AUTO_INJECT_CAPABILITY_TYPES', ())
    agent = Agent('test')
    result = await agent.run('hi')
    assert isinstance(result.output, str)


async def test_tool_search_toolset_discovers_from_builtin_return_part():
    """Discovery metadata on a `NativeToolSearchReturnPart` from a native provider search
    is picked up so the local path recovers state on cross-provider handover."""

    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset)

    messages: list[ModelMessage] = [
        ModelResponse(
            parts=[
                NativeToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'calculate_mortgage'}]},
                )
            ]
        )
    ]
    # `parse_discovered_tools` extracts discovery from the native return part; mirror that
    # into `discovered_tool_names`, which `get_tools` reads to flip visibility.
    assert parse_discovered_tools(messages) == {'calculate_mortgage'}
    ctx = _build_run_context(None, messages=messages, discovered_tool_names={'calculate_mortgage'})

    tools = await searchable.get_tools(ctx)
    assert tools['calculate_mortgage'].tool_def.defer_loading is False
    assert tools['stock_price'].tool_def.defer_loading is True


async def test_tool_search_toolset_custom_search_fn_filters_unknown_names():
    """Names returned by `search_fn` that aren't in the deferred set are discarded."""

    def custom_search(ctx: RunContext, queries: Sequence[str], tools: Sequence[ToolDefinition]) -> list[str]:
        return ['stock_price', 'not_a_real_tool', 'crypto_price']

    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset, search_fn=custom_search)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['anything']}, ctx, tools[_SEARCH_TOOLS_NAME])
    assert result == {
        'discovered_tools': [
            {'name': 'stock_price'},
            {'name': 'crypto_price'},
        ]
    }


async def test_tool_search_toolset_custom_search_fn_no_matches():
    """Custom search function returning no names produces the 'no matches' message."""

    def custom_search(ctx: RunContext, queries: Sequence[str], tools: Sequence[ToolDefinition]) -> list[str]:
        return []

    toolset = _create_function_toolset()
    searchable = ToolSearchToolset(wrapped=toolset, search_fn=custom_search)
    ctx = _build_run_context(None)

    tools = await searchable.get_tools(ctx)
    result = await searchable.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['anything']}, ctx, tools[_SEARCH_TOOLS_NAME])
    assert result == {
        'discovered_tools': [],
        'message': 'No matching tools found. The tools you need may not be available.',
    }


async def test_tool_search_capability_strategy_callable_registers_custom_builtin():
    """A callable strategy still registers a `ToolSearchTool` builtin with `strategy='custom'`
    so provider adapters that support a custom-callable native surface (e.g. Anthropic's
    `tool_reference` result blocks, OpenAI's `execution='client'`) can use it; models
    without support drop it as optional and fall back to the local `search_tools` tool."""

    def noop(ctx: RunContext, queries: Sequence[str], tools: Sequence[ToolDefinition]) -> list[str]:  # pragma: no cover
        return []

    cap = ToolSearch(strategy=noop)
    builtins = list(cap.get_native_tools())
    assert len(builtins) == 1
    tool = builtins[0]
    assert isinstance(tool, ToolSearchTool)
    assert tool.strategy == 'custom'


async def test_tool_search_capability_strategy_named_registers_builtin():
    """Named native strategies register a non-optional `ToolSearchTool` — the request
    must error on models that can't honor the choice rather than silently substituting
    a local algorithm for bm25/regex."""
    cap = ToolSearch(strategy='regex')
    builtins = list(cap.get_native_tools())
    assert len(builtins) == 1
    tool = builtins[0]
    assert isinstance(tool, ToolSearchTool)
    assert tool.strategy == 'regex'
    assert tool.optional is False


async def test_tool_search_capability_strategy_none_optional_builtin():
    """The default (`None`) strategy registers an optional builtin so the local
    token-matching fallback takes over on models without native support."""
    cap = ToolSearch()
    builtins = list(cap.get_native_tools())
    assert len(builtins) == 1
    tool = builtins[0]
    assert isinstance(tool, ToolSearchTool)
    assert tool.strategy is None
    assert tool.optional is True


async def test_tool_search_capability_wraps_with_tool_search_toolset():
    """`strategy='keywords'` wraps with `ToolSearchToolset` so the corpus is
    exposed and `search_tools` carries the user's customizations. The toolset's
    `search_fn` is set to the built-in keyword-overlap algorithm so the local
    dispatch routes through `_run_search_fn` (same path as a custom callable),
    enabling client-executed-native wire on supporting providers."""
    toolset = _create_function_toolset()
    cap = ToolSearch(strategy='keywords')
    wrapped = cap.get_wrapper_toolset(toolset)
    assert isinstance(wrapped, ToolSearchToolset)
    assert wrapped.search_fn is not None


async def test_tool_search_capability_named_strategy_wraps_with_tool_search_toolset():
    """Named native strategies (bm25/regex) still wrap with `ToolSearchToolset` so
    the corpus is exposed; `prepare_request` raises on unsupported models because the
    builtin is registered with `optional=False`."""
    toolset = _create_function_toolset()
    cap = ToolSearch(strategy='bm25')
    wrapped = cap.get_wrapper_toolset(toolset)
    assert isinstance(wrapped, ToolSearchToolset)
    assert wrapped.search_fn is None


async def test_tool_search_named_strategy_raises_on_unsupported_model():
    """Named native strategies error on models that don't support `ToolSearchTool`
    — there's no legal fallback for `strategy='bm25'` on e.g. GPT-4."""

    m = TestModel()
    with pytest.raises(UserError, match='not supported by this model'):
        m.prepare_request(
            None,
            ModelRequestParameters(function_tools=[], native_tools=[ToolSearchTool(strategy='bm25')]),
        )


@pytest.mark.parametrize('strategy', ['bm25', 'regex'])
async def test_tool_search_named_strategy_agent_run_raises_on_unsupported_model(strategy: str):
    """End-to-end: `ToolSearch(strategy='bm25'|'regex')` on a model without native
    tool-search support must raise `UserError` rather than silently substituting the
    local keyword-overlap algorithm. The capability promises that named-native strategies
    error on adapters that can't honor the choice; previously the toolset always
    registered the local `search_tools` function as a fallback, which masked the
    error by letting `_resolve_builtin_tool_swap` drop the optional-False builtin."""
    agent = Agent(TestModel(), capabilities=[ToolSearch(strategy=cast(Any, strategy))])

    @agent.tool_plain(defer_loading=True)
    def get_weather(city: str) -> str:  # pragma: no cover
        return f'Weather in {city}'

    with pytest.raises(UserError, match=r'ToolSearchTool.*not supported by this model'):
        await agent.run('what should I wear?')


async def test_tool_search_keywords_agent_run_falls_back_on_unsupported_model():
    """Inverse of the named-strategy test: `strategy='keywords'` has a local
    implementation, so the request must fall back silently on a model without native
    tool-search support — running the agent should not raise."""
    agent = Agent(TestModel(), capabilities=[ToolSearch(strategy='keywords')])

    @agent.tool_plain(defer_loading=True)
    def get_weather(city: str) -> str:  # pragma: no cover
        return f'Weather in {city}'

    # `TestModel` doesn't support `ToolSearchTool`; with a local fallback available
    # this should run without error.
    result = await agent.run('hello')
    assert result.output


@pytest.mark.parametrize('strategy', ['bm25', 'regex'])
async def test_tool_search_named_strategy_skips_local_search_tools_emission(strategy: str):
    """Named-native strategies (`'bm25'`/`'regex'`) construct the toolset with
    `enable_fallback=False`; `get_tools` then skips emitting the local `search_tools`
    function tool entirely. Two effects fall out:

    * On *supported* providers (Anthropic), the wire carries only the native
      `tool_search_tool_*` builtin — no redundant local function tool that could
      confuse the model or waste a tool slot.
    * On *unsupported* providers, `_resolve_builtin_tool_swap` has no fallback to count
      against the (non-optional) builtin and raises `UserError` as promised."""
    toolset = _create_function_toolset()
    cap = ToolSearch(strategy=cast(Any, strategy))
    wrapped = cap.get_wrapper_toolset(toolset)
    assert isinstance(wrapped, ToolSearchToolset)
    assert wrapped.enable_fallback is False

    ctx = _build_run_context(None)
    tools = await wrapped.get_tools(ctx)
    # `search_tools` is omitted entirely — the deferred corpus is still exposed by name
    # (carrying `with_native='tool_search'`) so the swap logic can route discovery.
    assert _SEARCH_TOOLS_NAME not in tools
    corpus_names = {name for name, t in tools.items() if t.tool_def.with_native == 'tool_search'}
    assert corpus_names == {'calculate_mortgage', 'stock_price', 'crypto_price'}


async def test_tool_search_keywords_ignores_builtin_support():
    """`strategy='keywords'` never tries to use a native builtin — the swap is a
    no-op even on models that support `ToolSearchTool`."""

    class ToolSearchTestModel(TestModel):
        @classmethod
        def supported_native_tools(cls):
            return frozenset({ToolSearchTool})

    m = ToolSearchTestModel()
    search_tool = ToolDefinition(name=_SEARCH_TOOLS_NAME, description='local', parameters_json_schema={})
    _, prepared = m.prepare_request(
        None,
        ModelRequestParameters(function_tools=[search_tool], native_tools=[]),
    )
    assert prepared.native_tools == []
    assert [t.name for t in prepared.function_tools] == [_SEARCH_TOOLS_NAME]


def test_with_native_undiscovered_drops_on_unsupported_model():
    """In `prepare_request`, `with_native` corpus members with `defer_loading=True`
    (still undiscovered) drop on a model that doesn't support the builtin — the model has
    no way to call them and the local `search_tools` fallback handles discovery."""

    m = TestModel()
    # `optional=True` models the default auto path where the builtin is a best-effort
    # upgrade; on a model that doesn't support it, both the builtin and its undiscovered
    # corpus drop so the local `ToolSearch` fallback handles discovery.
    search_builtin = ToolSearchTool(optional=True)
    corpus_tool = ToolDefinition(name='deferred_tool', with_native='tool_search', defer_loading=True)

    _, prepared = m.prepare_request(
        None,
        ModelRequestParameters(
            function_tools=[corpus_tool],
            native_tools=[search_builtin],
        ),
    )
    assert prepared.native_tools == []
    assert prepared.function_tools == []


def test_with_native_discovered_kept_on_unsupported_model():
    """A discovered corpus member (`defer_loading=False`) stays in the request even when
    the builtin is unsupported — the model can call it directly by name on the local path."""

    m = TestModel()
    corpus_tool = ToolDefinition(name='deferred_tool', with_native='tool_search', defer_loading=False)

    _, prepared = m.prepare_request(
        None,
        ModelRequestParameters(
            function_tools=[corpus_tool],
            native_tools=[ToolSearchTool(optional=True)],
        ),
    )
    assert prepared.native_tools == []
    assert [t.name for t in prepared.function_tools] == ['deferred_tool']


def test_with_native_kept_on_supporting_model():
    """On a supporting model, managed tools are kept so the adapter can emit them
    with provider-specific wire-format tweaks."""

    class ToolSearchTestModel(TestModel):
        @classmethod
        def supported_native_tools(cls):
            return frozenset({ToolSearchTool})

    m = ToolSearchTestModel()
    corpus_tool = ToolDefinition(name='deferred_tool', with_native='tool_search')
    _, prepared = m.prepare_request(
        None,
        ModelRequestParameters(
            function_tools=[corpus_tool],
            native_tools=[ToolSearchTool()],
        ),
    )
    assert [t.name for t in prepared.function_tools] == ['deferred_tool']
    assert any(isinstance(t, ToolSearchTool) for t in prepared.native_tools)


def test_optional_builtin_dropped_with_empty_corpus():
    """An `optional` builtin is silently dropped when no managed corpus is in the request."""

    class ToolSearchTestModel(TestModel):
        @classmethod
        def supported_native_tools(cls):
            return frozenset({ToolSearchTool})

    m = ToolSearchTestModel()
    _, prepared = m.prepare_request(
        None,
        ModelRequestParameters(function_tools=[], native_tools=[ToolSearchTool(optional=True)]),
    )
    assert prepared.native_tools == []


def test_narrow_type_promotes_builtin_call_to_tool_search() -> None:
    """Direct construction of `NativeToolCallPart` with `tool_kind='tool-search'`
    promotes to `NativeToolSearchCallPart` via the narrowing registry."""
    base = NativeToolCallPart(
        tool_name='tool_search',
        args={'queries': ['mortgage']},
        tool_call_id='c1',
        tool_kind='tool-search',
        provider_name='anthropic',
        provider_details={'strategy': 'bm25'},
    )
    narrowed = NativeToolCallPart.narrow_type(base)
    assert isinstance(narrowed, NativeToolSearchCallPart)
    assert narrowed.args == {'queries': ['mortgage']}
    assert narrowed.tool_call_id == 'c1'
    assert narrowed.provider_name == 'anthropic'
    assert narrowed.provider_details == {'strategy': 'bm25'}

    already_narrowed = NativeToolSearchCallPart(args={'queries': ['x']}, tool_call_id='c2')
    assert NativeToolCallPart.narrow_type(already_narrowed) is already_narrowed


def test_narrow_type_promotes_builtin_return_to_tool_search() -> None:
    """Direct construction of `NativeToolReturnPart` with `tool_kind='tool-search'`
    promotes to `NativeToolSearchReturnPart` via the narrowing registry."""
    base = NativeToolReturnPart(
        tool_name='tool_search',
        content={'discovered_tools': [{'name': 'foo'}]},
        tool_call_id='c1',
        tool_kind='tool-search',
        provider_name='anthropic',
    )
    narrowed = NativeToolReturnPart.narrow_type(base)
    assert isinstance(narrowed, NativeToolSearchReturnPart)
    assert narrowed.content == {'discovered_tools': [{'name': 'foo'}]}

    already_narrowed = NativeToolSearchReturnPart(content={'discovered_tools': []}, tool_call_id='c2')
    assert NativeToolReturnPart.narrow_type(already_narrowed) is already_narrowed


def test_narrow_type_unknown_tool_kind_returns_input_unchanged() -> None:
    """Unknown `tool_kind` values aren't promoted (future builtins not yet typed)."""
    base = NativeToolCallPart(
        tool_name='something_unregistered',
        args={},
        tool_call_id='c1',
        tool_kind=cast('ToolPartKind', 'custom_kind'),  # forward-compat: discriminator unknown to the current registry
    )
    assert NativeToolCallPart.narrow_type(base) is base


def test_narrow_type_no_tool_kind_returns_input_unchanged() -> None:
    """User-defined tools sharing a framework `tool_name` aren't promoted when `tool_kind` is unset.

    Protects users whose own tool happens to be called `tool_search` / `search_tools` from
    having their parts promoted to typed subclasses that would fail shape validation against
    the typed `args` `TypedDict`.
    """
    builtin_collision = NativeToolCallPart(tool_name='tool_search', args={'foo': 'bar'}, tool_call_id='c1')
    assert builtin_collision.tool_kind is None
    assert NativeToolCallPart.narrow_type(builtin_collision) is builtin_collision

    local_collision = ToolCallPart(tool_name='search_tools', args={'query': 'x'}, tool_call_id='c2')
    assert local_collision.tool_kind is None
    assert ToolCallPart.narrow_type(local_collision) is local_collision


def test_model_response_dict_round_trip_promotes_typed_subclasses() -> None:
    """Pydantic deserialization of a dict-shaped `ModelResponse` promotes
    `tool_search` builtin parts to typed subclasses via the discriminator."""

    raw: dict[str, Any] = {
        'kind': 'response',
        'parts': [
            {
                'part_kind': 'builtin-tool-call',
                'tool_name': 'tool_search',
                'tool_kind': 'tool-search',
                'args': {'queries': ['mortgage']},
                'tool_call_id': 'c1',
                'provider_name': 'anthropic',
            },
            {
                'part_kind': 'builtin-tool-return',
                'tool_name': 'tool_search',
                'tool_kind': 'tool-search',
                'content': {'discovered_tools': [{'name': 'foo'}]},
                'tool_call_id': 'c1',
                'provider_name': 'anthropic',
            },
            {
                'part_kind': 'builtin-tool-call',
                'tool_name': 'web_search',
                'args': {'query': 'x'},
                'tool_call_id': 'c2',
            },
            # User-defined builtin call colliding with a framework tool_name. Without
            # `tool_kind`, dispatch should NOT promote — args don't match `ToolSearchArgs`.
            {
                'part_kind': 'builtin-tool-call',
                'tool_name': 'tool_search',
                'args': {'foo': 'bar'},
                'tool_call_id': 'c3',
            },
        ],
    }
    resp = message(ModelMessagesTypeAdapter.validate_python([raw]), ModelResponse)
    assert isinstance(resp.parts[0], NativeToolSearchCallPart)
    assert isinstance(resp.parts[1], NativeToolSearchReturnPart)
    # Unrecognized `tool_name` (and unset `tool_kind`) falls through to the base class.
    assert isinstance(resp.parts[2], NativeToolCallPart)
    assert not isinstance(resp.parts[2], NativeToolSearchCallPart)
    # User-defined collision on `tool_name='tool_search'` without `tool_kind` stays base.
    assert type(resp.parts[3]) is NativeToolCallPart
    assert resp.parts[3].args == {'foo': 'bar'}


def test_model_response_instance_round_trip_promotes_typed_subclasses() -> None:
    """Re-validation of a `ModelResponse` instance preserves typed builtin parts."""

    resp = ModelResponse(
        parts=[
            NativeToolSearchCallPart(args={'queries': ['x']}, tool_call_id='c1'),
            NativeToolSearchReturnPart(
                content={'discovered_tools': [{'name': 'foo'}]},
                tool_call_id='c1',
            ),
            NativeToolCallPart(tool_name='web_search', args={}, tool_call_id='c2'),
        ]
    )
    revalidated = message(ModelMessagesTypeAdapter.validate_python([resp]), ModelResponse)
    assert isinstance(revalidated.parts[0], NativeToolSearchCallPart)
    assert isinstance(revalidated.parts[1], NativeToolSearchReturnPart)
    assert isinstance(revalidated.parts[2], NativeToolCallPart)


async def test_tool_search_toolset_protects_user_collision_on_builtin_tool_name() -> None:
    """A user-emitted `NativeToolReturnPart` with `tool_name='tool_search'` (no typed
    subclass, no `tool_kind`) is left alone — discoveries are only surfaced from typed
    `NativeToolSearchReturnPart` instances. This is the typed-trust contract: the
    framework constructs typed subclasses; user collisions on names alone don't get
    treated as our search payload."""

    history: list[ModelMessage] = [
        ModelResponse(
            parts=[
                # Framework-emitted: typed subclass surfaces discoveries.
                NativeToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'calculate_mortgage'}]},
                    tool_call_id='c1',
                ),
                # User collision on the name with a base part — `tool_kind=None`, not a typed
                # subclass: NOT surfaced.
                NativeToolReturnPart(
                    tool_name='tool_search',
                    content={'discovered_tools': [{'name': 'should_not_surface'}]},
                    tool_call_id='c2',
                ),
            ],
        ),
    ]

    discovered = parse_discovered_tools(history)
    assert 'calculate_mortgage' in discovered
    assert 'should_not_surface' not in discovered


async def test_local_tool_search_stream_emits_typed_call_part_from_first_event() -> None:
    """Streaming counterpart to the non-streaming typed-parts test. The model streams a
    `search_tools` call name + args delta-by-delta; `ModelResponsePartsManager` materializes
    the call part as the typed `ToolSearchCallPart` from the first `PartStartEvent` rather
    than only after a post-stream pass. This relies on the parts manager receiving
    `model_request_parameters` (set on `StreamedResponse.__post_init__`) so it can look up
    `ToolDefinition.tool_kind` for the called tool name.

    Forces the local-fallback path by using a model that doesn't claim native
    `ToolSearchTool` support — otherwise the swap drops `search_tools` from
    `function_tools` (Rule 1) on the assumption the model handles tool search
    server-side via the native wire shape.
    """

    class NoNativeToolSearchModel(FunctionModel):
        """A `FunctionModel` that drops `ToolSearchTool` from its supported builtins so the
        framework routes through the local `search_tools` function tool rather than the
        native wire shape."""

        @classmethod
        def supported_native_tools(cls) -> frozenset[type[AbstractNativeTool]]:
            return frozenset(super().supported_native_tools()) - {ToolSearchTool}

    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain(defer_loading=True)
    def calculate_mortgage(principal: float, rate: float, years: int) -> str:  # pragma: no cover
        return f'${principal * rate * years:.2f}'

    call_count = 0

    async def stream_function(_messages: list[ModelMessage], _info: AgentInfo) -> AsyncIterator[DeltaToolCalls | str]:
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            yield {0: DeltaToolCall(name='search_tools', tool_call_id='c1')}
            yield {0: DeltaToolCall(json_args='{"queries":')}
            yield {0: DeltaToolCall(json_args='["mortgage"]}')}
        else:
            yield 'done'

    agent = Agent(
        NoNativeToolSearchModel(stream_function=stream_function), toolsets=[toolset], capabilities=[ToolSearch()]
    )

    typed_at_start: list[bool] = []

    async def event_stream_handler(_ctx: RunContext, stream: AsyncIterable[AgentStreamEvent]) -> None:
        async for event in stream:
            if (
                isinstance(event, PartStartEvent)
                and isinstance(event.part, ToolCallPart)
                and event.part.tool_name == 'search_tools'
            ):
                typed_at_start.append(isinstance(event.part, ToolSearchCallPart))

    await agent.run('find a mortgage tool', event_stream_handler=event_stream_handler)

    # The first PartStartEvent for the search_tools call already carries the typed identity.
    assert typed_at_start, 'expected a PartStartEvent for search_tools during streaming'
    assert all(typed_at_start), f'expected typed `ToolSearchCallPart` from first event; got {typed_at_start}'


async def test_local_tool_search_dispatch_produces_typed_parts() -> None:
    """End-to-end typed identity for the local `search_tools` path: the model emits a
    base `ToolCallPart`, the framework promotes it to `ToolSearchCallPart` via the
    declared `ToolDefinition.tool_kind`, dispatches to `ToolSearchToolset`, and constructs
    a typed `ToolSearchReturnPart`. Both halves of the exchange carry the typed identity
    so multi-turn discovery parsing and cross-provider replay see typed parts everywhere.

    Reported by Devin's review of commit 53eb27b06 for the return side: previously the
    framework constructed a base `ToolReturnPart` (no `tool_kind`), and neither
    `parse_discovered_tools`' isinstance check nor the legacy-metadata reader caught
    it, so previously-discovered tools reverted to hidden on every subsequent turn.
    """
    toolset: FunctionToolset = FunctionToolset()

    @toolset.tool_plain(defer_loading=True)
    def calculate_mortgage(principal: float, rate: float, years: int) -> str:  # pragma: no cover
        return f'${principal * rate * years:.2f}'

    call_count = 0

    def model_function(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return ModelResponse(parts=[ToolCallPart(tool_name='search_tools', args={'queries': ['mortgage']})])
        return ModelResponse(parts=[TextPart(content='done')])

    agent = Agent(FunctionModel(model_function), toolsets=[toolset], capabilities=[ToolSearch()])
    result = await agent.run('find a mortgage tool')

    # The framework-promoted call part is typed (via `_narrow_tool_call_parts` post-hook).
    search_calls = [
        part
        for msg in result.all_messages()
        if isinstance(msg, ModelResponse)
        for part in msg.parts
        if isinstance(part, ToolCallPart) and part.tool_name == 'search_tools'
    ]
    assert len(search_calls) == 1
    assert isinstance(search_calls[0], ToolSearchCallPart)
    assert search_calls[0].tool_kind == 'tool-search'

    # The framework-constructed return part is typed (via `_call_tool` dispatch hook).
    search_returns = [
        part
        for msg in result.all_messages()
        if isinstance(msg, ModelRequest)
        for part in msg.parts
        if isinstance(part, ToolReturnPart) and part.tool_name == 'search_tools'
    ]
    assert len(search_returns) == 1
    assert isinstance(search_returns[0], ToolSearchReturnPart)
    assert search_returns[0].tool_kind == 'tool-search'
    # And the typed content carries the discovery.
    content = search_returns[0].content
    assert {m['name'] for m in content['discovered_tools']} == {'calculate_mortgage'}


async def test_tool_search_toolset_replays_main_branch_legacy_shape() -> None:
    """Histories serialized on `main` (before this PR's typed-content shape) carry the
    discovered names on `ToolReturnPart.metadata['discovered_tools']` rather than on a
    typed `content`. They must continue to replay cleanly on the typed-parts shape so
    upgrading users don't lose discovered-tool state on the next turn.

    This is the wire shape on the `main` branch as of the merge-base.
    """

    history: list[ModelMessage] = [
        ModelResponse(
            parts=[ToolCallPart(tool_name='search_tools', args={'queries': ['mortgage']}, tool_call_id='c1')]
        ),
        ModelRequest(
            parts=[
                ToolReturnPart(
                    tool_name='search_tools',
                    # `main`-branch shape: structured discoveries live on `metadata`,
                    # `content` is just the user-visible string the model sees.
                    content='Found 1 tool: calculate_mortgage',
                    tool_call_id='c1',
                    metadata={'discovered_tools': ['calculate_mortgage']},
                ),
            ],
        ),
    ]
    discovered = parse_discovered_tools(history)
    assert discovered == {'calculate_mortgage'}


def test_synthetic_injection_translates_builtin_to_local_tool_search_parts() -> None:
    """Cross-provider replay end-to-end: a `NativeToolSearch*Part` carried over from
    a prior native turn is translated into the local-shape typed parts so a non-native
    adapter can replay it as a normal `search_tools` function-call exchange. The
    `parse_discovered_tools` then surfaces the discoveries via the
    discriminated-union dispatch."""
    history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='Find me a mortgage tool.')]),
        ModelResponse(
            parts=[
                NativeToolSearchCallPart(
                    args={'queries': ['mortgage']},
                    tool_call_id='c1',
                    provider_name='anthropic',
                    provider_details={'strategy': 'bm25'},
                ),
                NativeToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'calculate_mortgage'}]},
                    tool_call_id='c1',
                    provider_name='anthropic',
                ),
            ],
        ),
    ]

    translated = synthesize_local_tool_search_messages(history)

    # The user prompt request passes through unchanged.
    assert translated[0] is history[0]

    # The response now carries a local `ToolSearchCallPart` (typed `ToolCallPart` subclass),
    # and the return part has been lifted onto a fresh trailing `ModelRequest`.
    response = message(translated, ModelResponse, index=1)
    assert len(response.parts) == 1
    call_part = message_part(translated, ToolSearchCallPart, message_index=1)
    # Subclass of `ToolCallPart`, NOT `NativeToolSearchCallPart`.
    assert isinstance(call_part, ToolCallPart)
    assert not isinstance(call_part, NativeToolSearchCallPart)
    assert call_part.tool_name == 'search_tools'
    assert call_part.args == {'queries': ['mortgage']}

    return_part = message_part(translated, ToolSearchReturnPart, message_index=2)
    assert isinstance(return_part, ToolReturnPart)
    assert not isinstance(return_part, NativeToolSearchReturnPart)
    assert return_part.tool_name == 'search_tools'
    assert return_part.content == {'discovered_tools': [{'name': 'calculate_mortgage'}]}

    # And the toolset's parser surfaces the discovery off the translated history.
    discovered = parse_discovered_tools(translated)
    assert discovered == {'calculate_mortgage'}


def test_synthesize_local_promotes_base_tool_return_with_tool_kind_in_request() -> None:
    """`synthesize_local_tool_search_messages` also reaches into existing `ModelRequest`
    parts: a base `ToolReturnPart` carrying `tool_kind='tool-search'` (e.g. one
    constructed manually before going through the discriminator) is promoted to its
    typed `ToolSearchReturnPart` subclass in place. Mirrors the response-side
    promotion so cross-provider history stays uniformly typed regardless of where
    the parts originated."""

    history: list[ModelMessage] = [
        ModelResponse(parts=[ToolCallPart(tool_name='search_tools', args={'queries': ['a']}, tool_call_id='c1')]),
        ModelRequest(
            parts=[
                ToolReturnPart(
                    tool_name='search_tools',
                    content={'discovered_tools': [{'name': 'foo'}]},
                    tool_call_id='c1',
                    tool_kind='tool-search',
                ),
            ],
        ),
    ]
    translated = synthesize_local_tool_search_messages(history)
    part = message_part(translated, ToolSearchReturnPart, message_index=1)
    assert part.content == {'discovered_tools': [{'name': 'foo'}]}


async def test_tool_search_toolset_uses_custom_parameter_description() -> None:
    """`ToolSearch(parameter_description=...)` flows through to the local `search_tools`
    function tool's `queries` parameter description on the wire — verifies the
    custom-description branch in `_build_search_args_schema` rebuilds the JSON
    schema rather than reusing the default."""
    cap = ToolSearch(parameter_description='custom queries hint')
    base_toolset = _create_function_toolset()
    wrapped = cap.get_wrapper_toolset(base_toolset)
    ctx = _build_run_context(None)
    tools = await wrapped.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]
    schema = search_tool.tool_def.parameters_json_schema
    assert schema['properties']['queries']['description'] == 'custom queries hint'


def test_prepare_messages_translates_on_non_native_model() -> None:
    """`Model.prepare_messages` is the centralized hook that runs before the adapter's
    message-prep on every request. On a model whose profile doesn't include
    `ToolSearchTool` in `supported_native_tools`, the hook translates any prior
    server-side tool-search exchange into the local-shape typed parts so the adapter
    sees a normal `search_tools` function-call exchange.

    The single `ModelResponse(call+return)` carrying the inline server-side result
    splits into `ModelResponse(call) + ModelRequest(return)`."""
    # Default `TestModel` excludes `ToolSearchTool` from `supported_native_tools`.
    model = TestModel()
    assert ToolSearchTool not in model.profile.get('supported_native_tools', SUPPORTED_NATIVE_TOOLS)

    history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='Find me a mortgage tool.')]),
        ModelResponse(
            parts=[
                NativeToolSearchCallPart(
                    args={'queries': ['mortgage']},
                    tool_call_id='c1',
                    provider_name='anthropic',
                ),
                NativeToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'calculate_mortgage'}]},
                    tool_call_id='c1',
                    provider_name='anthropic',
                ),
            ],
        ),
    ]

    prepared = model.prepare_messages(history)

    # Original 2 messages became 3: user prompt, response with local call,
    # request carrying the lifted return.
    assert len(prepared) == 3
    assert prepared[0] is history[0]

    response = message(prepared, ModelResponse, index=1)
    assert len(response.parts) == 1
    call_part = message_part(prepared, ToolSearchCallPart, message_index=1)
    assert not isinstance(call_part, NativeToolSearchCallPart)
    assert call_part.tool_name == 'search_tools'

    return_part = message_part(prepared, ToolSearchReturnPart, message_index=2)
    assert not isinstance(return_part, NativeToolSearchReturnPart)
    assert return_part.tool_name == 'search_tools'
    assert return_part.content == {'discovered_tools': [{'name': 'calculate_mortgage'}]}


def test_prepare_messages_passes_through_on_native_model() -> None:
    """A model whose profile *does* include `ToolSearchTool` in
    `supported_native_tools` keeps the prior exchange as-is — the native adapter
    knows how to ship the typed builtin parts back on the wire."""

    class NativeToolSearchTestModel(TestModel):
        @classmethod
        def supported_native_tools(cls):
            return frozenset({ToolSearchTool})

    model = NativeToolSearchTestModel()
    assert ToolSearchTool in model.profile.get('supported_native_tools', SUPPORTED_NATIVE_TOOLS)

    history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='Find me a mortgage tool.')]),
        ModelResponse(
            parts=[
                NativeToolSearchCallPart(
                    args={'queries': ['mortgage']},
                    tool_call_id='c1',
                    provider_name='anthropic',
                ),
                NativeToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'calculate_mortgage'}]},
                    tool_call_id='c1',
                    provider_name='anthropic',
                ),
            ],
        ),
    ]

    prepared = model.prepare_messages(history)

    assert prepared is history


def test_narrow_type_local_promotes_with_tool_kind_set() -> None:
    """A `ToolCallPart` with `tool_kind='tool-search'` promotes to `ToolSearchCallPart`.

    Promotion is keyed on `tool_kind`, not `tool_name` — a framework-emitted call carries
    `tool_kind='tool-search'` so it round-trips as the typed subclass.
    """

    part = ToolCallPart(
        tool_name='search_tools',
        args={'queries': ['mortgage']},
        tool_call_id='c1',
        tool_kind='tool-search',
    )
    narrowed = ToolCallPart.narrow_type(part)
    assert isinstance(narrowed, ToolSearchCallPart)
    assert narrowed.args == {'queries': ['mortgage']}


def test_narrow_type_local_passthrough_when_already_narrowed() -> None:
    """Narrowing an already-typed `ToolSearchCallPart` returns the input instance."""
    part = ToolSearchCallPart(args={'queries': ['x']}, tool_call_id='c1')

    assert ToolCallPart.narrow_type(part) is part


def test_pydantic_validation_accepts_search_tools_collision_when_tool_kind_unset() -> None:
    """A user-defined tool literally named `search_tools` with arbitrary args is safe.

    Dispatch is by `tool_kind`, not `tool_name`, so the absence of `tool_kind` keeps
    the part as a base `ToolReturnPart` regardless of args shape — no accidental
    auto-promotion to `ToolSearchReturnPart`, no spurious shape-validation failure.
    """

    raw = [
        {
            'kind': 'request',
            'parts': [
                {
                    'part_kind': 'tool-return',
                    'tool_name': 'search_tools',
                    # Arbitrary user-tool shape.
                    'content': {'unrelated': 'data', 'definitely_not_discovered_tools': 42},
                    'tool_call_id': 'c1',
                },
            ],
        },
    ]
    [req] = ModelMessagesTypeAdapter.validate_python(raw)
    [part] = req.parts
    assert type(part) is ToolReturnPart
    assert part.tool_kind is None
    assert part.content == {'unrelated': 'data', 'definitely_not_discovered_tools': 42}


def test_pydantic_validation_promotes_local_tool_return_with_tool_kind_set() -> None:
    """A serialized `tool-return` carrying `tool_kind='tool-search'` and a typed-shape
    `discovered_tools` payload is promoted to `ToolSearchReturnPart` by Pydantic's
    discriminated-union dispatch — the discriminator routes (part_kind, tool_kind)
    to the typed tag so deserialization rebuilds the typed subclass directly."""

    raw = [
        {
            'kind': 'request',
            'parts': [
                {
                    'part_kind': 'tool-return',
                    'tool_name': 'search_tools',
                    'tool_kind': 'tool-search',
                    'content': {'discovered_tools': [{'name': 'foo'}]},
                    'tool_call_id': 'c1',
                },
            ],
        },
    ]
    part = message_part(ModelMessagesTypeAdapter.validate_python(raw), ToolSearchReturnPart)
    assert part.content == {'discovered_tools': [{'name': 'foo'}]}


def test_pydantic_validation_accepts_search_tools_string_content_collision() -> None:
    """A user tool literally named `search_tools` returning plain text deserializes cleanly.

    Without `tool_kind`, the part stays a base `ToolReturnPart` — the str content survives
    intact. This is the user-tool-collision-tolerance contract: dispatch never promotes
    based on `tool_name` alone.
    """

    raw = [
        {
            'kind': 'request',
            'parts': [
                {
                    'part_kind': 'tool-return',
                    'tool_name': 'search_tools',
                    'content': 'hello world',
                    'tool_call_id': 'c1',
                },
            ],
        },
    ]
    [request] = ModelMessagesTypeAdapter.validate_python(raw)
    [part] = request.parts
    assert type(part) is ToolReturnPart
    assert part.tool_kind is None
    assert part.content == 'hello world'


def test_synthesize_local_from_native_call_str_args_passthrough() -> None:
    """Streaming partial-args (`str`) are passed through unchanged when translating."""

    part = NativeToolSearchCallPart(args='{"queries":', tool_call_id='c1')
    result = synthesize_local_from_native_call(part)
    assert result.args == '{"queries":'
    assert result.tool_call_id == 'c1'


def test_synthesize_local_from_native_call_none_args_falls_through() -> None:
    """`None` args remain `None` after translation."""

    part = NativeToolSearchCallPart(args=None, tool_call_id='c1')
    result = synthesize_local_from_native_call(part)
    assert result.args is None


def test_synthesize_messages_response_with_only_call_part_no_lift() -> None:
    """A response with only a `NativeToolSearchCallPart` (no return — streaming case)
    translates the call but doesn't synthesize a trailing `ModelRequest`."""

    history: list[ModelMessage] = [
        ModelResponse(parts=[NativeToolSearchCallPart(args={'queries': ['x']}, tool_call_id='c1')]),
    ]
    result = synthesize_local_tool_search_messages(history)
    assert len(result) == 1
    response = message(result, ModelResponse)
    assert len(response.parts) == 1
    assert isinstance(response.parts[0], ToolSearchCallPart)


def test_synthesize_messages_response_with_only_return_part_no_response_kept() -> None:
    """A response with only a `NativeToolSearchReturnPart` (no remaining parts) — the
    response is dropped since it'd be empty, and the return is lifted onto a fresh request."""

    history: list[ModelMessage] = [
        ModelResponse(
            parts=[
                NativeToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'foo'}]},
                    tool_call_id='c1',
                ),
            ],
        ),
    ]
    result = synthesize_local_tool_search_messages(history)
    assert len(result) == 1
    request = message(result, ModelRequest)
    assert len(request.parts) == 1
    message_part(result, ToolSearchReturnPart)


def test_synthesize_messages_request_with_unrelated_tool_return_passthrough() -> None:
    """A `ToolReturnPart` with `tool_name != 'search_tools'` doesn't get promoted —
    the request is returned unchanged."""

    request = ModelRequest(parts=[ToolReturnPart(tool_name='get_weather', content='sunny', tool_call_id='c1')])
    result = synthesize_local_tool_search_messages([request])
    assert len(result) == 1
    assert result[0] is request


def test_synthesize_messages_response_with_search_then_downstream_tool_call_splits_4_messages() -> None:
    """Native turn with `[Text, BuiltinSearchCall, BuiltinSearchReturn, ToolCall(weather)]`
    must split into a coherent local-shape sequence: response[Text, ToolSearchCall],
    request[ToolSearchReturn], response[ToolCall(weather)], (passthrough) request[ToolReturn(weather)].

    Currently we keep the downstream `ToolCall(weather)` on the same response as the
    `ToolSearchCall`, which is incoherent (model "called weather before seeing search
    results") and produces consecutive `ModelRequest`s after the lifted return —
    Devin's observation. Splitting at every `NativeToolSearchReturn` boundary fixes
    both: the timeline reads correctly and the lifted return doesn't collide with the
    next request.
    """
    history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='Look up something then call it.')]),
        ModelResponse(
            parts=[
                TextPart(content='Searching first.'),
                NativeToolSearchCallPart(
                    args={'queries': ['weather']},
                    tool_call_id='search1',
                    provider_name='anthropic',
                ),
                NativeToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'get_weather'}]},
                    tool_call_id='search1',
                    provider_name='anthropic',
                ),
                ToolCallPart(tool_name='get_weather', args={'city': 'NYC'}, tool_call_id='wx1'),
            ],
        ),
        ModelRequest(parts=[ToolReturnPart(tool_name='get_weather', content='sunny', tool_call_id='wx1')]),
    ]

    result = synthesize_local_tool_search_messages(history)

    # 5 messages: user request, response[Text, ToolSearchCall], request[ToolSearchReturn],
    # response[ToolCall(weather)], request[ToolReturn(weather)] (the original).
    assert len(result) == 5

    first_req = message(result, ModelRequest)
    assert first_req is history[0]

    # First synthetic response: text + search call only — NOT the downstream weather call.
    first_resp = message(result, ModelResponse, index=1)
    assert len(first_resp.parts) == 2
    assert isinstance(first_resp.parts[0], TextPart)
    assert isinstance(first_resp.parts[1], ToolSearchCallPart)
    # No `ToolCallPart(weather)` snuck onto this response.
    assert not any(isinstance(p, ToolCallPart) and not isinstance(p, ToolSearchCallPart) for p in first_resp.parts)

    # Lifted search return as a fresh request.
    search_return_req = message(result, ModelRequest, index=2)
    assert len(search_return_req.parts) == 1
    assert isinstance(search_return_req.parts[0], ToolSearchReturnPart)

    # Second synthetic response: weather call only.
    second_resp = message(result, ModelResponse, index=3)
    assert len(second_resp.parts) == 1
    weather_call = message_part(result, ToolCallPart, message_index=3)
    assert weather_call.tool_name == 'get_weather'

    # Original weather-return request flows naturally — no consecutive `ModelRequest`s.
    last_req = message(result, ModelRequest, index=4)
    assert last_req is history[2]


def test_synthesize_messages_devins_consecutive_request_repro() -> None:
    """Regression: synthesis must not produce two consecutive `ModelRequest`s.

    Reproduces Devin's bug report exactly: native search exchange immediately followed
    by a regular tool call within the same `ModelResponse`, then a `ModelRequest` for
    the regular tool's return. The proper splitter inserts a synthetic `ModelResponse`
    between the lifted search return and the original tool-return request, so message
    roles alternate correctly for adapters with strict user/assistant alternation.
    """
    history: list[ModelMessage] = [
        ModelResponse(
            parts=[
                NativeToolSearchCallPart(args={'queries': ['x']}, tool_call_id='s1'),
                NativeToolSearchReturnPart(content={'discovered_tools': []}, tool_call_id='s1'),
                ToolCallPart(tool_name='get_weather', args={}, tool_call_id='w1'),
            ],
        ),
        ModelRequest(parts=[ToolReturnPart(tool_name='get_weather', content='ok', tool_call_id='w1')]),
    ]

    result = synthesize_local_tool_search_messages(history)

    # Walk and verify no two consecutive entries are both `ModelRequest`.
    for i in range(len(result) - 1):
        if isinstance(result[i], ModelRequest):
            assert not isinstance(result[i + 1], ModelRequest), f'Consecutive ModelRequests at index {i}: {result}'


def test_synthesize_messages_multiple_search_rounds_in_one_response() -> None:
    """Two server-side search rounds inside a single native `ModelResponse` split into
    two response/request pairs, preserving order and not bundling them onto one response.
    """
    history: list[ModelMessage] = [
        ModelResponse(
            parts=[
                NativeToolSearchCallPart(args={'queries': ['a']}, tool_call_id='s1'),
                NativeToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'tool_a'}]},
                    tool_call_id='s1',
                ),
                NativeToolSearchCallPart(args={'queries': ['b']}, tool_call_id='s2'),
                NativeToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'tool_b'}]},
                    tool_call_id='s2',
                ),
            ],
        ),
    ]

    result = synthesize_local_tool_search_messages(history)

    # 4 messages: response[call_a], request[return_a], response[call_b], request[return_b].
    assert len(result) == 4
    call_part_1 = message_part(result, ToolSearchCallPart)
    assert call_part_1.tool_call_id == 's1'

    return_part_1 = message_part(result, ToolSearchReturnPart, message_index=1)
    assert return_part_1.tool_call_id == 's1'

    call_part_2 = message_part(result, ToolSearchCallPart, message_index=2)
    assert call_part_2.tool_call_id == 's2'

    return_part_2 = message_part(result, ToolSearchReturnPart, message_index=3)
    assert return_part_2.tool_call_id == 's2'


def test_synthesize_messages_metadata_kept_on_first_split_only() -> None:
    """Splitting one native `ModelResponse` into multiple responses must not duplicate
    its identity-level metadata (`provider_response_id`, usage). The first split keeps
    the original identity; subsequent splits get fresh/blank fields so downstream
    consumers don't double-count usage or find two responses for the same API call.
    """

    history: list[ModelMessage] = [
        ModelResponse(
            parts=[
                NativeToolSearchCallPart(args={'queries': ['x']}, tool_call_id='s1'),
                NativeToolSearchReturnPart(content={'discovered_tools': []}, tool_call_id='s1'),
                ToolCallPart(tool_name='get_weather', args={}, tool_call_id='w1'),
            ],
            usage=RequestUsage(input_tokens=100, output_tokens=50),
            provider_response_id='msg_real_anthropic_id',
            provider_name='anthropic',
            model_name='claude-sonnet-4-5',
        ),
    ]

    result = synthesize_local_tool_search_messages(history)

    # Two responses out (split around the search return).
    responses = [m for m in result if isinstance(m, ModelResponse)]
    assert len(responses) == 2

    # First split keeps full metadata.
    assert responses[0].provider_response_id == 'msg_real_anthropic_id'
    assert responses[0].usage.input_tokens == 100
    assert responses[0].usage.output_tokens == 50

    # Second split gets cleared identity to avoid double-counting / duplicate lookup.
    assert responses[1].provider_response_id is None
    assert responses[1].usage.input_tokens == 0
    assert responses[1].usage.output_tokens == 0


def test_prepare_messages_then_clean_history_merges_consecutive_requests() -> None:
    """Regression: the bare `[SearchCall, SearchReturn]` response shape — common when a model
    finishes a turn right after server-side search results — splits into `Response + Request`,
    which collides with the next `ModelRequest` in the history. `_clean_message_history` must
    run *after* `prepare_messages` so the splitter's synthetic `Request([SearchReturn])` and
    the original `Request([UserPromptPart])` merge into a single `ModelRequest`, preserving
    strict user/assistant alternation for adapters that require it.
    """

    history: list[ModelMessage] = [
        ModelResponse(
            parts=[
                NativeToolSearchCallPart(args={'queries': ['x']}, tool_call_id='s1'),
                NativeToolSearchReturnPart(content={'discovered_tools': []}, tool_call_id='s1'),
            ],
        ),
        ModelRequest(parts=[UserPromptPart(content='follow-up')]),
    ]

    # Mirror `_make_request`'s post-fix order: synthesize first, then clean. Without the cleanup
    # pass, the synthesizer produces `Response, Request, Request` (a synthetic request for the
    # search return next to the original request); the clean pass merges those two requests so
    # adapters with strict user/assistant alternation see `Response, Request`.
    after_synthesis = synthesize_local_tool_search_messages(history)
    assert [type(m).__name__ for m in after_synthesis] == ['ModelResponse', 'ModelRequest', 'ModelRequest']

    cleaned = _clean_message_history(after_synthesis)
    assert [type(m).__name__ for m in cleaned] == ['ModelResponse', 'ModelRequest']

    # The merged request carries both the synthetic search return and the original user prompt,
    # with the tool return part sorted ahead of the user prompt.
    last = message(cleaned, ModelRequest, index=-1)
    assert isinstance(last.parts[0], ToolSearchReturnPart)
    assert isinstance(last.parts[1], UserPromptPart)


def test_narrow_type_local_return_passthrough_when_already_narrowed() -> None:
    """Narrowing an already-typed `ToolSearchReturnPart` returns the input instance."""
    part = ToolSearchReturnPart(content={'discovered_tools': []}, tool_call_id='c1')
    assert ToolReturnPart.narrow_type(part) is part


def test_narrow_type_local_return_promotes_with_tool_kind_set() -> None:
    """A base `ToolReturnPart` with `tool_kind='tool-search'` and a valid typed-content
    payload is promoted to `ToolSearchReturnPart` by `narrow_type`. Mirror of the
    builtin-side promotion test, exercising the local (function-tool) variant."""
    base = ToolReturnPart(
        tool_name='search_tools',
        content={'discovered_tools': [{'name': 'foo'}]},
        tool_call_id='c1',
        tool_kind='tool-search',
    )
    narrowed = ToolReturnPart.narrow_type(base)
    assert isinstance(narrowed, ToolSearchReturnPart)
    assert narrowed.content == {'discovered_tools': [{'name': 'foo'}]}


def test_narrow_type_no_tool_kind_returns_input_unchanged_for_local_and_builtin_returns() -> None:
    """`narrow_type` is a no-op when `tool_kind` is `None` — the user-tool default —
    on both `ToolReturnPart` and `NativeToolReturnPart`. This is the early-exit
    branch that keeps user tools untouched without consulting the registry."""
    local = ToolReturnPart(tool_name='foo', content='bar', tool_call_id='c1')
    assert ToolReturnPart.narrow_type(local) is local
    builtin = NativeToolReturnPart(tool_name='foo', content='bar', tool_call_id='c1')
    assert NativeToolReturnPart.narrow_type(builtin) is builtin


def test_model_request_part_discriminator_recognizes_tool_search_return_instance() -> None:
    """The request-part discriminator returns the typed tag when called with a
    `ToolSearchReturnPart` instance.

    Pydantic's discriminated-union fast path bypasses the discriminator when the input
    already matches one of the tagged variants by isinstance, so this exercises the
    function directly rather than via `ModelMessagesTypeAdapter`.
    """

    part = ToolSearchReturnPart(content={'discovered_tools': []}, tool_call_id='c1')
    assert _model_request_part_discriminator(part) == 'tool-search-return'


def test_model_response_part_discriminator_recognizes_local_call_dict_dispatch() -> None:
    """A dict-shaped `ToolCallPart` with `tool_kind='tool-search'` gets dispatched to
    `ToolSearchCallPart` via the discriminator (covers the `'tool-call'` branch)."""

    raw = [
        {
            'kind': 'response',
            'parts': [
                {
                    'part_kind': 'tool-call',
                    'tool_name': 'search_tools',
                    'tool_kind': 'tool-search',
                    'args': {'queries': ['x']},
                    'tool_call_id': 'c1',
                },
            ],
        },
    ]
    message_part(ModelMessagesTypeAdapter.validate_python(raw), ToolSearchCallPart)


def test_model_response_part_discriminator_passthrough_for_unknown_part_kind() -> None:
    """Instance dispatch falls through to `getattr(v, 'part_kind', ...)` for other types."""

    resp = ModelResponse(parts=[TextPart(content='hello')])
    message_part(ModelMessagesTypeAdapter.validate_python([resp]), TextPart)


def test_model_response_part_discriminator_recognizes_typed_instances() -> None:
    """The response-part discriminator returns the typed tag for each typed-instance branch.

    Pydantic's discriminated-union fast path bypasses the discriminator when the input
    already matches one of the tagged variants by isinstance, so the instance branches
    in `_model_response_part_discriminator` are only reachable by calling the function
    directly. This locks in the contract for any future caller (or pydantic version
    that changes its short-circuit behavior).
    """

    builtin_call = NativeToolSearchCallPart(args={'queries': ['x']}, tool_call_id='c1', provider_name='anthropic')
    assert _model_response_part_discriminator(builtin_call) == 'builtin-tool-search-call'

    builtin_return = NativeToolSearchReturnPart(
        content={'discovered_tools': []},
        tool_call_id='c1',
        provider_name='anthropic',
    )
    assert _model_response_part_discriminator(builtin_return) == 'builtin-tool-search-return'

    local_call = ToolSearchCallPart(args={'queries': ['x']}, tool_call_id='c1')
    assert _model_response_part_discriminator(local_call) == 'tool-search-call'


def test_discriminator_unknown_tool_kind_falls_through_to_part_kind() -> None:
    """Dict-form parts with an unregistered `tool_kind` fall through to the bare `part_kind`.

    Exercises the registry-miss branch in both discriminator functions: `_TYPED_PART_TAGS`
    doesn't contain `(part_kind, 'unknown-kind')`, so the discriminator returns the bare
    `part_kind` rather than a typed-subclass tag.

    Calls the discriminator directly because constructing a valid ModelMessage with
    `tool_kind='unknown-kind'` would fail Pydantic's `ToolPartKind` Literal validation
    upstream — the registry-miss branch is internal logic, not a deserialization path
    that any well-formed input would take.
    """

    return_raw = {
        'part_kind': 'tool-return',
        'tool_name': 'something',
        'tool_kind': 'unknown-kind',
        'content': 'hello',
        'tool_call_id': 'c1',
    }
    assert _model_request_part_discriminator(return_raw) == 'tool-return'

    call_raw = {
        'part_kind': 'tool-call',
        'tool_name': 'something',
        'tool_kind': 'unknown-kind',
        'args': {'x': 1},
        'tool_call_id': 'c1',
    }
    assert _model_response_part_discriminator(call_raw) == 'tool-call'


def test_typed_call_part_accessors_return_typed_shapes() -> None:
    """`typed_args` and `queries` on typed call parts read the parsed args.

    Covers both the local-fallback (`ToolSearchCallPart`) and native server-side
    (`NativeToolSearchCallPart`) variants — they're symmetric.
    """

    local_call = ToolSearchCallPart(args={'queries': ['weather', 'github']}, tool_call_id='c1')
    assert local_call.typed_args == {'queries': ['weather', 'github']}
    assert local_call.queries == ['weather', 'github']

    builtin_call = NativeToolSearchCallPart(args={'queries': ['weather']}, tool_call_id='c2', provider_name='anthropic')
    assert builtin_call.typed_args == {'queries': ['weather']}
    assert builtin_call.queries == ['weather']


def test_typed_call_part_typed_args_returns_none_for_unparsed_args() -> None:
    """`typed_args` returns `None` when args haven't been finalized yet.

    Covers the streaming-partial path: `args=None`, partial JSON strings, and
    non-dict JSON values all yield `None` (the contract for streaming-not-yet-ready
    or unexpected shapes). Exercises both typed call part subclasses.
    """

    for cls in (ToolSearchCallPart, NativeToolSearchCallPart):
        kwargs: dict[str, Any] = {'tool_call_id': 'c1'}
        if cls is NativeToolSearchCallPart:
            kwargs['provider_name'] = 'anthropic'

        none_part = cls(args=None, **kwargs)
        assert none_part.typed_args is None
        assert none_part.queries == []

        # Partial JSON string raises during parsing → None.
        partial_part = cls(args='{"queries": ["wea', **kwargs)
        assert partial_part.typed_args is None
        assert partial_part.queries == []

        # Valid JSON that parses to a non-dict (e.g. a bare string) → None.
        scalar_part = cls(args='"just a string"', **kwargs)
        assert scalar_part.typed_args is None
        assert scalar_part.queries == []

        # Valid JSON dict → typed_args populated.
        complete_part = cls(args='{"queries": ["x"]}', **kwargs)
        assert complete_part.typed_args == {'queries': ['x']}
        assert complete_part.queries == ['x']


def test_builtin_tool_search_return_part_message_accessor() -> None:
    """`message` on `NativeToolSearchReturnPart` reads `content.get('message')`.

    The native server-side path doesn't currently populate `message` (Anthropic emits
    its own error/result blocks), so this accessor exists for symmetry with the local
    return part. Exercise it directly to lock in the contract.
    """

    with_message = NativeToolSearchReturnPart(
        content={'discovered_tools': [], 'message': 'no matches'},
        tool_call_id='c1',
        provider_name='anthropic',
    )
    assert with_message.message == 'no matches'

    without_message = NativeToolSearchReturnPart(
        content={'discovered_tools': [{'name': 'foo'}]},
        tool_call_id='c2',
        provider_name='anthropic',
    )
    assert without_message.message is None


async def test_tool_search_toolset_async_search_fn_is_awaited() -> None:
    """Custom search functions can be `async`; the toolset awaits them."""

    async def async_match(_ctx: RunContext, _queries: Sequence[str], tools: Sequence[ToolDefinition]) -> Sequence[str]:
        return [t.name for t in tools]

    ts = ToolSearchToolset(wrapped=_create_function_toolset(), search_fn=async_match)
    ctx = _build_run_context(None)
    tools = await ts.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]
    result = await ts.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['weather']}, ctx, search_tool)
    return_value = cast(dict[str, Any], result)
    discovered_names = {match['name'] for match in return_value['discovered_tools']}
    # `_create_function_toolset` registers a fixed set of deferred tools; verify the
    # async search function received the corpus and returned discoverable names.
    assert 'calculate_mortgage' in discovered_names


def test_anthropic_custom_replay_blocks_returns_message_on_empty_discovered() -> None:
    """When the typed return carries empty `discovered_tools` and a `message`, the
    helper returns `([], message)`. The `_map_message` flow then renders the message
    as a single text block (Anthropic rejects empty `tool_result.content`)."""
    pytest.importorskip('anthropic')

    empty = ToolSearchReturnPart(
        content={'discovered_tools': [], 'message': 'No matches; try other keywords.'},
        tool_call_id='c1',
    )
    refs, message = _build_custom_tool_search_replay_blocks(empty, tool_search_active=True, available_tool_names=set())
    assert refs == []
    assert message == 'No matches; try other keywords.'


def test_anthropic_custom_replay_blocks_skips_non_typed_returns() -> None:
    """A base `ToolReturnPart` (not a typed `ToolSearchReturnPart`) is left alone:
    helper returns `(None, None)` so the caller falls through to default text formatting.
    This is the typed-trust contract — the framework only re-shapes typed parts."""
    pytest.importorskip('anthropic')

    base_part = ToolReturnPart(
        tool_name='search_tools',
        content={'discovered_tools': [{'name': 'foo'}]},
        tool_call_id='c1',
    )
    refs, message = _build_custom_tool_search_replay_blocks(
        base_part, tool_search_active=True, available_tool_names={'foo'}
    )
    assert refs is None and message is None


def test_anthropic_replay_filters_stale_tool_references() -> None:
    """Anthropic rejects `tool_reference` blocks pointing at tools not in the request's
    `tools` list (e.g. an MCP whose connection failed this turn, dropping its tools
    from the corpus). Both replay paths — custom-callable `tool_result.content` and
    native `tool_search_tool_search_result.tool_references` — must filter against the
    set of tools the current turn will actually send."""
    pytest.importorskip('anthropic')

    discovered: list[ToolSearchMatch] = [
        {'name': 'still_here'},
        {'name': 'gone_this_turn'},
    ]
    content: ToolSearchReturnContent = {'discovered_tools': discovered}

    custom_part = ToolSearchReturnPart(content=content, tool_call_id='c1')
    refs, _ = _build_custom_tool_search_replay_blocks(
        custom_part, tool_search_active=True, available_tool_names={'still_here'}
    )
    assert refs == [{'tool_name': 'still_here', 'type': 'tool_reference'}]

    native_part = NativeToolSearchReturnPart(
        provider_name='anthropic',
        tool_call_id='srv_ok',
        content=content,
    )
    block = _build_tool_search_replay_block(native_part, 'srv_ok', available_tool_names={'still_here'})
    assert block == {
        'tool_use_id': 'srv_ok',
        'type': 'tool_search_tool_result',
        'content': {
            'type': 'tool_search_tool_search_result',
            'tool_references': [{'tool_name': 'still_here', 'type': 'tool_reference'}],
        },
    }


def test_anthropic_finalize_streamed_tool_search_call_part_with_canonical_dict_args() -> None:
    """Already-canonical `ToolSearchArgs` dict passes through unchanged — the typed
    contract guarantees `queries`, so re-running normalization would corrupt the data."""
    pytest.importorskip('anthropic')

    part = NativeToolSearchCallPart(
        args={'queries': ['mortgage']},
        tool_call_id='c1',
        provider_name='anthropic',
        provider_details={'strategy': 'bm25'},
    )
    result = _finalize_streamed_tool_search_call_part(part)
    assert result.args == {'queries': ['mortgage']}


def test_anthropic_finalize_streamed_tool_search_call_part_with_none_args() -> None:
    """`args=None` finalizes to a normalized empty `queries` list."""
    pytest.importorskip('anthropic')

    part = NativeToolSearchCallPart(args=None, tool_call_id='c1', provider_name='anthropic')
    result = _finalize_streamed_tool_search_call_part(part)
    assert isinstance(result.args, dict) and 'queries' in result.args


async def test_anthropic_map_message_empty_search_renders_message_text_block():
    """When custom-callable tool search returns no matches, `_map_message` emits the
    typed return as a single text-content `tool_result` block (not the default text
    fallthrough). Anthropic rejects empty `tool_result.content` arrays — this is the
    spec-compliant path for the custom-search empty-results case."""
    pytest.importorskip('anthropic')

    model = AnthropicModel(
        'claude-sonnet-4-6', provider=AnthropicProvider(anthropic_client=MockAnthropic.create_mock(()))
    )
    history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='find me a mortgage tool')]),
        ModelResponse(
            parts=[ToolCallPart(tool_name='search_tools', args={'queries': ['mortgage']}, tool_call_id='c1')]
        ),
        ModelRequest(
            parts=[
                ToolSearchReturnPart(
                    content={
                        'discovered_tools': [],
                        'message': 'No matching tools found. Try other keywords.',
                    },
                    tool_call_id='c1',
                ),
            ],
        ),
    ]
    params = ModelRequestParameters(
        function_tools=[],
        native_tools=[ToolSearchTool(strategy='custom')],
        allow_text_output=True,
    )
    _system, anthropic_messages = await model._map_message(history, params, AnthropicModelSettings())  # pyright: ignore[reportPrivateUsage]

    # Find the tool_result block across all user messages.
    tool_results: list[dict[str, Any]] = [
        c
        for m in anthropic_messages
        if m['role'] == 'user' and isinstance(m['content'], list)
        for c in cast(list[Any], m['content'])
        if isinstance(c, dict) and cast(dict[str, Any], c).get('type') == 'tool_result'
    ]
    [tool_result] = tool_results
    assert tool_result['content'] == [{'text': 'No matching tools found. Try other keywords.', 'type': 'text'}]
    assert tool_result['is_error'] is False


async def test_anthropic_map_message_replays_tool_search_call_without_queries():
    """A `NativeToolSearchCallPart` with `args=None` (streaming partial state, or a
    history fragment that never carried args) falls through to forwarding the empty
    `args_as_dict()` to the wire `input`. Covers the `else: wire_input = args_dict`
    branch where the cross-provider `queries` slot isn't populated."""
    pytest.importorskip('anthropic')

    model = AnthropicModel(
        'claude-sonnet-4-6', provider=AnthropicProvider(anthropic_client=MockAnthropic.create_mock(()))
    )
    history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='hello')]),
        ModelResponse(
            parts=[
                NativeToolSearchCallPart(
                    args=None,
                    tool_call_id='srv_1',
                    provider_name='anthropic',
                    provider_details={'strategy': 'bm25'},
                ),
                # Pair the call with a return so the orphan-drop pass keeps the call on the wire —
                # this test only exercises the `args=None` code path, not orphan handling.
                NativeToolSearchReturnPart(
                    content={'discovered_tools': []},
                    tool_call_id='srv_1',
                    provider_name='anthropic',
                ),
            ],
        ),
    ]
    params = ModelRequestParameters(
        function_tools=[],
        native_tools=[ToolSearchTool(strategy='bm25')],
        allow_text_output=True,
    )
    _system, anthropic_messages = await model._map_message(history, params, AnthropicModelSettings())  # pyright: ignore[reportPrivateUsage]

    [assistant_msg] = [m for m in anthropic_messages if m['role'] == 'assistant']
    assistant_content = cast(list[Any], assistant_msg['content'])
    server_tool_uses: list[dict[str, Any]] = [
        c for c in assistant_content if isinstance(c, dict) and cast(dict[str, Any], c).get('type') == 'server_tool_use'
    ]
    [server_tool_use] = server_tool_uses
    assert server_tool_use['input'] == {}


def test_openai_normalize_tool_search_args_empty_dict_returns_empty_queries() -> None:
    """An empty `arguments={}` payload (the streaming-mid first-event case) normalizes
    to `{'queries': []}` — that's "not yet populated", not "unrecognized"."""
    pytest.importorskip('openai')

    assert _normalize_tool_search_args({}) == {'queries': []}


def test_openai_normalize_tool_search_args_raises_on_unrecognized_shape() -> None:
    """Any non-empty payload that matches neither the `queries: list` nor `paths: list`
    shape raises `UnexpectedModelBehavior` so OpenAI SDK schema drift surfaces loudly
    at the parse boundary rather than silently degrading to an empty result."""
    pytest.importorskip('openai')

    # Non-dict input shouldn't happen given the SDK types arguments as a dict, but if it
    # ever does we want a loud failure rather than a silent fallback.
    with pytest.raises(UnexpectedModelBehavior, match='Unrecognized tool_search arguments shape'):
        _normalize_tool_search_args(None)
    with pytest.raises(UnexpectedModelBehavior, match='Unrecognized tool_search arguments shape'):
        _normalize_tool_search_args('')
    # Dict missing both recognized keys.
    with pytest.raises(UnexpectedModelBehavior, match='Unrecognized tool_search arguments shape'):
        _normalize_tool_search_args({'something_else': 'x'})
    # Dict with `paths` present but of a non-list type.
    with pytest.raises(UnexpectedModelBehavior, match='Unrecognized tool_search arguments shape'):
        _normalize_tool_search_args({'paths': 'not a list'})


# --- Cross-provider local→native promotion ---
#
# The local-fallback path emits typed `ToolSearchCallPart` / `ToolSearchReturnPart`
# (subclasses of the regular `ToolCallPart` / `ToolReturnPart`). When a follow-up
# turn runs on a provider that natively supports tool search, the adapter should
# render those local-shape parts back into the provider's native wire format so the
# previously discovered tools get unlocked from `defer_loading=True` without forcing
# the model to re-search. This must work regardless of the current turn's `strategy`
# (default native, named native, or custom callable) — the gate is "current request
# has any tool search active", not "strategy is custom".


async def test_anthropic_promotes_local_search_history_with_default_native_strategy() -> None:
    """Local-shape `ToolSearch*Part` from a prior cross-provider turn must render
    into Anthropic's native tool_search wire when the current turn is the default
    server-executed strategy (`ToolSearchTool()` / `strategy=None`).

    The wire shape uses Anthropic's "client-side flavor" of tool search per empirical
    research: a standard `tool_use` for the local `search_tools` function tool
    paired with a `tool_result` whose `content` is a `tool_reference` array
    (NOT a string of stringified discoveries). Anthropic's server unlocks the
    discovered tools' schemas from `defer_loading=true` once it sees the
    `tool_reference` block.

    Currently fails because `_build_custom_tool_search_replay_blocks` is gated on
    `strategy='custom'`, so the default-strategy case falls through and the return
    is rendered as a plain `tool_result` carrying stringified content — the
    discovered tools stay hidden and the model has to re-search.
    """
    pytest.importorskip('anthropic')

    model = AnthropicModel(
        'claude-sonnet-4-6',
        provider=AnthropicProvider(anthropic_client=MockAnthropic.create_mock(())),
    )

    history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='find a weather tool')]),
        ModelResponse(
            parts=[
                ToolSearchCallPart(args={'queries': ['weather']}, tool_call_id='c1'),
            ],
        ),
        ModelRequest(
            parts=[
                ToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'get_weather'}]},
                    tool_call_id='c1',
                ),
            ],
        ),
    ]
    # Default native strategy (NOT 'custom') — currently the gate that activates the
    # tool_reference replay re-formatting only fires for `strategy='custom'`. The
    # discovered tool ships on the wire with `defer_loading=True`; the replay
    # reference unlocks its schema server-side.
    params = ModelRequestParameters(
        function_tools=[ToolDefinition(name='get_weather', defer_loading=True)],
        native_tools=[ToolSearchTool()],
        allow_text_output=True,
    )

    _system, anthropic_messages = await model._map_message(history, params, AnthropicModelSettings())  # pyright: ignore[reportPrivateUsage]

    tool_results: list[dict[str, Any]] = [
        c
        for m in anthropic_messages
        if m['role'] == 'user' and isinstance(m['content'], list)
        for c in cast(list[Any], m['content'])
        if isinstance(c, dict) and cast(dict[str, Any], c).get('type') == 'tool_result'
    ]
    [tool_result] = tool_results
    # Promotion target: the result content must be a `tool_reference` array, not a
    # stringified discovery JSON. Anthropic uses this shape to unlock deferred tools.
    assert tool_result['content'] == [{'type': 'tool_reference', 'tool_name': 'get_weather'}]


async def test_anthropic_promotes_local_search_history_with_named_native_strategy() -> None:
    """Same promotion as above but with an explicit named native strategy
    (`strategy='bm25'`). Confirms the gate is "any tool search active", not "custom"
    or "default" — whenever the provider supports native tool search and the current
    request carries it, the historical local-shape parts get the native wire.
    """
    pytest.importorskip('anthropic')

    model = AnthropicModel(
        'claude-sonnet-4-6',
        provider=AnthropicProvider(anthropic_client=MockAnthropic.create_mock(())),
    )

    history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='find a calc tool')]),
        ModelResponse(
            parts=[
                ToolSearchCallPart(args={'queries': ['calc']}, tool_call_id='c2'),
            ],
        ),
        ModelRequest(
            parts=[
                ToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'calculate'}]},
                    tool_call_id='c2',
                ),
            ],
        ),
    ]
    params = ModelRequestParameters(
        function_tools=[ToolDefinition(name='calculate', defer_loading=True)],
        native_tools=[ToolSearchTool(strategy='bm25')],
        allow_text_output=True,
    )

    _system, anthropic_messages = await model._map_message(history, params, AnthropicModelSettings())  # pyright: ignore[reportPrivateUsage]

    tool_results: list[dict[str, Any]] = [
        c
        for m in anthropic_messages
        if m['role'] == 'user' and isinstance(m['content'], list)
        for c in cast(list[Any], m['content'])
        if isinstance(c, dict) and cast(dict[str, Any], c).get('type') == 'tool_result'
    ]
    [tool_result] = tool_results
    assert tool_result['content'] == [{'type': 'tool_reference', 'tool_name': 'calculate'}]


async def test_openai_promotes_local_search_history_with_default_native_strategy() -> None:
    """Local-shape `ToolSearch*Part` from a prior cross-provider turn must render
    into OpenAI's client-executed tool-search replay items when tool-search replay is active.

    The wire shape uses `tool_search_call` + `tool_search_output` items with
    `execution='client'`; request-building decides when that replay shape also requires
    the current top-level OpenAI `tool_search` registration to use client execution.
    """
    pytest.importorskip('openai')

    model = OpenAIResponsesModel(
        'gpt-5.4-mini',
        provider=OpenAIProvider(openai_client=MockOpenAIResponses.create_mock(())),
    )

    discovered_tool = ToolDefinition(
        name='get_weather',
        description='Get the weather for a city.',
        parameters_json_schema={
            'type': 'object',
            'properties': {'city': {'type': 'string'}},
            'required': ['city'],
        },
    )

    history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='find a weather tool')]),
        ModelResponse(
            parts=[
                ToolSearchCallPart(args={'queries': ['weather']}, tool_call_id='oc1'),
            ],
            provider_name='openai',
        ),
        ModelRequest(
            parts=[
                ToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'get_weather'}]},
                    tool_call_id='oc1',
                ),
            ],
        ),
    ]
    # Default native strategy — `ToolSearchTool()` with no `strategy='custom'`.
    # Discovered tool needs to be in `function_tools` so the replay can pair the
    # `tool_search_output.tools[]` schema by name.
    params = ModelRequestParameters(
        function_tools=[discovered_tool],
        native_tools=[ToolSearchTool()],
        allow_text_output=True,
    )

    _system, openai_messages = await model._map_messages(history, OpenAIResponsesModelSettings(), params)  # pyright: ignore[reportPrivateUsage]

    # The local search call should render as a `tool_search_call` item with
    # `execution='client'`, and the local return should render as a paired
    # `tool_search_output` carrying the `get_weather` schema.
    tool_search_calls = [
        item
        for item in openai_messages
        if isinstance(item, dict) and cast(dict[str, Any], item).get('type') == 'tool_search_call'
    ]
    tool_search_outputs = [
        item
        for item in openai_messages
        if isinstance(item, dict) and cast(dict[str, Any], item).get('type') == 'tool_search_output'
    ]
    function_calls = [
        item
        for item in openai_messages
        if isinstance(item, dict) and cast(dict[str, Any], item).get('type') == 'function_call'
    ]
    function_outputs = [
        item
        for item in openai_messages
        if isinstance(item, dict) and cast(dict[str, Any], item).get('type') == 'function_call_output'
    ]

    assert len(tool_search_calls) == 1, (
        f'expected 1 tool_search_call, got {len(tool_search_calls)}; full output: {openai_messages}'
    )
    assert len(tool_search_outputs) == 1
    assert tool_search_calls[0].get('execution') == 'client'
    assert tool_search_outputs[0].get('execution') == 'client'

    # Output carries the discovered tool's full schema for OpenAI to "rediscover".
    output_tools = cast(list[dict[str, Any]], tool_search_outputs[0].get('tools'))
    assert len(output_tools) == 1
    assert output_tools[0]['name'] == 'get_weather'
    assert output_tools[0]['type'] == 'function'

    # The `search_tools` exchange must NOT also surface as a regular function_call /
    # function_call_output — that would double-count the discovery.
    assert not any(
        cast(dict[str, Any], call).get('name') == _SEARCH_TOOLS_NAME
        for call in cast(list[ResponseFunctionToolCallParam], function_calls)
    )
    assert not function_outputs


# --- `strategy='keywords'` on natively-supporting providers ---
#
# `'keywords'` is a strategy CHOICE: "use the keyword-overlap algorithm". The execution
# mode (server-side / client-executed-native / local fallback) is auto-derived from
# the algorithm's needs and the provider's capabilities. On Anthropic and OpenAI,
# native tool search is available and the keyword algorithm runs LOCALLY but the
# wire ships in the provider's native tool-search shape so the prompt cache stays
# warm across discovery rounds (deferred tools don't get re-added to the request's
# tool definitions on each turn).


def test_tool_search_strategy_keywords_registers_builtin_for_client_execution() -> None:
    """`ToolSearch(strategy='keywords')` must register `ToolSearchTool(strategy='custom',
    optional=True)` so the client-executed native path engages on supporting providers.

    Currently fails because `get_native_tools` returns `[]` for `'keywords'`,
    forcing the local-fallback path on every provider — losing the cache benefit
    that the client-executed native path provides on Anthropic and OpenAI.
    """
    cap: ToolSearch = ToolSearch(strategy='keywords')
    builtins = cap.get_native_tools()
    assert len(builtins) == 1
    [builtin] = builtins
    assert isinstance(builtin, ToolSearchTool)
    # `strategy='custom'` marks the builtin as "the algorithm runs on our side"; the
    # adapter then wires it as Anthropic's tool_use+tool_reference flavor or OpenAI's
    # `execution='client'`. `optional=True` so it gets dropped on providers that
    # don't support it (toolset's local `search_tools` function tool is the fallback).
    assert builtin.strategy == 'custom'
    assert builtin.optional is True


async def test_openai_promotes_mixed_native_and_local_history_a_b_c_chain() -> None:
    """Multi-hop chain: Anthropic-native turn 1 → local turn 2 (Google etc.) → OpenAI turn 3.

    The persisted history at turn 3 carries BOTH a `NativeToolSearch*Part` from the
    Anthropic turn AND a `ToolSearch*Part` from the local turn. OpenAI's adapter must
    promote both into native `tool_search_call`+`tool_search_output` items so the
    discovered tools' schemas stay unlocked across the entire chain — the model
    shouldn't have to re-search anything it discovered earlier.
    """
    pytest.importorskip('openai')

    model = OpenAIResponsesModel(
        'gpt-5.4-mini',
        provider=OpenAIProvider(openai_client=MockOpenAIResponses.create_mock(())),
    )

    weather = ToolDefinition(
        name='get_weather',
        description='Get the weather for a city.',
        parameters_json_schema={'type': 'object', 'properties': {'city': {'type': 'string'}}, 'required': ['city']},
    )
    calc = ToolDefinition(
        name='calculate_mortgage',
        description='Calculate monthly mortgage payment.',
        parameters_json_schema={'type': 'object', 'properties': {'p': {'type': 'number'}}, 'required': ['p']},
    )

    # Turn 1 on Anthropic: native bm25, discovers `get_weather`.
    # Turn 2 on Google: local function tool, discovers `calculate_mortgage`.
    # Turn 3 on OpenAI: should promote BOTH discoveries to native wire.
    history: list[ModelMessage] = [
        ModelRequest(parts=[UserPromptPart(content='find a weather tool')]),
        # Anthropic-native (turn 1) — `NativeToolSearch*Part`.
        ModelResponse(
            parts=[
                NativeToolSearchCallPart(
                    args={'queries': ['weather']},
                    tool_call_id='ant_1',
                    provider_name='anthropic',
                    provider_details={'strategy': 'bm25'},
                ),
                NativeToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'get_weather'}]},
                    tool_call_id='ant_1',
                    provider_name='anthropic',
                ),
            ],
            provider_name='anthropic',
        ),
        ModelRequest(parts=[UserPromptPart(content='now find a mortgage one')]),
        # Local fallback (turn 2 on Google or similar) — `ToolSearch*Part`.
        ModelResponse(parts=[ToolSearchCallPart(args={'queries': ['mortgage']}, tool_call_id='loc_1')]),
        ModelRequest(
            parts=[
                ToolSearchReturnPart(
                    content={'discovered_tools': [{'name': 'calculate_mortgage'}]},
                    tool_call_id='loc_1',
                ),
            ],
        ),
        ModelRequest(parts=[UserPromptPart(content='now compute both')]),
    ]

    params = ModelRequestParameters(
        function_tools=[weather, calc],
        native_tools=[ToolSearchTool()],
        allow_text_output=True,
    )

    _system, openai_messages = await model._map_messages(history, OpenAIResponsesModelSettings(), params)  # pyright: ignore[reportPrivateUsage]

    tool_search_calls = [
        item
        for item in openai_messages
        if isinstance(item, dict) and cast(dict[str, Any], item).get('type') == 'tool_search_call'
    ]
    tool_search_outputs = [
        item
        for item in openai_messages
        if isinstance(item, dict) and cast(dict[str, Any], item).get('type') == 'tool_search_output'
    ]

    # Both prior discoveries should surface as native tool_search exchanges with execution=client.
    # The local-fallback one promotes via the new gating; the Anthropic-native one is left as-is
    # because its provider_name doesn't match self.system (foreign-provider builtin parts are
    # filtered out from the OpenAI wire, but get_weather still needs to be discoverable — that's
    # handled by the toolset re-emitting it as a regular function tool in this turn's `tools[]`).
    assert len(tool_search_calls) >= 1, (
        f'expected at least one promoted tool_search_call (local→native), got {len(tool_search_calls)}; '
        f'output: {openai_messages}'
    )
    assert len(tool_search_outputs) >= 1
    # The local discovery (`calculate_mortgage`) made it into the promoted output.
    output_tools_names = {
        cast(dict[str, Any], t).get('name')
        for output in tool_search_outputs
        for t in cast(list[Any], cast(dict[str, Any], output).get('tools', []))
    }
    assert 'calculate_mortgage' in output_tools_names, (
        f'local-fallback discovery should be promoted; got tools: {output_tools_names}'
    )


def test_keywords_search_fn_returns_empty_for_no_tokens() -> None:
    """The shared keyword algorithm returns `[]` when the queries tokenize to nothing
    (whitespace / punctuation only), instead of raising. Callers (`_run_search_fn`
    in the toolset) translate that into the empty-discoveries `_empty_return` shape.
    """

    ctx = _build_run_context(None)
    assert keywords_search_fn(ctx, ['   '], []) == []
    # Punctuation-only queries also produce no tokens — `_SEARCH_TOKEN_RE` matches
    # `[a-z0-9]+` only.
    assert keywords_search_fn(ctx, ['!!!'], []) == []


async def test_tool_search_strategy_keywords_runs_keyword_algorithm_via_search_fn() -> None:
    """When `strategy='keywords'` activates the client-executed native path, the local
    `search_tools` function (still in `function_tools` for client-execution) must run
    the built-in keyword-overlap algorithm — not error out with no `search_fn` set.

    Verifies end-to-end: the toolset's `search_fn` is wired to a callable that
    matches keywords against the corpus, returning matching tool names.
    """
    cap: ToolSearch = ToolSearch(strategy='keywords')
    base = _create_function_toolset()
    # `get_wrapper_toolset` is what the framework calls when injecting the capability.
    ts = cap.get_wrapper_toolset(base)
    assert isinstance(ts, ToolSearchToolset)
    # Internal `search_fn` is set so `_run_search_fn` (not `_run_keywords_search`) handles
    # the dispatch — but the algorithm is still keyword overlap.
    assert ts.search_fn is not None

    ctx = _build_run_context(None)
    tools = await ts.get_tools(ctx)
    search_tool = tools[_SEARCH_TOOLS_NAME]
    result = await ts.call_tool(_SEARCH_TOOLS_NAME, {'queries': ['mortgage']}, ctx, search_tool)
    return_value = cast(dict[str, Any], result)
    discovered_names = {match['name'] for match in return_value['discovered_tools']}
    assert 'calculate_mortgage' in discovered_names


# --- Capability-gated tool-search forces client-executed promotion ---
#
# Provider-side tool search can't honor capability gating — it would reveal corpus tools
# whose owning capability hasn't been loaded yet. When any function tool has both
# `with_native='tool_search'` and a `capability_id`, `_resolve_native_tool_swap` either
# raises (named-native strategies have no local equivalent) or promotes `strategy=None` to
# `'custom'` (client-executed), keeping `search_tools` on the wire as the callback.


def _capability_owned_corpus_tool() -> ToolDefinition:
    return ToolDefinition(
        name='lookup_refund_policy',
        parameters_json_schema={'type': 'object', 'properties': {}},
        with_native=ToolSearchTool.kind,
        capability_id='refunds',
        defer_loading=True,
        metadata={DEFERRED_CAPABILITY_TOOL_METADATA_KEY: True},
    )


def _local_search_tools_def() -> ToolDefinition:
    return ToolDefinition(name=_SEARCH_TOOLS_NAME, parameters_json_schema={}, unless_native=ToolSearchTool.kind)


@pytest.mark.parametrize('strategy', ['bm25', 'regex'])
def test_capability_gated_tool_search_raises_on_named_native_strategy(strategy: str) -> None:
    """Named-native strategies have no local equivalent — silently substituting `keywords`
    would change the user's chosen algorithm, so we raise."""

    class M(TestModel):
        @classmethod
        def supported_native_tools(cls) -> frozenset[type[AbstractNativeTool]]:
            return frozenset({ToolSearchTool})

    params = ModelRequestParameters(
        function_tools=[_capability_owned_corpus_tool()],
        native_tools=[ToolSearchTool(strategy=cast(Any, strategy), optional=True)],
    )
    with pytest.raises(UserError, match=rf'strategy={strategy!r}.*incompatible with deferred-loading'):
        M().prepare_request(None, params)


def test_capability_gated_tool_search_promotes_default_strategy_to_custom() -> None:
    """Promotion: `strategy=None` → `'custom'` and `search_tools` stays on the wire as the
    client-executed callback."""

    class M(TestModel):
        @classmethod
        def supported_native_tools(cls) -> frozenset[type[AbstractNativeTool]]:
            return frozenset({ToolSearchTool})

    params = ModelRequestParameters(
        function_tools=[_local_search_tools_def(), _capability_owned_corpus_tool()],
        native_tools=[ToolSearchTool(strategy=None, optional=True)],
    )
    _, prepared = M().prepare_request(None, params)

    [native] = prepared.native_tools
    assert isinstance(native, ToolSearchTool) and native.strategy == 'custom'
    assert _SEARCH_TOOLS_NAME in [t.name for t in prepared.function_tools]


def test_capability_gated_tool_search_skips_other_natives_and_leaves_custom_strategy_unchanged() -> None:
    """Promotion must skip past non-`ToolSearchTool` entries in `supported_natives` and
    leave an already-`'custom'` strategy (set by `ToolSearch(strategy='keywords'|callable)`)
    untouched — but still report `True` so `search_tools` stays on the wire."""

    class M(TestModel):
        @classmethod
        def supported_native_tools(cls) -> frozenset[type[AbstractNativeTool]]:
            return frozenset({ToolSearchTool, WebSearchTool})

    params = ModelRequestParameters(
        function_tools=[_local_search_tools_def(), _capability_owned_corpus_tool()],
        # WebSearchTool listed first so the promotion loop must `continue` past it.
        native_tools=[WebSearchTool(), ToolSearchTool(strategy='custom', optional=True)],
    )
    _, prepared = M().prepare_request(None, params)

    [tool_search] = [t for t in prepared.native_tools if isinstance(t, ToolSearchTool)]
    assert tool_search.strategy == 'custom'
    assert _SEARCH_TOOLS_NAME in [t.name for t in prepared.function_tools]


def test_capability_gated_tool_search_leaves_non_capability_corpus_alone() -> None:
    """No `capability_id` → no promotion: `strategy=None` reaches the adapter unchanged and
    rule 1 still strips the local `search_tools` since the native is supported."""

    class M(TestModel):
        @classmethod
        def supported_native_tools(cls) -> frozenset[type[AbstractNativeTool]]:
            return frozenset({ToolSearchTool})

    plain_deferred = ToolDefinition(
        name='deferred_tool',
        parameters_json_schema={'type': 'object', 'properties': {}},
        with_native=ToolSearchTool.kind,
        defer_loading=True,
    )
    params = ModelRequestParameters(
        function_tools=[_local_search_tools_def(), plain_deferred],
        native_tools=[ToolSearchTool(strategy=None, optional=True)],
    )
    _, prepared = M().prepare_request(None, params)

    [native] = prepared.native_tools
    assert isinstance(native, ToolSearchTool) and native.strategy is None
    assert _SEARCH_TOOLS_NAME not in [t.name for t in prepared.function_tools]


# --- Namespace synthesis for any tool-search corpus member ---
#
# OpenAI rejects replayed tool-search-discovered function calls without a `namespace`. For
# cross-provider replay there's no captured namespace, so the adapter synthesizes one from
# the tool name. The gate is `with_native='tool_search'` (any corpus member), not just
# capability-owned tools — plain `defer_loading=True` tools also need this on replay.


def test_tool_search_namespace_synthesis_returns_tool_name_for_corpus_member() -> None:
    """A function tool with `with_native='tool_search'` and no `capability_id` still gets a
    synthesized namespace — the gate is corpus membership, not capability ownership."""
    pytest.importorskip('openai')

    plain_corpus_tool = ToolDefinition(
        name='lookup_refund_policy',
        parameters_json_schema={'type': 'object', 'properties': {}},
        with_native=ToolSearchTool.kind,
    )
    params = ModelRequestParameters(function_tools=[plain_corpus_tool])
    assert _tool_search_namespace_for_synthesis('lookup_refund_policy', params) == 'lookup_refund_policy'


def test_tool_search_namespace_synthesis_returns_none_for_unrelated_function_tool() -> None:
    """A regular function tool (no `with_native`) must not be tagged — synthesizing a
    namespace there would inject a field the API didn't request."""
    pytest.importorskip('openai')

    regular_tool = ToolDefinition(name='get_weather', parameters_json_schema={'type': 'object', 'properties': {}})
    params = ModelRequestParameters(function_tools=[regular_tool])
    assert _tool_search_namespace_for_synthesis('get_weather', params) is None