pydantic--pydantic-ai/tests/test_logfire.py

from __future__ import annotations as _annotations

from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Literal

import pytest
from dirty_equals import IsJson, IsList
from pydantic import BaseModel
from typing_extensions import NotRequired, Self, TypedDict

from pydantic_ai import Agent, ModelMessage, ModelRequest, ModelResponse, TextPart, ToolCallPart, UserPromptPart
from pydantic_ai._utils import get_traceparent
from pydantic_ai._warnings import PydanticAIDeprecationWarning
from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.capabilities.instrumentation import Instrumentation
from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UnexpectedModelBehavior
from pydantic_ai.models.function import AgentInfo, FunctionModel
from pydantic_ai.models.instrumented import InstrumentationSettings
from pydantic_ai.models.test import TestModel
from pydantic_ai.output import PromptedOutput, TextOutput
from pydantic_ai.tools import DeferredToolRequests, RunContext
from pydantic_ai.toolsets.abstract import ToolsetTool
from pydantic_ai.toolsets.function import FunctionToolset
from pydantic_ai.toolsets.wrapper import WrapperToolset
from pydantic_ai.usage import RequestUsage

from ._inline_snapshot import snapshot
from .conftest import IsDatetime, IsInt, IsStr, strip_logfire_metrics

try:
    import logfire
    from logfire.testing import CaptureLogfire
except ImportError:  # pragma: lax no cover
    logfire_installed = False
else:
    logfire_installed = True


class SpanSummary(TypedDict):
    id: int
    name: str
    message: str
    children: NotRequired[list[SpanSummary]]


@dataclass(init=False)
class LogfireSummary:
    traces: list[SpanSummary]
    attributes: dict[int, dict[str, Any]]

    def __init__(self, capfire: CaptureLogfire):
        spans = capfire.exporter.exported_spans_as_dict()
        spans.sort(key=lambda s: s['start_time'])
        self.traces = []
        span_lookup: dict[tuple[str, str], SpanSummary] = {}
        self.attributes = {}
        id_counter = 0
        for span in spans:
            tid = span['context']['trace_id'], span['context']['span_id']
            span_lookup[tid] = span_summary = SpanSummary(
                id=id_counter, name=span['name'], message=span['attributes']['logfire.msg']
            )
            # `logfire.metrics` is a logfire-version-dependent span decoration (added in 4.3x): newer
            # logfire attaches the aggregated `gen_ai.client.token.usage` metric to spans, older does not.
            # Strip it so these assertions hold across the supported logfire range; the token usage itself
            # is still covered by the stable `gen_ai.usage.*` attributes and `get_collected_metrics()`.
            self.attributes[id_counter] = {k: v for k, v in span['attributes'].items() if k != 'logfire.metrics'}
            id_counter += 1
            if parent := span['parent']:
                parent_span = span_lookup[(parent['trace_id'], parent['span_id'])]
                parent_span.setdefault('children', []).append(span_summary)
            else:
                self.traces.append(span_summary)


@pytest.fixture
def get_logfire_summary(capfire: CaptureLogfire) -> Callable[[], LogfireSummary]:
    def get_summary() -> LogfireSummary:
        return LogfireSummary(capfire)

    return get_summary


