pydantic--pydantic-ai/tests/test_usage_limits.py

import asyncio
import functools
import operator
import re
from collections.abc import AsyncIterator
from datetime import timezone
from decimal import Decimal

import pytest
from genai_prices import Usage as GenaiPricesUsage, calc_price
from pydantic import BaseModel

from pydantic_ai import (
    Agent,
    ModelMessage,
    ModelRequest,
    ModelResponse,
    RetryPromptPart,
    RunContext,
    TextPart,
    ToolCallPart,
    ToolReturnPart,
    UsageLimitExceeded,
    UserPromptPart,
)
from pydantic_ai.exceptions import ModelRetry
from pydantic_ai.models.function import AgentInfo, FunctionModel
from pydantic_ai.models.test import TestModel
from pydantic_ai.output import ToolOutput
from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits

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

pytestmark = pytest.mark.anyio


def test_genai_prices():
    usage = GenaiPricesUsage(input_tokens=100, output_tokens=50)
    assert calc_price(usage, model_ref='gpt-4o').total_price == snapshot(Decimal('0.00075'))


def test_request_token_limit() -> None:
    test_agent = Agent(TestModel())

    with pytest.raises(UsageLimitExceeded, match=re.escape('Exceeded the input_tokens_limit of 5 (input_tokens=59)')):
        test_agent.run_sync(
            'Hello, this prompt exceeds the request tokens limit.', usage_limits=UsageLimits(input_tokens_limit=5)
        )


def test_response_token_limit() -> None:
    test_agent = Agent(
        TestModel(custom_output_text='Unfortunately, this response exceeds the response tokens limit by a few!')
    )

    with pytest.raises(UsageLimitExceeded, match=re.escape('Exceeded the output_tokens_limit of 5 (output_tokens=11)')):
        test_agent.run_sync('Hello', usage_limits=UsageLimits(output_tokens_limit=5))


def test_total_token_limit() -> None:
    test_agent = Agent(TestModel(custom_output_text='This utilizes 4 tokens!'))

    with pytest.raises(UsageLimitExceeded, match=re.escape('Exceeded the total_tokens_limit of 50 (total_tokens=55)')):
        test_agent.run_sync('Hello', usage_limits=UsageLimits(total_tokens_limit=50))


def test_retry_limit() -> None:
    test_agent = Agent(TestModel())

    @test_agent.tool_plain
    async def foo(x: str) -> str:
        return x

    @test_agent.tool_plain
    async def bar(y: str) -> str:
        return y

    with pytest.raises(UsageLimitExceeded, match=re.escape('The next request would exceed the request_limit of 1')):
        test_agent.run_sync('Hello', usage_limits=UsageLimits(request_limit=1))


async def test_streamed_text_limits() -> None:
    m = TestModel()

    test_agent = Agent(m)
    assert test_agent.name is None

    @test_agent.tool_plain
    async def ret_a(x: str) -> str:
        return f'{x}-apple'

    succeeded = False

    with pytest.raises(
        UsageLimitExceeded, match=re.escape('Exceeded the output_tokens_limit of 10 (output_tokens=11)')
    ):
        async with test_agent.run_stream('Hello', usage_limits=UsageLimits(output_tokens_limit=10)) as result:
            assert test_agent.name == 'test_agent'
            assert not result.is_complete
            assert result.all_messages() == snapshot(
                [
                    ModelRequest(
                        parts=[UserPromptPart(content='Hello', timestamp=IsNow(tz=timezone.utc))],
                        timestamp=IsNow(tz=timezone.utc),
                        run_id=IsStr(),
                        conversation_id=IsStr(),
                    ),
                    ModelResponse(
                        parts=[
                            ToolCallPart(
                                tool_name='ret_a',
                                args={'x': 'a'},
                                tool_call_id=IsStr(),
                            )
                        ],
                        usage=RequestUsage(input_tokens=51),
                        model_name='test',
                        timestamp=IsNow(tz=timezone.utc),
                        provider_name='test',
                        run_id=IsStr(),
                        conversation_id=IsStr(),
                    ),
                    ModelRequest(
                        parts=[
                            ToolReturnPart(
                                tool_name='ret_a',
                                content='a-apple',
                                timestamp=IsNow(tz=timezone.utc),
                                tool_call_id=IsStr(),
                            )
                        ],
                        timestamp=IsNow(tz=timezone.utc),
                        run_id=IsStr(),
                        conversation_id=IsStr(),
                    ),
                ]
            )
            assert result.usage == snapshot(
                RunUsage(
                    requests=2,
                    input_tokens=103,
                    output_tokens=5,
                    tool_calls=1,
                )
            )
            succeeded = True
            async for _ in result.stream_text(debounce_by=None):
                pass

    assert succeeded


