Repository: pydantic/pydantic-ai
File: tests/test_agent.py (Python)

Code before:
        from pydantic_ai.exceptions import ApprovalRequired

        raise ApprovalRequired()

    @sequential_toolset.tool
    def call_second():
        nonlocal integer_holder
        assert integer_holder == 2

    agent = Agent(
        FunctionModel(call_tools_sequential), toolsets=[sequential_toolset], output_type=[str, DeferredToolRequests]
    )

    user_prompt = 'call a lot of tools'

    if mode == 'contextmanager':
        with agent.sequential_tool_calls():
            result = agent.run_sync(user_prompt)
    else:
        result = agent.run_sync(user_prompt)

    assert isinstance(result.output, DeferredToolRequests)
    assert len(result.output.approvals) == 1
    assert result.output.approvals[0].tool_name == 'requires_approval'
    # When no metadata is provided, the tool_call_id should not be in metadata dict
    tool_call_id = result.output.approvals[0].tool_call_id
    assert tool_call_id not in result.output.metadata
    assert result.output.metadata == {}
    assert integer_holder == 2


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

    test_model = TestModel()
    server1 = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
    server2 = MCPServerStdio('python', ['-m', 'tests.mcp_server'], sampling_model=test_model)
    toolset = CombinedToolset([server1, PrefixedToolset(server2, 'prefix')])
    agent = Agent(None, toolsets=[toolset])

    with pytest.raises(UserError, match='No sampling model provided and no model set on the agent.'):
        agent.set_mcp_sampling_model()
    assert server1.sampling_model is None
    assert server2.sampling_model is test_model

    agent.model = test_model
    agent.set_mcp_sampling_model()
    assert server1.sampling_model is test_model

Review comment (suggestion):
I think we can revert this to the snapshot -- the `metadata` type doesn't allow `tool_call_id: None` anyway

Code after:
                        tool_name='foo',
                        content='Hello from foo',
                        tool_call_id=IsStr(),
                        timestamp=IsDatetime(),
                    )
                ],
                run_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='Done')],
                usage=RequestUsage(input_tokens=63, output_tokens=5),
                model_name='function:respond:',
                timestamp=IsDatetime(),
                run_id=IsStr(),
            ),
        ]
    )


def test_prepare_output_tools():
    @dataclass
    class AgentDeps:
        plan_presented: bool = False

    async def present_plan(ctx: RunContext[AgentDeps], plan: str) -> str:
        """
        Present the plan to the user.
        """
        ctx.deps.plan_presented = True
        return plan

    async def run_sql(ctx: RunContext[AgentDeps], purpose: str, query: str) -> str:
        """
        Run an SQL query.
        """
        return 'SQL query executed successfully'

    async def only_if_plan_presented(
        ctx: RunContext[AgentDeps], tool_defs: list[ToolDefinition]
    ) -> list[ToolDefinition]:
        return tool_defs if ctx.deps.plan_presented else []

    agent = Agent(
        model='test',
        deps_type=AgentDeps,
        tools=[present_plan],
        output_type=[ToolOutput(run_sql, name='run_sql')],
        prepare_output_tools=only_if_plan_presented,
    )

    result = agent.run_sync('Hello', deps=AgentDeps())