def deprecated_instrumentation_settings(version: Literal[2, 3, 4], **kwargs: Any) -> InstrumentationSettings:
    with pytest.warns(
        PydanticAIDeprecationWarning,
        match=r'Instrumentation format versions 2, 3, and 4 are deprecated',
    ):
        return InstrumentationSettings(version=version, **kwargs)


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize(
    'instrument',
    [
        True,
        False,
        deprecated_instrumentation_settings(version=2),
        deprecated_instrumentation_settings(version=3),
    ],
)
def test_logfire(
    get_logfire_summary: Callable[[], LogfireSummary],
    instrument: InstrumentationSettings | bool,
    capfire: CaptureLogfire,
) -> None:
    class InstrumentedToolset(WrapperToolset):
        async def __aenter__(self) -> Self:
            with logfire.span('toolset_enter'):  # pyright: ignore[reportPossiblyUnboundVariable]
                await super().__aenter__()
                return self

        async def __aexit__(self, *args: Any) -> bool | None:
            with logfire.span('toolset_exit'):  # pyright: ignore[reportPossiblyUnboundVariable]
                return await super().__aexit__(*args)

        async def call_tool(
            self, name: str, tool_args: dict[str, Any], ctx: RunContext[Any], tool: ToolsetTool[Any]
        ) -> Any:
            with logfire.span('toolset_call_tool {name}', name=name):  # pyright: ignore[reportPossiblyUnboundVariable]
                return await super().call_tool(name, tool_args, ctx, tool)

    toolset = FunctionToolset()

    @toolset.tool_plain
    async def my_ret(x: int) -> str:
        return str(x + 1)

    if instrument:
        toolset = InstrumentedToolset(toolset)

    capabilities = (
        [
            Instrumentation(
                settings=instrument if isinstance(instrument, InstrumentationSettings) else InstrumentationSettings()
            )
        ]
        if instrument
        else []
    )
    my_agent = Agent(
        model=TestModel(),
        toolsets=[toolset],
        capabilities=capabilities,
        metadata={'env': 'test'},
    )

    result = my_agent.run_sync('Hello')
    assert result.output == snapshot('{"my_ret":"1"}')

    summary = get_logfire_summary()
    if instrument is False:
        assert summary.traces == []
        return

    if instrument is True or (isinstance(instrument, InstrumentationSettings) and instrument.version >= 3):
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'invoke_agent my_agent',
                    'message': 'my_agent run',
                    'children': [
                        {'id': 1, 'name': 'toolset_enter', 'message': 'toolset_enter'},
                        {'id': 2, 'name': 'chat test', 'message': 'chat test'},
                        {
                            'id': 3,
                            'name': 'execute_tool my_ret',
                            'message': 'running tool: my_ret',
                            'children': [
                                {'id': 4, 'name': 'toolset_call_tool {name}', 'message': 'toolset_call_tool my_ret'}
                            ],
                        },
                        {'id': 5, 'name': 'chat test', 'message': 'chat test'},
                        {'id': 6, 'name': 'toolset_exit', 'message': 'toolset_exit'},
                    ],
                }
            ]
        )
    else:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'agent run',
                    'message': 'my_agent run',
                    'children': [
                        {'id': 1, 'name': 'toolset_enter', 'message': 'toolset_enter'},
                        {'id': 2, 'name': 'chat test', 'message': 'chat test'},
                        {
                            'id': 3,
                            'name': 'running tool',
                            'message': 'running tool: my_ret',
                            'children': [
                                {'id': 4, 'name': 'toolset_call_tool {name}', 'message': 'toolset_call_tool my_ret'}
                            ],
                        },
                        {'id': 5, 'name': 'chat test', 'message': 'chat test'},
                        {'id': 6, 'name': 'toolset_exit', 'message': 'toolset_exit'},
                    ],
                }
            ]
        )

    if isinstance(instrument, InstrumentationSettings) and instrument.version == 2:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'agent run',
                    'message': 'my_agent run',
                    'children': [
                        {'id': 1, 'name': 'toolset_enter', 'message': 'toolset_enter'},
                        {'id': 2, 'name': 'chat test', 'message': 'chat test'},
                        {
                            'id': 3,
                            'name': 'running tool',
                            'message': 'running tool: my_ret',
                            'children': [
                                {'id': 4, 'name': 'toolset_call_tool {name}', 'message': 'toolset_call_tool my_ret'}
                            ],
                        },
                        {'id': 5, 'name': 'chat test', 'message': 'chat test'},
                        {'id': 6, 'name': 'toolset_exit', 'message': 'toolset_exit'},
                    ],
                }
            ]
        )
    else:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'invoke_agent my_agent',
                    'message': 'my_agent run',
                    'children': [
                        {'id': 1, 'name': 'toolset_enter', 'message': 'toolset_enter'},
                        {'id': 2, 'name': 'chat test', 'message': 'chat test'},
                        {
                            'id': 3,
                            'name': 'execute_tool my_ret',
                            'message': 'running tool: my_ret',
                            'children': [
                                {'id': 4, 'name': 'toolset_call_tool {name}', 'message': 'toolset_call_tool my_ret'}
                            ],
                        },
                        {'id': 5, 'name': 'chat test', 'message': 'chat test'},
                        {'id': 6, 'name': 'toolset_exit', 'message': 'toolset_exit'},
                    ],
                }
            ]
        )

    assert summary.attributes[0] == snapshot(
        {
            'model_name': 'test',
            'agent_name': 'my_agent',
            'gen_ai.agent.name': 'my_agent',
            'gen_ai.agent.call.id': IsStr(),
            'gen_ai.conversation.id': IsStr(),
            'gen_ai.operation.name': 'invoke_agent',
            'logfire.msg': 'my_agent run',
            'logfire.span_type': 'span',
            'final_result': '{"my_ret":"1"}',
            'gen_ai.aggregated_usage.input_tokens': 103,
            'gen_ai.aggregated_usage.output_tokens': 12,
            'pydantic_ai.all_messages': IsJson(
                snapshot(
                    [
                        {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]},
                        {
                            'role': 'assistant',
                            'parts': [
                                {
                                    'type': 'tool_call',
                                    'id': IsStr(),
                                    'name': 'my_ret',
                                    'arguments': {'x': 0},
                                }
                            ],
                        },
                        {
                            'role': 'user',
                            'parts': [
                                {
                                    'type': 'tool_call_response',
                                    'id': IsStr(),
                                    'name': 'my_ret',
                                    'result': '1',
                                }
                            ],
                        },
                        {'role': 'assistant', 'parts': [{'type': 'text', 'content': '{"my_ret":"1"}'}]},
                    ]
                )
            ),
            'metadata': '{"env":"test"}',
            'logfire.json_schema': IsJson(
                snapshot(
                    {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'metadata': {'type': 'array'},
                            'final_result': {'type': 'object'},
                        },
                    }
                )
            ),
        }
    )
    chat_span_attributes = next(
        attrs for attrs in summary.attributes.values() if attrs.get('gen_ai.operation.name', None) == 'chat'
    )
    if hasattr(capfire, 'get_collected_metrics'):  # pragma: no branch
        assert capfire.get_collected_metrics() == snapshot(
            [
                {
                    'name': 'gen_ai.client.token.usage',
                    'description': 'Measures number of input and output tokens used',
                    'unit': '{token}',
                    'data': {
                        'data_points': [
                            {
                                'attributes': {
                                    'gen_ai.provider.name': 'test',
                                    'gen_ai.system': 'test',
                                    'gen_ai.operation.name': 'chat',
                                    'gen_ai.request.model': 'test',
                                    'gen_ai.response.model': 'test',
                                    'gen_ai.token.type': 'input',
                                },
                                'start_time_unix_nano': IsInt(),
                                'time_unix_nano': IsInt(),
                                'count': 2,
                                'sum': 103,
                                'scale': 12,
                                'zero_count': 0,
                                'positive': {
                                    'offset': 23234,
                                    'bucket_counts': IsList(length=...),
                                },
                                'negative': {'offset': 0, 'bucket_counts': [0]},
                                'flags': 0,
                                'min': 51,
                                'max': 52,
                                'exemplars': IsList(length=...),
                            },
                            {
                                'attributes': {
                                    'gen_ai.provider.name': 'test',
                                    'gen_ai.system': 'test',
                                    'gen_ai.operation.name': 'chat',
                                    'gen_ai.request.model': 'test',
                                    'gen_ai.response.model': 'test',
                                    'gen_ai.token.type': 'output',
                                },
                                'start_time_unix_nano': IsInt(),
                                'time_unix_nano': IsInt(),
                                'count': 2,
                                'sum': 12,
                                'scale': 7,
                                'zero_count': 0,
                                'positive': {
                                    'offset': 255,
                                    'bucket_counts': IsList(length=...),
                                },
                                'negative': {'offset': 0, 'bucket_counts': [0]},
                                'flags': 0,
                                'min': 4,
                                'max': 8,
                                'exemplars': IsList(length=...),
                            },
                        ],
                        'aggregation_temporality': 1,
                    },
                }
            ]
        )

    messages_attributes = {
        k: chat_span_attributes.pop(k)
        for k in ['events', 'gen_ai.input.messages', 'gen_ai.output.messages']
        if k in chat_span_attributes
    }
    assert messages_attributes == snapshot(
        {
            'gen_ai.input.messages': IsJson(
                snapshot([{'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]}])
            ),
            'gen_ai.output.messages': IsJson(
                snapshot(
                    [
                        {
                            'role': 'assistant',
                            'parts': [
                                {
                                    'type': 'tool_call',
                                    'id': IsStr(),
                                    'name': 'my_ret',
                                    'arguments': {'x': 0},
                                }
                            ],
                        }
                    ]
                )
            ),
        }
    )

    assert chat_span_attributes == snapshot(
        {
            'gen_ai.operation.name': 'chat',
            'gen_ai.provider.name': 'test',
            'gen_ai.system': 'test',
            'gen_ai.request.model': 'test',
            'model_request_parameters': IsJson(
                snapshot(
                    {
                        'function_tools': [
                            {
                                'name': 'my_ret',
                                'description': None,
                                'parameters_json_schema': {
                                    'additionalProperties': False,
                                    'properties': {'x': {'type': 'integer'}},
                                    'required': ['x'],
                                    'type': 'object',
                                },
                                'outer_typed_dict_key': None,
                                'strict': None,
                                'sequential': False,
                                'kind': 'function',
                                'metadata': None,
                                'timeout': None,
                                'defer_loading': False,
                                'unless_native': None,
                                'with_native': None,
                                'tool_kind': None,
                                'return_schema': None,
                                'include_return_schema': None,
                                'capability_id': None,
                            }
                        ],
                        'native_tools': [],
                        'output_mode': 'text',
                        'output_tools': [],
                        'output_object': None,
                        'prompted_output_template': None,
                        'allow_text_output': True,
                        'allow_image_output': False,
                        'instruction_parts': None,
                        'thinking': None,
                    }
                )
            ),
            'gen_ai.tool.definitions': '[{"type":"function","name":"my_ret","parameters":{"additionalProperties":false,"properties":{"x":{"type":"integer"}},"required":["x"],"type":"object"}}]',
            'logfire.json_schema': IsJson(),
            'gen_ai.conversation.id': IsStr(),
            'logfire.span_type': 'span',
            'gen_ai.agent.name': 'my_agent',
            'gen_ai.agent.call.id': IsStr(),
            'logfire.msg': 'chat test',
            'gen_ai.response.model': 'test',
            'gen_ai.usage.input_tokens': 51,
            'gen_ai.usage.output_tokens': 4,
        }
    )


def _test_logfire_metadata_values_callable_dict(ctx: RunContext[Any]) -> dict[str, str]:
    return {'model_name': ctx.model.model_name}


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize(
    ('metadata', 'expected'),
    [
        pytest.param({'env': 'test'}, '{"env":"test"}', id='dict'),
        pytest.param(_test_logfire_metadata_values_callable_dict, '{"model_name":"test"}', id='callable-dict'),
    ],
)
def test_logfire_metadata_values(
    get_logfire_summary: Callable[[], LogfireSummary],
    metadata: dict[str, Any] | Callable[[RunContext[Any]], dict[str, Any]],
    expected: dict[str, Any],
) -> None:
    agent = Agent(
        model=TestModel(),
        capabilities=[Instrumentation(settings=deprecated_instrumentation_settings(version=2))],
        metadata=metadata,
    )
    agent.run_sync('Hello')

    summary = get_logfire_summary()
    assert summary.attributes[0]['metadata'] == expected


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_logfire_metadata_override(get_logfire_summary: Callable[[], LogfireSummary]) -> None:
    agent = Agent(
        model=TestModel(),
        capabilities=[Instrumentation(settings=deprecated_instrumentation_settings(version=2))],
        metadata={'env': 'base'},
    )
    with agent.override(metadata={'env': 'override'}):
        agent.run_sync('Hello')

    summary = get_logfire_summary()
    assert summary.attributes[0]['metadata'] == '{"env":"override"}'


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.anyio
async def test_logfire_streaming_records_time_to_first_chunk(capfire: CaptureLogfire) -> None:
    """A streaming agent run records `gen_ai.client.operation.time_to_first_chunk` on the
    model-request span and as a histogram metric (value is non-deterministic, so assert shape)."""
    agent = Agent(
        model=TestModel(),
        capabilities=[Instrumentation(settings=InstrumentationSettings())],
    )
    async with agent.run_stream('Hello') as result:
        async for _ in result.stream_text(delta=True):
            pass

    chat_spans = [
        s for s in capfire.exporter.exported_spans_as_dict() if s['attributes'].get('gen_ai.operation.name') == 'chat'
    ]
    assert chat_spans
    for span in chat_spans:
        ttft = span['attributes'].get('gen_ai.client.operation.time_to_first_chunk')
        assert isinstance(ttft, float)

    # Pin the histogram emission through the agent-flow path (capability handler -> req_ctx ->
    # finish), not just the metric name, so a regression that drops the value between the handler
    # and `finish` can't slip through.
    ttft_metrics = [
        m for m in capfire.get_collected_metrics() if m['name'] == 'gen_ai.client.operation.time_to_first_chunk'
    ]
    assert len(ttft_metrics) == 1
    assert ttft_metrics[0]['unit'] == 's'
    assert len(ttft_metrics[0]['data']['data_points']) == 1


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize(
    'instrument',
    [deprecated_instrumentation_settings(version=2), deprecated_instrumentation_settings(version=3)],
)
def test_instructions_with_structured_output(
    get_logfire_summary: Callable[[], LogfireSummary], instrument: InstrumentationSettings
) -> None:
    @dataclass
    class MyOutput:
        content: str

    my_agent = Agent(
        model=TestModel(),
        instructions='Here are some instructions',
        capabilities=[Instrumentation(settings=instrument)],
    )

    result = my_agent.run_sync('Hello', output_type=MyOutput)
    assert result.output == MyOutput(content='a')

    summary = get_logfire_summary()
    chat_span_attributes = summary.attributes[1]
    if instrument.version == 2:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'agent run',
                    'message': 'my_agent run',
                    'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
                }
            ]
        )
    else:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'invoke_agent my_agent',
                    'message': 'my_agent run',
                    'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
                }
            ]
        )

    assert summary.attributes[0] == snapshot(
        {
            'model_name': 'test',
            'agent_name': 'my_agent',
            'gen_ai.agent.name': 'my_agent',
            'gen_ai.agent.call.id': IsStr(),
            'gen_ai.conversation.id': IsStr(),
            'gen_ai.operation.name': 'invoke_agent',
            'logfire.msg': 'my_agent run',
            'logfire.span_type': 'span',
            'final_result': '{"content":"a"}',
            'gen_ai.aggregated_usage.input_tokens': 51,
            'gen_ai.aggregated_usage.output_tokens': 5,
            'pydantic_ai.all_messages': IsJson(
                snapshot(
                    [
                        {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]},
                        {
                            'role': 'assistant',
                            'parts': [
                                {
                                    'type': 'tool_call',
                                    'id': IsStr(),
                                    'name': 'final_result',
                                    'arguments': {'content': 'a'},
                                }
                            ],
                        },
                        {
                            'role': 'user',
                            'parts': [
                                {
                                    'type': 'tool_call_response',
                                    'id': IsStr(),
                                    'name': 'final_result',
                                    'result': 'Final result processed.',
                                }
                            ],
                        },
                    ]
                )
            ),
            'gen_ai.system_instructions': '[{"type":"text","content":"Here are some instructions"}]',
            'logfire.json_schema': IsJson(
                snapshot(
                    {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'gen_ai.system_instructions': {'type': 'array'},
                            'final_result': {'type': 'object'},
                        },
                    }
                )
            ),
        }
    )

    assert chat_span_attributes['gen_ai.input.messages'] == IsJson(
        snapshot([{'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]}])
    )
    assert chat_span_attributes['gen_ai.output.messages'] == IsJson(
        snapshot(
            [
                {
                    'role': 'assistant',
                    'parts': [
                        {
                            'type': 'tool_call',
                            'id': IsStr(),
                            'name': 'final_result',
                            'arguments': {'content': 'a'},
                        }
                    ],
                }
            ]
        )
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_instructions_with_structured_output_exclude_content(get_logfire_summary: Callable[[], LogfireSummary]) -> None:
    @dataclass
    class MyOutput:
        content: str

    settings: InstrumentationSettings = InstrumentationSettings(include_content=False)

    my_agent = Agent(
        model=TestModel(), instructions='Here are some instructions', capabilities=[Instrumentation(settings=settings)]
    )

    result = my_agent.run_sync('Hello', output_type=MyOutput)
    assert result.output == snapshot(MyOutput(content='a'))

    summary = get_logfire_summary()
    assert summary.attributes[0] == snapshot(
        {
            'model_name': 'test',
            'agent_name': 'my_agent',
            'gen_ai.agent.name': 'my_agent',
            'gen_ai.agent.call.id': IsStr(),
            'gen_ai.conversation.id': IsStr(),
            'gen_ai.operation.name': 'invoke_agent',
            'logfire.msg': 'my_agent run',
            'logfire.span_type': 'span',
            'gen_ai.aggregated_usage.input_tokens': 51,
            'gen_ai.aggregated_usage.output_tokens': 5,
            'pydantic_ai.all_messages': IsJson(
                snapshot(
                    [
                        {'role': 'user', 'parts': [{'type': 'text'}]},
                        {
                            'role': 'assistant',
                            'parts': [
                                {
                                    'type': 'tool_call',
                                    'id': IsStr(),
                                    'name': 'final_result',
                                }
                            ],
                        },
                        {
                            'role': 'user',
                            'parts': [
                                {
                                    'type': 'tool_call_response',
                                    'id': IsStr(),
                                    'name': 'final_result',
                                }
                            ],
                        },
                    ]
                )
            ),
            'logfire.json_schema': IsJson(
                snapshot(
                    {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'final_result': {'type': 'object'},
                        },
                    }
                )
            ),
        }
    )
    chat_span_attributes = summary.attributes[1]
    assert chat_span_attributes['gen_ai.input.messages'] == IsJson(
        snapshot(
            [
                {'role': 'user', 'parts': [{'type': 'text'}]},
            ]
        )
    )
    assert chat_span_attributes['gen_ai.output.messages'] == IsJson(
        snapshot(
            [
                {
                    'role': 'assistant',
                    'parts': [
                        {
                            'type': 'tool_call',
                            'id': IsStr(),
                            'name': 'final_result',
                        }
                    ],
                }
            ]
        )
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize(
    'settings',
    [deprecated_instrumentation_settings(version=2), deprecated_instrumentation_settings(version=3)],
)
def test_prompted_output_schema_instructions_do_not_set_variable_instructions(
    get_logfire_summary: Callable[[], LogfireSummary],
    settings: InstrumentationSettings,
) -> None:
    class City(BaseModel):
        name: str
        population: int

    my_agent = Agent(
        model=TestModel(custom_output_text='{"name":"Paris","population":2148000}'),
        instructions='Be helpful',
        capabilities=[Instrumentation(settings=settings)],
        output_type=PromptedOutput(City, template='Return JSON matching this schema:\n{schema}'),
    )

    result = my_agent.run_sync('Tell me about Paris')
    assert result.output == snapshot(City(name='Paris', population=2148000))

    summary = get_logfire_summary()
    agent_run_attrs = summary.attributes[0]
    assert 'Be helpful' in agent_run_attrs['gen_ai.system_instructions']
    assert 'pydantic_ai.variable_instructions' not in agent_run_attrs


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize(
    'settings',
    [deprecated_instrumentation_settings(version=2), deprecated_instrumentation_settings(version=3)],
)
def test_stable_instructions_across_tool_calls_do_not_set_variable_instructions(
    get_logfire_summary: Callable[[], LogfireSummary],
    settings: InstrumentationSettings,
) -> None:
    @dataclass
    class MyOutput:
        content: str

    my_agent = Agent(
        model=TestModel(),
        instructions='Be helpful',
        capabilities=[Instrumentation(settings=settings)],
    )
    instruction_calls = 0

    @my_agent.tool_plain
    def my_tool() -> str:
        nonlocal instruction_calls
        instruction_calls += 1
        return 'tool result'

    result = my_agent.run_sync('Hello', output_type=MyOutput)
    assert result.output == MyOutput(content='a')
    # Ensure multi-step execution occurred so instructions were compared across requests
    assert instruction_calls >= 1

    summary = get_logfire_summary()
    assert 'pydantic_ai.variable_instructions' not in summary.attributes[0]


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize(
    'settings',
    [deprecated_instrumentation_settings(version=2), deprecated_instrumentation_settings(version=3)],
)
def test_dynamic_instructions_toggling_from_none_sets_variable_instructions(
    get_logfire_summary: Callable[[], LogfireSummary],
    settings: InstrumentationSettings,
) -> None:
    @dataclass
    class MyOutput:
        content: str

    my_agent = Agent(
        model=TestModel(),
        capabilities=[Instrumentation(settings=settings)],
    )
    instruction_calls = 0

    @my_agent.instructions
    def instructions(_: RunContext[object]) -> str | None:
        nonlocal instruction_calls
        instruction_calls += 1
        if instruction_calls == 1:
            return None
        return 'This is a later step'

    @my_agent.tool_plain
    def my_tool() -> str:
        return 'This is a tool call'

    result = my_agent.run_sync('Hello', output_type=MyOutput)
    assert result.output == MyOutput(content='a')
    # Ensure multi-step execution occurred so instructions actually toggled
    assert instruction_calls >= 2

    summary = get_logfire_summary()
    assert summary.attributes[0]['pydantic_ai.variable_instructions'] is True


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize('version', [2, 3])
def test_instructions_with_structured_output_exclude_content_v2_v3(
    get_logfire_summary: Callable[[], LogfireSummary],
    version: Literal[2, 3],
) -> None:
    @dataclass
    class MyOutput:
        content: str

    settings: InstrumentationSettings = deprecated_instrumentation_settings(include_content=False, version=version)

    my_agent = Agent(
        model=TestModel(), instructions='Here are some instructions', capabilities=[Instrumentation(settings=settings)]
    )

    result = my_agent.run_sync('Hello', output_type=MyOutput)
    assert result.output == MyOutput(content='a')

    summary = get_logfire_summary()

    if version == 2:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'agent run',
                    'message': 'my_agent run',
                    'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
                }
            ]
        )
    else:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'invoke_agent my_agent',
                    'message': 'my_agent run',
                    'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
                }
            ]
        )

    # Version 2 and 3 have identical snapshots for this test case
    assert summary.attributes[0] == snapshot(
        {
            'model_name': 'test',
            'agent_name': 'my_agent',
            'gen_ai.agent.name': 'my_agent',
            'gen_ai.agent.call.id': IsStr(),
            'gen_ai.conversation.id': IsStr(),
            'gen_ai.operation.name': 'invoke_agent',
            'logfire.msg': 'my_agent run',
            'logfire.span_type': 'span',
            'gen_ai.aggregated_usage.input_tokens': 51,
            'gen_ai.aggregated_usage.output_tokens': 5,
            'pydantic_ai.all_messages': IsJson(
                snapshot(
                    [
                        {'role': 'user', 'parts': [{'type': 'text'}]},
                        {
                            'role': 'assistant',
                            'parts': [
                                {
                                    'type': 'tool_call',
                                    'id': IsStr(),
                                    'name': 'final_result',
                                }
                            ],
                        },
                        {
                            'role': 'user',
                            'parts': [
                                {
                                    'type': 'tool_call_response',
                                    'id': IsStr(),
                                    'name': 'final_result',
                                }
                            ],
                        },
                    ]
                )
            ),
            'logfire.json_schema': IsJson(
                snapshot(
                    {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'final_result': {'type': 'object'},
                        },
                    }
                )
            ),
        }
    )
    chat_span_attributes = summary.attributes[1]
    assert chat_span_attributes == snapshot(
        {
            'gen_ai.operation.name': 'chat',
            'gen_ai.provider.name': 'test',
            'gen_ai.system': 'test',
            'gen_ai.request.model': 'test',
            'model_request_parameters': IsJson(
                snapshot(
                    {
                        'function_tools': [],
                        'native_tools': [],
                        'output_mode': 'tool',
                        'output_object': None,
                        'output_tools': [
                            {
                                'name': 'final_result',
                                'parameters_json_schema': {
                                    'properties': {'content': {'type': 'string'}},
                                    'required': ['content'],
                                    'title': 'MyOutput',
                                    'type': 'object',
                                },
                                'description': 'The final response which ends this conversation',
                                'outer_typed_dict_key': None,
                                'strict': None,
                                'sequential': False,
                                'kind': 'output',
                                'metadata': None,
                                'timeout': None,
                                'defer_loading': False,
                                'unless_native': None,
                                'with_native': None,
                                'tool_kind': None,
                                'return_schema': None,
                                'include_return_schema': None,
                                'capability_id': None,
                            }
                        ],
                        'prompted_output_template': None,
                        'allow_text_output': False,
                        'allow_image_output': False,
                        'instruction_parts': [
                            {'content': 'Here are some instructions', 'dynamic': False, 'part_kind': 'instruction'}
                        ],
                        'thinking': None,
                    }
                )
            ),
            'gen_ai.tool.definitions': '[{"type":"function","name":"final_result","description":"The final response which ends this conversation","parameters":{"properties":{"content":{"type":"string"}},"required":["content"],"title":"MyOutput","type":"object"}}]',
            'logfire.span_type': 'span',
            'gen_ai.conversation.id': IsStr(),
            'logfire.msg': 'chat test',
            'gen_ai.agent.name': 'my_agent',
            'gen_ai.agent.call.id': IsStr(),
            'gen_ai.input.messages': IsJson(snapshot([{'role': 'user', 'parts': [{'type': 'text'}]}])),
            'gen_ai.output.messages': IsJson(
                snapshot(
                    [
                        {
                            'role': 'assistant',
                            'parts': [
                                {
                                    'type': 'tool_call',
                                    'id': IsStr(),
                                    'name': 'final_result',
                                }
                            ],
                        }
                    ]
                )
            ),
            'logfire.json_schema': IsJson(
                snapshot(
                    {
                        'type': 'object',
                        'properties': {
                            'gen_ai.input.messages': {'type': 'array'},
                            'gen_ai.output.messages': {'type': 'array'},
                            'model_request_parameters': {'type': 'object'},
                        },
                    }
                )
            ),
            'gen_ai.usage.input_tokens': 51,
            'gen_ai.usage.output_tokens': 5,
            'gen_ai.response.model': 'test',
        }
    )


