test_pydantic_ai.py

import asyncio
import json
import pytest
from unittest.mock import MagicMock

from typing import Annotated
from pydantic import Field

import sentry_sdk
from sentry_sdk._types import BLOB_DATA_SUBSTITUTE
from sentry_sdk.consts import SPANDATA
from sentry_sdk.integrations.pydantic_ai import PydanticAIIntegration
from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages
from sentry_sdk.integrations.pydantic_ai.spans.utils import _set_usage_data
from pydantic_ai import Agent
from pydantic_ai.messages import BinaryContent, ImageUrl, UserPromptPart
from pydantic_ai.usage import RequestUsage
from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior


@pytest.fixture
def test_agent():
    """Create a test agent with model settings."""
    return Agent(
        "test",
        name="test_agent",
        system_prompt="You are a helpful test assistant.",
    )


@pytest.fixture
def test_agent_with_settings():
    """Create a test agent with explicit model settings."""
    from pydantic_ai import ModelSettings

    return Agent(
        "test",
        name="test_agent_settings",
        system_prompt="You are a test assistant with settings.",
        model_settings=ModelSettings(
            temperature=0.7,
            max_tokens=100,
            top_p=0.9,
        ),
    )


@pytest.mark.asyncio
async def test_agent_run_async(sentry_init, capture_events, test_agent):
    """
    Test that the integration creates spans for async agent runs.
    """
    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    result = await test_agent.run("Test input")

    assert result is not None
    assert result.output is not None

    (transaction,) = events
    spans = transaction["spans"]

    # Verify transaction (the transaction IS the invoke_agent span)
    assert transaction["transaction"] == "invoke_agent test_agent"
    assert transaction["contexts"]["trace"]["origin"] == "auto.ai.pydantic_ai"

    # The transaction itself should have invoke_agent data
    assert transaction["contexts"]["trace"]["op"] == "gen_ai.invoke_agent"

    # Find child span types (invoke_agent is the transaction, not a child span)
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    assert len(chat_spans) == 1

    # Check chat span
    chat_span = chat_spans[0]
    assert "chat" in chat_span["description"]
    assert chat_span["data"]["gen_ai.operation.name"] == "chat"
    assert chat_span["data"]["gen_ai.response.streaming"] is False
    assert "gen_ai.request.messages" in chat_span["data"]
    assert "gen_ai.usage.input_tokens" in chat_span["data"]
    assert "gen_ai.usage.output_tokens" in chat_span["data"]


@pytest.mark.asyncio
async def test_agent_run_async_usage_data(sentry_init, capture_events, test_agent):
    """
    Test that the invoke_agent span includes token usage and model data.
    """
    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    result = await test_agent.run("Test input")

    assert result is not None
    assert result.output is not None

    (transaction,) = events

    # Verify transaction (the transaction IS the invoke_agent span)
    assert transaction["transaction"] == "invoke_agent test_agent"

    # The invoke_agent span should have token usage data
    trace_data = transaction["contexts"]["trace"].get("data", {})
    assert "gen_ai.usage.input_tokens" in trace_data, (
        "Missing input_tokens on invoke_agent span"
    )
    assert "gen_ai.usage.output_tokens" in trace_data, (
        "Missing output_tokens on invoke_agent span"
    )
    assert "gen_ai.usage.total_tokens" in trace_data, (
        "Missing total_tokens on invoke_agent span"
    )
    assert "gen_ai.response.model" in trace_data, (
        "Missing response.model on invoke_agent span"
    )

    # Verify the values are reasonable
    assert trace_data["gen_ai.usage.input_tokens"] > 0
    assert trace_data["gen_ai.usage.output_tokens"] > 0
    assert trace_data["gen_ai.usage.total_tokens"] > 0
    assert trace_data["gen_ai.response.model"] == "test"  # Test model name


def test_agent_run_sync(sentry_init, capture_events, test_agent):
    """
    Test that the integration creates spans for sync agent runs.
    """
    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    result = test_agent.run_sync("Test input")

    assert result is not None
    assert result.output is not None

    (transaction,) = events
    spans = transaction["spans"]

    # Verify transaction
    assert transaction["transaction"] == "invoke_agent test_agent"
    assert transaction["contexts"]["trace"]["origin"] == "auto.ai.pydantic_ai"

    # Find span types
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    assert len(chat_spans) == 1

    # Verify streaming flag is False for sync
    for chat_span in chat_spans:
        assert chat_span["data"]["gen_ai.response.streaming"] is False


@pytest.mark.asyncio
async def test_agent_run_stream(sentry_init, capture_events, test_agent):
    """
    Test that the integration creates spans for streaming agent runs.
    """
    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    async with test_agent.run_stream("Test input") as result:
        # Consume the stream
        async for _ in result.stream_output():
            pass

    (transaction,) = events
    spans = transaction["spans"]

    # Verify transaction
    assert transaction["transaction"] == "invoke_agent test_agent"
    assert transaction["contexts"]["trace"]["origin"] == "auto.ai.pydantic_ai"

    # Find chat spans
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    assert len(chat_spans) == 1

    # Verify streaming flag is True for streaming
    for chat_span in chat_spans:
        assert chat_span["data"]["gen_ai.response.streaming"] is True
        assert "gen_ai.request.messages" in chat_span["data"]
        assert "gen_ai.usage.input_tokens" in chat_span["data"]
        # Streaming responses should still have output data
        assert (
            "gen_ai.response.text" in chat_span["data"]
            or "gen_ai.response.model" in chat_span["data"]
        )


@pytest.mark.asyncio
async def test_agent_run_stream_events(sentry_init, capture_events, test_agent):
    """
    Test that run_stream_events creates spans (it uses run internally, so non-streaming).
    """
    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    # Consume all events
    async for _ in test_agent.run_stream_events("Test input"):
        pass

    (transaction,) = events

    # Verify transaction
    assert transaction["transaction"] == "invoke_agent test_agent"

    # Find chat spans
    spans = transaction["spans"]
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    assert len(chat_spans) == 1

    for chat_span in chat_spans:
        assert chat_span["data"]["gen_ai.response.streaming"] is True


@pytest.mark.asyncio
async def test_agent_with_tools(sentry_init, capture_events, test_agent):
    """
    Test that tool execution creates execute_tool spans.
    """

    @test_agent.tool_plain
    def add_numbers(a: int, b: int) -> int:
        """Add two numbers together."""
        return a + b

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    result = await test_agent.run("What is 5 + 3?")

    assert result is not None

    (transaction,) = events
    spans = transaction["spans"]

    # Find child span types (invoke_agent is the transaction, not a child span)
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"]

    # Should have tool spans
    assert len(tool_spans) == 1

    # Check tool span
    tool_span = tool_spans[0]
    assert "execute_tool" in tool_span["description"]
    assert tool_span["data"]["gen_ai.operation.name"] == "execute_tool"
    assert tool_span["data"]["gen_ai.tool.type"] == "function"
    assert tool_span["data"]["gen_ai.tool.name"] == "add_numbers"
    assert "gen_ai.tool.input" in tool_span["data"]
    assert "gen_ai.tool.output" in tool_span["data"]

    # Check chat spans have available_tools
    for chat_span in chat_spans:
        assert "gen_ai.request.available_tools" in chat_span["data"]
        available_tools_str = chat_span["data"]["gen_ai.request.available_tools"]
        # Available tools is serialized as a string
        assert "add_numbers" in available_tools_str


