pydantic--pydantic-ai/tests/test_toolsets.py

from __future__ import annotations

import re
import sys
from collections import defaultdict
from dataclasses import dataclass, replace
from datetime import timezone
from typing import Any, TypeVar
from unittest.mock import AsyncMock

import anyio
import pytest
from pydantic import ValidationError
from typing_extensions import Self

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

from pydantic_ai import (
    AbstractToolset,
    Agent,
    CombinedToolset,
    FilteredToolset,
    FunctionToolset,
    PrefixedToolset,
    PreparedToolset,
    ToolCallPart,
    ToolsetTool,
    WrapperToolset,
    capture_run_messages,
)
from pydantic_ai._run_context import RunContext
from pydantic_ai.exceptions import ModelRetry, ToolRetryError, UnexpectedModelBehavior, UserError
from pydantic_ai.messages import (
    InstructionPart,
    ModelRequest,
    ModelResponse,
    RetryPromptPart,
    TextPart,
    ToolReturnPart,
    UserPromptPart,
)
from pydantic_ai.models.test import TestModel
from pydantic_ai.tool_manager import ToolManager
from pydantic_ai.tools import Tool, ToolDefinition
from pydantic_ai.toolsets._dynamic import DynamicToolset
from pydantic_ai.usage import RequestUsage, RunUsage

from ._inline_snapshot import snapshot
from .conftest import IsDatetime, IsNow, IsStr

pytestmark = pytest.mark.anyio

T = TypeVar('T')


def build_run_context(deps: T, run_step: int = 0, max_retries: int = 0) -> RunContext[T]:
    return RunContext(
        deps=deps,
        model=TestModel(),
        usage=RunUsage(),
        prompt=None,
        messages=[],
        run_step=run_step,
        max_retries=max_retries,
    )


class MockToolsetWithInstructions(AbstractToolset[Any]):
    """A test toolset that returns custom instructions."""

    def __init__(self, instructions: str | None = None, id: str | None = None):
        self.custom_instructions = instructions
        self._id = id

    @property
    def id(self) -> str | None:
        return self._id

    async def get_instructions(self, ctx: RunContext[Any]) -> str | None:
        return self.custom_instructions

    async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
        return {}

    async def call_tool(
        self, name: str, tool_args: dict[str, Any], ctx: RunContext[Any], tool: ToolsetTool[Any]
    ) -> Any:
        return None


async def test_mock_toolset_with_instructions_interface():
    """Test that MockToolsetWithInstructions correctly implements AbstractToolset interface."""
    toolset = MockToolsetWithInstructions(instructions='test instructions', id='my-id')
    ctx = build_run_context(None)

    assert toolset.id == 'my-id'
    assert await toolset.get_tools(ctx) == {}
    assert await toolset.call_tool('any_tool', {}, ctx, None) is None  # type: ignore[arg-type]


async def test_function_toolset_instructions():
    """Test that FunctionToolset returns None for instructions by default."""
    toolset = FunctionToolset()
    ctx = build_run_context(None)
    instructions = await toolset.get_instructions(ctx)
    assert instructions is None