async def test_stream_text_enforces_output_token_limit_mid_stream() -> None:
    # Regression: `_stream_response_text` previously iterated `self._raw_stream_response`
    # directly, bypassing the usage-checking wrapper in `AgentStream.__aiter__`, so
    # `UsageLimitExceeded` would not raise during `stream_text()` even when the output
    # token limit was exceeded mid-stream.
    async def stream_function(_messages: list[ModelMessage], _info: AgentInfo) -> AsyncIterator[str]:
        yield 'one'
        yield 'two'
        yield 'three'

    agent = Agent(FunctionModel(stream_function=stream_function))

    collected: list[str] = []
    with pytest.raises(UsageLimitExceeded, match=re.escape('Exceeded the output_tokens_limit of 2')):
        async with agent.run_stream('hi', usage_limits=UsageLimits(output_tokens_limit=2)) as result:
            async for text in result.stream_text(delta=True, debounce_by=None):
                collected.append(text)

    assert 0 < len(collected) < 3


def test_usage_so_far() -> None:
    test_agent = Agent(TestModel())

    with pytest.raises(
        UsageLimitExceeded, match=re.escape('Exceeded the total_tokens_limit of 105 (total_tokens=163)')
    ):
        test_agent.run_sync(
            'Hello, this prompt exceeds the request tokens limit.',
            usage_limits=UsageLimits(total_tokens_limit=105),
            usage=RunUsage(input_tokens=50, output_tokens=50),
        )