@pytest.mark.parametrize(
    "handled_tool_call_exceptions",
    [False, True],
)
@pytest.mark.asyncio
async def test_agent_with_tool_model_retry(
    sentry_init, capture_events, test_agent, handled_tool_call_exceptions
):
    """
    Test that a handled exception is captured when a tool raises ModelRetry.
    """

    retries = 0

    @test_agent.tool_plain
    def add_numbers(a: int, b: int) -> float:
        """Add two numbers together, but raises an exception on the first attempt."""
        nonlocal retries
        if retries == 0:
            retries += 1
            raise ModelRetry(message="Try again with the same arguments.")
        return a + b

    sentry_init(
        integrations=[
            PydanticAIIntegration(
                handled_tool_call_exceptions=handled_tool_call_exceptions
            )
        ],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    result = await test_agent.run("What is 5 + 3?")

    assert result is not None

    if handled_tool_call_exceptions:
        (error, transaction) = events
    else:
        (transaction,) = events
    spans = transaction["spans"]

    if handled_tool_call_exceptions:
        assert error["level"] == "error"
        assert error["exception"]["values"][0]["mechanism"]["handled"]

    # Find child span types (invoke_agent is the transaction, not a child span)
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"]

    # Should have tool spans
    assert len(tool_spans) >= 1

    # Check tool spans
    model_retry_tool_span = tool_spans[0]
    assert "execute_tool" in model_retry_tool_span["description"]
    assert model_retry_tool_span["data"]["gen_ai.operation.name"] == "execute_tool"
    assert model_retry_tool_span["data"]["gen_ai.tool.type"] == "function"
    assert model_retry_tool_span["data"]["gen_ai.tool.name"] == "add_numbers"
    assert "gen_ai.tool.input" in model_retry_tool_span["data"]

    tool_span = tool_spans[1]
    assert "execute_tool" in tool_span["description"]
    assert tool_span["data"]["gen_ai.operation.name"] == "execute_tool"
    assert tool_span["data"]["gen_ai.tool.type"] == "function"
    assert tool_span["data"]["gen_ai.tool.name"] == "add_numbers"
    assert "gen_ai.tool.input" in tool_span["data"]
    assert "gen_ai.tool.output" in tool_span["data"]

    # Check chat spans have available_tools
    for chat_span in chat_spans:
        assert "gen_ai.request.available_tools" in chat_span["data"]
        available_tools_str = chat_span["data"]["gen_ai.request.available_tools"]
        # Available tools is serialized as a string
        assert "add_numbers" in available_tools_str


@pytest.mark.parametrize(
    "handled_tool_call_exceptions",
    [False, True],
)
@pytest.mark.asyncio
async def test_agent_with_tool_validation_error(
    sentry_init, capture_events, test_agent, handled_tool_call_exceptions
):
    """
    Test that a handled exception is captured when a tool has unsatisfiable constraints.
    """

    @test_agent.tool_plain
    def add_numbers(a: Annotated[int, Field(gt=0, lt=0)], b: int) -> int:
        """Add two numbers together."""
        return a + b

    sentry_init(
        integrations=[
            PydanticAIIntegration(
                handled_tool_call_exceptions=handled_tool_call_exceptions
            )
        ],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    result = None
    with pytest.raises(UnexpectedModelBehavior):
        result = await test_agent.run("What is 5 + 3?")

    assert result is None

    if handled_tool_call_exceptions:
        (error, model_behaviour_error, transaction) = events
    else:
        (
            model_behaviour_error,
            transaction,
        ) = events
    spans = transaction["spans"]

    if handled_tool_call_exceptions:
        assert error["level"] == "error"
        assert error["exception"]["values"][0]["mechanism"]["handled"]

    # Find child span types (invoke_agent is the transaction, not a child span)
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"]

    # Should have tool spans
    assert len(tool_spans) >= 1

    # Check tool spans
    model_retry_tool_span = tool_spans[0]
    assert "execute_tool" in model_retry_tool_span["description"]
    assert model_retry_tool_span["data"]["gen_ai.operation.name"] == "execute_tool"
    assert model_retry_tool_span["data"]["gen_ai.tool.type"] == "function"
    assert model_retry_tool_span["data"]["gen_ai.tool.name"] == "add_numbers"
    assert "gen_ai.tool.input" in model_retry_tool_span["data"]

    # Check chat spans have available_tools
    for chat_span in chat_spans:
        assert "gen_ai.request.available_tools" in chat_span["data"]
        available_tools_str = chat_span["data"]["gen_ai.request.available_tools"]
        # Available tools is serialized as a string
        assert "add_numbers" in available_tools_str


@pytest.mark.asyncio
async def test_agent_with_tools_streaming(sentry_init, capture_events, test_agent):
    """
    Test that tool execution works correctly with streaming.
    """

    @test_agent.tool_plain
    def multiply(a: int, b: int) -> int:
        """Multiply two numbers."""
        return a * b

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    async with test_agent.run_stream("What is 7 times 8?") as result:
        async for _ in result.stream_output():
            pass

    (transaction,) = events
    spans = transaction["spans"]

    # Find span types
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"]

    # Should have tool spans
    assert len(tool_spans) >= 1

    # Verify streaming flag is True
    for chat_span in chat_spans:
        assert chat_span["data"]["gen_ai.response.streaming"] is True

    # Check tool span
    tool_span = tool_spans[0]
    assert tool_span["data"]["gen_ai.tool.name"] == "multiply"
    assert "gen_ai.tool.input" in tool_span["data"]
    assert "gen_ai.tool.output" in tool_span["data"]


@pytest.mark.asyncio
async def test_model_settings(sentry_init, capture_events, test_agent_with_settings):
    """
    Test that model settings are captured in spans.
    """
    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    events = capture_events()

    await test_agent_with_settings.run("Test input")

    (transaction,) = events
    spans = transaction["spans"]

    # Find chat span
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    assert len(chat_spans) == 1

    chat_span = chat_spans[0]
    # Check that model settings are captured
    assert chat_span["data"].get("gen_ai.request.temperature") == 0.7
    assert chat_span["data"].get("gen_ai.request.max_tokens") == 100
    assert chat_span["data"].get("gen_ai.request.top_p") == 0.9


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "send_default_pii, include_prompts",
    [
        (True, True),
        (True, False),
        (False, True),
        (False, False),
    ],
)
async def test_system_prompt_attribute(
    sentry_init, capture_events, send_default_pii, include_prompts
):
    """
    Test that system prompts are included as the first message.
    """
    agent = Agent(
        "test",
        name="test_system",
        system_prompt="You are a helpful assistant specialized in testing.",
    )

    sentry_init(
        integrations=[PydanticAIIntegration(include_prompts=include_prompts)],
        traces_sample_rate=1.0,
        send_default_pii=send_default_pii,
    )

    events = capture_events()

    await agent.run("Hello")

    (transaction,) = events
    spans = transaction["spans"]

    # The transaction IS the invoke_agent span, check for messages in chat spans instead
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    assert len(chat_spans) == 1

    chat_span = chat_spans[0]

    if send_default_pii and include_prompts:
        system_instructions = chat_span["data"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]
        assert json.loads(system_instructions) == [
            {
                "type": "text",
                "content": "You are a helpful assistant specialized in testing.",
            }
        ]
    else:
        assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in chat_span["data"]


@pytest.mark.asyncio
async def test_error_handling(sentry_init, capture_events):
    """
    Test error handling in agent execution.
    """
    # Use a simpler test that doesn't cause tool failures
    # as pydantic-ai has complex error handling for tool errors
    agent = Agent(
        "test",
        name="test_error",
    )

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    events = capture_events()

    # Simple run that should succeed
    await agent.run("Hello")

    # At minimum, we should have a transaction
    assert len(events) == 1
    transaction = [e for e in events if e.get("type") == "transaction"][0]
    assert transaction["transaction"] == "invoke_agent test_error"
    # Transaction should complete successfully (status key may not exist if no error)
    trace_status = transaction["contexts"]["trace"].get("status")
    assert trace_status != "error"  # Could be None or some other status


@pytest.mark.asyncio
async def test_without_pii(sentry_init, capture_events, test_agent):
    """
    Test that PII is not captured when send_default_pii is False.
    """
    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=False,
    )

    events = capture_events()

    await test_agent.run("Sensitive input")

    (transaction,) = events
    spans = transaction["spans"]

    # Find child spans (invoke_agent is the transaction, not a child span)
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]

    # Verify that messages and response text are not captured
    for span in chat_spans:
        assert "gen_ai.request.messages" not in span["data"]
        assert "gen_ai.response.text" not in span["data"]