def test_instrument_all():
    agent = Agent()

    def resolve():
        return agent._resolve_instrumentation_settings()  # type: ignore[reportPrivateUsage]

    Agent.instrument_all(False)
    assert resolve() is None

    Agent.instrument_all()
    settings = resolve()
    assert settings is not None
    assert settings.version == InstrumentationSettings().version

    options = InstrumentationSettings(version=5)
    Agent.instrument_all(options)
    assert resolve() is options

    Agent.instrument_all(False)
    assert resolve() is None


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.anyio
async def test_aggregated_usage_attribute_names_default(capfire: CaptureLogfire) -> None:
    """Agent run spans use aggregated usage attribute names by default."""

    def model_function(messages: list[ModelRequest | ModelResponse], info: AgentInfo) -> ModelResponse:
        # Return a response with usage that includes extra details (cache tokens)
        # to test that all gen_ai.usage.* attributes are translated
        return ModelResponse(
            parts=[TextPart('Hello!')],
            usage=RequestUsage(input_tokens=10, output_tokens=5, cache_read_tokens=2),
        )

    settings = InstrumentationSettings()
    agent = Agent(model=FunctionModel(model_function), capabilities=[Instrumentation(settings=settings)])

    await agent.run('Hello')

    spans = strip_logfire_metrics(capfire.exporter.exported_spans_as_dict(parse_json_attributes=True))

    # Verify that agent run span uses aggregated_usage attribute names
    agent_run_span = next(s for s in spans if s['name'] == 'invoke_agent agent')
    assert agent_run_span['attributes'] == snapshot(
        {
            'model_name': 'function:model_function:',
            'agent_name': 'agent',
            'gen_ai.agent.name': 'agent',
            'gen_ai.agent.call.id': IsStr(),
            'gen_ai.conversation.id': IsStr(),
            'gen_ai.operation.name': 'invoke_agent',
            'logfire.msg': 'agent run',
            'logfire.span_type': 'span',
            'final_result': 'Hello!',
            'gen_ai.aggregated_usage.input_tokens': 10,
            'gen_ai.aggregated_usage.output_tokens': 5,
            'gen_ai.aggregated_usage.cache_read.input_tokens': 2,
            'gen_ai.aggregated_usage.details.cache_read_tokens': 2,
            'pydantic_ai.all_messages': [
                {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]},
                {'role': 'assistant', 'parts': [{'type': 'text', 'content': 'Hello!'}]},
            ],
            'logfire.json_schema': {
                'type': 'object',
                'properties': {'pydantic_ai.all_messages': {'type': 'array'}, 'final_result': {'type': 'object'}},
            },
        }
    )

    # Verify that model/chat span still uses standard attribute names
    chat_span = next(s for s in spans if 'chat' in s['name'])
    assert chat_span['attributes']['gen_ai.usage.input_tokens'] == 10
    assert chat_span['attributes']['gen_ai.usage.output_tokens'] == 5


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.anyio
async def test_aggregated_usage_attribute_names_can_be_disabled(capfire: CaptureLogfire) -> None:
    def model_function(messages: list[ModelRequest | ModelResponse], info: AgentInfo) -> ModelResponse:
        return ModelResponse(parts=[TextPart('Hello!')], usage=RequestUsage(input_tokens=10, output_tokens=5))

    settings = InstrumentationSettings(use_aggregated_usage_attribute_names=False)
    agent = Agent(model=FunctionModel(model_function), capabilities=[Instrumentation(settings=settings)])

    await agent.run('Hello')

    spans = capfire.exporter.exported_spans_as_dict(parse_json_attributes=True)
    agent_run_span = next(s for s in spans if s['name'] == 'invoke_agent agent')
    assert agent_run_span['attributes']['gen_ai.usage.input_tokens'] == 10
    assert agent_run_span['attributes']['gen_ai.usage.output_tokens'] == 5
    assert 'gen_ai.aggregated_usage.input_tokens' not in agent_run_span['attributes']


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.anyio
async def test_feedback(capfire: CaptureLogfire) -> None:
    from logfire.experimental.annotations import record_feedback

    my_agent = Agent(model=TestModel(), capabilities=[Instrumentation(settings=InstrumentationSettings())])

    async with my_agent.iter('Hello') as agent_run:
        async for _ in agent_run:
            pass
        result = agent_run.result
        assert result
        traceparent = get_traceparent(result)
        assert traceparent == get_traceparent(agent_run)
    assert traceparent == snapshot('00-00000000000000000000000000000001-0000000000000001-01')
    record_feedback(traceparent, 'factuality', 0.1, comment='the agent lied', extra={'foo': 'bar'})

    assert strip_logfire_metrics(capfire.exporter.exported_spans_as_dict(parse_json_attributes=True)) == snapshot(
        [
            {
                'name': 'chat test',
                'context': {'trace_id': 1, 'span_id': 3, 'is_remote': False},
                'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                'start_time': 2000000000,
                'end_time': 3000000000,
                'attributes': {
                    'gen_ai.operation.name': 'chat',
                    'gen_ai.provider.name': 'test',
                    'gen_ai.system': 'test',
                    'gen_ai.request.model': 'test',
                    'model_request_parameters': {
                        'function_tools': [],
                        'native_tools': [],
                        'output_mode': 'text',
                        'output_object': None,
                        'output_tools': [],
                        'prompted_output_template': None,
                        'allow_text_output': True,
                        'allow_image_output': False,
                        'instruction_parts': None,
                        'thinking': None,
                    },
                    'logfire.span_type': 'span',
                    'logfire.msg': 'chat test',
                    'gen_ai.conversation.id': IsStr(),
                    'gen_ai.agent.name': 'agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.usage.input_tokens': 51,
                    'gen_ai.usage.output_tokens': 4,
                    'gen_ai.response.model': 'test',
                    'gen_ai.input.messages': [
                        {
                            'parts': [
                                {
                                    'type': 'text',
                                    'content': 'Hello',
                                },
                            ],
                            'role': 'user',
                        },
                    ],
                    'gen_ai.output.messages': [
                        {
                            'role': 'assistant',
                            'parts': [{'type': 'text', 'content': 'success (no tool calls)'}],
                        }
                    ],
                    'logfire.json_schema': {
                        'type': 'object',
                        'properties': {
                            'gen_ai.input.messages': {'type': 'array'},
                            'gen_ai.output.messages': {'type': 'array'},
                            'model_request_parameters': {'type': 'object'},
                        },
                    },
                },
            },
            {
                'name': 'invoke_agent agent',
                'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                'parent': None,
                'start_time': 1000000000,
                'end_time': 4000000000,
                'attributes': {
                    'model_name': 'test',
                    'agent_name': 'agent',
                    'gen_ai.agent.name': 'agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.conversation.id': IsStr(),
                    'gen_ai.operation.name': 'invoke_agent',
                    'logfire.msg': 'agent run',
                    'logfire.span_type': 'span',
                    'final_result': 'success (no tool calls)',
                    'gen_ai.aggregated_usage.input_tokens': 51,
                    'gen_ai.aggregated_usage.output_tokens': 4,
                    'pydantic_ai.all_messages': [
                        {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]},
                        {'role': 'assistant', 'parts': [{'type': 'text', 'content': 'success (no tool calls)'}]},
                    ],
                    'logfire.json_schema': {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'final_result': {'type': 'object'},
                        },
                    },
                },
            },
            {
                'name': 'feedback: factuality',
                'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
                'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': True},
                'start_time': 5000000000,
                'end_time': 5000000000,
                'attributes': {
                    'logfire.span_type': 'annotation',
                    'logfire.level_num': 9,
                    'logfire.msg_template': 'feedback: factuality',
                    'logfire.msg': 'feedback: factuality = 0.1',
                    'code.filepath': 'test_logfire.py',
                    'code.function': 'test_feedback',
                    'code.lineno': 123,
                    'logfire.feedback.name': 'factuality',
                    'factuality': 0.1,
                    'foo': 'bar',
                    'logfire.feedback.comment': 'the agent lied',
                    'logfire.json_schema': {
                        'type': 'object',
                        'properties': {
                            'logfire.feedback.name': {},
                            'factuality': {},
                            'foo': {},
                            'logfire.feedback.comment': {},
                            'logfire.span_type': {},
                        },
                    },
                },
            },
        ]
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize('include_content,tool_error', [(True, False), (True, True), (False, False), (False, True)])
def test_include_tool_args_span_attributes(
    get_logfire_summary: Callable[[], LogfireSummary],
    include_content: bool,
    tool_error: bool,
) -> None:
    """Test that tool arguments are included/excluded in span attributes based on instrumentation settings."""

    instrumentation_settings = InstrumentationSettings(include_content=include_content)
    test_model = TestModel(seed=42)
    my_agent = Agent(model=test_model, capabilities=[Instrumentation(settings=instrumentation_settings)])

    @my_agent.tool_plain
    async def add_numbers(x: int, y: int) -> int:
        """Add two numbers together."""
        if tool_error:
            raise ModelRetry('Tool error')
        return x + y

    try:
        result = my_agent.run_sync('Add 42 and 42')
        assert result.output == snapshot('{"add_numbers":84}')
    except UnexpectedModelBehavior:
        if not tool_error:
            raise  # pragma: no cover

    summary = get_logfire_summary()

    tool_attributes = next(
        attributes for attributes in summary.attributes.values() if attributes.get('gen_ai.tool.name') == 'add_numbers'
    )

    if include_content:
        if tool_error:
            assert tool_attributes == snapshot(
                {
                    'gen_ai.operation.name': 'execute_tool',
                    'gen_ai.tool.name': 'add_numbers',
                    'gen_ai.tool.call.id': IsStr(),
                    'gen_ai.tool.call.arguments': '{"x":42,"y":42}',
                    'logfire.msg': 'running tool: add_numbers',
                    'logfire.json_schema': IsJson(
                        snapshot(
                            {
                                'type': 'object',
                                'properties': {
                                    'gen_ai.tool.call.arguments': {'type': 'object'},
                                    'gen_ai.tool.call.result': {'type': 'object'},
                                    'gen_ai.tool.name': {},
                                    'gen_ai.tool.call.id': {},
                                },
                            }
                        )
                    ),
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.span_type': 'span',
                    'gen_ai.agent.name': 'my_agent',
                    'gen_ai.tool.call.result': """\
Tool error

Fix the errors and try again.\
""",
                    'gen_ai.agent.call.id': IsStr(),
                    'logfire.level_num': 17,
                }
            )
        else:
            assert tool_attributes == snapshot(
                {
                    'gen_ai.operation.name': 'execute_tool',
                    'gen_ai.tool.name': 'add_numbers',
                    'gen_ai.tool.call.id': IsStr(),
                    'gen_ai.tool.call.arguments': '{"x":42,"y":42}',
                    'logfire.msg': 'running tool: add_numbers',
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.json_schema': IsJson(
                        snapshot(
                            {
                                'type': 'object',
                                'properties': {
                                    'gen_ai.tool.call.arguments': {'type': 'object'},
                                    'gen_ai.tool.call.result': {'type': 'object'},
                                    'gen_ai.tool.name': {},
                                    'gen_ai.tool.call.id': {},
                                },
                            }
                        )
                    ),
                    'gen_ai.agent.name': 'my_agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'logfire.span_type': 'span',
                    'gen_ai.tool.call.result': '84',
                }
            )
    else:
        if tool_error:
            assert tool_attributes == snapshot(
                {
                    'gen_ai.operation.name': 'execute_tool',
                    'gen_ai.tool.name': 'add_numbers',
                    'gen_ai.tool.call.id': IsStr(),
                    'logfire.msg': 'running tool: add_numbers',
                    'logfire.json_schema': IsJson(
                        snapshot(
                            {
                                'type': 'object',
                                'properties': {
                                    'gen_ai.tool.name': {},
                                    'gen_ai.tool.call.id': {},
                                },
                            }
                        )
                    ),
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.span_type': 'span',
                    'gen_ai.agent.name': 'my_agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'logfire.level_num': 17,
                }
            )
        else:
            assert tool_attributes == snapshot(
                {
                    'gen_ai.operation.name': 'execute_tool',
                    'gen_ai.tool.name': 'add_numbers',
                    'gen_ai.tool.call.id': IsStr(),
                    'logfire.msg': 'running tool: add_numbers',
                    'logfire.json_schema': IsJson(
                        snapshot(
                            {
                                'type': 'object',
                                'properties': {
                                    'gen_ai.tool.name': {},
                                    'gen_ai.tool.call.id': {},
                                },
                            }
                        )
                    ),
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.span_type': 'span',
                    'gen_ai.agent.name': 'my_agent',
                    'gen_ai.agent.call.id': IsStr(),
                }
            )