async def test_function_toolset():
    @dataclass
    class PrefixDeps:
        prefix: str | None = None

    toolset = FunctionToolset[PrefixDeps]()

    async def prepare_add_prefix(ctx: RunContext[PrefixDeps], tool_def: ToolDefinition) -> ToolDefinition | None:
        if ctx.deps.prefix is None:
            return tool_def

        return replace(tool_def, name=f'{ctx.deps.prefix}_{tool_def.name}')

    @toolset.tool_plain(prepare=prepare_add_prefix)
    def add(a: int, b: int) -> int:
        """Add two numbers"""
        return a + b

    no_prefix_context = build_run_context(PrefixDeps())
    no_prefix_toolset = await ToolManager[PrefixDeps](toolset).for_run_step(no_prefix_context)
    assert no_prefix_toolset.tool_defs == snapshot(
        [
            ToolDefinition(
                name='add',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                    'required': ['a', 'b'],
                    'type': 'object',
                },
                description='Add two numbers',
                return_schema={'type': 'integer'},
            )
        ]
    )
    assert await no_prefix_toolset.handle_call(ToolCallPart(tool_name='add', args={'a': 1, 'b': 2})) == 3

    foo_context = build_run_context(PrefixDeps(prefix='foo'))
    foo_toolset = await ToolManager[PrefixDeps](toolset).for_run_step(foo_context)
    assert foo_toolset.tool_defs == snapshot(
        [
            ToolDefinition(
                name='foo_add',
                description='Add two numbers',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                    'required': ['a', 'b'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            )
        ]
    )
    assert await foo_toolset.handle_call(ToolCallPart(tool_name='foo_add', args={'a': 1, 'b': 2})) == 3

    @toolset.tool_plain
    def subtract(a: int, b: int) -> int:
        """Subtract two numbers"""
        return a - b  # pragma: lax no cover

    bar_context = build_run_context(PrefixDeps(prefix='bar'))
    bar_toolset = await ToolManager[PrefixDeps](toolset).for_run_step(bar_context)
    assert bar_toolset.tool_defs == snapshot(
        [
            ToolDefinition(
                name='bar_add',
                description='Add two numbers',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                    'required': ['a', 'b'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            ),
            ToolDefinition(
                name='subtract',
                description='Subtract two numbers',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                    'required': ['a', 'b'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            ),
        ]
    )
    assert await bar_toolset.handle_call(ToolCallPart(tool_name='bar_add', args={'a': 1, 'b': 2})) == 3


async def test_toolset_tool_function_signature_property():
    toolset = FunctionToolset()

    @toolset.tool_plain
    def add(a: int, b: int) -> int:
        return a + b

    managed_toolset = await ToolManager(toolset).for_run_step(build_run_context(None))
    assert managed_toolset.tools is not None

    td = managed_toolset.tools['add'].tool_def
    sig = td.function_signature
    assert sig is not None
    assert list(sig.params) == ['a', 'b']
    assert td.render_signature('...') == snapshot("""\
def add(*, a: int, b: int) -> int:
    ...\
""")

    assert await managed_toolset.handle_call(ToolCallPart(tool_name='add', args={'a': 1, 'b': 2})) == 3


async def test_function_toolset_with_defaults():
    defaults_toolset = FunctionToolset(require_parameter_descriptions=True)

    with pytest.raises(
        UserError,
        match=re.escape('Missing parameter descriptions for'),
    ):

        @defaults_toolset.tool_plain
        def add(a: int, b: int) -> int:
            """Add two numbers"""
            return a + b  # pragma: no cover


async def test_abstract_toolset_instructions_default():
    """Test that the default instructions method returns None."""
    toolset = MockToolsetWithInstructions(instructions=None)
    ctx = build_run_context(None)
    instructions = await toolset.get_instructions(ctx)
    assert instructions is None


async def test_abstract_toolset_instructions_custom():
    """Test that instructions can return custom instructions."""
    custom_instructions = 'Use these tools carefully and always validate inputs.'
    toolset = MockToolsetWithInstructions(instructions=custom_instructions)
    ctx = build_run_context(None)
    instructions = await toolset.get_instructions(ctx)
    assert instructions == custom_instructions


async def test_abstract_toolset_instructions_empty_string():
    """Test that instructions can return an empty string."""
    toolset = MockToolsetWithInstructions(instructions='')
    ctx = build_run_context(None)
    instructions = await toolset.get_instructions(ctx)
    assert instructions == ''


async def test_function_toolset_with_defaults_overridden():
    defaults_toolset = FunctionToolset(require_parameter_descriptions=True)

    @defaults_toolset.tool_plain(require_parameter_descriptions=False)
    def subtract(a: int, b: int) -> int:
        """Subtract two numbers"""
        return a - b  # pragma: no cover


async def test_prepared_toolset_sync_prepare_func():
    """`PreparedToolset` accepts a synchronous prepare function (no await needed)."""
    base_toolset = FunctionToolset()

    @base_toolset.tool_plain
    def add(a: int, b: int) -> int:
        """Add two numbers"""
        return a + b  # pragma: no cover

    def prepare_keep_first(ctx: RunContext, tool_defs: list[ToolDefinition]) -> list[ToolDefinition]:
        return tool_defs[:1]

    prepared_toolset = PreparedToolset(base_toolset, prepare_keep_first)

    tools = await prepared_toolset.get_tools(build_run_context(None))
    assert list(tools.keys()) == ['add']


async def test_prepared_toolset_user_error_none_result():
    """`PreparedToolset` requires [] when a prepare function intentionally exposes no tools."""
    base_toolset = FunctionToolset()

    @base_toolset.tool_plain
    def add(a: int, b: int) -> int:
        """Add two numbers"""
        return a + b  # pragma: no cover

    async def prepare_returns_none(ctx: RunContext, tool_defs: list[ToolDefinition]) -> list[ToolDefinition] | None:
        return None

    prepared_toolset = PreparedToolset(base_toolset, prepare_returns_none)  # pyright: ignore[reportArgumentType]

    with pytest.raises(UserError, match="Prepare function 'prepare_returns_none' returned `None`"):
        await prepared_toolset.get_tools(build_run_context(None))


async def test_prepared_toolset_user_error_add_new_tools():
    """Test that PreparedToolset raises UserError when prepare function tries to add new tools."""
    context = build_run_context(None)
    base_toolset = FunctionToolset()

    @base_toolset.tool_plain
    def add(a: int, b: int) -> int:
        """Add two numbers"""
        return a + b  # pragma: no cover

    async def prepare_add_new_tool(ctx: RunContext, tool_defs: list[ToolDefinition]) -> list[ToolDefinition]:
        # Try to add a new tool that wasn't in the original set
        new_tool = ToolDefinition(
            name='new_tool',
            description='A new tool',
            parameters_json_schema={
                'additionalProperties': False,
                'properties': {'x': {'type': 'integer'}},
                'required': ['x'],
                'type': 'object',
            },
        )
        return tool_defs + [new_tool]

    prepared_toolset = PreparedToolset(base_toolset, prepare_add_new_tool)

    with pytest.raises(
        UserError,
        match=re.escape(
            'Prepare function cannot add or rename tools. Use `FunctionToolset.add_function()` or `RenamedToolset` instead.'
        ),
    ):
        await ToolManager(prepared_toolset).for_run_step(context)


async def test_prepared_toolset_user_error_change_tool_names():
    """Test that PreparedToolset raises UserError when prepare function tries to change tool names."""
    context = build_run_context(None)
    base_toolset = FunctionToolset()

    @base_toolset.tool_plain
    def add(a: int, b: int) -> int:
        """Add two numbers"""
        return a + b  # pragma: no cover

    @base_toolset.tool_plain
    def subtract(a: int, b: int) -> int:
        """Subtract two numbers"""
        return a - b  # pragma: no cover

    async def prepare_change_names(ctx: RunContext, tool_defs: list[ToolDefinition]) -> list[ToolDefinition]:
        # Try to change the name of an existing tool
        modified_tool_defs: list[ToolDefinition] = []
        for tool_def in tool_defs:
            if tool_def.name == 'add':
                modified_tool_defs.append(replace(tool_def, name='modified_add'))
            else:
                modified_tool_defs.append(tool_def)
        return modified_tool_defs

    prepared_toolset = PreparedToolset(base_toolset, prepare_change_names)

    with pytest.raises(
        UserError,
        match=re.escape(
            'Prepare function cannot add or rename tools. Use `FunctionToolset.add_function()` or `RenamedToolset` instead.'
        ),
    ):
        await ToolManager(prepared_toolset).for_run_step(context)


async def test_comprehensive_toolset_composition():
    """Test that all toolsets can be composed together and work correctly."""

    @dataclass
    class TestDeps:
        user_role: str = 'user'
        enable_advanced: bool = True

    # Create first FunctionToolset with basic math operations
    math_toolset = FunctionToolset[TestDeps]()

    @math_toolset.tool_plain
    def add(a: int, b: int) -> int:
        """Add two numbers"""
        return a + b

    @math_toolset.tool_plain
    def subtract(a: int, b: int) -> int:
        """Subtract two numbers"""
        return a - b  # pragma: no cover

    @math_toolset.tool_plain
    def multiply(a: int, b: int) -> int:
        """Multiply two numbers"""
        return a * b  # pragma: no cover

    # Create second FunctionToolset with string operations
    string_toolset = FunctionToolset[TestDeps]()

    @string_toolset.tool_plain
    def concat(s1: str, s2: str) -> str:
        """Concatenate two strings"""
        return s1 + s2

    @string_toolset.tool_plain
    def uppercase(text: str) -> str:
        """Convert text to uppercase"""
        return text.upper()  # pragma: no cover

    @string_toolset.tool_plain
    def reverse(text: str) -> str:
        """Reverse a string"""
        return text[::-1]  # pragma: no cover

    # Create third FunctionToolset with advanced operations
    advanced_toolset = FunctionToolset[TestDeps]()

    @advanced_toolset.tool_plain
    def power(base: int, exponent: int) -> int:
        """Calculate base raised to the power of exponent"""
        return base**exponent  # pragma: no cover

    # Step 1: Prefix each FunctionToolset individually
    prefixed_math = PrefixedToolset(math_toolset, 'math')
    prefixed_string = PrefixedToolset(string_toolset, 'str')
    prefixed_advanced = PrefixedToolset(advanced_toolset, 'adv')

    # Step 2: Combine the prefixed toolsets
    combined_prefixed_toolset = CombinedToolset([prefixed_math, prefixed_string, prefixed_advanced])

    # Step 3: Filter tools based on user role and advanced flag, now using prefixed names
    def filter_tools(ctx: RunContext[TestDeps], tool_def: ToolDefinition) -> bool:
        # Only allow advanced tools if enable_advanced is True
        if tool_def.name.startswith('adv_') and not ctx.deps.enable_advanced:
            return False
        # Only allow string operations for admin users (simulating role-based access)
        if tool_def.name.startswith('str_') and ctx.deps.user_role != 'admin':
            return False
        return True

    filtered_toolset = FilteredToolset[TestDeps](combined_prefixed_toolset, filter_tools)

    # Step 4: Apply prepared toolset to modify descriptions (add user role annotation)
    async def prepare_add_context(ctx: RunContext[TestDeps], tool_defs: list[ToolDefinition]) -> list[ToolDefinition]:
        # Annotate each tool description with the user role
        role = ctx.deps.user_role
        return [replace(td, description=f'{td.description} (role: {role})') for td in tool_defs]

    prepared_toolset = PreparedToolset(filtered_toolset, prepare_add_context)

    # Step 5: Test the fully composed toolset
    # Test with regular user context
    regular_deps = TestDeps(user_role='user', enable_advanced=True)
    regular_context = build_run_context(regular_deps)
    final_toolset = await ToolManager[TestDeps](prepared_toolset).for_run_step(regular_context)
    # Tool definitions should have role annotation
    assert final_toolset.tool_defs == snapshot(
        [
            ToolDefinition(
                name='math_add',
                description='Add two numbers (role: user)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                    'required': ['a', 'b'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            ),
            ToolDefinition(
                name='math_subtract',
                description='Subtract two numbers (role: user)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                    'required': ['a', 'b'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            ),
            ToolDefinition(
                name='math_multiply',
                description='Multiply two numbers (role: user)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                    'required': ['a', 'b'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            ),
            ToolDefinition(
                name='adv_power',
                description='Calculate base raised to the power of exponent (role: user)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'base': {'type': 'integer'}, 'exponent': {'type': 'integer'}},
                    'required': ['base', 'exponent'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            ),
        ]
    )
    # Call a tool and check result
    result = await final_toolset.handle_call(ToolCallPart(tool_name='math_add', args={'a': 5, 'b': 3}))
    assert result == 8

    # Test with admin user context (should have string tools)
    admin_deps = TestDeps(user_role='admin', enable_advanced=True)
    admin_context = build_run_context(admin_deps)
    admin_final_toolset = await ToolManager[TestDeps](prepared_toolset).for_run_step(admin_context)
    assert admin_final_toolset.tool_defs == snapshot(
        [
            ToolDefinition(
                name='math_add',
                description='Add two numbers (role: admin)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                    'required': ['a', 'b'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            ),
            ToolDefinition(
                name='math_subtract',
                description='Subtract two numbers (role: admin)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                    'required': ['a', 'b'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            ),
            ToolDefinition(
                name='math_multiply',
                description='Multiply two numbers (role: admin)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                    'required': ['a', 'b'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            ),
            ToolDefinition(
                name='str_concat',
                description='Concatenate two strings (role: admin)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'s1': {'type': 'string'}, 's2': {'type': 'string'}},
                    'required': ['s1', 's2'],
                    'type': 'object',
                },
                return_schema={'type': 'string'},
            ),
            ToolDefinition(
                name='str_uppercase',
                description='Convert text to uppercase (role: admin)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'text': {'type': 'string'}},
                    'required': ['text'],
                    'type': 'object',
                },
                return_schema={'type': 'string'},
            ),
            ToolDefinition(
                name='str_reverse',
                description='Reverse a string (role: admin)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'text': {'type': 'string'}},
                    'required': ['text'],
                    'type': 'object',
                },
                return_schema={'type': 'string'},
            ),
            ToolDefinition(
                name='adv_power',
                description='Calculate base raised to the power of exponent (role: admin)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'base': {'type': 'integer'}, 'exponent': {'type': 'integer'}},
                    'required': ['base', 'exponent'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            ),
        ]
    )
    result = await admin_final_toolset.handle_call(
        ToolCallPart(tool_name='str_concat', args={'s1': 'Hello', 's2': 'World'})
    )
    assert result == 'HelloWorld'

    # Test with advanced features disabled
    basic_deps = TestDeps(user_role='user', enable_advanced=False)
    basic_context = build_run_context(basic_deps)
    basic_final_toolset = await ToolManager[TestDeps](prepared_toolset).for_run_step(basic_context)
    assert basic_final_toolset.tool_defs == snapshot(
        [
            ToolDefinition(
                name='math_add',
                description='Add two numbers (role: user)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                    'required': ['a', 'b'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            ),
            ToolDefinition(
                name='math_subtract',
                description='Subtract two numbers (role: user)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                    'required': ['a', 'b'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            ),
            ToolDefinition(
                name='math_multiply',
                description='Multiply two numbers (role: user)',
                parameters_json_schema={
                    'additionalProperties': False,
                    'properties': {'a': {'type': 'integer'}, 'b': {'type': 'integer'}},
                    'required': ['a', 'b'],
                    'type': 'object',
                },
                return_schema={'type': 'integer'},
            ),
        ]
    )


async def test_context_manager():
    try:
        from fastmcp.client.transports import StdioTransport

        from pydantic_ai.mcp import MCPToolset
    except ImportError:  # pragma: lax no cover
        pytest.skip('mcp is not installed')

    transport = StdioTransport(command='python', args=['-m', 'tests.mcp_server'])
    server1 = MCPToolset(transport)
    server2 = MCPToolset(transport)
    toolset = CombinedToolset([server1, PrefixedToolset(server2, 'prefix')])

    async with toolset:
        assert server1.is_running
        assert server2.is_running

        async with toolset:
            assert server1.is_running
            assert server2.is_running


class InitializationError(Exception):
    pass


async def test_context_manager_failed_initialization():
    """Test if MCP servers stop if any MCP server fails to initialize."""
    try:
        from fastmcp.client.transports import StdioTransport

        from pydantic_ai.mcp import MCPToolset
    except ImportError:  # pragma: lax no cover
        pytest.skip('mcp is not installed')

    server1 = MCPToolset(StdioTransport(command='python', args=['-m', 'tests.mcp_server']))
    server2 = AsyncMock()
    server2.__aenter__.side_effect = InitializationError

    toolset = CombinedToolset([server1, server2])

    with pytest.raises(InitializationError):
        async with toolset:
            pass

    assert server1.is_running is False


async def test_tool_manager_reuse_self():
    """Test the retry logic with failed_tools and for_run_step method."""

    run_context = build_run_context(None, run_step=1)

    tool_manager = await ToolManager(FunctionToolset()).for_run_step(run_context)

    same_tool_manager = await tool_manager.for_run_step(ctx=run_context)

    assert tool_manager is same_tool_manager

    step_2_context = build_run_context(None, run_step=2)

    updated_tool_manager = await tool_manager.for_run_step(ctx=step_2_context)

    assert tool_manager != updated_tool_manager


async def test_tool_manager_retry_logic():
    """Test the retry logic with failed_tools and for_run_step method."""

    @dataclass
    class TestDeps:
        pass

    # Create a toolset with tools that can fail
    toolset = FunctionToolset[TestDeps](max_retries=2)
    call_count: defaultdict[str, int] = defaultdict(int)

    @toolset.tool_plain
    def failing_tool(x: int) -> int:
        """A tool that always fails"""
        call_count['failing_tool'] += 1
        raise ModelRetry('This tool always fails')

    @toolset.tool_plain
    def other_tool(x: int) -> int:
        """A tool that works"""
        call_count['other_tool'] += 1
        return x * 2

    # Create initial context and tool manager
    initial_context = build_run_context(TestDeps())
    tool_manager = await ToolManager[TestDeps](toolset).for_run_step(initial_context)

    # Initially no failed tools
    assert tool_manager.failed_tools == set()
    assert initial_context.retries == {}

    # Call the failing tool - should add to failed_tools
    with pytest.raises(ToolRetryError):
        await tool_manager.handle_call(ToolCallPart(tool_name='failing_tool', args={'x': 1}))

    assert tool_manager.failed_tools == {'failing_tool'}
    assert call_count['failing_tool'] == 1

    # Call the working tool - should not add to failed_tools
    result = await tool_manager.handle_call(ToolCallPart(tool_name='other_tool', args={'x': 3}))
    assert result == 6
    assert tool_manager.failed_tools == {'failing_tool'}  # unchanged
    assert call_count['other_tool'] == 1

    # Test for_run_step - should create new tool manager with updated retry counts
    new_context = build_run_context(TestDeps(), run_step=1)
    new_tool_manager = await tool_manager.for_run_step(new_context)

    # The new tool manager should have retry count for the failed tool
    assert new_tool_manager.ctx is not None
    assert new_tool_manager.ctx.retries == {'failing_tool': 1}
    assert new_tool_manager.failed_tools == set()  # reset for new run step

    # Call the failing tool again in the new manager - should have retry=1
    with pytest.raises(ToolRetryError):
        await new_tool_manager.handle_call(ToolCallPart(tool_name='failing_tool', args={'x': 1}))

    # Call the failing tool another time in the new manager
    with pytest.raises(ToolRetryError):
        await new_tool_manager.handle_call(ToolCallPart(tool_name='failing_tool', args={'x': 1}))

    # Call the failing tool a third time in the new manager
    with pytest.raises(ToolRetryError):
        await new_tool_manager.handle_call(ToolCallPart(tool_name='failing_tool', args={'x': 1}))

    assert new_tool_manager.failed_tools == {'failing_tool'}
    assert call_count['failing_tool'] == 4

    # Create another run step
    another_context = build_run_context(TestDeps(), run_step=2)
    another_tool_manager = await new_tool_manager.for_run_step(another_context)

    # Should now have retry count of 2 for failing_tool
    assert another_tool_manager.ctx is not None
    assert another_tool_manager.ctx.retries == {'failing_tool': 2}
    assert another_tool_manager.failed_tools == set()

    # Call the failing tool _again_, now we should finally hit the limit
    with pytest.raises(UnexpectedModelBehavior, match="Tool 'failing_tool' exceeded max retries count of 2"):
        await another_tool_manager.handle_call(ToolCallPart(tool_name='failing_tool', args={'x': 1}))


async def test_handle_call_wrap_validation_errors_false():
    """`handle_call(wrap_validation_errors=False)` propagates raw errors and leaves retry-budget state untouched.

    Used by sandboxed callers (e.g. code-mode dispatch) that want validation and
    `ModelRetry` failures to surface at the sandbox `await` site as the original
    exception type, without consuming the agent's retry budget for the wrapping call.
    Mirrors the `wrap_validation_errors` flag on the output-tool methods.
    """

    toolset = FunctionToolset(max_retries=2)

    @toolset.tool_plain
    def needs_int(x: int) -> int:
        return x * 2

    @toolset.tool_plain
    def retrying() -> int:
        raise ModelRetry('please retry')

    tool_manager = await ToolManager(toolset).for_run_step(build_run_context(None))

    # Sanity: a valid call still works in raw mode (no path differences for happy paths).
    assert (
        await tool_manager.handle_call(
            ToolCallPart(tool_name='needs_int', args={'x': 5}),
            wrap_validation_errors=False,
        )
        == 10
    )

    # Pydantic ValidationError on bad args propagates raw, not as ToolRetryError.
    with pytest.raises(ValidationError):
        await tool_manager.handle_call(
            ToolCallPart(tool_name='needs_int', args={'x': 'not an int'}),
            wrap_validation_errors=False,
        )
    assert tool_manager.failed_tools == set()

    # ModelRetry from the tool body propagates raw too.
    with pytest.raises(ModelRetry, match='please retry'):
        await tool_manager.handle_call(
            ToolCallPart(tool_name='retrying', args={}),
            wrap_validation_errors=False,
        )
    assert tool_manager.failed_tools == set()

    # Default (wrap=True) still wraps as ToolRetryError and tracks failed tools.
    with pytest.raises(ToolRetryError):
        await tool_manager.handle_call(ToolCallPart(tool_name='needs_int', args={'x': 'not an int'}))
    assert tool_manager.failed_tools == {'needs_int'}

    with pytest.raises(ToolRetryError):
        await tool_manager.handle_call(ToolCallPart(tool_name='retrying', args={}))
    assert tool_manager.failed_tools == {'needs_int', 'retrying'}


async def test_toolset_max_retries_inherits_from_agent():
    """Agent(retries=...) should propagate to user-provided toolsets that don't set max_retries explicitly."""
    attempts: list[int] = []
    toolset = FunctionToolset()

    @toolset.tool_plain
    def always_fails(x: int) -> int:
        """A tool that always fails."""
        attempts.append(x)
        raise ModelRetry('Always fails')

    agent = Agent('test', toolsets=[toolset], retries={'tools': 0, 'output': 0})

    with capture_run_messages() as messages:
        with pytest.raises(UnexpectedModelBehavior, match='exceeded max retries count of 0'):
            await agent.run('call always_fails', model=TestModel())

    # retries=0 means the tool is called once and then fails immediately.
    assert len(attempts) == 1
    assert messages == snapshot(
        [
            ModelRequest(
                parts=[UserPromptPart(content='call always_fails', timestamp=IsDatetime())],
                timestamp=IsNow(tz=timezone.utc),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[ToolCallPart(tool_name='always_fails', args={'x': 0}, tool_call_id=IsStr())],
                usage=RequestUsage(input_tokens=52, output_tokens=4),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )


async def test_toolset_explicit_max_retries_overrides_agent():
    """FunctionToolset(max_retries=X) should take precedence over Agent(retries=Y)."""
    toolset = FunctionToolset(max_retries=2)
    attempts: list[int] = []

    @toolset.tool_plain
    def always_fails(x: int) -> int:
        """A tool that always fails."""
        attempts.append(x)
        raise ModelRetry('Always fails')

    agent = Agent('test', toolsets=[toolset], retries={'tools': 0, 'output': 0})

    with capture_run_messages() as messages:
        with pytest.raises(UnexpectedModelBehavior, match='exceeded max retries count of 2'):
            await agent.run('call always_fails', model=TestModel())

    # Initial call + 2 retries = 3 attempts.
    assert len(attempts) == 3
    assert [type(m).__name__ for m in messages] == snapshot(
        ['ModelRequest', 'ModelResponse', 'ModelRequest', 'ModelResponse', 'ModelRequest', 'ModelResponse']
    )
    retry_parts = [p for m in messages for p in getattr(m, 'parts', []) if isinstance(p, RetryPromptPart)]
    assert [p.content for p in retry_parts] == snapshot(['Always fails', 'Always fails'])


async def test_tool_explicit_retries_overrides_toolset_and_agent():
    """Tool(retries=X) should take precedence over both toolset and agent defaults."""
    attempts: list[int] = []

    def always_fails(x: int) -> int:
        """A tool that always fails."""
        attempts.append(x)
        raise ModelRetry('Always fails')

    toolset = FunctionToolset(tools=[Tool(always_fails, max_retries=3)])
    agent = Agent('test', toolsets=[toolset], retries={'tools': 0, 'output': 0})

    with capture_run_messages() as messages:
        with pytest.raises(UnexpectedModelBehavior, match='exceeded max retries count of 3'):
            await agent.run('call always_fails', model=TestModel())

    # Initial call + 3 retries = 4 attempts.
    assert len(attempts) == 4
    retry_parts = [p for m in messages for p in getattr(m, 'parts', []) if isinstance(p, RetryPromptPart)]
    assert [p.content for p in retry_parts] == snapshot(['Always fails', 'Always fails', 'Always fails'])


async def test_prepare_function_sees_agent_max_retries():
    """Prepare functions should see the agent's default max_retries on ctx when the toolset doesn't set one."""
    captured_max_retries: list[int] = []
    captured_last_attempt: list[bool] = []

    async def capture_prepare(ctx: RunContext, tool_def: ToolDefinition) -> ToolDefinition:
        captured_max_retries.append(ctx.max_retries)
        captured_last_attempt.append(ctx.last_attempt)
        return tool_def

    toolset = FunctionToolset()

    @toolset.tool_plain(prepare=capture_prepare)
    def my_tool(x: int) -> int:
        """A tool."""
        return x

    agent = Agent('test', toolsets=[toolset], retries={'tools': 3, 'output': 3})
    result = await agent.run('call my_tool', model=TestModel())

    assert captured_max_retries[0] == 3
    assert captured_last_attempt[0] is False
    assert result.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[UserPromptPart(content='call my_tool', timestamp=IsDatetime())],
                timestamp=IsNow(tz=timezone.utc),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[ToolCallPart(tool_name='my_tool', args={'x': 0}, tool_call_id=IsStr())],
                usage=RequestUsage(input_tokens=52, output_tokens=4),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='my_tool',
                        content=0,
                        tool_call_id=IsStr(),
                        timestamp=IsDatetime(),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='{"my_tool":0}')],
                usage=RequestUsage(input_tokens=53, output_tokens=7),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )


async def test_toolset_tool_max_retries_none_uses_tool_retries_not_output_retries():
    """When a user toolset leaves `max_retries=None` and `retries != output_retries`, the fallback
    must resolve to the agent's **tool** retry count, not the output retry count.
    Regression: `ctx.max_retries` previously carried `max_output_retries` during `get_tools`."""
    attempts: list[int] = []
    toolset = FunctionToolset()

    @toolset.tool_plain
    def always_fails(x: int) -> int:
        """A tool that always fails."""
        attempts.append(x)
        raise ModelRetry('Always fails')

    agent = Agent('test', toolsets=[toolset], retries={'tools': 1, 'output': 5})

    with capture_run_messages() as messages:
        with pytest.raises(UnexpectedModelBehavior, match='exceeded max retries count of 1'):
            await agent.run('call always_fails', model=TestModel())

    # retries=1 means initial call + 1 retry = 2 attempts, not 6 (which would be output_retries).
    assert len(attempts) == 2
    assert [type(m).__name__ for m in messages] == snapshot(
        ['ModelRequest', 'ModelResponse', 'ModelRequest', 'ModelResponse']
    )
    retry_parts = [p for m in messages for p in getattr(m, 'parts', []) if isinstance(p, RetryPromptPart)]
    assert [p.content for p in retry_parts] == snapshot(['Always fails'])


async def test_prepare_function_sees_tool_retries_not_output_retries():
    """Prepare functions must see the agent's **tool** retry count on `ctx.max_retries`,
    not the output retry count. Regression for a non-output toolset previously receiving the
    output-tool preparation context."""
    captured: list[int] = []

    async def capture_prepare(ctx: RunContext, tool_def: ToolDefinition) -> ToolDefinition:
        captured.append(ctx.max_retries)
        return tool_def

    toolset = FunctionToolset()

    @toolset.tool_plain(prepare=capture_prepare)
    def my_tool(x: int) -> int:
        """A tool."""
        return x

    agent = Agent('test', toolsets=[toolset], retries={'tools': 1, 'output': 5})
    result = await agent.run('call my_tool', model=TestModel())

    assert captured[0] == 1
    assert [type(m).__name__ for m in result.all_messages()] == snapshot(
        ['ModelRequest', 'ModelResponse', 'ModelRequest', 'ModelResponse']
    )


async def test_tool_manager_multiple_failed_tools():
    """Test retry logic when multiple tools fail in the same run step."""

    @dataclass
    class TestDeps:
        pass

    toolset = FunctionToolset[TestDeps]()

    @toolset.tool_plain
    def tool_a(x: int) -> int:
        """Tool A that fails"""
        raise ModelRetry('Tool A fails')

    @toolset.tool_plain
    def tool_b(x: int) -> int:
        """Tool B that fails"""
        raise ModelRetry('Tool B fails')

    @toolset.tool_plain
    def tool_c(x: int) -> int:
        """Tool C that works"""
        return x * 3

    # Create tool manager with max_retries=1, matching what _agent_graph.py sets in a real run
    context = build_run_context(TestDeps(), max_retries=1)
    tool_manager = await ToolManager[TestDeps](toolset).for_run_step(context)

    # Call tool_a - should fail and be added to failed_tools
    with pytest.raises(ToolRetryError):
        await tool_manager.handle_call(ToolCallPart(tool_name='tool_a', args={'x': 1}))
    assert tool_manager.failed_tools == {'tool_a'}

    # Call tool_b - should also fail and be added to failed_tools
    with pytest.raises(ToolRetryError):
        await tool_manager.handle_call(ToolCallPart(tool_name='tool_b', args={'x': 1}))
    assert tool_manager.failed_tools == {'tool_a', 'tool_b'}

    # Call tool_c - should succeed and not be added to failed_tools
    result = await tool_manager.handle_call(ToolCallPart(tool_name='tool_c', args={'x': 2}))
    assert result == 6
    assert tool_manager.failed_tools == {'tool_a', 'tool_b'}  # unchanged

    # Create next run step - should have retry counts for both failed tools
    new_context = build_run_context(TestDeps(), run_step=1, max_retries=1)
    new_tool_manager = await tool_manager.for_run_step(new_context)

    assert new_tool_manager.ctx is not None
    assert new_tool_manager.ctx.retries == {'tool_a': 1, 'tool_b': 1}
    assert new_tool_manager.failed_tools == set()  # reset for new run step


async def test_tool_manager_sequential_tool_call():
    toolset = FunctionToolset()

    @toolset.tool_plain(sequential=True)
    def tool_a(x: int) -> int: ...  # pragma: no cover

    @toolset.tool_plain(sequential=False)
    def tool_b(x: int) -> int: ...  # pragma: no cover

    tool_manager = ToolManager(toolset)

    prepared_tool_manager = await tool_manager.for_run_step(build_run_context(None))

    # A `sequential=True` tool is a per-tool barrier; it no longer forces the whole batch serial.
    assert prepared_tool_manager.is_sequential(ToolCallPart(tool_name='tool_a', args={'x': 1}))
    assert not prepared_tool_manager.is_sequential(ToolCallPart(tool_name='tool_b', args={'x': 1}))

    # The run-scoped mode defaults to parallel and isn't affected by per-tool barriers.
    assert prepared_tool_manager.get_parallel_execution_mode() == 'parallel'
    with ToolManager.parallel_execution_mode('sequential'):
        assert prepared_tool_manager.get_parallel_execution_mode() == 'sequential'


async def test_visit_and_replace():
    toolset1 = FunctionToolset(id='toolset1')
    toolset2 = FunctionToolset(id='toolset2')

    run_ctx = build_run_context(None)

    active_dynamic_toolset = DynamicToolset(toolset_func=lambda ctx: toolset2)
    active_dynamic_toolset = await active_dynamic_toolset.for_run(run_ctx)
    assert isinstance(active_dynamic_toolset, DynamicToolset)
    # for_run with per_run_step=True defers factory evaluation; for_run_step evaluates in-place
    await active_dynamic_toolset.for_run_step(run_ctx)
    assert active_dynamic_toolset._toolset is not None  # pyright: ignore[reportPrivateUsage]
    assert active_dynamic_toolset._toolset is toolset2  # pyright: ignore[reportPrivateUsage]

    inactive_dynamic_toolset = DynamicToolset(toolset_func=lambda ctx: FunctionToolset())

    toolset = CombinedToolset(
        [
            WrapperToolset(toolset1),
            active_dynamic_toolset,
            inactive_dynamic_toolset,
        ]
    )
    visited_toolset = toolset.visit_and_replace(lambda toolset: WrapperToolset(toolset))

    expected_dynamic = DynamicToolset(
        toolset_func=active_dynamic_toolset.toolset_func,
        per_run_step=active_dynamic_toolset.per_run_step,
        id=active_dynamic_toolset._id,  # pyright: ignore[reportPrivateUsage]
    )
    expected_dynamic._toolset = WrapperToolset(toolset2)  # pyright: ignore[reportPrivateUsage]

    assert visited_toolset == CombinedToolset(
        [
            WrapperToolset(WrapperToolset(toolset1)),
            expected_dynamic,
            WrapperToolset(inactive_dynamic_toolset),
        ]
    )


async def test_dynamic_toolset():
    class EnterableToolset(AbstractToolset):
        entered_count = 0
        exited_count = 0

        @property
        def id(self) -> str | None:
            return None  # pragma: no cover

        @property
        def depth_count(self) -> int:
            return self.entered_count - self.exited_count

        async def __aenter__(self) -> Self:
            self.entered_count += 1
            return self

        async def __aexit__(self, *args: Any) -> bool | None:
            self.exited_count += 1
            return None

        async def get_tools(self, ctx: RunContext) -> dict[str, ToolsetTool]:
            return {}

        async def call_tool(self, name: str, tool_args: dict[str, Any], ctx: RunContext, tool: ToolsetTool) -> Any:
            return None  # pragma: no cover

    def toolset_factory(ctx: RunContext) -> AbstractToolset:
        return EnterableToolset()

    original_toolset = DynamicToolset(toolset_func=toolset_factory)

    run_context = build_run_context(None)

    def get_inner_toolset(toolset: DynamicToolset | None) -> EnterableToolset | None:
        assert toolset is not None
        inner_toolset = toolset._toolset  # pyright: ignore[reportPrivateUsage]
        assert isinstance(inner_toolset, EnterableToolset) or inner_toolset is None
        return inner_toolset

    # for_run returns a new per-run copy; per_run_step=True defers factory evaluation
    toolset = await original_toolset.for_run(run_context)
    assert isinstance(toolset, DynamicToolset)
    assert toolset is not original_toolset
    assert toolset._toolset is None  # pyright: ignore[reportPrivateUsage]

    async with toolset:
        # for_run_step evaluates the factory and manages transitions in-place
        step_toolset = await toolset.for_run_step(run_context)
        assert step_toolset is toolset  # returns self after in-place update

        assert (inner_toolset := get_inner_toolset(toolset))
        assert inner_toolset.depth_count == 1

        tools = await toolset.get_tools(run_context)

        # Test that the visitor applies when the toolset is initialized
        def initialized_visitor(visited_toolset: AbstractToolset) -> None:
            assert visited_toolset is inner_toolset

        toolset.apply(initialized_visitor)

    assert get_inner_toolset(toolset) is None

    def uninitialized_visitor(visited_toolset: AbstractToolset) -> None:
        assert visited_toolset is original_toolset

    original_toolset.apply(uninitialized_visitor)

    assert tools == {}


async def test_dynamic_toolset_enter_failure_does_not_exit_unentered_toolset():
    """If the inner toolset's __aenter__ raises, DynamicToolset.__aexit__ must not
    try to exit a toolset that was never entered.

    Reproduces https://github.com/pydantic/pydantic-ai/issues/3542: a per-run-step
    factory produces a fresh toolset each step; if __aenter__ on the new one fails,
    the old logic still stored it and the outer context manager then called
    __aexit__ on an unentered toolset (MCPServer raised "__aexit__ called more
    times than __aenter__").
    """

    class FlakyToolset(AbstractToolset):
        enter_count = 0
        exit_count = 0
        fail_on_enter = False

        @property
        def id(self) -> str | None:
            return None  # pragma: no cover

        async def __aenter__(self) -> Self:
            if self.fail_on_enter:
                raise RuntimeError('enter failed')
            self.enter_count += 1
            return self

        async def __aexit__(self, *args: Any) -> bool | None:
            self.exit_count += 1
            return None

        async def get_tools(self, ctx: RunContext) -> dict[str, ToolsetTool]:
            return {}  # pragma: no cover

        async def call_tool(self, name: str, tool_args: dict[str, Any], ctx: RunContext, tool: ToolsetTool) -> Any:
            return None  # pragma: no cover

    first = FlakyToolset()
    second = FlakyToolset()
    second.fail_on_enter = True
    produced = iter([first, second])

    def factory(ctx: RunContext) -> AbstractToolset:
        return next(produced)

    dynamic = DynamicToolset(toolset_func=factory)
    run_context = build_run_context(None)

    toolset = await dynamic.for_run(run_context)
    assert isinstance(toolset, DynamicToolset)
    async with toolset:
        await toolset.for_run_step(run_context)
        assert first.enter_count == 1
        with pytest.raises(RuntimeError, match='enter failed'):
            await toolset.for_run_step(run_context)
        # After the failed transition, _toolset should be None so __aexit__ is a no-op.
        assert toolset._toolset is None  # pyright: ignore[reportPrivateUsage]

    # Old toolset was exited exactly once during the transition; the failed one
    # was never entered so must never be exited.
    assert first.exit_count == 1
    assert second.exit_count == 0


async def test_dynamic_toolset_aenter_failure_does_not_exit_unentered_toolset():
    """If the initial outer __aenter__ fails, __aexit__ must not try to exit it."""

    class FailingEnterToolset(AbstractToolset):
        exit_count = 0

        @property
        def id(self) -> str | None:
            return None  # pragma: no cover

        async def __aenter__(self) -> Self:
            raise RuntimeError('enter failed')

        async def __aexit__(self, *args: Any) -> bool | None:  # pragma: no cover
            self.exit_count += 1
            return None

        async def get_tools(self, ctx: RunContext) -> dict[str, ToolsetTool]:
            return {}  # pragma: no cover

        async def call_tool(self, name: str, tool_args: dict[str, Any], ctx: RunContext, tool: ToolsetTool) -> Any:
            return None  # pragma: no cover

    inner = FailingEnterToolset()
    dynamic = DynamicToolset(toolset_func=lambda ctx: inner, per_run_step=False)
    run_context = build_run_context(None)

    toolset = await dynamic.for_run(run_context)
    assert isinstance(toolset, DynamicToolset)
    with pytest.raises(RuntimeError, match='enter failed'):
        async with toolset:
            pass  # pragma: no cover

    assert toolset._toolset is None  # pyright: ignore[reportPrivateUsage]
    assert inner.exit_count == 0


async def test_dynamic_toolset_old_aexit_failure_does_not_store_new_toolset():
    """If the old toolset's __aexit__ raises during a per-run-step transition,
    the new toolset must not be stored (and thus not exited) since it was never entered.
    """

    class FailingExitToolset(AbstractToolset):
        entered = False
        exited = False

        @property
        def id(self) -> str | None:
            return None  # pragma: no cover

        async def __aenter__(self) -> Self:
            self.entered = True
            return self

        async def __aexit__(self, *args: Any) -> bool | None:
            self.exited = True
            raise RuntimeError('exit failed')

        async def get_tools(self, ctx: RunContext) -> dict[str, ToolsetTool]:
            return {}  # pragma: no cover

        async def call_tool(self, name: str, tool_args: dict[str, Any], ctx: RunContext, tool: ToolsetTool) -> Any:
            return None  # pragma: no cover

    class TrackedToolset(AbstractToolset):
        entered = False
        exited = False

        @property
        def id(self) -> str | None:
            return None  # pragma: no cover

        async def __aenter__(self) -> Self:  # pragma: no cover
            self.entered = True
            return self

        async def __aexit__(self, *args: Any) -> bool | None:  # pragma: no cover
            self.exited = True
            return None

        async def get_tools(self, ctx: RunContext) -> dict[str, ToolsetTool]:
            return {}  # pragma: no cover

        async def call_tool(self, name: str, tool_args: dict[str, Any], ctx: RunContext, tool: ToolsetTool) -> Any:
            return None  # pragma: no cover

    first = FailingExitToolset()
    second = TrackedToolset()
    produced = iter([first, second])

    def factory(ctx: RunContext) -> AbstractToolset:
        return next(produced)

    dynamic = DynamicToolset(toolset_func=factory)
    run_context = build_run_context(None)

    toolset = await dynamic.for_run(run_context)
    # Suppress the transition failure so we can inspect state afterwards; the
    # outer __aexit__ must not then try to exit the never-entered second toolset.
    with pytest.raises(RuntimeError, match='exit failed'):
        async with toolset:
            await toolset.for_run_step(run_context)
            assert first.entered
            await toolset.for_run_step(run_context)  # raises on old.__aexit__
            pass  # pragma: no cover

    assert first.exited
    assert not second.entered
    assert not second.exited


async def test_dynamic_toolset_empty():
    def no_toolset_func(ctx: RunContext) -> None:
        return None  # pragma: no cover

    original_toolset = DynamicToolset(toolset_func=no_toolset_func)

    run_context = build_run_context(None)

    # for_run evaluates the factory; factory returns None so _toolset stays None
    toolset = await original_toolset.for_run(run_context)
    assert isinstance(toolset, DynamicToolset)
    assert toolset._toolset is None  # pyright: ignore[reportPrivateUsage]

    async with toolset:
        tools = await toolset.get_tools(run_context)

        assert tools == {}

        assert toolset._toolset is None  # pyright: ignore[reportPrivateUsage]


def test_dynamic_toolset_id():
    """Test that DynamicToolset can have an id set."""

    def toolset_func(ctx: RunContext) -> FunctionToolset:
        return FunctionToolset()  # pragma: no cover

    # No id by default
    toolset_no_id = DynamicToolset(toolset_func=toolset_func)
    assert toolset_no_id.id is None

    # Explicit id
    toolset_with_id = DynamicToolset(toolset_func=toolset_func, id='my_dynamic_toolset')
    assert toolset_with_id.id == 'my_dynamic_toolset'


async def test_wrapper_toolsets_delegate_instructions():
    """Test that wrapper toolsets properly delegate instructions calls."""
    base_instructions = 'Follow the base toolset instructions carefully.'
    base_toolset = MockToolsetWithInstructions(instructions=base_instructions)
    ctx = build_run_context(None)

    # Test PrefixedToolset delegation
    prefixed_toolset = base_toolset.prefixed('test')
    assert await prefixed_toolset.get_instructions(ctx) == base_instructions

    # Test FilteredToolset delegation
    def allow_all_filter(ctx: RunContext[Any], tool_def: ToolDefinition) -> bool:
        return True

    assert allow_all_filter(ctx, ToolDefinition(name='test', description='', parameters_json_schema={})) is True
    filtered_toolset = base_toolset.filtered(allow_all_filter)
    assert await filtered_toolset.get_instructions(ctx) == base_instructions

    # Test RenamedToolset delegation
    rename_map = {'old_name': 'new_name'}
    renamed_toolset = base_toolset.renamed(rename_map)
    assert await renamed_toolset.get_instructions(ctx) == base_instructions

    # Test ApprovalRequiredToolset delegation
    approval_toolset = base_toolset.approval_required()
    assert await approval_toolset.get_instructions(ctx) == base_instructions

    # Test PreparedToolset delegation
    async def prepare_func(ctx: RunContext[Any], tools: list[ToolDefinition]) -> list[ToolDefinition]:
        return tools

    assert await prepare_func(ctx, []) == []
    prepared_toolset = base_toolset.prepared(prepare_func)
    assert await prepared_toolset.get_instructions(ctx) == base_instructions


async def test_combined_toolset_instructions():
    """Test that CombinedToolset aggregates instructions from all contained toolsets."""
    instructions1 = 'Instructions from toolset 1.'
    instructions2 = 'Instructions from toolset 2.'

    toolset1 = MockToolsetWithInstructions(instructions=instructions1, id='toolset1')
    toolset2 = MockToolsetWithInstructions(instructions=instructions2, id='toolset2')
    toolset3 = MockToolsetWithInstructions(instructions=None, id='toolset3')  # No instructions

    combined = CombinedToolset([toolset1, toolset2, toolset3])
    ctx = build_run_context(None)

    # CombinedToolset aggregates non-None instructions from all contained toolsets as a list
    instructions = await combined.get_instructions(ctx)
    assert instructions == ['Instructions from toolset 1.', 'Instructions from toolset 2.']


async def test_combined_toolset_instructions_all_none():
    """Test that CombinedToolset returns None when all toolsets have no instructions."""
    toolset1 = MockToolsetWithInstructions(instructions=None, id='toolset1')
    toolset2 = MockToolsetWithInstructions(instructions=None, id='toolset2')

    combined = CombinedToolset([toolset1, toolset2])
    ctx = build_run_context(None)

    instructions = await combined.get_instructions(ctx)
    assert instructions is None


async def test_combined_toolset_instructions_empty():
    """Test that CombinedToolset returns None when no toolsets are provided."""
    combined = CombinedToolset([])
    ctx = build_run_context(None)

    instructions = await combined.get_instructions(ctx)
    assert instructions is None


def test_agent_toolset_decorator_id():
    """Test that @agent.toolset decorator requires explicit id or defaults to None."""
    from pydantic_ai.models.test import TestModel

    agent = Agent(TestModel())

    @agent.toolset
    def my_tools(ctx: RunContext) -> FunctionToolset:
        return FunctionToolset()  # pragma: no cover

    @agent.toolset(id='custom_id')
    def other_tools(ctx: RunContext) -> FunctionToolset:
        return FunctionToolset()  # pragma: no cover

    # The toolsets are DynamicToolsets with None or explicit ids
    toolsets = agent.toolsets
    assert len(toolsets) == 3  # FunctionToolset for agent tools + 2 dynamic toolsets

    # First is the agent's own FunctionToolset
    assert isinstance(toolsets[0], FunctionToolset)

    # Second toolset without explicit id should have None
    assert isinstance(toolsets[1], DynamicToolset)
    assert toolsets[1].id is None

    # Third toolset should have explicit id
    assert isinstance(toolsets[2], DynamicToolset)
    assert toolsets[2].id == 'custom_id'


async def test_function_toolset_get_instructions_string():
    """FunctionToolset with a string instruction returns it via get_instructions."""
    toolset = FunctionToolset(instructions='Always use tool X for math.')

    ctx = build_run_context(None)
    result = await toolset.get_instructions(ctx)
    assert result == [InstructionPart(content='Always use tool X for math.', dynamic=False)]


async def test_function_toolset_get_instructions_function():
    """FunctionToolset with a function instruction calls it via get_instructions."""
    toolset = FunctionToolset(instructions=lambda: 'Use search for lookups.')

    ctx = build_run_context(None)
    result = await toolset.get_instructions(ctx)
    assert result == [InstructionPart(content='Use search for lookups.', dynamic=True)]


async def test_function_toolset_get_instructions_with_ctx():
    """FunctionToolset instruction function can access RunContext."""

    def my_instructions(ctx: RunContext[str]) -> str:
        return f'Deps are: {ctx.deps}'

    toolset = FunctionToolset[str](instructions=my_instructions)

    ctx = build_run_context('hello')
    result = await toolset.get_instructions(ctx)
    assert result == [InstructionPart(content='Deps are: hello', dynamic=True)]


async def test_function_toolset_get_instructions_async():
    """FunctionToolset with an async instruction function works."""

    async def my_instructions() -> str:
        return 'Async instructions here.'

    toolset = FunctionToolset(instructions=my_instructions)

    ctx = build_run_context(None)
    result = await toolset.get_instructions(ctx)
    assert result == [InstructionPart(content='Async instructions here.', dynamic=True)]


async def test_function_toolset_get_instructions_multiple():
    """FunctionToolset with a sequence of instructions returns them as a list."""
    toolset = FunctionToolset(instructions=['First instruction.', lambda: 'Second instruction.'])

    ctx = build_run_context(None)
    result = await toolset.get_instructions(ctx)
    assert result == [
        InstructionPart(content='First instruction.', dynamic=False),
        InstructionPart(content='Second instruction.', dynamic=True),
    ]


async def test_function_toolset_get_instructions_none_default():
    """FunctionToolset without instructions returns None."""
    toolset = FunctionToolset()

    ctx = build_run_context(None)
    result = await toolset.get_instructions(ctx)
    assert result is None


async def test_function_toolset_instructions_decorator():
    """The @toolset.instructions decorator registers instruction functions."""
    toolset = FunctionToolset()

    @toolset.instructions
    def my_instructions() -> str:
        return 'Use tool Y for data processing.'

    ctx = build_run_context(None)
    result = await toolset.get_instructions(ctx)
    assert result == [InstructionPart(content='Use tool Y for data processing.', dynamic=True)]


async def test_function_toolset_instructions_decorator_with_ctx():
    """The @toolset.instructions decorator works with RunContext."""
    toolset = FunctionToolset[int]()

    @toolset.instructions
    def my_instructions(ctx: RunContext[int]) -> str:
        return f'Dep value: {ctx.deps}'

    ctx = build_run_context(42)
    result = await toolset.get_instructions(ctx)
    assert result == [InstructionPart(content='Dep value: 42', dynamic=True)]


async def test_function_toolset_instructions_decorator_combined_with_constructor():
    """Constructor instructions and decorator instructions are combined."""
    toolset = FunctionToolset(instructions='From constructor.')

    @toolset.instructions
    def extra() -> str:
        return 'From decorator.'

    ctx = build_run_context(None)
    result = await toolset.get_instructions(ctx)
    assert result == [
        InstructionPart(content='From constructor.', dynamic=False),
        InstructionPart(content='From decorator.', dynamic=True),
    ]


async def test_function_toolset_instructions_none_filtered():
    """Instructions returning None are excluded."""
    toolset = FunctionToolset(instructions=[lambda: None, 'Only this.'])

    ctx = build_run_context(None)
    result = await toolset.get_instructions(ctx)
    assert result == [InstructionPart(content='Only this.', dynamic=False)]


async def test_function_toolset_empty_string_instructions():
    """Empty string instructions are filtered out, returning None."""
    toolset = FunctionToolset(instructions='')

    ctx = build_run_context(None)
    result = await toolset.get_instructions(ctx)
    assert result is None


async def test_function_toolset_whitespace_only_instructions():
    """Whitespace-only instructions are filtered out, returning None."""
    toolset = FunctionToolset(instructions='   \n\n  ')

    ctx = build_run_context(None)
    result = await toolset.get_instructions(ctx)
    assert result is None


async def test_wrapper_toolset_passes_through_instructions():
    """WrapperToolset delegates get_instructions to wrapped toolset."""
    inner = FunctionToolset(instructions='Inner instructions.')
    wrapped = inner.prefixed('my')

    ctx = build_run_context(None)
    result = await wrapped.get_instructions(ctx)
    assert result == [InstructionPart(content='Inner instructions.', dynamic=False)]


async def test_combined_toolset_aggregates_instructions():
    """CombinedToolset gathers instructions from all children."""
    ts1 = FunctionToolset(instructions='Toolset 1 instructions.')
    ts2 = FunctionToolset(instructions='Toolset 2 instructions.')
    combined = CombinedToolset([ts1, ts2])

    ctx = build_run_context(None)
    result = await combined.get_instructions(ctx)
    assert result == [
        InstructionPart(content='Toolset 1 instructions.', dynamic=False),
        InstructionPart(content='Toolset 2 instructions.', dynamic=False),
    ]


async def test_combined_toolset_skips_none_instructions():
    """CombinedToolset skips toolsets that return None for instructions."""
    ts1 = FunctionToolset(instructions='Only from ts1.')
    ts2 = FunctionToolset()  # No instructions
    combined = CombinedToolset([ts1, ts2])

    ctx = build_run_context(None)
    result = await combined.get_instructions(ctx)
    assert result == [InstructionPart(content='Only from ts1.', dynamic=False)]


async def test_combined_toolset_all_none_returns_none():
    """CombinedToolset returns None when all children return None."""
    ts1 = FunctionToolset()
    ts2 = FunctionToolset()
    combined = CombinedToolset([ts1, ts2])

    ctx = build_run_context(None)
    result = await combined.get_instructions(ctx)
    assert result is None


async def test_combined_toolset_with_nested_list_instructions():
    """CombinedToolset flattens list[str] results from child CombinedToolsets (covers combined.py list branch)."""
    ts1 = FunctionToolset(instructions='Instruction A.')
    ts2 = FunctionToolset(instructions='Instruction B.')
    inner = CombinedToolset([ts1, ts2])  # returns list[str]

    ts3 = FunctionToolset(instructions='Instruction C.')
    outer = CombinedToolset([inner, ts3])
    ctx = build_run_context(None)

    result = await outer.get_instructions(ctx)
    assert result == [
        InstructionPart(content='Instruction A.', dynamic=False),
        InstructionPart(content='Instruction B.', dynamic=False),
        InstructionPart(content='Instruction C.', dynamic=False),
    ]


async def test_combined_toolset_cancels_siblings_on_get_tools_failure():
    """When one child's get_tools fails, siblings are cancelled instead of leaking as orphan tasks."""
    sibling_completed = False

    class FailingToolset(WrapperToolset[Any]):
        async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
            raise RuntimeError('boom')

    class SlowToolset(WrapperToolset[Any]):
        async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
            nonlocal sibling_completed
            await anyio.sleep(0.1)
            sibling_completed = True  # pragma: no cover
            return await self.wrapped.get_tools(ctx)  # pragma: no cover

    inner = FunctionToolset()
    combined = CombinedToolset([FailingToolset(inner), SlowToolset(inner)])
    ctx = build_run_context(None)

    with pytest.raises(RuntimeError, match='boom'):
        await combined.get_tools(ctx)

    await anyio.sleep(0.2)
    assert sibling_completed is False


async def test_combined_toolset_get_tools_preserves_exception_cause():
    """Unwrapping the single-failure exception must preserve the original `__cause__` chain."""
    original = ValueError('underlying')

    class FailingToolset(WrapperToolset[Any]):
        async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
            raise RuntimeError('wrapper') from original

    inner = FunctionToolset()
    combined = CombinedToolset([FailingToolset(inner)])
    ctx = build_run_context(None)

    with pytest.raises(RuntimeError, match='wrapper') as exc_info:
        await combined.get_tools(ctx)

    assert exc_info.value.__cause__ is original


async def test_combined_toolset_get_tools_raises_group_on_multiple_failures():
    """When multiple children fail concurrently, their errors surface as an ExceptionGroup."""

    @dataclass
    class RaisingToolset(WrapperToolset[Any]):
        message: str = ''

        async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
            await anyio.sleep(0)
            raise RuntimeError(self.message)

    inner = FunctionToolset()
    combined = CombinedToolset(
        [RaisingToolset(wrapped=inner, message='first'), RaisingToolset(wrapped=inner, message='second')]
    )
    ctx = build_run_context(None)

    with pytest.raises(BaseExceptionGroup) as exc_info:
        await combined.get_tools(ctx)

    messages = {str(e) for e in exc_info.value.exceptions}
    assert messages == {'first', 'second'}


async def test_dynamic_toolset_instructions_before_resolution():
    """DynamicToolset returns None for instructions before get_tools resolves it."""
    dynamic = DynamicToolset(lambda ctx: FunctionToolset(instructions='Dynamic instructions.'))

    ctx = build_run_context(None)
    # Before get_tools is called, _toolset is None
    result = await dynamic.get_instructions(ctx)
    assert result is None


async def test_dynamic_toolset_instructions_after_resolution():
    """DynamicToolset delegates instructions after for_run_step resolves it."""

    def make_toolset(ctx: RunContext) -> FunctionToolset:
        return FunctionToolset(instructions='Dynamic instructions.')

    dynamic = DynamicToolset(make_toolset)

    ctx = build_run_context(None)
    # for_run_step triggers factory resolution for per_run_step=True
    await dynamic.for_run_step(ctx)
    result = await dynamic.get_instructions(ctx)
    assert result == [InstructionPart(content='Dynamic instructions.', dynamic=False)]


async def test_toolset_instructions_in_agent():
    """Toolset instructions appear in the model request when added to an agent."""
    from pydantic_ai import Agent

    toolset = FunctionToolset(instructions='Always use my_tool correctly.')

    @toolset.tool_plain
    def my_tool() -> str:
        """A simple tool."""
        return 'done'

    agent = Agent(TestModel(), toolsets=[toolset])
    result = await agent.run('Hello')
    first_message = result.all_messages()[0]
    assert first_message.instructions == 'Always use my_tool correctly.'  # type: ignore[union-attr]


async def test_dynamic_toolset_instructions_on_first_request():
    """Instructions from a DynamicToolset are present on the very first model request."""
    from pydantic_ai import Agent

    def make_toolset(ctx: RunContext) -> FunctionToolset:
        ts = FunctionToolset(instructions='Dynamic tool instructions.')

        @ts.tool_plain
        def my_dynamic_tool() -> str:
            """A tool inside the dynamic toolset."""
            return 'done'

        return ts

    agent = Agent(TestModel(), toolsets=[DynamicToolset(make_toolset)])
    result = await agent.run('Hello')
    first_message = result.all_messages()[0]
    assert first_message.instructions == 'Dynamic tool instructions.'  # type: ignore[union-attr]


async def test_resume_without_prompt_dynamic_toolset_instructions_resolve_once_for_request_step():
    """Resuming from a trailing ModelResponse should resolve dynamic toolsets only once for the next request step."""
    run_steps: list[int] = []

    def make_toolset(ctx: RunContext) -> FunctionToolset:
        run_steps.append(ctx.run_step)
        return FunctionToolset(instructions=f'Dynamic instructions at step {ctx.run_step}.')

    agent = Agent(TestModel(custom_output_text='done'), toolsets=[DynamicToolset(make_toolset)])
    result = await agent.run(message_history=[ModelResponse(parts=[TextPart(content='previous')])])

    # The resume pre-check and request preparation should use the same run_step.
    assert run_steps == [1]

    requests = [m for m in result.all_messages() if isinstance(m, ModelRequest)]
    assert requests[-1].instructions == 'Dynamic instructions at step 1.'


async def test_resume_without_prompt_dynamic_toolset_with_tool_calls_resolve_once_for_request_step():
    """Resuming from a trailing ModelResponse with ToolCallParts exercises the _handle_tool_calls path.

    This is the more common resume scenario and ensures dynamic toolsets are resolved only once even
    when the code path goes through _handle_tool_calls (which calls for_run_step).
    """
    run_steps: list[int] = []

    def make_toolset(ctx: RunContext) -> FunctionToolset:
        run_steps.append(ctx.run_step)
        toolset = FunctionToolset(instructions=f'Dynamic instructions at step {ctx.run_step}.')

        @toolset.tool_plain
        def greet(name: str) -> str:
            """Greet someone by name."""
            return f'Hello {name}!'

        return toolset

    agent = Agent(TestModel(custom_output_text='done'), toolsets=[DynamicToolset(make_toolset)])
    result = await agent.run(
        message_history=[ModelResponse(parts=[ToolCallPart(tool_name='greet', args={'name': 'Alice'})])]
    )

    # The toolset factory is evaluated multiple times:
    # - step 1: UserPromptNode pre-check (aligned with upcoming request step)
    # - step 0: CallToolsNode._handle_tool_calls (current run_step before increment)
    # - step 1: ModelRequestNode._prepare_request (after run_step increment)
    # The important thing is that the first and last evaluations use step 1.
    assert run_steps[0] == 1
    assert run_steps[-1] == 1

    requests = [m for m in result.all_messages() if isinstance(m, ModelRequest)]
    # The last request should have instructions prepared at step 1
    assert requests[-1].instructions == 'Dynamic instructions at step 1.'


async def test_toolset_instructions_combined_with_agent_instructions():
    """Toolset instructions are appended after agent-level instructions."""
    from pydantic_ai import Agent

    toolset = FunctionToolset(instructions='Use search for lookups.')

    @toolset.tool_plain
    def search() -> str:
        """Search for information."""
        return 'results'

    agent = Agent(TestModel(), instructions='You are a helpful assistant.', toolsets=[toolset])
    result = await agent.run('Hello')
    first_message = result.all_messages()[0]
    assert first_message.instructions == 'You are a helpful assistant.\n\nUse search for lookups.'  # type: ignore[union-attr]


async def test_multiple_toolset_instructions_in_agent():
    """Multiple toolsets with instructions are all included."""
    from pydantic_ai import Agent

    ts1 = FunctionToolset(instructions='Use calculator for math.')

    @ts1.tool_plain
    def calculator() -> str:
        """Evaluate a math expression."""
        return '4'

    ts2 = FunctionToolset(instructions='Use search for lookups.')

    @ts2.tool_plain
    def search() -> str:
        """Search for information."""
        return 'results'

    agent = Agent(TestModel(), toolsets=[ts1, ts2])
    result = await agent.run('Hello')
    first_message = result.all_messages()[0]
    assert first_message.instructions == 'Use calculator for math.\n\nUse search for lookups.'  # type: ignore[union-attr]


async def test_toolset_instructions_alone_satisfy_validation():
    """Toolset instructions alone (no user prompt, no agent instructions, no history) are sufficient to run."""
    from pydantic_ai import Agent

    toolset = FunctionToolset(instructions='Always use my_tool correctly.')

    @toolset.tool_plain
    def my_tool() -> str:
        """A simple tool."""
        return 'done'

    agent = Agent(TestModel(), toolsets=[toolset])
    result = await agent.run()
    first_message = result.all_messages()[0]
    assert first_message.instructions == 'Always use my_tool correctly.'  # type: ignore[union-attr]
    assert first_message.parts == []


async def test_no_input_raises_without_toolset_instructions():
    """Without any prompt, instructions, or history, the agent raises UserError."""
    from pydantic_ai import Agent

    agent = Agent(TestModel())
    with pytest.raises(UserError, match='No message history, user prompt, or instructions provided'):
        await agent.run()


class StatefulToolset(AbstractToolset):
    """A custom stateful toolset for testing for_run/for_run_step."""

    def __init__(self, *, call_count: int = 0, id: str | None = 'stateful'):
        self.call_count = call_count
        self._id = id

    @property
    def id(self) -> str | None:
        return self._id  # pragma: no cover

    async def for_run(self, ctx: RunContext) -> AbstractToolset:
        return StatefulToolset(call_count=0, id=self._id)

    async def for_run_step(self, ctx: RunContext) -> AbstractToolset:
        return StatefulToolset(call_count=self.call_count + 1, id=self._id)

    async def get_tools(self, ctx: RunContext) -> dict[str, ToolsetTool]:
        return {}  # pragma: no cover

    async def call_tool(self, name: str, tool_args: dict[str, Any], ctx: RunContext, tool: ToolsetTool) -> Any:
        self.call_count += 1  # pragma: no cover
        return self.call_count  # pragma: no cover


async def test_for_run_returns_fresh_instance():
    """Custom stateful toolset with for_run returning fresh instance."""
    original = StatefulToolset(call_count=5)
    ctx = build_run_context(None)

    run_toolset = await original.for_run(ctx)
    assert run_toolset is not original
    assert isinstance(run_toolset, StatefulToolset)
    assert run_toolset.call_count == 0
    assert original.call_count == 5  # original unchanged


async def test_for_run_step_default_returns_self():
    """Default for_run_step returns self for toolsets that don't override it."""
    toolset = FunctionToolset()
    ctx = build_run_context(None)

    step_toolset = await toolset.for_run_step(ctx)
    assert step_toolset is toolset


async def test_for_run_step_returns_new_instance():
    """StatefulToolset.for_run_step returns a new instance with bumped step counter."""
    toolset = StatefulToolset(call_count=3)
    ctx = build_run_context(None)

    step_toolset = await toolset.for_run_step(ctx)
    assert step_toolset is not toolset
    assert isinstance(step_toolset, StatefulToolset)
    assert step_toolset.call_count == 4
    assert toolset.call_count == 3  # original unchanged


async def test_wrapper_propagates_for_run():
    """Wrapper toolsets correctly propagate for_run to the wrapped toolset."""
    inner = StatefulToolset(call_count=10)
    wrapper = WrapperToolset(inner)
    ctx = build_run_context(None)

    run_wrapper = await wrapper.for_run(ctx)
    assert run_wrapper is not wrapper  # different because inner changed
    assert isinstance(run_wrapper, WrapperToolset)
    inner_after = run_wrapper.wrapped
    assert isinstance(inner_after, StatefulToolset)
    assert inner_after.call_count == 0  # fresh


async def test_wrapper_propagates_for_run_no_change():
    """Wrapper returns self when wrapped toolset returns self from for_run."""
    inner = FunctionToolset()  # FunctionToolset.for_run returns self
    wrapper = WrapperToolset(inner)
    ctx = build_run_context(None)

    run_wrapper = await wrapper.for_run(ctx)
    assert run_wrapper is wrapper


async def test_combined_propagates_for_run():
    """CombinedToolset propagates for_run to all children."""
    stateful = StatefulToolset(call_count=7)
    static = FunctionToolset()
    combined = CombinedToolset([stateful, static])
    ctx = build_run_context(None)

    run_combined = await combined.for_run(ctx)
    assert run_combined is not combined
    assert isinstance(run_combined, CombinedToolset)
    assert isinstance(run_combined.toolsets[0], StatefulToolset)
    assert run_combined.toolsets[0].call_count == 0
    assert run_combined.toolsets[1] is static  # unchanged


async def test_combined_for_run_always_fresh():
    """CombinedToolset.for_run always returns a new instance for per-run isolation."""
    static1 = FunctionToolset(id='a')
    static2 = FunctionToolset(id='b')
    combined = CombinedToolset([static1, static2])
    ctx = build_run_context(None)

    run_combined = await combined.for_run(ctx)
    assert run_combined is not combined
    assert isinstance(run_combined, CombinedToolset)
    # Children are unchanged (their for_run returns self)
    assert run_combined.toolsets[0] is static1
    assert run_combined.toolsets[1] is static2


async def test_wrapper_propagates_for_run_step_no_change():
    """Wrapper returns self when wrapped toolset returns self from for_run_step."""
    inner = FunctionToolset()  # FunctionToolset.for_run_step returns self
    wrapper = WrapperToolset(inner)
    ctx = build_run_context(None)

    step_wrapper = await wrapper.for_run_step(ctx)
    assert step_wrapper is wrapper


async def test_wrapper_propagates_for_run_step():
    """Wrapper creates new wrapper when wrapped toolset returns new instance from for_run_step."""
    inner = StatefulToolset(call_count=10)
    wrapper = WrapperToolset(inner)
    ctx = build_run_context(None)

    step_wrapper = await wrapper.for_run_step(ctx)
    assert step_wrapper is not wrapper
    assert isinstance(step_wrapper, WrapperToolset)
    inner_after = step_wrapper.wrapped
    assert isinstance(inner_after, StatefulToolset)
    assert inner_after.call_count == 11  # bumped by for_run_step


async def test_combined_propagates_for_run_step_no_change():
    """CombinedToolset returns self when no children change from for_run_step."""
    static1 = FunctionToolset(id='a')
    static2 = FunctionToolset(id='b')
    combined = CombinedToolset([static1, static2])
    ctx = build_run_context(None)

    step_combined = await combined.for_run_step(ctx)
    assert step_combined is combined


async def test_combined_propagates_for_run_step():
    """CombinedToolset creates new combined when a child returns new instance from for_run_step."""
    stateful = StatefulToolset(call_count=7)
    static = FunctionToolset()
    combined = CombinedToolset([stateful, static])
    ctx = build_run_context(None)

    step_combined = await combined.for_run_step(ctx)
    assert step_combined is not combined
    assert isinstance(step_combined, CombinedToolset)
    assert isinstance(step_combined.toolsets[0], StatefulToolset)
    assert step_combined.toolsets[0].call_count == 8  # bumped by for_run_step
    assert step_combined.toolsets[1] is static  # unchanged


async def test_dynamic_toolset_for_run_step_manages_transitions():
    """DynamicToolset with per_run_step=True manages internal transitions via for_run_step."""
    call_count = 0

    def factory(ctx: RunContext) -> FunctionToolset:
        nonlocal call_count
        call_count += 1
        return FunctionToolset(id=f'step-{call_count}')

    original = DynamicToolset(toolset_func=factory, per_run_step=True)
    ctx = build_run_context(None)

    # for_run creates a fresh copy without evaluating factory
    run_toolset = await original.for_run(ctx)
    assert isinstance(run_toolset, DynamicToolset)
    assert run_toolset._toolset is None  # pyright: ignore[reportPrivateUsage]
    assert call_count == 0

    async with run_toolset:
        # for_run_step evaluates the factory
        step1 = await run_toolset.for_run_step(ctx)
        assert step1 is run_toolset  # returns self after in-place update
        assert call_count == 1
        assert run_toolset._toolset is not None  # pyright: ignore[reportPrivateUsage]

        # Second for_run_step re-evaluates (new toolset each time)
        step2 = await run_toolset.for_run_step(ctx)
        assert step2 is run_toolset
        assert call_count == 2

    assert original._toolset is None  # pyright: ignore[reportPrivateUsage]


async def test_dynamic_toolset_for_run_step_same_instance_skips_transition():
    """DynamicToolset skips transition when factory returns the same instance."""
    stable_toolset = FunctionToolset(id='stable')

    def factory(ctx: RunContext) -> FunctionToolset:
        return stable_toolset

    original = DynamicToolset(toolset_func=factory, per_run_step=True)
    ctx = build_run_context(None)

    run_toolset = await original.for_run(ctx)
    assert isinstance(run_toolset, DynamicToolset)

    async with run_toolset:
        # First step: evaluates factory, sets _toolset
        step1 = await run_toolset.for_run_step(ctx)
        assert step1 is run_toolset
        assert run_toolset._toolset is stable_toolset  # pyright: ignore[reportPrivateUsage]

        # Second step: factory returns same instance, early return without transition
        step2 = await run_toolset.for_run_step(ctx)
        assert step2 is run_toolset
        assert run_toolset._toolset is stable_toolset  # pyright: ignore[reportPrivateUsage]


async def test_dynamic_toolset_for_run_step_factory_returns_none():
    """DynamicToolset handles factory returning None after previously returning a toolset."""
    call_count = 0

    def factory(ctx: RunContext) -> FunctionToolset | None:
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return FunctionToolset(id='first')
        return None

    original = DynamicToolset(toolset_func=factory, per_run_step=True)
    ctx = build_run_context(None)

    run_toolset = await original.for_run(ctx)
    assert isinstance(run_toolset, DynamicToolset)

    async with run_toolset:
        # First step: factory returns a toolset
        await run_toolset.for_run_step(ctx)
        assert run_toolset._toolset is not None  # pyright: ignore[reportPrivateUsage]

        # Second step: factory returns None — old toolset exited, new is None
        await run_toolset.for_run_step(ctx)
        assert run_toolset._toolset is None  # pyright: ignore[reportPrivateUsage]

        # Tools should be empty when _toolset is None
        tools = await run_toolset.get_tools(ctx)
        assert tools == {}


async def test_dynamic_toolset_per_run_step_false_for_run_evaluates():
    """DynamicToolset with per_run_step=False evaluates factory in for_run."""
    call_count = 0

    def factory(ctx: RunContext) -> FunctionToolset:
        nonlocal call_count
        call_count += 1
        return FunctionToolset()

    original = DynamicToolset(toolset_func=factory, per_run_step=False)
    ctx = build_run_context(None)

    run_toolset = await original.for_run(ctx)
    assert isinstance(run_toolset, DynamicToolset)
    assert call_count == 1
    assert run_toolset._toolset is not None  # pyright: ignore[reportPrivateUsage]

    # for_run_step returns self (no re-evaluation)
    step_toolset = await run_toolset.for_run_step(ctx)
    assert step_toolset is run_toolset
    assert call_count == 1


async def test_concurrent_runs_dont_share_state():
    """Multiple concurrent runs don't share state on stateful toolsets."""
    import asyncio

    from pydantic_ai import Agent
    from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart
    from pydantic_ai.models.function import AgentInfo, FunctionModel

    call_counts: list[int] = []

    class CountingToolset(AbstractToolset):
        def __init__(self) -> None:
            self.count = 0

        @property
        def id(self) -> str | None:
            return 'counting'  # pragma: no cover

        async def for_run(self, ctx: RunContext) -> AbstractToolset:
            return CountingToolset()

        async def get_tools(self, ctx: RunContext) -> dict[str, ToolsetTool]:
            self.count += 1
            call_counts.append(self.count)
            return {}

        async def call_tool(self, name: str, tool_args: dict[str, Any], ctx: RunContext, tool: ToolsetTool) -> Any:
            pass  # pragma: no cover

    def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        return ModelResponse(parts=[TextPart('Done')])

    agent = Agent(FunctionModel(respond), toolsets=[CountingToolset()])

    # Run two concurrent agent runs
    results = await asyncio.gather(agent.run('Hello'), agent.run('World'))

    assert results[0].output == 'Done'
    assert results[1].output == 'Done'
    # Each run should have its own count (1), not share state (1, 2)
    assert all(c == 1 for c in call_counts)


def test_include_return_schemas_toolset():
    """IncludeReturnSchemasToolset sets include_return_schema=True on wrapped tools."""

    def my_tool(x: int) -> int:
        return x

    toolset = FunctionToolset(tools=[my_tool])
    test_model = TestModel()
    agent = Agent(test_model, toolsets=[toolset.include_return_schemas()])
    agent.run_sync('test')

    params = test_model.last_model_request_parameters
    assert params is not None
    td = next(td for td in params.function_tools if td.name == 'my_tool')
    assert td.include_return_schema is True
    assert 'Return schema' in (td.description or '')


def test_set_metadata_toolset():
    """SetMetadataToolset merges metadata onto all wrapped tools."""

    def my_tool(x: int) -> int:
        return x

    toolset = FunctionToolset(tools=[my_tool])
    test_model = TestModel()
    agent = Agent(test_model, toolsets=[toolset.with_metadata(code_mode=True, priority=1)])
    agent.run_sync('test')

    params = test_model.last_model_request_parameters
    assert params is not None
    td = next(td for td in params.function_tools if td.name == 'my_tool')
    assert td.metadata is not None
    assert td.metadata['code_mode'] is True
    assert td.metadata['priority'] == 1


async def test_filtered_toolset_async_filter():
    """FilteredToolset supports async filter functions."""

    def tool_a(x: int) -> int:
        return x

    def tool_b(x: str) -> str:
        return x  # pragma: no cover

    async def async_filter(ctx: RunContext, td: ToolDefinition) -> bool:
        return td.name == 'tool_a'

    toolset = FunctionToolset(tools=[tool_a, tool_b])
    test_model = TestModel()
    agent = Agent(test_model, toolsets=[toolset.filtered(async_filter)])
    await agent.run('test')

    params = test_model.last_model_request_parameters
    assert params is not None
    tool_names = [td.name for td in params.function_tools]
    assert tool_names == ['tool_a']


def test_set_tool_metadata_capability():
    """SetToolMetadata capability merges metadata onto selected tools."""
    from pydantic_ai.capabilities import SetToolMetadata

    def tool_a(x: int) -> int:
        return x

    def tool_b(x: str) -> str:
        return x

    test_model = TestModel()
    agent = Agent(
        test_model,
        tools=[tool_a, tool_b],
        capabilities=[SetToolMetadata(tools=['tool_a'], code_mode=True)],
    )
    agent.run_sync('test')

    params = test_model.last_model_request_parameters
    assert params is not None
    td_a = next(td for td in params.function_tools if td.name == 'tool_a')
    td_b = next(td for td in params.function_tools if td.name == 'tool_b')
    assert td_a.metadata is not None
    assert td_a.metadata['code_mode'] is True
    # tool_b should not have the metadata
    assert td_b.metadata is None or 'code_mode' not in td_b.metadata


def test_set_tool_metadata_capability_with_async_selector():
    """SetToolMetadata with async callable selector."""
    from pydantic_ai.capabilities import SetToolMetadata

    def tool_a(x: int) -> int:
        return x

    def tool_b(x: str) -> str:
        return x

    async def only_tool_a(ctx: RunContext, td: ToolDefinition) -> bool:
        return td.name == 'tool_a'

    test_model = TestModel()
    agent = Agent(
        test_model,
        tools=[tool_a, tool_b],
        capabilities=[SetToolMetadata(tools=only_tool_a, tagged=True)],
    )
    agent.run_sync('test')

    params = test_model.last_model_request_parameters
    assert params is not None
    td_a = next(td for td in params.function_tools if td.name == 'tool_a')
    td_b = next(td for td in params.function_tools if td.name == 'tool_b')
    assert td_a.metadata is not None
    assert td_a.metadata['tagged'] is True
    assert td_b.metadata is None or 'tagged' not in td_b.metadata


def test_set_tool_metadata_capability_with_bare_string_selector():
    """SetToolMetadata with a bare string selector matches by exact name, not substring."""
    from pydantic_ai.capabilities import SetToolMetadata

    def my_tool(x: int) -> int:
        return x

    def my(x: str) -> str:
        return x

    test_model = TestModel()
    agent = Agent(
        test_model,
        tools=[my_tool, my],
        capabilities=[SetToolMetadata(tools='my_tool', tagged=True)],
    )
    agent.run_sync('test')

    params = test_model.last_model_request_parameters
    assert params is not None
    td_my_tool = next(td for td in params.function_tools if td.name == 'my_tool')
    td_my = next(td for td in params.function_tools if td.name == 'my')
    assert td_my_tool.metadata is not None
    assert td_my_tool.metadata['tagged'] is True
    # 'my' should NOT match — bare string does exact match, not substring
    assert td_my.metadata is None or 'tagged' not in td_my.metadata


def test_set_tool_metadata_capability_with_sync_callable_selector():
    """SetToolMetadata with sync callable selector."""
    from pydantic_ai.capabilities import SetToolMetadata

    def tool_a(x: int) -> int:
        return x

    def tool_b(x: str) -> str:
        return x

    test_model = TestModel()
    agent = Agent(
        test_model,
        tools=[tool_a, tool_b],
        capabilities=[SetToolMetadata(tools=lambda ctx, td: td.name == 'tool_a', flagged=True)],
    )
    agent.run_sync('test')

    params = test_model.last_model_request_parameters
    assert params is not None
    td_a = next(td for td in params.function_tools if td.name == 'tool_a')
    td_b = next(td for td in params.function_tools if td.name == 'tool_b')
    assert td_a.metadata is not None
    assert td_a.metadata['flagged'] is True
    assert td_b.metadata is None or 'flagged' not in td_b.metadata


def test_set_tool_metadata_capability_with_nested_dict_selector():
    """SetToolMetadata with nested dict selector exercises deep metadata matching."""
    from pydantic_ai.capabilities import SetToolMetadata
    from pydantic_ai.tools import Tool

    def tool_a(x: int) -> int:
        return x

    def tool_b(x: str) -> str:
        return x

    def tool_c(x: float) -> float:
        return x

    test_model = TestModel()
    agent = Agent(
        test_model,
        tools=[
            Tool(tool_a, metadata={'config': {'env': 'prod', 'region': 'us'}}),
            Tool(tool_b, metadata={'config': {'env': 'staging'}}),
            Tool(tool_c, metadata={'other': 'value'}),  # missing 'config' key entirely
        ],
        capabilities=[SetToolMetadata(tools={'config': {'env': 'prod'}}, verified=True)],
    )
    agent.run_sync('test')

    params = test_model.last_model_request_parameters
    assert params is not None
    td_a = next(td for td in params.function_tools if td.name == 'tool_a')
    td_b = next(td for td in params.function_tools if td.name == 'tool_b')
    td_c = next(td for td in params.function_tools if td.name == 'tool_c')
    # tool_a matches: config.env == 'prod' (deep inclusion, extra 'region' key is fine)
    assert td_a.metadata is not None
    assert td_a.metadata['verified'] is True
    # tool_b doesn't match: config.env == 'staging'
    assert td_b.metadata is not None
    assert 'verified' not in td_b.metadata
    # tool_c doesn't match: 'config' key missing entirely
    assert td_c.metadata is not None
    assert 'verified' not in td_c.metadata


def test_set_tool_metadata_capability_with_dict_selector():
    """SetToolMetadata with dict selector matches tools by metadata."""
    from pydantic_ai.capabilities import SetToolMetadata
    from pydantic_ai.tools import Tool

    def tool_a(x: int) -> int:
        return x

    def tool_b(x: str) -> str:
        return x

    test_model = TestModel()
    agent = Agent(
        test_model,
        tools=[
            Tool(tool_a, metadata={'env': 'prod'}),
            Tool(tool_b, metadata={'env': 'staging'}),
        ],
        capabilities=[SetToolMetadata(tools={'env': 'prod'}, audited=True)],
    )
    agent.run_sync('test')

    params = test_model.last_model_request_parameters
    assert params is not None
    td_a = next(td for td in params.function_tools if td.name == 'tool_a')
    td_b = next(td for td in params.function_tools if td.name == 'tool_b')
    # tool_a matched the dict selector, gets audited=True merged
    assert td_a.metadata is not None
    assert td_a.metadata['audited'] is True
    assert td_a.metadata['env'] == 'prod'
    # tool_b didn't match
    assert td_b.metadata is not None
    assert 'audited' not in td_b.metadata


async def test_custom_toolset_returning_plain_str_instructions():
    """A custom AbstractToolset returning a plain str from get_instructions is treated as dynamic."""
    from pydantic_ai import Agent

    class PlainStrInstructionsToolset(FunctionToolset):
        """A toolset that overrides get_instructions to return a plain str instead of InstructionPart."""

        async def get_instructions(self, ctx: RunContext) -> str | None:  # type: ignore[override]
            return 'Custom toolset instruction.'

    agent = Agent(TestModel(), toolsets=[PlainStrInstructionsToolset()])
    result = await agent.run('Hello')
    first_message = result.all_messages()[0]
    assert first_message.instructions == 'Custom toolset instruction.'  # type: ignore[union-attr]


async def test_toolset_empty_instructions_filtered():
    """Empty and whitespace-only instructions from toolsets are filtered out."""
    from pydantic_ai import Agent
    from pydantic_ai.messages import InstructionPart

    class EmptyInstructionsToolset(FunctionToolset):
        async def get_instructions(self, ctx: RunContext) -> list[str | InstructionPart] | None:  # type: ignore[override]
            return [
                '',
                '   ',
                InstructionPart(content='', dynamic=True),
                InstructionPart(content='  ', dynamic=False),
                'valid instruction',
                InstructionPart(content='another valid', dynamic=True),
            ]

    agent = Agent(TestModel(), toolsets=[EmptyInstructionsToolset()])
    result = await agent.run('Hello')
    first_message = result.all_messages()[0]
    assert first_message.instructions == 'valid instruction\n\nanother valid'  # type: ignore[union-attr]


def test_apply_walks_combined_and_wrapper_toolsets():
    """`apply()` walks through `CombinedToolset` children and unwraps `WrapperToolset` (e.g. `PrefixedToolset`)."""
    inner1 = FunctionToolset()
    inner2 = FunctionToolset()
    combined = CombinedToolset([inner1, PrefixedToolset(inner2, 'p')])

    visited: list[AbstractToolset] = []
    combined.apply(visited.append)
    assert inner1 in visited
    assert inner2 in visited