diff --git a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py index 90048e41fc..2581ed41bb 100644 --- a/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py +++ b/sentry_sdk/integrations/pydantic_ai/patches/agent_run.py @@ -15,7 +15,7 @@ raise DidNotEnable("pydantic-ai not installed") if TYPE_CHECKING: - from typing import Any, Callable, Optional + from typing import Any, Callable, Optional, Union class _StreamingContextManagerWrapper: @@ -37,7 +37,7 @@ def __init__( self.model_settings = model_settings self.is_streaming = is_streaming self._isolation_scope: "Any" = None - self._span: "Optional[sentry_sdk.tracing.Span]" = None + self._span: "Optional[Union[sentry_sdk.tracing.Span, sentry_sdk.traces.StreamedSpan]]" = None self._result: "Any" = None async def __aenter__(self) -> "Any": diff --git a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py index 800ec16e36..dfe898d139 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/ai_client.py @@ -8,6 +8,8 @@ truncate_and_annotate_messages, ) from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.traces import StreamedSpan +from sentry_sdk.tracing_utils import has_span_streaming_enabled from sentry_sdk.utils import safe_serialize from ..consts import SPAN_ORIGIN @@ -27,7 +29,7 @@ ) if TYPE_CHECKING: - from typing import Any, Dict, List + from typing import Any, Dict, List, Union from pydantic_ai.messages import ModelMessage, SystemPromptPart # type: ignore @@ -97,7 +99,9 @@ def _get_system_instructions( return permanent_instructions, current_instructions -def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> None: +def _set_input_messages( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", messages: "Any" +) -> None: """Set input messages data on a span.""" if not _should_send_prompts(): return @@ -107,14 +111,24 @@ def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> Non permanent_instructions, current_instructions = _get_system_instructions(messages) if len(permanent_instructions) > 0 or len(current_instructions) > 0: - span.set_data( - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - json.dumps( - _transform_system_instructions( - permanent_instructions, current_instructions - ) - ), - ) + if isinstance(span, StreamedSpan): + span.set_attribute( + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + json.dumps( + _transform_system_instructions( + permanent_instructions, current_instructions + ) + ), + ) + else: + span.set_data( + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + json.dumps( + _transform_system_instructions( + permanent_instructions, current_instructions + ) + ), + ) try: formatted_messages = [] @@ -198,7 +212,9 @@ def _set_input_messages(span: "sentry_sdk.tracing.Span", messages: "Any") -> Non pass -def _set_output_data(span: "sentry_sdk.tracing.Span", response: "Any") -> None: +def _set_output_data( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", response: "Any" +) -> None: """Set output data on a span.""" if not _should_send_prompts(): return @@ -206,7 +222,11 @@ def _set_output_data(span: "sentry_sdk.tracing.Span", response: "Any") -> None: if not response: return - span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) + set_on_span = ( + span.set_attribute if isinstance(span, StreamedSpan) else span.set_data + ) + set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) + try: # Extract text from ModelResponse if hasattr(response, "parts"): @@ -230,7 +250,7 @@ def _set_output_data(span: "sentry_sdk.tracing.Span", response: "Any") -> None: set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, texts) if tool_calls: - span.set_data( + set_on_span( SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, safe_serialize(tool_calls) ) @@ -241,7 +261,7 @@ def _set_output_data(span: "sentry_sdk.tracing.Span", response: "Any") -> None: def ai_client_span( messages: "Any", agent: "Any", model: "Any", model_settings: "Any" -) -> "sentry_sdk.tracing.Span": +) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": """Create a span for an AI client call (model request). Args: @@ -257,20 +277,31 @@ def ai_client_span( model_name = _get_model_name(model_obj) or "unknown" - span = sentry_sdk.start_span( - op=OP.GEN_AI_CHAT, - name=f"chat {model_name}", - origin=SPAN_ORIGIN, - ) + span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + if span_streaming: + span = sentry_sdk.traces.start_span( + name=f"chat {model_name}", + attributes={ + "sentry.op": OP.GEN_AI_CHAT, + "sentry.origin": SPAN_ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "chat", + SPANDATA.GEN_AI_RESPONSE_STREAMING: get_is_streaming(), + }, + ) + else: + span = sentry_sdk.start_span( + op=OP.GEN_AI_CHAT, + name=f"chat {model_name}", + origin=SPAN_ORIGIN, + ) - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") + # Set streaming flag from contextvar + span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, get_is_streaming()) _set_agent_data(span, agent) _set_model_data(span, model, model_settings) - # Set streaming flag from contextvar - span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, get_is_streaming()) - # Add available tools if agent is available agent_obj = agent or get_current_agent() _set_available_tools(span, agent_obj) @@ -283,7 +314,7 @@ def ai_client_span( def update_ai_client_span( - span: "sentry_sdk.tracing.Span", model_response: "Any" + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", model_response: "Any" ) -> None: """Update the AI client span with response data.""" if not span: diff --git a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py index 5b4cd1ac94..7648c1418a 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py @@ -2,13 +2,15 @@ import sentry_sdk from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.traces import StreamedSpan +from sentry_sdk.tracing_utils import has_span_streaming_enabled from sentry_sdk.utils import safe_serialize from ..consts import SPAN_ORIGIN from ..utils import _set_agent_data, _should_send_prompts if TYPE_CHECKING: - from typing import Any, Optional + from typing import Any, Optional, Union from pydantic_ai._tool_manager import ToolDefinition # type: ignore @@ -18,7 +20,7 @@ def execute_tool_span( tool_args: "Any", agent: "Any", tool_definition: "Optional[ToolDefinition]" = None, -) -> "sentry_sdk.tracing.Span": +) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": """Create a span for tool execution. Args: @@ -27,17 +29,33 @@ def execute_tool_span( agent: The agent executing the tool tool_definition: The definition of the tool, if available """ - span = sentry_sdk.start_span( - op=OP.GEN_AI_EXECUTE_TOOL, - name=f"execute_tool {tool_name}", - origin=SPAN_ORIGIN, - ) + span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + if span_streaming: + span = sentry_sdk.traces.start_span( + name=f"execute_tool {tool_name}", + attributes={ + "sentry.op": OP.GEN_AI_EXECUTE_TOOL, + "sentry.origin": SPAN_ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "execute_tool", + SPANDATA.GEN_AI_TOOL_NAME: tool_name, + }, + ) + + set_on_span = span.set_attribute + else: + span = sentry_sdk.start_span( + op=OP.GEN_AI_EXECUTE_TOOL, + name=f"execute_tool {tool_name}", + origin=SPAN_ORIGIN, + ) - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool") - span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name) + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool") + span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name) + + set_on_span = span.set_data if tool_definition is not None and hasattr(tool_definition, "description"): - span.set_data( + set_on_span( SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool_definition.description, ) @@ -45,15 +63,22 @@ def execute_tool_span( _set_agent_data(span, agent) if _should_send_prompts() and tool_args is not None: - span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args)) + set_on_span(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args)) return span -def update_execute_tool_span(span: "sentry_sdk.tracing.Span", result: "Any") -> None: +def update_execute_tool_span( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", result: "Any" +) -> None: """Update the execute tool span with the result.""" if not span: return - if _should_send_prompts() and result is not None: + if not _should_send_prompts() or result is None: + return + + if isinstance(span, StreamedSpan): + span.set_attribute(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) + else: span.set_data(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result)) diff --git a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py index e9fdc90fe0..e4021f3bfa 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py @@ -8,6 +8,8 @@ truncate_and_annotate_messages, ) from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.traces import StreamedSpan +from sentry_sdk.tracing_utils import has_span_streaming_enabled from ..consts import SPAN_ORIGIN from ..utils import ( @@ -22,7 +24,7 @@ ) if TYPE_CHECKING: - from typing import Any + from typing import Any, Union try: from pydantic_ai.messages import BinaryContent, ImageUrl # type: ignore @@ -37,20 +39,31 @@ def invoke_agent_span( model: "Any", model_settings: "Any", is_streaming: bool = False, -) -> "sentry_sdk.tracing.Span": +) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]": """Create a span for invoking the agent.""" # Determine agent name for span name = "agent" if agent and getattr(agent, "name", None): name = agent.name - span = get_start_span_function()( - op=OP.GEN_AI_INVOKE_AGENT, - name=f"invoke_agent {name}", - origin=SPAN_ORIGIN, - ) + span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + if span_streaming: + span = sentry_sdk.traces.start_span( + name=f"invoke_agent {name}", + attributes={ + "sentry.op": OP.GEN_AI_INVOKE_AGENT, + "sentry.origin": SPAN_ORIGIN, + SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent", + }, + ) + else: + span = get_start_span_function()( + op=OP.GEN_AI_INVOKE_AGENT, + name=f"invoke_agent {name}", + origin=SPAN_ORIGIN, + ) - span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") + span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent") _set_agent_data(span, agent) _set_model_data(span, model, model_settings) @@ -135,7 +148,10 @@ def invoke_agent_span( return span -def update_invoke_agent_span(span: "sentry_sdk.tracing.Span", result: "Any") -> None: +def update_invoke_agent_span( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + result: "Any", +) -> None: """Update and close the invoke agent span.""" if not span or not result: return @@ -154,7 +170,12 @@ def update_invoke_agent_span(span: "sentry_sdk.tracing.Span", result: "Any") -> try: response = result.response if hasattr(response, "model_name") and response.model_name: - span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) + if isinstance(span, StreamedSpan): + span.set_attribute( + SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name + ) + else: + span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response.model_name) except Exception: # If response access fails, continue without setting model name pass diff --git a/sentry_sdk/integrations/pydantic_ai/spans/utils.py b/sentry_sdk/integrations/pydantic_ai/spans/utils.py index 1564d2d9d0..330496c6b2 100644 --- a/sentry_sdk/integrations/pydantic_ai/spans/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/spans/utils.py @@ -7,6 +7,7 @@ from sentry_sdk.ai.consts import DATA_URL_BASE64_REGEX from sentry_sdk.ai.utils import get_modality_from_mime_type from sentry_sdk.consts import SPANDATA +from sentry_sdk.traces import StreamedSpan if TYPE_CHECKING: from typing import Any, Dict, Union @@ -46,7 +47,8 @@ def _serialize_binary_content_item(item: "Any") -> "Dict[str, Any]": def _set_usage_data( - span: "sentry_sdk.tracing.Span", usage: "Union[RequestUsage, RunUsage]" + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + usage: "Union[RequestUsage, RunUsage]", ) -> None: """Set token usage data on a span. @@ -60,24 +62,26 @@ def _set_usage_data( if usage is None: return + set_on_span = ( + span.set_attribute if isinstance(span, StreamedSpan) else span.set_data + ) + if hasattr(usage, "input_tokens") and usage.input_tokens is not None: - span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens) + set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens) # Pydantic AI uses cache_read_tokens (not input_tokens_cached) if hasattr(usage, "cache_read_tokens") and usage.cache_read_tokens is not None: - span.set_data( - SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, usage.cache_read_tokens - ) + set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED, usage.cache_read_tokens) # Pydantic AI uses cache_write_tokens (not input_tokens_cache_write) if hasattr(usage, "cache_write_tokens") and usage.cache_write_tokens is not None: - span.set_data( + set_on_span( SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE, usage.cache_write_tokens, ) if hasattr(usage, "output_tokens") and usage.output_tokens is not None: - span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens) + set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, usage.output_tokens) if hasattr(usage, "total_tokens") and usage.total_tokens is not None: - span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens) + set_on_span(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, usage.total_tokens) diff --git a/sentry_sdk/integrations/pydantic_ai/utils.py b/sentry_sdk/integrations/pydantic_ai/utils.py index 896ae20b9d..340dcf8953 100644 --- a/sentry_sdk/integrations/pydantic_ai/utils.py +++ b/sentry_sdk/integrations/pydantic_ai/utils.py @@ -4,10 +4,11 @@ import sentry_sdk from sentry_sdk.consts import SPANDATA from sentry_sdk.scope import should_send_default_pii +from sentry_sdk.traces import StreamedSpan from sentry_sdk.utils import event_from_exception, safe_serialize if TYPE_CHECKING: - from typing import Any, Optional + from typing import Any, Optional, Union # Store the current agent context in a contextvar for re-entrant safety @@ -68,7 +69,9 @@ def _should_send_prompts() -> bool: return getattr(integration, "include_prompts", False) -def _set_agent_data(span: "sentry_sdk.tracing.Span", agent: "Any") -> None: +def _set_agent_data( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", agent: "Any" +) -> None: """Set agent-related data on a span. Args: @@ -82,7 +85,10 @@ def _set_agent_data(span: "sentry_sdk.tracing.Span", agent: "Any") -> None: agent_obj = get_current_agent() if agent_obj and hasattr(agent_obj, "name") and agent_obj.name: - span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name) + if isinstance(span, StreamedSpan): + span.set_attribute(SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name) + else: + span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_obj.name) def _get_model_name(model_obj: "Any") -> "Optional[str]": @@ -111,7 +117,9 @@ def _get_model_name(model_obj: "Any") -> "Optional[str]": def _set_model_data( - span: "sentry_sdk.tracing.Span", model: "Any", model_settings: "Any" + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", + model: "Any", + model_settings: "Any", ) -> None: """Set model-related data on a span. @@ -128,15 +136,19 @@ def _set_model_data( if not model_obj and agent_obj and hasattr(agent_obj, "model"): model_obj = agent_obj.model + set_on_span = ( + span.set_attribute if isinstance(span, StreamedSpan) else span.set_data + ) + if model_obj: # Set system from model if hasattr(model_obj, "system"): - span.set_data(SPANDATA.GEN_AI_SYSTEM, model_obj.system) + set_on_span(SPANDATA.GEN_AI_SYSTEM, model_obj.system) # Set model name model_name = _get_model_name(model_obj) if model_name: - span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) + set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model_name) # Extract model settings settings = model_settings @@ -157,17 +169,19 @@ def _set_model_data( for setting_name, spandata_key in settings_map.items(): value = settings.get(setting_name) if value is not None: - span.set_data(spandata_key, value) + set_on_span(spandata_key, value) else: # Fallback for object-style settings for setting_name, spandata_key in settings_map.items(): if hasattr(settings, setting_name): value = getattr(settings, setting_name) if value is not None: - span.set_data(spandata_key, value) + set_on_span(spandata_key, value) -def _set_available_tools(span: "sentry_sdk.tracing.Span", agent: "Any") -> None: +def _set_available_tools( + span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", agent: "Any" +) -> None: """Set available tools data on a span from an agent's function toolset. Args: @@ -197,9 +211,14 @@ def _set_available_tools(span: "sentry_sdk.tracing.Span", agent: "Any") -> None: tools.append(tool_info) if tools: - span.set_data( - SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) - ) + if isinstance(span, StreamedSpan): + span.set_attribute( + SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) + ) + else: + span.set_data( + SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools) + ) except Exception: # If we can't extract tools, just skip it pass diff --git a/tests/integrations/pydantic_ai/test_pydantic_ai.py b/tests/integrations/pydantic_ai/test_pydantic_ai.py index cccf5d49bc..611a158967 100644 --- a/tests/integrations/pydantic_ai/test_pydantic_ai.py +++ b/tests/integrations/pydantic_ai/test_pydantic_ai.py @@ -55,6 +55,7 @@ def inner(): return inner +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_agent_run_async( @@ -63,6 +64,7 @@ async def test_agent_run_async( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that the integration creates spans for async agent runs. @@ -72,11 +74,59 @@ async def test_agent_run_async( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + result = await test_agent.run( + ["Message demonstrating the absence of truncation.", "Test input"] + ) + + assert result is not None + assert result.output is not None + + sentry_sdk.flush() + spans = [item.payload for item in items] + + assert spans[1]["name"] == "invoke_agent test_agent" + assert spans[1]["attributes"]["sentry.origin"] == "auto.ai.pydantic_ai" + + assert spans[1]["attributes"]["sentry.op"] == "gen_ai.invoke_agent" + + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + assert len(chat_spans) >= 1 + + # Check chat span + chat_span = chat_spans[0] + assert "chat" in chat_span["name"] + assert chat_span["attributes"]["gen_ai.operation.name"] == "chat" + assert chat_span["attributes"]["gen_ai.response.streaming"] is False + assert json.loads( + chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + ) == [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Message demonstrating the absence of truncation.", + }, + { + "type": "text", + "text": "Test input", + }, + ], + } + ] + assert "gen_ai.usage.input_tokens" in chat_span["attributes"] + assert "gen_ai.usage.output_tokens" in chat_span["attributes"] + elif stream_gen_ai_spans: items = capture_items("transaction", "span") result = await test_agent.run( @@ -158,6 +208,7 @@ async def test_agent_run_async( assert "gen_ai.usage.output_tokens" in chat_span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_agent_run_async_model_error( @@ -165,11 +216,13 @@ async def test_agent_run_async_model_error( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) def failing_model(messages, info): @@ -180,7 +233,21 @@ def failing_model(messages, info): name="test_agent", ) - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("event", "span") + + with pytest.raises(RuntimeError, match="model exploded"): + await agent.run("Test input") + + (error,) = (item.payload for item in items if item.type == "event") + assert error["level"] == "error" + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert len(spans) == 2 + + assert spans[0]["status"] == "error" + elif stream_gen_ai_spans: items = capture_items("event", "transaction", "span") with pytest.raises(RuntimeError, match="model exploded"): @@ -208,6 +275,7 @@ def failing_model(messages, info): assert spans[0]["status"] == "internal_error" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_agent_run_sync( sentry_init, @@ -215,6 +283,7 @@ def test_agent_run_sync( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that the integration creates spans for sync agent runs. @@ -224,11 +293,36 @@ def test_agent_run_sync( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + result = test_agent.run_sync( + ["Message demonstrating the absence of truncation.", "Test input"] + ) + + assert result is not None + assert result.output is not None + + sentry_sdk.flush() + spans = [item.payload for item in items] + + assert spans[1]["name"] == "invoke_agent test_agent" + assert spans[1]["attributes"]["sentry.origin"] == "auto.ai.pydantic_ai" + + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.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["attributes"]["gen_ai.response.streaming"] is False + elif stream_gen_ai_spans: items = capture_items("transaction", "span") result = test_agent.run_sync( @@ -279,17 +373,20 @@ def test_agent_run_sync( assert chat_span["data"]["gen_ai.response.streaming"] is False +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) def test_agent_run_sync_model_error( sentry_init, capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) def failing_model(messages, info): @@ -300,7 +397,21 @@ def failing_model(messages, info): name="test_agent", ) - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("event", "span") + + with pytest.raises(RuntimeError, match="model exploded"): + agent.run_sync("Test input") + + (error,) = (item.payload for item in items if item.type == "event") + assert error["level"] == "error" + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert len(spans) == 2 + + assert spans[0]["status"] == "error" + elif stream_gen_ai_spans: items = capture_items("event", "transaction", "span") with pytest.raises(RuntimeError, match="model exploded"): @@ -328,6 +439,7 @@ def failing_model(messages, info): assert spans[0]["status"] == "internal_error" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_agent_run_stream( @@ -336,6 +448,7 @@ async def test_agent_run_stream( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that the integration creates spans for streaming agent runs. @@ -345,11 +458,59 @@ async def test_agent_run_stream( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + async with test_agent.run_stream( + ["Message demonstrating the absence of truncation.", "Test input"] + ) as result: + # Consume the stream + async for _ in result.stream_output(): + pass + + sentry_sdk.flush() + spans = [item.payload for item in items] + + assert spans[1]["name"] == "invoke_agent test_agent" + assert spans[1]["attributes"]["sentry.origin"] == "auto.ai.pydantic_ai" + + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.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["attributes"]["gen_ai.response.streaming"] is True + assert json.loads( + chat_span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + ) == [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Message demonstrating the absence of truncation.", + }, + { + "type": "text", + "text": "Test input", + }, + ], + } + ] + assert "gen_ai.usage.input_tokens" in chat_span["attributes"] + # Streaming responses should still have output data + assert ( + "gen_ai.response.text" in chat_span["attributes"] + or "gen_ai.response.model" in chat_span["attributes"] + ) + elif stream_gen_ai_spans: items = capture_items("transaction", "span") async with test_agent.run_stream( @@ -430,6 +591,7 @@ async def test_agent_run_stream( ) +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_agent_run_stream_events( @@ -438,6 +600,7 @@ async def test_agent_run_stream_events( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that run_stream_events creates spans (it uses run internally, so non-streaming). @@ -447,12 +610,41 @@ async def test_agent_run_stream_events( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Consume all events test_agent = get_test_agent() - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + if PYDANTIC_AI_VERSION > (2,): + async with test_agent.run_stream_events( + ["Message demonstrating the absence of truncation.", "Test input"] + ) as stream_events: + async for _ in stream_events: + pass + else: + async for _ in test_agent.run_stream_events( + ["Message demonstrating the absence of truncation.", "Test input"] + ): + pass + + sentry_sdk.flush() + spans = [item.payload for item in items] + + assert spans[-1]["name"] == "invoke_agent test_agent" + + chat_spans = [ + s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" + ] + assert len(chat_spans) >= 1 + + # run_stream_events uses run() internally, so streaming should be False + for chat_span in chat_spans: + assert chat_span["attributes"]["gen_ai.response.streaming"] is False + elif stream_gen_ai_spans: items = capture_items("transaction", "span") if PYDANTIC_AI_VERSION > (2,): @@ -509,6 +701,7 @@ async def test_agent_run_stream_events( assert chat_span["data"]["gen_ai.response.streaming"] is False +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_agent_with_tools( @@ -517,6 +710,7 @@ async def test_agent_with_tools( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that tool execution creates execute_tool spans. @@ -526,6 +720,7 @@ async def test_agent_with_tools( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() @@ -535,13 +730,14 @@ def add_numbers(a: int, b: int) -> int: """Add two numbers together.""" return a + b - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") result = await test_agent.run("What is 5 + 3?") assert result is not None + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] # Find child span types (invoke_agent is the transaction, not a child span) @@ -606,6 +802,7 @@ def add_numbers(a: int, b: int) -> int: assert "add_numbers" in available_tools_str +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.parametrize( "handled_tool_call_exceptions", @@ -619,6 +816,7 @@ async def test_agent_with_tool_model_retry( get_test_agent, handled_tool_call_exceptions, stream_gen_ai_spans, + span_streaming, ): """ Test that a handled exception is captured when a tool raises ModelRetry. @@ -632,6 +830,7 @@ async def test_agent_with_tool_model_retry( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) retries = 0 @@ -647,7 +846,7 @@ def add_numbers(a: int, b: int) -> float: raise ModelRetry(message="Try again with the same arguments.") return a + b - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("event", "transaction", "span") result = await test_agent.run("What is 5 + 3?") @@ -659,6 +858,7 @@ def add_numbers(a: int, b: int) -> float: assert error["level"] == "error" assert error["exception"]["values"][0]["mechanism"]["handled"] + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] # Find child span types (invoke_agent is the transaction, not a child span) chat_spans = [ @@ -745,6 +945,7 @@ def add_numbers(a: int, b: int) -> float: assert "add_numbers" in available_tools_str +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.parametrize( "handled_tool_call_exceptions", @@ -758,6 +959,7 @@ async def test_agent_with_tool_validation_error( get_test_agent, handled_tool_call_exceptions, stream_gen_ai_spans, + span_streaming, ): """ Test that a handled exception is captured when a tool has unsatisfiable constraints. @@ -771,6 +973,7 @@ async def test_agent_with_tool_validation_error( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() @@ -780,7 +983,7 @@ def add_numbers(a: Annotated[int, Field(gt=0, lt=0)], b: int) -> int: """Add two numbers together.""" return a + b - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("event", "transaction", "span") result = None @@ -798,6 +1001,7 @@ def add_numbers(a: Annotated[int, Field(gt=0, lt=0)], b: int) -> int: assert error["level"] == "error" assert error["exception"]["values"][0]["mechanism"]["handled"] + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] chat_spans = [ s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" @@ -874,6 +1078,7 @@ def add_numbers(a: Annotated[int, Field(gt=0, lt=0)], b: int) -> int: assert "add_numbers" in available_tools_str +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_agent_with_tools_streaming( @@ -882,6 +1087,7 @@ async def test_agent_with_tools_streaming( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that tool execution works correctly with streaming. @@ -891,6 +1097,7 @@ async def test_agent_with_tools_streaming( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() @@ -900,13 +1107,14 @@ def multiply(a: int, b: int) -> int: """Multiply two numbers.""" return a * b - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") async with test_agent.run_stream("What is 7 times 8?") as result: async for _ in result.stream_output(): pass + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] # Find span types @@ -959,6 +1167,7 @@ def multiply(a: int, b: int) -> int: assert "gen_ai.tool.output" in tool_span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_model_settings( @@ -967,6 +1176,7 @@ async def test_model_settings( capture_items, get_test_agent_with_settings, stream_gen_ai_spans, + span_streaming, ): """ Test that model settings are captured in spans. @@ -975,15 +1185,17 @@ async def test_model_settings( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent_with_settings = get_test_agent_with_settings() - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await test_agent_with_settings.run("Test input") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] # Find chat span @@ -1016,6 +1228,7 @@ async def test_model_settings( assert chat_span["data"].get("gen_ai.request.top_p") == 0.9 +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio @pytest.mark.parametrize( @@ -1034,6 +1247,7 @@ async def test_system_prompt_attribute( send_default_pii, include_prompts, stream_gen_ai_spans, + span_streaming, ): """ Test that system prompts are included as the first message. @@ -1049,13 +1263,15 @@ async def test_system_prompt_attribute( traces_sample_rate=1.0, send_default_pii=send_default_pii, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await agent.run("Hello") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] # The transaction IS the invoke_agent span, check for messages in chat spans instead @@ -1104,6 +1320,7 @@ async def test_system_prompt_attribute( assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in chat_span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_error_handling( @@ -1111,6 +1328,7 @@ async def test_error_handling( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): """ Test error handling in agent execution. @@ -1126,9 +1344,21 @@ async def test_error_handling( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + # Simple run that should succeed + await agent.run("Hello") + + sentry_sdk.flush() + spans = [item.payload for item in items] + + assert spans[1]["is_segment"] is True + assert spans[1]["status"] != "error" # Could be None or some other status + elif stream_gen_ai_spans: items = capture_items("transaction", "span") # Simple run that should succeed @@ -1136,6 +1366,11 @@ async def test_error_handling( # At minimum, we should have a transaction transaction = next(item.payload for item in items if item.type == "transaction") + + 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 else: events = capture_events() @@ -1146,12 +1381,13 @@ async def test_error_handling( 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 + 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.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_without_pii( @@ -1160,6 +1396,7 @@ async def test_without_pii( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that PII is not captured when send_default_pii is False. @@ -1169,14 +1406,16 @@ async def test_without_pii( traces_sample_rate=1.0, send_default_pii=False, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") test_agent = get_test_agent() await test_agent.run("Sensitive input") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] # Find child spans (invoke_agent is the transaction, not a child span) @@ -1206,6 +1445,7 @@ async def test_without_pii( assert "gen_ai.response.text" not in span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_without_pii_tools( @@ -1214,6 +1454,7 @@ async def test_without_pii_tools( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that tool input/output are not captured when send_default_pii is False. @@ -1223,6 +1464,7 @@ async def test_without_pii_tools( traces_sample_rate=1.0, send_default_pii=False, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() @@ -1232,11 +1474,12 @@ def sensitive_tool(data: str) -> str: """A tool with sensitive data.""" return f"Processed: {data}" - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await test_agent.run("Use sensitive tool with private data") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] # Find tool spans @@ -1267,6 +1510,7 @@ def sensitive_tool(data: str) -> str: assert "gen_ai.tool.output" not in tool_span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_multiple_agents_concurrent( @@ -1275,6 +1519,7 @@ async def test_multiple_agents_concurrent( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that multiple agents can run concurrently without interfering. @@ -1283,6 +1528,7 @@ async def test_multiple_agents_concurrent( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() @@ -1290,7 +1536,21 @@ async def test_multiple_agents_concurrent( async def run_agent(input_text): return await test_agent.run(input_text) - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + # Run 3 agents concurrently + results = await asyncio.gather(*[run_agent(f"Input {i}") for i in range(3)]) + + assert len(results) == 3 + + sentry_sdk.flush() + spans = [item.payload for item in items] + for span in spans: + if span["is_segment"] is False: + continue + assert span["name"] == "invoke_agent test_agent" + elif stream_gen_ai_spans: items = capture_items("transaction", "span") # Run 3 agents concurrently @@ -1320,6 +1580,7 @@ async def run_agent(input_text): assert len(transaction["spans"]) >= 1 +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_message_history( @@ -1327,6 +1588,7 @@ async def test_message_history( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): """ Test that full conversation history is captured in chat spans. @@ -1341,6 +1603,7 @@ async def test_message_history( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Second message with history @@ -1356,7 +1619,7 @@ async def test_message_history( ), ] - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") # First message @@ -1364,12 +1627,8 @@ async def test_message_history( await agent.run("What is my name?", message_history=history) - # We should have 2 transactions - events = [item.payload for item in items if item.type == "transaction"] - - # Check the second transaction has the full history - second_transaction = events[1] - spans = second_transaction["spans"] + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] chat_spans = [ s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" ] @@ -1404,6 +1663,7 @@ async def test_message_history( assert len(messages_data) > 1 +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_gen_ai_system( @@ -1412,6 +1672,7 @@ async def test_gen_ai_system( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that gen_ai.system is set from the model. @@ -1420,15 +1681,17 @@ async def test_gen_ai_system( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await test_agent.run("Test input") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] # Find chat span @@ -1459,6 +1722,7 @@ async def test_gen_ai_system( assert chat_span["data"]["gen_ai.system"] == "test" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_include_prompts_false( @@ -1467,6 +1731,7 @@ async def test_include_prompts_false( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that prompts are not captured when include_prompts=False. @@ -1476,15 +1741,17 @@ async def test_include_prompts_false( traces_sample_rate=1.0, send_default_pii=True, # Even with PII enabled, prompts should not be captured stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await test_agent.run("Sensitive prompt") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] # Find child spans (invoke_agent is the transaction, not a child span) @@ -1513,6 +1780,7 @@ async def test_include_prompts_false( assert "gen_ai.response.text" not in span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_include_prompts_true( @@ -1521,6 +1789,7 @@ async def test_include_prompts_true( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that prompts are captured when include_prompts=True (default). @@ -1530,15 +1799,17 @@ async def test_include_prompts_true( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await test_agent.run("Test prompt") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] # Find child spans (invoke_agent is the transaction, not a child span) @@ -1567,6 +1838,7 @@ async def test_include_prompts_true( assert "gen_ai.request.messages" in chat_span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_include_prompts_false_with_tools( @@ -1575,6 +1847,7 @@ async def test_include_prompts_false_with_tools( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that tool input/output are not captured when include_prompts=False. @@ -1584,6 +1857,7 @@ async def test_include_prompts_false_with_tools( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() @@ -1593,11 +1867,12 @@ def test_tool(value: int) -> int: """A test tool.""" return value * 2 - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await test_agent.run("Use the test tool with value 5") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] # Find tool spans @@ -1628,6 +1903,7 @@ def test_tool(value: int) -> int: assert "gen_ai.tool.output" not in tool_span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_include_prompts_requires_pii( @@ -1636,6 +1912,7 @@ async def test_include_prompts_requires_pii( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that include_prompts requires send_default_pii=True. @@ -1645,15 +1922,17 @@ async def test_include_prompts_requires_pii( traces_sample_rate=1.0, send_default_pii=False, # PII disabled stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await test_agent.run("Test prompt") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] # Find child spans (invoke_agent is the transaction, not a child span) @@ -1840,6 +2119,7 @@ async def run_and_check_context(agent, agent_name): # ==================== Additional Coverage Tests ==================== +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_invoke_agent_with_list_user_prompt( @@ -1847,6 +2127,7 @@ async def test_invoke_agent_with_list_user_prompt( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): """ Test that invoke_agent span handles list user prompts correctly. @@ -1861,15 +2142,38 @@ async def test_invoke_agent_with_list_user_prompt( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + # Use a list as user prompt + await agent.run(["First part", "Second part"]) + + sentry_sdk.flush() + spans = [item.payload for item in items] + + if "gen_ai.request.messages" in spans[0]["attributes"]: + messages_str = spans[0]["attributes"]["gen_ai.request.messages"] + assert "First part" in messages_str + assert "Second part" in messages_str + elif stream_gen_ai_spans: items = capture_items("transaction", "span") # Use a list as user prompt await agent.run(["First part", "Second part"]) (transaction,) = [item.payload for item in items if item.type == "transaction"] + + # 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 else: events = capture_events() @@ -1878,16 +2182,17 @@ async def test_invoke_agent_with_list_user_prompt( (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 + # 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.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio @pytest.mark.parametrize( @@ -1906,6 +2211,7 @@ async def test_invoke_agent_with_instructions( send_default_pii, include_prompts, stream_gen_ai_spans, + span_streaming, ): """ Test that invoke_agent span handles instructions correctly. @@ -1927,13 +2233,15 @@ async def test_invoke_agent_with_instructions( traces_sample_rate=1.0, send_default_pii=send_default_pii, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await agent.run("Test input") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] # The transaction IS the invoke_agent span, check for messages in chat spans instead @@ -2069,6 +2377,7 @@ async def test_model_settings_object_style(sentry_init, capture_items): assert transaction is not None +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_usage_data_partial( @@ -2076,6 +2385,7 @@ async def test_usage_data_partial( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): """ Test that usage data is correctly handled when only some fields are present. @@ -2089,13 +2399,15 @@ async def test_usage_data_partial( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await agent.run("Test input") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] chat_spans = [ @@ -2119,6 +2431,7 @@ async def test_usage_data_partial( assert chat_span is not None +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_agent_data_from_scope( @@ -2126,6 +2439,7 @@ async def test_agent_data_from_scope( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): """ Test that agent data can be retrieved from Sentry scope when not passed directly. @@ -2140,9 +2454,20 @@ async def test_agent_data_from_scope( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + # The integration automatically sets agent in scope during execution + await agent.run("Test input") + + sentry_sdk.flush() + spans = [item.payload for item in items] + + assert spans[1]["name"] == "invoke_agent test_scope_agent" + elif stream_gen_ai_spans: items = capture_items("transaction", "span") # The integration automatically sets agent in scope during execution @@ -2150,6 +2475,9 @@ async def test_agent_data_from_scope( # Verify agent name is capture (transaction,) = (item.payload for item in items if item.type == "transaction") + + # Verify agent name is captured + assert transaction["transaction"] == "invoke_agent test_scope_agent" else: events = capture_events() @@ -2159,10 +2487,11 @@ async def test_agent_data_from_scope( # Verify agent name is capture (transaction,) = events - # Verify agent name is captured - assert transaction["transaction"] == "invoke_agent test_scope_agent" + # Verify agent name is captured + assert transaction["transaction"] == "invoke_agent test_scope_agent" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_available_tools_without_description( @@ -2171,6 +2500,7 @@ async def test_available_tools_without_description( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that available tools are captured even when description is missing. @@ -2179,6 +2509,7 @@ async def test_available_tools_without_description( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() @@ -2188,11 +2519,12 @@ def tool_without_desc(x: int) -> int: # No docstring = no description return x * 2 - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await test_agent.run("Use the tool with 5") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] chat_spans = [ @@ -2219,6 +2551,7 @@ def tool_without_desc(x: int) -> int: assert "tool_without_desc" in tools_str +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_output_with_tool_calls( @@ -2227,6 +2560,7 @@ async def test_output_with_tool_calls( capture_items, get_test_agent, stream_gen_ai_spans, + span_streaming, ): """ Test that tool calls in model response are captured correctly. @@ -2236,6 +2570,7 @@ async def test_output_with_tool_calls( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) test_agent = get_test_agent() @@ -2245,11 +2580,12 @@ def calc_tool(value: int) -> int: """Calculate something.""" return value + 10 - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await test_agent.run("Use calc_tool with 5") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] chat_spans = [ @@ -2284,6 +2620,7 @@ def calc_tool(value: int) -> int: assert "gen_ai.operation.name" in chat_span["data"] +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_message_formatting_with_different_parts( @@ -2291,6 +2628,7 @@ async def test_message_formatting_with_different_parts( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): """ Test that different message part types are handled correctly in ai_client span. @@ -2307,6 +2645,7 @@ async def test_message_formatting_with_different_parts( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) # Create message history with different part types @@ -2320,11 +2659,12 @@ async def test_message_formatting_with_different_parts( ), ] - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await agent.run("What did I say?", message_history=history) + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] chat_spans = [ @@ -2415,6 +2755,7 @@ async def test_update_ai_client_span_with_none_response(sentry_init, capture_ite assert transaction is not None +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_agent_without_name( @@ -2422,6 +2763,7 @@ async def test_agent_without_name( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): """ Test that agent without a name is handled correctly. @@ -2433,15 +2775,28 @@ async def test_agent_without_name( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + await agent.run("Test input") + + sentry_sdk.flush() + spans = [item.payload for item in items] + + assert "invoke_agent" in spans[1]["name"] + elif stream_gen_ai_spans: items = capture_items("transaction", "span") await agent.run("Test input") # Should still create transaction, just with default name (transaction,) = (item.payload for item in items if item.type == "transaction") + + # Transaction name should be "invoke_agent agent" or similar default + assert "invoke_agent" in transaction["transaction"] else: events = capture_events() @@ -2452,8 +2807,8 @@ async def test_agent_without_name( # 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"] + # Transaction name should be "invoke_agent agent" or similar default + assert "invoke_agent" in transaction["transaction"] @pytest.mark.asyncio @@ -2606,6 +2961,7 @@ async def test_set_usage_data_with_partial_fields(sentry_init, capture_items): assert transaction is not None +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_message_parts_with_tool_return( @@ -2613,6 +2969,7 @@ async def test_message_parts_with_tool_return( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): """ Test that ToolReturnPart messages are handled correctly. @@ -2634,14 +2991,16 @@ def test_tool(x: int) -> int: traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") # Run with history containing tool return await agent.run("Use test_tool with 5") + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] chat_spans = [ @@ -3563,6 +3922,7 @@ def _find_binary_content(messages_data, expected_modality, expected_mime_type): return False +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_binary_content_encoding_image( @@ -3570,6 +3930,7 @@ async def test_binary_content_encoding_image( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): """Test that BinaryContent with image data is properly encoded in messages.""" sentry_init( @@ -3577,9 +3938,36 @@ async def test_binary_content_encoding_image( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + with sentry_sdk.traces.start_span( + name="test", attributes={"sentry.op": "test"} + ): + span = sentry_sdk.traces.start_span( + name="custom span", attributes={"sentry.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() + + sentry_sdk.flush() + spans = [item.payload for item in items] + + span_data = spans[0]["attributes"] + messages_data = _get_messages_from_span(span_data) + assert _find_binary_content(messages_data, "image", "image/png") + elif stream_gen_ai_spans: items = capture_items("transaction", "span") with sentry_sdk.start_transaction(op="test", name="test"): @@ -3596,6 +3984,9 @@ async def test_binary_content_encoding_image( span.finish() (event,) = (item.payload for item in items if item.type == "transaction") + span_data = event["spans"][0]["data"] + messages_data = _get_messages_from_span(span_data) + assert _find_binary_content(messages_data, "image", "image/png") else: events = capture_events() @@ -3613,11 +4004,12 @@ async def test_binary_content_encoding_image( 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") + 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.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_binary_content_encoding_mixed_content( @@ -3625,6 +4017,7 @@ async def test_binary_content_encoding_mixed_content( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): """Test that BinaryContent mixed with text content is properly handled.""" sentry_init( @@ -3632,9 +4025,47 @@ async def test_binary_content_encoding_mixed_content( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + with sentry_sdk.traces.start_span( + name="test", attributes={"sentry.op": "test"} + ): + span = sentry_sdk.traces.start_span( + name="custom span", attributes={"sentry.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() + + sentry_sdk.flush() + spans = [item.payload for item in items] + + span_data = spans[0]["attributes"] + 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") + elif stream_gen_ai_spans: items = capture_items("transaction", "span") with sentry_sdk.start_transaction(op="test", name="test"): @@ -3653,6 +4084,18 @@ async def test_binary_content_encoding_mixed_content( span.finish() (event,) = (item.payload for item in items if item.type == "transaction") + 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") else: events = capture_events() @@ -3672,21 +4115,21 @@ async def test_binary_content_encoding_mixed_content( 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") + 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.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_binary_content_in_agent_run( @@ -3694,6 +4137,7 @@ async def test_binary_content_in_agent_run( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): """Test that BinaryContent in actual agent run is properly captured in spans.""" agent = Agent("test", name="test_binary_agent") @@ -3703,17 +4147,19 @@ async def test_binary_content_in_agent_run( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) binary_content = BinaryContent( data=b"fake_image_data_for_testing", media_type="image/png" ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await agent.run(["Analyze this image:", binary_content]) + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] chat_spans = [ s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" @@ -3744,6 +4190,7 @@ async def test_binary_content_in_agent_run( ) +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_set_usage_data_with_cache_tokens( @@ -3751,15 +4198,42 @@ async def test_set_usage_data_with_cache_tokens( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): """Test that cache_read_tokens and cache_write_tokens are tracked.""" sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming: + items = capture_items("span") + + with sentry_sdk.traces.start_span( + name="test", attributes={"sentry.op": "test"} + ): + span = sentry_sdk.traces.start_span( + name="custom span", attributes={"sentry.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() + + sentry_sdk.flush() + spans = [item.payload for item in items] + + assert spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED] == 80 + assert ( + spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE] == 20 + ) + elif stream_gen_ai_spans: items = capture_items("transaction", "span") with sentry_sdk.start_transaction(op="test", name="test"): @@ -3774,6 +4248,9 @@ async def test_set_usage_data_with_cache_tokens( span.finish() (event,) = (item.payload for item in items if item.type == "transaction") + (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 else: events = capture_events() @@ -3789,12 +4266,12 @@ async def test_set_usage_data_with_cache_tokens( 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 + (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.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.parametrize( "url,image_url_kwargs,expected_content", @@ -3845,6 +4322,7 @@ def test_image_url_base64_content_in_span( image_url_kwargs, expected_content, stream_gen_ai_spans, + span_streaming, ): from sentry_sdk.integrations.pydantic_ai.spans.ai_client import ai_client_span @@ -3853,10 +4331,11 @@ def test_image_url_base64_content_in_span( traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) found_image = False - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") with sentry_sdk.start_transaction(op="test", name="test"): @@ -3869,6 +4348,7 @@ def test_image_url_base64_content_in_span( span = ai_client_span([mock_msg], None, None, None) span.finish() + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] chat_spans = [ s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" @@ -3912,6 +4392,7 @@ def test_image_url_base64_content_in_span( assert found_image, "Image content item should be found in messages data" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio @pytest.mark.parametrize( @@ -3951,25 +4432,28 @@ async def test_invoke_agent_image_url( image_url_kwargs, expected_content, stream_gen_ai_spans, + span_streaming, ): sentry_init( integrations=[PydanticAIIntegration()], traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) agent = Agent("test", name="test_image_url_agent") image_url = ImageUrl(url=url, **image_url_kwargs) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") await agent.run([image_url, "Describe this image"]) found_image = False + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] chat_spans = [ s for s in spans if s["attributes"].get("sentry.op", "") == "gen_ai.chat" @@ -4006,6 +4490,7 @@ async def test_invoke_agent_image_url( assert found_image, "Image content item should be found in messages data" +@pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio async def test_tool_description_in_execute_tool_span( @@ -4013,6 +4498,7 @@ async def test_tool_description_in_execute_tool_span( capture_events, capture_items, stream_gen_ai_spans, + span_streaming, ): """ Test that tool description from the tool's docstring is included in execute_tool spans. @@ -4033,14 +4519,16 @@ def multiply_numbers(a: int, b: int) -> int: traces_sample_rate=1.0, send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - if stream_gen_ai_spans: + if span_streaming or stream_gen_ai_spans: items = capture_items("transaction", "span") result = await agent.run("What is 5 times 3?") assert result is not None + sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] tool_spans = [