fix(pydantic-ai): Determine if response is streamed based on the method

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -57,7 +57,7 @@ async def __aenter__(self) -> "Any":
 
         # Push agent to contextvar stack after span is successfully created and entered
         # This ensures proper pairing with pop_agent() in __aexit__ even if exceptions occur
-        push_agent(self.agent, self.is_streaming)
+        push_agent(self.agent)
 
         # Enter the original context manager
         result = await self.original_ctx_manager.__aenter__()
@@ -120,7 +120,7 @@ async def wrapper(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
             ) as span:
                 # Push agent to contextvar stack after span is successfully created and entered
                 # This ensures proper pairing with pop_agent() in finally even if exceptions occur
-                push_agent(self, is_streaming)
+                push_agent(self)
 
                 try:
                     result = await original_func(self, *args, **kwargs)

sentry_sdk/integrations/pydantic_ai/patches/graph_nodes.py

@@ -1,7 +1,9 @@
 from contextlib import asynccontextmanager
 from functools import wraps
 
+from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.traces import StreamedSpan
 
 from ..spans import (
     ai_client_span,
@@ -59,9 +61,19 @@
 
     @wraps(original_model_request_run)
     async def wrapped_model_request_run(self: "Any", ctx: "Any") -> "Any":
+        did_stream = getattr(self, "_did_stream", None)
+        cached_result = getattr(self, "_result", None)
+        if did_stream or cached_result is not None:
+            return await original_model_request_run(self, ctx)
+
         messages, model, model_settings = _extract_span_data(self, ctx)
 
         with ai_client_span(messages, None, model, model_settings) as span:
+            if isinstance(span, StreamedSpan):
+                span.set_attribute(SPANDATA.GEN_AI_RESPONSE_STREAMING, False)
+            else:
+                span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, False)
+
             result = await original_model_request_run(self, ctx)
 
             # Extract response from result if available
@@ -85,10 +97,20 @@
         @asynccontextmanager
         @wraps(original_stream_method)
         async def wrapped_model_request_stream(self: "Any", ctx: "Any") -> "Any":
+            did_stream = getattr(self, "_did_stream", None)
+            if did_stream:
+                async with original_stream_method(self, ctx) as stream:
+                    yield stream
+
             messages, model, model_settings = _extract_span_data(self, ctx)
 
             # Create chat span for streaming request
             with ai_client_span(messages, None, model, model_settings) as span:
+                if isinstance(span, StreamedSpan):
+                    span.set_attribute(SPANDATA.GEN_AI_RESPONSE_STREAMING, False)
+                else:
+                    span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, False)
+
                 # Call the original stream method
                 async with original_stream_method(self, ctx) as stream:
                     yield stream

sentry_sdk/integrations/pydantic_ai/spans/ai_client.py

@@ -23,7 +23,6 @@
     _set_model_data,
     _should_send_prompts,
     get_current_agent,
-    get_is_streaming,
 )
 from .utils import (
     _serialize_binary_content_item,
@@ -314,7 +313,6 @@ def ai_client_span(
                 "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:
@@ -325,8 +323,6 @@ def ai_client_span(
         )
 
         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)

sentry_sdk/integrations/pydantic_ai/utils.py

@@ -18,10 +18,10 @@
 )
 
 
-def push_agent(agent: "Any", is_streaming: bool = False) -> None:
-    """Push an agent context onto the stack along with its streaming flag."""
+def push_agent(agent: "Any") -> None:
+    """Push an agent context onto the stack."""
     stack = _agent_context_stack.get().copy()
-    stack.append({"agent": agent, "is_streaming": is_streaming})
+    stack.append(agent)
     _agent_context_stack.set(stack)
 
 
@@ -37,18 +37,10 @@ def get_current_agent() -> "Any":
     """Get the current agent from the contextvar stack."""
     stack = _agent_context_stack.get()
     if stack:
-        return stack[-1]["agent"]
+        return stack[-1]
     return None
 
 
-def get_is_streaming() -> bool:
-    """Get the streaming flag from the contextvar stack."""
-    stack = _agent_context_stack.get()
-    if stack:
-        return stack[-1].get("is_streaming", False)
-    return False
-
-
 def _should_send_prompts() -> bool:
     """
     Check if prompts should be sent to Sentry.