async def test_multi_agent_usage_no_incr():
    delegate_agent = Agent(TestModel(), output_type=int)

    controller_agent1 = Agent(TestModel())
    run_1_usages: list[RunUsage] = []

    @controller_agent1.tool
    async def delegate_to_other_agent1(ctx: RunContext, sentence: str) -> int:
        delegate_result = await delegate_agent.run(sentence)
        delegate_usage = delegate_result.usage
        run_1_usages.append(delegate_usage)
        assert delegate_usage == snapshot(RunUsage(requests=1, input_tokens=51, output_tokens=4))
        return delegate_result.output

    result1 = await controller_agent1.run('foobar')
    assert result1.output == snapshot('{"delegate_to_other_agent1":0}')
    run_1_usages.append(result1.usage)
    assert result1.usage == snapshot(RunUsage(requests=2, input_tokens=103, output_tokens=13, tool_calls=1))
    assert result1.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[UserPromptPart(content='foobar', timestamp=IsDatetime())],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[
                    ToolCallPart(
                        tool_name='delegate_to_other_agent1',
                        args={'sentence': 'a'},
                        tool_call_id='pyd_ai_tool_call_id__delegate_to_other_agent1',
                    )
                ],
                usage=RequestUsage(input_tokens=51, output_tokens=5),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='delegate_to_other_agent1',
                        content=0,
                        tool_call_id='pyd_ai_tool_call_id__delegate_to_other_agent1',
                        timestamp=IsDatetime(),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='{"delegate_to_other_agent1":0}')],
                usage=RequestUsage(input_tokens=52, output_tokens=8),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )

    controller_agent2 = Agent(TestModel())

    @controller_agent2.tool
    async def delegate_to_other_agent2(ctx: RunContext, sentence: str) -> int:
        delegate_result = await delegate_agent.run(sentence, usage=ctx.usage)
        delegate_usage = delegate_result.usage
        assert delegate_usage == snapshot(RunUsage(requests=2, input_tokens=102, output_tokens=9))
        return delegate_result.output

    result2 = await controller_agent2.run('foobar')
    assert result2.output == snapshot('{"delegate_to_other_agent2":0}')
    assert result2.usage == snapshot(RunUsage(requests=3, input_tokens=154, output_tokens=17, tool_calls=1))
    assert result2.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[UserPromptPart(content='foobar', timestamp=IsDatetime())],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[
                    ToolCallPart(
                        tool_name='delegate_to_other_agent2',
                        args={'sentence': 'a'},
                        tool_call_id='pyd_ai_tool_call_id__delegate_to_other_agent2',
                    )
                ],
                usage=RequestUsage(input_tokens=51, output_tokens=5),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='delegate_to_other_agent2',
                        content=0,
                        tool_call_id='pyd_ai_tool_call_id__delegate_to_other_agent2',
                        timestamp=IsDatetime(),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='{"delegate_to_other_agent2":0}')],
                usage=RequestUsage(input_tokens=52, output_tokens=8),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )

    # confirm the usage from result2 is the sum of the usage from result1
    assert result2.usage == functools.reduce(operator.add, run_1_usages)

    result1_usage = result1.usage
    result1_usage.details = {'custom1': 10, 'custom2': 20, 'custom3': 0}
    assert result1_usage.opentelemetry_attributes() == {
        'gen_ai.usage.input_tokens': 103,
        'gen_ai.usage.output_tokens': 13,
        'gen_ai.usage.details.custom1': 10,
        'gen_ai.usage.details.custom2': 20,
    }


def test_opentelemetry_attributes_excludes_first_class_token_details():
    """`details` entries named like a first-class token attribute must never be emitted under `details.*`.

    Adapters stash `input_tokens`/`output_tokens` in `details` for different reasons (Anthropic's
    streaming carry-forward and pre-compaction raw counts, Cohere's billed units), but the name
    collides with the first-class `gen_ai.usage.{input,output}_tokens` attributes. Emitting the value
    under both makes consumers like Langfuse sum them and double-count tokens and cost, regardless of
    whether the two values happen to match. They stay accessible on `RequestUsage.details`; only the
    ambiguous OTel emission is dropped. Not reachable through the public API since it depends on an
    adapter leaving these keys in `details`, so pinned directly on the OTel attribute mapping.
    """
    usage = RequestUsage(
        input_tokens=100,
        output_tokens=50,
        # A matching value (Anthropic exact-copy case) and a differing one (Cohere billed-units /
        # Anthropic compaction case) are both dropped: the colliding name is what makes them ambiguous.
        details={'input_tokens': 100, 'output_tokens': 42, 'reasoning_tokens': 10},
    )
    assert usage.opentelemetry_attributes() == {
        'gen_ai.usage.input_tokens': 100,
        'gen_ai.usage.output_tokens': 50,
        'gen_ai.usage.details.reasoning_tokens': 10,
    }