class WeatherInfo(BaseModel):
    temperature: float
    description: str


def get_weather_info(city: str) -> WeatherInfo:
    return WeatherInfo(temperature=28.7, description='sunny')


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize(
    'instrument',
    [
        True,
        False,
        deprecated_instrumentation_settings(version=2),
        deprecated_instrumentation_settings(version=3),
    ],
)
def test_logfire_output_function_v2_v3(
    get_logfire_summary: Callable[[], LogfireSummary],
    instrument: InstrumentationSettings | bool,
) -> None:
    def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        assert info.output_tools is not None
        args_json = '{"city": "Mexico City"}'
        return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)])

    capabilities = (
        [
            Instrumentation(
                settings=instrument if isinstance(instrument, InstrumentationSettings) else InstrumentationSettings()
            )
        ]
        if instrument
        else []
    )
    my_agent = Agent(model=FunctionModel(call_tool), capabilities=capabilities)
    result = my_agent.run_sync('Mexico City', output_type=get_weather_info)
    assert result.output == WeatherInfo(temperature=28.7, description='sunny')

    summary = get_logfire_summary()

    if isinstance(instrument, InstrumentationSettings) and instrument.version == 2:
        [output_function_attributes] = [
            attributes
            for attributes in summary.attributes.values()
            if attributes.get('gen_ai.tool.name') == 'final_result'
            and 'output function' in attributes.get('logfire.msg', '')
        ]
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'agent run',
                    'message': 'my_agent run',
                    'children': [
                        {'id': 1, 'name': 'chat function:call_tool:', 'message': 'chat function:call_tool:'},
                        {
                            'id': 2,
                            'name': 'running output function',
                            'message': 'running output function: final_result',
                        },
                    ],
                }
            ]
        )
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'final_result',
                'logfire.msg': 'running output function: final_result',
                'gen_ai.tool.call.id': IsStr(),
                'gen_ai.conversation.id': IsStr(),
                'logfire.json_schema': IsJson(
                    snapshot(
                        {
                            'type': 'object',
                            'properties': {
                                'tool_arguments': {'type': 'object'},
                                'tool_response': {'type': 'object'},
                                'gen_ai.tool.name': {},
                                'gen_ai.tool.call.id': {},
                            },
                        }
                    )
                ),
                'tool_arguments': '"Mexico City"',
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
                'tool_response': '{"temperature":28.7,"description":"sunny"}',
            }
        )

    elif instrument is True or isinstance(instrument, InstrumentationSettings) and instrument.version == 3:
        [output_function_attributes] = [
            attributes
            for attributes in summary.attributes.values()
            if attributes.get('gen_ai.tool.name') == 'final_result'
            and 'output function' in attributes.get('logfire.msg', '')
        ]
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'invoke_agent my_agent',
                    'message': 'my_agent run',
                    'children': [
                        {'id': 1, 'name': 'chat function:call_tool:', 'message': 'chat function:call_tool:'},
                        {
                            'id': 2,
                            'name': 'execute_tool final_result',
                            'message': 'running output function: final_result',
                        },
                    ],
                }
            ]
        )
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'final_result',
                'logfire.msg': 'running output function: final_result',
                'gen_ai.tool.call.id': IsStr(),
                'gen_ai.conversation.id': IsStr(),
                'gen_ai.tool.call.arguments': '"Mexico City"',
                'logfire.json_schema': IsJson(
                    snapshot(
                        {
                            'type': 'object',
                            'properties': {
                                'gen_ai.tool.call.arguments': {'type': 'object'},
                                'gen_ai.tool.call.result': {'type': 'object'},
                                'gen_ai.tool.name': {},
                                'gen_ai.tool.call.id': {},
                            },
                        }
                    )
                ),
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
                'gen_ai.tool.call.result': '{"temperature":28.7,"description":"sunny"}',
            }
        )
    else:
        assert summary.traces == snapshot([])
        assert summary.attributes == snapshot({})


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize('include_content', [True, False])
def test_output_type_function_logfire_attributes(
    get_logfire_summary: Callable[[], LogfireSummary],
    include_content: bool,
) -> None:
    def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        assert info.output_tools is not None
        args_json = '{"city": "Mexico City"}'
        return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)])

    instrumentation_settings = InstrumentationSettings(include_content=include_content)
    my_agent = Agent(model=FunctionModel(call_tool), capabilities=[Instrumentation(settings=instrumentation_settings)])

    result = my_agent.run_sync('Mexico City', output_type=get_weather_info)
    assert result.output == WeatherInfo(temperature=28.7, description='sunny')

    summary = get_logfire_summary()

    # Find the output function span attributes
    [output_function_attributes] = [
        attributes
        for attributes in summary.attributes.values()
        if attributes.get('gen_ai.tool.name') == 'final_result'
        and 'output function' in attributes.get('logfire.msg', '')
    ]

    if include_content:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'final_result',
                'gen_ai.tool.call.id': IsStr(),
                'gen_ai.conversation.id': IsStr(),
                'logfire.msg': 'running output function: final_result',
                'logfire.json_schema': IsJson(
                    snapshot(
                        {
                            'type': 'object',
                            'properties': {
                                'gen_ai.tool.call.arguments': {'type': 'object'},
                                'gen_ai.tool.call.result': {'type': 'object'},
                                'gen_ai.tool.name': {},
                                'gen_ai.tool.call.id': {},
                            },
                        }
                    )
                ),
                'gen_ai.tool.call.arguments': '"Mexico City"',
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
                'gen_ai.tool.call.result': '{"temperature":28.7,"description":"sunny"}',
            }
        )
    else:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'final_result',
                'gen_ai.tool.call.id': IsStr(),
                'logfire.msg': 'running output function: final_result',
                'gen_ai.conversation.id': IsStr(),
                'logfire.json_schema': '{"type":"object","properties":{"gen_ai.tool.name":{},"gen_ai.tool.call.id":{}}}',
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
            }
        )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize('include_content', [True, False])
def test_output_type_function_with_run_context_logfire_attributes(
    get_logfire_summary: Callable[[], LogfireSummary],
    include_content: bool,
) -> None:
    def get_weather_with_ctx(ctx: RunContext, city: str) -> WeatherInfo:
        assert ctx is not None
        return WeatherInfo(temperature=28.7, description='sunny')

    def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        assert info.output_tools is not None
        args_json = '{"city": "Mexico City"}'
        return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)])

    instrumentation_settings = InstrumentationSettings(include_content=include_content)
    my_agent = Agent(model=FunctionModel(call_tool), capabilities=[Instrumentation(settings=instrumentation_settings)])

    result = my_agent.run_sync('Mexico City', output_type=get_weather_with_ctx)
    assert result.output == WeatherInfo(temperature=28.7, description='sunny')

    summary = get_logfire_summary()

    # Find the output function span attributes
    [output_function_attributes] = [
        attributes
        for attributes in summary.attributes.values()
        if attributes.get('gen_ai.tool.name') == 'final_result'
        and 'output function' in attributes.get('logfire.msg', '')
    ]

    if include_content:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'final_result',
                'gen_ai.tool.call.id': IsStr(),
                'gen_ai.conversation.id': IsStr(),
                'logfire.msg': 'running output function: final_result',
                'logfire.json_schema': IsJson(
                    snapshot(
                        {
                            'type': 'object',
                            'properties': {
                                'gen_ai.tool.call.arguments': {'type': 'object'},
                                'gen_ai.tool.call.result': {'type': 'object'},
                                'gen_ai.tool.name': {},
                                'gen_ai.tool.call.id': {},
                            },
                        }
                    )
                ),
                'gen_ai.tool.call.arguments': '"Mexico City"',
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
                'gen_ai.tool.call.result': '{"temperature":28.7,"description":"sunny"}',
            }
        )
    else:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'final_result',
                'gen_ai.tool.call.id': IsStr(),
                'logfire.msg': 'running output function: final_result',
                'gen_ai.conversation.id': IsStr(),
                'logfire.json_schema': '{"type":"object","properties":{"gen_ai.tool.name":{},"gen_ai.tool.call.id":{}}}',
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
            }
        )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize('include_content', [True, False])
def test_output_type_function_with_retry_logfire_attributes(
    get_logfire_summary: Callable[[], LogfireSummary],
    include_content: bool,
) -> None:
    def get_weather_with_retry(city: str) -> WeatherInfo:
        if city != 'Mexico City':
            from pydantic_ai import ModelRetry

            raise ModelRetry('City not found, I only know Mexico City')
        return WeatherInfo(temperature=28.7, description='sunny')

    def call_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        assert info.output_tools is not None

        if len(messages) == 1:
            args_json = '{"city": "New York City"}'
        else:
            args_json = '{"city": "Mexico City"}'

        return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)])

    instrumentation_settings = InstrumentationSettings(include_content=include_content)
    my_agent = Agent(model=FunctionModel(call_tool), capabilities=[Instrumentation(settings=instrumentation_settings)])

    result = my_agent.run_sync('New York City', output_type=get_weather_with_retry)
    assert result.output == WeatherInfo(temperature=28.7, description='sunny')

    summary = get_logfire_summary()

    output_function_attributes = [
        attributes
        for attributes in summary.attributes.values()
        if attributes.get('gen_ai.tool.name') == 'final_result'
        and 'output function' in attributes.get('logfire.msg', '')
    ]

    if include_content:
        assert output_function_attributes == snapshot(
            [
                {
                    'gen_ai.operation.name': 'execute_tool',
                    'gen_ai.tool.name': 'final_result',
                    'gen_ai.agent.name': 'my_agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.msg': 'running output function: final_result',
                    'gen_ai.tool.call.id': IsStr(),
                    'gen_ai.tool.call.arguments': '"New York City"',
                    'logfire.json_schema': '{"type":"object","properties":{"gen_ai.tool.call.arguments":{"type":"object"},"gen_ai.tool.call.result":{"type":"object"},"gen_ai.tool.name":{},"gen_ai.tool.call.id":{}}}',
                    'logfire.span_type': 'span',
                    'logfire.level_num': 17,
                },
                {
                    'gen_ai.operation.name': 'execute_tool',
                    'gen_ai.tool.name': 'final_result',
                    'gen_ai.agent.name': 'my_agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.msg': 'running output function: final_result',
                    'gen_ai.tool.call.id': IsStr(),
                    'gen_ai.tool.call.arguments': '"Mexico City"',
                    'logfire.json_schema': '{"type":"object","properties":{"gen_ai.tool.call.arguments":{"type":"object"},"gen_ai.tool.call.result":{"type":"object"},"gen_ai.tool.name":{},"gen_ai.tool.call.id":{}}}',
                    'logfire.span_type': 'span',
                    'gen_ai.tool.call.result': '{"temperature":28.7,"description":"sunny"}',
                },
            ]
        )
    else:
        assert output_function_attributes == snapshot(
            [
                {
                    'gen_ai.operation.name': 'execute_tool',
                    'gen_ai.tool.name': 'final_result',
                    'gen_ai.agent.name': 'my_agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.msg': 'running output function: final_result',
                    'gen_ai.tool.call.id': IsStr(),
                    'logfire.json_schema': '{"type":"object","properties":{"gen_ai.tool.name":{},"gen_ai.tool.call.id":{}}}',
                    'logfire.span_type': 'span',
                    'logfire.level_num': 17,
                },
                {
                    'gen_ai.operation.name': 'execute_tool',
                    'gen_ai.tool.name': 'final_result',
                    'gen_ai.agent.name': 'my_agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.msg': 'running output function: final_result',
                    'gen_ai.tool.call.id': IsStr(),
                    'logfire.json_schema': '{"type":"object","properties":{"gen_ai.tool.name":{},"gen_ai.tool.call.id":{}}}',
                    'logfire.span_type': 'span',
                },
            ]
        )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize('include_content', [True, False])
