fix: using contextvar to store agent name in a re-entrant way

Author: constantinius

sentry_sdk/integrations/langchain.py

@@ -1,3 +1,4 @@
+import contextvars
 import itertools
 from collections import OrderedDict
 from functools import wraps
@@ -72,6 +73,40 @@
 }
 
 
+# Contextvar to track agent names in a stack for re-entrant agent support
+_agent_stack = contextvars.ContextVar("langchain_agent_stack", default=None)
+
+
+def _push_agent(agent_name):
+    # type: (Optional[str]) -> None
+    """Push an agent name onto the stack."""
+    stack = _agent_stack.get()
+    if stack is None:
+        stack = []
+    stack.append(agent_name)
+    _agent_stack.set(stack)
+
+
+def _pop_agent():
+    # type: () -> Optional[str]
+    """Pop an agent name from the stack and return it."""
+    stack = _agent_stack.get()
+    if stack and len(stack) > 0:
+        agent_name = stack.pop()
+        _agent_stack.set(stack)
+        return agent_name
+    return None
+
+
+def _get_current_agent():
+    # type: () -> Optional[str]
+    """Get the current agent name (top of stack) without removing it."""
+    stack = _agent_stack.get()
+    if stack and len(stack) > 0:
+        return stack[-1]
+    return None
+
+
 class LangchainIntegration(Integration):
     identifier = "langchain"
     origin = f"auto.ai.{identifier}"
@@ -276,11 +311,7 @@ def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
             elif "openai" in ai_type:
                 span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai")
 
-            agent_name = (
-                sentry_sdk.get_current_scope()
-                ._contexts.get("langchain_agent", {})
-                .get("agent_name")
-            )
+            agent_name = _get_current_agent()
             if agent_name:
                 span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
 
@@ -436,11 +467,7 @@ def on_tool_start(self, serialized, input_str, *, run_id, **kwargs):
             if tool_description is not None:
                 span.set_data(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool_description)
 
-            agent_name = (
-                sentry_sdk.get_current_scope()
-                ._contexts.get("langchain_agent", {})
-                .get("agent_name")
-            )
+            agent_name = _get_current_agent()
             if agent_name:
                 span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
 
@@ -772,9 +799,7 @@ def new_invoke(self, *args, **kwargs):
             name=f"invoke_agent {agent_name}" if agent_name else "invoke_agent",
             origin=LangchainIntegration.origin,
         ) as span:
-            sentry_sdk.get_current_scope().set_context(
-                "langchain_agent", {"agent_name": agent_name, "tools": tools}
-            )
+            _push_agent(agent_name)
             if agent_name:
                 span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
 
@@ -813,7 +838,7 @@ def new_invoke(self, *args, **kwargs):
             ):
                 set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
 
-            sentry_sdk.get_current_scope().remove_context("langchain_agent")
+            _pop_agent()
 
             return result
 
@@ -840,9 +865,7 @@ def new_stream(self, *args, **kwargs):
         )
         span.__enter__()
 
-        sentry_sdk.get_current_scope().set_context(
-            "langchain_agent", {"agent_name": agent_name, "tools": tools}
-        )
+        _push_agent(agent_name)
 
         if agent_name:
             span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
@@ -893,7 +916,7 @@ def new_iterator():
             ):
                 set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
 
-            sentry_sdk.get_current_scope().remove_context("langchain_agent")
+            _pop_agent()
 
             span.__exit__(None, None, None)
 
@@ -914,7 +937,7 @@ async def new_iterator_async():
             ):
                 set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
 
-            sentry_sdk.get_current_scope().remove_context("langchain_agent")
+            _pop_agent()
 
             span.__exit__(None, None, None)