async def test_multi_agent_usage_sync():
    """As in `test_multi_agent_usage_async`, with a sync tool."""
    controller_agent = Agent(TestModel())

    @controller_agent.tool
    def delegate_to_other_agent(ctx: RunContext, sentence: str) -> int:
        new_usage = RunUsage(requests=5, input_tokens=2, output_tokens=3)
        ctx.usage.incr(new_usage)
        return 0

    result = await controller_agent.run('foobar')
    assert result.output == snapshot('{"delegate_to_other_agent":0}')
    assert result.usage == snapshot(RunUsage(requests=7, input_tokens=105, output_tokens=16, tool_calls=1))
    assert result.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[UserPromptPart(content='foobar', timestamp=IsDatetime())],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[
                    ToolCallPart(
                        tool_name='delegate_to_other_agent',
                        args={'sentence': 'a'},
                        tool_call_id='pyd_ai_tool_call_id__delegate_to_other_agent',
                    )
                ],
                usage=RequestUsage(input_tokens=51, output_tokens=5),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='delegate_to_other_agent',
                        content=0,
                        tool_call_id='pyd_ai_tool_call_id__delegate_to_other_agent',
                        timestamp=IsDatetime(),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='{"delegate_to_other_agent":0}')],
                usage=RequestUsage(input_tokens=52, output_tokens=8),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )


def test_request_usage_basics():
    usage = RequestUsage()
    assert usage.output_audio_tokens == 0
    assert usage.requests == 1


def test_add_usages():
    usage = RunUsage(
        requests=2,
        input_tokens=10,
        output_tokens=20,
        output_audio_tokens=70,
        cache_read_tokens=30,
        cache_write_tokens=40,
        input_audio_tokens=50,
        cache_audio_read_tokens=60,
        tool_calls=3,
        details={
            'custom1': 10,
            'custom2': 20,
        },
    )
    assert usage + usage == snapshot(
        RunUsage(
            requests=4,
            input_tokens=20,
            output_tokens=40,
            cache_write_tokens=80,
            cache_read_tokens=60,
            input_audio_tokens=100,
            output_audio_tokens=140,
            cache_audio_read_tokens=120,
            tool_calls=6,
            details={'custom1': 20, 'custom2': 40},
        )
    )
    assert usage + RunUsage() == usage
    assert RunUsage() + RunUsage() == RunUsage()


def test_output_audio_tokens_increment():
    """Test that output_audio_tokens is correctly incremented in _incr_usage_tokens."""
    usage1 = RequestUsage(
        input_tokens=10,
        output_tokens=20,
        output_audio_tokens=15,
    )
    usage2 = RequestUsage(
        input_tokens=5,
        output_tokens=10,
        output_audio_tokens=8,
    )
    result = usage1 + usage2
    assert result.output_audio_tokens == 23
    assert result.input_tokens == 15
    assert result.output_tokens == 30

    # Also test through RunUsage.incr with RequestUsage
    run_usage = RunUsage(requests=1, output_audio_tokens=10)
    run_usage.incr(RequestUsage(output_audio_tokens=5))
    assert run_usage.output_audio_tokens == 15


def test_add_usages_with_none_detail_value():
    """Test that None values in details are skipped when incrementing usage."""
    usage = RunUsage(
        requests=1,
        input_tokens=10,
        output_tokens=20,
        details={'reasoning_tokens': 5},
    )

    # Create a usage with None in details (simulating model response with missing detail)
    incr_usage = RunUsage(
        requests=1,
        input_tokens=5,
        output_tokens=10,
    )
    # Manually set a None value in details to simulate edge case from model responses
    incr_usage.details = {'reasoning_tokens': None, 'other_tokens': 10}  # type: ignore[dict-item]

    result = usage + incr_usage
    assert result == snapshot(
        RunUsage(
            requests=2,
            input_tokens=15,
            output_tokens=30,
            details={'reasoning_tokens': 5, 'other_tokens': 10},
        )
    )


def test_add_request_usages_does_not_mutate_original():
    """Test that __add__ does not mutate the original object's details dict (issue #4605)."""
    u1 = RequestUsage(input_tokens=10, details={'reasoning_tokens': 5})
    u2 = RequestUsage(input_tokens=20, details={'reasoning_tokens': 3})

    result = u1 + u2

    # The result should have the summed details
    assert result.details == {'reasoning_tokens': 8}
    # The original must NOT be mutated
    assert u1.details == {'reasoning_tokens': 5}
    # They must be independent dict objects
    assert u1.details is not result.details