@pytest.mark.asyncio
async def test_without_pii_tools(sentry_init, capture_events, test_agent):
    """
    Test that tool input/output are not captured when send_default_pii is False.
    """

    @test_agent.tool_plain
    def sensitive_tool(data: str) -> str:
        """A tool with sensitive data."""
        return f"Processed: {data}"

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=False,
    )

    events = capture_events()

    await test_agent.run("Use sensitive tool with private data")

    (transaction,) = events
    spans = transaction["spans"]

    # Find tool spans
    tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"]

    # If tool was executed, verify input/output are not captured
    for tool_span in tool_spans:
        assert "gen_ai.tool.input" not in tool_span["data"]
        assert "gen_ai.tool.output" not in tool_span["data"]


@pytest.mark.asyncio
async def test_multiple_agents_concurrent(sentry_init, capture_events, test_agent):
    """
    Test that multiple agents can run concurrently without interfering.
    """
    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    events = capture_events()

    async def run_agent(input_text):
        return await test_agent.run(input_text)

    # Run 3 agents concurrently
    results = await asyncio.gather(*[run_agent(f"Input {i}") for i in range(3)])

    assert len(results) == 3
    assert len(events) == 3

    # Verify each transaction is separate
    for i, transaction in enumerate(events):
        assert transaction["type"] == "transaction"
        assert transaction["transaction"] == "invoke_agent test_agent"
        # Each should have its own spans
        assert len(transaction["spans"]) == 1


@pytest.mark.asyncio
async def test_message_history(sentry_init, capture_events):
    """
    Test that full conversation history is captured in chat spans.
    """
    agent = Agent(
        "test",
        name="test_history",
    )

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    # First message
    await agent.run("Hello, I'm Alice")

    # Second message with history
    from pydantic_ai import messages

    history = [
        messages.ModelRequest(
            parts=[messages.UserPromptPart(content="Hello, I'm Alice")]
        ),
        messages.ModelResponse(
            parts=[messages.TextPart(content="Hello Alice! How can I help you?")],
            model_name="test",
        ),
    ]

    await agent.run("What is my name?", message_history=history)

    # We should have 2 transactions
    assert len(events) == 2

    # Check the second transaction has the full history
    second_transaction = events[1]
    spans = second_transaction["spans"]
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]

    if chat_spans:
        chat_span = chat_spans[0]
        if "gen_ai.request.messages" in chat_span["data"]:
            messages_data = chat_span["data"]["gen_ai.request.messages"]
            # Should have multiple messages including history
            assert len(messages_data) > 1


@pytest.mark.asyncio
async def test_gen_ai_system(sentry_init, capture_events, test_agent):
    """
    Test that gen_ai.system is set from the model.
    """
    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    events = capture_events()

    await test_agent.run("Test input")

    (transaction,) = events
    spans = transaction["spans"]

    # Find chat span
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    assert len(chat_spans) == 1

    chat_span = chat_spans[0]
    # gen_ai.system should be set from the model (TestModel -> 'test')
    assert "gen_ai.system" in chat_span["data"]
    assert chat_span["data"]["gen_ai.system"] == "test"


@pytest.mark.asyncio
async def test_include_prompts_false(sentry_init, capture_events, test_agent):
    """
    Test that prompts are not captured when include_prompts=False.
    """
    sentry_init(
        integrations=[PydanticAIIntegration(include_prompts=False)],
        traces_sample_rate=1.0,
        send_default_pii=True,  # Even with PII enabled, prompts should not be captured
    )

    events = capture_events()

    await test_agent.run("Sensitive prompt")

    (transaction,) = events
    spans = transaction["spans"]

    # Find child spans (invoke_agent is the transaction, not a child span)
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]

    # Verify that messages and response text are not captured
    for span in chat_spans:
        assert "gen_ai.request.messages" not in span["data"]
        assert "gen_ai.response.text" not in span["data"]


@pytest.mark.asyncio
async def test_include_prompts_true(sentry_init, capture_events, test_agent):
    """
    Test that prompts are captured when include_prompts=True (default).
    """
    sentry_init(
        integrations=[PydanticAIIntegration(include_prompts=True)],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    await test_agent.run("Test prompt")

    (transaction,) = events
    spans = transaction["spans"]

    # Find child spans (invoke_agent is the transaction, not a child span)
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]

    # Verify that messages are captured in chat spans
    assert len(chat_spans) == 1
    for chat_span in chat_spans:
        assert "gen_ai.request.messages" in chat_span["data"]


@pytest.mark.asyncio
async def test_include_prompts_false_with_tools(
    sentry_init, capture_events, test_agent
):
    """
    Test that tool input/output are not captured when include_prompts=False.
    """

    @test_agent.tool_plain
    def test_tool(value: int) -> int:
        """A test tool."""
        return value * 2

    sentry_init(
        integrations=[PydanticAIIntegration(include_prompts=False)],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    await test_agent.run("Use the test tool with value 5")

    (transaction,) = events
    spans = transaction["spans"]

    # Find tool spans
    tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"]

    # If tool was executed, verify input/output are not captured
    for tool_span in tool_spans:
        assert "gen_ai.tool.input" not in tool_span["data"]
        assert "gen_ai.tool.output" not in tool_span["data"]


@pytest.mark.asyncio
async def test_include_prompts_requires_pii(sentry_init, capture_events, test_agent):
    """
    Test that include_prompts requires send_default_pii=True.
    """
    sentry_init(
        integrations=[PydanticAIIntegration(include_prompts=True)],
        traces_sample_rate=1.0,
        send_default_pii=False,  # PII disabled
    )

    events = capture_events()

    await test_agent.run("Test prompt")

    (transaction,) = events
    spans = transaction["spans"]

    # Find child spans (invoke_agent is the transaction, not a child span)
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]

    # Even with include_prompts=True, if PII is disabled, messages should not be captured
    for span in chat_spans:
        assert "gen_ai.request.messages" not in span["data"]
        assert "gen_ai.response.text" not in span["data"]


