simplify context stack

Author: alexander-alderman-webb

sentry_sdk/integrations/pydantic_ai/patches/agent_run.py

@@ -58,7 +58,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__()
@@ -113,7 +113,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/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,7 +37,7 @@ 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