fix: correct the propagation of the "is_streaming" flag

Author: constantinius

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -48,7 +48,7 @@ async def __aenter__(self):
         self._span.__enter__()
 
         # Push span and agent to contextvar stack
-        push_invoke_agent_span(self._span, self.agent)
+        push_invoke_agent_span(self._span, self.agent, self.is_streaming)
 
         # Enter the original context manager
         result = await self.original_ctx_manager.__aenter__()
@@ -105,7 +105,7 @@ async def wrapper(self, *args, **kwargs):
             # Create invoke_agent span
             with invoke_agent_span(user_prompt, self, model, model_settings) as span:
                 # Push span and agent to contextvar stack
-                push_invoke_agent_span(span, self)
+                push_invoke_agent_span(span, self, is_streaming)
                 try:
                     result = await original_func(self, *args, **kwargs)
 

sentry_sdk/integrations/pydantic_ai/spans/ai_client.py

@@ -10,6 +10,8 @@
     _set_model_data,
     _should_send_prompts,
     _get_model_name,
+    get_current_agent,
+    get_is_streaming,
 )
 
 from typing import TYPE_CHECKING
@@ -216,20 +218,11 @@ def ai_client_span(messages, agent, model, model_settings):
     _set_agent_data(span, agent)
     _set_model_data(span, model, model_settings)
 
-    # Set streaming flag
-    agent_data = sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {}
-    is_streaming = agent_data.get("_streaming", False)
-    span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, is_streaming)
+    # 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
-    if not agent_obj:
-        # Try to get from Sentry scope
-        agent_data = (
-            sentry_sdk.get_current_scope()._contexts.get("pydantic_ai_agent") or {}
-        )
-        agent_obj = agent_data.get("_agent")
-
+    agent_obj = agent or get_current_agent()
     _set_available_tools(span, agent_obj)
 
     # Set input messages (full conversation history)

sentry_sdk/integrations/pydantic_ai/utils.py

@@ -17,11 +17,11 @@
 # type: ContextVar[list[dict[str, Any]]]
 
 
-def push_invoke_agent_span(span, agent):
-    # type: (sentry_sdk.tracing.Span, Any) -> None
-    """Push an invoke_agent span onto the stack along with its agent."""
+def push_invoke_agent_span(span, agent, is_streaming=False):
+    # type: (sentry_sdk.tracing.Span, Any, bool) -> None
+    """Push an invoke_agent span onto the stack along with its agent and streaming flag."""
     stack = _invoke_agent_span_stack.get().copy()
-    stack.append({"span": span, "agent": agent})
+    stack.append({"span": span, "agent": agent, "is_streaming": is_streaming})
     _invoke_agent_span_stack.set(stack)
 
 
@@ -52,6 +52,15 @@ def get_current_agent():
     return None
 
 
+def get_is_streaming():
+    # type: () -> bool
+    """Get the streaming flag from the contextvar stack."""
+    stack = _invoke_agent_span_stack.get()
+    if stack:
+        return stack[-1].get("is_streaming", False)
+    return False
+
+
 def _should_send_prompts():
     # type: () -> bool
     """