@pytest.mark.asyncio
async def test_mcp_tool_execution_spans(sentry_init, capture_events):
    """
    Test that MCP (Model Context Protocol) tool calls create execute_tool spans.

    Tests MCP tools accessed through CombinedToolset, which is how they're typically
    used in practice (when an agent combines regular functions with MCP servers).
    """
    pytest.importorskip("mcp")

    from unittest.mock import MagicMock
    from pydantic_ai.mcp import MCPServerStdio
    from pydantic_ai import Agent
    from pydantic_ai.toolsets.combined import CombinedToolset
    import sentry_sdk

    # Create mock MCP server
    mock_server = MCPServerStdio(
        command="python",
        args=["-m", "test_server"],
    )

    # Mock the server's internal methods
    mock_server._client = MagicMock()
    mock_server._is_initialized = True
    mock_server._server_info = MagicMock()

    # Mock tool call response
    async def mock_send_request(request, response_type):
        from mcp.types import CallToolResult, TextContent

        return CallToolResult(
            content=[TextContent(type="text", text="MCP tool executed successfully")],
            isError=False,
        )

    mock_server._client.send_request = mock_send_request

    # Mock context manager methods
    async def mock_aenter():
        return mock_server

    async def mock_aexit(*args):
        pass

    mock_server.__aenter__ = mock_aenter
    mock_server.__aexit__ = mock_aexit

    # Mock _map_tool_result_part
    async def mock_map_tool_result_part(part):
        return part.text if hasattr(part, "text") else str(part)

    mock_server._map_tool_result_part = mock_map_tool_result_part

    # Create a CombinedToolset with the MCP server
    # This simulates how MCP servers are typically used in practice
    from pydantic_ai.toolsets.function import FunctionToolset

    function_toolset = FunctionToolset()
    combined = CombinedToolset([function_toolset, mock_server])

    # Create agent
    agent = Agent(
        "test",
        name="test_mcp_agent",
    )

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    # Simulate MCP tool execution within a transaction through CombinedToolset
    with sentry_sdk.start_transaction(
        op="ai.run", name="invoke_agent test_mcp_agent"
    ) as transaction:
        # Set up the agent context
        scope = sentry_sdk.get_current_scope()
        scope._contexts["pydantic_ai_agent"] = {
            "_agent": agent,
        }

        # Create a mock tool that simulates an MCP tool from CombinedToolset
        from pydantic_ai._run_context import RunContext
        from pydantic_ai.result import RunUsage
        from pydantic_ai.models.test import TestModel
        from pydantic_ai.toolsets.combined import _CombinedToolsetTool

        ctx = RunContext(
            deps=None,
            model=TestModel(),
            usage=RunUsage(),
            retry=0,
            tool_name="test_mcp_tool",
        )

        tool_name = "test_mcp_tool"

        # Create a tool that points to the MCP server
        # This simulates how CombinedToolset wraps tools from different sources
        tool = _CombinedToolsetTool(
            toolset=combined,
            tool_def=MagicMock(name=tool_name),
            max_retries=0,
            args_validator=MagicMock(),
            source_toolset=mock_server,
            source_tool=MagicMock(),
        )

        try:
            await combined.call_tool(tool_name, {"query": "test"}, ctx, tool)
        except Exception:
            # MCP tool might raise if not fully mocked, that's okay
            pass

    events_list = events
    if len(events_list) == 0:
        pytest.skip("No events captured, MCP test setup incomplete")

    (transaction,) = events_list
    transaction["spans"]

    # Note: This test manually calls combined.call_tool which doesn't go through
    # ToolManager._call_tool (which is what the integration patches).
    # In real-world usage, MCP tools are called through agent.run() which uses ToolManager.
    # This synthetic test setup doesn't trigger the integration's tool patches.
    # We skip this test as it doesn't represent actual usage patterns.
    pytest.skip(
        "MCP test needs to be rewritten to use agent.run() instead of manually calling toolset methods"
    )


@pytest.mark.asyncio
async def test_context_cleanup_after_run(sentry_init, test_agent):
    """
    Test that the pydantic_ai_agent context is properly cleaned up after agent execution.
    """
    import sentry_sdk

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    # Verify context is not set before run
    scope = sentry_sdk.get_current_scope()
    assert "pydantic_ai_agent" not in scope._contexts

    # Run the agent
    await test_agent.run("Test input")

    # Verify context is cleaned up after run
    assert "pydantic_ai_agent" not in scope._contexts


def test_context_cleanup_after_run_sync(sentry_init, test_agent):
    """
    Test that the pydantic_ai_agent context is properly cleaned up after sync agent execution.
    """
    import sentry_sdk

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    # Verify context is not set before run
    scope = sentry_sdk.get_current_scope()
    assert "pydantic_ai_agent" not in scope._contexts

    # Run the agent synchronously
    test_agent.run_sync("Test input")

    # Verify context is cleaned up after run
    assert "pydantic_ai_agent" not in scope._contexts


@pytest.mark.asyncio
async def test_context_cleanup_after_streaming(sentry_init, test_agent):
    """
    Test that the pydantic_ai_agent context is properly cleaned up after streaming execution.
    """
    import sentry_sdk

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    # Verify context is not set before run
    scope = sentry_sdk.get_current_scope()
    assert "pydantic_ai_agent" not in scope._contexts

    # Run the agent with streaming
    async with test_agent.run_stream("Test input") as result:
        async for _ in result.stream_output():
            pass

    # Verify context is cleaned up after streaming completes
    assert "pydantic_ai_agent" not in scope._contexts


@pytest.mark.asyncio
async def test_context_cleanup_on_error(sentry_init, test_agent):
    """
    Test that the pydantic_ai_agent context is cleaned up even when an error occurs.
    """
    import sentry_sdk

    # Create an agent with a tool that raises an error
    @test_agent.tool_plain
    def failing_tool() -> str:
        """A tool that always fails."""
        raise ValueError("Tool error")

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    # Verify context is not set before run
    scope = sentry_sdk.get_current_scope()
    assert "pydantic_ai_agent" not in scope._contexts

    # Run the agent - this may or may not raise depending on pydantic-ai's error handling
    try:
        await test_agent.run("Use the failing tool")
    except Exception:
        pass

    # Verify context is cleaned up even if there was an error
    assert "pydantic_ai_agent" not in scope._contexts


@pytest.mark.asyncio
async def test_context_isolation_concurrent_agents(sentry_init, test_agent):
    """
    Test that concurrent agent executions maintain isolated contexts.
    """
    import sentry_sdk

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    # Create a second agent
    agent2 = Agent(
        "test",
        name="test_agent_2",
        system_prompt="Second test agent.",
    )

    async def run_and_check_context(agent, agent_name):
        """Run an agent and verify its context during and after execution."""
        # Before execution, context should not exist in the outer scope
        outer_scope = sentry_sdk.get_current_scope()

        # Run the agent
        await agent.run(f"Input for {agent_name}")

        # After execution, verify context is cleaned up
        # Note: Due to isolation_scope, we can't easily check the inner scope here,
        # but we can verify the outer scope remains clean
        assert "pydantic_ai_agent" not in outer_scope._contexts

        return agent_name

    # Run both agents concurrently
    results = await asyncio.gather(
        run_and_check_context(test_agent, "agent1"),
        run_and_check_context(agent2, "agent2"),
    )

    assert results == ["agent1", "agent2"]

    # Final check: outer scope should be clean
    final_scope = sentry_sdk.get_current_scope()
    assert "pydantic_ai_agent" not in final_scope._contexts


# ==================== Additional Coverage Tests ====================


@pytest.mark.asyncio
async def test_invoke_agent_with_list_user_prompt(sentry_init, capture_events):
    """
    Test that invoke_agent span handles list user prompts correctly.
    """
    agent = Agent(
        "test",
        name="test_list_prompt",
    )

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    # Use a list as user prompt
    await agent.run(["First part", "Second part"])

    (transaction,) = events

    # Check that the invoke_agent transaction has messages data
    # The invoke_agent is the transaction itself
    if "gen_ai.request.messages" in transaction["contexts"]["trace"]["data"]:
        messages_str = transaction["contexts"]["trace"]["data"][
            "gen_ai.request.messages"
        ]
        assert "First part" in messages_str
        assert "Second part" in messages_str


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "send_default_pii, include_prompts",
    [
        (True, True),
        (True, False),
        (False, True),
        (False, False),
    ],
)
async def test_invoke_agent_with_instructions(
    sentry_init, capture_events, send_default_pii, include_prompts
):
    """
    Test that invoke_agent span handles instructions correctly.
    """
    from pydantic_ai import Agent

    # Create agent with instructions (can be string or list)
    agent = Agent(
        "test",
        name="test_instructions",
    )

    # Add instructions via _instructions attribute (internal API)
    agent._instructions = ["Instruction 1", "Instruction 2"]
    agent._system_prompts = ["System prompt"]

    sentry_init(
        integrations=[PydanticAIIntegration(include_prompts=include_prompts)],
        traces_sample_rate=1.0,
        send_default_pii=send_default_pii,
    )

    events = capture_events()

    await agent.run("Test input")

    (transaction,) = events
    spans = transaction["spans"]

    # The transaction IS the invoke_agent span, check for messages in chat spans instead
    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    assert len(chat_spans) == 1

    chat_span = chat_spans[0]

    if send_default_pii and include_prompts:
        system_instructions = chat_span["data"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]
        assert json.loads(system_instructions) == [
            {"type": "text", "content": "System prompt"},
            {"type": "text", "content": "Instruction 1\nInstruction 2"},
        ]
    else:
        assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in chat_span["data"]