def test_output_type_function_with_custom_tool_name_logfire_attributes(
    get_logfire_summary: Callable[[], LogfireSummary],
    include_content: bool,
) -> None:
    def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        assert info.output_tools is not None
        args_json = '{"city": "Mexico City"}'
        return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)])

    from pydantic_ai.output import ToolOutput

    instrumentation_settings = InstrumentationSettings(include_content=include_content)
    my_agent = Agent(model=FunctionModel(call_tool), capabilities=[Instrumentation(settings=instrumentation_settings)])

    result = my_agent.run_sync('Mexico City', output_type=ToolOutput(get_weather_info, name='get_weather'))
    assert result.output == WeatherInfo(temperature=28.7, description='sunny')

    summary = get_logfire_summary()

    # Find the output function span attributes with custom tool name
    [output_function_attributes] = [
        attributes
        for attributes in summary.attributes.values()
        if attributes.get('gen_ai.tool.name') == 'get_weather'
        and 'output function' in attributes.get('logfire.msg', '')
    ]

    if include_content:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'get_weather',
                'gen_ai.tool.call.id': IsStr(),
                'gen_ai.conversation.id': IsStr(),
                'logfire.msg': 'running output function: get_weather',
                'logfire.json_schema': IsJson(
                    snapshot(
                        {
                            'type': 'object',
                            'properties': {
                                'gen_ai.tool.call.arguments': {'type': 'object'},
                                'gen_ai.tool.call.result': {'type': 'object'},
                                'gen_ai.tool.name': {},
                                'gen_ai.tool.call.id': {},
                            },
                        }
                    )
                ),
                'gen_ai.tool.call.arguments': '"Mexico City"',
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
                'gen_ai.tool.call.result': '{"temperature":28.7,"description":"sunny"}',
            }
        )
    else:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'get_weather',
                'gen_ai.tool.call.id': IsStr(),
                'logfire.msg': 'running output function: get_weather',
                'gen_ai.conversation.id': IsStr(),
                'logfire.json_schema': IsJson(
                    snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}})
                ),
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
            }
        )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize('include_content', [True, False])
def test_output_type_bound_instance_method_logfire_attributes(
    get_logfire_summary: Callable[[], LogfireSummary],
    include_content: bool,
) -> None:
    class Weather(BaseModel):
        temperature: float
        description: str

        def get_weather(self, city: str):
            return self

    weather = Weather(temperature=28.7, description='sunny')

    def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        assert info.output_tools is not None
        args_json = '{"city": "Mexico City"}'
        return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)])

    instrumentation_settings = InstrumentationSettings(include_content=include_content)
    my_agent = Agent(model=FunctionModel(call_tool), capabilities=[Instrumentation(settings=instrumentation_settings)])

    result = my_agent.run_sync('Mexico City', output_type=weather.get_weather)
    assert result.output == Weather(temperature=28.7, description='sunny')

    summary = get_logfire_summary()

    # Find the output function span attributes
    [output_function_attributes] = [
        attributes
        for attributes in summary.attributes.values()
        if attributes.get('gen_ai.tool.name') == 'final_result'
        and 'output function' in attributes.get('logfire.msg', '')
    ]

    if include_content:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'final_result',
                'gen_ai.tool.call.id': IsStr(),
                'gen_ai.conversation.id': IsStr(),
                'logfire.msg': 'running output function: final_result',
                'logfire.json_schema': IsJson(
                    snapshot(
                        {
                            'type': 'object',
                            'properties': {
                                'gen_ai.tool.call.arguments': {'type': 'object'},
                                'gen_ai.tool.call.result': {'type': 'object'},
                                'gen_ai.tool.name': {},
                                'gen_ai.tool.call.id': {},
                            },
                        }
                    )
                ),
                'gen_ai.tool.call.arguments': '"Mexico City"',
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
                'gen_ai.tool.call.result': '{"temperature":28.7,"description":"sunny"}',
            }
        )
    else:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'final_result',
                'gen_ai.tool.call.id': IsStr(),
                'logfire.msg': 'running output function: final_result',
                'gen_ai.conversation.id': IsStr(),
                'logfire.json_schema': IsJson(
                    snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}})
                ),
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
            }
        )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize('include_content', [True, False])
def test_output_type_bound_instance_method_with_run_context_logfire_attributes(
    get_logfire_summary: Callable[[], LogfireSummary],
    include_content: bool,
) -> None:
    class Weather(BaseModel):
        temperature: float
        description: str

        def get_weather(self, ctx: RunContext, city: str):
            assert ctx is not None
            return self

    weather = Weather(temperature=28.7, description='sunny')

    def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        assert info.output_tools is not None
        args_json = '{"city": "Mexico City"}'
        return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)])

    instrumentation_settings = InstrumentationSettings(include_content=include_content)
    my_agent = Agent(model=FunctionModel(call_tool), capabilities=[Instrumentation(settings=instrumentation_settings)])

    result = my_agent.run_sync('Mexico City', output_type=weather.get_weather)
    assert result.output == Weather(temperature=28.7, description='sunny')

    summary = get_logfire_summary()

    # Find the output function span attributes
    [output_function_attributes] = [
        attributes
        for attributes in summary.attributes.values()
        if attributes.get('gen_ai.tool.name') == 'final_result'
        and 'output function' in attributes.get('logfire.msg', '')
    ]

    if include_content:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'final_result',
                'gen_ai.tool.call.id': IsStr(),
                'gen_ai.conversation.id': IsStr(),
                'logfire.msg': 'running output function: final_result',
                'logfire.json_schema': IsJson(
                    snapshot(
                        {
                            'type': 'object',
                            'properties': {
                                'gen_ai.tool.call.arguments': {'type': 'object'},
                                'gen_ai.tool.call.result': {'type': 'object'},
                                'gen_ai.tool.name': {},
                                'gen_ai.tool.call.id': {},
                            },
                        }
                    )
                ),
                'gen_ai.tool.call.arguments': '"Mexico City"',
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
                'gen_ai.tool.call.result': '{"temperature":28.7,"description":"sunny"}',
            }
        )
    else:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'final_result',
                'gen_ai.tool.call.id': IsStr(),
                'logfire.msg': 'running output function: final_result',
                'gen_ai.conversation.id': IsStr(),
                'logfire.json_schema': IsJson(
                    snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}})
                ),
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
            }
        )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize('include_content', [True, False])
def test_output_type_async_function_logfire_attributes(
    get_logfire_summary: Callable[[], LogfireSummary],
    include_content: bool,
) -> None:
    """Test logfire attributes for async output function types."""

    async def get_weather_async(city: str) -> WeatherInfo:
        return WeatherInfo(temperature=28.7, description='sunny')

    def call_tool(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        assert info.output_tools is not None
        args_json = '{"city": "Mexico City"}'
        return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, args_json)])

    instrumentation_settings = InstrumentationSettings(include_content=include_content)
    my_agent = Agent(model=FunctionModel(call_tool), capabilities=[Instrumentation(settings=instrumentation_settings)])

    result = my_agent.run_sync('Mexico City', output_type=get_weather_async)
    assert result.output == WeatherInfo(temperature=28.7, description='sunny')

    summary = get_logfire_summary()

    # Find the output function span attributes
    [output_function_attributes] = [
        attributes
        for attributes in summary.attributes.values()
        if attributes.get('gen_ai.tool.name') == 'final_result'
        and 'output function' in attributes.get('logfire.msg', '')
    ]

    if include_content:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'final_result',
                'gen_ai.tool.call.id': IsStr(),
                'gen_ai.conversation.id': IsStr(),
                'logfire.msg': 'running output function: final_result',
                'logfire.json_schema': IsJson(
                    snapshot(
                        {
                            'type': 'object',
                            'properties': {
                                'gen_ai.tool.call.arguments': {'type': 'object'},
                                'gen_ai.tool.call.result': {'type': 'object'},
                                'gen_ai.tool.name': {},
                                'gen_ai.tool.call.id': {},
                            },
                        }
                    )
                ),
                'gen_ai.tool.call.arguments': '"Mexico City"',
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
                'gen_ai.tool.call.result': '{"temperature":28.7,"description":"sunny"}',
            }
        )
    else:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'final_result',
                'gen_ai.tool.call.id': IsStr(),
                'logfire.msg': 'running output function: final_result',
                'gen_ai.conversation.id': IsStr(),
                'logfire.json_schema': IsJson(
                    snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}, 'gen_ai.tool.call.id': {}}})
                ),
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
            }
        )


def upcase_text(text: str) -> str:
    """Convert text to uppercase."""
    return text.upper()


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize('include_content', [True, False])
def test_text_output_function_logfire_attributes(
    get_logfire_summary: Callable[[], LogfireSummary],
    include_content: bool,
) -> None:
    """Test logfire attributes for TextOutput functions (TextOutputProcessor)."""

    def call_text_response(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        # Return a plain text response (not a tool call)
        from pydantic_ai import TextPart

        return ModelResponse(parts=[TextPart(content='hello world')])

    instrumentation_settings = InstrumentationSettings(include_content=include_content)
    my_agent = Agent(
        model=FunctionModel(call_text_response), capabilities=[Instrumentation(settings=instrumentation_settings)]
    )

    result = my_agent.run_sync('Say hello', output_type=TextOutput(upcase_text))
    assert result.output == 'HELLO WORLD'

    summary = get_logfire_summary()

    # Find the text output function span attributes
    [text_function_attributes] = [
        attributes
        for attributes in summary.attributes.values()
        if 'running output function: upcase_text' in attributes.get('logfire.msg', '')
    ]

    if include_content:
        assert text_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'upcase_text',
                'logfire.msg': 'running output function: upcase_text',
                'gen_ai.conversation.id': IsStr(),
                'logfire.json_schema': IsJson(
                    snapshot(
                        {
                            'type': 'object',
                            'properties': {
                                'gen_ai.tool.call.arguments': {'type': 'object'},
                                'gen_ai.tool.call.result': {'type': 'object'},
                                'gen_ai.tool.name': {},
                            },
                        }
                    )
                ),
                'gen_ai.tool.call.arguments': '"hello world"',
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
                'gen_ai.tool.call.result': 'HELLO WORLD',
            }
        )
    else:
        assert text_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'upcase_text',
                'logfire.msg': 'running output function: upcase_text',
                'logfire.json_schema': IsJson(snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}}})),
                'gen_ai.conversation.id': IsStr(),
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'my_agent',
                'gen_ai.agent.call.id': IsStr(),
            }
        )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize('include_content', [True, False])
def test_prompted_output_function_logfire_attributes(
    get_logfire_summary: Callable[[], LogfireSummary],
    include_content: bool,
) -> None:
    """Test that spans are created for PromptedOutput functions with appropriate attributes."""

    def upcase_text(text: str) -> str:
        return text.upper()

    call_count = 0

    def call_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        nonlocal call_count
        call_count += 1
        # Simulate the model returning JSON that will be parsed and used to call the function
        return ModelResponse(parts=[TextPart(content='{"text": "hello world"}')])

    instrumentation_settings = InstrumentationSettings(include_content=include_content)
    agent = Agent(
        model=FunctionModel(call_tool),
        capabilities=[Instrumentation(settings=instrumentation_settings)],
        output_type=PromptedOutput(upcase_text),
    )

    result = agent.run_sync('test')

    # Check that the function was called and returned the expected result
    assert result.output == 'HELLO WORLD'
    assert call_count == 1

    summary = get_logfire_summary()

    # Find the output function span attributes
    [output_function_attributes] = [
        attributes
        for attributes in summary.attributes.values()
        if attributes.get('logfire.msg', '').startswith('running output function: upcase_text')
    ]

    if include_content:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'upcase_text',
                'logfire.msg': 'running output function: upcase_text',
                'gen_ai.conversation.id': IsStr(),
                'logfire.json_schema': IsJson(
                    snapshot(
                        {
                            'type': 'object',
                            'properties': {
                                'gen_ai.tool.call.arguments': {'type': 'object'},
                                'gen_ai.tool.call.result': {'type': 'object'},
                                'gen_ai.tool.name': {},
                            },
                        }
                    )
                ),
                'gen_ai.tool.call.arguments': '"hello world"',
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'agent',
                'gen_ai.agent.call.id': IsStr(),
                'gen_ai.tool.call.result': 'HELLO WORLD',
            }
        )
    else:
        assert output_function_attributes == snapshot(
            {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'upcase_text',
                'logfire.msg': 'running output function: upcase_text',
                'logfire.json_schema': IsJson(snapshot({'type': 'object', 'properties': {'gen_ai.tool.name': {}}})),
                'gen_ai.conversation.id': IsStr(),
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'agent',
                'gen_ai.agent.call.id': IsStr(),
            }
        )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize('include_content', [True, False])