def test_add_run_usages_does_not_mutate_original():
    """Test that __add__ does not mutate the original object's details dict (issue #4605)."""
    r1 = RunUsage(requests=1, input_tokens=10, details={'reasoning_tokens': 50})
    r2 = RunUsage(requests=1, input_tokens=20, details={'reasoning_tokens': 30})

    result = r1 + r2

    assert result.details == {'reasoning_tokens': 80}
    assert r1.details == {'reasoning_tokens': 50}
    assert r1.details is not result.details


def test_add_usage_repeated_calls_stable():
    """Test that repeated __add__ calls return consistent results (issue #4605).

    This simulates `AgentStream.usage` being read multiple times:
        return self._initial_run_ctx_usage + self._raw_stream_response.usage
    """
    initial = RunUsage(requests=1, input_tokens=500, details={})
    stream = RequestUsage(input_tokens=500, output_tokens=200, details={'reasoning_tokens': 150})

    results = [initial + stream for _ in range(3)]

    # All calls must return the same values
    for r in results:
        assert r.details == {'reasoning_tokens': 150}
    # The initial usage must remain unchanged
    assert initial.details == {}


async def test_tool_call_limit() -> None:
    test_agent = Agent(TestModel())

    @test_agent.tool_plain
    async def ret_a(x: str) -> str:
        return f'{x}-apple'

    with pytest.raises(
        UsageLimitExceeded,
        match=re.escape('The next tool call(s) would exceed the tool_calls_limit of 0 (tool_calls=1).'),
    ):
        await test_agent.run('Hello', usage_limits=UsageLimits(tool_calls_limit=0))

    result = await test_agent.run('Hello', usage_limits=UsageLimits(tool_calls_limit=1))
    assert result.usage == snapshot(RunUsage(requests=2, input_tokens=103, output_tokens=14, tool_calls=1))
    assert result.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[UserPromptPart(content='Hello', timestamp=IsDatetime())],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[ToolCallPart(tool_name='ret_a', args={'x': 'a'}, tool_call_id='pyd_ai_tool_call_id__ret_a')],
                usage=RequestUsage(input_tokens=51, output_tokens=5),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='ret_a',
                        content='a-apple',
                        tool_call_id='pyd_ai_tool_call_id__ret_a',
                        timestamp=IsDatetime(),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='{"ret_a":"a-apple"}')],
                usage=RequestUsage(input_tokens=52, output_tokens=9),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )


async def test_output_tool_not_counted() -> None:
    """Test that output tools are not counted in tool_calls usage metric."""
    test_agent = Agent(TestModel())

    @test_agent.tool_plain
    async def regular_tool(x: str) -> str:
        return f'{x}-processed'

    class MyOutput(BaseModel):
        result: str

    result_regular = await test_agent.run('test')
    assert result_regular.usage == snapshot(RunUsage(requests=2, input_tokens=103, output_tokens=14, tool_calls=1))
    assert result_regular.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[UserPromptPart(content='test', timestamp=IsDatetime())],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[
                    ToolCallPart(
                        tool_name='regular_tool', args={'x': 'a'}, tool_call_id='pyd_ai_tool_call_id__regular_tool'
                    )
                ],
                usage=RequestUsage(input_tokens=51, output_tokens=5),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='regular_tool',
                        content='a-processed',
                        tool_call_id='pyd_ai_tool_call_id__regular_tool',
                        timestamp=IsDatetime(),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='{"regular_tool":"a-processed"}')],
                usage=RequestUsage(input_tokens=52, output_tokens=9),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )

    test_agent_with_output = Agent(TestModel(), output_type=ToolOutput(MyOutput))

    @test_agent_with_output.tool_plain
    async def another_regular_tool(x: str) -> str:
        return f'{x}-processed'

    result_output = await test_agent_with_output.run('test')

    assert result_output.usage == snapshot(RunUsage(requests=2, input_tokens=103, output_tokens=15, tool_calls=1))
    assert result_output.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[UserPromptPart(content='test', timestamp=IsDatetime())],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[
                    ToolCallPart(
                        tool_name='another_regular_tool',
                        args={'x': 'a'},
                        tool_call_id='pyd_ai_tool_call_id__another_regular_tool',
                    )
                ],
                usage=RequestUsage(input_tokens=51, output_tokens=5),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='another_regular_tool',
                        content='a-processed',
                        tool_call_id='pyd_ai_tool_call_id__another_regular_tool',
                        timestamp=IsDatetime(),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[
                    ToolCallPart(
                        tool_name='final_result', args={'result': 'a'}, tool_call_id='pyd_ai_tool_call_id__final_result'
                    )
                ],
                usage=RequestUsage(input_tokens=52, output_tokens=10),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='final_result',
                        content='Final result processed.',
                        tool_call_id='pyd_ai_tool_call_id__final_result',
                        timestamp=IsDatetime(),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )


async def test_output_tool_allowed_at_limit() -> None:
    """Test that output tools can be called even when at the tool_calls_limit."""

    class MyOutput(BaseModel):
        result: str

    def call_output_after_regular(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        if len(messages) == 1:
            return ModelResponse(
                parts=[
                    ToolCallPart('regular_tool', {'x': 'test'}, 'call_1'),
                ],
                usage=RequestUsage(input_tokens=10, output_tokens=5),
            )
        else:
            return ModelResponse(
                parts=[
                    ToolCallPart('final_result', {'result': 'success'}, 'call_2'),
                ],
                usage=RequestUsage(input_tokens=10, output_tokens=5),
            )

    test_agent = Agent(FunctionModel(call_output_after_regular), output_type=ToolOutput(MyOutput))

    @test_agent.tool_plain
    async def regular_tool(x: str) -> str:
        return f'{x}-processed'

    result = await test_agent.run('test', usage_limits=UsageLimits(tool_calls_limit=1))

    assert result.output.result == 'success'
    assert result.usage == snapshot(RunUsage(requests=2, input_tokens=20, output_tokens=10, tool_calls=1))
    assert result.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[UserPromptPart(content='test', timestamp=IsDatetime())],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[ToolCallPart(tool_name='regular_tool', args={'x': 'test'}, tool_call_id='call_1')],
                usage=RequestUsage(input_tokens=10, output_tokens=5),
                model_name='function:call_output_after_regular:',
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='regular_tool',
                        content='test-processed',
                        tool_call_id='call_1',
                        timestamp=IsDatetime(),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[ToolCallPart(tool_name='final_result', args={'result': 'success'}, tool_call_id='call_2')],
                usage=RequestUsage(input_tokens=10, output_tokens=5),
                model_name='function:call_output_after_regular:',
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='final_result',
                        content='Final result processed.',
                        tool_call_id='call_2',
                        timestamp=IsDatetime(),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )


async def test_failed_tool_calls_not_counted() -> None:
    """Test that failed tool calls (raising ModelRetry) are not counted in usage or against limits."""
    test_agent = Agent(TestModel())

    call_count = 0

    @test_agent.tool_plain
    async def flaky_tool(x: str) -> str:
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            raise ModelRetry('Temporary failure, please retry')
        return f'{x}-success'

    result = await test_agent.run('test', usage_limits=UsageLimits(tool_calls_limit=1))
    assert call_count == 2
    assert result.usage == snapshot(RunUsage(requests=3, input_tokens=176, output_tokens=29, tool_calls=1))
    assert result.all_messages() == snapshot(
        [
            ModelRequest(
                parts=[UserPromptPart(content='test', timestamp=IsDatetime())],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[
                    ToolCallPart(
                        tool_name='flaky_tool', args={'x': 'a'}, tool_call_id='pyd_ai_tool_call_id__flaky_tool'
                    )
                ],
                usage=RequestUsage(input_tokens=51, output_tokens=5),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    RetryPromptPart(
                        content='Temporary failure, please retry',
                        tool_name='flaky_tool',
                        tool_call_id='pyd_ai_tool_call_id__flaky_tool',
                        timestamp=IsDatetime(),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[ToolCallPart(tool_name='flaky_tool', args={'x': 'a'}, tool_call_id=IsStr())],
                usage=RequestUsage(input_tokens=62, output_tokens=10),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='flaky_tool',
                        content='a-success',
                        tool_call_id=IsStr(),
                        timestamp=IsDatetime(),
                    )
                ],
                timestamp=IsDatetime(),
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
            ModelResponse(
                parts=[TextPart(content='{"flaky_tool":"a-success"}')],
                usage=RequestUsage(input_tokens=63, output_tokens=14),
                model_name='test',
                timestamp=IsDatetime(),
                provider_name='test',
                run_id=IsStr(),
                conversation_id=IsStr(),
            ),
        ]
    )


async def test_parallel_tool_calls_limit_enforced():
    """Parallel tool calls must not exceed the limit and should raise immediately."""
    executed_tools: list[str] = []

    model_call_count = 0

    def test_model_function(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        nonlocal model_call_count
        model_call_count += 1

        if model_call_count == 1:
            # First response: 5 parallel tool calls (within limit)
            return ModelResponse(
                parts=[
                    ToolCallPart('tool_a', {}, 'call_1'),
                    ToolCallPart('tool_b', {}, 'call_2'),
                    ToolCallPart('tool_c', {}, 'call_3'),
                    ToolCallPart('tool_a', {}, 'call_4'),
                    ToolCallPart('tool_b', {}, 'call_5'),
                ]
            )
        else:
            assert model_call_count == 2
            # Second response: 3 parallel tool calls (would exceed limit of 6)
            return ModelResponse(
                parts=[
                    ToolCallPart('tool_c', {}, 'call_6'),
                    ToolCallPart('tool_a', {}, 'call_7'),
                    ToolCallPart('tool_b', {}, 'call_8'),
                ]
            )

    test_model = FunctionModel(test_model_function)
    agent = Agent(test_model)

    @agent.tool_plain
    async def tool_a() -> str:
        await asyncio.sleep(0.01)
        executed_tools.append('a')
        return 'result a'

    @agent.tool_plain
    async def tool_b() -> str:
        await asyncio.sleep(0.01)
        executed_tools.append('b')
        return 'result b'

    @agent.tool_plain
    async def tool_c() -> str:
        await asyncio.sleep(0.01)
        executed_tools.append('c')
        return 'result c'

    # Run with tool call limit of 6; expecting an error when trying to execute 3 more tools
    with pytest.raises(
        UsageLimitExceeded,
        match=re.escape('The next tool call(s) would exceed the tool_calls_limit of 6 (tool_calls=8).'),
    ):
        await agent.run('Use tools', usage_limits=UsageLimits(tool_calls_limit=6))

    # Only the first batch of 5 tools should have executed
    assert len(executed_tools) == 5


def test_usage_unknown_provider():
    assert RequestUsage.extract({}, provider='unknown', provider_url='', provider_fallback='') == RequestUsage()


def test_usage_limits_explicit_zero():
    """Explicit 0 token limits round-trip correctly (regression: zero is not coerced to None)."""
    limits = UsageLimits(input_tokens_limit=0)
    assert limits.input_tokens_limit == 0

    limits = UsageLimits(output_tokens_limit=0)
    assert limits.output_tokens_limit == 0

    limits = UsageLimits()
    assert limits.input_tokens_limit is None

    limits = UsageLimits(input_tokens_limit=100)
    assert limits.input_tokens_limit == 100