@pytest.mark.asyncio
async def test_model_name_extraction_with_callable(sentry_init, capture_events):
    """
    Test model name extraction when model has a callable name() method.
    """
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.utils import _get_model_name

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    # Test the utility function directly
    mock_model = MagicMock()
    # Remove model_name attribute so it checks name() next
    del mock_model.model_name
    mock_model.name = lambda: "custom-model-name"

    # Get model name - should call the callable name()
    result = _get_model_name(mock_model)

    # Should return the result from callable
    assert result == "custom-model-name"


@pytest.mark.asyncio
async def test_model_name_extraction_fallback_to_str(sentry_init, capture_events):
    """
    Test model name extraction falls back to str() when no name attribute exists.
    """
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.utils import _get_model_name

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    # Test the utility function directly
    mock_model = MagicMock()
    # Remove name and model_name attributes
    del mock_model.name
    del mock_model.model_name

    # Get model name - should fall back to str()
    result = _get_model_name(mock_model)

    # Should return string representation
    assert result is not None
    assert isinstance(result, str)


@pytest.mark.asyncio
async def test_model_settings_object_style(sentry_init, capture_events):
    """
    Test that object-style model settings (non-dict) are handled correctly.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.utils import _set_model_data

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create mock settings object (not a dict)
        mock_settings = MagicMock()
        mock_settings.temperature = 0.8
        mock_settings.max_tokens = 200
        mock_settings.top_p = 0.95
        mock_settings.frequency_penalty = 0.5
        mock_settings.presence_penalty = 0.3

        # Set model data with object-style settings
        _set_model_data(span, None, mock_settings)

        span.finish()

    # Should not crash and should set the settings
    assert transaction is not None


@pytest.mark.asyncio
async def test_usage_data_partial(sentry_init, capture_events):
    """
    Test that usage data is correctly handled when only some fields are present.
    """
    agent = Agent(
        "test",
        name="test_usage",
    )

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    events = capture_events()

    await agent.run("Test input")

    (transaction,) = events
    spans = transaction["spans"]

    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    assert len(chat_spans) == 1

    # Check that usage data fields exist (they may or may not be set depending on TestModel)
    chat_span = chat_spans[0]
    # At minimum, the span should have been created
    assert chat_span is not None


@pytest.mark.asyncio
async def test_agent_data_from_scope(sentry_init, capture_events):
    """
    Test that agent data can be retrieved from Sentry scope when not passed directly.
    """

    agent = Agent(
        "test",
        name="test_scope_agent",
    )

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    events = capture_events()

    # The integration automatically sets agent in scope during execution
    await agent.run("Test input")

    (transaction,) = events

    # Verify agent name is captured
    assert transaction["transaction"] == "invoke_agent test_scope_agent"


@pytest.mark.asyncio
async def test_available_tools_without_description(
    sentry_init, capture_events, test_agent
):
    """
    Test that available tools are captured even when description is missing.
    """

    @test_agent.tool_plain
    def tool_without_desc(x: int) -> int:
        # No docstring = no description
        return x * 2

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    events = capture_events()

    await test_agent.run("Use the tool with 5")

    (transaction,) = events
    spans = transaction["spans"]

    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]
    if chat_spans:
        chat_span = chat_spans[0]
        if "gen_ai.request.available_tools" in chat_span["data"]:
            tools_str = chat_span["data"]["gen_ai.request.available_tools"]
            assert "tool_without_desc" in tools_str


@pytest.mark.asyncio
async def test_output_with_tool_calls(sentry_init, capture_events, test_agent):
    """
    Test that tool calls in model response are captured correctly.
    """

    @test_agent.tool_plain
    def calc_tool(value: int) -> int:
        """Calculate something."""
        return value + 10

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    await test_agent.run("Use calc_tool with 5")

    (transaction,) = events
    spans = transaction["spans"]

    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]

    # At least one chat span should exist
    assert len(chat_spans) == 2

    # Check if tool calls are captured in response
    for chat_span in chat_spans:
        # Tool calls may or may not be in response depending on TestModel behavior
        # Just verify the span was created and has basic data
        assert "gen_ai.operation.name" in chat_span["data"]


@pytest.mark.asyncio
async def test_message_formatting_with_different_parts(sentry_init, capture_events):
    """
    Test that different message part types are handled correctly in ai_client span.
    """
    from pydantic_ai import Agent, messages

    agent = Agent(
        "test",
        name="test_message_parts",
    )

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    # Create message history with different part types
    history = [
        messages.ModelRequest(parts=[messages.UserPromptPart(content="Hello")]),
        messages.ModelResponse(
            parts=[
                messages.TextPart(content="Hi there!"),
            ],
            model_name="test",
        ),
    ]

    await agent.run("What did I say?", message_history=history)

    (transaction,) = events
    spans = transaction["spans"]

    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]

    # Should have chat spans
    assert len(chat_spans) == 1

    # Check that messages are captured
    chat_span = chat_spans[0]
    if "gen_ai.request.messages" in chat_span["data"]:
        messages_data = chat_span["data"]["gen_ai.request.messages"]
        # Should contain message history
        assert messages_data is not None


@pytest.mark.asyncio
async def test_update_invoke_agent_span_with_none_output(sentry_init, capture_events):
    """
    Test that update_invoke_agent_span handles None output gracefully.
    """
    import sentry_sdk
    from sentry_sdk.integrations.pydantic_ai.spans.invoke_agent import (
        update_invoke_agent_span,
    )

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Update with None output - should not raise
        update_invoke_agent_span(span, None)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_update_ai_client_span_with_none_response(sentry_init, capture_events):
    """
    Test that update_ai_client_span handles None response gracefully.
    """
    import sentry_sdk
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import (
        update_ai_client_span,
    )

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Update with None response - should not raise
        update_ai_client_span(span, None)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_agent_without_name(sentry_init, capture_events):
    """
    Test that agent without a name is handled correctly.
    """
    # Create agent without explicit name
    agent = Agent("test")

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    events = capture_events()

    await agent.run("Test input")

    (transaction,) = events

    # Should still create transaction, just with default name
    assert transaction["type"] == "transaction"
    # Transaction name should be "invoke_agent agent" or similar default
    assert "invoke_agent" in transaction["transaction"]


@pytest.mark.asyncio
async def test_model_response_without_parts(sentry_init, capture_events):
    """
    Test handling of model response without parts attribute.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_output_data

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create mock response without parts
        mock_response = MagicMock()
        mock_response.model_name = "test-model"
        del mock_response.parts  # Remove parts attribute

        # Should not raise, just skip formatting
        _set_output_data(span, mock_response)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_input_messages_error_handling(sentry_init, capture_events):
    """
    Test that _set_input_messages handles errors gracefully.
    """
    import sentry_sdk
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Pass invalid messages that would cause an error
        invalid_messages = [object()]  # Plain object without expected attributes

        # Should not raise, error is caught internally
        _set_input_messages(span, invalid_messages)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_available_tools_error_handling(sentry_init, capture_events):
    """
    Test that _set_available_tools handles errors gracefully.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.utils import _set_available_tools

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create mock agent with invalid toolset
        mock_agent = MagicMock()
        mock_agent._function_toolset.tools.items.side_effect = Exception("Error")

        # Should not raise, error is caught internally
        _set_available_tools(span, mock_agent)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_set_usage_data_with_none_usage(sentry_init, capture_events):
    """
    Test that _set_usage_data handles None usage gracefully.
    """
    import sentry_sdk
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_usage_data

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Pass None usage - should not raise
        _set_usage_data(span, None)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_set_usage_data_with_partial_fields(sentry_init, capture_events):
    """
    Test that _set_usage_data handles usage with only some fields.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_usage_data

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create usage object with only some fields
        mock_usage = MagicMock()
        mock_usage.input_tokens = 100
        mock_usage.output_tokens = None  # Missing
        mock_usage.total_tokens = 100

        # Should only set the non-None fields
        _set_usage_data(span, mock_usage)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_message_parts_with_tool_return(sentry_init, capture_events):
    """
    Test that ToolReturnPart messages are handled correctly.
    """
    from pydantic_ai import Agent

    agent = Agent(
        "test",
        name="test_tool_return",
    )

    @agent.tool_plain
    def test_tool(x: int) -> int:
        """Test tool."""
        return x * 2

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    # Run with history containing tool return
    await agent.run("Use test_tool with 5")

    (transaction,) = events
    spans = transaction["spans"]

    chat_spans = [s for s in spans if s["op"] == "gen_ai.chat"]

    # Should have chat spans
    assert len(chat_spans) == 2


