fix: update agent context management in invoke_agent spans

Author: constantinius

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -3,7 +3,7 @@
 import sentry_sdk
 
 from ..spans import invoke_agent_span, update_invoke_agent_span
-from ..utils import _capture_exception, pop_agent
+from ..utils import _capture_exception, pop_agent, push_agent
 
 from typing import TYPE_CHECKING
 from pydantic_ai.agent import Agent  # type: ignore
@@ -51,6 +51,10 @@ async def __aenter__(self):
         )
         self._span.__enter__()
 
+        # 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)
+
         # Enter the original context manager
         result = await self.original_ctx_manager.__aenter__()
         self._result = result
@@ -107,6 +111,10 @@ async def wrapper(self, *args, **kwargs):
             with invoke_agent_span(
                 user_prompt, self, model, model_settings, is_streaming
             ) 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)
+
                 try:
                     result = await original_func(self, *args, **kwargs)
 

sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py

@@ -8,7 +8,6 @@
     _set_available_tools,
     _set_model_data,
     _should_send_prompts,
-    push_agent,
 )
 
 from typing import TYPE_CHECKING
@@ -33,10 +32,6 @@ def invoke_agent_span(user_prompt, agent, model, model_settings, is_streaming=Fa
 
     span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
 
-    # Push 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_agent(agent, is_streaming)
-
     _set_agent_data(span, agent)
     _set_model_data(span, model, model_settings)
     _set_available_tools(span, agent)