def test_output_type_text_output_function_with_retry_logfire_attributes(
    get_logfire_summary: Callable[[], LogfireSummary],
    include_content: bool,
) -> None:
    def get_weather_with_retry(ctx: RunContext, city: str) -> WeatherInfo:
        assert ctx is not None
        if city != 'Mexico City':
            from pydantic_ai import ModelRetry

            raise ModelRetry('City not found, I only know Mexico City')
        return WeatherInfo(temperature=28.7, description='sunny')

    def call_tool(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        assert info.output_tools is not None

        if len(messages) == 1:
            city = 'New York City'
        else:
            city = 'Mexico City'

        return ModelResponse(parts=[TextPart(content=city)])

    instrumentation_settings = InstrumentationSettings(include_content=include_content)
    my_agent = Agent(model=FunctionModel(call_tool), capabilities=[Instrumentation(settings=instrumentation_settings)])

    result = my_agent.run_sync('New York City', output_type=TextOutput(get_weather_with_retry))
    assert result.output == WeatherInfo(temperature=28.7, description='sunny')

    summary = get_logfire_summary()

    text_function_attributes = [
        attributes
        for attributes in summary.attributes.values()
        if 'running output function: get_weather_with_retry' in attributes.get('logfire.msg', '')
    ]

    if include_content:
        assert text_function_attributes == snapshot(
            [
                {
                    'gen_ai.operation.name': 'execute_tool',
                    'gen_ai.tool.name': 'get_weather_with_retry',
                    'gen_ai.agent.name': 'my_agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.msg': 'running output function: get_weather_with_retry',
                    'gen_ai.tool.call.arguments': '"New York City"',
                    'logfire.json_schema': '{"type":"object","properties":{"gen_ai.tool.call.arguments":{"type":"object"},"gen_ai.tool.call.result":{"type":"object"},"gen_ai.tool.name":{}}}',
                    'logfire.span_type': 'span',
                    'logfire.level_num': 17,
                },
                {
                    'gen_ai.operation.name': 'execute_tool',
                    'gen_ai.tool.name': 'get_weather_with_retry',
                    'gen_ai.agent.name': 'my_agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.msg': 'running output function: get_weather_with_retry',
                    'gen_ai.tool.call.arguments': '"Mexico City"',
                    'logfire.json_schema': '{"type":"object","properties":{"gen_ai.tool.call.arguments":{"type":"object"},"gen_ai.tool.call.result":{"type":"object"},"gen_ai.tool.name":{}}}',
                    'logfire.span_type': 'span',
                    'gen_ai.tool.call.result': '{"temperature":28.7,"description":"sunny"}',
                },
            ]
        )
    else:
        assert text_function_attributes == snapshot(
            [
                {
                    'gen_ai.operation.name': 'execute_tool',
                    'gen_ai.tool.name': 'get_weather_with_retry',
                    'gen_ai.agent.name': 'my_agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.msg': 'running output function: get_weather_with_retry',
                    'logfire.json_schema': '{"type":"object","properties":{"gen_ai.tool.name":{}}}',
                    'logfire.span_type': 'span',
                    'logfire.level_num': 17,
                },
                {
                    'gen_ai.operation.name': 'execute_tool',
                    'gen_ai.tool.name': 'get_weather_with_retry',
                    'gen_ai.agent.name': 'my_agent',
                    'gen_ai.agent.call.id': IsStr(),
                    'gen_ai.conversation.id': IsStr(),
                    'logfire.msg': 'running output function: get_weather_with_retry',
                    'logfire.json_schema': '{"type":"object","properties":{"gen_ai.tool.name":{}}}',
                    'logfire.span_type': 'span',
                },
            ]
        )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize(
    'instrument',
    [deprecated_instrumentation_settings(version=2), deprecated_instrumentation_settings(version=3)],
)
def test_static_function_instructions_in_agent_run_span(
    get_logfire_summary: Callable[[], LogfireSummary], instrument: InstrumentationSettings
) -> None:
    @dataclass
    class MyOutput:
        content: str

    my_agent = Agent(model=TestModel(), capabilities=[Instrumentation(settings=instrument)])

    @my_agent.instructions
    def instructions():
        return 'Here are some instructions'

    result = my_agent.run_sync('Hello', output_type=MyOutput)
    assert result.output == MyOutput(content='a')

    summary = get_logfire_summary()
    chat_span_attributes = summary.attributes[1]
    if instrument.version == 2:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'agent run',
                    'message': 'my_agent run',
                    'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
                }
            ]
        )
    else:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'invoke_agent my_agent',
                    'message': 'my_agent run',
                    'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
                }
            ]
        )

    assert summary.attributes[0] == snapshot(
        {
            'model_name': 'test',
            'agent_name': 'my_agent',
            'gen_ai.agent.name': 'my_agent',
            'gen_ai.agent.call.id': IsStr(),
            'gen_ai.conversation.id': IsStr(),
            'gen_ai.operation.name': 'invoke_agent',
            'logfire.msg': 'my_agent run',
            'logfire.span_type': 'span',
            'final_result': '{"content":"a"}',
            'gen_ai.aggregated_usage.input_tokens': 51,
            'gen_ai.aggregated_usage.output_tokens': 5,
            'pydantic_ai.all_messages': IsJson(
                snapshot(
                    [
                        {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]},
                        {
                            'role': 'assistant',
                            'parts': [
                                {
                                    'type': 'tool_call',
                                    'id': IsStr(),
                                    'name': 'final_result',
                                    'arguments': {'content': 'a'},
                                }
                            ],
                        },
                        {
                            'role': 'user',
                            'parts': [
                                {
                                    'type': 'tool_call_response',
                                    'id': IsStr(),
                                    'name': 'final_result',
                                    'result': 'Final result processed.',
                                }
                            ],
                        },
                    ]
                )
            ),
            'gen_ai.system_instructions': '[{"type":"text","content":"Here are some instructions"}]',
            'logfire.json_schema': IsJson(
                snapshot(
                    {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'gen_ai.system_instructions': {'type': 'array'},
                            'final_result': {'type': 'object'},
                        },
                    }
                )
            ),
        }
    )

    assert chat_span_attributes['gen_ai.input.messages'] == IsJson(
        snapshot([{'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]}])
    )
    assert chat_span_attributes['gen_ai.output.messages'] == IsJson(
        snapshot(
            [
                {
                    'role': 'assistant',
                    'parts': [
                        {
                            'type': 'tool_call',
                            'id': IsStr(),
                            'name': 'final_result',
                            'arguments': {'content': 'a'},
                        }
                    ],
                }
            ]
        )
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_instructions_from_history_when_model_request_fails_before_instrumentation(
    get_logfire_summary: Callable[[], LogfireSummary],
) -> None:
    class FailBeforeModelRequest(AbstractCapability[Any]):
        async def before_model_request(self, ctx: RunContext[Any], request_context: Any) -> Any:
            raise RuntimeError('boom')

    my_agent = Agent(
        model=TestModel(),
        capabilities=[Instrumentation(settings=InstrumentationSettings()), FailBeforeModelRequest()],
    )

    with pytest.raises(RuntimeError, match='boom'):
        my_agent.run_sync(
            'Hello',
            message_history=[
                ModelRequest(
                    parts=[UserPromptPart(content='Hi')],
                    instructions='Instructions from history',
                    timestamp=IsDatetime(),
                ),
                ModelResponse(parts=[TextPart(content='Hello')]),
            ],
        )

    summary = get_logfire_summary()
    assert summary.attributes[0]['gen_ai.system_instructions'] == snapshot(
        '[{"type":"text","content":"Instructions from history"}]'
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize(
    'instrument',
    [deprecated_instrumentation_settings(version=2), deprecated_instrumentation_settings(version=3)],
)
def test_dynamic_function_instructions_in_agent_run_span(
    get_logfire_summary: Callable[[], LogfireSummary], instrument: InstrumentationSettings
) -> None:
    @dataclass
    class MyOutput:
        content: str

    my_agent = Agent(model=TestModel(), capabilities=[Instrumentation(settings=instrument)])

    @my_agent.instructions
    def instructions(ctx: RunContext):
        return f'This is step {ctx.run_step}'

    @my_agent.tool_plain
    def my_tool() -> str:
        return 'This is a tool call'

    result = my_agent.run_sync('Hello', output_type=MyOutput)
    assert result.output == MyOutput(content='a')

    summary = get_logfire_summary()
    chat_span_attributes = summary.attributes[1]
    if instrument.version == 2:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'agent run',
                    'message': 'my_agent run',
                    'children': [
                        {'id': 1, 'name': 'chat test', 'message': 'chat test'},
                        {'id': 2, 'name': 'running tool', 'message': 'running tool: my_tool'},
                        {'id': 3, 'name': 'chat test', 'message': 'chat test'},
                    ],
                }
            ]
        )
    else:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'invoke_agent my_agent',
                    'message': 'my_agent run',
                    'children': [
                        {'id': 1, 'name': 'chat test', 'message': 'chat test'},
                        {'id': 2, 'name': 'execute_tool my_tool', 'message': 'running tool: my_tool'},
                        {'id': 3, 'name': 'chat test', 'message': 'chat test'},
                    ],
                }
            ]
        )

    assert summary.attributes[0] == snapshot(
        {
            'model_name': 'test',
            'agent_name': 'my_agent',
            'gen_ai.agent.name': 'my_agent',
            'gen_ai.agent.call.id': IsStr(),
            'gen_ai.conversation.id': IsStr(),
            'gen_ai.operation.name': 'invoke_agent',
            'logfire.msg': 'my_agent run',
            'logfire.span_type': 'span',
            'final_result': '{"content":"a"}',
            'gen_ai.aggregated_usage.input_tokens': 107,
            'gen_ai.aggregated_usage.output_tokens': 9,
            'pydantic_ai.all_messages': IsJson(
                snapshot(
                    [
                        {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]},
                        {
                            'role': 'assistant',
                            'parts': [
                                {
                                    'type': 'tool_call',
                                    'id': 'pyd_ai_tool_call_id__my_tool',
                                    'name': 'my_tool',
                                    'arguments': {},
                                }
                            ],
                        },
                        {
                            'role': 'user',
                            'parts': [
                                {
                                    'type': 'tool_call_response',
                                    'id': 'pyd_ai_tool_call_id__my_tool',
                                    'name': 'my_tool',
                                    'result': 'This is a tool call',
                                }
                            ],
                        },
                        {
                            'role': 'assistant',
                            'parts': [
                                {
                                    'type': 'tool_call',
                                    'id': IsStr(),
                                    'name': 'final_result',
                                    'arguments': {'content': 'a'},
                                }
                            ],
                        },
                        {
                            'role': 'user',
                            'parts': [
                                {
                                    'type': 'tool_call_response',
                                    'id': IsStr(),
                                    'name': 'final_result',
                                    'result': 'Final result processed.',
                                }
                            ],
                        },
                    ]
                )
            ),
            'gen_ai.system_instructions': '[{"type":"text","content":"This is step 2"}]',
            'pydantic_ai.variable_instructions': True,
            'logfire.json_schema': IsJson(
                snapshot(
                    {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'gen_ai.system_instructions': {'type': 'array'},
                            'pydantic_ai.variable_instructions': {},
                            'final_result': {'type': 'object'},
                        },
                    }
                )
            ),
        }
    )

    assert chat_span_attributes['gen_ai.input.messages'] == IsJson(
        snapshot([{'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]}])
    )
    assert chat_span_attributes['gen_ai.output.messages'] == IsJson(
        snapshot(
            [
                {
                    'role': 'assistant',
                    'parts': [
                        {
                            'type': 'tool_call',
                            'id': IsStr(),
                            'name': 'my_tool',
                            'arguments': {},
                        }
                    ],
                }
            ]
        )
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize(
    'instrument',
    [deprecated_instrumentation_settings(version=2), deprecated_instrumentation_settings(version=3)],
)
def test_function_instructions_with_history_in_agent_run_span(
    get_logfire_summary: Callable[[], LogfireSummary], instrument: InstrumentationSettings
) -> None:
    @dataclass
    class MyOutput:
        content: str

    my_agent = Agent(model=TestModel(), capabilities=[Instrumentation(settings=instrument)])

    @my_agent.instructions
    def instructions(ctx: RunContext):
        return 'Instructions for the current agent run'

    result = my_agent.run_sync(
        'Hello',
        message_history=[
            ModelRequest(
                parts=[UserPromptPart(content='Hi')],
                instructions='Instructions from a previous agent run',
                timestamp=IsDatetime(),
            ),
            ModelResponse(parts=[TextPart(content='Hello')]),
        ],
        output_type=MyOutput,
    )
    assert result.output == MyOutput(content='a')

    summary = get_logfire_summary()
    chat_span_attributes = summary.attributes[1]
    if instrument.version == 2:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'agent run',
                    'message': 'my_agent run',
                    'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
                }
            ]
        )
    else:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'invoke_agent my_agent',
                    'message': 'my_agent run',
                    'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
                }
            ]
        )

    assert summary.attributes[0] == snapshot(
        {
            'model_name': 'test',
            'agent_name': 'my_agent',
            'gen_ai.agent.name': 'my_agent',
            'gen_ai.agent.call.id': IsStr(),
            'gen_ai.conversation.id': IsStr(),
            'gen_ai.operation.name': 'invoke_agent',
            'logfire.msg': 'my_agent run',
            'logfire.span_type': 'span',
            'final_result': '{"content":"a"}',
            'gen_ai.aggregated_usage.input_tokens': 52,
            'gen_ai.aggregated_usage.output_tokens': 6,
            'pydantic_ai.all_messages': IsJson(
                snapshot(
                    [
                        {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hi'}]},
                        {'role': 'assistant', 'parts': [{'type': 'text', 'content': 'Hello'}]},
                        {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]},
                        {
                            'role': 'assistant',
                            'parts': [
                                {
                                    'type': 'tool_call',
                                    'id': IsStr(),
                                    'name': 'final_result',
                                    'arguments': {'content': 'a'},
                                }
                            ],
                        },
                        {
                            'role': 'user',
                            'parts': [
                                {
                                    'type': 'tool_call_response',
                                    'id': IsStr(),
                                    'name': 'final_result',
                                    'result': 'Final result processed.',
                                }
                            ],
                        },
                    ]
                )
            ),
            'pydantic_ai.new_message_index': 2,
            'gen_ai.system_instructions': '[{"type":"text","content":"Instructions for the current agent run"}]',
            'logfire.json_schema': IsJson(
                snapshot(
                    {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'pydantic_ai.new_message_index': {},
                            'gen_ai.system_instructions': {'type': 'array'},
                            'final_result': {'type': 'object'},
                        },
                    }
                )
            ),
        }
    )

    assert chat_span_attributes['gen_ai.input.messages'] == IsJson(
        snapshot(
            [
                {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hi'}]},
                {'role': 'assistant', 'parts': [{'type': 'text', 'content': 'Hello'}]},
                {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]},
            ]
        )
    )
    assert chat_span_attributes['gen_ai.output.messages'] == IsJson(
        snapshot(
            [
                {
                    'role': 'assistant',
                    'parts': [
                        {
                            'type': 'tool_call',
                            'id': IsStr(),
                            'name': 'final_result',
                            'arguments': {'content': 'a'},
                        }
                    ],
                }
            ]
        )
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.parametrize(
    'instrument',
    [deprecated_instrumentation_settings(version=2), deprecated_instrumentation_settings(version=3)],
)
async def test_run_stream(
    get_logfire_summary: Callable[[], LogfireSummary], instrument: InstrumentationSettings
) -> None:
    my_agent = Agent(model=TestModel(), capabilities=[Instrumentation(settings=instrument)])

    @my_agent.instructions
    def instructions(ctx: RunContext):
        return 'Instructions for the current agent run'

    async with my_agent.run_stream('Hello') as stream:
        async for _ in stream.stream_output():
            pass

    summary = get_logfire_summary()
    chat_span_attributes = summary.attributes[1]
    if instrument.version == 2:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'agent run',
                    'message': 'my_agent run',
                    'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
                }
            ]
        )
    else:
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'invoke_agent my_agent',
                    'message': 'my_agent run',
                    'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
                }
            ]
        )

    assert summary.attributes[0] == snapshot(
        {
            'model_name': 'test',
            'agent_name': 'my_agent',
            'gen_ai.agent.name': 'my_agent',
            'gen_ai.agent.call.id': IsStr(),
            'gen_ai.conversation.id': IsStr(),
            'gen_ai.operation.name': 'invoke_agent',
            'logfire.msg': 'my_agent run',
            'logfire.span_type': 'span',
            'final_result': 'success (no tool calls)',
            'gen_ai.aggregated_usage.input_tokens': 51,
            'gen_ai.aggregated_usage.output_tokens': 4,
            'pydantic_ai.all_messages': IsJson(
                snapshot(
                    [
                        {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]},
                        {'role': 'assistant', 'parts': [{'type': 'text', 'content': 'success (no tool calls)'}]},
                    ]
                )
            ),
            'gen_ai.system_instructions': '[{"type":"text","content":"Instructions for the current agent run"}]',
            'logfire.json_schema': IsJson(
                snapshot(
                    {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'gen_ai.system_instructions': {'type': 'array'},
                            'final_result': {'type': 'object'},
                        },
                    }
                )
            ),
        }
    )

    assert chat_span_attributes['gen_ai.input.messages'] == IsJson(
        snapshot([{'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]}])
    )
    assert chat_span_attributes['gen_ai.output.messages'] == IsJson(
        snapshot(
            [
                {
                    'role': 'assistant',
                    'parts': [{'type': 'text', 'content': 'success (no tool calls)'}],
                }
            ]
        )
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_run_stream_sync(get_logfire_summary: Callable[[], LogfireSummary]) -> None:
    my_agent = Agent(model=TestModel(), capabilities=[Instrumentation()])

    @my_agent.instructions
    def instructions(ctx: RunContext):
        return 'Instructions for the current agent run'

    with my_agent.run_stream_sync('Hello') as stream:
        for _ in stream.stream_output():
            pass

    summary = get_logfire_summary()
    # The `chat` span is correctly nested under the agent run span (not orphaned).
    assert summary.traces == snapshot(
        [
            {
                'id': 0,
                'name': 'invoke_agent my_agent',
                'message': 'my_agent run',
                'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
            }
        ]
    )
    # The agent run span's end attributes (all_messages, final_result, usage) are all populated.
    assert summary.attributes[0] == snapshot(
        {
            'model_name': 'test',
            'agent_name': 'my_agent',
            'gen_ai.agent.name': 'my_agent',
            'gen_ai.agent.call.id': IsStr(),
            'gen_ai.conversation.id': IsStr(),
            'gen_ai.operation.name': 'invoke_agent',
            'logfire.msg': 'my_agent run',
            'logfire.span_type': 'span',
            'final_result': 'success (no tool calls)',
            'gen_ai.aggregated_usage.input_tokens': 51,
            'gen_ai.aggregated_usage.output_tokens': 4,
            'pydantic_ai.all_messages': IsJson(
                snapshot(
                    [
                        {'role': 'user', 'parts': [{'type': 'text', 'content': 'Hello'}]},
                        {'role': 'assistant', 'parts': [{'type': 'text', 'content': 'success (no tool calls)'}]},
                    ]
                )
            ),
            'gen_ai.system_instructions': '[{"type":"text","content":"Instructions for the current agent run"}]',
            'logfire.json_schema': IsJson(
                snapshot(
                    {
                        'type': 'object',
                        'properties': {
                            'pydantic_ai.all_messages': {'type': 'array'},
                            'gen_ai.system_instructions': {'type': 'array'},
                            'final_result': {'type': 'object'},
                        },
                    }
                )
            ),
        }
    )


def _get_tool_span(capfire: CaptureLogfire) -> dict[str, Any]:
    """Get the completed tool span from exported spans."""
    spans = strip_logfire_metrics(capfire.exporter.exported_spans_as_dict(parse_json_attributes=True))
    tool_span = next(
        s for s in spans if s['attributes'].get('logfire.span_type') == 'span' and 'tool' in s['name'].lower()
    )
    return tool_span


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_deferral_call_deferred_v2(capfire: CaptureLogfire) -> None:
    """Test that CallDeferred on v2 marks span as ERROR with deferral attributes."""
    agent = Agent(
        TestModel(),
        output_type=[str, DeferredToolRequests],
        capabilities=[Instrumentation(settings=deprecated_instrumentation_settings(version=2))],
    )

    @agent.tool_plain
    def my_tool(x: int) -> int:
        raise CallDeferred(metadata={'task_id': 'task-123'})

    agent.run_sync('Hello')

    assert _get_tool_span(capfire) == snapshot(
        {
            'name': 'running tool',
            'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
            'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
            'start_time': 4000000000,
            'end_time': 6000000000,
            'attributes': {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'my_tool',
                'gen_ai.tool.call.id': 'pyd_ai_tool_call_id__my_tool',
                'tool_arguments': {'x': 0},
                'logfire.msg': 'running tool: my_tool',
                'logfire.json_schema': {
                    'type': 'object',
                    'properties': {
                        'tool_arguments': {'type': 'object'},
                        'tool_response': {'type': 'object'},
                        'gen_ai.tool.name': {},
                        'gen_ai.tool.call.id': {},
                    },
                },
                'gen_ai.conversation.id': IsStr(),
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'agent',
                'gen_ai.agent.call.id': IsStr(),
                'pydantic_ai.tool.deferral.name': 'CallDeferred',
                'pydantic_ai.tool.deferral.metadata': {'task_id': 'task-123'},
                'logfire.level_num': 17,
            },
            'events': [
                {
                    'name': 'exception',
                    'timestamp': 5000000000,
                    'attributes': {
                        'exception.type': 'pydantic_ai.exceptions.CallDeferred',
                        'exception.message': '',
                        'exception.stacktrace': 'pydantic_ai.exceptions.CallDeferred',
                        'exception.escaped': 'True',
                    },
                }
            ],
        }
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_deferral_approval_required_v2(capfire: CaptureLogfire) -> None:
    """Test that ApprovalRequired on v2 marks span as ERROR with deferral attributes."""
    agent = Agent(
        TestModel(),
        output_type=[str, DeferredToolRequests],
        capabilities=[Instrumentation(settings=deprecated_instrumentation_settings(version=2))],
    )

    @agent.tool_plain
    def my_tool(x: int) -> int:
        raise ApprovalRequired(metadata={'task_id': 'task-123'})

    agent.run_sync('Hello')

    assert _get_tool_span(capfire) == snapshot(
        {
            'name': 'running tool',
            'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
            'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
            'start_time': 4000000000,
            'end_time': 6000000000,
            'attributes': {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'my_tool',
                'gen_ai.tool.call.id': 'pyd_ai_tool_call_id__my_tool',
                'tool_arguments': {'x': 0},
                'logfire.msg': 'running tool: my_tool',
                'logfire.json_schema': {
                    'type': 'object',
                    'properties': {
                        'tool_arguments': {'type': 'object'},
                        'tool_response': {'type': 'object'},
                        'gen_ai.tool.name': {},
                        'gen_ai.tool.call.id': {},
                    },
                },
                'gen_ai.conversation.id': IsStr(),
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'agent',
                'gen_ai.agent.call.id': IsStr(),
                'pydantic_ai.tool.deferral.name': 'ApprovalRequired',
                'pydantic_ai.tool.deferral.metadata': {'task_id': 'task-123'},
                'logfire.level_num': 17,
            },
            'events': [
                {
                    'name': 'exception',
                    'timestamp': 5000000000,
                    'attributes': {
                        'exception.type': 'pydantic_ai.exceptions.ApprovalRequired',
                        'exception.message': '',
                        'exception.stacktrace': 'pydantic_ai.exceptions.ApprovalRequired',
                        'exception.escaped': 'True',
                    },
                }
            ],
        }
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_deferral_call_deferred_v5(capfire: CaptureLogfire) -> None:
    """Test that CallDeferred on v5 leaves span as UNSET with deferral attributes."""
    agent = Agent(
        TestModel(),
        output_type=[str, DeferredToolRequests],
        capabilities=[Instrumentation(settings=InstrumentationSettings(version=5))],
    )

    @agent.tool_plain
    def my_tool(x: int) -> int:
        raise CallDeferred(metadata={'task_id': 'task-123'})

    agent.run_sync('Hello')

    assert _get_tool_span(capfire) == snapshot(
        {
            'name': 'execute_tool my_tool',
            'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
            'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
            'start_time': 4000000000,
            'end_time': 5000000000,
            'attributes': {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'my_tool',
                'gen_ai.tool.call.id': 'pyd_ai_tool_call_id__my_tool',
                'gen_ai.tool.call.arguments': {'x': 0},
                'logfire.msg': 'running tool: my_tool',
                'logfire.json_schema': {
                    'type': 'object',
                    'properties': {
                        'gen_ai.tool.call.arguments': {'type': 'object'},
                        'gen_ai.tool.call.result': {'type': 'object'},
                        'gen_ai.tool.name': {},
                        'gen_ai.tool.call.id': {},
                    },
                },
                'gen_ai.conversation.id': IsStr(),
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'agent',
                'gen_ai.agent.call.id': IsStr(),
                'pydantic_ai.tool.deferral.name': 'CallDeferred',
                'pydantic_ai.tool.deferral.metadata': {'task_id': 'task-123'},
            },
        }
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_deferral_approval_required_v5(capfire: CaptureLogfire) -> None:
    """Test that ApprovalRequired on v5 leaves span as UNSET with deferral attributes."""
    agent = Agent(
        TestModel(),
        output_type=[str, DeferredToolRequests],
        capabilities=[Instrumentation(settings=InstrumentationSettings(version=5))],
    )

    @agent.tool_plain
    def my_tool(x: int) -> int:
        raise ApprovalRequired(metadata={'task_id': 'task-123'})

    agent.run_sync('Hello')

    assert _get_tool_span(capfire) == snapshot(
        {
            'name': 'execute_tool my_tool',
            'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
            'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
            'start_time': 4000000000,
            'end_time': 5000000000,
            'attributes': {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'my_tool',
                'gen_ai.tool.call.id': 'pyd_ai_tool_call_id__my_tool',
                'gen_ai.tool.call.arguments': {'x': 0},
                'logfire.msg': 'running tool: my_tool',
                'logfire.json_schema': {
                    'type': 'object',
                    'properties': {
                        'gen_ai.tool.call.arguments': {'type': 'object'},
                        'gen_ai.tool.call.result': {'type': 'object'},
                        'gen_ai.tool.name': {},
                        'gen_ai.tool.call.id': {},
                    },
                },
                'gen_ai.conversation.id': IsStr(),
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'agent',
                'gen_ai.agent.call.id': IsStr(),
                'pydantic_ai.tool.deferral.name': 'ApprovalRequired',
                'pydantic_ai.tool.deferral.metadata': {'task_id': 'task-123'},
            },
        }
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_deferral_no_metadata(capfire: CaptureLogfire) -> None:
    """Test that deferral without metadata doesn't set the metadata attribute."""
    agent = Agent(
        TestModel(),
        output_type=[str, DeferredToolRequests],
        capabilities=[Instrumentation(settings=InstrumentationSettings(version=5))],
    )

    @agent.tool_plain
    def my_tool(x: int) -> int:
        raise CallDeferred()

    agent.run_sync('Hello')

    tool_span = _get_tool_span(capfire)

    assert tool_span == snapshot(
        {
            'name': 'execute_tool my_tool',
            'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
            'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
            'start_time': 4000000000,
            'end_time': 5000000000,
            'attributes': {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'my_tool',
                'gen_ai.tool.call.id': 'pyd_ai_tool_call_id__my_tool',
                'gen_ai.tool.call.arguments': {'x': 0},
                'logfire.msg': 'running tool: my_tool',
                'logfire.json_schema': {
                    'type': 'object',
                    'properties': {
                        'gen_ai.tool.call.arguments': {'type': 'object'},
                        'gen_ai.tool.call.result': {'type': 'object'},
                        'gen_ai.tool.name': {},
                        'gen_ai.tool.call.id': {},
                    },
                },
                'gen_ai.conversation.id': IsStr(),
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'agent',
                'gen_ai.agent.call.id': IsStr(),
                'pydantic_ai.tool.deferral.name': 'CallDeferred',
            },
        }
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_deferral_non_serializable_metadata(capfire: CaptureLogfire) -> None:
    """Test that non-JSON-serializable metadata falls back to repr() representation."""

    class CustomObj:
        def __repr__(self) -> str:
            return '<CustomObj>'

    agent = Agent(
        TestModel(),
        output_type=[str, DeferredToolRequests],
        capabilities=[Instrumentation(settings=InstrumentationSettings(version=5))],
    )

    @agent.tool_plain
    def my_tool(x: int) -> int:
        raise CallDeferred(metadata={'obj': CustomObj()})

    agent.run_sync('Hello')

    tool_span = _get_tool_span(capfire)

    assert tool_span == snapshot(
        {
            'name': 'execute_tool my_tool',
            'context': {'trace_id': 1, 'span_id': 5, 'is_remote': False},
            'parent': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
            'start_time': 4000000000,
            'end_time': 5000000000,
            'attributes': {
                'gen_ai.operation.name': 'execute_tool',
                'gen_ai.tool.name': 'my_tool',
                'gen_ai.tool.call.id': 'pyd_ai_tool_call_id__my_tool',
                'gen_ai.tool.call.arguments': {'x': 0},
                'logfire.msg': 'running tool: my_tool',
                'logfire.json_schema': {
                    'type': 'object',
                    'properties': {
                        'gen_ai.tool.call.arguments': {'type': 'object'},
                        'gen_ai.tool.call.result': {'type': 'object'},
                        'gen_ai.tool.name': {},
                        'gen_ai.tool.call.id': {},
                    },
                },
                'gen_ai.conversation.id': IsStr(),
                'logfire.span_type': 'span',
                'gen_ai.agent.name': 'agent',
                'gen_ai.agent.call.id': IsStr(),
                'pydantic_ai.tool.deferral.name': 'CallDeferred',
                'pydantic_ai.tool.deferral.metadata': "{'obj': <CustomObj>}",
            },
        }
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_deferral_model_retry_still_errors_v5(capfire: CaptureLogfire) -> None:
    """Test that ModelRetry on v5 still records the span as an error.

    The deferral fix (CallDeferred/ApprovalRequired → UNSET on v5) must not affect
    ModelRetry, which wraps as ToolRetryError and should always be an error span.
    """
    agent = Agent(
        TestModel(),
        capabilities=[Instrumentation(settings=InstrumentationSettings(version=5))],
    )

    @agent.tool_plain
    def my_tool(x: int) -> str:
        raise ModelRetry('please try again with different input')

    with pytest.raises(UnexpectedModelBehavior):
        agent.run_sync('Hello')

    tool_span = _get_tool_span(capfire)

    # ToolRetryError should still be recorded as an error on v5 — only deferrals get UNSET
    assert tool_span['attributes'].get('logfire.level_num') == 17
    # No deferral attributes should be set — this is a retry, not a deferral
    assert 'pydantic_ai.tool.deferral.name' not in tool_span['attributes']
    assert 'pydantic_ai.tool.deferral.metadata' not in tool_span['attributes']


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_deferral_unexpected_exception_still_errors_v5(capfire: CaptureLogfire) -> None:
    """Test that unexpected exceptions on v5 still record the span as an error.

    The deferral fix must not affect general exception handling — only
    CallDeferred and ApprovalRequired get UNSET status on v5.
    """
    agent = Agent(
        TestModel(),
        capabilities=[Instrumentation(settings=InstrumentationSettings(version=5))],
    )

    @agent.tool_plain
    def my_tool(x: int) -> str:
        raise ValueError('something went wrong')

    with pytest.raises(ValueError, match='something went wrong'):
        agent.run_sync('Hello')

    tool_span = _get_tool_span(capfire)

    # ValueError path should still record error regardless of instrumentation version
    assert tool_span['attributes'].get('logfire.level_num') == 17
    # No deferral attributes should be set
    assert 'pydantic_ai.tool.deferral.name' not in tool_span['attributes']
    assert 'pydantic_ai.tool.deferral.metadata' not in tool_span['attributes']


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.anyio
async def test_agent_description(capfire: CaptureLogfire) -> None:
    agent = Agent(
        model=TestModel(),
        name='my_agent',
        description='An agent that greets users',
        capabilities=[Instrumentation(settings=InstrumentationSettings())],
    )
    assert agent.description == 'An agent that greets users'

    await agent.run('Hello')

    spans = capfire.exporter.exported_spans_as_dict()
    agent_run_span = next(s for s in spans if s['name'] == 'invoke_agent my_agent')
    assert agent_run_span['attributes']['gen_ai.agent.description'] == 'An agent that greets users'

    agent.description = 'Updated description'
    assert agent.description == 'Updated description'


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
@pytest.mark.anyio
async def test_agent_description_absent_when_none(capfire: CaptureLogfire) -> None:
    agent = Agent(
        model=TestModel(), name='my_agent', capabilities=[Instrumentation(settings=InstrumentationSettings())]
    )
    assert agent.description is None

    await agent.run('Hello')

    spans = capfire.exporter.exported_spans_as_dict()
    agent_run_span = next(s for s in spans if s['name'] == 'invoke_agent my_agent')
    assert 'gen_ai.agent.description' not in agent_run_span['attributes']


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_instrumentation_capability_with_model_settings(
    get_logfire_summary: Callable[[], LogfireSummary],
) -> None:
    """Test that Instrumentation capability correctly records model settings like temperature."""
    agent = Agent(
        model=TestModel(),
        model_settings={'temperature': 0.5, 'max_tokens': 100},
        capabilities=[Instrumentation(settings=InstrumentationSettings())],
    )
    agent.run_sync('Hello')

    summary = get_logfire_summary()
    chat_attrs = summary.attributes[1]
    assert chat_attrs['gen_ai.request.temperature'] == 0.5
    assert chat_attrs['gen_ai.request.max_tokens'] == 100


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_instrumentation_capability_serialization() -> None:
    """`Instrumentation` is constructible from a YAML/JSON spec via the serializable subset of
    `InstrumentationSettings` kwargs; non-serializable OTel providers fall back to the globals."""
    from pydantic_ai.capabilities.instrumentation import Instrumentation
    from pydantic_ai.models.instrumented import InstrumentationSettings

    assert Instrumentation.get_serialization_name() == 'Instrumentation'

    with pytest.warns(
        PydanticAIDeprecationWarning, match=r'Instrumentation format versions 2, 3, and 4 are deprecated'
    ):
        cap = Instrumentation.from_spec(version=2, include_content=False)
    assert isinstance(cap, Instrumentation)
    assert isinstance(cap.settings, InstrumentationSettings)
    assert cap.settings.version == 2
    assert cap.settings.include_content is False

    # Empty kwargs form: `Instrumentation: {}` in YAML.
    cap_default = Instrumentation.from_spec()
    assert cap_default.settings.version == InstrumentationSettings().version


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_instrumentation_capability_explicit(
    get_logfire_summary: Callable[[], LogfireSummary],
) -> None:
    """Test using Instrumentation as an explicit capability (not via instrument=True)."""
    from pydantic_ai.capabilities.instrumentation import Instrumentation

    instrumentation = Instrumentation(settings=InstrumentationSettings())
    agent = Agent(model=TestModel(), capabilities=[instrumentation])

    result = agent.run_sync('Hello')
    assert result.output == snapshot('success (no tool calls)')

    summary = get_logfire_summary()
    assert summary.traces == snapshot(
        [
            {
                'id': 0,
                'name': 'invoke_agent agent',
                'message': 'agent run',
                'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
            }
        ]
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_instrument_all_skipped_when_capability_already_present(
    get_logfire_summary: Callable[[], LogfireSummary],
) -> None:
    """`Agent.instrument_all(...)` must not stack a second `Instrumentation` capability when the
    user already added one via `capabilities=[...]` — otherwise spans would be emitted twice.

    Guarded by the `has_capability_type(..., InstrumentationCap)` check in `Agent.iter`.
    """
    from pydantic_ai.capabilities.instrumentation import Instrumentation

    Agent.instrument_all(True)
    try:
        agent = Agent(model=TestModel(), capabilities=[Instrumentation(settings=InstrumentationSettings())])
        result = agent.run_sync('Hello')
        assert result.output == snapshot('success (no tool calls)')

        summary = get_logfire_summary()
        # One `agent run` span, not two.
        assert summary.traces == snapshot(
            [
                {
                    'id': 0,
                    'name': 'invoke_agent agent',
                    'message': 'agent run',
                    'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
                }
            ]
        )
    finally:
        Agent.instrument_all(False)


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_agent_with_user_provided_instrumented_model(
    get_logfire_summary: Callable[[], LogfireSummary],
) -> None:
    """If a user explicitly hands an `InstrumentedModel` to `Agent(model=...)`, its
    settings drive the agent's instrumentation — the model is unwrapped and the
    `Instrumentation` capability emits the spans.
    """
    from pydantic_ai.models.instrumented import InstrumentedModel

    settings = InstrumentationSettings()
    agent = Agent(model=InstrumentedModel(TestModel(), settings))

    result = agent.run_sync('Hello')
    assert result.output == snapshot('success (no tool calls)')

    summary = get_logfire_summary()
    assert summary.traces == snapshot(
        [
            {
                'id': 0,
                'name': 'invoke_agent agent',
                'message': 'agent run',
                'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
            }
        ]
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_agent_instrument_setter(
    get_logfire_summary: Callable[[], LogfireSummary],
) -> None:
    """`agent.instrument = settings` configures instrumentation post-construction.

    This is the path `logfire.instrument_pydantic_ai(agent)` uses on its `Agent` branch.
    """
    agent = Agent(model=TestModel())
    assert agent.instrument is None
    settings = InstrumentationSettings()
    agent.instrument = settings
    assert agent.instrument is settings

    result = agent.run_sync('Hello')
    assert result.output == snapshot('success (no tool calls)')

    summary = get_logfire_summary()
    assert summary.traces == snapshot(
        [
            {
                'id': 0,
                'name': 'invoke_agent agent',
                'message': 'agent run',
                'children': [{'id': 1, 'name': 'chat test', 'message': 'chat test'}],
            }
        ]
    )


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_instrumentation_capability_template_description(
    capfire: CaptureLogfire,
) -> None:
    """Test that TemplateStr descriptions are rendered in agent run spans."""
    # `TemplateStr` rendering requires the `[spec]` extra (pydantic-handlebars); not in default v2 deps.
    pytest.importorskip('pydantic_handlebars')
    from dataclasses import dataclass

    from pydantic_ai._template import TemplateStr
    from pydantic_ai.capabilities.instrumentation import Instrumentation

    @dataclass
    class MyDeps:
        name: str

    instrumentation = Instrumentation(settings=InstrumentationSettings())
    agent = Agent(
        model=TestModel(),
        capabilities=[instrumentation],
        description=TemplateStr('Agent for {{name}}'),
        deps_type=MyDeps,
    )

    result = agent.run_sync('Hello', deps=MyDeps(name='testing'))
    assert result.output == snapshot('success (no tool calls)')

    spans = strip_logfire_metrics(capfire.exporter.exported_spans_as_dict(parse_json_attributes=True))
    agent_span = spans[-1]  # outermost span is the agent run
    assert agent_span['attributes']['gen_ai.agent.description'] == snapshot('Agent for testing')


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
async def test_instrumentation_capability_with_noop_tracer() -> None:
    """When the configured tracer provider is a no-op, model-request spans skip
    attribute population entirely. Regression coverage for the non-recording
    branch in `Instrumentation.wrap_model_request`."""
    from opentelemetry.trace import NoOpTracerProvider

    agent = Agent(
        model=TestModel(),
        capabilities=[Instrumentation(settings=InstrumentationSettings(tracer_provider=NoOpTracerProvider()))],
    )
    result = await agent.run('hello')
    assert result.output == snapshot('success (no tool calls)')


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
async def test_instrument_combines_with_outermost_and_innermost_capabilities() -> None:
    """Auto-prepending `Instrumentation` must not wrap a `CombinedCapability` that
    already contains both an outermost and innermost cap — the wrap would force
    `_effective_ordering` to merge those positions and raise "Conflicting positions".
    """
    from pydantic_ai.capabilities import AbstractCapability, CapabilityOrdering

    class OutermostCap(AbstractCapability[Any]):
        def get_ordering(self) -> CapabilityOrdering:
            return CapabilityOrdering(position='outermost')

        async def before_run(self, ctx: RunContext[Any]) -> None:
            pass

    class InnermostCap(AbstractCapability[Any]):
        def get_ordering(self) -> CapabilityOrdering:
            return CapabilityOrdering(position='innermost')

        async def before_run(self, ctx: RunContext[Any]) -> None:
            pass

    agent = Agent(
        model=TestModel(),
        capabilities=[OutermostCap(), InnermostCap(), Instrumentation(settings=InstrumentationSettings())],
    )
    result = await agent.run('hello')
    assert result.output == snapshot('success (no tool calls)')


@pytest.mark.skipif(not logfire_installed, reason='logfire not installed')
def test_output_function_call_deferred_recorded_as_error(
    get_logfire_summary: Callable[[], LogfireSummary],
) -> None:
    """An output function raising `CallDeferred` is recorded as a regular error on the
    `wrap_output_process` span — that hook reserves the deferral-attribute path for
    real tool executions (`wrap_tool_execute`)."""

    def defer_text(text: str) -> str:
        from pydantic_ai.exceptions import CallDeferred

        raise CallDeferred()

    def call_text_response(_: list[ModelMessage], info: AgentInfo) -> ModelResponse:
        return ModelResponse(parts=[TextPart(content='hi')])

    from pydantic_ai.exceptions import CallDeferred

    my_agent = Agent(
        model=FunctionModel(call_text_response), capabilities=[Instrumentation(settings=InstrumentationSettings())]
    )
    with pytest.raises(CallDeferred):
        my_agent.run_sync('anything', output_type=TextOutput(defer_text))

    summary = get_logfire_summary()
    [span_attrs] = [attrs for attrs in summary.attributes.values() if attrs.get('gen_ai.tool.name') == 'defer_text']
    # The span was recorded with ERROR status — the standard exception path,
    # not the deferral-attribute path that `wrap_tool_execute` uses.
    assert span_attrs.get('logfire.level_num', 0) >= 17  # error level
    assert 'pydantic_ai.tool.deferral.name' not in span_attrs