@pytest.mark.asyncio
async def test_message_parts_with_list_content(sentry_init, capture_events):
    """
    Test that message parts with list content are handled correctly.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create message with list content
        mock_msg = MagicMock()
        mock_part = MagicMock()
        mock_part.content = ["item1", "item2", {"complex": "item"}]
        mock_msg.parts = [mock_part]
        mock_msg.instructions = None

        messages = [mock_msg]

        # Should handle list content
        _set_input_messages(span, messages)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_output_data_with_text_and_tool_calls(sentry_init, capture_events):
    """
    Test that _set_output_data handles both text and tool calls in response.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_output_data

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create mock response with both TextPart and ToolCallPart
        from pydantic_ai import messages

        text_part = messages.TextPart(content="Here's the result")
        tool_call_part = MagicMock()
        tool_call_part.tool_name = "test_tool"
        tool_call_part.args = {"x": 5}

        mock_response = MagicMock()
        mock_response.model_name = "test-model"
        mock_response.parts = [text_part, tool_call_part]

        # Should handle both text and tool calls
        _set_output_data(span, mock_response)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_output_data_error_handling(sentry_init, capture_events):
    """
    Test that _set_output_data handles errors in formatting gracefully.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_output_data

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create mock response that will cause error
        mock_response = MagicMock()
        mock_response.model_name = "test-model"
        mock_response.parts = [MagicMock(side_effect=Exception("Error"))]

        # Should catch error and not crash
        _set_output_data(span, mock_response)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_message_with_system_prompt_part(sentry_init, capture_events):
    """
    Test that SystemPromptPart is handled with correct role.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages
    from pydantic_ai import messages

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create message with SystemPromptPart
        system_part = messages.SystemPromptPart(content="You are a helpful assistant")

        mock_msg = MagicMock()
        mock_msg.parts = [system_part]
        mock_msg.instructions = None

        msgs = [mock_msg]

        # Should handle system prompt
        _set_input_messages(span, msgs)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_message_with_instructions(sentry_init, capture_events):
    """
    Test that messages with instructions field are handled correctly.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create message with instructions
        mock_msg = MagicMock()
        mock_msg.instructions = "System instructions here"
        mock_part = MagicMock()
        mock_part.content = "User message"
        mock_msg.parts = [mock_part]

        msgs = [mock_msg]

        # Should extract system prompt from instructions
        _set_input_messages(span, msgs)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_set_input_messages_without_prompts(sentry_init, capture_events):
    """
    Test that _set_input_messages respects _should_send_prompts().
    """
    import sentry_sdk
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages

    sentry_init(
        integrations=[PydanticAIIntegration(include_prompts=False)],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Even with messages, should not set them
        messages = ["test"]
        _set_input_messages(span, messages)

        span.finish()

    # Should not crash and should not set messages
    assert transaction is not None


@pytest.mark.asyncio
async def test_set_output_data_without_prompts(sentry_init, capture_events):
    """
    Test that _set_output_data respects _should_send_prompts().
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_output_data

    sentry_init(
        integrations=[PydanticAIIntegration(include_prompts=False)],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Even with response, should not set output data
        mock_response = MagicMock()
        mock_response.model_name = "test"
        _set_output_data(span, mock_response)

        span.finish()

    # Should not crash and should not set output
    assert transaction is not None


@pytest.mark.asyncio
async def test_get_model_name_with_exception_in_callable(sentry_init, capture_events):
    """
    Test that _get_model_name handles exceptions in name() callable.
    """
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.utils import _get_model_name

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    # Create model with callable name that raises exception
    mock_model = MagicMock()
    mock_model.name = MagicMock(side_effect=Exception("Error"))

    # Should fall back to str()
    result = _get_model_name(mock_model)

    # Should return something (str fallback)
    assert result is not None


@pytest.mark.asyncio
async def test_get_model_name_with_string_model(sentry_init, capture_events):
    """
    Test that _get_model_name handles string models.
    """
    from sentry_sdk.integrations.pydantic_ai.utils import _get_model_name

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    # Pass a string as model
    result = _get_model_name("gpt-4")

    # Should return the string
    assert result == "gpt-4"


@pytest.mark.asyncio
async def test_get_model_name_with_none(sentry_init, capture_events):
    """
    Test that _get_model_name handles None model.
    """
    from sentry_sdk.integrations.pydantic_ai.utils import _get_model_name

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    # Pass None
    result = _get_model_name(None)

    # Should return None
    assert result is None


@pytest.mark.asyncio
async def test_set_model_data_with_system(sentry_init, capture_events):
    """
    Test that _set_model_data captures system from model.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.utils import _set_model_data

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create model with system
        mock_model = MagicMock()
        mock_model.system = "openai"
        mock_model.model_name = "gpt-4"

        # Set model data
        _set_model_data(span, mock_model, None)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_set_model_data_from_agent_scope(sentry_init, capture_events):
    """
    Test that _set_model_data retrieves model from agent in scope when not passed.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.utils import _set_model_data

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        # Set agent in scope
        scope = sentry_sdk.get_current_scope()
        mock_agent = MagicMock()
        mock_agent.model = MagicMock()
        mock_agent.model.model_name = "test-model"
        mock_agent.model_settings = {"temperature": 0.5}
        scope._contexts["pydantic_ai_agent"] = {"_agent": mock_agent}

        span = sentry_sdk.start_span(op="test_span")

        # Pass None for model, should get from scope
        _set_model_data(span, None, None)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_set_model_data_with_none_settings_values(sentry_init, capture_events):
    """
    Test that _set_model_data skips None values in settings.
    """
    import sentry_sdk
    from sentry_sdk.integrations.pydantic_ai.utils import _set_model_data

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create settings with None values
        settings = {
            "temperature": 0.7,
            "max_tokens": None,  # Should be skipped
            "top_p": None,  # Should be skipped
        }

        # Set model data
        _set_model_data(span, None, settings)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_should_send_prompts_without_pii(sentry_init, capture_events):
    """
    Test that _should_send_prompts returns False when PII disabled.
    """
    from sentry_sdk.integrations.pydantic_ai.utils import _should_send_prompts

    sentry_init(
        integrations=[PydanticAIIntegration(include_prompts=True)],
        traces_sample_rate=1.0,
        send_default_pii=False,  # PII disabled
    )

    # Should return False
    result = _should_send_prompts()
    assert result is False


@pytest.mark.asyncio
async def test_set_agent_data_without_agent(sentry_init, capture_events):
    """
    Test that _set_agent_data handles None agent gracefully.
    """
    import sentry_sdk
    from sentry_sdk.integrations.pydantic_ai.utils import _set_agent_data

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Pass None agent, with no agent in scope
        _set_agent_data(span, None)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_set_agent_data_from_scope(sentry_init, capture_events):
    """
    Test that _set_agent_data retrieves agent from scope when not passed.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.utils import _set_agent_data

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        # Set agent in scope
        scope = sentry_sdk.get_current_scope()
        mock_agent = MagicMock()
        mock_agent.name = "test_agent_from_scope"
        scope._contexts["pydantic_ai_agent"] = {"_agent": mock_agent}

        span = sentry_sdk.start_span(op="test_span")

        # Pass None for agent, should get from scope
        _set_agent_data(span, None)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_set_agent_data_without_name(sentry_init, capture_events):
    """
    Test that _set_agent_data handles agent without name attribute.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.utils import _set_agent_data

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create agent without name
        mock_agent = MagicMock()
        mock_agent.name = None  # No name

        # Should not set agent name
        _set_agent_data(span, mock_agent)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_set_available_tools_without_toolset(sentry_init, capture_events):
    """
    Test that _set_available_tools handles agent without toolset.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.utils import _set_available_tools

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create agent without _function_toolset
        mock_agent = MagicMock()
        del mock_agent._function_toolset

        # Should handle gracefully
        _set_available_tools(span, mock_agent)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_set_available_tools_with_schema(sentry_init, capture_events):
    """
    Test that _set_available_tools extracts tool schema correctly.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.utils import _set_available_tools

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        span = sentry_sdk.start_span(op="test_span")

        # Create agent with toolset containing schema
        mock_agent = MagicMock()
        mock_tool = MagicMock()
        mock_schema = MagicMock()
        mock_schema.description = "Test tool description"
        mock_schema.json_schema = {"type": "object", "properties": {}}
        mock_tool.function_schema = mock_schema

        mock_agent._function_toolset.tools = {"test_tool": mock_tool}

        # Should extract schema
        _set_available_tools(span, mock_agent)

        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_execute_tool_span_creation(sentry_init, capture_events):
    """
    Test direct creation of execute_tool span.
    """
    import sentry_sdk
    from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import (
        execute_tool_span,
        update_execute_tool_span,
    )

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        # Create execute_tool span
        with execute_tool_span("test_tool", {"arg": "value"}, None, "function") as span:
            # Update with result
            update_execute_tool_span(span, {"result": "success"})

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_execute_tool_span_with_mcp_type(sentry_init, capture_events):
    """
    Test execute_tool span with MCP tool type.
    """
    import sentry_sdk
    from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import (
        execute_tool_span,
    )

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        # Create execute_tool span with mcp type
        with execute_tool_span("mcp_tool", {"arg": "value"}, None, "mcp") as span:
            # Verify type is set
            assert span is not None

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_execute_tool_span_without_prompts(sentry_init, capture_events):
    """
    Test that execute_tool span respects _should_send_prompts().
    """
    import sentry_sdk
    from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import (
        execute_tool_span,
        update_execute_tool_span,
    )

    sentry_init(
        integrations=[PydanticAIIntegration(include_prompts=False)],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        # Create execute_tool span
        with execute_tool_span("test_tool", {"arg": "value"}, None, "function") as span:
            # Update with result - should not set input/output
            update_execute_tool_span(span, {"result": "success"})

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_execute_tool_span_with_none_args(sentry_init, capture_events):
    """
    Test execute_tool span with None args.
    """
    import sentry_sdk
    from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import execute_tool_span

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        # Create execute_tool span with None args
        with execute_tool_span("test_tool", None, None, "function") as span:
            assert span is not None

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_update_execute_tool_span_with_none_span(sentry_init, capture_events):
    """
    Test that update_execute_tool_span handles None span gracefully.
    """
    from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import (
        update_execute_tool_span,
    )

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    # Update with None span - should not raise
    update_execute_tool_span(None, {"result": "success"})

    # Should not crash
    assert True


@pytest.mark.asyncio
async def test_update_execute_tool_span_with_none_result(sentry_init, capture_events):
    """
    Test that update_execute_tool_span handles None result gracefully.
    """
    import sentry_sdk
    from sentry_sdk.integrations.pydantic_ai.spans.execute_tool import (
        execute_tool_span,
        update_execute_tool_span,
    )

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        # Create execute_tool span
        with execute_tool_span("test_tool", {"arg": "value"}, None, "function") as span:
            # Update with None result
            update_execute_tool_span(span, None)

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_tool_execution_without_span_context(sentry_init, capture_events):
    """
    Test that tool execution patch handles case when no span context exists.
    This tests the code path where current_span is None in _patch_tool_execution.
    """
    # Import the patching function

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    # Create a simple agent with no tools (won't have function_toolset)
    agent = Agent("test", name="test_no_span")

    # Call without span context (no transaction active)
    # The patches should handle this gracefully
    try:
        # This will fail because we're not in a transaction, but it should not crash
        await agent.run("test")
    except Exception:
        # Expected to fail, that's okay
        pass

    # Should not crash
    assert True


@pytest.mark.asyncio
async def test_invoke_agent_span_with_callable_instruction(sentry_init, capture_events):
    """
    Test that invoke_agent_span skips callable instructions correctly.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.spans.invoke_agent import invoke_agent_span

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        # Create mock agent with callable instruction
        mock_agent = MagicMock()
        mock_agent.name = "test_agent"
        mock_agent._system_prompts = []

        # Add both string and callable instructions
        mock_callable = lambda: "Dynamic instruction"
        mock_agent._instructions = ["Static instruction", mock_callable]

        # Create span
        span = invoke_agent_span("Test prompt", mock_agent, None, None)
        span.finish()

    # Should not crash (callable should be skipped)
    assert transaction is not None


@pytest.mark.asyncio
async def test_invoke_agent_span_with_string_instructions(sentry_init, capture_events):
    """
    Test that invoke_agent_span handles string instructions (not list).
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.spans.invoke_agent import invoke_agent_span

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        # Create mock agent with string instruction
        mock_agent = MagicMock()
        mock_agent.name = "test_agent"
        mock_agent._system_prompts = []
        mock_agent._instructions = "Single instruction string"

        # Create span
        span = invoke_agent_span("Test prompt", mock_agent, None, None)
        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_ai_client_span_with_streaming_flag(sentry_init, capture_events):
    """
    Test that ai_client_span reads streaming flag from scope.
    """
    import sentry_sdk
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import ai_client_span

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        # Set streaming flag in scope
        scope = sentry_sdk.get_current_scope()
        scope._contexts["pydantic_ai_agent"] = {"_streaming": True}

        # Create ai_client span
        span = ai_client_span([], None, None, None)
        span.finish()

    # Should not crash
    assert transaction is not None


@pytest.mark.asyncio
async def test_ai_client_span_gets_agent_from_scope(sentry_init, capture_events):
    """
    Test that ai_client_span gets agent from scope when not passed.
    """
    import sentry_sdk
    from unittest.mock import MagicMock
    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import ai_client_span

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
    )

    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
        # Set agent in scope
        scope = sentry_sdk.get_current_scope()
        mock_agent = MagicMock()
        mock_agent.name = "test_agent"
        mock_agent._function_toolset = MagicMock()
        mock_agent._function_toolset.tools = {}
        scope._contexts["pydantic_ai_agent"] = {"_agent": mock_agent}

        # Create ai_client span without passing agent
        span = ai_client_span([], None, None, None)
        span.finish()

    # Should not crash
    assert transaction is not None


def _get_messages_from_span(span_data):
    """Helper to extract and parse messages from span data."""
    messages_data = span_data["gen_ai.request.messages"]
    return (
        json.loads(messages_data) if isinstance(messages_data, str) else messages_data
    )


def _find_binary_content(messages_data, expected_modality, expected_mime_type):
    """Helper to find and verify binary content in messages."""
    for msg in messages_data:
        if "content" not in msg:
            continue
        for content_item in msg["content"]:
            if content_item.get("type") == "blob":
                assert content_item["modality"] == expected_modality
                assert content_item["mime_type"] == expected_mime_type
                assert content_item["content"] == BLOB_DATA_SUBSTITUTE
                return True
    return False


@pytest.mark.asyncio
async def test_binary_content_encoding_image(sentry_init, capture_events):
    """Test that BinaryContent with image data is properly encoded in messages."""
    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    with sentry_sdk.start_transaction(op="test", name="test"):
        span = sentry_sdk.start_span(op="test_span")
        binary_content = BinaryContent(
            data=b"fake_image_data_12345", media_type="image/png"
        )
        user_part = UserPromptPart(content=["Look at this image:", binary_content])
        mock_msg = MagicMock()
        mock_msg.parts = [user_part]
        mock_msg.instructions = None

        _set_input_messages(span, [mock_msg])
        span.finish()

    (event,) = events
    span_data = event["spans"][0]["data"]
    messages_data = _get_messages_from_span(span_data)
    assert _find_binary_content(messages_data, "image", "image/png")


@pytest.mark.asyncio
async def test_binary_content_encoding_mixed_content(sentry_init, capture_events):
    """Test that BinaryContent mixed with text content is properly handled."""
    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    with sentry_sdk.start_transaction(op="test", name="test"):
        span = sentry_sdk.start_span(op="test_span")
        binary_content = BinaryContent(
            data=b"fake_image_bytes", media_type="image/jpeg"
        )
        user_part = UserPromptPart(
            content=["Here is an image:", binary_content, "What do you see?"]
        )
        mock_msg = MagicMock()
        mock_msg.parts = [user_part]
        mock_msg.instructions = None

        _set_input_messages(span, [mock_msg])
        span.finish()

    (event,) = events
    span_data = event["spans"][0]["data"]
    messages_data = _get_messages_from_span(span_data)

    # Verify both text and binary content are present
    found_text = any(
        content_item.get("type") == "text"
        for msg in messages_data
        if "content" in msg
        for content_item in msg["content"]
    )
    assert found_text, "Text content should be found"
    assert _find_binary_content(messages_data, "image", "image/jpeg")


@pytest.mark.asyncio
async def test_binary_content_in_agent_run(sentry_init, capture_events):
    """Test that BinaryContent in actual agent run is properly captured in spans."""
    agent = Agent("test", name="test_binary_agent")

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()
    binary_content = BinaryContent(
        data=b"fake_image_data_for_testing", media_type="image/png"
    )
    await agent.run(["Analyze this image:", binary_content])

    (transaction,) = events
    chat_spans = [s for s in transaction["spans"] if s["op"] == "gen_ai.chat"]
    assert len(chat_spans) == 1

    chat_span = chat_spans[0]
    if "gen_ai.request.messages" in chat_span["data"]:
        messages_str = str(chat_span["data"]["gen_ai.request.messages"])
        assert any(keyword in messages_str for keyword in ["blob", "image", "base64"])


@pytest.mark.asyncio
async def test_set_usage_data_with_cache_tokens(sentry_init, capture_events):
    """Test that cache_read_tokens and cache_write_tokens are tracked."""
    sentry_init(integrations=[PydanticAIIntegration()], traces_sample_rate=1.0)

    events = capture_events()

    with sentry_sdk.start_transaction(op="test", name="test"):
        span = sentry_sdk.start_span(op="test_span")
        usage = RequestUsage(
            input_tokens=100,
            output_tokens=50,
            cache_read_tokens=80,
            cache_write_tokens=20,
        )
        _set_usage_data(span, usage)
        span.finish()

    (event,) = events
    (span_data,) = event["spans"]
    assert span_data["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80
    assert span_data["data"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "url, image_url_kwargs, expected_content",
    [
        pytest.param(
            "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
            {},
            BLOB_DATA_SUBSTITUTE,
            id="base64_data_url",
        ),
        pytest.param(
            "https://example.com/image.png",
            {},
            "https://example.com/image.png",
            id="http_url_no_redaction",
        ),
        pytest.param(
            "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
            {},
            "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
            id="http_url_with_base64_query_param",
        ),
        pytest.param(
            "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
            {"media_type": "image/png"},
            "https://example.com/api?data=iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
            id="http_url_with_base64_query_param_and_media_type",
        ),
        pytest.param(
            "data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciLz4=",
            {},
            BLOB_DATA_SUBSTITUTE,
            id="complex_mime_type",
        ),
        pytest.param(
            "data:image/png;name=file.png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs",
            {},
            BLOB_DATA_SUBSTITUTE,
            id="optional_parameters",
        ),
        pytest.param(
            "data:text/plain;charset=utf-8;name=hello.txt;base64,SGVsbG8sIFdvcmxkIQ==",
            {},
            BLOB_DATA_SUBSTITUTE,
            id="multiple_optional_parameters",
        ),
    ],
)
async def test_invoke_agent_image_url(
    sentry_init, capture_events, url, image_url_kwargs, expected_content
):
    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    agent = Agent("test", name="test_image_url_agent")

    events = capture_events()
    image_url = ImageUrl(url=url, **image_url_kwargs)
    await agent.run([image_url, "Describe this image"])

    (transaction,) = events

    found_image = False

    chat_spans = [s for s in transaction["spans"] if s["op"] == "gen_ai.chat"]
    for chat_span in chat_spans:
        messages_data = _get_messages_from_span(chat_span["data"])
        for msg in messages_data:
            if "content" not in msg:
                continue
            for content_item in msg["content"]:
                if content_item.get("type") == "image":
                    assert content_item["content"] == expected_content
                    found_image = True

    assert found_image, "Image content item should be found in messages data"


@pytest.mark.asyncio
async def test_tool_description_in_execute_tool_span(sentry_init, capture_events):
    """
    Test that tool description from the tool's docstring is included in execute_tool spans.
    """
    agent = Agent(
        "test",
        name="test_agent",
        system_prompt="You are a helpful test assistant.",
    )

    @agent.tool_plain
    def multiply_numbers(a: int, b: int) -> int:
        """Multiply two numbers and return the product."""
        return a * b

    sentry_init(
        integrations=[PydanticAIIntegration()],
        traces_sample_rate=1.0,
        send_default_pii=True,
    )

    events = capture_events()

    result = await agent.run("What is 5 times 3?")
    assert result is not None

    (transaction,) = events
    spans = transaction["spans"]

    tool_spans = [s for s in spans if s["op"] == "gen_ai.execute_tool"]
    assert len(tool_spans) >= 1

    tool_span = tool_spans[0]
    assert tool_span["data"]["gen_ai.tool.name"] == "multiply_numbers"
    assert SPANDATA.GEN_AI_TOOL_DESCRIPTION in tool_span["data"]
    assert "Multiply two numbers" in tool_span["data"][SPANDATA.GEN_AI_TOOL_DESCRIPTION]