fix: reshuffling where the _invoke_agent_span_stack context variable is set

Author: constantinius

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -43,13 +43,14 @@ async def __aenter__(self):
 
         # Create invoke_agent span (will be closed in __aexit__)
         self._span = invoke_agent_span(
-            self.user_prompt, self.agent, self.model, self.model_settings
+            self.user_prompt,
+            self.agent,
+            self.model,
+            self.model_settings,
+            self.is_streaming,
         )
         self._span.__enter__()
 
-        # Push span and agent to contextvar stack
-        push_invoke_agent_span(self._span, self.agent, self.is_streaming)
-
         # Enter the original context manager
         result = await self.original_ctx_manager.__aenter__()
         self._result = result
@@ -103,9 +104,9 @@ async def wrapper(self, *args, **kwargs):
             model_settings = kwargs.get("model_settings")
 
             # 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, is_streaming)
+            with invoke_agent_span(
+                user_prompt, self, model, model_settings, is_streaming
+            ) as span:
                 try:
                     result = await original_func(self, *args, **kwargs)
 

sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py

@@ -8,6 +8,7 @@
     _set_available_tools,
     _set_model_data,
     _should_send_prompts,
+    push_invoke_agent_span,
 )
 
 from typing import TYPE_CHECKING
@@ -16,8 +17,8 @@
     from typing import Any
 
 
-def invoke_agent_span(user_prompt, agent, model, model_settings):
-    # type: (Any, Any, Any, Any) -> sentry_sdk.tracing.Span
+def invoke_agent_span(user_prompt, agent, model, model_settings, is_streaming=False):
+    # type: (Any, Any, Any, Any, bool) -> sentry_sdk.tracing.Span
     """Create a span for invoking the agent."""
     # Determine agent name for span
     name = "agent"
@@ -32,6 +33,10 @@ def invoke_agent_span(user_prompt, agent, model, model_settings):
 
     span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
 
+    # Push span and agent to contextvar stack immediately after span creation
+    # This ensures the agent is available in get_current_agent() before _set_model_data is called
+    push_invoke_agent_span(span, agent, is_streaming)
+
     _set_agent_data(span, agent)
     _set_model_data(span, model, model_settings)
     _set_available_tools(span, agent)